一架梯子,一头程序猿,仰望星空!
LangChain教程(Python版本) > 内容正文

LangChain 自定义模型


自定义 LLM

目前AI模型领域百家争鸣,LangChain官方也没有对接好所有模型,有时候你需要自定义模型,接入LangChain框架。

本章将介绍如何创建自定义 LLM 包装器,方便你使用自己的模型或者LangChain官方不支持的模型。

自定义 LLM 需要实现以下必要的内容:

  1. _call 方法,它需要接受一个字符串、一些可选的停用词,并返回一个字符串。

它还可以实现第二个可选的内容:

  1. _identifying_params 属性,用于帮助打印此类。应该返回一个字典。

我们来实现一个非常简单的自定义 LLM,它只是返回输入的前 N 个字符。

from typing import Any, List, Mapping, Optional

from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM

API 参考:

# 自定义模型类,实现基础的LLM接口
class CustomLLM(LLM):
    n: int

    @property
    def _llm_type(self) -> str:
                # 模型名称
        return "custom"
    # 在这里实现模型调用
    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        run_manager: Optional[CallbackManagerForLLMRun] = None,
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")
        return prompt[: self.n]

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """返回模型信息,类似其他java的tostring函数,返回对象信息"""
        return {"n": self.n}

我们现在可以像使用任何其他 LLM 一样使用它。

# 实例化模型
llm = CustomLLM(n=10)
# 调用模型
llm("This is a foobar thing")
# 模型返回结果
'This is a '

我们还可以打印 LLM 并查看其自定义打印。

print(llm)
CustomLLM    Params: {'n': 10}


关联主题