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