"""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'" )