LLM代码实现-Qwen（挂载知识库）

为什么要挂载知识库？

LLM 在回答用户的问题时可能会产生幻觉，或者由于训练数据中不包含用户想要的内容而无法回答，通常情况下我们可以选择微调模型或者外挂知识库来缓解这类问题。微调模型的对数据和算力都有一定的要求，而知识库的门槛会更低一些，所以通常情况下会选择外挂知识库高效地来解决这类问题。

挂载知识库其实相当于引入外部知识，为了扩展语言模型以减少歧义，从大型文本数据库中检索相关文档。通常将输入序列分割成块并检索与用户输入的 query 相似的文档，然后将所选文档放在输入文本之前作为前置知识以改进模型的预测。使得模型可以更容易、更准确地访问专业知识。

挂载知识库的流程

文档 -> 文档向量化 -> 文档检索 -> 对话交互

我们可以借助 langchain 来实现这个流程：

1. 文档

我们使用中医药书籍来作为知识库，下载链接（700 本中医药古籍文本），下载完成后运行以下代码以合并数据。

import os


dir_path = "./TCM-Ancient-Books-master"
for index, filename in enumerate(os.listdir(dir_path)):
    if not filename.endswith(".txt"):
        continue

    file_path = os.path.join(dir_path, filename)
    with open(file_path, "r", encoding='gb18030', errors='ignore') as f:
        text = f.read().replace("\n", "")

    mode = "a" if index else "w"
    with open("./knowledge.txt", mode, encoding="utf-8") as f:
        f.write(text + "\n")

2. 文档向量化

在文档检索的过程中如果利用字符串直接匹配文本相似度效率是很低的，尤其是知识库体量非常大的时候，因此一般会先对文档进行切割和向量化以提升检索的速度。将每个文档转换为数值向量，以便计算文档之间的相似度或进行聚类分析。

from langchain.document_loaders import UnstructuredFileLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS

def load_knowledge():
    filepath = "./knowledge.txt"
    loader = UnstructuredFileLoader(filepath)
    docs = loader.load()

    # chunk-size是文本最大的字符数。chunk-overlap是前后两个chunk的重叠部分最大字数
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=40)
    docs = text_splitter.split_documents(docs)

    # 这里需要下载一个中文的文本向量化模型
    embeddings = HuggingFaceEmbeddings(model_name="./text2vec-large-chinese",
                                       model_kwargs={'device': 'cuda'})
    
    # 指定向量化文档加载/保存路径
    save_path = "./med_faiss_store.faiss"
    if not os.path.exists(save_path):
        vector_store = FAISS.from_documents(docs, embeddings)
        vector_store.save_local(save_path)

    else:
        vector_store = FAISS.load_local(save_path, embeddings=embeddings)

    return vector_store

3. 文档检索

利用以下代码根据用户的输入按照相关性对向量库中的文本文本进行排序并取出排名靠前的 5 条知识，并将知识库中的知识库和用户的输入拼接在一起作为新的 prompt。

docs = vector_store.similarity_search(patient_history)     # 计算相似度，并把相似度高的chunk放在前面
knowledge = [doc.page_content for doc in docs[:5]]  # 提取chunk的文本内容
prompt = f"知识库：{knowledge}\n问题如下：\n{patient_input}"

4. 对话交互

最后是把 prompt 送到模型中得到输出，与模型的交互在上一篇文章中有详细的介绍，这里就不赘述了，直接给出完整的代码（运行时要注意模型路径、知识库路径是否正确）。

from langchain.document_loaders import UnstructuredFileLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import GenerationConfig
import torch
import os


def load_model(model_path):
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_path, 
                                                 device_map="auto", 
                                                 trust_remote_code=True
                                                 ).eval()

    return tokenizer, model


def load_knowledge():
    filepath = "./knowledge.txt"
    loader = UnstructuredFileLoader(filepath)
    docs = loader.load()

    # chunk-size是文本最大的字符数。chunk-overlap是前后两个chunk的重叠部分最大字数
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=40)
    docs = text_splitter.split_documents(docs)

    embeddings = HuggingFaceEmbeddings(model_name="./text2vec-large-chinese",
                                       model_kwargs={'device': 'cuda'}.

    # 指定向量化文档加载/保存路径
    save_path = "./med_faiss_store.faiss"
    if not os.path.exists(save_path):
        vector_store = FAISS.from_documents(docs, embeddings)
        vector_store.save_local(save_path)

    else:
        vector_store = FAISS.load_local(save_path, embeddings=embeddings)

    return vector_store


def clear_screen():
    os.system('clear')
    return [], ""


def chat_qwen(model, tokenizer, SYSTEM_PROMPT, with_knowledge=False):
    vector_store = load_knowledge()
    history, patient_history = clear_screen()

    while True:
        patient_input = input("user：")
        patient_history += patient_input
        if patient_input.lower() == "clc":
            history, patient_history = clear_screen()
            continue

        if with_knowledge:
            docs = vector_store.similarity_search(patient_history)     # 计算相似度，并把相似度高的chunk放在前面
            knowledge = [doc.page_content for doc in docs[:5]]  # 提取chunk的文本内容
            prompt = f"知识库：{knowledge}\n问题如下：\n{inputs}"

        else:
            prompt = inputs

        response, _ = model.chat(tokenizer, prompt, history=history, system=SYSTEM_PROMPT)
        history.append((patient_input, response))    # history 不包括系统提示和知识库信息
        print("assistant：", response, end="\n\n")


if __name__ == '__main__':
    model_path = "/root/autodl-tmp/LLM_MODEL/Qwen-1_8B-Chat"

    SYSTEM_PROMPT = ""

    tokenizer, model = load_model(model_path)
    chat_qwen(model, tokenizer, SYSTEM_PROMPT, with_knowledge=True)