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许可证。