code/src/config.py

"""Central configuration.

Per the repo convention, every constant/hyperparameter lives here — nothing
hardcoded inline elsewhere. Values resolve in three layers, later wins:

    dataclass defaults  <  --config some.yaml  <  explicit CLI flags

so the workflows from the README both work:

    python src/train.py --config configs/dev.yaml
    python src/train.py --config configs/full.yaml --lr 1e-4
"""
from __future__ import annotations

import argparse
from dataclasses import dataclass, fields, MISSING


@dataclass
class Config:
    # ---- Data ---------------------------------------------------------------
    # Comma-separated roots of preprocessed clip trees, each laid out as
    # <root>/<WORD>/<split>/*.npy (see src/extract_word_clips.py roi). Any
    # number of sources merge transparently — e.g. LRS2-derived clips,
    # LRS3-derived clips, and custom recordings can each be their own root.
    data_roots: str = "data/lrs2,data/lrs3,data/custom"
    # Optional explicit list of words to keep (comma-separated). Empty = use all
    # words found across the sources. Useful to train on a small subset first.
    words: str = ""

    num_frames: int = 29          # frames per clip
    image_size: int = 88          # mouth ROI size (square, grayscale)

    # ---- Model --------------------------------------------------------------
    gru_hidden: int = 256
    gru_layers: int = 2
    dropout: float = 0.4

    # ---- Optimization -------------------------------------------------------
    epochs: int = 40
    batch_size: int = 32
    lr: float = 3e-4
    weight_decay: float = 1e-4
    warmup_epochs: int = 2
    label_smoothing: float = 0.1
    grad_clip: float = 5.0

    # ---- Runtime ------------------------------------------------------------
    num_workers: int = 4
    seed: int = 1337
    amp: bool = True              # mixed-precision (only used on CUDA)
    device: str = "auto"         # "auto" | "cuda" | "cpu"
    out_dir: str = "checkpoints"
    resume: str = ""             # path to a checkpoint to resume from
    log_every: int = 20          # log training loss every N steps

    # ---- Smoke test ---------------------------------------------------------
    dummy: bool = False          # use synthetic data to verify the pipeline
    dummy_classes: int = 10
    dummy_samples: int = 256

    # ------------------------------------------------------------------------
    @staticmethod
    def add_args(parser: argparse.ArgumentParser) -> None:
        for f in fields(Config):
            default = f.default if f.default is not MISSING else None
            if f.type == "bool" or isinstance(default, bool):
                # allow --amp / --no-amp style toggles
                parser.add_argument(f"--{f.name}", dest=f.name,
                                    action=argparse.BooleanOptionalAction,
                                    default=default)
            else:
                parser.add_argument(f"--{f.name}", type=type(default),
                                    default=default)

    @classmethod
    def from_args(cls, argv=None) -> "Config":
        # First pass: pick up --config only, so its YAML can become the new
        # defaults before the full parse (letting explicit CLI flags win).
        pre = argparse.ArgumentParser(add_help=False)
        pre.add_argument("--config", default="")
        known, _ = pre.parse_known_args(argv)

        yaml_defaults = {}
        if known.config:
            import yaml
            with open(known.config) as f:
                yaml_defaults = yaml.safe_load(f) or {}
            valid = {f.name for f in fields(cls)}
            unknown = set(yaml_defaults) - valid
            if unknown:
                raise SystemExit(
                    f"{known.config}: unknown config keys {sorted(unknown)}")

        parser = argparse.ArgumentParser(description="Lip-reading baseline")
        parser.add_argument("--config", default="",
                            help="YAML file of Config overrides")
        cls.add_args(parser)
        parser.set_defaults(**yaml_defaults)
        ns = parser.parse_args(argv)
        kwargs = vars(ns)
        kwargs.pop("config", None)
        return cls(**kwargs)