from optimum.onnxruntime import ORTModelForFeatureExtraction
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-m3") # Better to use the original tokenizer (very lightweight)
model = ORTModelForFeatureExtraction.from_pretrained("MahradHosseini/bge-m3-onnx-int8")

questions = ["What is your opening hour?", "Where are your offices?"]

input_q = tokenizer(
    questions,
    padding=True,
    truncation=True,
    return_tensors="np"
)
print(f"Question input keys: {list(input_q.keys())}, shapes: {[v.shape for v in input_q.values()]}")

output_q = self.model(**input_q)
print(f"Question output keys: {list(output_q.keys())}, shapes: {[v.shape for v in output_q.values()]}")

question_embeddings = {
    "dense_vecs": mean_pooling(output_q["last_hidden_state"], input_q["attention_mask"]),
}
print(f"Embedded {len(question_embeddings['dense_vecs'])} questions from {self.data_file}")

def mean_pooling(last_hidden_state, attention_mask):
    # last_hidden_state: [batch_size, seq_len, hidden_size]
    # attention_mask: [batch_size, seq_len]
    input_mask_expanded = np.expand_dims(attention_mask, -1).astype(np.float32)
    return np.sum(last_hidden_state * input_mask_expanded, axis=1) / np.clip(
        input_mask_expanded.sum(axis=1), a_min=1e-9, a_max=None
    )

Conversion Details

HuggingFace Optimum was used to convert the base model to ONNX and then to quantize to int8.

0. fresh env with safe versions

pip install "optimum[exporters]" onnx onnxruntime

1. export WITHOUT optimisation

optimum-cli export onnx -m BAAI/bge-m3 bge-m3-onnx

2. int8 dynamic quantisation, AVX2 preset, per-channel

optimum-cli onnxruntime quantize
--onnx_model bge-m3-onnx
--avx2 --per_channel
-o bge-m3-int8

aris-bb/bge-m3-onnx-int8

Model Card for BGE-M3 ONNX Int8

Model Description

Uses

How to Get Started with the Model

Conversion Details

0. fresh env with safe versions

1. export WITHOUT optimisation

2. int8 dynamic quantisation, AVX2 preset, per-channel

aris-bb/
bge-m3-onnx-int8