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