lib/embeddings.py

"""
lib/embeddings.py

Embedding abstraction layer. Deliberately implemented as TF-IDF (word
1-2grams) + TruncatedSVD ("classic LSA") rather than a downloaded
sentence-transformer model. This is a constraint-driven design choice, not
a quality shortcut:

  - The ranking step must run with NO network access. A sentence-transformer
    needs its weights present on disk already, which means the repo must
    either vendor ~80-400MB of model binaries or pull them at precompute
    time over the network -- one more thing to break reproducibility on a
    judge's machine at Stage 3.
  - TF-IDF + SVD is 100% stdlib-adjacent (scikit-learn only), fits in a few
    hundred KB on disk, fits the "small semantic *support* signal, not the
    main decision" role described in the brief exactly, and is bit-for-bit
    reproducible with a fixed random_state.
  - It still recovers vocabulary-level synonymy (a candidate who built a
    "recommendation system" without ever writing "RAG" still scores
    reasonably against the JD's ideal-candidate text), which is the
    specific Tier-5 trap the JD explicitly warns about.

This module exposes a tiny interface (`fit`, `transform`, `similarity_to`)
so a future sentence-transformer backend could be swapped in later without
touching callers -- this is the "embedding abstraction" piece.
"""
from __future__ import annotations
import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
from sklearn.preprocessing import normalize


class TfidfSvdEmbedder:
    def __init__(self, n_components: int = 100, random_state: int = 42):
        self.vectorizer = TfidfVectorizer(
            max_features=40_000,
            ngram_range=(1, 2),
            min_df=2,
            max_df=0.95,
            stop_words="english",
            sublinear_tf=True,
        )
        self.svd = TruncatedSVD(n_components=n_components, random_state=random_state)
        self._fitted = False

    def fit(self, texts: list[str]) -> "TfidfSvdEmbedder":
        n = len(texts)
        # Adjust min_df/max_df for small corpora to avoid sklearn error
        md = min(2, max(1, n // 3)) if n < 20 else 3
        mdf = 0.95 if n >= 20 else 1.0
        self.vectorizer.min_df = md
        self.vectorizer.max_df = mdf
        tfidf = self.vectorizer.fit_transform(texts)
        self.svd.fit(tfidf)
        self._fitted = True
        return self

    def transform(self, texts: list[str]) -> np.ndarray:
        assert self._fitted, "call fit() first"
        tfidf = self.vectorizer.transform(texts)
        emb = self.svd.transform(tfidf)
        return normalize(emb)  # unit-norm rows -> dot product == cosine sim

    def similarity_to_query(self, doc_embeddings: np.ndarray, query_text: str) -> np.ndarray:
        q = self.transform([query_text])[0]
        return doc_embeddings @ q