Gte Small

🚀 gte-small

General Text Embeddings (GTE) 模型是由阿里巴巴达摩院训练的一系列模型，主要基于 BERT 框架，目前提供三种不同大小的模型，包括 GTE-large、GTE-base 和 GTE-small。这些模型在大规模的相关文本对语料库上进行训练，覆盖了广泛的领域和场景，可应用于文本嵌入的各种下游任务，如信息检索、语义文本相似度、文本重排序等。Towards General Text Embeddings with Multi-stage Contrastive Learning

🚀 快速开始

General Text Embeddings (GTE) 模型主要基于 BERT 框架，在大规模相关文本对语料库上训练，适用于信息检索、语义文本相似度、文本重排序等多种文本嵌入下游任务。

代码示例

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

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-small")
model = AutoModel.from_pretrained("thenlper/gte-small")

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# (Optionally) normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

结合 sentence-transformers 使用

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('thenlper/gte-large')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

✨ 主要特性

• 多尺寸模型：提供三种不同大小的模型，包括 GTE-large、GTE-base 和 GTE-small，可根据不同需求选择。
• 广泛适用性：在大规模的相关文本对语料库上进行训练，覆盖了广泛的领域和场景，可应用于文本嵌入的各种下游任务，包括信息检索、语义文本相似度、文本重排序等。

📦 安装指南

文档未提及具体安装命令，故跳过此章节。

💻 使用示例

基础用法

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

def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

input_texts = [
    "what is the capital of China?",
    "how to implement quick sort in python?",
    "Beijing",
    "sorting algorithms"
]

tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-small")
model = AutoModel.from_pretrained("thenlper/gte-small")

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# (Optionally) normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:1] @ embeddings[1:].T) * 100
print(scores.tolist())

高级用法

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('thenlper/gte-large')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

📚 详细文档

指标

我们在 MTEB 基准测试中比较了 GTE 模型与其他流行文本嵌入模型的性能。如需更详细的比较结果，请参考 MTEB 排行榜。

局限性

此模型仅适用于英文文本，并且任何长文本将被截断为最多 512 个标记。

引用

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

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

🔧 技术细节

文档未提供具体的技术实现细节，故跳过此章节。

📄 许可证

本项目采用 MIT 许可证。