| --- |
| license: mit |
| base_model: |
| - Qwen/Qwen3-Embedding-0.6B |
| --- |
| |
| onnx version. |
|
|
|
|
| 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 |
| |
