Spaces:
Sleeping
Sleeping
| """ | |
| Embedding module using a quantized ONNX MiniLM model and a pure Python WordPiece tokenizer. | |
| Calculates cosine similarity and caches column embeddings for performance. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import urllib.request | |
| import numpy as np | |
| import onnxruntime as ort | |
| import re | |
| from typing import Sequence | |
| # Default Hugging Face URLs for the quantized all-MiniLM-L6-v2 model and vocab | |
| MODEL_URL = "https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/onnx/model_quantized.onnx" | |
| VOCAB_URL = "https://huggingface.co/onnx-community/all-MiniLM-L6-v2-ONNX/resolve/main/vocab.txt" | |
| class WordPieceTokenizer: | |
| """Pure Python implementation of a WordPiece tokenizer.""" | |
| def __init__(self, vocab_path: str): | |
| self.vocab: dict[str, int] = {} | |
| with open(vocab_path, "r", encoding="utf-8") as f: | |
| for i, line in enumerate(f): | |
| token = line.strip() | |
| self.vocab[token] = i | |
| self.unk_token = "[UNK]" | |
| self.cls_token = "[CLS]" | |
| self.sep_token = "[SEP]" | |
| self.unk_id = self.vocab.get(self.unk_token, 100) | |
| self.cls_id = self.vocab.get(self.cls_token, 101) | |
| self.sep_id = self.vocab.get(self.sep_token, 102) | |
| def tokenize_word(self, word: str) -> list[str]: | |
| """Tokenize a single word into WordPiece subwords.""" | |
| if word in self.vocab: | |
| return [word] | |
| tokens = [] | |
| start = 0 | |
| is_bad = False | |
| while start < len(word): | |
| end = len(word) | |
| cur_substr = None | |
| while start < end: | |
| substr = word[start:end] | |
| if start > 0: | |
| substr = "##" + substr | |
| if substr in self.vocab: | |
| cur_substr = substr | |
| break | |
| end -= 1 | |
| if cur_substr is None: | |
| is_bad = True | |
| break | |
| tokens.append(cur_substr) | |
| start = end | |
| if is_bad: | |
| return [self.unk_token] | |
| return tokens | |
| def encode(self, text: str, max_length: int = 128) -> dict[str, np.ndarray]: | |
| """Encode text into model inputs (input_ids, attention_mask, token_type_ids).""" | |
| text = text.lower() | |
| # Basic word and punctuation splitter | |
| words = re.findall(r"\w+|[^\w\s]", text, re.UNICODE) | |
| tokens = [] | |
| for word in words: | |
| tokens.extend(self.tokenize_word(word)) | |
| # Truncate | |
| if len(tokens) > max_length - 2: | |
| tokens = tokens[:max_length - 2] | |
| # Build token IDs | |
| input_ids = [self.cls_id] + [self.vocab.get(t, self.unk_id) for t in tokens] + [self.sep_id] | |
| attention_mask = [1] * len(input_ids) | |
| token_type_ids = [0] * len(input_ids) | |
| # Pad to max_length | |
| padding_len = max_length - len(input_ids) | |
| if padding_len > 0: | |
| input_ids.extend([0] * padding_len) | |
| attention_mask.extend([0] * padding_len) | |
| token_type_ids.extend([0] * padding_len) | |
| # Convert to numpy arrays of type int64 (as expected by ONNX model) | |
| return { | |
| "input_ids": np.array([input_ids], dtype=np.int64), | |
| "attention_mask": np.array([attention_mask], dtype=np.int64), | |
| "token_type_ids": np.array([token_type_ids], dtype=np.int64), | |
| } | |
| def cosine_similarity(v1: np.ndarray, v2: np.ndarray) -> float: | |
| """Calculate the cosine similarity between two 1D vectors.""" | |
| dot = np.dot(v1, v2) | |
| norm1 = np.linalg.norm(v1) | |
| norm2 = np.linalg.norm(v2) | |
| if norm1 == 0 or norm2 == 0: | |
| return 0.0 | |
| return float(dot / (norm1 * norm2)) | |
| class EmbeddingModel: | |
| """ONNX-based text embedding generator with built-in WordPiece tokenization.""" | |
| def __init__(self, cache_dir: str | None = None): | |
| if cache_dir is None: | |
| # First try zero-llm-engine/models directory, fallback to temp dir /tmp/models | |
| base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) | |
| cache_dir = os.path.join(base_dir, "zero-llm-engine", "models", "onnx_cache") | |
| if not os.path.exists(cache_dir): | |
| try: | |
| os.makedirs(cache_dir, exist_ok=True) | |
| except Exception: | |
| # Fallback if the folder is read-only (e.g. some system settings or Docker environments) | |
| cache_dir = os.path.join("/tmp", "onnx_cache") | |
| os.makedirs(cache_dir, exist_ok=True) | |
| self.cache_dir = cache_dir | |
| self.model_path = os.path.join(cache_dir, "model_quantized.onnx") | |
| self.vocab_path = os.path.join(cache_dir, "vocab.txt") | |
| self.session: ort.InferenceSession | None = None | |
| self.tokenizer: WordPieceTokenizer | None = None | |
| # In-memory embedding cache: maps text string -> np.ndarray embedding | |
| self._embedding_cache: dict[str, np.ndarray] = {} | |
| def ensure_model_files(self) -> None: | |
| """Download model and vocabulary files if they do not exist locally.""" | |
| if not os.path.exists(self.vocab_path): | |
| print(f"[Embeddings] Downloading vocabulary to {self.vocab_path}...") | |
| urllib.request.urlretrieve(VOCAB_URL, self.vocab_path) | |
| if not os.path.exists(self.model_path): | |
| print(f"[Embeddings] Downloading quantized ONNX model to {self.model_path}...") | |
| urllib.request.urlretrieve(MODEL_URL, self.model_path) | |
| def load_model(self) -> None: | |
| """Ensure files are downloaded and load the ONNX session and tokenizer.""" | |
| if self.session is not None and self.tokenizer is not None: | |
| return | |
| self.ensure_model_files() | |
| # Initialize tokenization & ONNX session | |
| self.tokenizer = WordPieceTokenizer(self.vocab_path) | |
| # Using CPU execution provider by default for maximum compatibility | |
| self.session = ort.InferenceSession(self.model_path, providers=["CPUExecutionProvider"]) | |
| def get_embedding(self, text: str) -> np.ndarray: | |
| """Generate a 1D mean-pooled normalized embedding vector for the text.""" | |
| self.load_model() | |
| text_key = text.lower().strip() | |
| if text_key in self._embedding_cache: | |
| return self._embedding_cache[text_key] | |
| assert self.tokenizer is not None | |
| assert self.session is not None | |
| # Tokenize and format inputs | |
| inputs = self.tokenizer.encode(text_key) | |
| # Run ONNX inference | |
| outputs = self.session.run(None, inputs) | |
| # The first output contains the token embeddings [batch_size, seq_len, hidden_dim] | |
| token_embeddings = outputs[0] | |
| attention_mask = inputs["attention_mask"] | |
| # Perform mean pooling over the active tokens | |
| input_mask_expanded = np.expand_dims(attention_mask, axis=-1) | |
| input_mask_expanded = np.broadcast_to(input_mask_expanded, token_embeddings.shape) | |
| sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1) | |
| sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None) | |
| # Calculate mean pooled embedding (vector shape: [hidden_dim]) | |
| mean_pooled = (sum_embeddings / sum_mask)[0] | |
| # L2 Normalize | |
| norm = np.linalg.norm(mean_pooled) | |
| if norm > 0: | |
| mean_pooled = mean_pooled / norm | |
| self._embedding_cache[text_key] = mean_pooled | |
| return mean_pooled | |
| def match_column(self, text: str, columns: Sequence[str], threshold: float = 0.4) -> tuple[str | None, float]: | |
| """Match query text to the best column name using cosine similarity.""" | |
| if not columns: | |
| return None, 0.0 | |
| query_emb = self.get_embedding(text) | |
| best_col = None | |
| best_sim = -1.0 | |
| for col in columns: | |
| col_emb = self.get_embedding(col) | |
| sim = cosine_similarity(query_emb, col_emb) | |
| if sim > best_sim: | |
| best_sim = sim | |
| best_col = col | |
| if best_sim >= threshold: | |
| return best_col, best_sim | |
| return None, best_sim | |