跳转到主要内容

大规模文本嵌入基准

项目描述

大规模文本嵌入基准

GitHub release GitHub release License Downloads

安装 | 使用 | 排行榜 | 文档 | 引用

安装

pip install mteb

示例用法

  • 使用Python脚本
import mteb
from sentence_transformers import SentenceTransformer

# Define the sentence-transformers model name
model_name = "average_word_embeddings_komninos"
# or directly from 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

# if nothing is specified default to saving the results in the results/{model_name} folder
  • 可以通过仅使用自定义编码函数来并行使用多个GPU,该函数将输入分配给多个GPU,例如这里这里

使用文档

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


任务选择

任务选择

可以通过提供数据集列表来选择任务,也可以

  • 根据其任务(例如,“聚类”或“分类”)
tasks = mteb.get_tasks(task_types=["Clustering", "Retrieval"]) # Only select clustering and retrieval tasks
  • 根据其类别,例如“s2s”(句子到句子)或“p2p”(段落到段落)
tasks = mteb.get_tasks(categories=["s2s", "p2p"]) # Only select sentence2sentence and paragraph2paragraph datasets
  • 根据其语言
tasks = mteb.get_tasks(languages=["eng", "deu"]) # Only select datasets which contain "eng" or "deu" (iso 639-3 codes)

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

import mteb

tasks = [
    mteb.get_task("AmazonReviewsClassification", languages = ["eng", "fra"]),
    mteb.get_task("BUCCBitextMining", languages = ["deu"]), # all subsets containing "deu"
]

# or you can select specific huggingface subsets like this:
from mteb.tasks import AmazonReviewsClassification, BUCCBitextMining

evaluation = mteb.MTEB(tasks=[
        AmazonReviewsClassification(hf_subsets=["en", "fr"]) # Only load "en" and "fr" subsets of Amazon Reviews
        BUCCBitextMining(hf_subsets=["de-en"]), # Only load "de-en" subset of BUCC
])
# for an example of a HF subset see "Subset" in the dataset viewer at: https://hugging-face.cn/datasets/mteb/bucc-bitext-mining
运行基准测试

运行基准测试

mteb附带一组预定义的基准测试。这些可以通过get_benchmark获取并按照类似其他任务集的方式运行。例如,要选择构成“总体MTEB英语排行榜”的56个英语数据集

import mteb
benchmark = mteb.get_benchmark("MTEB(eng)")
evaluation = mteb.MTEB(tasks=benchmark)

指定的基准不仅包含任务列表,还包含要运行的分割和语言。要获取所有可用基准的概述,请运行

import mteb
benchmarks = mteb.get_benchmarks()

通常我们使用基准的命名方案为MTEB(*),其中“*”表示基准的目标。在语言的情况下,我们使用三位语言代码。对于大型语言组,我们使用组表示法,例如,MTEB(Scandinavian)表示斯堪的纳维亚语。在MTEB中实现的如CoIR之类的外部基准使用其原始名称。当使用MTEB中的基准时,请引用mteb以及您可以通过以下方式访问的基准的引用

benchmark.citation
传递encode参数

传递encode参数

