bge-m3-onnx

Usage with ONNX Runtime (Python)

If you haven't already, you can install the ONNX Runtime Python library with pip:

pip install onnxruntime==1.17.0

For tokenization, you can for example use HF Transformers by installing it with pip:

pip install transformers==4.37.2

Clone this repository with Git LFS to get the ONNX model files.

You can then use the model to compute embeddings, as follows:

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)

Note: you can use following sparse token weight processor from FlagEmbedding to get same the output for the sparse representation from the ONNX model:

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())
)

Export ONNX weights

You can export ONNX weights with the provided custom BGE-M3 PyTorch model bgem3_model.py file and with the provided export_onnx.py ONNX weight export script which leverages HF Optimum. If needed, you can modify the bgem3_model.py model configuration to for example remove embedding normalization or to not output all three embedding representations. If you modify the number of output representations, you need to also modify the ONNX output config BGEM3OnnxConfig in export_onnx.py.

First, install needed Python requirements as follows:

pip install -r requirements.txt

Then you can export ONNX weights as follows:

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

You can read more about the optional optimization levels here