模型信息

gte-Qwen2-1.5B-instruct 是 gte（通用文本嵌入）模型系列的最新成员。该模型基于 Qwen2-1.5B 大型语言模型构建，并采用了与 gte-Qwen2-7B-instruct 相同的训练数据和策略。

该模型融合了多项关键技术进展：

• 整合双向注意力机制，增强上下文理解能力。
• 仅对查询端进行指令微调，提升效率。
• 基于跨领域、多场景的多语言文本语料库进行训练，结合弱监督和监督数据，确保模型在多种语言和下游任务中的广泛适用性。

模型规格

• 模型大小：1.5B
• 嵌入维度：1536
• 最大输入令牌数：32k

环境要求

transformers>=4.39.2
flash_attn>=2.5.6

使用方法

使用 Sentence Transformers

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Alibaba-NLP/gte-Qwen2-1.5B-instruct", trust_remote_code=True)
# 如需调整最大长度：
model.max_seq_length = 8192

queries = [
    "女性每日应摄入多少蛋白质",
    "峰会定义",
]
documents = [
    "根据CDC指南，19至70岁女性每日平均蛋白质需求量为46克。但如怀孕或备战马拉松等情况需增加摄入量。详见下表查看每日建议摄入量。",
    "峰会定义（英语学习者版）：1. 山脉最高点；2. 最高级别；3. 政府首脑间的会议。",
]

query_embeddings = model.encode(queries, prompt_name="query")
document_embeddings = model.encode(documents)

scores = (query_embeddings @ document_embeddings.T) * 100
print(scores.tolist())

参考 config_sentence_transformers.json 查看所有预设提示模板。也可使用 model.encode(queries, prompt="指令: ...\n查询: " 自定义提示。

使用 Transformers

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

def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
    left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
    return last_hidden_states[:, -1] if left_padding else last_hidden_states[torch.arange(last_hidden_states.shape[0], device=last_hidden_states.device), attention_mask.sum(dim=1)-1]

def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'指令: {task_description}\n查询: {query}'

task = '给定网页搜索查询，检索回答问题的相关段落'
queries = [
    get_detailed_instruct(task, '女性每日应摄入多少蛋白质'),
    get_detailed_instruct(task, '峰会定义')
]
documents = [
    "根据CDC指南，19至70岁女性每日平均蛋白质需求量为46克...",
    "峰会定义（英语学习者版）：1. 山脉最高点..."
]

tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/gte-Qwen2-1.5B-instruct', trust_remote_code=True)
model = AutoModel.from_pretrained('Alibaba-NLP/gte-Qwen2-1.5B-instruct', trust_remote_code=True)

batch_dict = tokenizer(queries + documents, max_length=8192, padding=True, truncation=True, return_tensors='pt')
outputs = model(**batch_dict)
embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())

使用 infinity_emb

通过 infinity (MIT许可) 部署：

docker run \
--gpus "0" -p "7997":"7997" \
michaelf34/infinity:0.0.68-trt-onnx \
v2 --model-id Alibaba-NLP/gte-Qwen2-1.5B-instruct --revision "refs/pr/20" --dtype bfloat16 --batch-size 16 --device cuda --engine torch --port 7997 --no-bettertransformer

评估结果

MTEB & C-MTEB

使用 scripts/eval_mteb.py 复现 gte-Qwen2-1.5B-instruct 在MTEB(英语)/C-MTEB(中文)的表现：

GTE 模型系列

gte系列包含编码器专用模型(BERT架构)和解码器专用模型(LLM架构)：

商业API服务

除开源模型外，GTE系列还提供阿里云商业API服务：

• 嵌入模型：提供text-embedding-v1/v2/v3三个版本
• 重排序模型：提供gte-rerank服务

注：商业API背后的模型与开源版本不完全相同。

社区支持

微调指导

可使用SWIFT框架微调GTE模型：

pip install ms-swift -U

# 详见：https://swift.readthedocs.io/en/latest/BestPractices/Embedding.html
swift sft --model Alibaba-NLP/gte-Qwen2-1.5B-instruct --task_type embedding ...

引用

如果觉得我们的研究或模型有帮助，请引用：

@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}
}