更新：Jina AI embedding 的 pooling 方法已调整为使用平均 pooling，结果也已相应更新。值得注意的是，使用 JinaAI-v2-base-en 和 bge-reranker-large 现在表现出 0.938202 的命中率 (Hit Rate) 和 0.868539 的平均倒数排名 (MRR - Mean Reciprocal Rank)，而与 CohereRerank 结合则表现出 0.932584 的命中率和 0.873689 的 MRR。

在构建检索增强生成 (RAG) 流水线时，一个关键组件是检索器 (Retriever)。我们可以选择多种 embedding 模型，包括 OpenAI、CohereAI 和开源的 sentence transformers。此外，CohereAI 和 sentence transformers 还提供了多种 reranker。

但在所有这些选项中，我们如何确定实现顶尖检索性能的最佳组合呢？我们如何知道哪种 embedding 模型最适合我们的数据？或者哪种 reranker 对我们的结果提升最大？

在这篇博客文章中，我们将使用 LlamaIndex 中的 检索评估 (Retrieval Evaluation) 模块来快速确定 embedding 和 reranker 模型的最佳组合。让我们开始吧！

我们首先来了解 检索评估 (Retrieval Evaluation) 中可用的指标。

理解检索评估中的指标

为了衡量我们检索系统的有效性，我们主要依赖两个广泛接受的指标：命中率 (Hit Rate) 和 平均倒数排名 (MRR - Mean Reciprocal Rank)。让我们深入研究这些指标，了解它们的意义和工作原理。

命中率 (Hit Rate)

命中率计算的是在排名前 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 页的数据，这部分排除了目录、参考文献和附录。

然后，我们将这些数据解析并转换为节点 (nodes)，节点代表了我们想要检索的数据块。我们使用了 512 作为 chunk_size。

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

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

生成问答对 (Question-Context Pairs)

