File size: 1,502 Bytes
6fa892b 049c77a 6fa892b 049c77a 6fa892b 049c77a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | ---
license: mit
base_model:
- Qwen/Qwen3-Embedding-0.6B
---
interface sample:
```
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
class QwenEmbeddingONNX:
def __init__(self, model_dir, model_path, max_length=512, providers=None):
self.max_length = max_length
self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
if providers is None:
providers = ['CPUExecutionProvider']
self.session = ort.InferenceSession(
model_path,
providers=providers
)
self.input_names = [i.name for i in self.session.get_inputs()]
def encode(self, texts, normalize=True):
if isinstance(texts, str):
texts = [texts]
inputs = self.tokenizer(
texts,
padding=True,
truncation=True,
max_length=self.max_length,
return_tensors="np",
)
input_feed = {
k: inputs[k].astype(np.int64)
for k in self.input_names
if k in inputs
}
outputs = self.session.run(None, input_feed)
embeddings = outputs[0]
# 如果输出是 [B, L, H]，取第一个 token
if embeddings.ndim == 3:
embeddings = embeddings[:, 0, :]
# L2 normalize
if normalize:
norm = np.linalg.norm(embeddings, axis=1, keepdims=True)
embeddings = embeddings / (norm + 1e-9)
return embeddings
``` |