gte-base-en-v1.5

我们推出 gte-v1.5 系列，这是升级版的 gte 嵌入模型，支持最高 8192 的上下文长度，同时进一步提升了模型性能。 该模型基于 transformer++ 编码器 骨干网络（BERT + RoPE + GLU）构建。

gte-v1.5 系列在同规模模型的 MTEB 基准测试中实现了最先进水平，并在 LoCo 长上下文检索测试中表现优异（请参阅评估）。

我们还推出了 gte-Qwen1.5-7B-instruct， 这是一款 SOTA 指令微调多语言嵌入模型，在 MTEB 排名第 2，在 C-MTEB 排名第 1。

<!-- 提供更长的模型摘要。 -->

• 开发机构： 阿里巴巴集团智能计算研究所
• 模型类型： 文本嵌入
• 论文： mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval

<!-- - 演示 [可选]: [需要更多信息] -->

模型列表

如何开始使用模型

使用以下代码开始使用模型。

# 需要 transformers>=4.36.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-base-en-v1.5'
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True)

# 对输入文本进行分词
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]
 
# （可选）对嵌入进行归一化
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

建议安装 xformers 并启用 unpadding 以加速，参阅 enable-unpadding-and-xformers。

配合 sentence-transformers 使用：

# 需要 sentence_transformers>=2.7.0

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

sentences = ['That is a happy person', 'That is a very happy person']

model = SentenceTransformer('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

配合 transformers.js 使用：

// npm i @xenova/transformers
import { pipeline, dot } from '@xenova/transformers';

// 创建特征提取管道
const extractor = await pipeline('feature-extraction', 'Alibaba-NLP/gte-base-en-v1.5', {
    quantized: false, // 注释此行以使用量化版本
});

// 生成句子嵌入
const sentences = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]
const output = await extractor(sentences, { normalize: true, pooling: 'cls' });

// 计算相似度分数
const [source_embeddings, ...document_embeddings ] = output.tolist();
const similarities = document_embeddings.map(x => 100 * dot(source_embeddings, x));
console.log(similarities); // [34.504930869007296, 64.03973265120138, 19.520042686034362]

配合 Infinity 使用： Infinity 是一个 MIT 许可的 OpenAI 兼容部署服务器。

docker run --gpus all -v $PWD/data:/app/.cache -p "7997":"7997" \
michaelf34/infinity:0.0.68 \
v2 --model-id Alibaba-NLP/gte-base-en-v1.5 --revision "4c742dc2b781e4ab062a4a77f4f7cbad4bdee970" --dtype bfloat16 --batch-size 32 --device cuda --engine torch --port 7997

训练详情

训练数据

• 掩码语言建模（MLM）：c4-en
• 弱监督对比预训练（CPT）：GTE 预训练数据
• 有监督对比微调：GTE 微调数据

训练流程

为使骨干模型支持 8192 的上下文长度，我们采用了多阶段训练策略。 模型首先在较短的序列长度上进行初步 MLM 预训练。 然后，我们对数据进行重采样，减少短文本的比例，继续进行 MLM 预训练。

整个训练过程如下：

• MLM-2048：lr 5e-4，mlm_probability 0.3，batch_size 4096，num_steps 70000，rope_base 10000
• MLM-8192：lr 5e-5，mlm_probability 0.3，batch_size 1024，num_steps 20000，rope_base 500000
• CPT：max_len 512，lr 2e-4，batch_size 32768，num_steps 100000
• 微调：待定

评估

MTEB

其他模型的结果来自 MTEB 排行榜。

gte 评估设置：mteb==1.2.0, fp16 auto mix precision, max_length=8192，并设置 ntk 缩放因子为 2（相当于 rope_base * 2）。

LoCo

引用

如果您发现我们的论文或模型有所帮助，请考虑按以下方式引用：

@misc{zhang2024mgte,
  title={mGTE: Generalized Long-Context Text Representation and Reranking Models for Multilingual Text Retrieval}, 
  author={Xin Zhang and Yanzhao Zhang and Dingkun Long and Wen Xie and Ziqi Dai and Jialong Tang and Huan Lin and Baosong Yang and Pengjun Xie and Fei Huang and Meishan Zhang and Wenjie Li and Min Zhang},
  year={2024},
  eprint={2407.19669},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2407.19669}, 
}
@misc{li2023gte,
  title={Towards General Text Embeddings with Multi-stage Contrastive Learning}, 
  author={Zehan Li and Xin Zhang and Yanzhao Zhang and Dingkun Long and Pengjun Xie and Meishan Zhang},
  year={2023},
  eprint={2308.03281},
  archivePrefix={arXiv},
  primaryClass={cs.CL},
  url={https://arxiv.org/abs/2308.03281}, 
}