如何为 RAG 应用选择最佳 Embedding 模型

向量 Embedding 是目前检索增强生成（RAG）应用程序的核心。它们捕获数据对象（如文本，图像等）的语义信息，并以数字数组表示。在时下的生成式 AI 应用中，这些向量 Embedding 通常由 Embedding 模型生成。如何为 RAG 应用程序选择合适的 Embedding 模型呢？总体来说，这取决于具体用例以及具体需求。接下来，让我们拆分步骤来分别来看。

确定具体用例

我们基于 RAG 应用程序需求考虑以下问题：

• 首先，通用模型是否足以满足需求？

• 其次，是否有特定的需求？例如模态（如仅文本或图片，关于多模态 Embedding 选择参见 如何选择合适的 Embedding 模型）、特定领域（如法律、医学等）

在大多数情况下，针对所需模态，通常选择通用模型。

选择通用模型

如何选择通用模型呢？HuggingFace 中Massive Text Embedding Benchmark（MTEB）排行榜罗列了当前各种专有和开源文本 Embedding 模型，对于每个 Embedding 模型，MTEB 列出了各种指标，包括模型参数、内存、Embedding 维度、最大 token 数量，以及其在检索、摘要等任务中的得分。

在为 RAG 应用程序选择 Embedding 模型的时候，需要考虑以下因素：

任务：在 MTEB 排行榜顶部，我们会看到各种任务选项卡。对于一个 RAG 应用程序，我们可能需要更关注“检索”任务，我们可以选择 Retrial 这个选项卡。

语言：基于 RAG 应用的数据集语言来选择对应语言的 Embedding 模型。

得分：表示模型在特定基准数据集或多个基准数据集上的表现。根据任务的不同，使用不同的评估指标。通常，这些指标的取值范围在 0 到 1 之间，值越高表示性能越好。

模型规模和内存用量：这些指标让我们了解模型运行时所需的计算资源。虽然检索性能随着模型规模的增加而提高，但需要注意的是，模型规模也直接影响延迟。此外，较大的模型可能会过拟合，其泛化性能低，从而在生产中表现不佳。因此，我们需要在生产环境中寻求性能和延迟之间的平衡。一般来说，我们可以先从一个小的、轻量级的模型开始，先快速构建 RAG 应用。在应用的基础流程正常运行之后，我们可以切换到更大、性能更高的模型，来对应用做进一步优化。

Embedding 维度：这是 Embedding 向量的长度。虽然更大的 Embedding 维度可以捕获数据中更细微的细节，但是效果不一定是最佳。例如，我们是否真的需要 8192 维度来处理文档数据呢？可能不需要。另一方面，较小的 Embedding 维度提供了更快的推理速度，并且在存储和内存方面更高效。因此，我们需要在捕获数据内容和执行效率之间寻求良好的平衡。

最大 token 数量：表示单个 Embedding 的最大 token 数量。对于常见的 RAG 应用程序，Embedding 较好的分块大小通常是单个段落，此时，最大 token 为 512 的 Embedding 模型应该足够。然而，在某些特殊情况下，我们可能需要 token 数量更大的模型来处理更长的文本。

在 RAG 应用中评估模型

虽然我们可以从 MTEB 排行榜找到通用模型，但我们需要谨慎对待其结果。谨记，这些结果是模型的自我报告，有可能某些模型产生了夸大其性能的分数，因为它们可能在训练数据中包含了 MTEB 数据集，毕竟这些是公开的数据集。另外，模型使用基准的数据集可能无法准确代表我们应用里使用的数据。因此，我们需要在自己的数据集上评估 Embedding 模型。

数据集

我们可以从 RAG 应用使用的数据中，生成一个小型标记数据集。我们以如下数据集为例。

创建 Embedding

接下来，我们采用pymilvus[model]对于上述数据集生成相应的向量 Embedding。关于 pymilvus[model] 使用，参见https://milvus.io/blog/introducing-pymilvus-integrations-with-embedding-models.md

def gen_embedding(model_name):    openai_ef = model.dense.OpenAIEmbeddingFunction(        model_name=model_name,         api_key=os.environ["OPENAI_API_KEY"]    )        docs_embeddings = openai_ef.encode_documents(df['description'].tolist())
    return docs_embeddings, openai_ef

然后，把生成的 Embedding 存入到 Milvus 的 collection。

def save_embedding(docs_embeddings, collection_name, dim):    data = [        {"id": i, "vector": docs_embeddings[i].data, "text": row.language}        for i, row in df.iterrows()    ]        if milvus_client.has_collection(collection_name=collection_name):        milvus_client.drop_collection(collection_name=collection_name)    milvus_client.create_collection(collection_name=collection_name, dimension=dim)        res = milvus_client.insert(collection_name=collection_name, data=data)

查询

我们定义查询函数，方便对于向量 Embedding 进行召回。

def query_results(query, collection_name, openai_ef):    query_embeddings = openai_ef.encode_queries(query)    res = milvus_client.search(        collection_name=collection_name,         data=query_embeddings,         limit=4,          output_fields=["text"],    )
    result = {}    for items in res:        for item in items:            result[item.get("entity").get("text")] = item.get('distance')
    return result

评估 Embedding 模型性能

我们采用 OpenAI 的两个 Embedding 模型，text-embedding-3-small 和 text-embedding-3-large，对于如下两个查询进行比较。有很多评估指标，例如准确率、召回率、MRR、MAP 等。在这里，我们采用准确率和召回率。

准确率（Precision） 评估检索结果中的真正相关内容的占比，即返回的结果中有多少与搜索查询相关。

Precision = TP / (TP + FP)

其中，检索结果中与查询真正相关的内容 True Positives（TP）, 而 False Positives（FP） 指的是检索结果中不相关的内容。

召回率 （Recall）评估从整个数据集中成功检索到相关内容的数量。

Recall = TP / (TP + FN)

其中，False Negatives （FN） 指的是所有未包含在最终结果集中的相关项目

对于这两个概念更详细的解释，参见 https://zilliz.com/learn/information-retrieval-metrics

查询 1：auto garbage collection

相关项： Java， Python， JavaScript， Golang

查询 2： suite for web backend server development

相关项： Java, JavaScript, PHP, Python (答案包含主观判断)

在这两个查询中，我们通过准确率和召回率对比了两个 Embedding 模型 text-embedding-3-small 和 text-embedding-3-large 。我们可以以此为起点，增加数据集中数据对象的数量以及查询的数量，如此才能更有效地评估 Embedding 模型。

总结

在检索增强生成（RAG）应用中，选择合适的向量 Embedding 模型至关重要。本文阐述了从实际业务需求出发，从 MTEB 选择通用模型之后，采用准确率和召回率对于模型基于特定业务的数据集进行测试，从而选择最合适的 Embedding 模型，进而有效地提高 RAG 应用的召回准确率。

完整代码可通过链接获取：

链接: https://pan.baidu.com/s/1CqewAUpWxrnbUuDUIun-Lw?pwd=1234 提取码: 1234

作者介绍

Zilliz 黄金写手：臧伟