适用于 AWS Graviton4 的 BGE-M3 ONNX

这是 BAAI/bge-m3 的 ONNX 优化版本，专门针对 AWS Graviton4 处理器进行了优化。

模型描述

BGE-M3 是 FlagEmbedding 项目的一个多功能嵌入模型，支持：

• 稠密检索：传统的句子嵌入
• 稀疏检索：具有可学习稀疏权重的词汇匹配
• 多向量检索：ColBERT 风格的词元级嵌入
• 多语言支持：超过 100 种语言

优化细节

• ONNX Opset：17
• 优化级别：O3（为 BGE-M3 启用 GELU 近似）
• 目标硬件：AWS Graviton4（支持 bfloat16 的 ARM64）
• 模型大小：约 2.2GB
• 量化版本可用：quantized/ 子目录中的 INT8 量化模型

使用方法

快速开始

from optimum.onnxruntime import ORTModelForCustomTasks
from transformers import AutoTokenizer

# Load from Hugging Face Hub
model = ORTModelForCustomTasks.from_pretrained(
    \"idomeneo/bge-m3-onnx-graviton4\",
    file_name=\"model_optimized.onnx\"
)
tokenizer = AutoTokenizer.from_pretrained(\"idomeneo/bge-m3-onnx-graviton4\")

# Tokenize and get embeddings
inputs = tokenizer(\"Your text here\", return_tensors=\"np\", padding=True, truncation=True)
outputs = model.forward(**inputs)

# Access different embedding types
dense_embeddings = outputs[\"dense_vecs\"]    # Shape: (batch_size, 1024)
sparse_embeddings = outputs[\"sparse_vecs\"]   # Shape: (batch_size, seq_len, 1)
colbert_embeddings = outputs[\"colbert_vecs\"] # Shape: (batch_size, seq_len, 1024)

启用 bfloat16 加速的 AWS Graviton4

import onnxruntime as ort
from optimum.onnxruntime import ORTModelForCustomTasks
from transformers import AutoTokenizer

# Enable bfloat16 acceleration for Graviton4
sess_options = ort.SessionOptions()
sess_options.add_session_config_entry(\"mlas.enable_gemm_fastmath_arm64_bfloat16\", \"1\")

model = ORTModelForCustomTasks.from_pretrained(
    \"idomeneo/bge-m3-onnx-graviton4\",
    file_name=\"model_optimized.onnx\",
    session_options=sess_options
)
tokenizer = AutoTokenizer.from_pretrained(\"idomeneo/bge-m3-onnx-graviton4\")

使用量化模型

INT8 量化模型可显著节省内存并加快推理速度，同时质量损失极小：

from optimum.onnxruntime import ORTModelForCustomTasks
from transformers import AutoTokenizer

# Load quantized model
model = ORTModelForCustomTasks.from_pretrained(
    \"idomeneo/bge-m3-onnx-graviton4\",
    subfolder=\"quantized\",
    file_name=\"model_optimized_quantized.onnx\"
)
tokenizer = AutoTokenizer.from_pretrained(\"idomeneo/bge-m3-onnx-graviton4\")

# Usage is identical to the standard model
inputs = tokenizer(\"Your text here\", return_tensors=\"np\", padding=True, truncation=True)
outputs = model.forward(**inputs)

量化细节

• 量化类型：静态 INT8 量化
• 校准数据集：GLUE SST-2（300 个样本）
• 按通道：已启用以提高精度
• 质量：在多样化测试用例中与原始模型的相似度为 99.98%
• 性能：在 Graviton4 处理器上推理速度提升约 2-4 倍

量化质量结果

对 20 个多样化示例的全面测试表明质量保持良好：

整体统计：

• 平均相似度：99.98%
• 最低相似度：99.95%
• 最高相似度：99.99%
• 标准差：0.01%
• 所有测试用例的相似度均保持在 99.95% 以上

性能

在 AWS Graviton4 实例上，此优化模型提供：

• 与 PyTorch 相比，推理速度提升高达 3 倍
• 减少内存占用
• 原生支持 bfloat16 加速

模型文件

标准模型

• model_optimized.onnx：启用 GELU 近似的 O3 优化 ONNX 模型
• model_optimized.onnx.data：外部权重文件
• config.json：模型配置
• tokenizer.json：快速分词器
• tokenizer_config.json：分词器配置
• sentencepiece.bpe.model：SentencePiece 模型
• special_tokens_map.json：特殊标记映射
• ort_config.json：ONNX Runtime 配置

量化模型（位于 quantized/ 子目录）

• model_optimized_quantized.onnx：INT8 量化模型
• model_optimized_quantized.onnx.data：量化权重
• 所有分词器文件与标准模型共享

引用

@article{bge-m3,
  title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation},
  author={Chen, Jianlv and Xiao, Shitao and Zhang, Peitian and Luo, Kun and Lian, Defu and Liu, Zheng},
  journal={arXiv preprint arXiv:2402.03216},
  year={2024}
}

许可证

MIT 许可证（继承自 BAAI/bge-m3）