Bge M3 Onnx

🚀 [BGE - M3 ONNX 模型]

[BGE - M3](https://huggingface.co/BAAI/bge - m3) 借助 HF Optimum 转换为 ONNX 权重，以实现兼容性，例如与 ONNX Runtime 兼容。此 ONNX 模型可一次性输出密集、稀疏和 ColBERT 嵌入表示，输出为按上述表示顺序排列的 numpy 数组列表。

🚀 快速开始

本部分将介绍如何使用转换后的 ONNX 模型进行嵌入计算，以及如何导出 ONNX 权重。

✨ 主要特性

• 多类型嵌入输出：该 ONNX 模型能够同时输出密集、稀疏和 ColBERT 嵌入表示。
• 规范化输出：默认情况下，密集和 ColBERT 嵌入会进行规范化处理。若需要非规范化输出，可修改 bgem3_model.py 中的代码，并使用 export_onnx.py 脚本重新导出 ONNX 模型。
• 图优化：模型应用了 “O2” 级别的图优化，若需要不同级别的优化或不进行优化，可使用 export_onnx.py 脚本并传入适当的优化参数。

📦 安装指南

安装 ONNX Runtime

若尚未安装 ONNX Runtime Python 库，可使用以下命令进行安装：

pip install onnxruntime==1.17.0

安装 HF Transformers

对于分词操作，可使用 HF Transformers，通过以下命令进行安装：

pip install transformers==4.37.2

克隆仓库

使用 [Git LFS](https://git - lfs.com/) 克隆此仓库以获取 ONNX 模型文件。

💻 使用示例

基础用法

以下代码展示了如何使用该 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 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.", 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)

高级用法

若要从 ONNX 模型获取稀疏表示的相同输出，可使用以下代码：

from collections import defaultdict


def process_token_weights(token_weights: np.ndarray, input_ids: list):
    # conver to dict
    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)
            # w = int(w)
            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 权重

可使用提供的自定义 BGE - M3 PyTorch 模型文件 bgem3_model.py 和 export_onnx.py ONNX 权重导出脚本（借助 HF Optimum）来导出 ONNX 权重。

若有需要，可修改 bgem3_model.py 模型配置，例如移除嵌入规范化或不输出所有三种嵌入表示。若修改了输出表示的数量，则还需修改 export_onnx.py 中的 ONNX 输出配置 BGEM3OnnxConfig。

安装依赖

首先，使用以下命令安装所需的 Python 依赖：

pip install -r requirements.txt

导出 ONNX 权重

然后，使用以下命令导出 ONNX 权重：

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

可在 此处 了解更多关于可选优化级别的信息。

📄 许可证

本项目采用 MIT 许可证。