Upload folder using huggingface_hub

3e363e2 verified 13 days ago

12 kB

"""Text + structured query encoder."""

from __future__ import annotations

from collections.abc import Sequence
from dataclasses import dataclass

import torch
from torch import Tensor, nn
from torch.nn import functional
from torch.utils.checkpoint import checkpoint

from graph.config import GraphConfig
from graph.text_features import TextFeaturiser, readable_text
from graph.types import GraphQuery
from graph.utils import stable_hash, tokenize


@dataclass(frozen=True)
class QueryEncoding:
    pooled: Tensor
    token_states: Tensor
    token_ids: Tensor
    attention_mask: Tensor


class QueryEncoder(nn.Module):
    """Encodes the query text and its structured fields.
    ``featuriser`` replaces the hashed token embedding with dense semantic vectors.
    The transformer above it is unchanged -- only the source of the token vectors moves
    -- because the defect is not the architecture but that a token's vector came from
    ``stable_hash`` into a table learned on synthetic text alone, leaving a real
    relation string such as ``/film/film/genre`` on a row that never carried meaning.
    """

    def __init__(
        self, config: GraphConfig, featuriser: TextFeaturiser | None = None
    ) -> None:
        super().__init__()
        self.config = config
        self.featuriser = featuriser
        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
        # Projections are trained; the featuriser itself stays frozen, so the semantics
        # imported from pretraining are not overwritten by the synthetic corpus.
        self.text_projection = (
            nn.Linear(featuriser.dimension, config.hidden_dim)
            if featuriser is not None
            else None
        )
        self.position_embedding = nn.Embedding(
            config.max_query_tokens, config.hidden_dim
        )
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=config.hidden_dim,
            nhead=config.num_reasoning_heads,
            dim_feedforward=config.hidden_dim * 4,
            dropout=config.dropout,
            batch_first=True,
            norm_first=True,
        )
        # Pre-LN is intentional for deep GRM optimization.  PyTorch's nested
        # tensor inference fast path is incompatible with norm_first=True; make
        # that choice explicit so construction does not emit a warning on every
        # model instantiation.
        self.transformer = nn.TransformerEncoder(
            encoder_layer,
            num_layers=config.num_reasoning_layers,
            enable_nested_tensor=False,
        )
        self.structure_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.norm = nn.LayerNorm(config.hidden_dim)

    def encode_text(self, text: str | None) -> tuple[Tensor, Tensor]:
        tokens = tokenize(text)[: self.config.max_query_tokens]
        device = self.token_embedding.weight.device
        ids = torch.tensor(
            [stable_hash(token, self.config.vocab_size) for token in tokens],
            dtype=torch.long,
            device=device,
        )
        mask = ids.new_ones(ids.shape, dtype=ids.dtype)
        return ids, mask

    def embed_tokens(self, texts: Sequence[str]) -> Tensor:
        """Dense vectors for query tokens, projected into the model width."""
        if self.featuriser is None or self.text_projection is None:
            raise RuntimeError("embed_tokens requires a featuriser")
        device = self.text_projection.weight.device
        vectors = self.featuriser.encode(list(texts)).to(device)
        return self.text_projection(vectors)

    def forward(
        self,
        token_ids: Tensor,
        attention_mask: Tensor | None = None,
        token_embeddings: Tensor | None = None,
    ) -> QueryEncoding:
        device = self.token_embedding.weight.device
        token_ids = token_ids.to(device)
        if token_ids.ndim == 1:
            token_ids = token_ids.unsqueeze(0)
        if attention_mask is None:
            attention_mask = token_ids.new_ones(token_ids.shape)
        elif attention_mask.ndim == 1:
            attention_mask = attention_mask.unsqueeze(0)
        attention_mask = attention_mask.to(device)
        positions = torch_arange_like(token_ids, token_ids.shape[1])
        base = (
            token_embeddings.to(device)
            if token_embeddings is not None
            else self.token_embedding(token_ids)
        )
        if base.ndim == 2:
            base = base.unsqueeze(0)
        states = base + self.position_embedding(positions)
        padding_mask = attention_mask == 0
        if self.config.training.gradient_checkpointing and self.training:
            states = checkpoint(
                lambda values: self.transformer(
                    values, src_key_padding_mask=padding_mask
                ),
                states,
                use_reentrant=False,
            )
        else:
            states = self.transformer(states, src_key_padding_mask=padding_mask)
        weights = attention_mask.to(states.dtype).unsqueeze(-1)
        pooled = self.norm((states * weights).sum(1) / weights.sum(1).clamp_min(1.0))
        return QueryEncoding(pooled, states, token_ids, attention_mask)

    def encode(self, query: GraphQuery) -> QueryEncoding:
        text = query.text or query.canonical_text()
        token_ids, mask = self.encode_text(text)
        token_embeddings = None
        if self.featuriser is not None:
            tokens = tokenize(text)[: self.config.max_query_tokens]
            token_embeddings = self.embed_tokens(tokens)
        encoding = self.forward(token_ids, mask, token_embeddings)
        structured = _structured_fields(query)
        if structured:
            if self.featuriser is not None:
                # The structured fields name relations and types too, so they get the
                # same semantic treatment rather than a second hashed table.
                structure = self.embed_tokens(
                    [readable_text(item) for item in structured]
                ).mean(0)
            else:
                structure_ids = (
                    Tensor(
                        [
                            stable_hash(item, self.config.vocab_size)
                            for item in structured
                        ]
                    )
                    .long()
                    .to(encoding.pooled.device)
                )
                structure = self.structure_embedding(structure_ids).mean(0)
            pooled = functional.normalize(
                encoding.pooled + structure.unsqueeze(0), dim=-1
            )
            return QueryEncoding(
                pooled,
                encoding.token_states,
                encoding.token_ids,
                encoding.attention_mask,
            )
        return encoding

    def encode_batch(self, queries: Sequence[GraphQuery]) -> list[QueryEncoding]:
        """:meth:`encode` for many queries in one transformer pass.
        Token sequences are right-padded and masked, so each row computes what
        :meth:`encode` computes for that query alone.  A query with no tokens is
        kept out of the padded batch: an all-masked row makes attention NaN, and
        :meth:`encode` gives it the pooled state of an empty sequence instead.
        """
        device = self.token_embedding.weight.device
        texts = [query.text or query.canonical_text() for query in queries]
        tokens = [tokenize(text)[: self.config.max_query_tokens] for text in texts]
        structured = [_structured_fields(query) for query in queries]
        if self.featuriser is not None:
            self.featuriser.prefetch(
                [token for row in tokens for token in row]
                + [readable_text(item) for row in structured for item in row]
            )
        hidden = self.config.hidden_dim
        rows = [index for index, row in enumerate(tokens) if row]
        length = max((len(tokens[index]) for index in rows), default=0)
        token_ids = torch.zeros((len(queries), length), dtype=torch.long)
        mask = torch.zeros((len(queries), length), dtype=torch.long)
        for index in rows:
            ids = [
                stable_hash(token, self.config.vocab_size) for token in tokens[index]
            ]
            token_ids[index, : len(ids)] = torch.tensor(ids, dtype=torch.long)
            mask[index, : len(ids)] = 1
        token_ids, mask = token_ids.to(device), mask.to(device)

        pooled = self.norm(torch.zeros((len(queries), hidden), device=device))
        token_states = torch.zeros((len(queries), length, hidden), device=device)
        if rows:
            selected = torch.tensor(rows, dtype=torch.long, device=device)
            embeddings = None
            if self.featuriser is not None:
                flat = self.embed_tokens([t for index in rows for t in tokens[index]])
                embeddings = flat.new_zeros((len(rows), length, hidden))
                row_index, column_index = mask[selected].nonzero(as_tuple=True)
                embeddings = embeddings.index_put((row_index, column_index), flat)
            encoding = self.forward(token_ids[selected], mask[selected], embeddings)
            pooled = pooled.index_copy(0, selected, encoding.pooled.to(pooled.dtype))
            token_states = token_states.index_copy(
                0, selected, encoding.token_states.to(token_states.dtype)
            )

        with_structure = [index for index, row in enumerate(structured) if row]
        if with_structure:
            items = [item for index in with_structure for item in structured[index]]
            owner = torch.tensor(
                [
                    position
                    for position, index in enumerate(with_structure)
                    for _ in structured[index]
                ],
                dtype=torch.long,
                device=device,
            )
            if self.featuriser is not None:
                vectors = self.embed_tokens([readable_text(item) for item in items])
            else:
                vectors = self.structure_embedding(
                    torch.tensor(
                        [stable_hash(item, self.config.vocab_size) for item in items],
                        dtype=torch.long,
                        device=device,
                    )
                )
            sums = vectors.new_zeros((len(with_structure), vectors.shape[-1]))
            sums = sums.index_add(0, owner, vectors)
            counts = torch.bincount(owner, minlength=len(with_structure))
            structure = sums / counts.to(sums.dtype).unsqueeze(-1)
            selected = torch.tensor(with_structure, dtype=torch.long, device=device)
            updated = functional.normalize(pooled[selected] + structure, dim=-1)
            pooled = pooled.index_copy(0, selected, updated.to(pooled.dtype))
        return [
            QueryEncoding(
                pooled[index : index + 1],
                token_states[index : index + 1, : len(tokens[index])],
                token_ids[index : index + 1, : len(tokens[index])],
                mask[index : index + 1, : len(tokens[index])],
            )
            for index in range(len(queries))
        ]


