知识库助手的构建之路：ChatGLM3-6B和LangChain的深度应用

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, Optional

import 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 = None
    model: AutoModelForCausalLM = None

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

    @property
    def _llm_type(self) -> str:
	    """ 返回模型类型 """
        return "ChatGLM3-6B"

文件读取与处理

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

import os
from typing import List

from langchain_community.document_loaders import TextLoader, UnstructuredMarkdownLoader, PyPDFLoader, Docx2txtLoader
from tqdm import tqdm


class 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_list

    def 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中，跳过当前文件
                continue
        return docs

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

from File_Proces import FileProces


class EmbeddingLLM():
    embeddings: HuggingFaceBgeEmbeddings = None
    vectordb = 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 PromptTemplate

def builder_chain(llm, embeddingLlm):
    # 构造Prompt模板
    template = """
    使用以下上下文来回答问题。如果你不知道答案，你需要澄清，不要试图编造答案。请简明扼要的回答。”。
    {context}
    问题: {question}
    回答:
    """

    # 实例化Template对象，context和question两个变量会被检索到的文档片段和用户提问填充
    qa_chain_prompt = PromptTemplate(input_variables=["context", "question"], template=template)

    # 构建检索链
    vectordb = embeddingLlm.vectordb

    qa_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_chain

if __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的问答效果