Jerry Liu • 2023-05-17

使用 LLM 进行检索和重排

摘要

这篇博文概述了我们在 LlamaIndex 中围绕 LLM 驱动的检索和重排创建的一些核心抽象，这有助于超越简单的 top-k 向量嵌入查找，从而增强文档检索能力。

与基于向量嵌入的检索相比，LLM 驱动的检索可以返回更相关的文档，但其缺点是延迟和成本要高得多。我们展示了如何使用基于向量嵌入的检索作为第一阶段，然后将第二阶段检索作为重排步骤，以达到一个折中的效果。我们在《了不起的盖茨比》和 Lyft 的 SEC 10-k 文件上提供了实验结果。

两阶段检索流程：1) Top-k 向量嵌入检索，然后 2) 基于 LLM 的重排

引言与背景

在过去几个月里，“在您的数据之上构建聊天机器人”的应用掀起了一股浪潮，这得益于 LlamaIndex 和 LangChain 等框架。这些应用很多都使用了检索增强生成（RAG）的标准堆栈。

使用向量存储来存储非结构化文档（知识语料库）
给定一个查询，使用一个检索模型从语料库中检索相关文档，并使用一个合成模型来生成响应。
检索模型根据与查询的向量嵌入相似度来获取 top-k 文档。

在这个堆栈中，检索模型并不是一个新颖的想法；基于 top-k 向量嵌入的语义搜索概念已经存在至少十年了，并且完全不涉及 LLM。

基于向量嵌入的检索有很多优点

计算点积非常快。查询时不需要任何模型调用。
即使不完美，向量嵌入也能相当好地编码文档和查询的语义。有一类查询，基于向量嵌入的检索能返回非常相关的结果。

然而，由于各种原因，基于向量嵌入的检索可能不准确，并返回与查询不相关的上下文，这反过来会降低整个 RAG 系统的质量，无论 LLM 的质量如何。

这也不是一个新问题：在现有的信息检索（IR）和推荐系统中，解决这个问题的一种方法是创建两阶段过程。第一阶段使用基于向量嵌入的检索，设置较高的 top-k 值以最大化召回率，同时接受较低的准确率。然后第二阶段使用计算成本稍高但准确率更高、召回率较低的过程（例如使用 BM25）来“重排”现有检索到的候选文档。

详细阐述基于向量嵌入的检索的缺点值得写一系列博文。这篇博文是对另一种检索方法的初步探索，以及它如何（可能）增强基于向量嵌入的检索方法。

LLM 检索和重排

在过去一周里，我们围绕“基于 LLM 的”检索和重排概念开发了一系列初步抽象。概括地说，这种方法使用 LLM 来决定哪些文档/文本块与给定查询相关。输入提示将包含一组候选文档，LLM 的任务是选择相关文档集并使用内部指标对其相关性进行评分。

示例提示如下所示

A list of documents is shown below. Each document has a number next to it along with a summary of the document. A question is also provided.
  Respond with the numbers of the documents you should consult to answer the question, in order of relevance, as well
  as the relevance score. The relevance score is a number from 1–10 based on how relevant you think the document is to the question.
  Do not include any documents that are not relevant to the question.
  Example format:
  Document 1:
  <summary of document 1>
  Document 2:
  <summary of document 2>
  …
  Document 10:
  <summary of document 10>
  Question: <question>
  Answer:
  Doc: 9, Relevance: 7
  Doc: 3, Relevance: 4
  Doc: 7, Relevance: 3
  Let's try this now:
  {context_str}
  Question: {query_str}
  Answer:

提示格式意味着每个文档的文本应该相对简洁。将文本输入到与每个文档对应的提示中有两种方法

您可以直接输入与文档对应的原始文本。如果文档对应一个片段大小的文本块，这种方法效果很好。
您可以输入每个文档的精简摘要。如果文档本身对应一篇长文本，则这种方法更可取。我们在幕后使用新的文档摘要索引来实现这一点，但您也可以选择自己来做。

给定一个文档集合，我们可以创建文档“批次”，并将每个批次发送到 LLM 输入提示。每个批次的输出将是该批次内的相关文档集 + 相关性得分。最终的检索响应将聚合所有批次中的相关文档。

您可以使用我们的抽象概念以两种形式：作为独立的检索模块 (ListIndexLLMRetriever) 或重排模块 (LLMRerank)。鉴于速度/成本考虑，这篇博文的其余部分主要关注重排模块。

LLM 检索器`(ListIndexLLMRetriever)`

此模块定义在列表索引之上，列表索引简单地将一组节点存储为平面列表。您可以构建列表索引，然后在其中使用 LLM 检索器来检索相关文档。

