本地知识库回答问题

import json
from langchain.llms.base import LLM
from transformers import AutoTokenizer, AutoModel, AutoConfig
from typing import List, Optional

class ChatGLM4(LLM):
    max_token: int = 8192
    do_sample: bool = True
    temperature: float = 0.8
    modelPath = "/home/cmcc/server/model/chatglm3-6b"
    top_p = 0.8
    tokenizer: object = AutoTokenizer.from_pretrained(modelPath, trust_remote_code=True)
    model: object = None
    history: List = []
    def __init__(self):
        super().__init__()
        print("--------------------------------")
        self.model = AutoModel.from_pretrained(self.modelPath, trust_remote_code=True).half().cuda().eval()

    @property
    def _llm_type(self) -> str:
        return "ChatGLM4"

    def _call(self, prompt: str, history: List = [], stop: Optional[List[str]] = ["<|user|>"]):
        if history is None:
            history = history
        history.append({"role": "user", "content": prompt})
        response, history = self.model.chat(self.tokenizer, prompt, history=history)
        return response

from langchain.chains import LLMChain, RetrievalQA
from langchain_community.document_loaders import TextLoader
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores.faiss import FAISS
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_text_splitters import RecursiveCharacterTextSplitter

from glm聊天客服 import ChatGLM4

embedding_model_dict = {
    "bge-large-zh-v1.5":"/home/cmcc/server/file/embFile/bge-large-zh-v1.5",
    "ebge-m3":"BAAI/bge-m3",
    "text2vec":"BAAI/text2vec-large-chinese",
    "text2vec2":"uer/sbert-base-chinese-nli",
    "text2vec3":"shibing624/text2vec-base-chinese"
}

def load_embedding_model(model_name="bge-large-zh-v1.5"):
    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
    )
embedding = load_embedding_model("bge-large-zh-v1.5")

llm = ChatGLM4()

loader = TextLoader("/home/cmcc/server/file/测试文档.txt", encoding="utf-8")

doc = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=400)

docs = text_splitter.split_documents(doc)

docsearch = FAISS.from_documents(docs, embedding)
retriever = docsearch.as_retriever()

template = """只根据以下情境回答问题：
{context}

问题: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
# model = ChatOpenAI()
output_parser = StrOutputParser()

setup_and_retrieval = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
)
chain = setup_and_retrieval | prompt | llm | output_parser

response = chain.invoke("小明和小红是什么关系")
print(response)