Spaces:
Sleeping
Sleeping
File size: 8,237 Bytes
b336134 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 | """
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
|