mteb

<h1 align="center">大规模文本嵌入基准测试</h1> <p align="center"> <a href="https://github.com/embeddings-benchmark/mteb/releases"> <img alt="GitHub 发布版本" src="https://yellow-cdn.veclightyear.com/835a84d5/7fc7b9e2-85e9-4e0a-bb04-e10c407bd304.svg"> </a> <a href="https://arxiv.org/abs/2210.07316"> <img alt="GitHub 发布版本" src="https://yellow-cdn.veclightyear.com/835a84d5/e1bc5618-0fe9-4a9d-8a94-bd9fdf5ac7f5.svg"> </a> <a href="https://github.com/embeddings-benchmark/mteb/blob/master/LICENSE"> <img alt="许可证" src="https://yellow-cdn.veclightyear.com/835a84d5/4ad028de-673a-4c34-9ef4-d9610f825686.svg?color=green"> </a> <a href="https://pepy.tech/project/mteb"> <img alt="下载量" src="https://static.pepy.tech/personalized-badge/mteb?period=total&units=international_system&left_color=grey&right_color=orange&left_text=Downloads"> </a> </p> <h4 align="center"> <p> <a href="#installation">安装</a> | <a href="#usage">使用</a> | <a href="https://huggingface.co/spaces/mteb/leaderboard">排行榜</a> | <a href="#documentation">文档</a> | <a href="#citing">引用</a> <p> </h4> <h3 align="center"> <a href="https://huggingface.co/spaces/mteb/leaderboard"><img style="float: middle; padding: 10px 10px 10px 10px;" width="60" height="55" src="https://yellow-cdn.veclightyear.com/835a84d5/680fc015-da20-401d-aa4b-8888a668c3b7.png" /></a> </h3>

安装

pip install mteb

使用

使用 Python 脚本（更多示例请参见 scripts/run_mteb_english.py 和 mteb/mtebscripts）：

import mteb
from sentence_transformers import SentenceTransformer

# 定义 sentence-transformers 模型名称
model_name = "average_word_embeddings_komninos"
# 或直接从 huggingface 加载：
# model_name = "sentence-transformers/all-MiniLM-L6-v2"

model = SentenceTransformer(model_name)
tasks = mteb.get_tasks(tasks=["Banking77Classification"])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder=f"results/{model_name}")

使用命令行界面

mteb available_tasks

mteb run -m sentence-transformers/all-MiniLM-L6-v2 \
    -t Banking77Classification  \
    --verbosity 3

# 如果未指定，默认将结果保存在 results/{model_name} 文件夹中

使用多个 GPU 并行处理可以通过自定义编码函数来实现，该函数将输入分配给多个 GPU，例如这里或这里。

高级用法

点击下面的每个部分查看详细信息。

<br /> <details> <summary> 数据集选择 </summary>

数据集选择

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

按任务类型选择（例如 "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"])

</details> <details> <summary> 传递 `encode` 参数 </summary>

传递 `encode` 参数

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

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

</details> <details> <summary> 选择评估分割 </summary>

选择评估分割

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

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

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

</details> <details> <summary> 使用自定义模型 </summary>

使用自定义模型

模型应实现以下接口，实现一个 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

</details> <details> <summary> 在自定义数据集上评估 </summary>

在自定义数据集上评估

要在自定义任务上进行评估，你可以在自定义任务上运行以下代码。关于如何在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)

</details> <br />

文档

文档
📋 任务	可用任务概览
📈 排行榜	基准测试的交互式排行榜
🤖 添加模型	有关如何向排行榜提交模型的信息
👩‍🔬 可复现的工作流程	有关如何使用MTEB复现和创建可复现工作流程的信息
👩‍💻 添加数据集	如何向MTEB添加新任务/数据集
👩‍💻 添加排行榜标签页	如何向MTEB添加新的排行榜标签页
🤝 贡献	如何为MTEB做出贡献并进行开发设置
🌐 MMTEB	一项开源努力，旨在扩展MTEB以涵盖广泛的语言

引用

MTEB在"MTEB: Massive Text Embedding Benchmark"中被介绍，欢迎引用：

@article{muennighoff2022mteb,
  doi = {10.48550/ARXIV.2210.07316},
  url = {https://arxiv.org/abs/2210.07316},
  author = {Muennighoff, Niklas and Tazi, Nouamane and Magne, Lo{\"\i}c and Reimers, Nils},
  title = {MTEB: Massive Text Embedding Benchmark},
  publisher = {arXiv},
  journal={arXiv preprint arXiv:2210.07316},  
  year = {2022}
}

你可能还想阅读并引用那些扩展MTEB和集成新数据集的出色工作：

Shitao Xiao, Zheng Liu, Peitian Zhang, Niklas Muennighoff. "C-Pack: Packaged Resources To Advance General Chinese Embedding" arXiv 2023
Michael Günther, Jackmin Ong, Isabelle Mohr, Alaeddine Abdessalem, Tanguy Abel, Mohammad Kalim Akram, Susana Guzman, Georgios Mastrapas, Saba Sturua, Bo Wang, Maximilian Werk, Nan Wang, Han Xiao. "Jina Embeddings 2: 8192-Token General-Purpose Text Embeddings for Long Documents" arXiv 2023
Silvan Wehrli, Bert Arnrich, Christopher Irrgang. "German Text Embedding Clustering Benchmark" arXiv 2024
Orion Weller, Benjamin Chang, Sean MacAvaney, Kyle Lo, Arman Cohan, Benjamin Van Durme, Dawn Lawrie, Luca Soldaini. "FollowIR: Evaluating and Teaching Information Retrieval Models to Follow Instructions" arXiv 2024
Dawei Zhu, Liang Wang, Nan Yang, Yifan Song, Wenhao Wu, Furu Wei, Sujian Li. "LongEmbed: Extending Embedding Models for Long Context Retrieval" arXiv 2024
Kenneth Enevoldsen, Márton Kardos, Niklas Muennighoff, Kristoffer Laigaard Nielbo. "The Scandinavian Embedding Benchmarks: Comprehensive Assessment of Multilingual and Monolingual Text Embedding" arXiv 2024