File size: 4,676 Bytes
6b6e83f 64ce412 6b6e83f 64ce412 6b6e83f 64ce412 6b6e83f 64ce412 6b6e83f 64ce412 6b6e83f 1a782a6 6b6e83f | 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 | """
embedding_engine.py
────────────────────────────────────────────────────────────────
Offline embedding strategy using TF-IDF + SVD (LSA) for
environments where sentence-transformers cannot be installed
due to disk/network constraints.
When sentence-transformers IS available (recommended for
production), swap EmbeddingEngine with SentenceTransformerEngine
below. The rest of the pipeline is identical.
"""
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import Normalizer
# ─────────────────────────────────────────────────────────────
# Baseline: TF-IDF + Latent Semantic Analysis (offline, fast)
# ─────────────────────────────────────────────────────────────
class TFIDFEmbeddingEngine:
"""
Produces 256-dim L2-normalised embeddings via TF-IDF → SVD.
Works completely offline with no extra downloads.
Optimised for English text processing.
"""
def __init__(self, n_components: int = 256):
self.n_components = n_components
self.pipeline = Pipeline([
("tfidf", TfidfVectorizer(
analyzer = "word", # word level
stop_words = "english", # filter out english stop words
max_features = 20_000,
sublinear_tf = True,
strip_accents = "unicode",
lowercase = True,
)),
("svd", TruncatedSVD(n_components=n_components, random_state=42)),
("norm", Normalizer(norm="l2")),
])
self.is_fitted = False
def fit(self, texts: list[str]) -> "TFIDFEmbeddingEngine":
self.pipeline.fit(texts)
self.is_fitted = True
return self
def encode(self, texts: list[str]) -> np.ndarray:
"""Return (N, n_components) float32 array."""
if not self.is_fitted:
raise RuntimeError("Call .fit() before .encode()")
return self.pipeline.transform(texts).astype(np.float32)
def encode_single(self, text: str) -> np.ndarray:
return self.encode([text])[0]
# ─────────────────────────────────────────────────────────────
# Upgrade path: sentence-transformers (recommended)
# ─────────────────────────────────────────────────────────────
class SentenceTransformerEngine:
"""
Drop-in replacement using all-MiniLM-L6-v2.
State-of-the-art English semantic search embedding model.
Install: pip install sentence-transformers
The model is downloaded once and cached locally (~80 MB).
Usage:
engine = SentenceTransformerEngine()
vectors = engine.encode(texts)
"""
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
try:
from sentence_transformers import SentenceTransformer
self.model = SentenceTransformer(model_name)
self.is_fitted = True
except ImportError:
raise ImportError(
"pip install sentence-transformers # then retry"
)
def fit(self, texts): # no-op: pretrained
return self
def encode(self, texts: list[str]) -> np.ndarray:
return self.model.encode(texts, convert_to_numpy=True,
show_progress_bar=False, batch_size=32)
def encode_single(self, text: str) -> np.ndarray:
return self.encode([text])[0]
def get_embedding_engine(prefer_transformer: bool = True):
"""Factory: use sentence-transformers if available, else TF-IDF+SVD."""
if prefer_transformer:
try:
engine = SentenceTransformerEngine()
print("[EmbeddingEngine] Using sentence-transformers (English).")
return engine
except (ImportError, Exception) as e:
print(f"[EmbeddingEngine] sentence-transformers unavailable ({e})."
" Falling back to TF-IDF+SVD.")
print("[EmbeddingEngine] Using TF-IDF + SVD (offline baseline).")
return TFIDFEmbeddingEngine()
|