LlamaIndex 评测：如何为 RAG 挑选最佳 Embedding 与 Reranker 模型组合

更新：Jina AI 嵌入的池化方法已调整为使用平均池化，结果也相应更新。值得注意的是，使用 bge-reranker-large 的 JinaAI-v2-base-en 现在显示出 0.938202 的命中率和 0.868539 的 MRR（平均倒数排名），而使用 CohereRerank 时，命中率为 0.932584，MRR 为 0.873689。

立即探索我们的免费和付费计划。

在构建检索增强生成（RAG）管道时，一个关键组件是检索器。我们有多种嵌入模型可供选择，包括 OpenAI、CohereAI 和开源的句子转换器。此外，CohereAI 和句子转换器还提供了几种重排器。

但面对这么多选项，我们如何确定最佳组合以实现顶级的检索性能？我们如何知道哪个嵌入模型最适合我们的数据？或者哪个重排器最能提升我们的结果？

在这篇博客文章中，我们将使用 LlamaIndex 的 Retrieval Evaluation 模块来快速确定嵌入和重排器模型的最佳组合。让我们开始吧！

首先，让我们从理解 Retrieval Evaluation 中可用的指标开始。

理解检索评估中的指标：

为了衡量我们检索系统的有效性，我们主要依赖两个广泛接受的指标：命中率和平均倒数排名（MRR）。让我们深入探讨这些指标，了解它们的重要性以及它们如何运作。

命中率：

命中率计算的是在 top-k 检索文档中找到正确答案的查询比例。简单来说，就是我们的系统在前几次猜测中猜对的频率。

平均倒数排名（MRR）：

对于每个查询，MRR 通过查看最高排名相关文档的排名来评估系统的准确性。具体来说，它是所有查询中这些排名的倒数的平均值。因此，如果第一个相关文档是排名第一的结果，倒数排名为 1；如果是第二，倒数排名为 1/2，依此类推。

现在我们已经确定了范围并熟悉了这些指标，是时候深入实验了。为了获得实践经验，您也可以使用我们的 Google Colab Notebook 进行操作。

设置环境

!pip install llama-index sentence-transformers cohere anthropic voyageai protobuf pypdf

设置密钥

openai_api_key = 'YOUR OPENAI API KEY'
cohere_api_key = 'YOUR COHEREAI API KEY'
anthropic_api_key = 'YOUR ANTHROPIC API KEY'
openai.api_key = openai_api_key

下载数据

我们将使用 Llama2 论文进行此实验。让我们下载这篇论文。

!wget --user-agent "Mozilla" "https://arxiv.org/pdf/2307.09288.pdf" -O "llama2.pdf"

加载数据

让我们加载数据。我们将使用从开头到第 36 页的内容进行实验，这排除了目录、参考文献和附录。

然后，这些数据被解析并转换为节点，节点代表我们希望检索的数据块。我们确实使用了 chunk_size 为 512。

documents = SimpleDirectoryReader(input_files=["llama2.pdf"]).load_data()

node_parser = SimpleNodeParser.from_defaults(chunk_size=512)
nodes = node_parser.get_nodes_from_documents(documents)

生成问题-上下文对：

为了评估目的，我们创建了一个问题-上下文对的数据集。这个数据集可以看作是一组问题及其在我们数据中对应的上下文。为了消除评估嵌入（OpenAI/ CohereAI）和重排器（CohereAI）时的偏差，我们使用 Anthropic LLM 来生成问题-上下文对。

让我们初始化一个提示模板来生成问题-上下文对。

# Prompt to generate questions
qa_generate_prompt_tmpl = """\
Context information is below.

---------------------
{context_str}
---------------------

Given the context information and not prior knowledge.
generate only questions based on the below query.

You are a Professor. Your task is to setup \
{num_questions_per_chunk} questions for an upcoming \
quiz/examination. The questions should be diverse in nature \
across the document. The questions should not contain options, not start with Q1/ Q2. \
Restrict the questions to the context information provided.\
"""

