config/model_config.py

from dataclasses import dataclass


@dataclass
class ModelConfig:
    # Model architecture
    n_embeds: int = 384
    n_heads: int = 6
    n_layers: int = 6
    dropout: float = 0.3

    # Training
    batch_size: int = 64
    block_size: int = 128
    max_iters: int = 20000
    eval_interval: int = 250
    eval_iters: int = 200
    learning_rate: float = 3e-4
    weight_decay: float = 0.1

    # Paths
    checkpoint_path: str = "checkpoints/model.pth"
    data_path: str = "data/input.txt"