source/src/strata/modeling/model.py

"""STRATA decoder language model."""

from __future__ import annotations

from pathlib import Path

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
from strata.modeling.modules import GraphMode, RMSNorm, StrataDecoderBlock
from strata.modeling.outputs import GraphObjectBlockOutput, PredicateBlockOutput, StrataCausalLMOutput, StrataModelOutput


class StrataModel(nn.Module):
    """Decoder backbone with local attention and predicate-memory blocks."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.config = config
        self.token_embeddings = nn.Embedding(config.vocab_size, config.d_model)
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.d_model
        )
        self.blocks = nn.ModuleList(
            [
                StrataDecoderBlock(
                    config,
                    has_predicate_block=(layer_index + 1) % config.predicate_block_every == 0,
                )
                for layer_index in range(config.num_layers)
            ]
        )
        # Only the deepest predicate block emits prediction heads (the ones the
        # losses/eval consume) unless the config restores all-block heads.
        predicate_indices = [
            i for i in range(config.num_layers) if (i + 1) % config.predicate_block_every == 0
        ]
        self._final_predicate_index = predicate_indices[-1] if predicate_indices else -1
        self.final_norm = RMSNorm(config.d_model)
        self.dropout = nn.Dropout(config.dropout)
        self.apply(self._init_weights)

    def forward(
        self,
        input_ids: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
        graph_attention_bias: torch.Tensor | None = None,
        predicate_memory_bias: torch.Tensor | None = None,
        mode: GraphMode = "causal_lm",
        return_edge_logits: bool = False,
        predicate_memory_intervention: str = "none",
        predicate_memory_residual_scale: torch.Tensor | None = None,
        graph_object: GraphObject | None = None,
        graph_object_intervention: str = "none",
        graph_object_residual_scale: torch.Tensor | float | int | None = None,
        return_graph_object_logits: bool = False,
    ) -> StrataModelOutput:
        if input_ids.ndim != 2:
            raise ValueError(f"input_ids must have shape [batch, seq], got {tuple(input_ids.shape)}")
        batch_size, seq_len = input_ids.shape
        if seq_len > self.config.max_position_embeddings:
            raise ValueError(
                f"sequence length {seq_len} exceeds max_position_embeddings "
                f"{self.config.max_position_embeddings}"
            )
        if mode not in {"causal_lm", "full_graph"}:
            raise ValueError("mode must be 'causal_lm' or 'full_graph'")
        if attention_mask is not None and attention_mask.shape != input_ids.shape:
            raise ValueError(
                f"attention_mask must match input_ids shape {tuple(input_ids.shape)}, "
                f"got {tuple(attention_mask.shape)}"
            )
        if predicate_memory_bias is not None and predicate_memory_bias.shape != (batch_size, seq_len, seq_len):
            raise ValueError(
                f"predicate_memory_bias must have shape ({batch_size}, {seq_len}, {seq_len}), "
                f"got {tuple(predicate_memory_bias.shape)}"
            )
        _validate_residual_scale(predicate_memory_residual_scale, batch_size, name="predicate_memory_residual_scale")
        _validate_residual_scale(graph_object_residual_scale, batch_size, name="graph_object_residual_scale")

        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        positions = positions.expand(batch_size, seq_len)
        hidden_states = self.token_embeddings(input_ids) + self.position_embeddings(positions)
        hidden_states = self.dropout(hidden_states)
        if attention_mask is not None:
            hidden_states = hidden_states * attention_mask.to(hidden_states.dtype).unsqueeze(-1)

        predicate_outputs: list[PredicateBlockOutput] = []
        graph_object_outputs: list[GraphObjectBlockOutput] = []
        replacement_gates: list[torch.Tensor] = []
        for layer_index, block in enumerate(self.blocks):
            emit_heads = self.config.emit_all_block_graph_heads or layer_index == self._final_predicate_index
            hidden_states, predicate_output, replacement_gate, graph_object_output = block(
                hidden_states,
                attention_mask=attention_mask,
                graph_attention_bias=graph_attention_bias,
                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=graph_object,
                graph_object_intervention=graph_object_intervention,
                graph_object_residual_scale=graph_object_residual_scale,
                return_graph_object_logits=return_graph_object_logits and emit_heads,
            )
            if predicate_output is not None:
                predicate_outputs.append(predicate_output)
            if graph_object_output is not None:
                graph_object_outputs.append(graph_object_output)
            if replacement_gate is not None:
                replacement_gates.append(replacement_gate.reshape(1))

        hidden_states = self.final_norm(hidden_states)
        if replacement_gates:
            gates = torch.cat(replacement_gates)
        else:
            gates = torch.empty(0, device=input_ids.device)
        return StrataModelOutput(
            last_hidden_state=hidden_states,
            predicate_outputs=tuple(predicate_outputs),
            graph_object_outputs=tuple(graph_object_outputs),
            attention_replacement_gates=gates,
        )

    def _init_weights(self, module: nn.Module) -> None:
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)


class StrataForCausalLM(nn.Module):
    """STRATA decoder with tied causal language-modeling head."""

    def __init__(self, config: StrataConfig) -> None:
        super().__init__()
        self.config = config
        self.model = StrataModel(config)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.token_embeddings.weight
        else:
            nn.init.normal_(
                self.lm_head.weight,
                mean=0.0,
                std=self.config.initializer_range,
            )

    def forward(
        self,
        input_ids: torch.Tensor,
        *,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        graph_attention_bias: torch.Tensor | None = None,
        predicate_memory_bias: torch.Tensor | None = None,
        mode: GraphMode = "causal_lm",
        return_edge_logits: bool = False,
        predicate_memory_intervention: str = "none",
        predicate_memory_residual_scale: torch.Tensor | None = None,
        graph_object: GraphObject | None = None,
        graph_object_intervention: str = "none",
        graph_object_residual_scale: torch.Tensor | float | int | None = None,
        return_graph_object_logits: bool = False,
    ) -> StrataCausalLMOutput:
        model_output = self.model(
            input_ids,
            attention_mask=attention_mask,
            graph_attention_bias=graph_attention_bias,
            predicate_memory_bias=predicate_memory_bias,
            mode=mode,
            return_edge_logits=return_edge_logits,
            predicate_memory_intervention=predicate_memory_intervention,
            predicate_memory_residual_scale=predicate_memory_residual_scale,
            graph_object=graph_object,
            graph_object_intervention=graph_object_intervention,
            graph_object_residual_scale=graph_object_residual_scale,
            return_graph_object_logits=return_graph_object_logits,
        )
        logits = self.lm_head(model_output.last_hidden_state)
        loss = None
        if labels is not None:
            if labels.shape != input_ids.shape:
                raise ValueError(
                    f"labels must match input_ids shape {tuple(input_ids.shape)}, "
                    f"got {tuple(labels.shape)}"
                )
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100,
            )
        return StrataCausalLMOutput(
            logits=logits,
            loss=loss,
            hidden_states=model_output.last_hidden_state,
            predicate_outputs=model_output.predicate_outputs,
            graph_object_outputs=model_output.graph_object_outputs,
            attention_replacement_gates=model_output.attention_replacement_gates,
        )

    def save_pretrained(self, output_dir: str | Path, *, exist_ok: bool = False) -> None:
        """Save config and weights to a run-scoped artifact directory."""

        destination = Path(output_dir)
        if destination.exists() and any(destination.iterdir()) and not exist_ok:
            raise FileExistsError(
                f"refusing to overwrite non-empty model directory: {destination}"
            )
        destination.mkdir(parents=True, exist_ok=True)
        self.config.to_json_file(destination / "config.json")
        torch.save(self.state_dict(), destination / "model.pt")

    @classmethod
    def from_pretrained(
        cls,
        model_dir: str | Path,
        *,
        map_location: str | torch.device | None = None,
    ) -> "StrataForCausalLM":
        """Load a STRATA checkpoint saved by :meth:`save_pretrained`."""

        source = Path(model_dir)
        config = StrataConfig.from_json_file(source / "config.json")
        model = cls(config)
        state_dict = torch.load(
            source / "model.pt",
            map_location=map_location,
            weights_only=True,
        )
        try:
            model.load_state_dict(state_dict)
        except RuntimeError:
            if not config.use_graph_object_memory:
                raise
            current = model.state_dict()
            compatible = {
                key: value
                for key, value in state_dict.items()
                if key in current and current[key].shape == value.shape
            }
            unexpected = sorted(key for key in state_dict if key not in current)
            missing = sorted(key for key in current if key not in compatible)
            mismatched = sorted(
                key
                for key, value in state_dict.items()
                if key in current and current[key].shape != value.shape
            )
            non_graph_missing = [key for key in missing if "graph_object" not in key]
            non_graph_mismatched = [key for key in mismatched if "graph_object" not in key]
            if unexpected or non_graph_missing or non_graph_mismatched:
                raise
            model.load_state_dict(compatible, strict=False)
        return model


def _validate_residual_scale(scale: torch.Tensor | float | int | None, batch_size: int, *, name: str) -> None:
    if scale is None or isinstance(scale, (float, int)):
        return
    scale_shape = tuple(scale.shape)
    if scale_shape in {(), (batch_size,), (batch_size, 1), (batch_size, 1, 1)}:
        return
    raise ValueError(
        f"{name} must be scalar or have shape ({batch_size},), "
        f"({batch_size}, 1), or ({batch_size}, 1, 1); got {scale_shape}"
    )