ChatGLM3-6B和LangChain构建知识库助手

安装依赖库

使用pip命令安装以下库：

pip install modelscope langchain==0.1.7 chromadb==0.5.0 sentence-transformers==2.7.0 unstructured==0.13.7 markdown==3.0.0 docx2txt==0.8  pypdf==4.2.0

依赖库简介：

ModelScope：一个用于机器学习模型管理和部署的库LangChain：一个用于构建语言模型应用的框架ChromaDB：一个用于高效存储和检索嵌入向量的数据库，支持相似性搜索Sentence Transformers：一个用于生成句子和文本嵌入的库，基于Transformer模型。Unstructured：一个用于处理非结构化数据的库，提供了多种工具来提取和转换数据Markdown：一个用于将Markdown文本转换为HTML的库docx2txt：一个用于从DOCX文件中提取文本的库PyPDF：一个用于处理PDF文件的库，支持读取、写入和修改PDF文档

下载大模型

from modelscope import snapshot_download# 下载指定模型
model_dir = snapshot_download('ZhipuAI/chatglm3-6b',  # 模型的名称cache_dir='/root/models',  # 缓存目录，用于存储下载的模型文件revision='master'  # 指定模型的版本，通常为'master'表示最新版本
)

下载向量模型

from modelscope import snapshot_download# 下载指定模型
model_dir = snapshot_download('AI-ModelScope/bge-large-zh',  # 模型的名称cache_dir='/root/models'  # 缓存目录，用于存储下载的模型文件
)

自定义LLM 类

在本地部署的ChatGLM3-6B基础上，构建LLM应用需要自定义一个LLM类，并将ChatGLM接入到LangChain框架中。通过自定义LLM类，可以实现与LangChain接口的完全一致调用方式，无需担心底层模型调用的不一致性。

创建ChatGLM_LLM.py文件，需从LangChain.llms.base.LLM类继承一个子类，并重写构造函数与_call函数

from typing import Any, List, Optionalimport torch
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from transformers import AutoTokenizer, AutoModelForCausalLM# 自定义GLM类
class ChatGLMLLM(LLM):"""自定义 LLM 类，用于加载和使用 ChatGLM 模型"""tokenizer: AutoTokenizer = Nonemodel: AutoModelForCausalLM = Nonedef __init__(self, model_path: str):""":param model_path: 从本地初始化模型"""super().__init__()print("-------------开始加载LLM模型-------------")# 加载分词器self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)# 加载模型self.model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True).to(torch.bfloat16).cuda()# 设置模型为评估模式self.model = self.model.eval()print("-------------模型LLM加载完毕-------------")# 定义_call方法：进行模型的推理def _call(self, prompt: str, stop: Optional[List[str]] = None,run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any):""" 重写调用函数 在给定的提示和输入下运行 LLM """response, history = self.model.chat(self.tokenizer, prompt, history=[])return response@propertydef _llm_type(self) -> str:""" 返回模型类型 """return "ChatGLM3-6B"

文件读取与处理

创建File_Proces.py文件，递归指定文件夹路径，返回所有满足条件的文件路径，并使用LangChain提供的各类Loader对象加载文件，解析出纯文本内容

import os
from typing import Listfrom langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader, PyPDFLoader, Docx2txtLoader
from tqdm import tqdmclass FileProces:def get_files(self, dir_path):"""根据路径，递归获取路径下的文件列表"""# 指定要筛选的文件扩展名列表extensions = [".pdf", ".docx", ".md", "txt"]file_list = []"""os.walk 函数将递归遍历指定文件夹filepath: 当前遍历到的文件路径，包含文件名。dirnames: 一个包含当前目录中所有子目录名的列表。filenames: 一个包含当前目录中所有文件名的列表"""for filepath, dirnames, filenames in os.walk(dir_path):for filename in filenames:# 检查文件后缀是否在指定的扩展名列表中if any(filename.endswith(ext) for ext in extensions):# 将符合条件的文件的绝对路径添加到列表中file_list.append(os.path.join(filepath, filename))return file_listdef get_text(self, file_lst):"""根据传入文件列表，分别获取每个文件对象的文本信息"""# docs 存放加载之后的纯文本对象docs = []# 定义文件类型对应的加载器字典loader_mapping = {'pdf': PyPDFLoader,'docx': Docx2txtLoader,'md': UnstructuredMarkdownLoader,'txt': TextLoader}# tqdm是一个 Python 库，用于在循环中展示进度条，帮助用户实时了解循环迭代的进度for file in tqdm(file_lst):# 文件名按照.分割成多个部分，选择列表中的最后一个元素file_type = file.split('.')[-1]# 根据文件类型选择对应的加载器loader = loader_mapping.get(file_type)if loader:loader_instance = loader(file)doc = loader_instance.load()docs.extend(doc)else:# 如果文件类型不在loader_mapping中，跳过当前文件continuereturn docsdef getALLDocs(self, tar_dir: List):"""传入目标目录，返回目录下所有文档列表"""# 处理文件all_docs = []for dir_path in tar_dir:# 得到目标文件路径列表file_list = self.get_files(dir_path)print("文档列表：", file_list)docs = self.get_text(file_list)# 合并所有文档all_docs.extend(docs)return all_docs