要传递参数给模型的encode函数,您可以使用编码关键字参数(encode_kwargs

evaluation.run(model, encode_kwargs={"batch_size": 32}
选择评估分割

选择评估分割

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

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

请注意,公共排行榜使用除MSMARCO外的所有数据集的测试分割,“dev”分割用于MSMARCO。

使用自定义模型

使用自定义模型

模型应实现以下接口,实现一个encode函数,该函数接收句子列表作为输入,并返回嵌入列表(嵌入可以是np.arraytorch.tensor等)。作为灵感,您可以查看用于通过SLURM脚本来运行各种模型的mteb/mtebscripts存储库

class MyModel():
    def encode(
        self, sentences: list[str], **kwargs: Any
    ) -> torch.Tensor | np.ndarray:
        """Encodes the given sentences using the encoder.

        Args:
            sentences: The sentences to encode.
            **kwargs: Additional arguments to pass to the encoder.

        Returns:
            The encoded sentences.
        """
        pass

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

如果您在检索或重排序任务中希望使用不同的编码函数对查询和语料库进行评估,可以为encode_queriesencode_corpus添加单独的方法。如果存在这些方法,它们将自动用于这些任务。您可以通过查看mteb/evaluation/evaluators/RetrievalEvaluator.py中的DRESModel来获取这些函数的示例。

class MyModel():
    def encode_queries(self, queries: list[str], **kwargs) -> list[np.ndarray] | list[torch.Tensor]:
        """
        Returns a list of embeddings for the given sentences.
        Args:
            queries: List of sentences to encode

        Returns:
            List of embeddings for the given sentences
        """
        pass

    def encode_corpus(self, corpus: list[str] | list[dict[str, str]], **kwargs) -> list[np.ndarray] | list[torch.Tensor]:
        """
        Returns a list of embeddings for the given sentences.
        Args:
            corpus: List of sentences to encode
                or list of dictionaries with keys "title" and "text"

        Returns:
            List of embeddings for the given sentences
        """
        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)
使用交叉编码器进行重排序

使用交叉编码器进行重排序

要使用交叉编码器进行重排序,您可以直接使用SentenceTransformers中的CrossEncoder。以下代码展示了具有两个阶段运行的示例,第二阶段读取第一阶段保存的结果。

from mteb import MTEB
import mteb
from sentence_transformers import CrossEncoder, SentenceTransformer

cross_encoder = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2")
dual_encoder = SentenceTransformer("all-MiniLM-L6-v2")

tasks = mteb.get_tasks(tasks=["NFCorpus"], languages=["eng"])

subset = "default" # subset name used in the NFCorpus dataset
eval_splits = ["test"]

evaluation = MTEB(tasks=tasks)
evaluation.run(
    dual_encoder,
    eval_splits=eval_splits,
    save_predictions=True,
    output_folder="results/stage1",
)
evaluation.run(
    cross_encoder,
    eval_splits=eval_splits,
    top_k=5,
    save_predictions=True,
    output_folder="results/stage2",
    previous_results=f"results/stage1/NFCorpus_{subset}_predictions.json",
)
保存检索任务的预测结果

保存检索任务的预测结果

要保存检索任务的预测结果,请在CLI中添加--save_predictions标志或在运行方法中将save_predictions=True设置为真。文件名将为"{task_name}_{subset}_predictions.json"格式。

Python

from mteb import MTEB
import mteb
from sentence_transformers import SentenceTransformer

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

tasks = mteb.get_tasks( tasks=["NFCorpus"], languages=["eng"])

evaluation = MTEB(tasks=tasks)
evaluation.run(
    model,
    eval_splits=["test"],
    save_predictions=True,
    output_folder="results",
)

命令行界面

mteb run -t NFCorpus -m all-MiniLM-L6-v2 --output_folder results --save_predictions
从结果存储库获取结果

从结果存储库获取结果

已经对MTEB内可用的任务运行了多个模型。这些结果可在结果存储库中找到。

为了使结果更容易访问,我们为从存储库中检索而设计了自定义功能。例如,如果您正在选择用于法律文件的法语和英语检索任务的最佳模型,可以使用以下代码检索相关任务并使用以下代码创建结果数据框

import mteb
from mteb.task_selection import results_to_dataframe

tasks = mteb.get_tasks(
    task_types=["Retrieval"], languages=["eng", "fra"], domains=["Legal"]
)

model_names = [
    "GritLM/GritLM-7B",
    "intfloat/multilingual-e5-small",
    "intfloat/multilingual-e5-base",
    "intfloat/multilingual-e5-large",
]
models = [mteb.get_model_meta(name) for name in model_names]

results = mteb.load_results(models=models, tasks=tasks)

df = results_to_dataframe(results)

文档

文档
📋 任务  可用的任务概述
📐 基准 可用的基准概述
📈 排行榜 基准的交互式排行榜
🤖 添加模型 有关如何将模型提交到排行榜的相关信息
👩‍🔬 可重复工作流程 有关如何使用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并集成新数据集的出色工作

对于使用MTEB进行基准测试的作品,您可以在排行榜上找到它们。

项目详情


发布历史 发布通知 | RSS源

下载文件

下载适合您平台的文件。如果您不确定选择哪个,请了解有关安装包的更多信息。

源分布

mteb-1.15.2.tar.gz (687.9 kB 查看散列值)

上传于 来源

构建分发

mteb-1.15.2-py3-none-any.whl (1.1 MB 查看哈希值)

上传于 Python 3

支持