llm = Anthropic(api_key=anthropic_api_key)
qa_dataset = generate_question_context_pairs(
    nodes, llm=llm, num_questions_per_chunk=2
)

用于过滤掉诸如 Here are 2 questions based on provided context 这样的句子的函数。


def filter_qa_dataset(qa_dataset):
    """
    Filters out queries from the qa_dataset that contain certain phrases and the corresponding
    entries in the relevant_docs, and creates a new EmbeddingQAFinetuneDataset object with
    the filtered data.

    :param qa_dataset: An object that has 'queries', 'corpus', and 'relevant_docs' attributes.
    :return: An EmbeddingQAFinetuneDataset object with the filtered queries, corpus and relevant_docs.
    """

    
    queries_relevant_docs_keys_to_remove = {
        k for k, v in qa_dataset.queries.items()
        if 'Here are 2' in v or 'Here are two' in v
    }

    
    filtered_queries = {
        k: v for k, v in qa_dataset.queries.items()
        if k not in queries_relevant_docs_keys_to_remove
    }
    filtered_relevant_docs = {
        k: v for k, v in qa_dataset.relevant_docs.items()
        if k not in queries_relevant_docs_keys_to_remove
    }

    
    return EmbeddingQAFinetuneDataset(
        queries=filtered_queries,
        corpus=qa_dataset.corpus,
        relevant_docs=filtered_relevant_docs
    )


qa_dataset = filter_qa_dataset(qa_dataset)

自定义检索器：

为了确定最佳检索器，我们结合使用了嵌入模型和重排序器。起初，我们建立一个基础VectorIndexRetriever。检索到节点后，我们引入重排序器进一步优化结果。值得注意的是，对于这个特定实验，我们将similarity_top_k设为10，并选用重排序后的前5名。但请根据您具体实验的需求随意调整此参数。我们在此展示使用OpenAIEmbedding的代码，请参考笔记本查看其他嵌入的代码。

embed_model = OpenAIEmbedding()
service_context = ServiceContext.from_defaults(llm=None, embed_model = embed_model)
vector_index = VectorStoreIndex(nodes, service_context=service_context)
vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k = 10)

class CustomRetriever(BaseRetriever):
    """Custom retriever that performs both Vector search and Knowledge Graph search"""

    def __init__(
        self,
        vector_retriever: VectorIndexRetriever,
    ) -&gt; None:
        """Init params."""

        self._vector_retriever = vector_retriever

    def _retrieve(self, query_bundle: QueryBundle) -&gt; List[NodeWithScore]:
        """Retrieve nodes given query."""

    retrieved_nodes = self._vector_retriever.retrieve(query_bundle)

    if reranker != 'None':
      retrieved_nodes = reranker.postprocess_nodes(retrieved_nodes, query_bundle)
       else:
          retrieved_nodes = retrieved_nodes[:5]
         
       return retrieved_nodes

    async def _aretrieve(self, query_bundle: QueryBundle) -&gt; List[NodeWithScore]:
        """Asynchronously retrieve nodes given query.

        Implemented by the user.

        """
        return self._retrieve(query_bundle)

    async def aretrieve(self, str_or_query_bundle: QueryType) -&gt; List[NodeWithScore]:
        if isinstance(str_or_query_bundle, str):
            str_or_query_bundle = QueryBundle(str_or_query_bundle)
        return await self._aretrieve(str_or_query_bundle)

custom_retriever = CustomRetriever(vector_retriever)

评估：

为了评估我们的检索器，我们计算了平均倒数排名（MRR）和命中率指标：

retriever_evaluator = RetrieverEvaluator.from_metric_names(
    ["mrr", "hit_rate"], retriever=custom_retriever
)
eval_results = await retriever_evaluator.aevaluate_dataset(qa_dataset)

结果：

我们测试了多种嵌入模型和重排序器。以下是考虑到的模型：

