Views
No views yet
1# Requires onnxruntime>=1.22.1
2# Requires tokenizers>=0.21.4
3
4import os
5import numpy as np
6import onnxruntime as ort
7from tokenizers import Tokenizer
8
9# Download the model in a folder called Qwen3-Embedding-0.6B-ONNX
10model_name_on_disk = "Qwen3-Embedding-0.6B-ONNX"
11
12documents = [
13 "The capital of China is Beijing.",
14 "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.",
15]
16
17def normalize_l2(array, axis=1, eps=1e-12):
18 norm = np.linalg.norm(array, ord=2, axis=axis, keepdims=True)
19 norm = np.maximum(norm, eps)
20 return array / norm
21
22tokenizer = Tokenizer.from_file(os.path.join(model_name_on_disk, "tokenizer.json"))
23tokenizer.enable_padding(direction='left')
24encodings = tokenizer.encode_batch(sentences)
25
26ids = np.array([e.ids for e in encodings], dtype=np.int64)
27mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
28position_ids = np.array([range(len(e.ids)) for e in encodings], dtype=np.int64)
29
30sess = ort.InferenceSession(os.path.join(model_name_on_disk, "onnx", "model.onnx"), graph_optimization_level=1)
31
32outputs = sess.run(None, {"input_ids": ids, "attention_mask": mask, "position_ids": position_ids})
33
34token_embeddings = outputs[0]
35last_token = token_embeddings[:, -1]
36embeddings = normalize_l2(last_token)
37