🚀 ONNX轉換all-MiniLM-L6-v2
本項目是將sentence-transformers/all-MiniLM-L6-v2轉換為ONNX模型。該模型可將句子和段落映射到384維的密集向量空間,適用於聚類或語義搜索等任務。此自定義模型會輸出last_hidden_state
和pooler_output
,而使用默認ONNX配置導出的sentence-transformers模型僅包含last_hidden_state
作為輸出。
🚀 快速開始
✨ 主要特性
📦 安裝指南
若要使用此模型,需先安裝optimum:
python -m pip install optimum
💻 使用示例
基礎用法
from optimum.onnxruntime.modeling_ort import ORTModelForCustomTasks
from transformers import AutoTokenizer
model = ORTModelForCustomTasks.from_pretrained("optimum/sbert-all-MiniLM-L6-with-pooler")
tokenizer = AutoTokenizer.from_pretrained("optimum/sbert-all-MiniLM-L6-with-pooler")
inputs = tokenizer("I love burritos!", return_tensors="pt")
pred = model(**inputs)
高級用法
使用transformers的pipeline API:
from transformers import pipeline
onnx_extractor = pipeline("feature-extraction", model=model, tokenizer=tokenizer)
text = "I love burritos!"
pred = onnx_extractor(text)
📚 詳細文檔
評估結果
若要對該模型進行自動評估,可參考 Sentence Embeddings Benchmark:https://seb.sbert.net
背景
本項目旨在使用自監督對比學習目標,在大規模句子級數據集上訓練句子嵌入模型。我們使用了預訓練的nreimers/MiniLM-L6-H384-uncased
模型,並在包含10億個句子對的數據集上進行微調。該項目是在Hugging Face組織的Community week using JAX/Flax for NLP & CV活動中開發的,作為Train the Best Sentence Embedding Model Ever with 1B Training Pairs項目的一部分。我們藉助高效的硬件基礎設施(7個TPU v3 - 8)進行訓練,並得到了Google的Flax、JAX和Cloud團隊成員在高效深度學習框架方面的支持。
預期用途
我們的模型旨在用作句子和短段落編碼器。給定輸入文本,它會輸出一個捕獲語義信息的向量。該句子向量可用於信息檢索、聚類或句子相似度任務。默認情況下,輸入文本超過256個詞塊時會被截斷。
訓練過程
預訓練
我們使用了預訓練的nreimers/MiniLM-L6-H384-uncased
模型。有關預訓練過程的詳細信息,請參考該模型的卡片。
微調
我們使用對比目標對模型進行微調。具體而言,我們計算批次中每個可能句子對的餘弦相似度,然後通過與真實對進行比較來應用交叉熵損失。
超參數
我們在TPU v3 - 8上訓練模型,共訓練100k步,批次大小為1024(每個TPU核心128)。我們使用了500步的學習率預熱,序列長度限制為128個標記。我們使用AdamW優化器,學習率為2e - 5。完整的訓練腳本可在當前倉庫中找到:train_script.py
。
訓練數據
我們使用多個數據集的組合來微調模型,句子對總數超過10億。我們根據加權概率對每個數據集進行採樣,具體配置詳見data_config.json
文件。
數據集 |
論文 |
訓練元組數 |
Reddit comments (2015 - 2018) |
論文 |
726,484,430 |
S2ORC 引用對(摘要) |
論文 |
116,288,806 |
WikiAnswers 重複問題對 |
論文 |
77,427,422 |
PAQ (問題,答案)對 |
論文 |
64,371,441 |
S2ORC 引用對(標題) |
論文 |
52,603,982 |
S2ORC (標題,摘要) |
論文 |
41,769,185 |
Stack Exchange (標題,正文)對 |
- |
25,316,456 |
Stack Exchange (標題 + 正文,答案)對 |
- |
21,396,559 |
Stack Exchange (標題,答案)對 |
- |
21,396,559 |
MS MARCO 三元組 |
論文 |
9,144,553 |
GOOAQ: Open Question Answering with Diverse Answer Types |
論文 |
3,012,496 |
Yahoo Answers (標題,答案) |
論文 |
1,198,260 |
Code Search |
- |
1,151,414 |
COCO 圖像描述 |
論文 |
828,395 |
SPECTER 引用三元組 |
論文 |
684,100 |
Yahoo Answers (問題,答案) |
論文 |
681,164 |
Yahoo Answers (標題,問題) |
論文 |
659,896 |
SearchQA |
論文 |
582,261 |
Eli5 |
論文 |
325,475 |
Flickr 30k |
論文 |
317,695 |
Stack Exchange 重複問題(標題) |
|
304,525 |
AllNLI (SNLI 和 MultiNLI) |
論文 SNLI,論文 MultiNLI |
277,230 |
Stack Exchange 重複問題(正文) |
|
250,519 |
Stack Exchange 重複問題(標題 + 正文) |
|
250,460 |
Sentence Compression |
論文 |
180,000 |
Wikihow |
論文 |
128,542 |
Altlex |
論文 |
112,696 |
Quora Question Triplets |
- |
103,663 |
Simple Wikipedia |
論文 |
102,225 |
Natural Questions (NQ) |
論文 |
100,231 |
SQuAD2.0 |
論文 |
87,599 |
TriviaQA |
- |
73,346 |
總計 |
|
1,170,060,424 |
📄 許可證
本項目採用Apache 2.0許可證。