ChatGLM的知识库/知识增强问答模块的简单设计
主要问题和需求
需求:
- 实验室在大模型的研究和应用上进展缓慢,需要做一些demo出来
- 需要做一些大模型应用的工作以便以后找工作
问题:
- langchain需要的cuda版本较高,实验室服务器cuda版本较低,docker中禁止更新cuda,管理员更新出现困难,不能直接使用全套的langchain设计
- 调用api的过程中出现了网络问题,不能一直调用api
- 之后要训练领域模型,虽然6b做领域模型很困难(灾难性遗忘),不过没辙,就是这么要求的,所以需要部署自己的模型
- langchain需要部署大模型api,服务器GPU资源有限,不能一直运行个api(demo懒得开两回,部署服务要等需求确定之后再部署,我是在开发的过程中分离开发代码,之后分离部署)
问答模块的能力(demo)
- 能使用大模型进行一定程度的问答
- 能使用知识库中的部分知识进行问答
链接:部分代码
主要思路
- 使用相似度抽取top-n的代码片,组装成模板进行问答
功能拆解
- 数据清洗(没做,回头丢给研一的调包做)
- 数据读取
- 文本滑动切片
- 文本向量化
- 知识库存储(两种方式存储)
- 大模型问答
功能实现
- 数据读取
- 直接使用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'])
- 文本滑动切片
- 自己实现的简单函数
# 单纯切分
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
)
- 文本向量化
- 为了以后支持各种模型的编码,我还是自己写了一个简单的备用
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
- 知识库存储
- 我使用了两种方法,一个是自己构的简易版本,一个是向量数据库,首先介绍我的方法,然后介绍向量库
- 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
- 使用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)