from llama_index import GPTListIndex
from llama_index.indices.list.retrievers import ListIndexLLMRetriever
index = GPTListIndex.from_documents(documents, service_context=service_context)
# high - level API
query_str = "What did the author do during his time in college?"
retriever = index.as_retriever(retriever_mode="llm")
nodes = retriever.retrieve(query_str)
# lower-level API
retriever = ListIndexLLMRetriever()
response_synthesizer = ResponseSynthesizer.from_args()
query_engine = RetrieverQueryEngine(retriever=retriever, response_synthesizer=response_synthesizer)
response = query_engine.query(query_str)

使用场景：这可能可以替代我们的向量存储索引。您使用 LLM 而非基于向量嵌入的查找来选择节点。

LLM 重排器 (LLMRerank)

此模块定义为我们的 NodePostprocessor 抽象的一部分，该抽象用于在初始检索通过后的第二阶段处理。

后处理器可以单独使用，也可以作为 RetrieverQueryEngine 调用的一部分使用。在下面的示例中，我们展示了如何在从向量索引进行初始检索调用后，将后处理器作为独立模块使用。

from llama_index.indices.query.schema import QueryBundle
query_bundle = QueryBundle(query_str)
# configure retriever
retriever = VectorIndexRetriever(
index=index,
similarity_top_k=vector_top_k,
)
retrieved_nodes = retriever.retrieve(query_bundle)
# configure reranker
reranker = LLMRerank(choice_batch_size=5, top_n=reranker_top_n, service_context=service_context)
retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)

局限性/注意事项

基于 LLM 的检索存在某些局限性和注意事项，尤其是这个初始版本。

基于 LLM 的检索比基于向量嵌入的检索慢几个数量级。对成千上万甚至数百万个向量嵌入进行搜索只需不到一秒钟。向 OpenAI 发送一个包含 4000 个 token 的 LLM 提示可能需要几分钟才能完成。
使用第三方 LLM API 会产生费用。
当前的文档批量处理方法可能不是最优的，因为它依赖于文档批次可以相互独立评分的假设。这缺乏对所有文档全局排名的视图。

使用 LLM 检索和排文库中每个节点可能非常昂贵。这就是为什么在第一阶段向量嵌入通过后，将 LLM 用作第二阶段重排步骤会很有帮助的原因。

初步实验结果

让我们看看 LLM 重排的效果如何！

我们展示了简单的 top-k 向量嵌入检索与两阶段检索流程（先进行第一阶段向量嵌入检索过滤，后进行第二阶段 LLM 重排）之间的一些比较。我们还展示了一些纯粹基于 LLM 的检索结果（尽管由于它的运行速度比前两种方法慢得多，我们没有展示太多结果）。

我们分析了来自两个非常不同来源的数据的结果：《了不起的盖茨比》和 2021 年 Lyft SEC 10-k 文件。我们只分析了“检索”部分的结果，而非生成部分，以便更好地隔离不同检索方法的性能。

结果以定性方式呈现。下一步肯定是在整个数据集上进行更全面的评估！

《了不起的盖茨比》

在我们的第一个示例中，我们将《了不起的盖茨比》作为 Document 对象加载，并在其上构建一个向量索引（块大小设置为 512）。

# LLM Predictor (gpt-3.5-turbo) + service context
llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo"))
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, chunk_size_limit=512)
# load documents
documents = SimpleDirectoryReader('../../../examples/gatsby/data').load_data()
index = GPTVectorStoreIndex.from_documents(documents, service_context=service_context)

然后我们定义一个 get_retrieved_nodes 函数——这个函数可以只在索引上进行基于向量嵌入的检索，或者基于向量嵌入的检索 + 重排。

def get_retrieved_nodes(
    query_str, vector_top_k=10, reranker_top_n=3, with_reranker=False
):
  query_bundle = QueryBundle(query_str)
  # configure retriever
  retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=vector_top_k,
  )
  retrieved_nodes = retriever.retrieve(query_bundle)
  if with_reranker:
    # configure reranker
    reranker = LLMRerank(choice_batch_size=5, top_n=reranker_top_n, service_context=service_context)
    retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)
  return retrieved_nodes

然后我们提出一些问题。对于基于向量嵌入的检索，我们设置 k=3。对于两阶段检索，我们将向量嵌入检索的 k 设置为 10，将基于 LLM 的重排的 n 设置为 3。

问题：“是谁开车撞了 Myrtle？”

（对于不熟悉《了不起的盖茨比》的读者来说，叙述者后来从盖茨比那里得知，实际上是黛西开车撞了人，但盖茨比为她承担了责任。）

下图显示了排名前列的检索上下文。我们看到，在基于向量嵌入的检索中，排名前两位的文本包含了车祸的语义信息，但没有提供关于谁是实际肇事者的详细信息。只有第三段文本包含了正确的答案。