Gte Base Zh

🚀 gte-base-zh

General Text Embeddings (GTE) 是由阿里巴巴達摩院基於 BERT 框架訓練的文本嵌入模型，在大規模相關文本對語料上進行訓練，適用於信息檢索、語義文本相似度、文本重排序等多種下游任務。

🚀 快速開始

安裝依賴

你可以使用以下命令安裝所需的庫：

pip install transformers torch sentence-transformers

代碼示例

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

input_texts = [
    "中國的首都是哪裡",
    "你喜歡去哪裡旅遊",
    "北京",
    "今天中午吃什麼"
]

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

# Tokenize the input texts
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]
 
# (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 = ['中國的首都是哪裡', '中國的首都是北京']

model = SentenceTransformer('thenlper/gte-base-zh')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

✨ 主要特性

• 多語言支持：目前提供了中文和英文不同大小的模型。
• 廣泛的適用性：在大規模相關文本對語料上訓練，可應用於信息檢索、語義文本相似度、文本重排序等多種下游任務。

📦 安裝指南

使用 pip 安裝所需的庫：

pip install transformers torch sentence-transformers

💻 使用示例

基礎用法

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

input_texts = [
    "中國的首都是哪裡",
    "你喜歡去哪裡旅遊",
    "北京",
    "今天中午吃什麼"
]

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

# Tokenize the input texts
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]
 
# (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 = ['中國的首都是哪裡', '中國的首都是北京']

model = SentenceTransformer('thenlper/gte-base-zh')
embeddings = model.encode(sentences)
print(cos_sim(embeddings[0], embeddings[1]))

📚 詳細文檔

模型列表

評估指標

在 MTEB（中文為 CMTEB）基準上，將 GTE 模型與其他流行的文本嵌入模型進行了性能比較。更多詳細的比較結果，請參考 MTEB 排行榜。

🔧 技術細節

General Text Embeddings (GTE) 模型由阿里巴巴達摩院訓練，主要基於 BERT 框架，在大規模相關文本對語料上進行訓練，涵蓋了廣泛的領域和場景，使其能夠應用於各種文本嵌入的下游任務。

📄 許可證

本項目採用 MIT 許可證。

⚠️ 重要提示

該模型僅適用於中文文本，任何長文本將被截斷為最多 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}
}