source/src/strata/modeling/modules.py

"""Plain-PyTorch neural modules for STRATA."""

from __future__ import annotations

import math
from typing import Literal

import torch
from torch import nn
from torch.nn import functional as F

from strata.modeling.config import StrataConfig
from strata.modeling.graph_object import GraphObject, apply_graph_object_intervention
from strata.modeling.interventions import apply_predicate_memory_intervention
from strata.modeling.outputs import GraphObjectBlockOutput, PredicateBlockOutput

GraphMode = Literal["causal_lm", "full_graph"]
ResidualScale = torch.Tensor | float | int | None


class RMSNorm(nn.Module):
    """Root-mean-square normalization without bias."""

    def __init__(self, d_model: int, eps: float = 1e-6) -> None:
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        variance = hidden_states.pow(2).mean(dim=-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
        return hidden_states * self.weight


class SwiGLU(nn.Module):
    """SwiGLU feed-forward block."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.up = nn.Linear(config.d_model, 2 * config.d_ff, bias=False)
        self.down = nn.Linear(config.d_ff, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate, value = self.up(hidden_states).chunk(2, dim=-1)
        return self.down(self.dropout(F.silu(gate) * value))


class LocalCausalSelfAttention(nn.Module):
    """Local causal multi-head attention.

    The implementation uses dense score tensors masked to a local causal band.
    It is intentionally simple and deterministic for the foundation codebase;
    optimized kernels can be introduced later behind the same interface.
    """

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.config = config
        self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False)
        self.out = nn.Linear(config.d_model, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
        graph_attention_bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch_size, seq_len, _ = hidden_states.shape
        qkv = self.qkv(hidden_states)
        query, key, value = qkv.chunk(3, dim=-1)
        query = _split_heads(query, self.config.num_heads)
        key = _split_heads(key, self.config.num_heads)
        value = _split_heads(value, self.config.num_heads)

        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(
            self.config.head_dim
        )
        mask = _local_causal_mask(
            seq_len,
            self.config.local_attention_window,
            device=hidden_states.device,
        )
        scores = scores.masked_fill(~mask.view(1, 1, seq_len, seq_len), -torch.inf)

        if attention_mask is not None:
            key_mask = attention_mask.to(torch.bool).view(batch_size, 1, 1, seq_len)
            scores = scores.masked_fill(~key_mask, -torch.inf)

        if graph_attention_bias is not None and self.config.use_attention_bias_from_graph:
            if graph_attention_bias.shape != (batch_size, seq_len, seq_len):
                raise ValueError(
                    "graph_attention_bias must have shape "
                    f"({batch_size}, {seq_len}, {seq_len}), got "
                    f"{tuple(graph_attention_bias.shape)}"
                )
            scores = scores + graph_attention_bias.unsqueeze(1)

        weights = torch.softmax(scores, dim=-1)
        weights = torch.nan_to_num(weights, nan=0.0)
        weights = self.dropout(weights)
        output = torch.matmul(weights, value)
        output = _merge_heads(output)
        output = self.out(output)
        if attention_mask is not None:
            output = output * attention_mask.to(output.dtype).unsqueeze(-1)
        return output


class LexicalValencyProposer(nn.Module):
    """Predict token-anchored linguistic candidates used by predicate memory."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.node_type = nn.Linear(config.d_model, config.node_type_vocab_size)
        self.chart_type = nn.Linear(config.d_model, config.chart_type_vocab_size)
        self.predicate_gate = nn.Linear(config.d_model, 1)
        self.candidate_key = nn.Linear(config.d_model, config.d_model, bias=False)
        self.candidate_value = nn.Linear(config.d_model, config.d_model, bias=False)

    def forward(
        self, hidden_states: torch.Tensor, *, emit_heads: bool = True
    ) -> tuple[PredicateBlockOutput, torch.Tensor, torch.Tensor]:
        # predicate_gate feeds the memory attention (forward path) and is always
        # computed; node/chart are prediction heads, computed only when emit_heads.
        predicate_gate = torch.sigmoid(self.predicate_gate(hidden_states))
        output = PredicateBlockOutput(
            node_type_logits=self.node_type(hidden_states) if emit_heads else None,
            chart_type_logits=self.chart_type(hidden_states) if emit_heads else None,
            predicate_gate=predicate_gate,
        )
        return output, self.candidate_key(hidden_states), self.candidate_value(hidden_states)


class PredicateMemoryAttention(nn.Module):
    """Causal token-to-predicate memory attention."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.config = config
        self.query = nn.Linear(config.d_model, config.d_model, bias=False)
        self.out = nn.Linear(config.d_model, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)
        self.replacement_logit = nn.Parameter(torch.tensor(0.0))
        self.edge_type = nn.Linear(config.d_model, config.graph_relation_types)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        candidate_key: torch.Tensor,
        candidate_value: torch.Tensor,
        predicate_gate: torch.Tensor,
        attention_mask: torch.Tensor | None,
        predicate_memory_bias: torch.Tensor | None,
        mode: GraphMode,
        return_edge_logits: bool,
        emit_heads: bool = True,
        predicate_memory_intervention: str = "none",
    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]:
        batch_size, seq_len, _ = hidden_states.shape
        candidate_key, candidate_value, predicate_gate = apply_predicate_memory_intervention(
            candidate_key=candidate_key,
            candidate_value=candidate_value,
            predicate_gate=predicate_gate,
            attention_mask=attention_mask,
            intervention=predicate_memory_intervention,
        )
        query = _split_heads(self.query(hidden_states), self.config.num_heads)
        key = _split_heads(candidate_key, self.config.num_heads)
        value = _split_heads(candidate_value, self.config.num_heads)

        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(
            self.config.head_dim
        )
        scores = scores + torch.log(predicate_gate.clamp_min(1e-6)).transpose(1, 2).view(
            batch_size, 1, 1, seq_len
        )
        if predicate_memory_bias is not None:
            if predicate_memory_bias.shape != (batch_size, seq_len, seq_len):
                raise ValueError(
                    "predicate_memory_bias must have shape "
                    f"({batch_size}, {seq_len}, {seq_len}), got "
                    f"{tuple(predicate_memory_bias.shape)}"
                )
            scores = scores + predicate_memory_bias.to(scores.dtype).unsqueeze(1)
        if mode == "causal_lm":
            causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=hidden_states.device).tril()
            scores = scores.masked_fill(~causal.view(1, 1, seq_len, seq_len), -torch.inf)

        if attention_mask is not None:
            key_mask = attention_mask.to(torch.bool).view(batch_size, 1, 1, seq_len)
            scores = scores.masked_fill(~key_mask, -torch.inf)

        weights = torch.softmax(scores, dim=-1)
        weights = torch.nan_to_num(weights, nan=0.0)
        weights = self.dropout(weights)
        memory_read = self.out(_merge_heads(torch.matmul(weights, value)))
        if attention_mask is not None:
            memory_read = memory_read * attention_mask.to(memory_read.dtype).unsqueeze(-1)

        edge_logits: torch.Tensor | None = None
        if return_edge_logits and emit_heads:
            pair_repr = hidden_states.unsqueeze(2) * candidate_key.unsqueeze(1)
            edge_logits = self.edge_type(pair_repr)
            if mode == "causal_lm":
                causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=hidden_states.device).tril()
                edge_logits = edge_logits.masked_fill(
                    ~causal.view(1, seq_len, seq_len, 1), -torch.inf
                )

        return memory_read, edge_logits, torch.sigmoid(self.replacement_logit)


class BoundedDeductionLayer(nn.Module):
    """Bounded differentiable closure over predicate-memory states."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.norm = RMSNorm(config.d_model)
        self.rule_mlp = nn.Sequential(
            nn.Linear(config.d_model, config.d_ff, bias=False),
            nn.SiLU(),
            nn.Linear(config.d_ff, config.d_model, bias=False),
        )
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, predicate_read: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.rule_mlp(self.norm(predicate_read)))


