Piccolo Base Zh

🚀 piccolo-base-zh

piccolo 是一箇中文通用文本嵌入模型，由商湯科技研究院通用模型組訓練完成。該模型借鑑了 E5 和 GTE 的訓練流程，採用兩階段訓練方式，目前提供了 piccolo-base-zh 和 piccolo-large-zh 兩個版本。

🚀 快速開始

在 sentence-transformer package 中可以很容易地調用 piccolo 模型。

# for s2s dataset, you can use piccolo as below
# 對於短對短數據集，下面是通用的使用方式
from sentence_transformers import SentenceTransformer
sentences = ["數據1", "數據2"]
model = SentenceTransformer('sensenova/piccolo-base-zh')
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

# for s2p dataset, we recommend to add instruction for passage retrieval
# 對於短對長數據集，我們推薦添加 instruction，來幫助模型更好地進行檢索。
from sentence_transformers import SentenceTransformer
queries = ['query_1', 'query_2']
passages = ["doc_1", "doc_2"]

model = SentenceTransformer('sensenova/piccolo-base-zh')
q_embeddings = model.encode(["查詢：" + q for q in queries], normalize_embeddings=True)
p_embeddings = model.encode(["結果：" + p for p in passages], normalize_embeddings=True)
scores = q_embeddings @ p_embeddings.T

✨ 主要特性

• 兩階段訓練：第一階段蒐集和爬取 4 億中文文本對（弱監督文本對數據），採用二元組的 softmax 對比學習損失優化模型；第二階段蒐集整理 2000 萬人工標註的中文文本對（精標數據），採用帶有難負樣本的三元組的 softmax 對比學習損失進一步優化。
• 多場景適用：在分類、聚類、重排序、檢索和語義文本相似度等多種任務上均有良好表現。

📦 安裝指南

文檔未提及安裝步驟，故跳過此章節。

💻 使用示例

基礎用法

# for s2s dataset, you can use piccolo as below
# 對於短對短數據集，下面是通用的使用方式
from sentence_transformers import SentenceTransformer
sentences = ["數據1", "數據2"]
model = SentenceTransformer('sensenova/piccolo-base-zh')
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)

高級用法

# for s2p dataset, we recommend to add instruction for passage retrieval
# 對於短對長數據集，我們推薦添加 instruction，來幫助模型更好地進行檢索。
from sentence_transformers import SentenceTransformer
queries = ['query_1', 'query_2']
passages = ["doc_1", "doc_2"]

model = SentenceTransformer('sensenova/piccolo-base-zh')
q_embeddings = model.encode(["查詢：" + q for q in queries], normalize_embeddings=True)
p_embeddings = model.encode(["結果：" + p for p in passages], normalize_embeddings=True)
scores = q_embeddings @ p_embeddings.T

📚 詳細文檔

模型對比

我們將 piccolo 與其他的開源 embedding 模型在 CMTEB 榜單上進行了比較，請參考 CMTEB 榜單。我們在 eval 文件夾中提供了復現結果的腳本。

訓練細節

預訓練

pretrain 通常不需要太大的 max length，推薦 128。小的 max length 用以提高 batch size，加快訓練速度，從而適應大規模數據。pretrain 損失我們採用二元組 contrastive loss，不加入 hard negative，直接採用 inbatch negative。在實際訓練中，我們使用了 32 張 40G A100 進行訓練，單卡的 batch size 為 1024。

微調

對於微調，我們通常將 max length 擴增到 512，以適應更大長度的文本輸入。微調時會更多地採樣 S2P 數據，以增強模型在檢索任務上的性能。微調損失採用三元組 contrastive loss，加入 hard negative，neg num 通常設置為 2 - 7，損失計算方式可以參考 GTE 裡的 improved contrastive loss。注意，我們給 query 和 passage 設置了不同的 max length，query 的 max length 始終保持在 64。

其他技巧

1. 減小顯存的方式：fp16 + gradient checkpointing + ZERO STAGE1（stage2 不支持雙塔結構下的 gradient checkpointing），相關 issue 見：https://github.com/microsoft/DeepSpeed/issues/988
2. 數據集採樣器：我們採用了 M3E 的 dataset sampler，用以保證每個 batch 裡的樣本均來自於一個 dataset，負樣本更有價值。
3. 指令：指令在我們的實驗中對檢索任務有非常大的性能提升，我們在每個訓練樣本前都加入'查詢: '和'結果: '這樣的指令。

🔧 技術細節

文檔中關於模型訓練的技術細節描述較為詳細，包括訓練階段、損失函數、顯存優化技巧等，滿足技術細節章節的要求。

📄 許可證

Piccolo 使用 MIT License，可免費用於商業目的。

致謝

piccolo 由商湯科技研究院通用模型組訓練完成。Jinkin 完成了代碼實現和模型訓練，Jinkin 和 CCCCxxx 共同完成了數據蒐集、整理和評測工作。項目由 Gaomengya 和 chaorenwu111 主導。同時，感謝 lux0933 以及 yangkai001 的交流與討論，他們提供了許多有用的建議。