大语言模型RAG-将本地大模型封装为langchain的chat model（三）

往期文章：
大语言模型RAG-技术概览 (一)
大语言模型RAG-langchain models (二)

上一期langchain还在0.1时代，这期使用的langchain v0.2已经与之前不兼容了。
本期介绍如何将开源模型封装为langchain的chat模型。 这是后面构建RAG应用的基础。

基于langchain_core.language_models.BaseChatModel创建类。这个类就是本地llm的langchain封装，需要定义以下方法或属性。

方法/属性	作用	是否必要
`_generate`	根据prompt生成聊天结果	必要
`_llm_type` (属性)	模型的唯一标识，用于生成日志	必要
`_identifying_params`（属性）	用于返回模型的一些参数，用于日志和调试	可选
`_stream`	用于流式输出	可选
`_agenerate`	本机异步的`_generate`	可选
`_astream`	异步的_stream	可选

示例如下，只实现比较常用的功能：

from typing import Any, AsyncIterator, Dict, Iterator, List, Optional
from threading import Thread
from langchain_core.callbacks import (AsyncCallbackManagerForLLMRun,CallbackManagerForLLMRun,
)
from langchain_core.language_models import BaseChatModel, SimpleChatModel
from langchain_core.messages import AIMessageChunk, BaseMessage, HumanMessage
from langchain_core.outputs import (
ChatGeneration, 
ChatGenerationChunk, 
ChatResult
)
from langchain_core.messages import (AIMessage,BaseMessage,HumanMessage,
)
from typing import AsyncIterator, Literal, Iterator, Optional, List, Dict, Any
from transformers import AutoTokenizer, AutoModelForCausalLM
from pydantic import BaseModel, Field, validator
import torchclass CustomChatModel(BaseChatModel):"""将本地模型封装为langchain的chat模型"""model_name: str  # 模型的本地地址device: str = "cpu"  # 计算设备n: int = 512  # 最大生成长度temperature: float = 1.0  # 温度model: Optional[AutoModelForCausalLM] = Nonetokenizer: Optional[AutoTokenizer] = Nonedef __init__(self, model_name: str, device: str, temperature: float = 1.0, **kwargs):"""我把模型加载放在初始化里面了。当然有更灵活的处理方式，这部分你可以随心所欲，只要在模型获得输入前加载好模型就OK。Params:model_name (str): 本地模型的路径device (str): 计算设备temperature (float): 模型温度"""super().__init__(model_name=model_name, device=device, **kwargs)self.model_name = model_nameself.device = deviceself.temperature = temperature# 加载本地模型。try:self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True, use_fast=False)self.model = AutoModelForCausalLM.from_pretrained(self.model_name, trust_remote_code=True).half().to(self.device)except Exception as e:raise RuntimeError(f"模型或分词器加载失败: {e}")def _generate(self, messages: List[BaseMessage], stop: Optional[List[str]] = ['.', '。', '?', '？', '!', '！'], run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> ChatResult:"""实现模型的文本生成逻辑这是langchain必须的类方法。Params:messages (BaseMessage): 输入模型的消息列表stop (List[str])： 这是一个字符列表，langchain会根据这个列表中的字符找到完整句子的结束位置。run_manager (CallbackManagerForLLMRun): 用于管理和追踪执行的过程，通常用于实现回调和异步处理。虽然这个功能很重要，但本文暂时不做示例（我没玩明白）Return：ChatResult类的输出，完全符合langchain风格。"""last_message = messages[-1].content  # 获取最后一条消息inputs = self.tokenizer.encode(last_message, return_tensors="pt").to(self.device)  # 将消息内容编码为张量outputs = self.model.generate(inputs, max_length=self.n + len(inputs[0]), temperature=self.temperature)  # 调用模型生成文本张量tokens = self.tokenizer.decode(outputs[0], skip_special_tokens=True)  # 将张量解码为文本# 截取合适的长度（找到完整句子的结束位置）end_positions = [tokens.find(c) for c in stop if tokens.find(c) != -1]if end_positions:end_pos = max(end_positions) + 1  # 包括结束符if end_pos < self.n:tokens = tokens[:end_pos]else:tokens = tokens[:self.n]else:tokens = tokens[:self.n]message = AIMessage(content=tokens)  # 将生成的文本封装为消息generation = ChatGeneration(message=message)  # 封装为ChatGenerationreturn ChatResult(generations=[generation])def _stream(self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any) -> Iterator[ChatGenerationChunk]:"""流式生成文本。这个方法不是必须的，但实际应用中流式输出能提高用户体验，从这个角度来说它是必须的。下面仅给出了最基本的流式输出的逻辑，实际工程中你可以自行添加你需要的逻辑。Params:messages (BaseMessage): 输入模型的消息列表stop (List[str])： 这是一个字符列表，langchain会根据这个列表中的字符找到完整句子的结束位置。run_manager (CallbackManagerForLLMRun): 用于管理和追踪执行的过程，通常用于实现回调和异步处理。Return:ChatGenerationChunk类的输出，完全符合langchain风格。"""last_message = messages[-1].contentinputs = self.tokenizer(last_message, return_tensors="pt").to(self.device)streamer = TextIteratorStreamer(tokenizer=self.tokenizer)  # transformers提供的标准接口，专为流式输出而生inputs.update({"streamer": streamer, "max_new_tokens": 512})  # 这是model.generate的参数，可以自由发挥thread = Thread(target=self.model.generate, kwargs=inputs)  # TextIteratorStreamer需要配合线程使用thread.start()for new_token in streamer:# 用迭代器的返回形式 。yield ChatGenerationChunk(message=AIMessage(content=new_token))@propertydef _llm_type(self) -> str:"""返回自定义模型的标记"""return "Custom"@propertydef _identifying_params(self) -> Dict[str, Any]:"""返回自定义的debug信息"""return {"model_path": self.model_name, "device": self.device}

本文使用chatglm3-6b模型为例，演示最基础使用方法。

import CustomChatModel
import os
from langchain_core.messages import HumanMessage# 模型实例化
model = LoadChatLLM(model_name="THUDM/chatglm-6b", device=“cuda”, temperature=float(1.0))def stream_generate_text(text):"""流式输出的方式打印到终端。原理就是每生成一个新token，就清空屏幕，然后把原先生成过的所有tokens都打印一次。"""generated_text = ''count = 0message = [HumanMessage(content=text)]for new_text in model._stream(message):generated_text += new_text.textcount += 1if count % 8 == 0:  # 避免刷新太频繁，每8个tokens刷新一次os.system("clear")print(generated_text)os.system("clear")print(generated_text)while True:ipt = input("请输入：")stream_generate_text(ipt)