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 r_max: int = 8 # hard cap on actions per token37 eta_max: float = 0.2538 clip_c: float = 1.039 clip_eps: float = 1e-640 readout_zero_init: bool = True41 readout_init_std: float = 0.02 # used only when readout_zero_init is False (bootstrap ablation)42 readout_scale: float = 1.0 # fixed (non-learned) multiplier on U(...); set to the backbone residual rms at bootstrap (D9)43 # --- controller ---------------------------------------------------------44 controller_hidden: int = 0 # 0 -> ffn_w45 ctrl_input: str = "z" # "z" (w-dim summary) or "h" (normalised backbone residual, d-dim; costs 2dw per token)46 ctrl_use_entropy: bool = False # add the backbone's own predictive entropy at budget 0 (costs one head evaluation per token)47 lambda_cost: float = 0.048 kappa_risk: float = 0.049 # --- dense-workspace ablation (writes all slots) --------------------------50 dense_writes: bool = False5152 @property53 def head_dim(self) -> int:54 assert self.d_model % self.n_heads == 055 return self.d_model // self.n_heads5657 @property58 def kv_dim(self) -> int:59 return self.n_kv_heads * self.head_dim6061 @property62 def editor_head_dim(self) -> int:63 assert self.w % self.n_editor_heads == 064 return self.w // self.n_editor_heads6566 @property67 def n_pairs(self) -> int:68 return self.n_slots // 26970 @property71 def n_actions(self) -> int:72 """Nonterminal actions (singles + fixed pairs); STOP is index n_actions."""73 return self.n_slots + self.n_pairs7475 @property76 def stop_action(self) -> int:77 return self.n_actions7879 def validate(self) -> None:80 assert self.n_heads % self.n_kv_heads == 0, "GQA requires n_heads % n_kv_heads == 0"81 assert self.n_slots % 2 == 0, "fixed pairs require an even slot count"82 assert self.w % self.n_editor_heads == 083 assert self.r_max >= 18485 def to_dict(self) -> dict[str, Any]:86 return dataclasses.asdict(self)8788 def sha256(self) -> str:89 return hashlib.sha256(json.dumps(self.to_dict(), sort_keys=True).encode()).hexdigest()9091 @classmethod92 def from_dict(cls, d: dict[str, Any]) -> "PrismeConfig":93 known = {f.name for f in dataclasses.fields(cls)}94 unknown = set(d) - known95 if unknown:96 raise ValueError(f"unknown config keys: {sorted(unknown)}")97 cfg = cls(**d)98 cfg.validate()99 return cfg100101102def load_config(path: str | Path) -> PrismeConfig:103 with open(path) as f:104 raw = yaml.safe_load(f)105 model = raw.get("model", raw)106 return PrismeConfig.from_dict(model)107