1"""Configuration dataclass for PRISME models (backbone + workspace + controller)."""2from __future__ import annotations34import dataclasses5import hashlib6import json7from dataclasses import dataclass, field8from pathlib import Path9from typing import Any1011import yaml121314@dataclass15class PrismeConfig:16 # --- backbone -----------------------------------------------------------17 name: str = "prisme_tiny"18 vocab_size: int = 3276819 d_model: int = 76820 n_layers: int = 1221 n_heads: int = 1222 n_kv_heads: int = 323 ffn_width: int = 204824 max_seq_len: int = 409625 rope_theta: float = 10000.026 norm_eps: float = 1e-527 tie_embeddings: bool = True28 init_std: float = 0.0229 # --- workspace ----------------------------------------------------------30 use_workspace: bool = True31 w: int = 256 # workspace width32 ffn_w: int = 768 # workspace FFN width33 n_slots: int = 16 # M34 n_editor_blocks: int = 2 # Lw (shared weights, applied sequentially)35 n_editor_heads: int = 436 pair_scheme: str = "adjacent" # fixed pairs (2i, 2i+1) or "offset" pairs (i, i+M/2) — ablation37 r_max: int = 8 # hard cap on actions per token38 eta_max: float = 0.2539 clip_c: float = 1.040 clip_eps: float = 1e-641 readout_zero_init: bool = True42 readout_init_std: float = 0.02 # used only when readout_zero_init is False (bootstrap ablation)43 readout_scale: float = 1.0 # fixed (non-learned) multiplier on U(...); set to the backbone residual rms at bootstrap (D9)44 # --- controller ---------------------------------------------------------45 controller_hidden: int = 0 # 0 -> ffn_w46 ctrl_input: str = "z" # "z" (w-dim summary) or "h" (normalised backbone residual, d-dim; costs 2dw per token)47 ctrl_use_entropy: bool = False # add the backbone's own predictive entropy at budget 0 (costs one head evaluation per token)48 lambda_cost: float = 0.049 kappa_risk: float = 0.050 # --- dense-workspace ablation (writes all slots) --------------------------51 dense_writes: bool = False5253 @property54 def head_dim(self) -> int:55 assert self.d_model % self.n_heads == 056 return self.d_model // self.n_heads5758 @property59 def kv_dim(self) -> int:60 return self.n_kv_heads * self.head_dim6162 @property63 def editor_head_dim(self) -> int:64 assert self.w % self.n_editor_heads == 065 return self.w // self.n_editor_heads6667 @property68 def n_pairs(self) -> int:69 return self.n_slots // 27071 @property72 def n_actions(self) -> int:73 """Nonterminal actions (singles + fixed pairs); STOP is index n_actions."""74 return self.n_slots + self.n_pairs7576 @property77 def stop_action(self) -> int:78 return self.n_actions7980 def validate(self) -> None:81 assert self.n_heads % self.n_kv_heads == 0, "GQA requires n_heads % n_kv_heads == 0"82 assert self.n_slots % 2 == 0, "fixed pairs require an even slot count"83 assert self.w % self.n_editor_heads == 084 assert self.r_max >= 18586 def to_dict(self) -> dict[str, Any]:87 return dataclasses.asdict(self)8889 def sha256(self) -> str:90 return hashlib.sha256(json.dumps(self.to_dict(), sort_keys=True).encode()).hexdigest()9192 @classmethod93 def from_dict(cls, d: dict[str, Any]) -> "PrismeConfig":94 known = {f.name for f in dataclasses.fields(cls)}95 unknown = set(d) - known96 if unknown:97 raise ValueError(f"unknown config keys: {sorted(unknown)}")98 cfg = cls(**d)99 cfg.validate()100 return cfg101102103def load_config(path: str | Path) -> PrismeConfig:104 with open(path) as f:105 raw = yaml.safe_load(f)106 model = raw.get("model", raw)107 return PrismeConfig.from_dict(model)108