class GraphObjectMemoryAttention(nn.Module):
    """Explicit graph-object memory read over token-anchored graph edges."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.config = config
        self.query = nn.Linear(config.d_model, config.d_model, bias=False)
        self.key = nn.Linear(config.d_model, config.d_model, bias=False)
        self.value = nn.Linear(config.d_model, config.d_model, bias=False)
        self.out = nn.Linear(config.d_model, config.d_model, bias=False)
        self.node_type = nn.Embedding(config.graph_object_node_types, config.d_model)
        self.relation_bias = nn.Embedding(config.graph_object_relation_types, config.num_heads)
        self.relation_value = nn.Embedding(config.graph_object_relation_types, config.d_model)
        self.use_relation_conditioned_messages = config.graph_object_relation_conditioned_messages
        self.use_predicate_slot_memory = config.graph_object_predicate_slot_memory
        if self.use_relation_conditioned_messages:
            self.message_pair = nn.Linear(2 * config.d_model, config.d_model, bias=False)
            self.message_gamma = nn.Embedding(config.graph_object_relation_types, config.d_model)
            self.message_beta = nn.Embedding(config.graph_object_relation_types, config.d_model)
        if self.use_predicate_slot_memory:
            self.slot_event = nn.Linear(config.d_model, config.d_model, bias=False)
            self.slot_filler = nn.Linear(config.d_model, config.d_model, bias=False)
            self.slot_role_key = nn.Embedding(config.graph_object_relation_types, config.d_model)
            self.slot_role_gamma = nn.Embedding(config.graph_object_relation_types, config.d_model)
            self.slot_role_beta = nn.Embedding(config.graph_object_relation_types, config.d_model)
        self.relation_aux = nn.Linear(config.d_model, config.graph_object_relation_types, bias=False)
        self.src_aux = nn.Linear(config.d_model, config.d_model, bias=False)
        self.dst_aux = nn.Linear(config.d_model, config.d_model, bias=False)
        self.gate_logit = nn.Parameter(torch.tensor(float(config.graph_object_gate_init)))
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        graph_object: GraphObject | None,
        attention_mask: torch.Tensor | None,
        mode: GraphMode,
        graph_object_intervention: str = "none",
        return_aux_logits: bool = False,
    ) -> tuple[torch.Tensor, GraphObjectBlockOutput | None]:
        if graph_object is None:
            return torch.zeros_like(hidden_states), None

        graph_object = apply_graph_object_intervention(
            graph_object,
            intervention=graph_object_intervention,
            relation_vocab_size=self.config.graph_object_relation_types,
        )
        if graph_object is None:
            return torch.zeros_like(hidden_states), None

        batch_size, seq_len, _ = hidden_states.shape
        node_type = graph_object["node_type"]
        node_mask = graph_object["node_mask"].to(torch.bool)
        edge_src = graph_object["edge_src"]
        edge_dst = graph_object["edge_dst"]
        edge_rel = graph_object["edge_rel"]
        edge_slot_mask = graph_object["edge_slot_mask"].to(torch.bool)
        if node_type.shape != (batch_size, seq_len):
            raise ValueError(f"graph node_type must have shape {(batch_size, seq_len)}, got {tuple(node_type.shape)}")
        if edge_src.ndim != 2 or edge_src.shape != edge_dst.shape or edge_src.shape != edge_rel.shape:
            raise ValueError("graph edge_src/edge_dst/edge_rel must have matching shape [batch, edges]")
        if edge_src.shape[0] != batch_size:
            raise ValueError(f"graph edge batch size {edge_src.shape[0]} != hidden batch size {batch_size}")

        safe_node_type = node_type.clamp_min(0).clamp_max(self.config.graph_object_node_types - 1)
        node_repr = hidden_states + self.node_type(safe_node_type) * node_mask.to(hidden_states.dtype).unsqueeze(-1)

        query = _split_heads(self.query(hidden_states), self.config.num_heads)
        edge_count = edge_src.shape[1]
        if edge_count == 0:
            return torch.zeros_like(hidden_states), None
        safe_src = edge_src.clamp(0, seq_len - 1)
        safe_dst = edge_dst.clamp(0, seq_len - 1)
        gather = lambda x, idx: x.gather(1, idx.unsqueeze(-1).expand(batch_size, edge_count, x.shape[-1]))
        src_repr = gather(node_repr, safe_src)
        dst_repr = gather(node_repr, safe_dst)
        safe_edge_type = edge_rel.clamp_min(0).clamp_max(self.config.graph_object_relation_types - 1)
        rel_repr = self.relation_value(safe_edge_type)
        if self.use_predicate_slot_memory:
            # Relation labels are addresses, not decorations: the key names the
            # event-role slot and the value writes the endpoint filler through a
            # role-specific projection. Untyped collapse therefore destroys the
            # ARG0/ARG1 address distinction even when endpoints/topology remain.
            role_key = self.slot_role_key(safe_edge_type)
            relation_gamma = torch.tanh(self.slot_role_gamma(safe_edge_type))
            relation_beta = self.slot_role_beta(safe_edge_type)
            slot_key_repr = self.slot_event(dst_repr) + role_key
            slot_value_repr = self.slot_filler(src_repr) * (1.0 + relation_gamma) + relation_beta
            slot_repr = slot_key_repr + slot_value_repr
        elif self.use_relation_conditioned_messages:
            pair_repr = torch.cat([src_repr, dst_repr], dim=-1)
            base_message = self.message_pair(pair_repr)
            relation_gamma = torch.tanh(self.message_gamma(safe_edge_type))
            relation_beta = self.message_beta(safe_edge_type)
            slot_repr = base_message * (1.0 + relation_gamma) + relation_beta + rel_repr
            slot_key_repr = slot_repr
            slot_value_repr = slot_repr
        else:
            slot_repr = src_repr + dst_repr + rel_repr
            slot_key_repr = slot_repr
            slot_value_repr = slot_repr
        aux_output: GraphObjectBlockOutput | None = None
        if return_aux_logits:
            src_query = self.src_aux(slot_repr)
            dst_query = self.dst_aux(slot_repr)
            token_keys = hidden_states.transpose(1, 2)
            aux_output = GraphObjectBlockOutput(
                relation_logits=self.relation_aux(slot_repr),
                src_logits=torch.matmul(src_query, token_keys) / math.sqrt(self.config.d_model),
                dst_logits=torch.matmul(dst_query, token_keys) / math.sqrt(self.config.d_model),
            )
        key = _split_heads(self.key(slot_key_repr), self.config.num_heads)
        value = _split_heads(self.value(slot_value_repr), self.config.num_heads)
        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.config.head_dim)

        rel_bias = self.relation_bias(safe_edge_type).transpose(1, 2).unsqueeze(2)
        scores = scores + rel_bias

        valid_edge = edge_slot_mask.view(batch_size, 1, edge_count).expand(batch_size, seq_len, edge_count)
        if mode == "causal_lm":
            positions = torch.arange(seq_len, device=hidden_states.device).view(1, seq_len, 1)
            endpoints_visible = (safe_src.view(batch_size, 1, edge_count) <= positions) & (
                safe_dst.view(batch_size, 1, edge_count) <= positions
            )
            valid_edge = valid_edge & endpoints_visible
        if attention_mask is not None:
            query_mask = attention_mask.to(torch.bool).view(batch_size, seq_len, 1)
            src_valid = attention_mask.gather(1, safe_src).to(torch.bool).view(batch_size, 1, edge_count)
            dst_valid = attention_mask.gather(1, safe_dst).to(torch.bool).view(batch_size, 1, edge_count)
            valid_edge = valid_edge & query_mask & src_valid & dst_valid
        scores = scores.masked_fill(~valid_edge.view(batch_size, 1, seq_len, edge_count), -torch.inf)

        weights = torch.softmax(scores, dim=-1)
        weights = torch.nan_to_num(weights, nan=0.0)
        weights = self.dropout(weights)
        edge_read = torch.matmul(weights, value)
        update = self.out(_merge_heads(edge_read))
        update = update * torch.sigmoid(self.gate_logit)
        if attention_mask is not None:
            update = update * attention_mask.to(update.dtype).unsqueeze(-1)
        return update, aux_output


class PredicateBlock(nn.Module):
    """Lexical proposal, predicate memory read, and bounded closure."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.norm = RMSNorm(config.d_model)
        self.proposer = LexicalValencyProposer(config)
        self.memory = PredicateMemoryAttention(config)
        self.deduction = BoundedDeductionLayer(config)
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None,
        predicate_memory_bias: torch.Tensor | None,
        mode: GraphMode,
        return_edge_logits: bool,
        emit_heads: bool = True,
        predicate_memory_intervention: str = "none",
        predicate_memory_residual_scale: ResidualScale = None,
    ) -> tuple[torch.Tensor, PredicateBlockOutput, torch.Tensor]:
        normalized = self.norm(hidden_states)
        proposed, candidate_key, candidate_value = self.proposer(normalized, emit_heads=emit_heads)
        memory_read, edge_logits, replacement_gate = self.memory(
            normalized,
            candidate_key=candidate_key,
            candidate_value=candidate_value,
            predicate_gate=proposed.predicate_gate,
            attention_mask=attention_mask,
            predicate_memory_bias=predicate_memory_bias,
            mode=mode,
            return_edge_logits=return_edge_logits,
            emit_heads=emit_heads,
            predicate_memory_intervention=predicate_memory_intervention,
        )
        proposed.edge_logits = edge_logits
        closure = self.deduction(memory_read)
        update = replacement_gate * memory_read + (1.0 - replacement_gate) * closure
        update = _apply_residual_scale(update, predicate_memory_residual_scale, name="predicate_memory_residual_scale")
        return hidden_states + self.dropout(update), proposed, replacement_gate