为了评估目的，我们创建了一个问答对数据集。这个数据集可以看作是我们数据中问题及其对应上下文的集合。为了消除对 embedding (OpenAI/ CohereAI) 和 Reranker (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' 等句子的函数

# function to clean the dataset
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.
    """

    # Extract keys from queries and relevant_docs that need to be removed
    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
    }

    # Filter queries and relevant_docs using dictionary comprehensions
    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
    }

    # Create a new instance of EmbeddingQAFinetuneDataset with the filtered data
    return EmbeddingQAFinetuneDataset(
        queries=filtered_queries,
        corpus=qa_dataset.corpus,
        relevant_docs=filtered_relevant_docs
    )

# filter out pairs with phrases `Here are 2 questions based on provided context`
qa_dataset = filter_qa_dataset(qa_dataset)

自定义检索器 (Custom Retriever)

为了找到最优的检索器，我们结合使用了 embedding 模型和 reranker。首先，我们建立了一个基础的 VectorIndexRetriever。检索到节点后，我们再引入 reranker 来进一步优化结果。值得注意的是，对于本次特定实验，我们将 similarity_top_k 设置为 10，并通过 reranker 选取了前 5 个结果。然而，您可以根据您的具体实验需求自由调整此参数。此处我们展示的是使用 OpenAIEmbedding 的代码，其他 embedding 的代码请参考 notebook。

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,
    ) -> None:
        """Init params."""

        self._vector_retriever = vector_retriever

    def _retrieve(self, query_bundle: QueryBundle) -> 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) -> 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) -> 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) 和命中率 (Hit Rate) 指标

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

结果

我们对各种 embedding 模型和 reranker 进行了测试。以下是我们考虑的模型：

Embedding 模型:

• OpenAI Embedding
• Voyage Embedding
• CohereAI Embedding (v2.0/ v3.0)
• Jina Embeddings (small/ base)
• BAAI/bge-large-en
• Google PaLM Embedding

Reranker:

• CohereAI
• bge-reranker-base
• bge-reranker-large

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

下表展示了基于命中率 (Hit Rate) 和平均倒数排名 (MRR) 指标的评估结果

分析

按 Embedding 分类性能

• OpenAI：展现顶级性能，特别是与 CohereRerank (0.926966 命中率, 0.86573 MRR) 和 bge-reranker-large (0.910112 命中率, 0.855805 MRR) 结合时，表明与 reranking 工具的兼容性很强。
• bge-large：与 reranker 结合后性能显著提升，其中与 CohereRerank 结合获得最佳结果 (0.876404 命中率, 0.822753 MRR)。
• llm-embedder：从 reranking 中获益匪浅，特别是与 CohereRerank 结合时 (0.882022 命中率, 0.830243 MRR)，带来了显著的性能提升。
• Cohere：Cohere 最新的 v3.0 embedding 优于 v2.0，并且通过集成原生的 CohereRerank，其指标显著提升，命中率为 0.88764，MRR 为 0.836049。
• Voyage：初始性能强劲，通过 CohereRerank 进一步放大 (0.91573 命中率, 0.851217 MRR)，表明对 reranking 的响应性很高。
• JinaAI：性能非常强劲，与 bge-reranker-large (0.938202 命中率, 0.868539 MRR) 和 CohereRerank (0.932584 命中率, 0.873689) 结合后性能显著提升，表明 reranking 大幅提高了其性能。
• Google-PaLM：该模型表现出强劲性能，使用 CohereRerank 后获得了可衡量的提升 (0.910112 命中率, 0.855712 MRR)。这表明 reranking 对其整体结果带来了明显的提升。

Reranker 的影响:

• WithoutReranker：这提供了每种 embedding 的基线性能。
• bge-reranker-base：总体上提高了各种 embedding 的命中率和 MRR。
• bge-reranker-large：这种 reranker 经常为 embedding 提供最高或接近最高的 MRR。对于几种 embedding 来说，它的性能与 CohereRerank 不相上下甚至更高。
• CohereRerank：持续提升所有 embedding 的性能，通常提供最佳或接近最佳的结果。

Reranker 的必要性:

• 数据清楚地表明了 reranker 在优化搜索结果方面的重要性。几乎所有 embedding 都从 reranking 中受益，显示出命中率和 MRR 的提升。
• Reranker，特别是 CohereRerank，已经证明了它们能够将任何 embedding 转换为具有竞争力的 embedding 的能力。

整体优势:

• 综合考虑命中率和 MRR，OpenAI + CohereRerank 以及 JinaAI-Base + bge-reranker-large/ CohereRerank 的组合脱颖而出，成为顶级竞争者。
• 然而，CohereRerank/ bge-reranker-large reranker 在各种 embedding 上的持续改进，使它们成为提升搜索质量的杰出选择，无论使用哪种 embedding。

总而言之，为了在命中率和 MRR 上达到最佳性能，OpenAI 或 JinaAI-Base embedding 与 CohereRerank/bge-reranker-large reranker 的组合表现出色。

请注意，我们的基准测试旨在提供一个可重现的脚本供您处理自己的数据。然而，请将这些数字视为估计值，并在解释时谨慎行事。

结论

在这篇博客文章中，我们展示了如何使用各种 embedding 和 reranker 评估和提升检索器性能。以下是我们的最终结论。

• Embeddings：OpenAI 和 JinaAI-Base embedding，特别是与 CohereRerank/bge-reranker-large reranker 搭配使用时，为命中率和 MRR 设定了黄金标准。
• Reranker：reranker 的影响，特别是 CohereRerank/bge-reranker-large 的影响，怎么强调都不为过。它们在提高许多 embedding 的 MRR 方面起着关键作用，显示了它们在改善搜索结果方面的重要性。
• 基础是关键：为初始搜索选择正确的 embedding 至关重要；即使是最好的 reranker，如果基础搜索结果不佳，也无济于事。
• 协同工作：为了充分发挥检索器的作用，找到 embedding 和 reranker 的正确组合非常重要。这项研究表明，仔细测试并找到最佳搭配是多么重要。