chatglm实现基于知识库问答的应用

背景

目前由于ChatGPT横空出世,互联网如雨后春笋冒出了非常多的类ChatGPT的大型语言模型。但是对于这些语言模型,我们应该如何将它应用到我们实际的生产中需要一个更加成熟的解决方案。

介绍

本文旨在通过介绍ChatGLM的使用来讲述如何将一个开源的语言模型应用于智能问答,知识库问答的场景中,通过一系列实操例子来理解整个应用思路。

前期准备

  • 一个开源语言模型,这里推荐ChatGLM-6B,开源的、支持中英双语的对话语言模型,并且要求的显存内存非常低,可以在个人PC中轻松部署。
  • python3.8+
  • milvus,向量索引库
  • pytorch以及运行ChatGLM-6B所需要的CUDA和NVIDIA驱动
  • 基于文档的知识库问答

    实现步骤

    1. 清洗知识库文档,将文档向量化并存入向量数据库
    2. 用户提问
    3. 将用户提问向量化并查询向量数据库得到匹配的N条知识
    4. 将匹配的知识构建prompt,并通过langchain处理用户的问题
    5. 调用llm搭配prompt回答用户的问题

    向量索引

    我们首先需要定义一个向量索引库,在这里我选用的是milvus作为向量索引库来实现我们的文档向量索引和相似度匹配的工作

    为了更方便的部署,这里我采用了docker-compose来启动milvus服务。

    大家可以在milvus的官方文档中看到最新版本的部署方式Install Milvus Standalone with Docker Compose
    嫌麻烦也可以直接复制使用下面的yaml文件

    version: '3.5'
    
    services:
      etcd:
        container_name: milvus-etcd
        image: quay.io/coreos/etcd:v3.5.0
        environment:
          - ETCD_AUTO_COMPACTION_MODE=revision
          - ETCD_AUTO_COMPACTION_RETENTION=1000
          - ETCD_QUOTA_BACKEND_BYTES=4294967296
          - ETCD_SNAPSHOT_COUNT=50000
        volumes:
          - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/etcd:/etcd
        command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    
      minio:
        container_name: milvus-minio
        image: minio/minio:RELEASE.2023-03-20T20-16-18Z
        environment:
          MINIO_ACCESS_KEY: minioadmin
          MINIO_SECRET_KEY: minioadmin
        volumes:
          - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/minio:/minio_data
        command: minio server /minio_data
        healthcheck:
          test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
          interval: 30s
          timeout: 20s
          retries: 3
    
      standalone:
        container_name: milvus-standalone
        image: milvusdb/milvus:v2.2.5
        command: ["milvus", "run", "standalone"]
        environment:
          ETCD_ENDPOINTS: etcd:2379
          MINIO_ADDRESS: minio:9000
        volumes:
          - ${DOCKER_VOLUME_DIRECTORY:-.}/volumes/milvus:/var/lib/milvus
        ports:
          - "19530:19530"
          - "9091:9091"
        depends_on:
          - "etcd"
          - "minio"
    
    networks:
      default:
        name: milvus
    

    当我们创建好docker-compose.yml文件之后就可以使用命令行docker-compose up -d来启动milvus服务。

    接下来就是文档预处理

    文档预处理

    当我们收集到足够的文档之后,我们需要对文档进行一些清洗,方便我们之后的向量匹配更加精准。
    这里,我们需要完成以下步骤:

    1. 连接milvus向量库
    2. 创建对应的connection
    3. 遍历读取文档
    4. 文档预处理
    5. 文档内容转向量
    6. 存入向量库

    为此,我们编写代码如下

    import os
    import re
    import jieba
    import torch
    import pandas as pd
    from pymilvus import utility
    from pymilvus import connections, CollectionSchema, FieldSchema, Collection, DataType
    from transformers import AutoTokenizer, AutoModel
    
    connections.connect(
        alias="default",
        host='localhost',
        port='19530'
    )
    
    # 定义集合名称和维度
    collection_name = "document"
    dimension = 768
    docs_folder = "./knowledge/"
    
    tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
    model = AutoModel.from_pretrained("bert-base-chinese")
    
    
    # 获取文本的向量
    def get_vector(text):
        input_ids = tokenizer(text, padding=True, truncation=True, return_tensors="pt")["input_ids"]
        with torch.no_grad():
            output = model(input_ids)[0][:, 0, :].numpy()
        return output.tolist()[0]
    
    
    def create_collection():
        # 定义集合字段
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True, description="primary id"),
            FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=50),
            FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=10000),
            FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=dimension),
        ]
    
        # 定义集合模式
        schema = CollectionSchema(fields=fields, description="collection schema")
    
        # 创建集合
    
        if utility.has_collection(collection_name):
        	# 如果你想继续添加新的文档可以直接 return。但你想要重新创建collection,就可以执行下面的代码
            # return
            utility.drop_collection(collection_name)
            collection = Collection(name=collection_name, schema=schema, using='default', shards_num=2)
            # 创建索引
            default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 2048}, "metric_type": "IP"}
            collection.create_index(field_name="vector", index_params=default_index)
            print(f"Collection {collection_name} created successfully")
        else:
            collection = Collection(name=collection_name, schema=schema, using='default', shards_num=2)
            # 创建索引
            default_index = {"index_type": "IVF_FLAT", "params": {"nlist": 2048}, "metric_type": "IP"}
            collection.create_index(field_name="vector", index_params=default_index)
            print(f"Collection {collection_name} created successfully")
    
    
    def init_knowledge():
        collection = Collection(collection_name)
        # 遍历指定目录下的所有文件,并导入到 Milvus 集合中
        docs = []
        for root, dirs, files in os.walk(docs_folder):
            for file in files:
                # 只处理以 .txt 结尾的文本文件
                if file.endswith(".txt"):
                    file_path = os.path.join(root, file)
                    with open(file_path, "r", encoding="utf-8") as f:
                        content = f.read()
                    # 对文本进行清洗处理
                    content = re.sub(r"\s+", " ", content)
                    title = os.path.splitext(file)[0]
                    # 分词
                    words = jieba.lcut(content)
                    # 将分词后的文本重新拼接成字符串
                    content = " ".join(words)
                    # 获取文本向量
                    vector = get_vector(title + content)
                    docs.append({"title": title, "content": content, "vector": vector})
    
        # 将文本内容和向量通过 DataFrame 一起导入集合中
        df = pd.DataFrame(docs)
        collection.insert(df)
        print("Documents inserted successfully")
    
    
    if __name__ == "__main__":
        create_collection()
        init_knowledge()
    
    

    可以看到,我们创建了一个名为document的collection。它包含四个字段idtitlecontentvector其中vector储存的是content转化的向量。(当然,我们只是简单的实现了一个最原始的向量索引,如果你想要之后的匹配更加精准更加高效,你可以考虑将大文档按照段落切割并分别转化为向量,并且相互关联上。

    于此同时,我们采用了jieba作为分词库,对文本进行清洗,还使用了正则去除了文档中不必要的一些特殊符号。这些操作可以让我们向量匹配更加精准。

    当这些步骤全部执行完毕之后,我们就可以进行用户提问匹配向量库的操作了。

    用户提问匹配知识库

    首先,我们需要将用户提供的查询向量转换为blob对象,以便与数据库中的向量进行比较。我们在上个步骤实现了get_vector方法来将文本转为向量,现在可以继续调用该方法来实现。

    其次我们需要将问题转化的向量用来查找向量库,并得出最为匹配的几个结果。编写代码如下:

    import torch
    from document_preprocess import get_vector
    from pymilvus import Collection
    
    collection = Collection("document")  # Get an existing collection.
    collection.load()
    DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
    
    
    # 定义查询函数
    def search_similar_text(input_text):
        # 将输入文本转换为向量
        input_vector = get_vector(input_text)
    	# 查询前三个最匹配的向量ID
        similarity = collection.search(
            data=[input_vector],
            anns_field="vector",
            param={"metric_type": "IP", "params": {"nprobe": 10}, "offset": 0},
            limit=3,
            expr=None,
            consistency_level="Strong"
        )
        ids = similarity[0].ids
        # 通过ID查询出对应的知识库文档
        res = collection.query(
            expr=f"id in {ids}",
            offset=0,
            limit=3,
            output_fields=["id", "content", "title"],
            consistency_level="Strong"
        )
        print(res)
        return res
    
    
    if __name__ == "__main__":
    	question = input('Please enter your question: ')
        search_similar_text(question)
    
    

    上面我们通过向量索引库计算查询出了与问题最为接近的文档并打印了出来,接下来就到了最终的获取模型回答的环节了。

    通过提示模板获取准确回答

    在这一步,我们需要加载ChatGLM的预训练模型,并获取回答。

    from transformers import AutoModel, AutoTokenizer
    from knowledge_query import search_similar_text
    
    
    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
    model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True).half().cuda()
    model = model.eval()
    
    
    def predict(input, max_length=2048, top_p=0.7, temperature=0.95, history=[]):
    	res = search_similar_text(input)
    	prompt_template = f"""基于以下已知信息,简洁和专业的来回答用户的问题。
    如果无法从中得到答案,请说 "当前会话仅支持解决一个类型的问题,请清空历史信息重试",不允许在答案中添加编造成分,答案请使用中文。
    
    已知内容:
    {res}
    
    问题:
    {input}
    """
    	query = prompt_template
    	for response, history in model.stream_chat(tokenizer, query, history, max_length=max_length, top_p=top_p,
    	                                           temperature=temperature):
    	    chatbot[-1] = (parse_text(input), parse_text(response))
    	
    	    yield chatbot, history
    

    上面使用了提示模板的方式,将我们查询出来的文档作为提示内容交给模型进行推理回答。到此,我们就简单实现了一个基于知识库的问答应用。

    如果你想在web上像chatgpt一样提问,也可以丰富一下上面的代码

    from transformers import AutoModel, AutoTokenizer
    import gradio as gr
    import mdtex2html
    
    from knowledge_query import search_similar_text
    
    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True)
    model = AutoModel.from_pretrained("THUDM/chatglm-6b-int4", trust_remote_code=True).half().cuda()
    model = model.eval()
    is_knowledge = True
    
    """Override Chatbot.postprocess"""
    
    
    def postprocess(self, y):
        if y is None:
            return []
        for i, (message, response) in enumerate(y):
            y[i] = (
                None if message is None else mdtex2html.convert((message)),
                None if response is None else mdtex2html.convert(response),
            )
        return y
    
    
    gr.Chatbot.postprocess = postprocess
    
    
    def parse_text(text):
        """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/"""
        lines = text.split("\n")
        lines = [line for line in lines if line != ""]
        count = 0
        for i, line in enumerate(lines):
            if "```" in line:
                count += 1
                items = line.split('`')
                if count % 2 == 1:
                    lines[i] = f'<pre><code class="language-{items[-1]}">'
                else:
                    lines[i] = f'<br></code></pre>'
            else:
                if i > 0:
                    if count % 2 == 1:
                        line = line.replace("`", "\`")
                        line = line.replace("<", "&lt;")
                        line = line.replace(">", "&gt;")
                        line = line.replace(" ", "&nbsp;")
                        line = line.replace("*", "&ast;")
                        line = line.replace("_", "&lowbar;")
                        line = line.replace("-", "-")
                        line = line.replace(".", ".")
                        line = line.replace("!", "!")
                        line = line.replace("(", "(")
                        line = line.replace(")", ")")
                        line = line.replace("$", "$")
                    lines[i] = "<br>"+line
        text = "".join(lines)
        return text
    
    
    def predict(input, chatbot, max_length, top_p, temperature, history):
        global is_knowledge
    
        chatbot.append((parse_text(input), ""))
        query = input
        if is_knowledge:
            res = search_similar_text(input)
            prompt_template = f"""基于以下已知信息,简洁和专业的来回答用户的问题。
    如果无法从中得到答案,请说 "当前会话仅支持解决一个类型的问题,请清空历史信息重试",不允许在答案中添加编造成分,答案请使用中文。
    
    已知内容:
    {res}
    
    问题:
    {input}
    """
            query = prompt_template
            is_knowledge = False
        for response, history in model.stream_chat(tokenizer, query, history, max_length=max_length, top_p=top_p,
                                                   temperature=temperature):
            chatbot[-1] = (parse_text(input), parse_text(response))
    
            yield chatbot, history
    
    
    def reset_user_input():
        return gr.update(value='')
    
    
    def reset_state():
        global is_knowledge
    
        is_knowledge = False
        return [], []
    
    
    with gr.Blocks() as demo:
        gr.HTML("""<h1 align="center">ChatGLM</h1>""")
    
        chatbot = gr.Chatbot()
        with gr.Row():
            with gr.Column(scale=4):
                with gr.Column(scale=12):
                    user_input = gr.Textbox(show_label=False, placeholder="Input...", lines=10).style(
                        container=False)
                with gr.Column(min_width=32, scale=1):
                    submitBtn = gr.Button("Submit", variant="primary")
            with gr.Column(scale=1):
                emptyBtn = gr.Button("Clear History")
                max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
                top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top P", interactive=True)
                temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)
    
        history = gr.State([])
    
        submitBtn.click(predict, [user_input, chatbot, max_length, top_p, temperature, history], [chatbot, history],
                        show_progress=True)
        submitBtn.click(reset_user_input, [], [user_input])
    
        emptyBtn.click(reset_state, outputs=[chatbot, history], show_progress=True)
    
    demo.queue().launch(share=False, inbrowser=True)
    
    

    ChatGLM中的web_demo代码简单改写,我们就得到了一个一模一样的前端应用,不同的是它现在可以基于我们的知识库来回答问题。

    结语

    本文章仅仅介绍了最简单的通过向量索引库加AI模型加提示工程来实现知识库问答的方案,其中向量索引和文档的处理非常原始与粗糙,想要实现更加精准的匹配还需要根据实际文档内容和场景来进行修改。

    相关代码已上传github knowledge_with_chatglm感兴趣的同学可以 clone 下来跑一跑

    物联沃分享整理
    物联沃-IOTWORD物联网 » chatglm实现基于知识库问答的应用

    发表评论