ChatGLM的知识库/知识增强问答模块的简单设计

主要问题和需求

需求：

1. 实验室在大模型的研究和应用上进展缓慢，需要做一些demo出来
2. 需要做一些大模型应用的工作以便以后找工作

问题：

1. langchain需要的cuda版本较高，实验室服务器cuda版本较低，docker中禁止更新cuda，管理员更新出现困难，不能直接使用全套的langchain设计
2. 调用api的过程中出现了网络问题，不能一直调用api
3. 之后要训练领域模型，虽然6b做领域模型很困难（灾难性遗忘），不过没辙，就是这么要求的，所以需要部署自己的模型
4. langchain需要部署大模型api，服务器GPU资源有限，不能一直运行个api（demo懒得开两回，部署服务要等需求确定之后再部署，我是在开发的过程中分离开发代码，之后分离部署）

问答模块的能力（demo）

1. 能使用大模型进行一定程度的问答
2. 能使用知识库中的部分知识进行问答
   链接：部分代码

主要思路

• 使用相似度抽取top-n的代码片，组装成模板进行问答

功能拆解

0. 数据清洗（没做，回头丢给研一的调包做）
1. 数据读取
2. 文本滑动切片
3. 文本向量化
4. 知识库存储（两种方式存储）
5. 大模型问答

功能实现

1. 数据读取

• 直接使用langchain的是最优解，格式直接与向量数据库Chroma匹配

from langchain_community.document_loaders import UnstructuredFileLoader
paths = ['path']
loaders = [UnstructuredFileLoader(path) for path in paths]
docs = []
for loader in loaders:
    docs.extend(loader.load())
print(docs)
print(docs[0].metadata['source'])

2. 文本滑动切片

• 自己实现的简单函数

# 单纯切分
def text_splitter(text,chunk_size=1000):
    return [text[i:i + chunk_size] for i in range(0, len(text), chunk_size)]

# 滑动窗口
def text_splitter_sliding_window(text, window_size=1000, step_size=500):
    text_length = len(text)
    segments = []
    for i in range(0, text_length - window_size + 1, step_size):
        segment = text[i:i + window_size]
        segments.append(segment)
    return segments

• 调包

from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1000,
    chunk_overlap = 500
)

3. 文本向量化

• 为了以后支持各种模型的编码，我还是自己写了一个简单的备用

from text2vec import SentenceModel
text2vec_model = SentenceModel(text2vec_model_path)

def text_2_vec(text2vec_model,texts,batch_size=1,device='cpu'):
    embeddings = text2vec_model.encode(texts,batch_size=batch_size,device=device)
    return embeddings

• langchain自带的文本向量化模块是调接口

def doc_embedding(source_text, dataset_path):
    # embeddings = OpenAIEmbeddings()
    embeddings = HuggingFaceEmbeddings(model_name='/root/autodl-fs/sentence-transformers/all-mpnet-base-v2')
    with open(source_text, 'r') as f:
        state_of_the_union = f.read()
    text_splitter = CharacterTextSplitter(chunk_size=256, chunk_overlap=0)
    pages = text_splitter.split_text(state_of_the_union)
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=256, chunk_overlap=100)
    texts = text_splitter.create_documents(pages)
    print(texts)
    db = DeepLake.from_documents(texts, embeddings, dataset_path=dataset_path, overwrite=True)
    # return db

4. 知识库存储

• 我使用了两种方法，一个是自己构的简易版本，一个是向量数据库，首先介绍我的方法，然后介绍向量库
• 1.我的方法包含以下步骤：
• • 文本向量化
• • 相似度检索
• • 拼接模板

# 这是直接调用的代码
def chat_with_data(prompt,embedding_vectors,history,prompt_template,data_list,text2vec_model,llm,show_input=True):
    index,_ = find_most_similar_data(prompt, embedding_vectors,text2vec_model)
    data_text = data_list[index]
    print('==知识库检索成功')
    response,history = chat(llm,prompt_template,data_text,prompt,history=history,show_input=show_input)
    return index,response,history

# 知识库向量化
window_size = 400 # 窗口大小
step_size = 200   # 步幅大小
# 两种切片方式
# data_list = text_splitter(data_texts,chunk_size=chunk_size)
data_sliding_list = text_splitter_sliding_window(data_texts, window_size=window_size, step_size=step_size)
for i in data_sliding_list:
    print(len(i))
embedding_vectors = text_2_vec(text2vec_model,data_sliding_list,batch_size=1)
print('==知识库向量化成功')

# 寻找最相似的向量
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
def similarity(input_embeddings, data_embeddings):
    one_dim_matrix = input_embeddings.reshape(1, -1)
    similarities = cosine_similarity(one_dim_matrix, data_embeddings)
    most_similar_index = np.argmax(similarities)
    most_similar_vector = data_embeddings[most_similar_index]
    return most_similar_index,similarities,most_similar_vector


def find_most_similar_vec(input_embeddings, data_embeddings):
    most_similar_index,similarities,most_similar_vector = similarity(input_embeddings, data_embeddings)
    return most_similar_index,similarities

# 根据知识替换模板 并进行对话
prompt_template = """【任务说明】:你现在的任务是根据相关知识正确回答问题，加油！
    
    这是可能相关的知识：
    {text}

    这是我的问题：
    {prompt}

    请根据相关知识进行回答 ->
    """
    
def get_prompt(prompt_template,text,prompt):
    return prompt_template.replace('{text}',text).replace('{prompt}',prompt)
def chat(llm,prompt,history=[],show_input=True):
    if show_input == True:
        print()
        print('输入格式：')
        print(prompt)
        print()
    rst,history = llm.call(prompt,history)
    return rst,history

2. 使用chromaDB进行向量持久化存储

from langchain_community.document_loaders import UnstructuredFileLoader
paths = ['./政府工作报告2023.txt']
loaders = [UnstructuredFileLoader(path) for path in paths]

docs = []
for loader in loaders:
    docs.extend(loader.load())
print(docs)

print(docs[0].metadata['source'])

# 数据清洗（换行符等）




# Split
from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 1000,
    chunk_overlap = 500
)


#切割文档
splits = text_splitter.split_documents(docs)
#查看切割后文档的数量
print(splits[0])
print(len(splits))

splits = splits[:3] # 少一点
print(splits)
#向量模型
from text2vec import SentenceModel
text2vec_model_path = 'E:/LLM/models/text2vec'
t2v_model = SentenceModel(text2vec_model_path)
def t2v(t2v_model,text):
    embeddings = t2v_model.encode(text)
    return embeddings
embeddings = [t2v(t2v_model,peice.page_content) for peice in splits]
embeddings = [numpy_array.tolist() for numpy_array in embeddings]
print(len(embeddings))

# 向量数据库设置
import chromadb
chroma_client = chromadb.PersistentClient(path="./chroma")
# 新建集合
print('create')
collection = chroma_client.create_collection(name="test_collection")
# 增加数据
print('add')
collection.add(
    embeddings=embeddings,
    documents=[ peice.page_content for peice in splits],
    metadatas=[{"source": peice.metadata['source'],"user": "user1"} for peice in splits],
    ids=["id" + str(i) for i in range(len(splits))]
)

# 查询
# Chroma 会返回 n 个最相似的结果
# 
print('query')
query_emd = t2v(t2v_model,"粮食产量怎么样？").tolist()
print(query_emd)
results = collection.query(
    query_embeddings=query_emd,
    n_results=2
)
print(results)