src/diffusion_lm/config.py

"""Typed experiment configuration loaded from YAML."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

import yaml


MODEL_BACKBONES = ("project", "hf-qwen3")


@dataclass(frozen=True)
class ModelConfig:
    """Architecture for the bidirectional token denoiser."""

    vocab_size: int
    mask_token_id: int
    max_seq_len: int
    d_model: int
    n_layers: int
    n_heads: int
    d_ff: int
    dropout: float = 0.1
    tie_embeddings: bool = True
    activation_checkpointing: bool = False
    use_flex_attention: bool = False
    compile_backbone: bool = False
    forbidden_output_token_ids: tuple[int, ...] = (0, 1, 2, 4)
    backbone: str = "project"
    pretrained_path: str | None = None
    # First id of the untrained tail of a pretrained embedding matrix; logits from that id
    # upward are masked without enumerating hundreds of forbidden ids.
    forbidden_output_from: int | None = None

    def __post_init__(self) -> None:
        if self.vocab_size <= 1:
            raise ValueError("vocab_size must be greater than one")
        if self.backbone not in MODEL_BACKBONES:
            raise ValueError(f"backbone must be one of {MODEL_BACKBONES}")
        if self.backbone != "project" and self.pretrained_path is None:
            raise ValueError("pretrained backbones require pretrained_path")
        if self.forbidden_output_from is not None and not (
            0 < self.forbidden_output_from <= self.vocab_size
        ):
            raise ValueError("forbidden_output_from must be inside the vocabulary")
        if not 0 <= self.mask_token_id < self.vocab_size:
            raise ValueError("mask_token_id must be inside the vocabulary")
        if self.max_seq_len <= 0:
            raise ValueError("max_seq_len must be positive")
        if self.d_model <= 0 or self.n_layers <= 0 or self.n_heads <= 0 or self.d_ff <= 0:
            raise ValueError("all model dimensions must be positive")
        if self.d_model % self.n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")
        if not isinstance(self.activation_checkpointing, bool):
            raise ValueError("activation_checkpointing must be a boolean")
        if not isinstance(self.use_flex_attention, bool):
            raise ValueError("use_flex_attention must be a boolean")
        if not isinstance(self.compile_backbone, bool):
            raise ValueError('compile_backbone must be a boolean')

        forbidden_ids = tuple(self.forbidden_output_token_ids)
        if any(
            isinstance(token_id, bool) or not isinstance(token_id, int)
            for token_id in forbidden_ids
        ):
            raise ValueError("forbidden_output_token_ids must contain integers")
        if len(set(forbidden_ids)) != len(forbidden_ids):
            raise ValueError("forbidden_output_token_ids must not contain duplicates")
        if any(not 0 <= token_id < self.vocab_size for token_id in forbidden_ids):
            raise ValueError("forbidden output token ids must be inside the vocabulary")
        if self.mask_token_id not in forbidden_ids:
            raise ValueError("mask_token_id must be a forbidden output token")
        if self.backbone == "project" and 3 in forbidden_ids:
            raise ValueError("EOS token id 3 must remain an allowed output token")
        object.__setattr__(self, "forbidden_output_token_ids", forbidden_ids)


@dataclass(frozen=True)
class TrainingConfig:
    """Single-device training settings for the first research iteration."""

    train_data: str
    tokenizer: str
    output_dir: str = "outputs/run"
    val_data: str | None = None
    batch_size: int = 32
    gradient_accumulation_steps: int = 1
    max_steps: int = 10_000
    learning_rate: float = 3e-4
    min_learning_rate: float = 3e-5
    warmup_steps: int = 500
    weight_decay: float = 0.1
    grad_clip: float = 1.0
    mask_eps: float = 1e-3
    seed: int = 1337
    device: str = "auto"
    precision: str = "auto"
    optimizer: str = "adamw"
    optimizer_min_8bit_size: int = 4096
    optimizer_embedding_32bit: bool = True
    require_fused_attention: bool = False
    save_inference_checkpoint: bool = False
    num_workers: int = 0
    log_interval: int = 10
    eval_interval: int = 500
    eval_batches: int = 20
    save_interval: int = 500
    keep_last_checkpoints: int = 3

    def __post_init__(self) -> None:
        if self.batch_size <= 0 or self.gradient_accumulation_steps <= 0:
            raise ValueError("batch sizes must be positive")
        if self.max_steps <= 0:
            raise ValueError("max_steps must be positive")
        if not 0.0 < self.learning_rate:
            raise ValueError("learning_rate must be positive")
        if not 0.0 <= self.min_learning_rate <= self.learning_rate:
            raise ValueError("min_learning_rate must be between zero and learning_rate")
        if not 0 <= self.warmup_steps < self.max_steps:
            raise ValueError("warmup_steps must be non-negative and less than max_steps")
        if self.weight_decay < 0.0:
            raise ValueError("weight_decay must be non-negative")
        if self.grad_clip <= 0.0:
            raise ValueError("grad_clip must be positive")
        if not 0.0 < self.mask_eps < 1.0:
            raise ValueError("mask_eps must be in (0, 1)")
        if self.num_workers < 0:
            raise ValueError("num_workers must be non-negative")
        if self.log_interval <= 0 or self.eval_interval <= 0 or self.save_interval <= 0:
            raise ValueError("log, eval, and save intervals must be positive")
        if self.eval_batches <= 0:
            raise ValueError("eval_batches must be positive")
        if self.keep_last_checkpoints < 0:
            raise ValueError("keep_last_checkpoints must be non-negative")
        if self.precision not in {"auto", "float32", "bfloat16", "float16"}:
            raise ValueError("precision must be auto, float32, bfloat16, or float16")
        if self.optimizer not in {"adamw", "adamw8bit"}:
            raise ValueError("optimizer must be adamw or adamw8bit")
        if self.optimizer_min_8bit_size <= 0:
            raise ValueError("optimizer_min_8bit_size must be positive")
        if not isinstance(self.optimizer_embedding_32bit, bool):
            raise ValueError("optimizer_embedding_32bit must be a boolean")
        if not isinstance(self.require_fused_attention, bool):
            raise ValueError("require_fused_attention must be a boolean")
        if not isinstance(self.save_inference_checkpoint, bool):
            raise ValueError("save_inference_checkpoint must be a boolean")


@dataclass(frozen=True)
class ExperimentConfig:
    model: ModelConfig
    training: TrainingConfig

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def load_config(path: str | Path) -> ExperimentConfig:
    """Load and validate an experiment YAML file."""

    config_path = Path(path)
    with config_path.open("r", encoding="utf-8") as handle:
        raw = yaml.safe_load(handle)
    if not isinstance(raw, dict) or "model" not in raw or "training" not in raw:
        raise ValueError("config must contain top-level model and training mappings")
    return ExperimentConfig(
        model=ModelConfig(**raw["model"]),
        training=TrainingConfig(**raw["training"]),
    )