def _structured_fields(query: GraphQuery) -> list[str]:
    """The structured query fields :meth:`QueryEncoder.encode` pools."""
    return [
        *query.start_entities,
        *query.start_entity_ids,
        *query.allowed_relations,
        *query.required_relation_sequence,
        *query.forbidden_relations,
        *query.required_types,
        *query.forbidden_types,
        *query.required_intermediate_types,
        query.operation,
        query.answer_mode,
        *query.target_entity_ids,
    ]


def torch_arange_like(reference: Tensor, length: int) -> Tensor:
    return (
        torch.arange(length, device=reference.device)
        .unsqueeze(0)
        .expand(reference.shape[0], -1)
    )

	"""Text + structured query encoder."""

	from __future__ import annotations

	from collections.abc import Sequence
	from dataclasses import dataclass

	import torch
	from torch import Tensor, nn
	from torch.nn import functional
	from torch.utils.checkpoint import checkpoint

	from graph.config import GraphConfig
	from graph.text_features import TextFeaturiser, readable_text
	from graph.types import GraphQuery
	from graph.utils import stable_hash, tokenize


	@dataclass(frozen=True)
	class QueryEncoding:
	pooled: Tensor
	token_states: Tensor
	token_ids: Tensor
	attention_mask: Tensor


	class QueryEncoder(nn.Module):
	"""Encodes the query text and its structured fields.

	``featuriser`` replaces the hashed token embedding with dense semantic vectors.
	The transformer above it is unchanged -- only the source of the token vectors moves
	-- because the defect is not the architecture but that a token's vector came from
	``stable_hash`` into a table learned on synthetic text alone, leaving a real
	relation string such as ``/film/film/genre`` on a row that never carried meaning.
	"""

	def __init__(
	self, config: GraphConfig, featuriser: TextFeaturiser \| None = None
	) -> None:
	super().__init__()
	self.config = config
	self.featuriser = featuriser
	self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
	# Projections are trained; the featuriser itself stays frozen, so the semantics
	# imported from pretraining are not overwritten by the synthetic corpus.
	self.text_projection = (
	nn.Linear(featuriser.dimension, config.hidden_dim)
	if featuriser is not None
	else None
	)
	self.position_embedding = nn.Embedding(
	config.max_query_tokens, config.hidden_dim
	)
	encoder_layer = nn.TransformerEncoderLayer(
	d_model=config.hidden_dim,
	nhead=config.num_reasoning_heads,
	dim_feedforward=config.hidden_dim * 4,
	dropout=config.dropout,
	batch_first=True,
	norm_first=True,
	)
	# Pre-LN is intentional for deep GRM optimization. PyTorch's nested
	# tensor inference fast path is incompatible with norm_first=True; make
	# that choice explicit so construction does not emit a warning on every
	# model instantiation.
	self.transformer = nn.TransformerEncoder(
	encoder_layer,
	num_layers=config.num_reasoning_layers,
	enable_nested_tensor=False,
	)
	self.structure_embedding = nn.Embedding(config.vocab_size, config.hidden_dim)
	self.norm = nn.LayerNorm(config.hidden_dim)

	def encode_text(self, text: str \| None) -> tuple[Tensor, Tensor]:
	tokens = tokenize(text)[: self.config.max_query_tokens]
	device = self.token_embedding.weight.device
	ids = torch.tensor(
	[stable_hash(token, self.config.vocab_size) for token in tokens],
	dtype=torch.long,
	device=device,
	)
	mask = ids.new_ones(ids.shape, dtype=ids.dtype)
	return ids, mask

	def embed_tokens(self, texts: Sequence[str]) -> Tensor:
	"""Dense vectors for query tokens, projected into the model width."""
	if self.featuriser is None or self.text_projection is None:
	raise RuntimeError("embed_tokens requires a featuriser")
	device = self.text_projection.weight.device
	vectors = self.featuriser.encode(list(texts)).to(device)
	return self.text_projection(vectors)

	def forward(
	self,
	token_ids: Tensor,
	attention_mask: Tensor \| None = None,
	token_embeddings: Tensor \| None = None,
	) -> QueryEncoding:
	device = self.token_embedding.weight.device
	token_ids = token_ids.to(device)
	if token_ids.ndim == 1:
	token_ids = token_ids.unsqueeze(0)
	if attention_mask is None:
	attention_mask = token_ids.new_ones(token_ids.shape)
	elif attention_mask.ndim == 1:
	attention_mask = attention_mask.unsqueeze(0)
	attention_mask = attention_mask.to(device)
	positions = torch_arange_like(token_ids, token_ids.shape[1])
	base = (
	token_embeddings.to(device)
	if token_embeddings is not None
	else self.token_embedding(token_ids)
	)
	if base.ndim == 2:
	base = base.unsqueeze(0)
	states = base + self.position_embedding(positions)
	padding_mask = attention_mask == 0
	if self.config.training.gradient_checkpointing and self.training:
	states = checkpoint(
	lambda values: self.transformer(
	values, src_key_padding_mask=padding_mask
	),
	states,
	use_reentrant=False,
	)
	else:
	states = self.transformer(states, src_key_padding_mask=padding_mask)
	weights = attention_mask.to(states.dtype).unsqueeze(-1)
	pooled = self.norm((states * weights).sum(1) / weights.sum(1).clamp_min(1.0))
	return QueryEncoding(pooled, states, token_ids, attention_mask)

	def encode(self, query: GraphQuery) -> QueryEncoding:
	text = query.text or query.canonical_text()
	token_ids, mask = self.encode_text(text)
	token_embeddings = None
	if self.featuriser is not None:
	tokens = tokenize(text)[: self.config.max_query_tokens]
	token_embeddings = self.embed_tokens(tokens)
	encoding = self.forward(token_ids, mask, token_embeddings)
	structured = _structured_fields(query)
	if structured:
	if self.featuriser is not None:
	# The structured fields name relations and types too, so they get the
	# same semantic treatment rather than a second hashed table.
	structure = self.embed_tokens(
	[readable_text(item) for item in structured]
	).mean(0)
	else:
	structure_ids = (
	Tensor(
	[
	stable_hash(item, self.config.vocab_size)
	for item in structured
	]
	)
	.long()
	.to(encoding.pooled.device)
	)
	structure = self.structure_embedding(structure_ids).mean(0)
	pooled = functional.normalize(
	encoding.pooled + structure.unsqueeze(0), dim=-1
	)
	return QueryEncoding(
	pooled,
	encoding.token_states,
	encoding.token_ids,
	encoding.attention_mask,
	)
	return encoding

	def encode_batch(self, queries: Sequence[GraphQuery]) -> list[QueryEncoding]:
	""":meth:`encode` for many queries in one transformer pass.

	Token sequences are right-padded and masked, so each row computes what
	:meth:`encode` computes for that query alone. A query with no tokens is
	kept out of the padded batch: an all-masked row makes attention NaN, and
	:meth:`encode` gives it the pooled state of an empty sequence instead.
	"""
	device = self.token_embedding.weight.device
	texts = [query.text or query.canonical_text() for query in queries]
	tokens = [tokenize(text)[: self.config.max_query_tokens] for text in texts]
	structured = [_structured_fields(query) for query in queries]
	if self.featuriser is not None:
	self.featuriser.prefetch(
	[token for row in tokens for token in row]
	+ [readable_text(item) for row in structured for item in row]
	)
	hidden = self.config.hidden_dim
	rows = [index for index, row in enumerate(tokens) if row]
	length = max((len(tokens[index]) for index in rows), default=0)
	token_ids = torch.zeros((len(queries), length), dtype=torch.long)
	mask = torch.zeros((len(queries), length), dtype=torch.long)
	for index in rows:
	ids = [
	stable_hash(token, self.config.vocab_size) for token in tokens[index]
	]
	token_ids[index, : len(ids)] = torch.tensor(ids, dtype=torch.long)
	mask[index, : len(ids)] = 1
	token_ids, mask = token_ids.to(device), mask.to(device)

	pooled = self.norm(torch.zeros((len(queries), hidden), device=device))
	token_states = torch.zeros((len(queries), length, hidden), device=device)
	if rows:
	selected = torch.tensor(rows, dtype=torch.long, device=device)
	embeddings = None
	if self.featuriser is not None:
	flat = self.embed_tokens([t for index in rows for t in tokens[index]])
	embeddings = flat.new_zeros((len(rows), length, hidden))
	row_index, column_index = mask[selected].nonzero(as_tuple=True)
	embeddings = embeddings.index_put((row_index, column_index), flat)
	encoding = self.forward(token_ids[selected], mask[selected], embeddings)
	pooled = pooled.index_copy(0, selected, encoding.pooled.to(pooled.dtype))
	token_states = token_states.index_copy(
	0, selected, encoding.token_states.to(token_states.dtype)
	)

	with_structure = [index for index, row in enumerate(structured) if row]
	if with_structure:
	items = [item for index in with_structure for item in structured[index]]
	owner = torch.tensor(
	[
	position
	for position, index in enumerate(with_structure)
	for _ in structured[index]
	],
	dtype=torch.long,
	device=device,
	)
	if self.featuriser is not None:
	vectors = self.embed_tokens([readable_text(item) for item in items])
	else:
	vectors = self.structure_embedding(
	torch.tensor(
	[stable_hash(item, self.config.vocab_size) for item in items],
	dtype=torch.long,
	device=device,
	)
	)
	sums = vectors.new_zeros((len(with_structure), vectors.shape[-1]))
	sums = sums.index_add(0, owner, vectors)
	counts = torch.bincount(owner, minlength=len(with_structure))
	structure = sums / counts.to(sums.dtype).unsqueeze(-1)
	selected = torch.tensor(with_structure, dtype=torch.long, device=device)
	updated = functional.normalize(pooled[selected] + structure, dim=-1)
	pooled = pooled.index_copy(0, selected, updated.to(pooled.dtype))
	return [
	QueryEncoding(
	pooled[index : index + 1],
	token_states[index : index + 1, : len(tokens[index])],
	token_ids[index : index + 1, : len(tokens[index])],
	mask[index : index + 1, : len(tokens[index])],
	)
	for index in range(len(queries))
	]


	def _structured_fields(query: GraphQuery) -> list[str]:
	"""The structured query fields :meth:`QueryEncoder.encode` pools."""
	return [
	*query.start_entities,
	*query.start_entity_ids,
	*query.allowed_relations,
	*query.required_relation_sequence,
	*query.forbidden_relations,
	*query.required_types,
	*query.forbidden_types,
	*query.required_intermediate_types,
	query.operation,
	query.answer_mode,
	*query.target_entity_ids,
	]


	def torch_arange_like(reference: Tensor, length: int) -> Tensor:
	return (
	torch.arange(length, device=reference.device)
	.unsqueeze(0)
	.expand(reference.shape[0], -1)
	)