config.py

"""
Shared configuration for the ru->en transformer.

Everything that more than one script needs to agree on lives here:
special-token strings, the model hyperparameter dataclass, and the
convention for where data files live on disk.

Nothing in here imports torch, so it is cheap to import from any script.

============================================================================
WHERE THIS FILE SITS IN THE PIPELINE (read this once, it explains the rest)
============================================================================
config.py has no dependencies on the rest of the project — it is pure
Python (dataclasses + pathlib). Every other file in the project starts with
`from config import ...`, which is why it is the right file to read first.

Who imports what from here:
  - tokenizer.py    imports SPECIAL_TOKENS, PAD/BOS/EOS/UNK_TOKEN, SRC/TGT_LANG, Paths
                     (needs the special-token STRINGS to train the BPE model)
  - dataset.py       imports PAD/BOS/EOS_TOKEN, SRC/TGT_LANG, Paths, human
                     (needs the strings to look up their integer ids via the
                     already-trained tokenizer: tok.token_to_id(PAD_TOKEN))
  - model.py         imports ModelConfig only
                     (the dataclass below IS the model's architecture)
  - train.py         imports ModelConfig, Paths, PAD/BOS/EOS_TOKEN, human
  - evaluate.py       imports ModelConfig, Paths, PAD/BOS/EOS_TOKEN, human
  - prepare_data.py  imports Paths, SRC_LANG, TGT_LANG, human
  - sanity_checks.py imports ModelConfig only

So: the STRINGS (PAD_TOKEN etc.) are shared vocabulary between tokenizer.py
(which assigns them ids while training the BPE model) and every script that
later needs to convert "<pad>" -> its integer id. The ModelConfig dataclass
is shared between train.py (which builds one from CLI flags) and model.py
(which consumes one to construct the actual nn.Module layers) — and it gets
serialized INTO every checkpoint .pt file so evaluate.py can reconstruct the
exact same architecture later without being told the flags again.
"""

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

# ---------------------------------------------------------------------------
# Special tokens
# ---------------------------------------------------------------------------
# These are just PYTHON STRINGS at this point — not yet integers. They become
# integers only once tokenizer.py trains a BPE model and assigns each string
# an id (see tokenizer.py's build_tokenizer(), which passes SPECIAL_TOKENS as
# the `special_tokens=` argument to the BPE trainer).
#
# The IDs are fixed by listing them first when training the BPE model, so
# <pad> == 0, <bos> == 1, <eos> == 2, <unk> == 3 in every script. Hard-coding
# the *order* (not the ID) here is what keeps tokenizer.py, dataset.py,
# train.py and evaluate.py consistent.
PAD_TOKEN = "<pad>"   # padding filler for batching unequal-length sentences -> id 0
BOS_TOKEN = "<bos>"   # "beginning of sequence" — prepended to every target sentence -> id 1
EOS_TOKEN = "<eos>"   # "end of sequence" — appended to source AND target -> id 2
UNK_TOKEN = "<unk>"   # fallback for anything the BPE vocab can't represent -> id 3
                      # (byte-level BPE means this almost never actually fires — see tokenizer.py)

# SPECIAL_TOKENS: the LIST, in this exact order. tokenizer.py trains the BPE
# model with special_tokens=SPECIAL_TOKENS, which is what pins
# PAD_TOKEN -> id 0 etc. Downstream code (dataset.py, train.py, evaluate.py)
# never hard-codes "0" for pad — it always asks the tokenizer object:
#     tok.token_to_id(PAD_TOKEN)
# so if you ever reorder this list, everything downstream still works
# correctly as long as you retrain the tokenizer. Only sanity_checks.py
# hard-codes bos=1, eos=2 (as literal ints, for a throwaway toy example),
# which is fine there because it never loads a real tokenizer.
SPECIAL_TOKENS = [PAD_TOKEN, BOS_TOKEN, EOS_TOKEN, UNK_TOKEN]

