Sbert All MiniLM L6 With Pooler

🚀 ONNX轉換all-MiniLM-L6-v2

本項目是將sentence-transformers/all-MiniLM-L6-v2轉換為ONNX模型。該模型可將句子和段落映射到384維的密集向量空間，適用於聚類或語義搜索等任務。此自定義模型會輸出last_hidden_state和pooler_output，而使用默認ONNX配置導出的sentence-transformers模型僅包含last_hidden_state作為輸出。

🚀 快速開始

✨ 主要特性

• 基於sentence-transformers的ONNX模型，可將文本映射到384維向量空間。
• 自定義模型輸出last_hidden_state和pooler_output。

📦 安裝指南

若要使用此模型，需先安裝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文件。

📄 許可證

本項目採用Apache 2.0許可證。