构建向量数据库

创建Embedding_LLM.py文件，加载词向量模型，把经过文件的读取与处理后得到的纯文本对象列表引入LangChain框架，对文本进行分块，进行向量化处理，然后构建向量数据库。

# 首先导入所需第三方库from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceBgeEmbeddingsfrom File_Proces import FileProcesclass EmbeddingLLM():embeddings: HuggingFaceBgeEmbeddings = Nonevectordb = None# 定义持久化路径persist_directory = './vectorDB/chroma'def __init__(self, model_path: str):""":param model_path: 加载词向量模型"""print("-------------开始加载词向量模型-------------")model_kwargs = {"device": "cuda"}encode_kwargs = {"normalize_embeddings": True}# 加载开源词向量模型self.embeddings = HuggingFaceBgeEmbeddings(model_name=model_path, model_kwargs=model_kwargs,encode_kwargs=encode_kwargs)print("-------------词向量模型加载完毕-------------")# 初始化向量数据库self.vectordb = Chroma(persist_directory=self.persist_directory,embedding_function=self.embeddings)def docsToEmbedding(self, docs):"""文档转向量，同时持久化到本地目录"""# 对文本进行分块text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)split_docs = text_splitter.split_documents(docs)# 保存向量数据vectordb = Chroma.from_documents(documents=split_docs,embedding=self.embeddings,persist_directory=self.persist_directory  # 持久化的目录)# 将加载的向量数据库持久化到磁盘上vectordb.persist()print("向量数据持久化完成")if __name__ == '__main__':# 初始化向量内嵌模型llm = EmbeddingLLM(model_path="./models/bge-large-zh")# 目标文件目录tar_dir = ["./data"]# 获取指定目录下的所有文档内容fileproces = FileProces()all_docs = fileproces.getALLDocs(tar_dir=tar_dir)# 开始向量化处理llm.docsToEmbedding(docs=all_docs)

执行日志如下：

-------------开始加载词向量模型-------------
-------------词向量模型加载完毕-------------
文档列表： ['./data/《Scrum指南》中文版.pdf', './data/README.md', './data/demo.txt']
100%|██████████| 3/3 [00:03<00:00,  1.17s/it]
向量数据持久化完成

本地文件存储情况如下：

构建检索问答链

LangChain通过提供RetrievalQA对象实现了对RAG全流程的封装。通过初始化一个RetrievalQA对象，并填入已构建的数据库和自定义LLM作为参数，就能方便地完成检索增强问答的全流程。

LangChain会自动根据用户提问进行检索、获取相关文档、生成适当的提示，并将其传递给LLM进行问答的整个过程。

创建Retrieval.py文件，创建一个构建检索问答链函数

from langchain.chains import RetrievalQA
from langchain_core.prompts import PromptTemplatedef builder_chain(llm, embeddingLlm):# 构造Prompt模板template = """使用以下上下文来回答问题。如果你不知道答案，你需要澄清，不要试图编造答案。请简明扼要的回答。”。{context}问题: {question}回答:"""# 实例化Template对象，context和question两个变量会被检索到的文档片段和用户提问填充qa_chain_prompt = PromptTemplate(input_variables=["context", "question"], template=template)# 构建检索链vectordb = embeddingLlm.vectordbqa_chain = RetrievalQA.from_chain_type(llm, retriever=vectordb.as_retriever(), return_source_documents=True,chain_type_kwargs={"prompt": qa_chain_prompt})return qa_chain

测试

from ChatGLM_LLM import ChatGLMLLM
from Embedding_LLM import EmbeddingLLM
from Retrieval import builder_chainif __name__ == '__main__':# 初始化大模型llm = ChatGLMLLM(model_path="./models/chatglm3-6b")# 初始化向量内嵌模型embeddingLlm = EmbeddingLLM(model_path="./models/bge-large-zh")qa_chain = builder_chain(llm, embeddingLlm)# 检索问答链回答效果question = "Scrum理论是什么"result = qa_chain({"query": question})print("检索问答的结果：")print(result["result"])# LLM回答效果result_2 = llm(question)print("大模型回答的结果：")print(result_2)

对比检索问答链和纯LLM的问答效果