class StrataDecoderBlock(nn.Module):
    """Local sequence computation plus optional STRATA predicate block."""

    def __init__(self, config: StrataConfig, *, has_predicate_block: bool) -> None:
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model)
        self.attention = LocalCausalSelfAttention(config)
        self.ffn_norm = RMSNorm(config.d_model)
        self.ffn = SwiGLU(config)
        self.predicate = PredicateBlock(config) if has_predicate_block else None
        self.graph_object = GraphObjectMemoryAttention(config) if has_predicate_block and config.use_graph_object_memory else None
        self.dropout = nn.Dropout(config.dropout)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None,
        graph_attention_bias: torch.Tensor | None,
        predicate_memory_bias: torch.Tensor | None,
        mode: GraphMode,
        return_edge_logits: bool,
        emit_heads: bool = True,
        predicate_memory_intervention: str = "none",
        predicate_memory_residual_scale: ResidualScale = None,
        graph_object: GraphObject | None = None,
        graph_object_intervention: str = "none",
        graph_object_residual_scale: ResidualScale = None,
        return_graph_object_logits: bool = False,
    ) -> tuple[torch.Tensor, PredicateBlockOutput | None, torch.Tensor | None, GraphObjectBlockOutput | None]:
        hidden_states = hidden_states + self.dropout(
            self.attention(
                self.attn_norm(hidden_states),
                attention_mask=attention_mask,
                graph_attention_bias=graph_attention_bias,
            )
        )
        hidden_states = hidden_states + self.dropout(self.ffn(self.ffn_norm(hidden_states)))
        if self.predicate is None:
            return hidden_states, None, None, None
        hidden_states, predicate_output, replacement_gate = self.predicate(
            hidden_states,
            attention_mask=attention_mask,
            predicate_memory_bias=predicate_memory_bias,
            mode=mode,
            return_edge_logits=return_edge_logits,
            emit_heads=emit_heads,
            predicate_memory_intervention=predicate_memory_intervention,
            predicate_memory_residual_scale=predicate_memory_residual_scale,
        )
        graph_object_output: GraphObjectBlockOutput | None = None
        if self.graph_object is not None and graph_object is not None:
            graph_update, graph_object_output = self.graph_object(
                hidden_states,
                graph_object=graph_object,
                attention_mask=attention_mask,
                mode=mode,
                graph_object_intervention=graph_object_intervention,
                return_aux_logits=return_graph_object_logits,
            )
            graph_update = _apply_residual_scale(graph_update, graph_object_residual_scale, name="graph_object_residual_scale")
            hidden_states = hidden_states + self.dropout(graph_update)
        return hidden_states, predicate_output, replacement_gate, graph_object_output


