Blame for code/graph/pat.py

"""PAT: query-conditioned, relation-aware Path-Aware Transformer."""

from __future__ import annotations

import itertools
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from enum import IntEnum

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

from graph.config import PATConfig
from graph.types import ConstraintState, GraphPath
from graph.utils import qk_normalize, soft_cap, stable_hash


class PATTokenType(IntEnum):
    QUERY = 0
    ENTITY = 1
    RELATION = 2
    EDGE = 3
    TYPE = 4
    PATH_START = 5
    PATH_STEP = 6
    PATH_END = 7
    HOP = 8
    CONSTRAINT = 9


@dataclass(frozen=True)
class PATInput:
    token_features: Tensor
    token_type_ids: Tensor
    hop_positions: Tensor
    attention_mask: Tensor


@dataclass(frozen=True)
class PATOutput:
    path_states: Tensor
    relation_sequence_states: Tensor
    transition_states: Tensor
    compatibility: Tensor
    token_states: Tensor
    endpoint_states: Tensor
    hop_states: Tensor
    constraint_compatibility: Tensor


class RelationAwareAttention(nn.Module):
    def __init__(
        self,
        hidden_dim: int,
        num_heads: int,
        num_token_types: int,
        dropout: float,
        logit_cap: float = 0.0,
        qk_norm: bool = False,
    ) -> None:
        super().__init__()
        if hidden_dim % num_heads:
            raise ValueError("PAT hidden_dim must be divisible by num_heads")
        self.logit_cap = logit_cap
        self.qk_norm = qk_norm
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.head_dim = hidden_dim // num_heads
        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
        self.pair_bias = nn.Parameter(
            torch.zeros(num_token_types, num_token_types, num_heads)
        )
        self.dropout = nn.Dropout(dropout)

    def forward(
        self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor
    ) -> Tensor:
        batch, length, _ = states.shape
        q = (
            self.q_proj(states)
            .view(batch, length, self.num_heads, self.head_dim)
            .transpose(1, 2)
        )
        k = (
            self.k_proj(states)
            .view(batch, length, self.num_heads, self.head_dim)
            .transpose(1, 2)
        )
        v = (
            self.v_proj(states)
            .view(batch, length, self.num_heads, self.head_dim)
            .transpose(1, 2)
        )
        if self.qk_norm:
            q, k = qk_normalize(q), qk_normalize(k)
        logits = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim**0.5)
        pair = self.pair_bias[
            token_type_ids.unsqueeze(-1), token_type_ids.unsqueeze(-2)
        ]
        logits = logits + pair.permute(0, 3, 1, 2)
        query_key = query_state.view(batch, self.num_heads, self.head_dim)
        logits = logits + torch.einsum("bhd,bhld->bhl", query_key, k).unsqueeze(-2) / (
            self.head_dim**0.5
        )
        logits = soft_cap(logits, self.logit_cap)
        valid = mask.to(torch.bool)
        logits = logits.masked_fill(
            ~valid[:, None, None, :], torch.finfo(logits.dtype).min
        )
        weights = torch.softmax(logits, dim=-1).to(v.dtype)
        weights = self.dropout(weights)
        result = (
            torch.matmul(weights, v)
            .transpose(1, 2)
            .contiguous()
            .view(batch, length, self.hidden_dim)
        )
        return self.out_proj(result)


class PATBlock(nn.Module):
    def __init__(
        self, config: PATConfig, logit_cap: float = 0.0, qk_norm: bool = False
    ) -> None:
        super().__init__()
        self.attn = RelationAwareAttention(
            config.hidden_dim,
            config.num_heads,
            config.num_token_types,
            config.dropout,
            logit_cap,
            qk_norm,
        )
        self.norm1 = nn.LayerNorm(config.hidden_dim)
        self.norm2 = nn.LayerNorm(config.hidden_dim)
        self.ff = nn.Sequential(
            nn.Linear(config.hidden_dim, config.hidden_dim * 4),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.hidden_dim * 4, config.hidden_dim),
        )
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self, states: Tensor, token_type_ids: Tensor, mask: Tensor, query_state: Tensor
    ) -> Tensor:
        attended = self.attn(self.norm1(states), token_type_ids, mask, query_state)
        states = states + self.dropout(attended)
        return states + self.dropout(self.ff(self.norm2(states)))


