使用 MTEB 进行评估

大规模文本嵌入基准 (MTEB) 是一个全面的基准测试套件,用于评估嵌入模型在各种 NLP 任务(如检索、分类、聚类、重排序和语义相似性)上的性能。

本指南将引导您使用 MTEB 和 SentenceTransformer 模型进行训练后评估。这 *不* 是为训练期间使用而设计的,因为这有在公开基准上过拟合的风险。对于训练期间的评估,请参阅 训练概览中的评估器部分。要将您的模型完全集成到 MTEB 中,您可以遵循 MTEB 文档中的 将模型添加到排行榜 指南。

安装

安装 MTEB 及其依赖项

pip install mteb

评估

您可以像这样在 MTEB 套件的单个任务上评估您的 SentenceTransformer 模型

import mteb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Example 1: Run a specific single task
tasks = mteb.get_tasks(tasks=["STS22.v2"], languages=["eng"])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder="results/")

有关可用任务的完整列表,您可以查看 MTEB 任务文档

您还可以根据任务类型、领域、语言等筛选可用的 MTEB 任务。例如,以下代码片段在医疗领域的英语检索任务上进行评估

import mteb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Example 2: Run all English retrieval tasks in the medical domain
tasks = mteb.get_tasks(
    task_types=["Retrieval"],
    domains=["Medical"],
    languages=["eng"]
)
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder="results/")

最后,对预定义的基准进行评估通常很有价值。例如,要运行 MTEB(eng, v2) 基准中的所有检索任务

import mteb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

# Example 3: Run the MTEB benchmark for English tasks
benchmark = mteb.get_benchmark("MTEB(eng, v2)")
evaluation = mteb.MTEB(tasks=benchmark)
results = evaluation.run(model, output_folder="results/")

有关支持的基准的完整列表,请访问 MTEB 基准文档

额外参数

在运行评估时,您可以使用 evaluation.run() 上的 encode_kwargs 参数将参数传递给 model.encode()。这允许您自定义嵌入的生成方式,例如设置 batch_sizetruncate_dimnormalize_embeddings。例如

...

results = evaluation.run(
    model,
    verbosity=2,
    output_folder="results/",
    encode_kwargs={"batch_size": 64, "normalize_embeddings": True}
)

此外,您的 SentenceTransformer 模型可能已配置为使用 prompts (提示)。如果这些提示在您模型的配置中定义了,MTEB 会自动检测并使用它们。对于特定于任务或特定于文档/查询的提示,您应该阅读 MTEB 文档中关于 使用提示运行 SentenceTransformer 模型 的部分。

结果处理

MTEB 会将所有结果缓存到磁盘,因此您可以重新运行 evaluation.run(),而无需重新下载数据集或重新计算分数。

import mteb
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("all-MiniLM-L6-v2")

tasks = mteb.get_tasks(tasks=["STS17", "STS22.v2"], languages=["eng"])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder="results/")

for task_results in results:
    # Print the aggregated main scores for each task
    print(f"{task_results.task_name}: {task_results.get_score():.4f} mean {task_results.task.metadata.main_score}")
    """
    STS17: 0.2881 mean cosine_spearman
    STS22.v2: 0.4925 mean cosine_spearman
    """

    # Or e.g. print the individual scores for each split or subset
    print(task_results.only_main_score().to_dict())

排行榜提交

要将您的模型添加到 MTEB 排行榜,您需要遵循 MTEB 文档中的 添加模型 部分。

在此过程中,您需要遵循以下步骤

  1. 将您的模型元数据(名称、语言、参数数量、框架、训练数据集等)添加到 MTEB 仓库

  2. 使用 MTEB 在您期望的任务上评估您的模型并保存结果。

  3. 将您的结果提交到 MTEB 结果仓库

一旦两者都被合并,一天后您就可以在官方排行榜上找到您的模型。