Gte Qwen2 1.5B Instruct

🚀 gte-Qwen2-1.5B-instruct

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

该模型有以下关键改进：

• 集成了双向注意力机制，增强了上下文理解能力。
• 仅在查询端进行指令微调，提高了效率。
• 在广泛的多语言文本语料库上进行全面训练，涵盖了不同领域和场景。这种训练结合了弱监督和监督数据，确保模型适用于多种语言和众多下游任务。

🚀 快速开始

模型信息

• 模型大小：15亿参数
• 嵌入维度：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 = [
    "how much protein should a female eat",
    "summit define",
]
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments.",
]

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="Instruct: ...\nQuery: " 来自定义提示。

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])
    if left_padding:
        return last_hidden_states[:, -1]
    else:
        sequence_lengths = attention_mask.sum(dim=1) - 1
        batch_size = last_hidden_states.shape[0]
        return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths]


def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery: {query}'


# 每个查询都必须附带一个描述任务的单句指令
task = 'Given a web search query, retrieve relevant passages that answer the query'
queries = [
    get_detailed_instruct(task, 'how much protein should a female eat'),
    get_detailed_instruct(task, 'summit define')
]
# 检索文档无需添加指令
documents = [
    "As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
    "Definition of summit for English Language Learners. : 1  the highest point of a mountain : the top of a mountain. : 2  the highest level. : 3  a meeting or series of meetings between the leaders of two or more governments."
]
input_texts = queries + documents

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)

max_length = 8192

# 对输入文本进行分词
batch_dict = tokenizer(input_texts, max_length=max_length, 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

✨ 主要特性

• 双向注意力机制：集成双向注意力机制，丰富了模型的上下文理解能力。
• 指令微调：仅在查询端进行指令微调，提高了效率。
• 多语言训练：在广泛的多语言文本语料库上进行全面训练，确保模型适用于多种语言和众多下游任务。

📦 安装指南

确保你的环境满足以下要求：

transformers>=4.39.2
flash_attn>=2.5.6

你可以使用以下命令安装所需的库：

pip install transformers>=4.39.2 flash_attn>=2.5.6

💻 使用示例

基础用法

上述 Sentence Transformers 和 Transformers 的使用示例展示了如何使用该模型进行文本嵌入和相似度计算。

高级用法

通过 infinity 进行使用，提供了更高效的部署方式。

📚 详细文档

评估

MTEB & C-MTEB

你可以使用 scripts/eval_mteb.py 来复现 gte-Qwen2-1.5B-instruct 在 MTEB（英文）/C-MTEB（中文）上的以下结果：

GTE 模型

GTE 系列模型一直发布两种类型的模型：仅编码器模型（基于 BERT 架构）和解码器模型（基于大语言模型架构）。

云 API 服务

除了开源的 GTE 系列模型外，GTE 系列模型还可以在阿里云上作为商业 API 服务使用。

• 嵌入模型：提供三种版本的文本嵌入模型：text-embedding-v1/v2/v3，其中 v3 是最新的 API 服务。
• 重排序模型：提供 gte-rerank 模型服务。

请注意，商业 API 背后的模型与开源模型并不完全相同。

社区支持

微调

GTE 模型可以使用第三方框架 SWIFT 进行微调。

pip install ms-swift -U

# 查看：https://swift.readthedocs.io/en/latest/BestPractices/Embedding.html
nproc_per_node=8
NPROC_PER_NODE=$nproc_per_node \
USE_HF=1 \
swift sft \
    --model Alibaba-NLP/gte-Qwen2-1.5B-instruct \
    --train_type lora \
    --dataset 'sentence-transformers/stsb' \
    --torch_dtype bfloat16 \
    --num_train_epochs 10 \
    --per_device_train_batch_size 2 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps $(expr 64 / $nproc_per_node) \
    --eval_steps 100 \
    --save_steps 100 \
    --eval_strategy steps \
    --use_chat_template false \
    --save_total_limit 5 \
    --logging_steps 5 \
    --output_dir output \
    --warmup_ratio 0.05 \
    --learning_rate 5e-6 \
    --deepspeed zero3 \
    --dataloader_num_workers 4 \
    --task_type embedding \
    --loss_type cosine_similarity \
    --dataloader_drop_last true

🔧 技术细节

• 双向注意力机制：模型集成了双向注意力机制，能够更好地捕捉文本的上下文信息，从而提高文本嵌入的质量。
• 指令微调：仅在查询端进行指令微调，减少了计算量，提高了效率。
• 多语言训练：在广泛的多语言文本语料库上进行训练，使用了弱监督和监督数据，确保模型在多种语言和下游任务上都能有良好的表现。

📄 许可证

如果你发现我们的论文或模型有帮助，请考虑引用：

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