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-o3")

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

Developed by: Mahrad Hosseini
Model type: Embedding Model
License: Apache-2.0
Finetuned from model: BAAI/BGE-M3

HuggingFace Optimum was used to convert the base model to ONNX with O3 optimization.

0. fresh env with safe versions

pip install "optimum[exporters]" onnx onnxruntime

1. export WITH optimisation

optimum-cli export onnx --model BAAI/bge-m3 --optimize O3 bge-m3-onnx-o3