BGE-M3 Sparse —— ONNX 稀疏词级嵌入模型

模型简介

本模型源自 BAAI 发布的 BGE-M3 多语言嵌入模型，提取其稀疏头（Sparse Head） 并转换为 ONNX 格式，专为离线本地推理场景优化。

BGE-M3 是首个支持稠密（Dense）、稀疏（Sparse） 和多向量（ColBERT） 三种检索方式的统一嵌入模型。本仓库提供的是其 ONNX 格式的稀疏嵌入推理资源，能够在无 GPU 的桌面 / 边缘设备上，以纯 CPU 方式高效运行。

核心价值：稀疏向量保留了词级别的精确匹配能力，对于案号（如 (2024)沪民终0001号）、法条编号（如 《民法典》第584条）、专有名词等需要严格命中的关键信息具有不可替代的召回优势。

模型文件

注意：model.onnx 与 model.onnx_data 必须放在同一目录下，ONNX Runtime 才能正确加载模型权重。

模型架构

• 基础模型：BGE-M3（BAAI），基于 XLM-RoBERTa
• 参数规模：~568M
• 支持语言：中文、英文，覆盖 100+ 语言
• 最大序列长度：512 tokens
• 输出格式：float32[1, vocab_size] 稀疏权重张量
• 后处理：ReLU 过滤（权重 > 0 的非零项），输出 {sparse_indices, sparse_values}

工作流程：

输入文本 → AutoTokenizer → input_ids + attention_mask
         → ONNX 推理（稀疏头） → 词级权重 [1, vocab_size]
         → ReLU 过滤 → 非零词汇的 (索引, 权重) 对

快速开始

环境要求

• onnxruntime >= 1.24.3
• 分词器依赖 HuggingFace transformers 或 JavaScript @xenova/transformers

Python 示例

import onnxruntime as ort
import numpy as np
from transformers import AutoTokenizer

# 1. 加载模型与分词器
model_path = "./bge-m3-sparse/model.onnx"
tokenizer = AutoTokenizer.from_pretrained("./bge-m3-sparse")

session = ort.InferenceSession(model_path)

# 2. 分词
text = "（2024）沪民终 0001 号 《民法典》第584条 合同违约"
tokens = tokenizer(text, max_length=512, truncation=True, return_tensors="np")

# 3. ONNX 推理
outputs = session.run(None, {
    "input_ids": tokens["input_ids"].astype(np.int64),
    "attention_mask": tokens["attention_mask"].astype(np.int64),
})

# 4. ReLU 过滤，提取非零稀疏向量
sparse_weights = outputs[0][0]  # shape: [vocab_size]
indices = np.where(sparse_weights > 0)[0]
values = sparse_weights[indices]

print(f"非零 Token 数: {len(indices)}")
print(f"稀疏向量维度: {indices[:5]} ...")   # 词表索引
print(f"对应权重: {values[:5]} ...")

TypeScript / Node.js 示例

import * as ort from "onnxruntime-node";
import { AutoTokenizer } from "@xenova/transformers";

// 1. 加载
const modelPath = "./bge-m3-sparse";
const tokenizer = await AutoTokenizer.from_pretrained(modelPath);
const session = await ort.InferenceSession.create(`${modelPath}/model.onnx`);

// 2. 分词
const text = "（2024）沪民终 0001 号";
const encoded = await tokenizer(text, {
  max_length: 512,
  truncation: true,
  padding: false,
});

// 3. 推理
const inputIds = new ort.Tensor("int64", BigInt64Array.from(encoded.input_ids.data), [1, encoded.input_ids.data.length]);
const attentionMask = new ort.Tensor("int64", BigInt64Array.from(encoded.attention_mask.data), [1, encoded.attention_mask.data.length]);

const outputs = await session.run({
  input_ids: inputIds,
  attention_mask: attentionMask,
});

// 4. 提取稀疏向量
const data = outputs[session.outputNames[0]].data as Float32Array;
const sparseIndices: number[] = [];
const sparseValues: number[] = [];
for (let i = 0; i < data.length; i++) {
  if (data[i] > 0) {
    sparseIndices.push(i);
    sparseValues.push(data[i]);
  }
}
console.log(`非零 Token 数: ${sparseIndices.length}`);

应用场景

法律文本精确检索

• 案号精确匹配：(2024)沪民终0001号 必须原样命中，不能被语义泛化
• 法条索引：《民法典》第 584 条、《刑法》第 232 条 等法条引用的精确召回
• 法律专有名词：不可抗力、违约金、连带责任 等术语的高精度检索

混合检索（Hybrid Search）

稀疏向量通常与稠密向量（Dense Embedding）搭配使用：

• Dense（语义召回）：处理同义词、近义表达，如 "赔偿损失" ↔ "损害赔偿"
• Sparse（精确召回）：处理案号、法条编号等必须原样命中的关键信息

建议与 LanceDB、Milvus 等支持混合检索的向量数据库配合使用。

技术细节

与原版 BGE-M3 的关系

后处理说明

模型输出为全词表（~250,000 tokens）的权重张量。实际使用时经过 ReLU 过滤：

1. 保留所有权重 > 0 的 token 索引
2. 得到紧凑的 (indices, values) 稀疏表示
3. 典型输出约 150-400 个非零 token（取决于文本长度和内容）

许可证

本模型遵循 MIT License，与原版 BGE-M3 保持一致。

引用

如果你在研究中使用了本模型，请引用原版 BGE-M3 论文：

@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}
}

致谢

• 原始模型：BAAI/bge-m3
• ONNX 转换工具：Optimum / 🤗 Transformers ONNX export
• 运行时：ONNX Runtime