dot_rd/config.py

from __future__ import annotations

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

import yaml


@dataclass(frozen=True)
class ModelConfig:
    base_model: str
    insertion_after: int = 15
    source_layers: tuple[int, ...] = (12, 13, 14, 15)
    max_loops: int = 8
    active_loops: int = 1
    initial_loop_scale: float = 0.01
    attention_implementation: str = "sdpa"

    def validate(self) -> None:
        if not self.base_model:
            raise ValueError("model.base_model is required")
        if self.insertion_after < 0:
            raise ValueError("model.insertion_after must be non-negative")
        if not self.source_layers:
            raise ValueError("model.source_layers must not be empty")
        if len(set(self.source_layers)) != len(self.source_layers):
            raise ValueError("model.source_layers must be unique")
        if min(self.source_layers) < 0:
            raise ValueError("model.source_layers must be non-negative")
        if self.max_loops < 1:
            raise ValueError("model.max_loops must be at least 1")
        if not 1 <= self.active_loops <= self.max_loops:
            raise ValueError("model.active_loops must be within [1, max_loops]")
        if not 0.0 <= self.initial_loop_scale < 1.0:
            raise ValueError("model.initial_loop_scale must be within [0, 1)")


@dataclass(frozen=True)
class TrainingConfig:
    data_path: str
    output_dir: str
    sequence_length: int = 512
    micro_batch_size: int = 1
    gradient_accumulation_steps: int = 4
    max_steps: int = 20
    learning_rate_core: float = 1.0e-5
    learning_rate_scales: float = 1.0e-3
    weight_decay: float = 0.1
    warmup_steps: int = 2
    max_grad_norm: float = 1.0
    log_every: int = 1
    save_every: int = 10
    seed: int = 1337
    gradient_checkpointing: bool = True
    keep_last_checkpoints: int = 2
    reasoning_loss_weight: float = 0.35
    answer_loss_weight: float = 1.0

    def validate(self) -> None:
        if not self.data_path:
            raise ValueError("training.data_path is required")
        if not self.output_dir:
            raise ValueError("training.output_dir is required")
        positive_ints = {
            "sequence_length": self.sequence_length,
            "micro_batch_size": self.micro_batch_size,
            "gradient_accumulation_steps": self.gradient_accumulation_steps,
            "max_steps": self.max_steps,
            "log_every": self.log_every,
            "save_every": self.save_every,
            "keep_last_checkpoints": self.keep_last_checkpoints,
        }
        for name, value in positive_ints.items():
            if value < 1:
                raise ValueError(f"training.{name} must be at least 1")
        if self.warmup_steps < 0:
            raise ValueError("training.warmup_steps must be non-negative")
        if self.learning_rate_core <= 0 or self.learning_rate_scales <= 0:
            raise ValueError("training learning rates must be positive")
        if self.max_grad_norm <= 0:
            raise ValueError("training.max_grad_norm must be positive")
        if self.reasoning_loss_weight <= 0 or self.answer_loss_weight <= 0:
            raise ValueError("training thinking loss weights must be positive")


@dataclass(frozen=True)
class DotRDConfig:
    model: ModelConfig
    training: TrainingConfig
    metadata: dict[str, Any] = field(default_factory=dict)

    def validate(self) -> None:
        self.model.validate()
        self.training.validate()

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


def load_config(path: str | Path) -> DotRDConfig:
    config_path = Path(path)
    payload = yaml.safe_load(config_path.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("configuration root must be an object")
    model_payload = dict(payload.get("model") or {})
    training_payload = dict(payload.get("training") or {})
    if "source_layers" in model_payload:
        model_payload["source_layers"] = tuple(model_payload["source_layers"])
    config = DotRDConfig(
        model=ModelConfig(**model_payload),
        training=TrainingConfig(**training_payload),
        metadata=dict(payload.get("metadata") or {}),
    )
    config.validate()
    return config