Gte Large

🚀 gte-large

General Text Embeddings (GTE) 模型，可將文本轉換為向量表示，在信息檢索、語義相似度計算等自然語言處理任務中表現出色。該模型基於多階段對比學習，在大規模相關文本對上進行訓練，具有廣泛的適用性和強大的泛化能力。

🚀 快速開始

代碼示例

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-large")
model = AutoModel.from_pretrained("thenlper/gte-large")

# 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 模型由阿里巴巴達摩院訓練，基於 BERT 框架，提供了三種不同大小的模型，分別是 GTE-large、GTE-base 和 GTE-small，可根據不同需求進行選擇。
• 廣泛適用性：在大規模相關文本對語料庫上進行訓練，涵蓋了廣泛的領域和場景，適用於各種文本嵌入的下游任務，包括信息檢索、語義文本相似度、文本重排序等。

📚 詳細文檔

指標對比

我們在 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 許可證。