许可证：MIT

BGE-M3 已通过HF Optimum转换为ONNX权重，以便兼容如ONNX Runtime等框架。

此ONNX模型可同时输出稠密、稀疏及ColBERT向量表示。输出结果为按前述表示顺序排列的numpy数组列表。

注：稠密与ColBERT向量默认遵循原始FlagEmbedding库的归一化处理。如需非归一化输出，可修改bgem3_model.py中的代码，并重新运行export_onnx.py脚本导出ONNX模型。

该ONNX模型已应用"O2"级图优化，优化级别详情参见此处。如需不同优化级别或无优化模型，可通过export_onnx.py脚本指定相应优化参数重新导出。

ONNX Runtime使用指南（Python）

若未安装，请通过pip安装ONNX Runtime：

pip install onnxruntime==1.17.0

分词处理可安装HF Transformers：

pip install transformers==4.37.2

通过Git LFS克隆本仓库获取ONNX模型文件。

使用示例：

import onnxruntime as ort
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3")
ort_session = ort.InferenceSession("model.onnx")

inputs = tokenizer("BGE M3是支持稠密检索、词汇匹配和多向量交互的嵌入模型。", padding="longest", return_tensors="np")
inputs_onnx = {k: ort.OrtValue.ortvalue_from_numpy(v) for k, v in inputs.items()}

outputs = ort_session.run(None, inputs_onnx)

注：稀疏表示输出可通过FlagEmbedding的权重处理器处理：

from collections import defaultdict

def process_token_weights(token_weights: np.ndarray, input_ids: list):
    result = defaultdict(int)
    unused_tokens = set([
        tokenizer.cls_token_id,
        tokenizer.eos_token_id,
        tokenizer.pad_token_id,
        tokenizer.unk_token_id,
    ])
    for w, idx in zip(token_weights, input_ids):
        if idx not in unused_tokens and w > 0:
            idx = str(idx)
            if w > result[idx]:
                result[idx] = w
    return result

token_weights = outputs[1].squeeze(-1)
lexical_weights = list(map(process_token_weights, token_weights, inputs["input_ids"].tolist()))

ONNX权重导出

使用提供的bgem3_model.py自定义PyTorch模型及基于HF Optimum的export_onnx.py脚本可导出ONNX权重。如需修改模型配置（如取消归一化或调整输出表示数量），需同步修改export_onnx.py中的BGEM3OnnxConfig输出配置。

安装依赖：

pip install -r requirements.txt

导出命令：

python export_onnx.py --output . --opset 17 --device cpu --optimize O2

优化级别说明详见官方文档