gte-reranker-modernbert-base

我们很高兴推出 gte-modernbert 系列模型，这些模型基于最新的 ModernBERT 预训练纯编码器基础模型构建。gte-modernbert 系列模型包括文本嵌入模型和重排序模型。

gte-modernbert 模型在多个文本嵌入和文本检索评估任务中展现出与当前开源社区同规模模型相当的竞争力。评估包括 MTEB、LoCO 和 COIR。

模型概览

• 开发机构：阿里巴巴集团通义实验室
• 模型类型：文本重排序器
• 主要语言：英语
• 模型参数：149M
• 最大输入长度：8192 个 token

模型列表

使用方法

[!TIP] 对于 transformers 和 sentence-transformers，如果您的 GPU 支持，并且已安装 flash_attn，高效的 Flash Attention 2 将自动启用。这不是强制要求的。

pip install flash_attn

配合 transformers 使用：

# 需要 transformers>=4.48.0
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name_or_path = "Alibaba-NLP/gte-reranker-modernbert-base"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
model = AutoModelForSequenceClassification.from_pretrained(
    model_name_or_path,
    torch_dtype=torch.float16,
)
model.eval()

pairs = [
    ["what is the capital of China?", "Beijing"],
    ["how to implement quick sort in python?", "Introduction of quick sort"],
    ["how to implement quick sort in python?", "The weather is nice today"],
]

with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)

# tensor([ 2.1387,  2.4609, -1.6729])

配合 sentence-transformers 使用：

开始之前，请安装 sentence-transformers 库：

pip install sentence-transformers

# 需要 transformers>=4.48.0
from sentence_transformers import CrossEncoder

model = CrossEncoder(
    "Alibaba-NLP/gte-reranker-modernbert-base",
    automodel_args={"torch_dtype": "auto"},
)

pairs = [
    ["what is the capital of China?", "Beijing"],
    ["how to implement quick sort in python?","Introduction of quick sort"],
    ["how to implement quick sort in python?", "The weather is nice today"],
]

scores = model.predict(pairs)
print(scores)
# [0.8945664  0.9213594  0.15742092]
# 注意：Sentence Transformers 默认对输出调用 Softmax，因此分数范围在 [0, 1] 内。

配合 transformers.js 使用：

import {
  AutoTokenizer,
  AutoModelForSequenceClassification,
} from "@huggingface/transformers";

const model_id = "Alibaba-NLP/gte-reranker-modernbert-base";
const model = await AutoModelForSequenceClassification.from_pretrained(
  model_id,
  { dtype: "fp32" }, // 支持的选项："fp32"、"fp16"、"q8"、"q4"、"q4f16"
);
const tokenizer = await AutoTokenizer.from_pretrained(model_id);

const pairs = [
  ["what is the capital of China?", "Beijing"],
  ["how to implement quick sort in python?", "Introduction of quick sort"],
  ["how to implement quick sort in python?", "The weather is nice today"],
];
const inputs = tokenizer(
  pairs.map((x) => x[0]),
  {
    text_pair: pairs.map((x) => x[1]),
    padding: true,
    truncation: true,
  },
);
const { logits } = await model(inputs);
console.log(logits.tolist()); // [[2.138258218765259], [2.4609625339508057], [-1.6775450706481934]]

此外，您还可以通过 Text Embeddings Inference (TEI) 部署 Alibaba-NLP/gte-reranker-modernbert-base：

• CPU

docker run --platform linux/amd64 \
  -p 8080:80 \
  -v $PWD/data:/data \
  --pull always \
  ghcr.io/huggingface/text-embeddings-inference:cpu-1.7 \
  --model-id Alibaba-NLP/gte-reranker-modernbert-base

• GPU

docker run --gpus all \
  -p 8080:80 \
  -v $PWD/data:/data \
  --pull always \
  ghcr.io/huggingface/text-embeddings-inference:1.7 \
  --model-id Alibaba-NLP/gte-reranker-modernbert-base

然后您可以通过 /rerank 路由向部署的 API 发送请求（详见 Text Embeddings Inference OpenAPI 规范）：

curl https://0.0.0.0:8080/rerank \
  -H "Content-Type: application/json" \
  -d '{
    "query": "What is the capital of China?",
    "raw_scores": false,
    "return_text": false,
    "texts": [ "Beijing" ],
    "truncate": true,
    "truncation_direction": "right"
  }'

训练详情

gte-modernbert 系列模型遵循先前 GTE 模型的训练方案，唯一的区别是将预训练语言模型基座从 GTE-MLM 替换为 ModernBert。更多训练详情请参阅我们的论文：mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval

评估结果

MTEB

其他模型的结果来自 MTEB 排行榜。由于 gte-modernbert 系列所有模型的参数规模均小于 1B，我们仅关注 MTEB 排行榜中参数规模小于 1B 的模型结果。

LoCO（长文档检索）

COIR（代码检索任务）

BEIR

招聘

我们在通义实验室有 研究实习生 和 全职研究员 的开放职位。 我们寻求在表示学习、LLM 驱动的信息检索、检索增强生成（RAG）和基于代理的系统方面具有专业知识的热情人才。 我们的团队位于充满活力的 北京 和 杭州。 如果您好奇心强、渴望通过工作产生有意义的影响，我们很乐意收到您的来信。请将您的简历和简要介绍发送至 <a href="mailto:dingkun.ldk@alibaba-inc.com">dingkun.ldk@alibaba-inc.com</a>。

引用

如果您发现我们的论文或模型有帮助，欢迎引用我们。

@inproceedings{zhang2024mgte,
  title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval},
  author={Zhang, Xin and Zhang, Yanzhao and Long, Dingkun and Xie, Wen and Dai, Ziqi and Tang, Jialong and Lin, Huan and Yang, Baosong and Xie, Pengjun and Huang, Fei and others},
  booktitle={Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing: Industry Track},
  pages={1393--1412},
  year={2024}
}

@article{li2023towards,
  title={Towards general text embeddings with multi-stage contrastive learning},
  author={Li, Zehan and Zhang, Xin and Zhang, Yanzhao and Long, Dingkun and Xie, Pengjun and Zhang, Meishan},
  journal={arXiv preprint arXiv:2308.03281},
  year={2023}
}