在这篇文章中,我将向你展示如何使用Huggingface的Transformers和Sentence Transformers库来通过重排序模型提升你的检索增强生成(RAG)管道。
什么是重排序?
在我们深入评估之前,我想先简要介绍一下重排序器是什么。重排序器通常按以下方式应用:
但你可能会问,为什么重排序模型会给出与我已经相当强大的嵌入模型不同的结果?为什么不在更早的阶段利用重排序器的语义理解能力呢?这个问题涉及多个方面,但一些关键点是,例如,我们这里使用的bge-reranker本质上采用交叉编码的方法同时处理查询和文档,因此可以显式地建模查询-文档交互。另一个主要区别是,重排序模型是在监督学习的方式下训练的,用于预测通过人工标注获得的相关性得分。这在实际操作中的意义也将在后面的评估部分展示。
我们的基线
对于基线,我们选择了尽可能简单的RAG管道,并仅关注检索部分。具体来说,我们:
完成这些步骤后,只需两行代码就可以实现简单的语义搜索,即:
query_embedding = model.encode([query])[0]
results = table.search(query_embedding).limit(INITIAL_RESULTS).to_pandas()
在这里,query将是用户提供的查询,例如问题“形状补全是什么?”。Limit在这种情况下是要检索的结果数量。在一个普通的RAG管道中,检索到的结果现在将直接作为上下文提供给大型语言模型(LLM),由后者综合生成答案。在很多情况下,这样做也是完全有效的,但是对于这篇文章,我们想探索重排序的好处。
实现重排序
使用像Huggingface Transformers这样的库,使用重排序模型简直是小菜一碟。为了利用重排序来改进我们的“RAG管道”,我们按照以下方式扩展了我们的方法:
在代码中,这看起来也相当简单,并且可以用几行代码来实现:
# Instantiate the reranker
from transformers import AutoModelForSequenceClassification, AutoTokenizer
reranker_tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
reranker_model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-v2-m3').to("mps")
reranker_model.eval()
# results = ... put code to query your vector database here...
# Note that in our case the results are a dataframe containing the text
# in the "chunk" column.
# Perform a reranking
# Form query-chunk-pairs
pairs = [[query, row['chunk']] for _, row in results.iterrows()]
# Calculate relevance scores
with torch.no_grad():
inputs = reranker_tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512).to("mps")
scores = reranker_model(**inputs, return_dict=True).logits.view(-1,).float()
# Add scores to the results DataFrame
results['rerank_score'] = scores.tolist()
# Sort results by rerank score and add new rank
reranked_results = results.sort_values('rerank_score', ascending=False).reset_index(drop=True)
正如你所看到的,主要机制只是向模型提供查询和可能相关的文本对。模型会输出一个相关性得分,然后我们可以用这个得分来重新排序我们的结果列表。但这值得吗?在哪些情况下值得花费额外的推理时间?
评估重排序器
为了评估我们的系统,我们需要定义一些测试查询。在我的情况下,我选择使用以下问题类别:
1. 事实性问题,如“什么是刚体运动?”
这些问题通常在文档中有一个特定的来源,并且措辞方式使得它们甚至可能通过文本搜索找到。
2. 改述的事实性问题,如“在某些点云分类方法的架构中,是什么机制使它们对点的顺序不变?”
如你所见,这些问题在提及某些术语时不太具体,并且需要例如识别点云分类和PointNet架构之间的关系。
3. 多来源问题,如“与论文中提出的方法相比,Co-Fusion方法是如何工作的?它们有什么相似之处和不同之处?”
这些问题需要检索多个来源,这些来源要么应该被列出,要么应该相互比较。
4. 针对摘要或表格的问题,如“手部分割实验使用了哪些网络和参数大小?”
这些问题针对文本和表格形式的摘要,如模型结果的比较表。它们用于测试重排序器是否更好地识别出检索文档中的总结部分是有用的。
我每个类别只定义了5个问题来获得一个大致的印象,并评估了有和没有重排序的情况下检索到的上下文。我选择的评估标准包括:
那么结果如何呢?
即使在概览中,我们也可以看到,不同类别的问题之间存在显著差异,特别是multi_source_question类别似乎进行了大量的重排序。当我们更仔细地查看各项指标的分布时,这一点得到了进一步的证实。
具体来说,在这一类别中的5个问题中,有3个问题的最终前10名结果几乎都是通过重排序步骤进入的。现在,我们需要找出为什么会这样。因此,我们查看了受重排序影响最显著(正面影响)的两个查询。
问题1:“Co-Fusion方法是如何工作的,与论文中提出的方法相比如何?它们有什么相似之处和不同之处?”
对于这个查询,第一印象是重排序器产生了两个主要效果。它将原本位于第6位的内容提升为了顶级结果。同时,它还将几个排名很低的结果拉进了前10名。当我们进一步检查这些内容时,发现了以下情况:
总的来说,这里出现的主要模式是,重排序器能够捕捉到语调的细微差别。具体来说,像“SLAM方法与论文中提出的方法密切相关,然而”这样的表述,再加上对Co-Fusion的潜在稀疏提及,其排名会比使用标准嵌入模型高得多。这可能是因为嵌入模型很可能没有捕捉到Co-Fusion是一种SLAM方法,而文本中的主要模式是关于SLAM的一般信息。因此,重排序器在这里可以给我们带来两方面的好处:
问题2:“根据每个实验的结果,提供对引言中设定的目标完成情况的总结”
同样,在这里我们发现很多低排名的来源通过重排序步骤被提升到了前10名。让我们再次探究一下为什么会这样:
结论
实现重排序并不是一项艰巨的任务,因为像Huggingface Transformers这样的包提供了易于使用的接口,可以将其集成到你的检索增强生成(RAG)管道中,而且像llama-index和langchain这样的主要RAG框架也默认支持它们。此外,还有基于API的重排序器,如Cohere的重排序器,你可以在你的应用程序中使用。