def _apply_residual_scale(update: torch.Tensor, scale: ResidualScale, *, name: str) -> torch.Tensor:
    if scale is None:
        return update
    if isinstance(scale, (float, int)):
        return update * float(scale)
    if scale.ndim == 0:
        return update * scale.to(update.dtype)
    batch_size = update.shape[0]
    if scale.shape == (batch_size,):
        shaped = scale.view(batch_size, 1, 1)
    elif scale.shape == (batch_size, 1):
        shaped = scale.view(batch_size, 1, 1)
    elif scale.shape == (batch_size, 1, 1):
        shaped = scale
    else:
        raise ValueError(
            f"{name} must be scalar or have shape ({batch_size},), "
            f"({batch_size}, 1), or ({batch_size}, 1, 1); got {tuple(scale.shape)}"
        )
    return update * shaped.to(update.dtype)


def _split_heads(tensor: torch.Tensor, num_heads: int) -> torch.Tensor:
    batch_size, seq_len, d_model = tensor.shape
    head_dim = d_model // num_heads
    return tensor.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)


def _merge_heads(tensor: torch.Tensor) -> torch.Tensor:
    batch_size, num_heads, seq_len, head_dim = tensor.shape
    return tensor.transpose(1, 2).contiguous().view(batch_size, seq_len, num_heads * head_dim)


def _local_causal_mask(seq_len: int, window: int, *, device: torch.device) -> torch.Tensor:
    positions = torch.arange(seq_len, device=device)
    query = positions.view(seq_len, 1)
    key = positions.view(1, seq_len)
    return (key <= query) & ((query - key) < window)