Stella Base Zh V2

🚀 stella模型

stella是一個通用的文本編碼模型，可用於文本編碼、檢索和語義匹配等任務。它有多種版本，能適應不同的應用場景和需求。

🚀 快速開始

安裝依賴

運行以下代碼前，請確保你已經安裝了必要的庫：

pip install torch numpy mteb sentence-transformers transformers sklearn

代碼示例

以下是使用stella-base-zh-v2模型進行文本編碼和相似度計算的示例代碼：

# 對於短對短數據集，下面是通用的使用方式
from sentence_transformers import SentenceTransformer

sentences = ["數據1", "數據2"]
model = SentenceTransformer('infgrad/stella-base-zh-v2')
print(model.max_seq_length)
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

你也可以直接使用transformers庫：

from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize

model = AutoModel.from_pretrained('infgrad/stella-base-zh-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh-v2')
sentences = ["數據1", "數據ABCDEFGH"]
batch_data = tokenizer(
    batch_text_or_text_pairs=sentences,
    padding="longest",
    return_tensors="pt",
    max_length=1024,
    truncation=True,
)
attention_mask = batch_data["attention_mask"]
model_output = model(**batch_data)
last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = normalize(vectors, norm="l2", axis=1, )
print(vectors.shape)  # 2,768

✨ 主要特性

• 多種模型選擇：提供了stella-large-zh-v2、stella-base-zh-v2、stella-large-zh和stella-base-zh等多種模型，以滿足不同的需求。
• 無需前綴文本：stella-base-zh-v2和stella-large-zh-v2模型使用簡單，在任何使用場景中都不需要加前綴文本。
• 支持向量可變維度：部分模型支持向量可變維度，如stella-mrl-large-zh-v3.5-1792d模型。
• 中英雙語支持：部分模型支持中英雙語，如puff系列模型。

📦 安裝指南

使用sentence-transformers庫加載模型：

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('infgrad/stella-base-zh-v2')

💻 使用示例

基礎用法

from sentence_transformers import SentenceTransformer

sentences = ["這是一個示例句子", "這是另一個示例句子"]
model = SentenceTransformer('infgrad/stella-base-zh-v2')
embeddings = model.encode(sentences)
print(embeddings)

高級用法

from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize
import torch

model = AutoModel.from_pretrained('infgrad/stella-base-zh-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh-v2')
sentences = ["高級用法示例句子", "用於演示更復雜的場景"]
batch_data = tokenizer(
    batch_text_or_text_pairs=sentences,
    padding="longest",
    return_tensors="pt",
    max_length=1024,
    truncation=True,
)
attention_mask = batch_data["attention_mask"]
model_output = model(**batch_data)
last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = normalize(vectors, norm="l2", axis=1, )
print(vectors.shape)

📚 詳細文檔

模型信息

模型參數

C-MTEB排行榜（中文）

長文本評測

現有數據集存在長度大於512的文本過少，且即便大於512，對於檢索而言也只需要前512的文本內容的問題，導致無法準確評估模型的長文本編碼能力。為了解決這個問題，整理了6份長文本測試集：

• CMRC2018，通用百科
• CAIL，法律閱讀理解
• DRCD，繁體百科，已轉簡體
• Military，軍工問答
• Squad，英文閱讀理解，已轉中文
• Multifieldqa_zh，清華的大模型長文本理解能力評測數據[9]

除了Military數據集，提供了其他5個測試數據的下載地址：https://drive.google.com/file/d/1WC6EWaCbVgz-vPMDFH4TwAMkLyh5WNcN/view?usp=sharing

評測指標為Recall@5，結果如下：

新聞動態

• [2024-04-06] 開源puff系列模型，專門針對檢索和語義匹配任務，更多的考慮泛化性和私有通用測試集效果，向量維度可變，中英雙語。
• [2024-02-27] 開源stella-mrl-large-zh-v3.5-1792d模型，支持向量可變維度。
• [2024-02-17] 開源stella v3系列、dialogue編碼模型和相關訓練數據。
• [2023-10-19] 開源stella-base-en-v2，使用簡單，不需要任何前綴文本。
• [2023-10-12] 開源stella-base-zh-v2和stella-large-zh-v2，效果更好且使用簡單，不需要任何前綴文本。
• [2023-09-11] 開源stella-base-zh和stella-large-zh

歡迎去本人主頁查看最新模型，並提出您的寶貴意見！

🔧 技術細節

訓練環境

• 硬件：單卡A100 - 80GB
• 環境：torch1.13.*；transformers-trainer + deepspeed + gradient-checkpointing

訓練參數

• 學習率：1e - 6
• batch_size：base模型為1024，額外增加20%的難負例；large模型為768，額外增加20%的難負例
• 數據量：第一版模型約100萬，其中用LLM構造的數據約有200K。LLM模型大小為13b。v2系列模型到了2000萬訓練數據。

📄 許可證

文檔中未提及許可證相關信息。

📋 ToDoList

• 評測的穩定性：評測過程中發現Clustering任務會和官方的結果不一致，大約有±0.0x的小差距，原因是聚類代碼沒有設置random_seed，差距可以忽略不計，不影響評測結論。
• 更高質量的長文本訓練和測試數據：訓練數據多是用13b模型構造的，肯定會存在噪聲。測試數據基本都是從mrc數據整理來的，所以問題都是factoid類型，不符合真實分佈。
• OOD的性能：雖然近期出現了很多向量編碼模型，但是對於不是那麼通用的domain，這一眾模型包括stella、openai和cohere，它們的效果均比不上BM25。

📖 參考資料

1. https://www.scidb.cn/en/detail?dataSetId=c6a3fe684227415a9db8e21bac4a15ab
2. https://github.com/wangyuxinwhy/uniem
3. https://github.com/CLUEbenchmark/SimCLUE
4. https://arxiv.org/abs/1612.00796
5. https://kexue.fm/archives/8847
6. https://huggingface.co/sensenova/piccolo-base-zh
7. https://kexue.fm/archives/7947
8. https://github.com/FlagOpen/FlagEmbedding
9. https://github.com/THUDM/LongBench

⚠️ 重要提示

因為長文本評測數據數量稀少，所以構造時也使用了train部分，如果自行評測，請注意模型的訓練數據以免數據洩露。

💡 使用建議

stella-base-zh和stella-large-zh用法和piccolo完全一致，即在檢索重排任務上給query和passage加上查詢: 和結果: 。對於短短匹配不需要做任何操作。stella-base-zh-v2和stella-large-zh-v2使用簡單，任何使用場景中都不需要加前綴文本。