大规模文本嵌入基准
项目描述
大规模文本嵌入基准
安装
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
使用文档
点击下面的每个部分以查看详细信息。
任务选择
任务选择
可以通过提供数据集列表来选择任务,也可以
- 根据其任务(例如,“聚类”或“分类”)
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.array
、torch.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_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]:
"""
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并集成新数据集的出色工作
- 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)。《斯堪的纳维亚嵌入基准:多语言和单语言文本嵌入的全面评估》arXiv 2024
对于使用MTEB进行基准测试的作品,您可以在排行榜上找到它们。
项目详情
下载文件
下载适合您平台的文件。如果您不确定选择哪个,请了解有关安装包的更多信息。
源分布
构建分发
mteb-1.15.2.tar.gz的哈希值
算法 | 哈希摘要 | |
---|---|---|
SHA256 | bc616af026275bad542c7063d2178b8f6135e96b97e73136c2cde29568df6683 |
|
MD5 | 00f6a5d19e8f222569534e2859c186b0 |
|
BLAKE2b-256 | 91426be4a3a3849e46e64df21e634b2431232fad383f5530eead150968b7ce45 |
mteb-1.15.2-py3-none-any.whl的哈希值
算法 | 哈希摘要 | |
---|---|---|
SHA256 | 4ff06dc76aea7a46ff225db67802fbb64bac23a93de9e30fb3fb308701a0535e |
|
MD5 | b34698b67d90f6e3701821f6a060a12b |
|
BLAKE2b-256 | 60abfb3fecdaaa6a093b005e6c7c8f95e091d83f558451216bd5bd789a1625a7 |