ChatGLM2+Langchain 构建本地知识库——构建本地知识库

参考视频：ChatGLM+Langchain构建本地知识库，只需6G显存，支持实时上传文档_哔哩哔哩_bilibili

实现原理

一、文档处理

读取文档并将文档分割成指定长度

# 读取文档
# from langchain.document_loaders import DirectoryLoader
# from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import TextLoader

#拆分文档
from langchain.text_splitter import CharacterTextSplitter


#加载文档
def load_document(document_path = "/home/david/20240207/ChatGLM2-6B/demo/date/novel.txt"):

    loader = TextLoader(document_path)
    documents = loader.load()
    #print(f'documents:{len(data)}')

    text_spliter = CharacterTextSplitter(chunk_size=512, chunk_overlap=64)
    spliter_docs = text_spliter.split_documents(documents)

    print(spliter_docs[:2]) #输入分割文档之后的前两个文档

    return spliter_docs

load_document()

$ python doc_pro.py
[Document(page_content='第一卷 菜鸟总动员\n\n\n------------\n\n第1章 好事上门\n\n    忽如一夜春风来，千树万树梨花开。\n\n    雪后方晴，倍受雾霾困挠的城市终于迎来了一个抬头见日的天气，连日降雪，道路两旁的街树上积了厚厚的一层，像玉树琼枝装点着城市，过往的行人终于卸下了成天不离的大口罩，舒一口胸中的浊气。\n\n    路牌，向右，滨海东路。向左，省警校。\n\n    一辆现代suv警车在红绿灯前稍停片刻，左转向，驶向省警校的方向。\n\n    那里被誉为全省警察的摇篮，每年向各地市县输送的各类警务人员有数百名之多，每年在最后一个学期开始之前，都有各地市的公安部门到应届毕业生里挑选实习人员，不过挂着省厅牌照的警车来此可是第一次，又驶几公里，已经看到了警校高耸的教学楼，是橄榄色的，在楼群中显得格外另类。\n\n    车驶进校园，停在教学楼下的时候，已经有学校的训导主任江晓原和校长王岚在迎接了，数人一行寒喧的场景，落在了三层一间窗户后的视线中，是一位其貌不扬的男生，他捅捅身边一位正在手机上玩连连看的同学，轻声道着：“来了。”', metadata={'source': '/home/david/20240207/ChatGLM2-6B/demo/date/novel.txt'}), Document(page_content='手机收起来了，是位胖胖的，腮帮有点鼓，五官往一块凑的男生，脸型浑圆，因为这长相被同班同学冠了个豆包的绰号，提醒他是同桌余罪，他小声道着：“余儿，这次省厅选拔，教导员让咱们高度重视，你说，这好事会不会落咱们头上？”\n\n    叫余罪的眼神很清澈，扫了眼这间大阶梯教室，乱哄哄地都在说话，省厅来本校招聘的消息早传出来来了，把小学员们刺激得，都开始憧憬未来的生活了。可学员里的阶级差别也很明显，一百多名学员，有不少是内部保送，还有不少就是本市户口，和后排这群偏远地市县来的，像两个泾渭分明的群体，连坐也很难坐到一起。\n\n    余罪一念至此，摇摇头道：“不会。有好事轮不着咱们，说不定早内定了。”\n\n    “可教导员说，这次是自愿报名，公开选拔，不至于这个上面还搞暗箱操作吧？”豆包狐疑地问。\n\n    “要没暗箱都不叫操作，留省城的机会都给你，你以为看cct.v呀？幸福那么容易？”余罪轻声道。\n\n    “可毕竟是招聘嘛，不至于都全黑了吧？”豆包抱着一线希望。\n\n    “就照顾个名额，也轮不着你呀？”余罪笑着道，看豆包不太相信，他凑了凑，小声又续道：“我猜没戏，相信兄弟我，还是相信组织吧？”', metadata={'source': '/home/david/20240207/ChatGLM2-6B/demo/date/novel.txt'})]

chunk_size: 对输入文本序列进行切分的最大长度。大语言模型一般会限制最大输入序列长度, 比如GPT - 3的最大输入长度是2048个token。
为了处理更长的文本, 需要切分成多个chunk, chunk_size控制每个chunk的最大长度。
chunk_overlap: 相邻两个chunk之间的重叠token数量。为了保证文本语义的连贯性, 相邻chunk会有一定的重叠。
chunk_overlap控制这个重叠区域的大小。
举例来说, 如果chunk_size设为1024, chunk_overlap设为128, 则对一个长度为2560的文本序列, 会切分成3个chunk:
    # chunk 1: 第1 - 1024 个token
    # chunk 2: 第897 - 1920 个token(与chunk1 重叠128个)
    # chunk 3: 第1793 - 2560
    # 个token(与chunk 2 重叠128个)
    这样的切分方式既满足了最大长度限制, 也保证了相邻chunk间语义的衔接。适当的chunk大小和重叠可以提升大语言模型处理长文本的流畅性和连贯性。

注意：

from langchain.document_loaders import DirectoryLoader

关于这个模块的引用，由于视频的制作时间是去年的九月份，至今已有半年，很多内容进行了更新。下面是有几点需要注意：

1、from langchain 需要改为 langchain_community，否则会发出警告：

