🚀 gte-large-zh
General Text Embeddings (GTE) 是一種通用文本嵌入模型。論文 Towards General Text Embeddings with Multi-stage Contrastive Learning 對其進行了詳細闡述。該模型由阿里巴巴達摩院訓練,主要基於 BERT 框架,目前提供了不同規模的中英文模型。GTE 模型在大規模相關文本對語料庫上進行訓練,覆蓋了廣泛的領域和場景,使其能夠應用於多種文本嵌入下游任務,如信息檢索、語義文本相似度和文本重排序等。
🚀 快速開始
模型列表
評估指標
我們在 MTEB(中文為 CMTEB)基準測試中,將 GTE 模型與其他流行的文本嵌入模型進行了性能比較。更多詳細的比較結果,請參考 MTEB 排行榜。
CMTEB 評估結果
模型 |
模型大小 (GB) |
嵌入維度 |
序列長度 |
平均得分 (35 個數據集) |
分類任務 (9 個數據集) |
聚類任務 (4 個數據集) |
成對分類任務 (2 個數據集) |
重排序任務 (4 個數據集) |
檢索任務 (8 個數據集) |
語義文本相似度任務 (8 個數據集) |
gte-large-zh |
0.65 |
1024 |
512 |
66.72 |
71.34 |
53.07 |
81.14 |
67.42 |
72.49 |
57.82 |
gte-base-zh |
0.20 |
768 |
512 |
65.92 |
71.26 |
53.86 |
80.44 |
67.00 |
71.71 |
55.96 |
stella-large-zh-v2 |
0.65 |
1024 |
1024 |
65.13 |
69.05 |
49.16 |
82.68 |
66.41 |
70.14 |
58.66 |
stella-large-zh |
0.65 |
1024 |
1024 |
64.54 |
67.62 |
48.65 |
78.72 |
65.98 |
71.02 |
58.3 |
bge-large-zh-v1.5 |
1.3 |
1024 |
512 |
64.53 |
69.13 |
48.99 |
81.6 |
65.84 |
70.46 |
56.25 |
stella-base-zh-v2 |
0.21 |
768 |
1024 |
64.36 |
68.29 |
49.4 |
79.96 |
66.1 |
70.08 |
56.92 |
stella-base-zh |
0.21 |
768 |
1024 |
64.16 |
67.77 |
48.7 |
76.09 |
66.95 |
71.07 |
56.54 |
piccolo-large-zh |
0.65 |
1024 |
512 |
64.11 |
67.03 |
47.04 |
78.38 |
65.98 |
70.93 |
58.02 |
piccolo-base-zh |
0.2 |
768 |
512 |
63.66 |
66.98 |
47.12 |
76.61 |
66.68 |
71.2 |
55.9 |
gte-small-zh |
0.1 |
512 |
512 |
60.04 |
64.35 |
48.95 |
69.99 |
66.21 |
65.50 |
49.72 |
bge-small-zh-v1.5 |
0.1 |
512 |
512 |
57.82 |
63.96 |
44.18 |
70.4 |
60.92 |
61.77 |
49.1 |
m3e-base |
0.41 |
768 |
512 |
57.79 |
67.52 |
47.68 |
63.99 |
59.54 |
56.91 |
50.47 |
text-embedding-ada-002(openai) |
- |
1536 |
8192 |
53.02 |
64.31 |
45.68 |
69.56 |
54.28 |
52.0 |
43.35 |
💻 使用示例
基礎用法
import torch.nn.functional as F
from torch import Tensor
from transformers import AutoTokenizer, AutoModel
input_texts = [
"中國的首都是哪裡",
"你喜歡去哪裡旅遊",
"北京",
"今天中午吃什麼"
]
tokenizer = AutoTokenizer.from_pretrained("thenlper/gte-large-zh")
model = AutoModel.from_pretrained("thenlper/gte-large-zh")
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')
outputs = model(**batch_dict)
embeddings = outputs.last_hidden_state[:, 0]
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-zh')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))
侷限性
本模型僅適用於中文文本,並且任何長文本將被截斷為最多 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 許可證。