class PathAwareTransformer(nn.Module):
    """PAT operates on structured vectors, never on a textual graph dump."""

    def __init__(
        self,
        config: PATConfig,
        gradient_checkpointing: bool = False,
        logit_cap: float = 0.0,
        qk_norm: bool = False,
    ) -> None:
        super().__init__()
        self.config = config
        self.gradient_checkpointing = gradient_checkpointing
        self.token_type_embedding = nn.Embedding(
            config.num_token_types, config.hidden_dim
        )
        self.path_position_embedding = nn.Embedding(
            config.max_path_length * 5 + 5, config.hidden_dim
        )
        self.hop_embedding = nn.Embedding(config.max_path_length + 2, config.hidden_dim)
        self.type_embedding = nn.Embedding(config.type_vocab_size, config.hidden_dim)
        self.attribute_embedding = nn.Embedding(
            config.attribute_vocab_size, config.hidden_dim
        )
        self.query_gate = nn.Linear(config.hidden_dim, config.hidden_dim)
        self.blocks = nn.ModuleList(
            [PATBlock(config, logit_cap, qk_norm) for _ in range(config.num_layers)]
        )
        self.final_norm = nn.LayerNorm(config.hidden_dim)
        self.query_projection = nn.Linear(config.hidden_dim, config.hidden_dim)

    def forward(self, inputs: PATInput, query_state: Tensor) -> PATOutput:
        token_types = inputs.token_type_ids.clamp(0, self.config.num_token_types - 1)
        positions = torch.arange(
            inputs.token_features.shape[1], device=inputs.token_features.device
        )
        positions = positions.clamp_max(self.config.max_path_length * 5 + 4).unsqueeze(
            0
        )
        hops = inputs.hop_positions.clamp(0, self.config.max_path_length + 1)
        states = inputs.token_features + self.token_type_embedding(token_types)
        if query_state.ndim == 1:
            query_state = query_state.unsqueeze(0)
        if query_state.shape[0] == 1 and states.shape[0] != 1:
            # contiguous(): MPS aborts on Linear over a zero-stride expanded input
            query_state = query_state.expand(states.shape[0], -1).contiguous()
        states = (
            states + self.path_position_embedding(positions) + self.hop_embedding(hops)
        )
        attention_query = (
            query_state
            if self.config.query_conditioning
            else torch.zeros_like(query_state)
        )
        if self.config.query_conditioning:
            states = states + self.query_gate(query_state).unsqueeze(1)
        for block in self.blocks:
            if self.gradient_checkpointing and self.training:
                states = checkpoint(
                    lambda values, layer=block: layer(
                        values, token_types, inputs.attention_mask, attention_query
                    ),
                    states,
                    use_reentrant=False,
                )
            else:
                states = block(
                    states, token_types, inputs.attention_mask, attention_query
                )
        states = self.final_norm(states)
        mask = inputs.attention_mask.to(states.dtype).unsqueeze(-1)
        path_states = (states * mask).sum(1) / mask.sum(1).clamp_min(1.0)
        path_states = functional.normalize(path_states, dim=-1)
        relation_mask = (
            token_types == int(PATTokenType.RELATION)
        ) & inputs.attention_mask.bool()
        max_relations = max(1, int(relation_mask.sum(dim=1).max().item()))
        # Each row's relation tokens, left-aligned and zero-padded: one scatter
        # instead of a gather-and-pad per path.
        rows, columns = relation_mask.nonzero(as_tuple=True)
        slots = relation_mask.long().cumsum(1)[rows, columns] - 1
        relation_sequence_states = states.new_zeros(
            (states.shape[0], max_relations, states.shape[-1])
        ).index_put((rows, slots), states[rows, columns])
        transition_states = relation_sequence_states
        compatibility = (
            path_states
            * functional.normalize(self.query_projection(attention_query), dim=-1)
        ).sum(-1)
        entity_mask = (
            token_types == int(PATTokenType.ENTITY)
        ) & inputs.attention_mask.bool()
        entity_positions = (
            torch.arange(states.shape[1], device=states.device)
            .unsqueeze(0)
            .expand(states.shape[0], -1)
        )
        last_entity = (
            torch.where(
                entity_mask, entity_positions, torch.zeros_like(entity_positions)
            )
            .max(dim=1)
            .values
        )
        batch_indices = torch.arange(states.shape[0], device=states.device)
        endpoint_states = states[batch_indices, last_entity]
        hop_states = relation_sequence_states.mean(1)
        constraint_compatibility = functional.cosine_similarity(
            path_states, self.query_projection(attention_query), dim=-1
        )
        return PATOutput(
            path_states,
            relation_sequence_states,
            transition_states,
            compatibility,
            states,
            endpoint_states,
            hop_states,
            constraint_compatibility,
        )

    def build_input(
        self,
        paths: Sequence[GraphPath],
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        query_state: Tensor,
        constraint_state: ConstraintState | None = None,
    ) -> PATInput:
        if query_state.ndim == 2:
            query_state = query_state[0]
        if not paths:
            raise ValueError("PAT requires at least one candidate path")
        feature_rows: list[Tensor] = []
        type_rows: list[Tensor] = []
        hop_rows: list[Tensor] = []
        for path in paths:
            if not path.steps:
                raise ValueError("PAT cannot encode an empty path")
            features = [torch.zeros_like(query_state)]
            types = [int(PATTokenType.PATH_START)]
            hops = [0]
            if constraint_state is not None:
                constraint_signal = float(
                    constraint_state.max_hops - constraint_state.min_hops
                )
                if constraint_state.exact_hops is not None:
                    constraint_signal += float(constraint_state.exact_hops)
                features.append(
                    query_state
                    * (1.0 + constraint_signal / max(1.0, self.config.max_path_length))
                )
                types.append(int(PATTokenType.CONSTRAINT))
                hops.append(0)
            first = path.steps[0]
            features.append(node_states[first.source_id])
            types.append(int(PATTokenType.ENTITY))
            hops.append(0)
            features.append(self._type_state(first.source_type, query_state))
            types.append(int(PATTokenType.TYPE))
            hops.append(0)
            for hop, step in enumerate(path.steps, start=1):
                features.append(
                    relation_states.get(step.relation_id, torch.zeros_like(query_state))
                )
                types.append(int(PATTokenType.RELATION))
                hops.append(hop)
                features.append(
                    self._attribute_state(step.edge_attributes, query_state)
                )
                types.append(int(PATTokenType.EDGE))
                hops.append(hop)
                features.append(node_states[step.target_id])
                types.append(int(PATTokenType.ENTITY))
                hops.append(hop)
                features.append(self._type_state(step.target_type, query_state))
                types.append(int(PATTokenType.TYPE))
                hops.append(hop)
                features.append(
                    query_state * (float(hop) / max(1.0, self.config.max_path_length))
                )
                types.append(int(PATTokenType.HOP))
                hops.append(hop)
            features.append(torch.zeros_like(query_state))
            types.append(int(PATTokenType.PATH_END))
            hops.append(len(path.steps))
            feature_rows.append(torch.stack(features))
            type_rows.append(
                torch.tensor(types, dtype=torch.long, device=query_state.device)
            )
            hop_rows.append(
                torch.tensor(hops, dtype=torch.long, device=query_state.device)
            )
        length = max(row.shape[0] for row in feature_rows)
        hidden = query_state.shape[-1]
        padded_features = torch.zeros(
            (len(paths), length, hidden),
            dtype=query_state.dtype,
            device=query_state.device,
        )
        padded_types = torch.full(
            (len(paths), length),
            int(PATTokenType.PATH_END),
            dtype=torch.long,
            device=query_state.device,
        )
        padded_hops = torch.zeros(
            (len(paths), length), dtype=torch.long, device=query_state.device
        )
        padded_mask = torch.zeros(
            (len(paths), length), dtype=torch.long, device=query_state.device
        )
        for i, (feature_row, type_row, hop_row) in enumerate(
            zip(feature_rows, type_rows, hop_rows, strict=True)
        ):
            size = feature_row.shape[0]
            padded_features[i, :size] = feature_row
            padded_types[i, :size] = type_row
            padded_hops[i, :size] = hop_row
            padded_mask[i, :size] = 1
        return PATInput(padded_features, padded_types, padded_hops, padded_mask)

    def build_batch_input(
        self,
        paths: Sequence[GraphPath],
        owners: Sequence[int],
        node_rows: Sequence[Mapping[str, int]],
        relation_rows: Sequence[Mapping[str, int]],
        node_table: Tensor,
        relation_table: Tensor,
        query_states: Tensor,
        constraints: Sequence[ConstraintState | None],
    ) -> PATInput:
        """:meth:`build_input` for paths drawn from many examples at once.

        Path ``i`` belongs to example ``owners[i]``, whose node and relation
        states are rows ``node_rows[owner][id]`` of ``node_table`` and
        ``relation_rows[owner][id]`` of ``relation_table``.  Every token is a row
        of one feature table plus a multiple of its example's query state, so the
        whole batch is assembled by a single gather instead of one small tensor
        operation per token.  Token for token the result matches
        :meth:`build_input`, padded to the batch's longest path.
        """
        if not paths:
            raise ValueError("PAT requires at least one candidate path")
        max_length = max(1.0, self.config.max_path_length)
        type_ids: dict[int, int] = {}
        attribute_sets: dict[tuple[int, ...], int] = {}
        node_offset = 1
        relation_offset = node_offset + node_table.shape[0]
        # Tokens reference types and attribute sets by position until their rows
        # are known; ("t", k) and ("a", k) are resolved after the table is built.
        token_rows: list[list[int | tuple[str, int]]] = []
        token_scales: list[list[float]] = []
        token_types: list[list[int]] = []
        token_hops: list[list[int]] = []

        def type_slot(type_name: str | None) -> tuple[str, int]:
            index = stable_hash(type_name or "unknown", self.config.type_vocab_size)
            return ("t", type_ids.setdefault(index, len(type_ids)))

        def attribute_slot(attributes: Mapping[str, str]) -> int | tuple[str, int]:
            if not attributes:
                return 0
            key = tuple(
                stable_hash(f"{name}={value}", self.config.attribute_vocab_size)
                for name, value in sorted(attributes.items())
            )
            return ("a", attribute_sets.setdefault(key, len(attribute_sets)))

        for path, owner in zip(paths, owners, strict=True):
            if not path.steps:
                raise ValueError("PAT cannot encode an empty path")
            nodes, relations = node_rows[owner], relation_rows[owner]
            rows: list[int | tuple[str, int]] = [0]
            scales = [0.0]
            types = [int(PATTokenType.PATH_START)]
            hops = [0]
            constraint = constraints[owner]
            if constraint is not None:
                signal = float(constraint.max_hops - constraint.min_hops)
                if constraint.exact_hops is not None:
                    signal += float(constraint.exact_hops)
                rows.append(0)
                scales.append(1.0 + signal / max_length)
                types.append(int(PATTokenType.CONSTRAINT))
                hops.append(0)
            first = path.steps[0]
            rows += [node_offset + nodes[first.source_id], type_slot(first.source_type)]
            scales += [0.0, 0.0]
            types += [int(PATTokenType.ENTITY), int(PATTokenType.TYPE)]
            hops += [0, 0]
            for hop, step in enumerate(path.steps, start=1):
                relation = relations.get(step.relation_id)
                rows += [
                    0 if relation is None else relation_offset + relation,
                    attribute_slot(step.edge_attributes),
                    node_offset + nodes[step.target_id],
                    type_slot(step.target_type),
                    0,
                ]
                scales += [0.0, 0.0, 0.0, 0.0, float(hop) / max_length]
                types += [
                    int(PATTokenType.RELATION),
                    int(PATTokenType.EDGE),
                    int(PATTokenType.ENTITY),
                    int(PATTokenType.TYPE),
                    int(PATTokenType.HOP),
                ]
                hops += [hop] * 5
            rows.append(0)
            scales.append(0.0)
            types.append(int(PATTokenType.PATH_END))
            hops.append(len(path.steps))
            token_rows.append(rows)
            token_scales.append(scales)
            token_types.append(types)
            token_hops.append(hops)

        device = query_states.device
        type_offset = relation_offset + relation_table.shape[0]
        attribute_offset = type_offset + len(type_ids)
        parts = [
            query_states.new_zeros((1, query_states.shape[-1])),
            node_table.to(query_states.dtype),
            relation_table.to(query_states.dtype),
        ]
        if type_ids:
            parts.append(
                self.type_embedding(
                    torch.tensor(list(type_ids), dtype=torch.long, device=device)
                ).to(query_states.dtype)
            )
        if attribute_sets:
            flat = [index for key in attribute_sets for index in key]
            offsets = torch.tensor(
                [0, *itertools.accumulate(len(key) for key in attribute_sets)][:-1],
                dtype=torch.long,
                device=device,
            )
            parts.append(
                functional.embedding_bag(
                    torch.tensor(flat, dtype=torch.long, device=device),
                    self.attribute_embedding.weight,
                    offsets,
                    mode="mean",
                ).to(query_states.dtype)
            )
        table = torch.cat(parts)

        length = max(len(rows) for rows in token_rows)
        count = len(paths)
        index = torch.zeros((count, length), dtype=torch.long)
        scale = torch.zeros((count, length), dtype=query_states.dtype)
        padded_types = torch.full(
            (count, length), int(PATTokenType.PATH_END), dtype=torch.long
        )
        padded_hops = torch.zeros((count, length), dtype=torch.long)
        padded_mask = torch.zeros((count, length), dtype=torch.long)
        for row, (rows, scales, types, hops) in enumerate(
            zip(token_rows, token_scales, token_types, token_hops, strict=True)
        ):
            size = len(rows)
            index[row, :size] = torch.tensor(
                [
                    slot
                    if isinstance(slot, int)
                    else (type_offset if slot[0] == "t" else attribute_offset) + slot[1]
                    for slot in rows
                ],
                dtype=torch.long,
            )
            scale[row, :size] = torch.tensor(scales, dtype=query_states.dtype)
            padded_types[row, :size] = torch.tensor(types, dtype=torch.long)
            padded_hops[row, :size] = torch.tensor(hops, dtype=torch.long)
            padded_mask[row, :size] = 1
        index, scale = index.to(device), scale.to(device)
        owner_index = torch.tensor(list(owners), dtype=torch.long, device=device)
        features = table[index] + scale.unsqueeze(-1) * query_states[
            owner_index
        ].unsqueeze(1)
        return PATInput(
            features,
            padded_types.to(device),
            padded_hops.to(device),
            padded_mask.to(device),
        )

    def _type_state(self, type_name: str | None, query_state: Tensor) -> Tensor:
        index = stable_hash(type_name or "unknown", self.config.type_vocab_size)
        return self.type_embedding(torch.tensor(index, device=query_state.device))

    def _attribute_state(
        self, attributes: Mapping[str, str], query_state: Tensor
    ) -> Tensor:
        if not attributes:
            return torch.zeros_like(query_state)
        indices = [
            stable_hash(f"{key}={value}", self.config.attribute_vocab_size)
            for key, value in sorted(attributes.items())
        ]
        ids = torch.tensor(indices, dtype=torch.long, device=query_state.device)
        return self.attribute_embedding(ids).mean(0)

    def encode_paths(
        self,
        paths: Sequence[GraphPath],
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        query_state: Tensor,
        constraint_state: ConstraintState | None = None,
    ) -> PATOutput:
        return self(
            self.build_input(
                paths, node_states, relation_states, query_state, constraint_state
            ),
            query_state if query_state.ndim == 2 else query_state.unsqueeze(0),
        )