chatglm3-6b+langchain本地知识库问答

创建环境

conda create -n chatglm3 python=3.10
#这里建议使用python版本3.10

克隆代码

git clone https://github.com/THUDM/ChatGLM3.git

 起服务

cd openai_api_demo
python api_server.py

 问答代码

import os.path

from langchain_community.document_loaders import TextLoader
from langchain_community.document_loaders import DirectoryLoader


from langchain_community.llms.chatglm3 import ChatGLM3

from langchain_text_splitters import CharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.chains import RetrievalQA



# 加载embedding
embedding_model_dict = {
    "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
    "ernie-base": "nghuyong/ernie-3.0-base-zh",
    "text2vec": "GanyedeNil/text2vec-large-chinese",
    "text2vec2": "use/sbert-base-chinese-nli",
    "text2vec3": "shibing624/text2vec-base-chinese",
}


# 加载文档、将其分割成块
def load_documents(directory="books"):
    loader = DirectoryLoader(directory)
    documents = loader.load()
    text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
    split_docs = text_splitter.split_documents(documents)
    print(split_docs[:2])
    return split_docs


# 加载嵌入模型
def load_embedding_mode(model_name="text2vec3"):
    encode_kwargs = {"normalize_embeddings": False}
    model_kwargs = {"device": "cuda:0"}
    return HuggingFaceEmbeddings(
        model_name=embedding_model_dict[model_name],
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )


# 把docs文件转换为向量
def store_chroma(docs, embeddings, persist_directory="VectorStore"):
    db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
    db.persist()
    return db


# 判断是否创建
embeddings = load_embedding_mode("text2vec3")
if not os.path.exists('VectorStore'):
    documents = load_documents()
    db = store_chroma(documents, embeddings)
else:
    db = Chroma(persist_directory='VectorStore', embedding_function=embeddings)

# 加载模型
llm = ChatGLM3()

retriever = db.as_retriever()
qa = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type='stuff',
    retriever=retriever
)
response = qa.invoke('自卑与超越讲的什么？')
print(response)