LangChainDeprecationWarning: Importing document loaders from langchain is deprecated. Importing from langchain will no longer be supported as of langchain==0.2.0. Please import from langchain-community instead:

`from langchain_community.document_loaders import DirectoryLoader`.

To install langchain-community run `pip install -U langchain-community`.

2、如果读取的文档是txt文件，DirectoryLoader 方法应该改为使用 TextLoader ，否则会报错：

zipfile.BadZipFile: File is not a zip file

具体使用方法可以查找 Langchain 官方文档：Document loaders | 🦜️🔗 Langchain

将分割后的文档向量化

from langchain_community.vectorstores import Chroma

#加载文档
def load_document(document_path = "/home/david/20240207/ChatGLM2-6B/demo/date/novel.txt"):

    loader = TextLoader(document_path)
#    loader = DirectoryLoader(document_path)
    documents = loader.load()

    text_spliter = CharacterTextSplitter(chunk_size=512, chunk_overlap=64)
    spliter_docs = text_spliter.split_documents(documents)

    print("文档读取完成~")

    return spliter_docs


def load_embedding_mode():
    encode_kwargs = {"normalize_embeddings":False}
    model_kwargs = {"device":"cuda:0"}
    print("embedding加载完成~")

    return HuggingFaceEmbeddings(
        # embedding 模型地址
        model_name = "/home/david/20240207/Langchain-Chatchat/model/m3e-base/",
        model_kwargs = model_kwargs,
        encode_kwargs = encode_kwargs
    )

def store_chroma(docs,embeddings,persist_directory="VectorStore"):
    # 创建数据库
    db = Chroma.from_documents(docs,embeddings,persist_directory=persist_directory)
    db.persist()
    print("数据库创建完成~")
    return db

documents = load_document()
embeddings = load_embedding_mode()
db = store_chroma(documents,embeddings)

$ python doc_pro.py
文档读取完成~
embedding加载完成~
数据库创建完成~

运行部分可以进行优化一下：

import os

embeddings = load_embedding_mode()
# 判断是否已经创建过数据库，如果没有，则创建新的数据库，存在则使用已经存在的
if not os.path.exists('VectorStore'):
    documents = load_document()
    db = store_chroma(documents,embeddings)
else:
    db = Chroma(persist_directory="VectorStore",embedding_function=embeddings)

二、本地知识库问答

import os
# 读取文档
from langchain_community.document_loaders import TextLoader

#拆分文档
from langchain.text_splitter import CharacterTextSplitter

from langchain.embeddings.huggingface import HuggingFaceEmbeddings
# 加载embedding

# 导入向量数据库
from langchain_community.vectorstores import Chroma
# 提取文档
from langchain.chains import RetrievalQA
# 加载模型
from langchain_community.llms import ChatGLM

#加载文档
def load_document(document_path = "/home/david/20240207/ChatGLM2-6B/demo/date/novel.txt"):

    loader = TextLoader(document_path)
    documents = loader.load()
    print(f'documents:{len(data)}')

    text_spliter = CharacterTextSplitter(chunk_size=512, chunk_overlap=64)
    spliter_docs = text_spliter.split_documents(documents)

    print("文档读取完成~")
    return spliter_docs

def load_embedding_mode():
    encode_kwargs = {"normalize_embeddings":False}
    model_kwargs = {"device":"cuda:0"}
    print("embedding加载完成~")

    return HuggingFaceEmbeddings(
        # embedding 模型地址
        model_name = "/home/david/20240207/Langchain-Chatchat/model/m3e-base/",
        model_kwargs = model_kwargs,
        encode_kwargs = encode_kwargs
    )

def store_chroma(docs,embeddings,persist_directory="VectorStore"):
    # 创建数据库
    db = Chroma.from_documents(docs,persist_directory=persist_directory)
    db.persist()
    print("数据库创建完成~")
    return db



embeddings = load_embedding_mode()
# 判断是否已经创建过数据库，如果没有，则创建新的数据库，存在则使用已经存在的
if not os.path.exists('VectorStore'):
    documents = load_document()
    db = store_chroma(documents,embeddings)
else:
    db = Chroma(persist_directory="VectorStore",embedding_function=embeddings)

# 加载模型 ChatGLM
llm = ChatGLM(
    endpoint = 'http://127.0.0.1:8000',
    max_token = 80000,
    top_p = 0.9
)

retriever =  db.as_retriever()
qa = RetrievalQA.from_chain_type(
    llm = llm,
    chain_type = 'stuff',
    retriever = retriever
)
# 提问文档相关的问题
#The function `run` was deprecated in LangChain 0.1.0 and will be removed in 0.2.0. Use invoke instead.
response = qa.invoke('请介绍一下余罪')
print(response)

我添加的文档是一篇小说，检测一下效果

$ python qa_demo.py
embedding加载完成~
{'query': '请介绍一下余罪', 'result': '余罪是一个警校学员，他的品行不端、手脚不净，身上每个部分都可能成为杀器。在警校里，他受到了极端的锻炼，让他的每一个部位都变得强壮。他的性格安静，不易被吓倒，甚至有时候会对敌人表现出一种戏谑的态度。在余罪被拘留后，他经历了许多危险和恐惧，包括被围攻、被追捕等。'}

效果还可以，至此，本地知识库问答基础功能实现了。