# Language direction. v1 is Russian -> English only.
# SRC_LANG/TGT_LANG select which columns of the opus-100 dataset get read in
# prepare_data.py's extract() function, and which file suffix (.ru / .en)
# each split is written to / read from via Paths.split_file() below.
SRC_LANG = "ru"
TGT_LANG = "en"


# ---------------------------------------------------------------------------
# Model configuration
# ---------------------------------------------------------------------------
@dataclass
class ModelConfig:
    """Architecture hyperparameters.

    This object is saved inside every checkpoint, so evaluate.py can rebuild
    the exact same architecture without being told the flags again. That
    matters a lot for the ablation runs, where different checkpoints have
    different depths / vocab sizes / positional encodings.

    CHAIN OF EVENTS for this class:
      1. train.py's parse_args() reads CLI flags like --d-model, --n-heads
      2. train.py's main() builds:  cfg = ModelConfig(d_model=args.d_model, ...)
      3. train.py passes cfg into model.py's build_model(cfg, ...), which
         reads cfg.vocab_size, cfg.d_model, cfg.n_heads, etc. to construct
         the actual nn.Embedding / nn.MultiheadAttention / etc. layers
         (see model.py: TransformerTranslator.__init__)
      4. train.py saves cfg.to_dict() into the checkpoint dict passed to
         torch.save(...) — see the "new best, saved" block in train.py:main()
      5. evaluate.py loads that checkpoint and calls
         ModelConfig.from_dict(ckpt["config"]) to recover this exact object,
         then calls model.py's build_model(cfg, ...) again — same
         architecture, now with trained weights loaded on top via
         model.load_state_dict(ckpt["model"])
    """

    vocab_size: int = 16000
    # ^ how many rows the shared embedding table has (model.py: self.embed).
    #   Must exactly match the tokenizer's vocab size — train.py actually
    #   overwrites this with tok.get_vocab_size() rather than trusting the
    #   CLI default, so it can never silently drift out of sync.
    d_model: int = 256
    # ^ width of every token's vector representation throughout the network
    #   (the "d" in every (B, L, d) tensor shape comment you'll see in model.py).
    n_heads: int = 8
    # ^ how many parallel attention heads each of the 3 attention calls uses.
    #   d_model must be divisible by n_heads (nn.MultiheadAttention enforces this).
    n_encoder_layers: int = 4
    # ^ how many EncoderLayer blocks are stacked (model.py: self.encoder_layers)
    n_decoder_layers: int = 4
    # ^ how many DecoderLayer blocks are stacked (model.py: self.decoder_layers)
    d_ff: int = 1024
    # ^ hidden width inside each layer's position-wise FeedForward sublayer
    #   (model.py: FeedForward does d_model -> d_ff -> d_model)
    dropout: float = 0.1
    # ^ used in attention, the FFN, and the positional encoding modules
    max_len: int = 512           # longest position the pos-encoding supports
    # ^ SinusoidalPositionalEncoding / LearnedPositionalEncoding precompute a
    #   table of this many rows (model.py). Sentences longer than this would
    #   index out of range, so train.py sets this to max(512, max_len + 2).
    pos_encoding: str = "sinusoidal"  # "sinusoidal" | "learned"
    # ^ selects between model.py's two positional-encoding classes via
    #   build_positional_encoding() — see model.py near the bottom of the
    #   "Positional encodings" section.
    tie_embeddings: bool = True  # share enc emb / dec emb / output projection
    # ^ if True, model.py's TransformerTranslator.__init__ does
    #   self.output_proj.weight = self.embed.weight — literally the same
    #   tensor object, not a copy. Only possible because tokenizer.py trains
    #   ONE shared vocabulary for both languages (see tokenizer.py's module
    #   docstring, "Why a shared vocabulary").

    def to_dict(self):
        # asdict(): dataclasses -> plain dict, so it is JSON/torch.save-friendly.
        # Called by train.py right before torch.save(...) to freeze this
        # config into the checkpoint file.
        return asdict(self)

    @staticmethod
    def from_dict(d):
        # Ignore unknown keys so old checkpoints still load after we add fields.
        # Called by evaluate.py as: cfg = ModelConfig.from_dict(ckpt["config"])
        # right after torch.load(...). The filtering matters if you add a new
        # field to this dataclass later — checkpoints saved BEFORE that change
        # won't have the new key in their dict, but ** unpacking below would
        # otherwise crash on an unexpected-keyword-argument if the checkpoint
        # dict had EXTRA stale keys from a since-removed field. `known` guards
        # against that direction (extra keys in `d` that ModelConfig no longer has).
        known = {f for f in ModelConfig.__dataclass_fields__}
        return ModelConfig(**{k: v for k, v in d.items() if k in known})


# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
@dataclass
class Paths:
    """Where each stage reads from and writes to.

    Every script constructs one of these near the top of main() as
    `paths = Paths(root=Path(args.root))`, then calls its methods/properties
    instead of hand-building path strings. This is the ONE place that knows
    the on-disk layout, so if you ever reorganize folders you edit only here.

    Layout produced by the pipeline:

        data/
            train.ru  train.en
            valid.ru  valid.en
            test.ru   test.en
            stats.json
        tokenizer/
            bpe-16000.json
        checkpoints/
            <run_name>.pt
        experiments/
            results.csv
    """

    root: Path = field(default_factory=lambda: Path("."))
    # ^ every property below is computed relative to this. Passing
    #   --root sanity to prepare_data.py, for instance, redirects the ENTIRE
    #   data/tokenizer/checkpoints/experiments tree under ./sanity/ instead
    #   of the repo root, which is how you get an isolated tiny debug dataset.

    @property
    def data(self) -> Path:
        # Used by: prepare_data.write_split() (writes here), and
        # dataset.py's TranslationDataset.__init__ (reads from here via
        # split_file() below).
        return self.root / "data"

    @property
    def tokenizer_dir(self) -> Path:
        # Used by tokenizer.py's main() to mkdir() before saving, and by
        # tokenizer_file() below to build the actual .json path.
        return self.root / "tokenizer"

    @property
    def checkpoints(self) -> Path:
        # Used by train.py (torch.save target) and evaluate.py /
        # run_ablations.py (torch.load source / --checkpoint path building).
        return self.root / "checkpoints"

    @property
    def experiments(self) -> Path:
        # Used by evaluate.py's append_result() as the default results.csv
        # location, and printed by run_ablations.py at the end of a sweep.
        return self.root / "experiments"

    def split_file(self, split: str, lang: str) -> Path:
        """e.g. split_file("train", "ru") -> data/train.ru"""
        # Called from:
        #   - prepare_data.write_split(): paths.split_file(split, SRC_LANG) /
        #     paths.split_file(split, TGT_LANG) to open the two output files
        #   - dataset.py's _read_lines(paths.split_file(split, SRC_LANG)) and
        #     the TGT_LANG counterpart, inside TranslationDataset.__init__
        #   - tokenizer.py's main(): builds the two training-file paths
        #     (train.ru, train.en) that get handed to build_tokenizer()
        return self.data / f"{split}.{lang}"

    def tokenizer_file(self, vocab_size: int) -> Path:
        # Called from tokenizer.py (save target), dataset.py's standalone
        # main() and train.py/evaluate.py (load source) — always keyed by
        # vocab_size so different --vocab-size ablation runs don't collide.
        return self.tokenizer_dir / f"bpe-{vocab_size}.json"


def human(n: int) -> str:
    """1234567 -> '1,234,567' — used in the sanity printouts."""
    # A pure formatting helper, imported by prepare_data.py, tokenizer.py,
    # dataset.py, train.py, evaluate.py — anywhere a script prints a count
    # (pair counts, vocab size, parameter count) and wants thousands separators.
    return f"{n:,}"