使用Huggingface Transformers优化RAG管道中的检索

2024年11月11日 由 alex 发表 90 0

在这篇文章中,我将向你展示如何使用Huggingface的Transformers和Sentence Transformers库来通过重排序模型提升你的检索增强生成(RAG)管道。


什么是重排序?

在我们深入评估之前,我想先简要介绍一下重排序器是什么。重排序器通常按以下方式应用:

  1. 在RAG管道的检索步骤中,使用简单的基于嵌入的检索方法来检索一组初始候选项。
  2. 使用重排序器对结果进行重新排序,以提供更适合用户查询的新结果顺序。


但你可能会问,为什么重排序模型会给出与我已经相当强大的嵌入模型不同的结果?为什么不在更早的阶段利用重排序器的语义理解能力呢?这个问题涉及多个方面,但一些关键点是,例如,我们这里使用的bge-reranker本质上采用交叉编码的方法同时处理查询和文档,因此可以显式地建模查询-文档交互。另一个主要区别是,重排序模型是在监督学习的方式下训练的,用于预测通过人工标注获得的相关性得分。这在实际操作中的意义也将在后面的评估部分展示。


我们的基线

对于基线,我们选择了尽可能简单的RAG管道,并仅关注检索部分。具体来说,我们:

  1. 选择一个大型的PDF文档。我选择了我的硕士论文,但你可以选择任何你喜欢的文档。
  2. 从PDF中提取文本,并将其分成每个约10个句子的等长段落。
  3. 为这些段落创建嵌入,并将它们插入到向量数据库中,本例中使用的是LanceDB。


完成这些步骤后,只需两行代码就可以实现简单的语义搜索,即:


query_embedding = model.encode([query])[0]
results = table.search(query_embedding).limit(INITIAL_RESULTS).to_pandas()


在这里,query将是用户提供的查询,例如问题“形状补全是什么?”。Limit在这种情况下是要检索的结果数量。在一个普通的RAG管道中,检索到的结果现在将直接作为上下文提供给大型语言模型(LLM),由后者综合生成答案。在很多情况下,这样做也是完全有效的,但是对于这篇文章,我们想探索重排序的好处。


实现重排序

使用像Huggingface Transformers这样的库,使用重排序模型简直是小菜一碟。为了利用重排序来改进我们的“RAG管道”,我们按照以下方式扩展了我们的方法:

  1. 与之前一样,首先通过标准的嵌入模型检索一定数量的初始结果。但是我们将结果的数量从10个增加到大约50个。
  2. 在检索到这么多初始来源后,我们应用一个重排序模型来重新排序这些来源。这是通过计算每个查询-来源对的相关性得分来完成的。
  3. 对于答案生成,我们通常会使用新的前x个结果。(在我们的例子中,我们使用前10个)


在代码中,这看起来也相当简单,并且可以用几行代码来实现:


# 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个问题来获得一个大致的印象,并评估了有和没有重排序的情况下检索到的上下文。我选择的评估标准包括:

  1. 重排序是否为上下文添加了重要信息。
  2. 重排序是否减少了上下文的冗余。
  3. 重排序是否使最相关的结果在列表中的位置更高(更好的优先级排序)。
  4. ……


那么结果如何呢?


2


即使在概览中,我们也可以看到,不同类别的问题之间存在显著差异,特别是multi_source_question类别似乎进行了大量的重排序。当我们更仔细地查看各项指标的分布时,这一点得到了进一步的证实。


3


具体来说,在这一类别中的5个问题中,有3个问题的最终前10名结果几乎都是通过重排序步骤进入的。现在,我们需要找出为什么会这样。因此,我们查看了受重排序影响最显著(正面影响)的两个查询。


问题1:“Co-Fusion方法是如何工作的,与论文中提出的方法相比如何?它们有什么相似之处和不同之处?”


4


对于这个查询,第一印象是重排序器产生了两个主要效果。它将原本位于第6位的内容提升为了顶级结果。同时,它还将几个排名很低的结果拉进了前10名。当我们进一步检查这些内容时,发现了以下情况:

  1. 重排序器成功地找出了一个与查询高度相关且描述了与论文中方法相对的SLAM方法的内容块。
  2. 重排序器还成功地包含了一个提到Co-Fusion作为能够处理动态对象的SLAM方法示例的内容块,并讨论了其局限性。


总的来说,这里出现的主要模式是,重排序器能够捕捉到语调的细微差别。具体来说,像“SLAM方法与论文中提出的方法密切相关,然而”这样的表述,再加上对Co-Fusion的潜在稀疏提及,其排名会比使用标准嵌入模型高得多。这可能是因为嵌入模型很可能没有捕捉到Co-Fusion是一种SLAM方法,而文本中的主要模式是关于SLAM的一般信息。因此,重排序器在这里可以给我们带来两方面的好处:

  1. 关注各自内容块中的细节,而不是追求平均语义内容。
  2. 更加关注用户将某种方法与论文中的方法进行比较的意图。


问题2:“根据每个实验的结果,提供对引言中设定的目标完成情况的总结”


5


同样,在这里我们发现很多低排名的来源通过重排序步骤被提升到了前10名。让我们再次探究一下为什么会这样:

  1. 重排序器再次成功地捕捉到了问题的细微意图,并将包含“因此推测……”这样表述的内容块重新排名为高度相关,这确实是正确的,因为接下来描述的是这些假设是否有效以及该方法是否能利用这些假设。
  2. 重排序器提供了很多以隐晦方式表述的实验结果,其中包括大量关于机器学习训练结果的表格概览,可能理解了这些部分的总结性质。


结论

实现重排序并不是一项艰巨的任务,因为像Huggingface Transformers这样的包提供了易于使用的接口,可以将其集成到你的检索增强生成(RAG)管道中,而且像llama-index和langchain这样的主要RAG框架也默认支持它们。此外,还有基于API的重排序器,如Cohere的重排序器,你可以在你的应用程序中使用。



文章来源:https://medium.com/towards-data-science/reranking-using-huggingface-transformers-for-optimizing-retrieval-in-rag-pipelines-fbfc6288c91f
欢迎关注ATYUN官方公众号
商务合作及内容投稿请联系邮箱:bd@atyun.com
评论 登录
热门职位
Maluuba
20000~40000/月
Cisco
25000~30000/月 深圳市
PilotAILabs
30000~60000/年 深圳市
写评论取消
回复取消