Blame for code/graph/text_features.py

"""Semantic text features for entity, relation and query strings.

Every string in the model was previously a **one-hot at a hashed index**:
``node_features`` and ``edge_features`` set a single 1.0 at
``stable_hash(name, feature_dim)``, and the query encoder embedded
``stable_hash(token, vocab_size)``.  Two relations with the same meaning and
different names therefore received *orthogonal* features, and a relation string never
seen in training landed on a row that carried no meaning.

That is measurable, and it was measured.  The shipped 100M checkpoint scores 0.150
exact match on the unseen-ontology holdout with PAT discrimination at 0.557 against a
0.5 floor, loses to a heuristic on HotpotQA bridge recall (0.095 against 0.518), and
loses to a released 34M model on real typed KGs -- a model whose one architectural
advantage on unfamiliar vocabulary is that it reads strings through a pretrained text
embedder.  See ``docs/graph/model_accuracy.md``.

This module supplies the missing pathway.  A featuriser maps strings to dense vectors
that preserve similarity, so ``authored_by`` and ``written_by`` land near each other
instead of orthogonally, and an unseen relation inherits the geometry of the words it
is made of.

Two backends, both deterministic and both usable without network access at inference:

:class:`HashedBagFeaturiser`
    The honest baseline and the default.  It is *not* a pretrained model: it hashes
    sub-word character n-grams into a fixed space, so it shares structure between
    strings with shared substrings (``authored_by`` and ``author_of`` overlap) while
    needing no dependency.  It is strictly better than one-hot -- which shares nothing
    -- and strictly weaker than a pretrained encoder, which knows that ``wrote`` and
    ``authored`` are related without sharing characters.

:class:`SentenceTransformerFeaturiser`
    A frozen pretrained sentence encoder.  This is the one expected to matter, and the
    one whose absence explains the transfer results.  The dependency is optional and
    the class fails loudly if it is missing rather than silently degrading.
"""

from __future__ import annotations

import hashlib
import re
from abc import ABC, abstractmethod
from collections import OrderedDict
from collections.abc import Sequence

import torch
from torch import Tensor

#: Split an identifier into word-ish pieces: ``/film/film/genre`` and
#: ``authored_by`` and ``AuthoredBy`` must all decompose into usable tokens.
_SPLIT = re.compile(r"[A-Za-z0-9]+")
_CAMEL = re.compile(r"(?<=[a-z0-9])(?=[A-Z])")


def readable_tokens(identifier: str) -> tuple[str, ...]:
    """Decompose an identifier into lowercase word tokens.

    Real knowledge-graph names are punctuation-heavy (``/film/film/genre``,
    ``_hypernym``) and synthetic ones are snake or camel case.  Both must reduce to the
    same kind of token stream, or the pretrained encoder sees noise.
    """
    spaced = _CAMEL.sub(" ", identifier)
    tokens = tuple(piece.lower() for piece in _SPLIT.findall(spaced))
    return tokens or ("unknown",)


def readable_text(identifier: str) -> str:
    """The string handed to a text encoder for an identifier."""
    return " ".join(readable_tokens(identifier))


#: Cached vectors kept per featuriser.  Synthetic entity names are mostly
#: unique, so an unbounded cache grows with every example trained on: a 12M
#: example epoch would hold hundreds of millions of vectors.
DEFAULT_CACHE_SIZE = 1_000_000

#: Strings per sentence-encoder forward.  Identifiers are a few tokens long, so
#: large batches stay small in memory.
ENCODE_BATCH_SIZE = 1024


class TextFeaturiser(ABC):
    """Maps strings to dense vectors of a fixed width, with an LRU cache."""

    def __init__(self, cache_size: int = DEFAULT_CACHE_SIZE) -> None:
        if cache_size <= 0:
            raise ValueError("cache_size must be positive")
        self._cache: OrderedDict[str, Tensor] = OrderedDict()
        self._cache_size = cache_size

    @property
    @abstractmethod
    def dimension(self) -> int:
        """Width of the produced vectors."""

    @abstractmethod
    def _encode(self, texts: Sequence[str]) -> Tensor:
        """Encode uncached strings; shape ``(len(texts), dimension)``."""

    def encode(self, texts: Sequence[str]) -> Tensor:
        """Encode ``texts``, reusing cached vectors.

        Entity and relation strings repeat heavily inside one graph and across a
        batch, so caching is what keeps this affordable rather than an optimisation.
        """
        if not texts:
            return torch.zeros((0, self.dimension))
        unique = list(dict.fromkeys(texts))
        missing = [text for text in unique if text not in self._cache]
        for text in unique:
            if text in self._cache:
                self._cache.move_to_end(text)
        if missing:
            encoded = self._encode(missing)
            if encoded.shape != (len(missing), self.dimension):
                raise RuntimeError(
                    f"{type(self).__name__} returned {tuple(encoded.shape)}, "
                    f"expected {(len(missing), self.dimension)}"
                )
            for text, vector in zip(
                missing, encoded.detach().float().cpu(), strict=True
            ):
                self._cache[text] = vector
        found = torch.stack([self._cache[text] for text in texts])
        while len(self._cache) > self._cache_size:
            self._cache.popitem(last=False)
        return found

    def prefetch(self, texts: Sequence[str]) -> None:
        """Encode every uncached string in one call.

        A training batch meets its strings in dozens of small ``encode`` calls,
        one per graph and query; filling the cache up front turns those into
        cache hits and the encoder into one batched forward.
        """
        missing = [text for text in dict.fromkeys(texts) if text not in self._cache]
        if missing:
            self.encode(missing)

    def prefetch_identifiers(self, identifiers: Sequence[str]) -> None:
        """:meth:`prefetch` for identifiers, made readable as in ``encode_identifiers``."""
        self.prefetch([readable_text(name) for name in identifiers])

    def encode_identifiers(self, identifiers: Sequence[str]) -> Tensor:
        """Encode identifiers after making them readable."""
        return self.encode([readable_text(name) for name in identifiers])


