前言
本文实现的RAG极其简易,仅用于本人研究,可能无法满足实际应用。代码较为透明,便于对过程、结果进行后续处理,感兴趣者可以参考。
一、模型加载
本文使用的计算卡为寒武纪MLU370-X8,大模型调用可参考小军军军军军军(小军军军军军军-CSDN博客)发布的系列文章。在此以Qwen1.5-7B-chat为例,加载过程如下:
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch_mlu
import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt
device = "mlu"
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen1.5-7B-Chat",#根据自己所使用的大模型及路径填写
torch_dtype=torch.float16,
device_map="balanced_low_0"
).eval()
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-7B-Chat")
完成model和tokenizer加载后,通过下列代码实现基础大模型输出功能:
def resp(input):
messages = [
{"role": "system", "content": "YOUR SYSTEM PROMPT"},
{"role": "user", "content": input}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(device)
generated_ids = model.generate(
model_inputs.input_ids,
max_new_tokens=512
)
generated_ids = [
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
二、数据准备
本人研究过程中使用的外挂数据为简单的列表形式,列表中的一项为一条数据,例如:
texts = [
' Increased economic policy uncertainty can lead to reduced investment, which may decrease PM2.5 pollution.',
' Higher economic policy uncertainty can hinder corporate innovation, potentially increasing PM2.5 pollution.',
' The overall impact of economic policy uncertainty on PM2.5 pollution is determined by the combined effect of these two channels.',
' Empirical evidence suggests that as economic policy uncertainty increases, PM2.5 pollution levels have significantly decreased.',
]
如需读取文件,可使用pandas、numpy等读取本地文件,并转化为列表格式。
三、知识库embedding
根据喜好从Huggingface上下载embedding模型,并加载。本文以BAAI/bge-large-en-v1.5为例进行演示:
embedding = HuggingFaceEmbeddings(model_name="BAAI/bge-large-en-v1.5")
embeddings_texts = embedding.embed_documents(texts)
embeddings_texts = np.array(embeddings_texts).astype("float32")
词嵌入后,数据准备中的列表将转化为array格式如下:
其中array中的每项对应一条数据,embedding后,array形状为(样本量,embedding维度),以本文案例为例,查询array形状为:
即embedding后,每条数据被转化为1024维词向量。
四、计算查询内容余弦相似度
定义余弦相似度计算函数,功能为输入两个文本向量,返回两向量间的余弦相似度:
def cosine_similarity(vec1, vec2):
dot_product = np.dot(vec1, vec2.T)
norm_vec1 = np.linalg.norm(vec1)
norm_vec2 = np.linalg.norm(vec2)
return dot_product / (norm_vec1 * norm_vec2)
计算效果如下:
可以看出,“This is an example.”和“This is an example, too.”句意相似,返回向量间余弦相似度为0.93,而“This is an example.”和“I like doing research!”句意差别较大,余弦相似度仅有0.49。
五、RAG基础功能实现
实现RAG基础功能步骤大体为:
- 计算query与数据库中所有数据之间的余弦相似度
- 选择top k个数据库中的数据,作为大模型的提示
- 输入内容,结合先前提示,得到大模型输出
其中步骤1代码为:
def get_query_similarity(query,embeddings_texts):
query1 = [query]
embeddings_query = embedding.embed_documents(query1)
embeddings_query = np.array(embeddings_query).astype("float32")
sims = []
for each in embeddings_texts:
similarity = cosine_similarity(each, embeddings_query[0])
sims.append(similarity)
return sims
步骤2代码为:
def get_query_res(query,embeddings_texts,topk,orig_texts):
similarities = get_query_similarity(query,embeddings_texts)
idx = find_positions_of_max_k(similarities,topk)
return_res = []
for each in idx:
return_res.append(orig_texts[each])#仅输出匹配的文本
#return_res.append([orig_texts[each],similarities[each]])#输出匹配的文本及余弦相似度
return return_res
其中find_positions_of_max_k功能为从列表中找出排名前k的值,并返回其位置,具体代码为:
def find_positions_of_max_k(numbers,topk):
nums_array = np.array(numbers)
indices = np.argsort(nums_array)
max_indices = indices[-topk:]
max_indices = max_indices[::-1]
return max_indices
步骤三代码为:
def RAG_Response(query,topk,embeddings_texts,orig_texts):
query_res = get_query_res(query=query,embeddings_texts = embeddings_texts,topk = 10,orig_texts=answer)
query_res_concat = ""
for each in query_res:
query_res_concat = query_res_concat+"\n"+each
INPUT = f'''Please answer the questions strictly based on the references I gave you after ### #, not on your own knowledge. Note that the references I give you may contain content unrelated to the question. Please answer strictly according to my questions, filtering and eliminating irrelevant references.
references:### \n {query_res_concat}
question:{query}
'''
res = resp(INPUT)
return res,query_res
RAG_Response函数可根据输入的query查询知识库中相应的文本,并将文本放入INPUT中作为大模型输入的一部分,返回大模型的回答以及参考的知识库中相应文本。代码中,embeddings_texts为使用embedding模型嵌入后的词向量数据,orig_texts为原始数据。
直接获取大模型回答,可使用如下代码:
response,references = RAG_Response(query=query,topk=10,embeddings_texts=embeddings_question,orig_texts=answer)
print(f"Output:{response}\n references:{references}")
结果如下:
实现对话功能则将上述代码放置于While True循环中即可。