mxbai-edge-colbert-v0-17m — ONNX 导出版本（ColBERT，ModernBERT 主干网络）

本仓库包含 mixedbread-ai/mxbai-edge-colbert-v0-17m 的 ONNX 导出版本，使用 PyLate + ColBERT 感知封装器生成。它保留了投影层堆栈和 ColBERT 标记（[Q] / [D] ），并包含用于 MaxSim 的跳过列表。

内容

• onnx/model.onnx — FP32 导出，opset 17（✅ 与 PyTorch 余弦相似度 1.0）
• onnx/model_quantized.onnx — 动态 INT8 量化（⚠️ 与 PyTorch 余弦相似度约 0.972；质量有所下降）
• tokenizer.json, tokenizer_config.json, special_tokens_map.json — 从带有标记的 PyLate 修改版分词器保存
• config.json — 模型配置
• conversion_metadata.json — 最小导出元数据
• skiplist.json — MaxSim 期间需要跳过的 token ID（此模型为 32 个标点符号 ID）

架构（已验证）

Transformer 256 → 投影层 512 → 投影层 48

特殊标记

• [Q] : 50368
• [D] : 50369

质量检查

• PyTorch vs ONNX FP32：余弦相似度 1.00000000，MSE 0.0
• PyTorch vs ONNX INT8：余弦相似度约 0.9719（观察到质量下降）

使用方法（Python，onnxruntime）

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

model_dir = "path/to/this/repo"
sess = ort.InferenceSession(f"{model_dir}/onnx/model.onnx", providers=["CPUExecutionProvider"])
tok = AutoTokenizer.from_pretrained(model_dir)

q = "[Q] what is colbert?"
enc = tok(q, return_tensors="np", padding="max_length", max_length=128, truncation=True)
out = sess.run(None, {"input_ids": enc["input_ids"], "attention_mask": enc["attention_mask"]})[0]
print(out.shape)  # (batch, seq_len, 48)

使用方法（Node，onnxruntime-node）

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

const modelDir = "path/to/this/repo";
const tokenizer = await AutoTokenizer.from_pretrained(modelDir);
const session = await ort.InferenceSession.create(`${modelDir}/onnx/model.onnx`);

const q = "[Q] what is colbert?";
const encoded = await tokenizer(q, { return_tensors: "np", padding: "max_length", max_length: 128, truncation: true });
const outputs = await session.run({ input_ids: encoded.input_ids, attention_mask: encoded.attention_mask });
console.log(outputs[session.outputNames[0]].dims); // [1, 128, 48]

const skiplist = new Set(JSON.parse(fs.readFileSync(`${modelDir}/skiplist.json`, "utf8")));

注意事项

• INT8 文件已提供，但存在可测量的精度漂移；建议使用 FP32（如果大小是关键因素，可导出 FP16）。
• 此处的跳过列表仅包含标点符号，因为模型未暴露额外的跳过词；如果您维护了更丰富的跳过列表，请在下游进行调整。

转换信息

• 工具：PyLate + 自定义 ColBERT 封装器
• Opset：17
• 日期：2025-02（详情见 conversion_metadata.json）