class HashedBagFeaturiser(TextFeaturiser):
    """Character n-gram bag, hashed into ``dimension`` and L2-normalised.

    The default, and deliberately not presented as a solution.  It shares structure
    between strings with shared substrings, which one-hot hashing does not, but it
    cannot know that ``wrote`` and ``authored`` are related.  It exists so the semantic
    pathway has a dependency-free baseline to be compared against -- otherwise any
    improvement from a pretrained encoder could not be separated from the improvement
    of simply having dense features at all.
    """

    def __init__(self, dimension: int = 256, ngram_range: tuple[int, int] = (3, 5)):
        super().__init__()
        if dimension <= 0:
            raise ValueError("dimension must be positive")
        low, high = ngram_range
        if low <= 0 or high < low:
            raise ValueError(f"invalid ngram_range {ngram_range}")
        self._dimension = dimension
        self._ngram_range = ngram_range

    @property
    def dimension(self) -> int:
        return self._dimension

    def _encode(self, texts: Sequence[str]) -> Tensor:
        low, high = self._ngram_range
        out = torch.zeros((len(texts), self._dimension), dtype=torch.float32)
        for row, text in enumerate(texts):
            padded = f" {text} "
            for size in range(low, high + 1):
                for start in range(max(0, len(padded) - size + 1)):
                    gram = padded[start : start + size]
                    digest = hashlib.blake2b(gram.encode("utf-8"), digest_size=8)
                    index = int.from_bytes(digest.digest(), "big") % self._dimension
                    out[row, index] += 1.0
        norms = out.norm(dim=-1, keepdim=True).clamp_min(1e-6)
        return out / norms


class SentenceTransformerFeaturiser(TextFeaturiser):
    """A frozen pretrained sentence encoder.

    Loaded lazily and never trained: the point is to import semantics the synthetic
    corpus cannot contain, and fine-tuning it on that corpus would undo exactly that.
    """

    def __init__(
        self,
        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
        *,
        device: str = "cpu",
        truncate_dim: int | None = None,
    ) -> None:
        super().__init__()
        # The encoder is frozen and runs once per unseen string; on CPU it was
        # 98% of a 100M training step, so use the GPU whenever there is one.
        if device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError as error:  # pragma: no cover - optional dependency
            raise RuntimeError(
                "SentenceTransformerFeaturiser needs sentence-transformers; "
                "install it or use HashedBagFeaturiser"
            ) from error
        # An explicit device context overrides an enclosing ``torch.device("meta")``:
        # parameter counting builds the model on meta, and from_pretrained refuses
        # to load real weights there.  The encoder is frozen and not a submodule,
        # so it always belongs on a real device.
        with torch.device(device):
            self._model = SentenceTransformer(
                model_name, device=device, truncate_dim=truncate_dim
            )
        self._model.eval()
        for parameter in self._model.parameters():
            parameter.requires_grad_(False)
        # Renamed in newer sentence-transformers; support both rather than pinning.
        measure = getattr(
            self._model,
            "get_embedding_dimension",
            getattr(self._model, "get_sentence_embedding_dimension", None),
        )
        if measure is None:  # pragma: no cover - depends on library version
            raise RuntimeError(
                "cannot determine the embedding dimension of "
                f"{model_name!r}; unsupported sentence-transformers version"
            )
        self._dimension = int(measure())

    @property
    def dimension(self) -> int:
        return self._dimension

    def _encode(self, texts: Sequence[str]) -> Tensor:
        with torch.no_grad():
            # The library default of 32 strings per forward split one training
            # batch's few thousand identifiers into hundreds of encoder calls.
            return self._model.encode(
                list(texts),
                batch_size=ENCODE_BATCH_SIZE,
                convert_to_tensor=True,
                normalize_embeddings=True,
                show_progress_bar=False,
            )


def build_featuriser(
    kind: str, dimension: int, *, device: str = "cpu", model_name: str | None = None
) -> TextFeaturiser:
    """Construct a featuriser by name, failing loudly on an unknown kind."""
    if kind == "hashed-bag":
        return HashedBagFeaturiser(dimension=dimension)
    if kind == "sentence-transformer":
        return SentenceTransformerFeaturiser(
            model_name or "sentence-transformers/all-MiniLM-L6-v2",
            device=device,
            truncate_dim=dimension,
        )
    raise ValueError(
        f"unknown text featuriser {kind!r}; expected 'hashed-bag' or "
        "'sentence-transformer'"
    )