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