嵌入模型：

重排序器：

值得一提的是，这些结果为该特定数据集和任务提供了坚实的性能洞察。然而，实际结果可能因数据特征、数据集大小以及chunk_size、similarity_top_k等其他变量而有所不同。

下表展示了基于命中率和平均倒数排名（MRR）指标的评估结果：

分析：

按嵌入性能：

OpenAI：展现出顶级性能，尤其是与CohereRerank（命中率0.926966，MRR 0.86573）和bge-reranker-large（命中率0.910112，MRR 0.855805）结合时，表明与重排序工具高度兼容。
bge-large：使用重排序器后性能显著提升，最佳结果来自CohereRerank（命中率0.876404，MRR 0.822753）。
llm-embedder：从重排序中获益匪浅，特别是与CohereRerank（命中率0.882022，MRR 0.830243）搭配时，性能大幅提升。
Cohere：Cohere最新的v3.0嵌入优于v2.0，并且结合原生CohereRerank后，其指标显著提高，拥有0.88764的命中率和0.836049的MRR。
Voyage：初始表现强劲，通过CohereRerank（命中率0.91573，MRR 0.851217）进一步增强，显示出对重排序的高度响应性。
JinaAI：表现非常强劲，与bge-reranker-large（命中率0.938202，MRR 0.868539）和CohereRerank（命中率0.932584，MRR 0.873689）结合时收益显著，表明重排序显著提升了其性能。
Google-PaLM：该模型表现出色，使用CohereRerank时（命中率0.910112，MRR 0.855712）有可测量的提升。这表明重排序对其整体结果有明确的提升作用。

重排序器的影响：

无重排序器：这为每个嵌入提供了基线性能。
bge-reranker-base：通常能提高所有嵌入的命中率和MRR。
bge-reranker-large：此重排序器经常为嵌入提供最高或接近最高的MRR。对于几个嵌入，其性能可与CohereRerank相媲美或超越。
CohereRerank：持续提升所有嵌入的性能，往往提供最佳或接近最佳的结果。

重排序器的必要性：

数据清楚地表明了重排序器在优化搜索结果中的重要性。几乎所有嵌入都能从重排序中受益，显示出更高的命中率和MRR。
重排序器，尤其是CohereRerank，已经证明了它们能够将任何嵌入转变为具有竞争力的嵌入。

整体优势：

当同时考虑命中率和MRR时，OpenAI + CohereRerank和JinaAI-Base + bge-reranker-large/ CohereRerank的组合成为最佳竞争者。
然而，CohereRerank/ bge-reranker-large重排序器在各种嵌入上带来的一致改进，使其成为提升搜索质量的突出选择，无论使用哪种嵌入。

总之，要在命中率和MRR两方面都达到最佳性能，OpenAI或JinaAI-Base嵌入与CohereRerank/bge-reranker-large重排序器的组合最为突出。

请注意，我们的基准测试旨在为您自己的数据提供一个可复现的脚本。尽管如此，请将这些数字视为估计值，并在解读时保持谨慎。

结论：

在这篇博客文章中，我们展示了如何使用各种嵌入和重排序器来评估和增强检索器性能。以下是我们最终的结论。

嵌入：OpenAI和JinaAI-Base嵌入，尤其是与CohereRerank/bge-reranker-large重排序器搭配时，为命中率和MRR设定了黄金标准。
重排序器：重排序器的影响，尤其是CohereRerank/bge-reranker-large，怎么强调都不为过。它们在提高许多嵌入的MRR方面发挥了关键作用，显示了它们在改善搜索结果方面的重要性。
基础是关键：为初始搜索选择合适的嵌入至关重要；如果基本搜索结果不好，即使是最好的重排序器也帮不上什么忙。
协同工作：要充分发挥检索器的优势，找到嵌入和重排序器的正确组合非常重要。这项研究显示了仔细测试并找到最佳配对的重要性。