Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import numpy as np | |
| import onnxruntime as ort | |
| from transformers import AutoTokenizer | |
| from typing import List, Tuple | |
| from src.ontology.models import OntologyRecord | |
| from src.runtime.device_manager import DeviceManager | |
| class ONNXEmbeddingEngine: | |
| def __init__(self, model_path: str = "JairoDanielMT/paraphrase-multilingual-MiniLM-L12-v2-onnx-int8", index_dir: str = "data/faiss_indices_onnx"): | |
| self.model_path = model_path | |
| self.index_dir = index_dir | |
| self.tokenizer = None | |
| self.session = None | |
| self.indices = {} | |
| self.records = {} | |
| self.provider = DeviceManager.get_optimal_device() | |
| self.batch_size = DeviceManager.get_optimal_batch_size() | |
| def _init_model(self): | |
| if not self.tokenizer: | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_path) | |
| if not self.session: | |
| from optimum.onnxruntime import ORTModelForFeatureExtraction | |
| self.model = ORTModelForFeatureExtraction.from_pretrained( | |
| self.model_path, | |
| file_name="model_quantized.onnx", | |
| provider=self.provider | |
| ) | |
| # We assign the underlying inference session to maintain compatibility with the rest of the code | |
| self.session = self.model.model | |
| def load_index(self): | |
| import faiss | |
| if not os.path.exists(self.index_dir): | |
| raise FileNotFoundError(f"Index directory not found at {self.index_dir}") | |
| for filename in os.listdir(self.index_dir): | |
| if filename.endswith(".index"): | |
| category = filename.replace(".index", "") | |
| index_path = os.path.join(self.index_dir, filename) | |
| mapping_path = os.path.join(self.index_dir, f"{category}_records.json") | |
| if os.path.exists(mapping_path): | |
| self.indices[category] = faiss.read_index(index_path) | |
| with open(mapping_path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| self.records[category] = [OntologyRecord(**item) for item in data] | |
| self._init_model() | |
| def build_index(self, ontology_dir: str): | |
| import faiss | |
| self._init_model() | |
| os.makedirs(self.index_dir, exist_ok=True) | |
| for filename in os.listdir(ontology_dir): | |
| if filename.endswith(".json"): | |
| category = filename.replace(".json", "") | |
| path = os.path.join(ontology_dir, filename) | |
| cat_records = [] | |
| texts = [] | |
| with open(path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| for item in data: | |
| record = OntologyRecord(**item) | |
| cat_records.append(record) | |
| texts.append(record.canonical) | |
| for alias in record.aliases: | |
| cat_records.append(record) | |
| texts.append(alias) | |
| if not texts: | |
| continue | |
| embeddings = self.encode(texts) | |
| dimension = embeddings.shape[1] | |
| index = faiss.IndexFlatL2(dimension) | |
| index.add(embeddings) | |
| index_path = os.path.join(self.index_dir, f"{category}.index") | |
| faiss.write_index(index, index_path) | |
| mapping_path = os.path.join(self.index_dir, f"{category}_records.json") | |
| with open(mapping_path, "w", encoding="utf-8") as f: | |
| json.dump([r.model_dump() for r in cat_records], f, ensure_ascii=False) | |
| self.indices[category] = index | |
| self.records[category] = cat_records | |
| def encode(self, texts: List[str], batch_size: int = 128) -> np.ndarray: | |
| self._init_model() | |
| all_embeddings = [] | |
| import tqdm | |
| for i in tqdm.tqdm(range(0, len(texts), batch_size), desc="Encoding batches (ONNX)"): | |
| batch_texts = texts[i:i + batch_size] | |
| inputs = self.tokenizer(batch_texts, padding=True, truncation=True, return_tensors="np") | |
| outputs = self.session.run(None, { | |
| "input_ids": inputs["input_ids"], | |
| "attention_mask": inputs["attention_mask"], | |
| "token_type_ids": inputs["token_type_ids"] | |
| }) | |
| token_embeddings = outputs[0] | |
| attention_mask = inputs["attention_mask"] | |
| # Mean pooling | |
| input_mask_expanded = np.broadcast_to(np.expand_dims(attention_mask, -1), token_embeddings.shape) | |
| sum_embeddings = np.sum(token_embeddings * input_mask_expanded, 1) | |
| sum_mask = np.clip(np.sum(input_mask_expanded, 1), a_min=1e-9, a_max=None) | |
| embeddings = sum_embeddings / sum_mask | |
| # Normalize | |
| norms = np.linalg.norm(embeddings, axis=1, keepdims=True) | |
| embeddings = embeddings / norms | |
| all_embeddings.append(embeddings.astype('float32')) | |
| return np.vstack(all_embeddings) | |
| def calibrate_score(self, l2_distance: float) -> float: | |
| cos_sim = 1.0 - (l2_distance / 2.0) | |
| k = 15.0 | |
| x0 = 0.8 | |
| conf = 1.0 / (1.0 + np.exp(-k * (cos_sim - x0))) | |
| return float(np.clip(conf, 0.0, 1.0)) | |
| def search(self, query: str, category: str = None, top_k: int = 5) -> List[Tuple[OntologyRecord, float]]: | |
| if not self.indices: | |
| self.load_index() | |
| query_vector = self.encode([query]) | |
| results = [] | |
| categories_to_search = [category] if category and category in self.indices else self.indices.keys() | |
| for cat in categories_to_search: | |
| distances, indices = self.indices[cat].search(query_vector, top_k) | |
| for dist, idx in zip(distances[0], indices[0]): | |
| if idx < len(self.records[cat]): | |
| conf = self.calibrate_score(float(dist)) | |
| results.append((self.records[cat][idx], conf)) | |
| results.sort(key=lambda x: x[1], reverse=True) | |
| return results[:top_k] | |