spb/wpai0001_uqo

1"""Typed configuration objects and deterministic hashing.23Every experiment is described by a pydantic model; ``config_hash`` is the SHA-256 of the4canonical JSON dump so that identical configurations always produce identical run IDs.5"""6from __future__ import annotations78import hashlib9import json10from typing import Any, Literal1112from pydantic import BaseModel, Field131415def canonical_json(obj: Any) -> str:16    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=True, default=str)171819def config_hash(model: BaseModel, length: int = 12) -> str:20    return hashlib.sha256(canonical_json(model.model_dump()).encode()).hexdigest()[:length]212223class OpSeqTaskConfig(BaseModel):24    """Micro-world A: compositional operator sequences over Z_p (see datasets/opseq.py)."""2526    family: Literal["opseq"] = "opseq"27    modulus: int = 1128    operators: list[str] = Field(default_factory=lambda: ["zib", "zob", "qua", "lar", "neg", "flip", "dro", "mik"])29    min_len: int = 130    max_len: int = 531    held_out_pairs: list[tuple[str, str]] = Field(default_factory=lambda: [("zob", "lar"), ("qua", "neg")])32    """Adjacent operator pairs never seen in training (compositional held-out)."""33    excluded_operators: list[str] = Field(default_factory=list)34    """Operators removed from a model's training distribution (used to create recipient weakness)."""35    trace_stride: int = 136    """Emit every ``trace_stride``-th intermediate result (plus the final one). 1 = full scratchpad;37    larger values force multi-step composition inside a single prediction (difficulty knob)."""383940class FSATaskConfig(BaseModel):41    """Micro-world B: run a fixed random DFA over an invented alphabet and output the final state."""4243    family: Literal["fsa"] = "fsa"44    n_states: int = 445    alphabet: list[str] = Field(default_factory=lambda: ["ka", "mo", "ti"])46    min_len: int = 147    max_len: int = 848    automaton_seed: int = 7495051class MorphTaskConfig(BaseModel):52    """Micro-world C: context-dependent morphological rule on pseudo-words."""5354    family: Literal["morph"] = "morph"55    min_stem: int = 356    max_stem: int = 657    rules: dict[str, str] = Field(default_factory=lambda: {"vowel": "-na", "consonant": "-ek", "nasal": "-ur"})585960TaskConfig = OpSeqTaskConfig | FSATaskConfig | MorphTaskConfig616263class ToyModelConfig(BaseModel):64    d_model: int = 12865    n_layers: int = 466    n_heads: int = 467    d_ff: int = 51268    max_seq_len: int = 3269    dropout: float = 0.070    tie_embeddings: bool = False717273class TrainConfig(BaseModel):74    seed: int = 075    n_train: int = 40_00076    n_dev: int = 2_00077    batch_size: int = 25678    steps: int = 3_00079    lr: float = 1e-380    weight_decay: float = 0.0181    warmup_steps: int = 10082    device: Literal["auto", "cpu", "mps", "cuda"] = "auto"83    eval_every: int = 500848586class ModelRoleConfig(BaseModel):87    """A trained toy model with a role in the donor/recipient design."""8889    role: Literal["donor", "recipient", "unrelated"]90    name: str91    model: ToyModelConfig = Field(default_factory=ToyModelConfig)92    train: TrainConfig = Field(default_factory=TrainConfig)93    task: OpSeqTaskConfig | FSATaskConfig | MorphTaskConfig = Field(default_factory=OpSeqTaskConfig)949596class PilotConfig(BaseModel):97    experiment_id: str = "pilot_tier0_opseq"98    description: str = "Tier-0 pilot: donor/recipient gap, hook no-op test, within-model causal localization"99    seeds: list[int] = Field(default_factory=lambda: [0, 1, 2])100    donor: ModelRoleConfig101    recipient: ModelRoleConfig102    target_operator: str = "qua"103    n_eval: int = 1_000104    n_bootstrap: int = 2_000105