当前位置：首页 > news >正文

创建自定义LLM类：更灵活地使用你的语言模型

news 2024/9/24 15:19:59

引言

在使用LangChain时，可能会希望集成自定义的语言模型（LLM）或使用不同的封装器。本文将指导你如何创建一个自定义LLM包装器，使你的模型与LangChain程序无缝集成。通过实现标准的LLM接口，你的模型能够自动支持LangChain的优化、异步支持、流式API等功能。

主要内容

1. 实现基础方法

为了打造一个自定义的LLM类，你需要实现以下两个基本方法：

_call: 处理输入字符串并返回输出字符串。
_llm_type: 返回一个字符串，仅用于日志记录。

2. 可选方法

可以根据需求实现以下可选方法：

_identifying_params: 返回用于标识模型的参数字典。
_acall: 异步实现_call。
_stream: 支持令牌逐个输出。
_astream: 异步实现_stream。

3. 实现示例

我们将实现一个简单的自定义LLM模型，返回输入的前n个字符。

from typing import Any, Dict, Iterator, List, Optionalfrom langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunkclass CustomLLM(LLM):"""A custom chat model that echoes the first `n` characters of the input."""n: int  # Echoes the first n charactersdef _call(self, prompt: str, stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> str:if stop is not None:raise ValueError("stop kwargs are not permitted.")return prompt[: self.n]def _stream(self, prompt: str, stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> Iterator[GenerationChunk]:for char in prompt[: self.n]:chunk = GenerationChunk(text=char)if run_manager:run_manager.on_llm_new_token(chunk.text, chunk=chunk)yield chunk@propertydef _identifying_params(self) -> Dict[str, Any]:return {"model_name": "CustomChatModel"}@propertydef _llm_type(self) -> str:return "custom"

代码示例

以下展示了如何测试自定义LLM类：

llm = CustomLLM(n=5)
print(llm.invoke("This is a foobar thing"))  # Output: 'This '# 异步调用示例
import asyncioasync def async_example():print(await llm.ainvoke("world"))  # Output: 'world'# 流式输出示例async for token in llm.astream("hello"):print(token, end="|")asyncio.run(async_example())