接上一篇的内容。
一、近似检索索引构建
对上一篇中的构建再次进行进一步的介绍。
因为我们对每个文本的嵌入,得到的向量都是768维,设置d=768;nlist是构建索引的一个超参数,用于指定IVFFlat索引中将所有向量分成的粗粒度簇的数目,考虑我们数据集的大小并经过若干次试验后发现100是一个比较合适的值。
为了提高索引构建的速度,使用faiss.index_cpu_to_gpu将该过程转到gpu上进行构建;之后直接调用train方法生成检索索引,这里使用全量的data用于训练(生成100个簇心);只用调用add方法,将data中的每个向量分配到训练生成的100个簇心中(分桶)。为了使用方便,我们还是选择在cpu上使用索引,使用index_gpu_to_cpu方法将索引重构为cpu上的索引;最后将索引写入cpu_index.ivfdata文件,后续可以直接读取使用,不用再新增索引。
## Using an IVF index
d=768
nlist = 100
quantizer = faiss.IndexFlatL2(d) # the other index
index_ivf = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
# here we specify METRIC_L2, by default it performs inner-product search
res = faiss.StandardGpuResources() # use a single GPU
# make it an IVF GPU index
gpu_index_ivf = faiss.index_cpu_to_gpu(res, 0, index_ivf)
assert not gpu_index_ivf.is_trained
gpu_index_ivf.train(data) # add vectors to the index
assert gpu_index_ivf.is_trained
gpu_index_ivf.add(data) # add vectors to the index
print(gpu_index_ivf.ntotal)
# 将 GPU 索引转移到 CPU 上
cpu_index_ivf = faiss.index_gpu_to_cpu(gpu_index_ivf)
# 存储索引
index_filename = "cpu_index.ivfdata" # 索引文件的路径和名称
faiss.write_index(cpu_index_ivf, index_filename)
IVFFlat本质上仍是暴力搜索,只是使用分桶的方式减少搜索范围,但考虑到我们的数据集本身就是10万量级的,不是一个很大的数据量,分桶后进行暴力搜索也能达到可接受的时间;且调整好搜索簇数,使用全精度的距离计算,IVFFlat也能提供近似暴力搜索下的top-k近邻,因此对我们的任务来说,我们认为使用IVFFlat是一个较好的选择。
在后续的统计中,对于一条query的查询耗时大约为0.03秒,与模型生成的速度相比可以忽略不计。
(除此IVFFlat之外,还调研了faiss库中提供的其他检索索引,见博客。。)
二、检索部分与模型的连接
这部分中的API连接部分(flask框架实现客户端和服务端的分离)由我们组的另一个同学负责,我在其中做了检索实现和query嵌入的工作。
1.初始化工作
从request中获取到本次查询的内容和模式,加载SentenceTransformers库中提供的moki-ai/m2e-base嵌入模型(与知识库的文本嵌入方式相同)用于对qeruy做嵌入,以便在知识库中进行搜索。
query = request.args.get('content')
mode = request.args.get('mode')
model = SentenceTransformer("moka-ai/m3e-base")
index_filename = "cpu_index.ivfdata" # 索引文件的路径和名称
# 加载索引
loaded_index = faiss.read_index(index_filename)
加载上节中生成的知识库索引cpu_index.ivfdata到变量loaded_index中。
加载在记录(一)中生成的知识库切分文本文件到变量text_document中,因为rag最终要返回给模型文本作为prompt,因此提前加载到内存中,方便检索到top-k索引后取出原始知识库文本。
text_documents = []
# 打开文本文件以供读取
with open('./chunk.txt', 'r', encoding='utf-8') as file:
# 逐行读取文件内容
for line in file:
# 去除行末的换行符
line = line.rstrip('\n')
text_documents.append(line)
2.在知识库中进行向量检索
使用嵌入模型对query做嵌入,之后转为1*d的numpy数组以符合search方法要求的参数格式。直接调用loaded_index的search在知识库索引中检索距query向量xq的L2距离最近的3个向量,返回值I中存储top-3个节点的索引。之后用这些索引到text_document取出原始文本(检索索引中向量的索引下标与切分后的知识库中文文本的索引下标是一致的),使用字符串拼接技术 拼成一个rag_answer
xq = model.encode(query)
# 将嵌入向量转换为 1*d 的 NumPy 数组
xq = np.reshape(xq, (1, -1))
k = 3 # we want to see 3 nearest neighbors
D, I = loaded_index.search(xq, k) # actual search
rag_answer = ''
I = I[0]
for index in I:
bef = mask[index]
rag_answer += text_documents[index][bef:] + '\n'
之后调用chatglm提供的API,将query和rag_answer两个中文文本传给大语言模型,完成rag和LLM的连接工作。
3.rag效果简单展示
对于给定的query="最近失恋了,心情不好"
在知识库中检索到的top-3如下:
将query与top-3拼接后,传给chatGLM,返回的结果如下:‘
面对失恋的痛楚,确实是一件让人心力交瘁的事情。你提到的“打开心扉,拥抱当下的生活”是一种积极的态度,虽然在情绪低落时可能难以做到。从你的对话中,我感受到你正在努力理解这份丧失感,并试图找到走出困境的方法。
具体到你的苦恼,每个人失恋的原因都不尽相同。可能是由于你们之间的某些差异,可能是环境或现实的原因,也可能只是简单的感情变化。你的感受是复杂的,既有可能难过自己的选择,也可能因为无法掌控这段关系的走向而感到挫败。
无论是因为什么原因导致恋爱失败,这都不是你个人的失败。它是两个人互动的结果,也是生活中不可避免的一部分。你现在需要的,可能正如你所说,是一点时间去疗伤。旅行、阅读、看电影、听音乐,这些都能帮助你暂时从悲伤中抽离,给你空间去整理思绪。
至于你的年龄或是是否分析过原因,这些都并不是最重要的。重要的是,你愿意面对自己的情感,承认自己的不舍和难过。当你准备好时,可以尝试分析这段关系,从中学习,但不要急于求成。
记住,每段关系都是一次学习和成长的机会。你不必为自己的情感眼光或是无法掌控恋情感到羞愧。生活总是充满了未知,而勇气和决心正是我们面对这些未知时所必需的。当你准备好了,我相信你会以更加成熟和自信的姿态,继续前行。祝你好运。
可以看出模型在一定程度上参考了检索到的top-3中的内容,生成了对话的回答。