Spaces:
Running
Running
| """ | |
| embedder.py — BGE-small-en-v1.5 via pure onnxruntime (no torch, no optimum). | |
| ONNX weights downloaded from Xenova/bge-small-en-v1.5 at build time. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| from typing import List | |
| import numpy as np | |
| from config import EMBED_BATCH_SIZE, EMBED_MODEL_NAME, EMBED_ONNX_DIR | |
| logger = logging.getLogger(__name__) | |
| _BGE_QUERY_PREFIX = "Represent this sentence: " | |
| class Embedder: | |
| def __init__(self): | |
| self._session = None | |
| self._tokenizer = None | |
| self._ready = False | |
| def load(self) -> None: | |
| if self._ready: | |
| return | |
| import onnxruntime as ort | |
| from transformers import AutoTokenizer | |
| # model.onnx is placed flat in EMBED_ONNX_DIR by download_models.py | |
| model_path = os.path.join(EMBED_ONNX_DIR, "model.onnx") | |
| if not os.path.isfile(model_path): | |
| raise FileNotFoundError( | |
| f"ONNX model not found at {model_path}. " | |
| "Run download_models.py first." | |
| ) | |
| logger.info("Loading ONNX embedder from %s", model_path) | |
| opts = ort.SessionOptions() | |
| opts.intra_op_num_threads = max(1, os.cpu_count() or 2) | |
| opts.inter_op_num_threads = max(1, os.cpu_count() or 2) | |
| self._session = ort.InferenceSession( | |
| model_path, | |
| sess_options=opts, | |
| providers=["CPUExecutionProvider"], | |
| ) | |
| self._tokenizer = AutoTokenizer.from_pretrained(EMBED_ONNX_DIR) | |
| self._ready = True | |
| logger.info("Embedder ready (pure ONNX/onnxruntime, no torch)") | |
| def encode_chunks(self, texts: List[str]) -> np.ndarray: | |
| return self._encode(texts) | |
| def encode_query(self, text: str) -> np.ndarray: | |
| return self._encode([_BGE_QUERY_PREFIX + text])[0] | |
| def _encode(self, texts: List[str]) -> np.ndarray: | |
| if not self._ready: | |
| raise RuntimeError("Embedder.load() has not been called") | |
| all_embeddings = [] | |
| for i in range(0, len(texts), EMBED_BATCH_SIZE): | |
| batch = texts[i: i + EMBED_BATCH_SIZE] | |
| enc = self._tokenizer( | |
| batch, padding=True, truncation=True, | |
| max_length=512, return_tensors="np", | |
| ) | |
| inputs = { | |
| "input_ids": enc["input_ids"].astype(np.int64), | |
| "attention_mask": enc["attention_mask"].astype(np.int64), | |
| } | |
| if "token_type_ids" in enc: | |
| inputs["token_type_ids"] = enc["token_type_ids"].astype(np.int64) | |
| outputs = self._session.run(None, inputs) | |
| token_emb = outputs[0] # (batch, seq, dim) | |
| mask = enc["attention_mask"][:, :, np.newaxis].astype(np.float32) | |
| summed = np.sum(token_emb * mask, axis=1) | |
| count = np.clip(mask.sum(axis=1), 1e-9, None) | |
| all_embeddings.append((summed / count).astype(np.float32)) | |
| return np.vstack(all_embeddings) | |
| def is_ready(self) -> bool: | |
| return self._ready | |
| embedder = Embedder() |