数据集选择

可以通过提供数据集列表来选择数据集，也可以：

• 按任务类型选择（例如 "Clustering" 或 "Classification"）

tasks = mteb.get_tasks(task_types=["Clustering", "Retrieval"]) # 仅选择聚类和检索任务

• 按类别选择，例如 "s2s"（句子到句子）或 "p2p"（段落到段落）

tasks = mteb.get_tasks(categories=["s2s", "p2p"]) # 仅选择句子到句子和段落到段落数据集

• 按语言选择

tasks = mteb.get_tasks(languages=["eng", "deu"]) # 仅选择包含 "eng" 或 "deu"（ISO 639-3 代码）的数据集

您还可以为多语言/跨语言任务指定要加载的语言，如下所示：

import mteb

tasks = [
    mteb.get_task("AmazonReviewsClassification", languages = ["eng", "fra"]),
    mteb.get_task("BUCCBitextMining", languages = ["deu"]), # 所有包含 "deu" 的子集
]

# 或者您可以像这样选择特定的 huggingface 子集：
from mteb.tasks import AmazonReviewsClassification, BUCCBitextMining

evaluation = mteb.MTEB(tasks=[
        AmazonReviewsClassification(hf_subsets=["en", "fr"]) # 仅加载 Amazon Reviews 的 "en" 和 "fr" 子集
        BUCCBitextMining(hf_subsets=["de-en"]), # 仅加载 BUCC 的 "de-en" 子集
])
# 有关 HF 子集的示例，请参见数据集查看器中的 "Subset"：https://huggingface.co/datasets/mteb/bucc-bitext-mining

还有一些预设可用于某些任务集合，例如选择构成 "Overall MTEB English leaderboard" 的 56 个英语数据集：

from mteb import MTEB_MAIN_EN
evaluation = mteb.MTEB(tasks=MTEB_MAIN_EN, task_langs=["en"])

传递 encode 参数

要向模型的 encode 函数传递参数，您可以使用编码关键字参数（encode_kwargs）：

evaluation.run(model, encode_kwargs={"batch_size": 32}

选择评估分割

您可以通过以下方式仅在所有任务的 test 分割上进行评估：

evaluation.run(model, eval_splits=["test"])

请注意，公开排行榜使用所有数据集的测试分割，但 MSMARCO 除外，它使用 "dev" 分割。

使用自定义模型

模型应实现以下接口，实现一个 encode 函数，该函数接受句子列表作为输入，并返回嵌入列表（嵌入可以是 np.array、torch.tensor 等）。为获得灵感，您可以查看 mteb/mtebscripts 仓库，该仓库用于通过 SLURM 脚本运行论文中的各种模型。

class MyModel():
    def encode(
        self, sentences: list[str], **kwargs: Any
    ) -> torch.Tensor | np.ndarray:
        """使用编码器对给定句子进行编码。

        参数：
            sentences: 要编码的句子。
            **kwargs: 传递给编码器的其他参数。

        返回：
            编码后的句子。
        """
        pass

model = MyModel()
tasks = mteb.get_task("Banking77Classification")
evaluation = MTEB(tasks=tasks)
evaluation.run(model)

如果您想在评估检索或重排任务时为查询和语料库使用不同的编码函数，可以添加单独的 encode_queries 和 encode_corpus 方法。如果这些方法存在，它们将自动用于这些任务。您可以参考 mteb/evaluation/evaluators/RetrievalEvaluator.py 中的 DRESModel 作为这些函数的示例。

class MyModel():
    def encode_queries(self, queries: list[str], **kwargs) -> list[np.ndarray] | list[torch.Tensor]:
        """
        返回给定句子的嵌入列表。
        参数：
            queries: 要编码的句子列表

        返回：
            给定句子的嵌入列表
        """
        pass

    def encode_corpus(self, corpus: list[str] | list[dict[str, str]], **kwargs) -> list[np.ndarray] | list[torch.Tensor]:
        """
        返回给定句子的嵌入列表。
        参数：
            corpus: 要编码的句子列表
                或包含 "title" 和 "text" 键的字典列表

        返回：
            给定句子的嵌入列表
        """
        pass

在自定义数据集上评估

要在自定义任务上进行评估，你可以在自定义任务上运行以下代码。关于如何在MTEB中创建新任务，请参见如何添加新任务。

from mteb import MTEB
from mteb.abstasks.AbsTaskReranking import AbsTaskReranking
from sentence_transformers import SentenceTransformer


class MyCustomTask(AbsTaskReranking):
    ...

model = SentenceTransformer("average_word_embeddings_komninos")
evaluation = MTEB(tasks=[MyCustomTask()])
evaluation.run(model)