使用 MTEB 进行评估

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

本指南将引导您使用 MTEB 和 SentenceTransformer 模型进行后训练评估。这适用于训练期间的评估,因为这有在公共基准上过拟合的风险。有关训练期间的评估,请参阅训练概述中的评估器部分。要将您的模型完全集成到 MTEB,您可以按照 MTEB 文档中的将模型添加到排行榜指南进行操作。

安装

安装 MTEB 及其依赖项

pip install mteb>=2.0.0

评估

您可以像这样在 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"])
results = mteb.evaluate(model, tasks)

.. 注意:

如果您正在评估现有模型,MTEB 团队建议您使用 mteb.get_model("{model_name}") 而不是 SentenceTransformer。这将加载 MTEB 中实现的模型,通常由模型开发者实现。这确保了可重现的结果,否则结果可能会因归一化、量化、提示或类似原因而异。如果模型未在 mteb 中实现,它将尝试使用 SentenceTransformer 加载模型。

有关可用任务的完整列表,您可以查看 MTEB 任务概述,例如 STS22.v2

您还可以根据任务类型、领域、语言等过滤可用的 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"]
)
results = mteb.evaluate(model, tasks)

最后,在预定义基准上进行评估通常很有价值。例如,要运行 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)")
results = mteb.evaluate(model, benchmark)

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

附加参数

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

...

results = mteb.evaluate(
    model,
    tasks,
    encode_kwargs={"batch_size": 64, "normalize_embeddings": True}
)

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

结果处理

MTEB 将所有结果缓存到磁盘,因此您可以重新运行 mteb.evaluate 而无需重新下载数据集或重新计算分数。默认情况下,这些结果存储在 ~/.cache/mteb 中,可以使用环境变量 MTEB_CACHE 进行配置。但是,您也可以使用 ResultCache 对象管理缓存

import mteb.cache import ResultCache
from sentence_transformers import SentenceTransformer

cache = ResultCache("my_mteb_results_folder")

model = SentenceTransformer("all-MiniLM-L6-v2")
tasks = mteb.get_tasks(tasks=["STS17", "STS22.v2"], languages=["eng"])
results = mteb.evaluate(model, tasks, cache=cache)

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())

您甚至可以通过从结果存储库下载现有结果来避免重新运行已经存在的结果

import mteb.cache import ResultCache

cache = ResultCache("my_mteb_results_folder")
cache.download_from_remote() # will take a while the first time

# will only rerun missing results
results = mteb.evaluate(
    tasks, 
    model, 
    cache=cache,
    overwrite_strategy="only-missing" # default
)

要了解有关如何加载和使用结果的更多信息,请查看 MTEB 文档

排行榜提交

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

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

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

  2. 使用 MTEB 评估您的模型在所需任务上的表现并保存结果。

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

两者合并后,一天后您将能够在官方排行榜上找到您的模型。