medialab-sciencespo/Qwen3-Embedding-0.6B-ONNX

# Requires onnxruntime>=1.22.1
# Requires tokenizers>=0.21.4

import os
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

# Download the model in a folder called Qwen3-Embedding-0.6B-ONNX
model_name_on_disk = "Qwen3-Embedding-0.6B-ONNX"

documents = [
    "The capital of China is Beijing.",
    "Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun.",
]

def normalize_l2(array, axis=1, eps=1e-12):
	norm = np.linalg.norm(array, ord=2, axis=axis, keepdims=True)
	norm = np.maximum(norm, eps)
	return array / norm

tokenizer = Tokenizer.from_file(os.path.join(model_name_on_disk, "tokenizer.json"))
tokenizer.enable_padding(direction='left')
encodings = tokenizer.encode_batch(sentences)

ids = np.array([e.ids for e in encodings], dtype=np.int64)
mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
position_ids = np.array([range(len(e.ids)) for e in encodings], dtype=np.int64)

sess = ort.InferenceSession(os.path.join(model_name_on_disk, "onnx", "model.onnx"), graph_optimization_level=1)

outputs = sess.run(None, {"input_ids": ids, "attention_mask": mask, "position_ids": position_ids})

token_embeddings = outputs[0]
last_token = token_embeddings[:, -1]
embeddings = normalize_l2(last_token)