gte-modernbert-base

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

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

模型概览

• 开发团队：通义实验室，阿里巴巴集团
• 模型类型：文本嵌入
• 主要语言：英语
• 模型大小：149M
• 最大输入长度：8192 tokens
• 输出维度：768

模型列表

使用方法

[!提示] 对于 transformers 和 sentence-transformers，如果您的 GPU 支持，高效的 Flash Attention 2 会在您安装 flash_attn 后自动使用。这不是强制要求的。

pip install flash_attn

配合 transformers 使用：

# Requires transformers>=4.48.0

import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

model_path = "Alibaba-NLP/gte-modernbert-base"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path)

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=8192, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = outputs.last_hidden_state[:, 0]
 
# (Optionally) normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())
# [[42.89073944091797, 71.30911254882812, 33.664554595947266]]

配合 sentence-transformers 使用：

# Requires transformers>=4.48.0
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

model = SentenceTransformer("Alibaba-NLP/gte-modernbert-base")
embeddings = model.encode(input_texts)
print(embeddings.shape)
# (4, 768)

similarities = cos_sim(embeddings[0], embeddings[1:])
print(similarities)
# tensor([[0.4289, 0.7131, 0.3366]])

配合 transformers.js 使用：

// npm i @huggingface/transformers
import { pipeline, matmul } from "@huggingface/transformers";

// Create a feature extraction pipeline
const extractor = await pipeline(
  "feature-extraction",
  "Alibaba-NLP/gte-modernbert-base",
  { dtype: "fp32" }, // Supported options: "fp32", "fp16", "q8", "q4", "q4f16"
);

// Embed queries and documents
const embeddings = await extractor(
  [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms",
  ],
  { pooling: "cls", normalize: true },
);

// Compute similarity scores
const similarities = (await matmul(embeddings.slice([0, 1]), embeddings.slice([1, null]).transpose(1, 0))).mul(100);
console.log(similarities.tolist()); // [[42.89077377319336, 71.30916595458984, 33.66455841064453]]

此外，您还可以通过 Text Embeddings Inference (TEI) 部署 Alibaba-NLP/gte-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-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-modernbert-base

随后，您可以通过 OpenAI 兼容的 v1/embeddings 接口向部署的 API 发送请求（更多关于 OpenAI Embeddings API 的信息）：

curl https://0.0.0.0:8080/v1/embeddings \
  -H "Content-Type: application/json" \
  -d '{
    "input": [
      "what is the capital of China?",
      "how to implement quick sort in python?",
      "Beijing",
      "sorting algorithms"
    ],
    "model": "Alibaba-NLP/gte-modernbert-base",
    "encoding_format": "float"
  }'

训练细节

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（长文档检索）(NDCG@10)

COIR（代码检索任务）(NDCG@10)

BEIR (NDCG@10)

招聘

通义实验室现招聘 研究实习生 和 全职研究员 加入我们的团队。我们寻找对表示学习、基于 LLM 的信息检索、检索增强生成（RAG）以及基于 Agent 的系统充满热情的专业人才。我们的团队位于充满活力的 北京 和 杭州。如果你对工作充满好奇并渴望通过自己的努力产生有意义的影响，我们很乐意收到你的简历。请将你的简历及简要介绍发送至 <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}
}