🚀 gte-Qwen2-1.5B-instruct

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

该模型包含了多项关键改进：

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

🚀 快速开始

模型信息

环境要求

transformers>=4.39.2
flash_attn>=2.5.6

✨ 主要特性

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

📦 安装指南

根据上述环境要求，使用以下命令安装所需的库：

pip install 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())

📚 详细文档

评估

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 背后的模型与开源模型并不完全相同。

📄 许可证

本项目采用 Apache 2.0 许可证。

🔧 技术细节

该模型基于 Qwen2-1.5B 大语言模型构建，并采用了与 gte-Qwen2-7B-instruct 模型相同的训练数据和策略。模型集成了双向注意力机制，增强了上下文理解能力，并仅在查询端进行指令微调，提高了效率。此外，模型在庞大的多语言文本语料库上进行全面训练，涵盖了不同领域和场景，结合了弱监督和监督数据，确保模型适用于多种语言和广泛的下游任务。

📄 引用

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

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