training_loop.py

"""
training_loop.py
================
Custom training loop for the MiniLM model.

This module is part of the project:
    "A bilingual PT+EN LLM with BPE tokenizer and training loop
     implemented from scratch, with didactic and documented code"

Author  : André Costa
License : MIT

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
THEORETICAL BACKGROUND
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

The training objective
-----------------------
Training an LLM is an optimization problem: we want to find the
weights θ that minimize the average loss over the corpus:

    L(θ) = -1/N Σ log P(t_i | t_1, ..., t_{i-1}; θ)

In other words: maximize the probability the model assigns to the
correct next token given the previous context. This is called
"Language Modeling" or "next-token prediction".

The standard metric is Perplexity (PPL):
    PPL = exp(L)

Intuitively, perplexity measures "how many words the model considers
equally likely at each step". PPL = 10 means the model is, on average,
as uncertain as if it were choosing between 10 equally probable options.

Stochastic Gradient Descent (SGD)
-----------------------------------
Instead of computing the gradient over the entire corpus (infeasible),
we use mini-batches: random samples of B sequences per step.

    θ ← θ - η × ∇_θ L(batch)

where η is the learning rate.

AdamW Optimizer (Loshchilov & Hutter, 2019)
---------------------------------------------
AdamW combines two insights:
    1. Adam: adaptive per-parameter learning rate using first and
       second order gradient moments
    2. Decoupled weight decay: L2 regularization applied directly
       to weights, without interfering with Adam's adaptation

    m_t = β1 × m_{t-1} + (1-β1) × g_t          (1st order moment)
    v_t = β2 × v_{t-1} + (1-β2) × g_t²          (2nd order moment)
    θ_t = θ_{t-1} - η × m̂_t / (√v̂_t + ε) - η × λ × θ_{t-1}

Typical values: β1=0.9, β2=0.95, ε=1e-8, λ=0.1

Cosine Learning Rate Schedule with Warmup
-------------------------------------------
The learning rate is not constant — it varies throughout training:

    Phase 1 — Linear warmup (first ~2% of steps):
        lr grows linearly from 0 to lr_max
        Avoids instability at the start when weights are random

    Phase 2 — Cosine decay:
        lr decays smoothly from lr_max to lr_min
        lr(t) = lr_min + 0.5 × (lr_max - lr_min) × (1 + cos(π × t/T))

    Cosine decay is preferable to linear because:
        - Decays slowly at the start (still much to learn)
        - Decays faster in the middle
        - Stabilizes near the end (fine-grained refinement)

Gradient Clipping
------------------
Limits the gradient norm to a maximum value (typically 1.0):
    if ||g|| > max_norm:
        g ← g × max_norm / ||g||

Prevents "gradient explosion" — situations where the gradient grows
uncontrollably, causing destructive weight updates.
Especially important at the start of training.

Gradient Accumulation
----------------------
Simulates larger batch sizes without increasing VRAM usage:
    - Instead of one step with batch=32, do 4 steps with batch=8
    - Accumulate gradients across the 4 steps (without optimizer.step())
    - Apply the update after the 4th step

    effective_batch_size = batch_size × accumulation_steps

Useful for the RTX 4060 Ti (16GB), where physical batch size is limited.

Mixed Precision Training (bf16)
---------------------------------
Uses bfloat16 (16 bits) instead of float32 to:
    - Reduce VRAM usage by half
    - Speed up computation (bf16 ops are ~2x faster on modern GPUs)

bf16 vs fp16:
    - fp16: range 6×10⁻⁵ to 65504 → risk of overflow/underflow
    - bf16: same range as fp32 → more stable, no grad scaling needed

The RTX 4060 Ti natively supports bf16 — always use it.

References:
    - Loshchilov, I., & Hutter, F. (2019). Decoupled weight decay
      regularization. ICLR 2019.
    - Loshchilov, I., & Hutter, F. (2017). SGDR: Stochastic gradient
      descent with warm restarts. ICLR 2017.
    - Micikevicius, P. et al. (2018). Mixed precision training. ICLR 2018.
"""

import os
import math
import time
import json
import logging
from pathlib import Path
from dataclasses import dataclass, field
from typing import Optional

import torch
import torch.nn as nn
from torch.utils.data import DataLoader

# Project modules
from transformer import MiniLM, ModelConfig
from data_pipeline import CorpusDataset


# ─────────────────────────────────────────────────────────────
# Training configuration
# ─────────────────────────────────────────────────────────────

@dataclass
class TrainingConfig:
    """
    Training hyperparameters and settings.

    Separating training configuration from model configuration
    allows experimenting with different optimization regimes using
    the same architecture, and vice versa.

    Fields:
        # Paths
        corpus_dir:         Directory of the pre-processed corpus.
        checkpoint_dir:     Where to save checkpoints during training.
        model_config_path:  Path to save/load the model config.

        # Optimization
        lr_max:             Maximum (peak) learning rate.
                            Typical values for LLMs: 3e-4 to 6e-4.
        lr_min:             Minimum learning rate (end of cosine decay).
                            Typically lr_max / 10.
        weight_decay:       Decoupled L2 regularization in AdamW.
        beta1, beta2:       Adam moments. β2=0.95 is more conservative
                            than the default 0.999 — more stable for LLMs.
        grad_clip:          Maximum gradient norm.

        # Batch and accumulation
        batch_size:         Sequences per GPU step.
        accum_steps:        Gradient accumulation steps.
                            effective_batch = batch_size × accum_steps.

        # Schedule
        warmup_steps:       Linear warmup steps.
        max_steps:          Total optimization steps.
                            None = train for 1 full epoch.

        # Logging and checkpoints
        log_interval:       How often (in steps) to log metrics.
        eval_interval:      How often (in steps) to evaluate on val set.
        save_interval:      How often (in steps) to save a checkpoint.
        eval_steps:         How many batches to use for evaluation.

        # Hardware
        dtype:              Data type for mixed precision.
                            "bfloat16" for RTX 4060 Ti (recommended).
        compile_model:      If True, uses torch.compile() for ~20% speedup.
        num_workers:        DataLoader workers for parallel data loading.
    """
    # Paths
    corpus_dir:        str   = "./corpus"
    checkpoint_dir:    str   = "./checkpoints"
    model_config_path: str   = "./model_config.json"

    # Optimization
    lr_max:       float = 3e-4
    lr_min:       float = 3e-5
    weight_decay: float = 0.1
    beta1:        float = 0.9
    beta2:        float = 0.95
    grad_clip:    float = 1.0

    # Batch
    batch_size:   int   = 8       # adjust according to available VRAM
    accum_steps:  int   = 4       # effective_batch = 32

    # Schedule
    warmup_steps: int           = 200
    max_steps:    Optional[int] = None   # None = 1 full epoch

    # Logging
    log_interval:  int = 10
    eval_interval: int = 200
    save_interval: int = 500
    eval_steps:    int = 50

    # Hardware
    dtype:         str  = "bfloat16"
    compile_model: bool = True
    num_workers:   int  = 4

    @property
    def effective_batch_size(self) -> int:
        """Effective batch size after gradient accumulation."""
        return self.batch_size * self.accum_steps

    def save(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.__dict__, f, indent=2)

    @classmethod
    def load(cls, path: str) -> "TrainingConfig":
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        config = cls()
        for key, value in data.items():
            setattr(config, key, value)
        return config


# ─────────────────────────────────────────────────────────────
# Learning Rate Schedule
# ─────────────────────────────────────────────────────────────

def get_lr(
    step: int,
    warmup_steps: int,
    max_steps: int,
    lr_max: float,
    lr_min: float,
) -> float:
    """
    Compute the learning rate for the current step.

    Implements the standard LLM schedule:
        - Linear warmup from 0 → lr_max over the first `warmup_steps`
        - Cosine decay from lr_max → lr_min until `max_steps`

    Cosine decay is derived from the work of Loshchilov & Hutter (2017)
    on SGDR (Stochastic Gradient Descent with Restarts).
    Here we use only half a cycle (no restarts).

    Args:
        step:         Current optimization step (starts at 0).
        warmup_steps: Duration of the linear warmup.
        max_steps:    Total training steps.
        lr_max:       Maximum learning rate (warmup peak).
        lr_min:       Minimum learning rate (cosine end).

    Returns:
        Learning rate for the current step.

    Example curve (warmup=100, max=1000, lr_max=3e-4, lr_min=3e-5):
        step=0:    lr = 0.0
        step=50:   lr = 1.5e-4  (midpoint of warmup)
        step=100:  lr = 3e-4    (peak)
        step=550:  lr = 1.65e-4 (midpoint of cosine)
        step=1000: lr = 3e-5    (end)
    """
    # Phase 1: linear warmup
    if step < warmup_steps:
        return lr_max * (step + 1) / warmup_steps

    # Beyond max_steps: hold lr_min
    if step >= max_steps:
        return lr_min

    # Phase 2: cosine decay
    # Normalize progress after warmup to [0, 1]
    progress = (step - warmup_steps) / (max_steps - warmup_steps)

    # Half-cosine decay formula
    cosine_decay = 0.5 * (1.0 + math.cos(math.pi * progress))

    return lr_min + cosine_decay * (lr_max - lr_min)


# ─────────────────────────────────────────────────────────────
# Metrics and logging
# ─────────────────────────────────────────────────────────────

class MetricsTracker:
    """
    Track and record training metrics.

    Maintains a full history of loss and perplexity for
    post-training analysis and learning curve generation.

    Perplexity (PPL) is the main metric for LLMs:
        PPL = exp(cross_entropy_loss)

    Interpretation:
        PPL = 1:    perfect model (impossible in practice)
        PPL = 10:   good for small models on general text
        PPL = 50:   reasonable for very small models
        PPL = 100+: model still learning / difficult corpus
    """

    def __init__(self, log_dir: str):
        """
        Initialize the tracker and configure the logger.

        Args:
            log_dir: Directory where logs and metrics will be saved.
        """
        os.makedirs(log_dir, exist_ok=True)
        self.log_dir = log_dir

        # Full history for post-training analysis
        self.history: list[dict] = []

        # Accumulators for moving average
        self._loss_accum  = 0.0
        self._accum_count = 0

        # Configure logger to write to both file and console
        self.logger = logging.getLogger("MiniLM")
        self.logger.setLevel(logging.INFO)

        # File handler
        fh = logging.FileHandler(os.path.join(log_dir, "training.log"))
        fh.setFormatter(logging.Formatter("%(asctime)s | %(message)s"))

        # Console handler
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter("%(message)s"))

        self.logger.addHandler(fh)
        self.logger.addHandler(ch)

    def update(self, loss: float) -> None:
        """Accumulate loss for average computation."""
        self._loss_accum  += loss
        self._accum_count += 1

    def log_step(
        self,
        step: int,
        lr: float,
        tokens_per_sec: float,
        split: str = "train",
    ) -> dict:
        """
        Record metrics for the current step.

        Args:
            step:           Current step.
            lr:             Current learning rate.
            tokens_per_sec: Token throughput per second.
            split:          "train" or "val".

        Returns:
            Dictionary with the recorded metrics.
        """
        avg_loss = self._loss_accum / max(self._accum_count, 1)
        ppl      = math.exp(min(avg_loss, 20))  # clamp to avoid overflow

        metrics = {
            "step":           step,
            "split":          split,
            "loss":           round(avg_loss, 4),
            "perplexity":     round(ppl, 2),
            "lr":             f"{lr:.2e}",
            "tokens_per_sec": int(tokens_per_sec),
        }

        self.history.append(metrics)

        # Format log line
        self.logger.info(
            f"step {step:>6} | {split:<5} | "
            f"loss {avg_loss:.4f} | ppl {ppl:.2f} | "
            f"lr {lr:.2e} | {tokens_per_sec:.0f} tok/s"
        )

        # Reset accumulators
        self._loss_accum  = 0.0
        self._accum_count = 0

        return metrics

    def save_history(self) -> None:
        """Save the full history to JSON."""
        path = os.path.join(self.log_dir, "metrics_history.json")
        with open(path, "w", encoding="utf-8") as f:
            json.dump(self.history, f, indent=2)


# ─────────────────────────────────────────────────────────────
# Checkpoint
# ─────────────────────────────────────────────────────────────

def save_checkpoint(
    model: MiniLM,
    optimizer: torch.optim.Optimizer,
    step: int,
    loss: float,
    config: TrainingConfig,
    model_config: ModelConfig,
    is_best: bool = False,
) -> None:
    """
    Save a full training state checkpoint.

    A checkpoint includes everything needed to resume training
    exactly where it left off:
        - Model weights (state_dict)
        - Optimizer state (accumulated Adam moments)
        - Current step and best loss (for comparison)
        - Model and training configurations

    Checkpoint strategy:
        - Saves a periodic checkpoint every `save_interval` steps
        - Keeps only the 3 most recent checkpoints (saves disk space)
        - Separately saves the "best checkpoint" (lowest val loss)

    Args:
        model:        Model to save.
        optimizer:    Optimizer with its internal state.
        step:         Current step.
        loss:         Current validation loss.
        config:       Training configuration.
        model_config: Architecture configuration.
        is_best:      If True, also saves as "best_model.pt".
    """
    os.makedirs(config.checkpoint_dir, exist_ok=True)

    checkpoint = {
        "step":         step,
        "loss":         loss,
        "model_state":  model.state_dict(),
        "optim_state":  optimizer.state_dict(),
        "model_config": model_config.__dict__,
        "train_config": {k: v for k, v in config.__dict__.items()
                         if not callable(v)},
    }

    # Periodic checkpoint
    ckpt_path = os.path.join(config.checkpoint_dir, f"ckpt_step_{step:07d}.pt")
    torch.save(checkpoint, ckpt_path)

    # Keep only the 3 most recent
    ckpts = sorted(Path(config.checkpoint_dir).glob("ckpt_step_*.pt"))
    for old_ckpt in ckpts[:-3]:
        old_ckpt.unlink()

    # Save as best model if applicable
    if is_best:
        best_path = os.path.join(config.checkpoint_dir, "best_model.pt")
        torch.save(checkpoint, best_path)
        print(f"  → New best model saved (loss={loss:.4f})")


def load_checkpoint(
    path: str,
    model: MiniLM,
    optimizer: Optional[torch.optim.Optimizer] = None,
) -> dict:
    """
    Load a saved checkpoint.

    Args:
        path:      Path to the checkpoint .pt file.
        model:     Model to load weights into.
        optimizer: Optimizer to load state into (optional).

    Returns:
        Dictionary with checkpoint metadata (step, loss, configs).
    """
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)

    model.load_state_dict(checkpoint["model_state"])

    if optimizer is not None and "optim_state" in checkpoint:
        optimizer.load_state_dict(checkpoint["optim_state"])

    print(f"Checkpoint loaded: step={checkpoint['step']}, "
          f"loss={checkpoint['loss']:.4f}")

    return checkpoint


# ─────────────────────────────────────────────────────────────
# Evaluation
# ─────────────────────────────────────────────────────────────

@torch.no_grad()
def evaluate(
    model: MiniLM,
    val_loader: DataLoader,
    device: torch.device,
    dtype: torch.dtype,
    eval_steps: int,
) -> float:
    """
    Evaluate the model on the validation set.

    Disables gradient computation (@torch.no_grad) to save memory
    and speed up evaluation — during evaluation we only need the
    forward pass, not the backward pass.

    Loss is computed over `eval_steps` random batches from the val
    set, which is sufficient for a reliable estimate without running
    the full val set (which would be slow).

    Args:
        model:       Model to evaluate.
        val_loader:  DataLoader for the validation set.
        device:      Device (cuda/cpu).
        dtype:       Data type for autocast.
        eval_steps:  How many batches to evaluate.

    Returns:
        Average validation loss.
    """
    model.eval()

    total_loss = 0.0
    steps_done = 0

    for batch in val_loader:
        if steps_done >= eval_steps:
            break

        # Prepare input and targets
        # input_ids: all tokens except the last
        # targets:   all tokens except the first (shift of 1)
        input_ids = batch[:, :-1].to(device)
        targets   = batch[:, 1:].to(device)

        # Forward pass with autocast
        with torch.autocast(device_type=device.type, dtype=dtype):
            _, loss = model(input_ids, targets)

        total_loss += loss.item()
        steps_done += 1

    model.train()
    return total_loss / max(steps_done, 1)


# ─────────────────────────────────────────────────────────────
# Trainer — main class
# ─────────────────────────────────────────────────────────────

class Trainer:
    """
    Orchestrates the full training of MiniLM.

    Responsibilities:
        - Set up device, dtype and compilation
        - Initialize model, optimizer and LR schedule
        - Run the training loop with gradient accumulation
        - Periodically evaluate on the val set
        - Save checkpoints and metrics
        - Resume training from a checkpoint

    Basic usage:
        >>> model_config = ModelConfig()
        >>> train_config = TrainingConfig()
        >>> trainer = Trainer(model_config, train_config)
        >>> trainer.train()

    Resuming training:
        >>> trainer = Trainer(model_config, train_config)
        >>> trainer.train(resume_from="./checkpoints/ckpt_step_0005000.pt")
    """

    def __init__(self, model_config: ModelConfig, train_config: TrainingConfig):
        """
        Initialize the Trainer.

        Args:
            model_config: Model architecture configuration.
            train_config: Training configuration.
        """
        self.model_config = model_config
        self.config       = train_config

        # ── Device ────────────────────────────────────────────────────────
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"Device: {self.device}")

        if self.device.type == "cuda":
            print(f"  GPU: {torch.cuda.get_device_name(0)}")
            print(f"  Total VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

        # ── Data type for mixed precision ──────────────────────────────────
        # bf16 for RTX 4060 Ti (Ampere+), fp16 for older GPUs
        if train_config.dtype == "bfloat16" and torch.cuda.is_bf16_supported():
            self.dtype = torch.bfloat16
            print("  Mixed precision: bfloat16 ✓")
        elif train_config.dtype == "float16":
            self.dtype = torch.float16
            print("  Mixed precision: float16 ✓")
        else:
            self.dtype = torch.float32
            print("  Mixed precision: disabled (float32)")

        # ── Model ──────────────────────────────────────────────────────────
        self.model = MiniLM(model_config).to(self.device)
        print(f"\nModel: {self.model.count_parameters()['total'] / 1e6:.1f}M parameters")

        # torch.compile() — JIT compilation for ~20% speedup
        # Requires PyTorch 2.0+ and may take a few minutes the first time
        if train_config.compile_model and hasattr(torch, "compile"):
            print("  Compiling model with torch.compile()...")
            self.model = torch.compile(self.model)
            print("  torch.compile() ✓")

        # ── Optimizer ──────────────────────────────────────────────────────
        # Weight decay should NOT be applied to:
        #   - Embeddings (weight decay collapses them)
        #   - Bias terms
        #   - Normalization parameters (RMSNorm.weight)
        decay_params    = []
        no_decay_params = []

        for name, param in self.model.named_parameters():
            if not param.requires_grad:
                continue
            if param.ndim < 2 or "norm" in name or "bias" in name:
                no_decay_params.append(param)
            else:
                decay_params.append(param)

        optimizer_groups = [
            {"params": decay_params,    "weight_decay": train_config.weight_decay},
            {"params": no_decay_params, "weight_decay": 0.0},
        ]

        self.optimizer = torch.optim.AdamW(
            optimizer_groups,
            lr=train_config.lr_max,
            betas=(train_config.beta1, train_config.beta2),
            eps=1e-8,
            fused=True if self.device.type == "cuda" else False,
            # fused=True: CUDA fused implementation, ~10% faster
        )

        # ── DataLoaders ────────────────────────────────────────────────────
        train_dataset = CorpusDataset(
            os.path.join(train_config.corpus_dir, "train")
        )
        val_dataset = CorpusDataset(
            os.path.join(train_config.corpus_dir, "val")
        )

        self.train_loader = DataLoader(
            train_dataset,
            batch_size=train_config.batch_size,
            shuffle=True,
            num_workers=train_config.num_workers,
            pin_memory=True,    # speeds up CPU→GPU transfer
            drop_last=True,     # discard incomplete batch at the end
        )

        self.val_loader = DataLoader(
            val_dataset,
            batch_size=train_config.batch_size,
            shuffle=False,
            num_workers=train_config.num_workers,
            pin_memory=True,
        )

        # ── Max steps ──────────────────────────────────────────────────────
        if train_config.max_steps is None:
            # 1 epoch = iterate through the full dataset once
            self.max_steps = len(self.train_loader) // train_config.accum_steps
        else:
            self.max_steps = train_config.max_steps

        print(f"  Max steps: {self.max_steps:,}")
        print(f"  Effective batch size: {train_config.effective_batch_size}")
        print(f"  Steps per epoch: {len(self.train_loader) // train_config.accum_steps:,}")

        # ── Metrics ────────────────────────────────────────────────────────
        self.metrics = MetricsTracker(train_config.checkpoint_dir)

        # ── Internal state ─────────────────────────────────────────────────
        self.step      = 0
        self.best_loss = float("inf")

    def _set_lr(self, step: int) -> float:
        """
        Update the learning rate for all optimizer parameter groups.

        Args:
            step: Current step.

        Returns:
            Computed learning rate.
        """
        lr = get_lr(
            step=step,
            warmup_steps=self.config.warmup_steps,
            max_steps=self.max_steps,
            lr_max=self.config.lr_max,
            lr_min=self.config.lr_min,
        )
        for group in self.optimizer.param_groups:
            group["lr"] = lr
        return lr

    def train(self, resume_from: Optional[str] = None) -> None:
        """
        Run the full training loop.

        Main loop:
            For each batch from train_loader:
                1. Forward pass → loss
                2. loss /= accum_steps (scale for accumulation)
                3. Backward pass (accumulate gradients)
                4. Every accum_steps:
                    a. Gradient clipping
                    b. Update weights (optimizer.step)
                    c. Zero gradients (optimizer.zero_grad)
                5. Log metrics periodically
                6. Evaluate on val set periodically
                7. Save checkpoint periodically

        Args:
            resume_from: Path to a checkpoint to resume from (optional).
        """
        # Resume from checkpoint if provided
        if resume_from is not None:
            ckpt = load_checkpoint(resume_from, self.model, self.optimizer)
            self.step = ckpt["step"]
            self.best_loss = ckpt.get("loss", float("inf"))
            print(f"Resuming from step {self.step}")

        self.model.train()
        self.metrics.logger.info("=" * 60)
        self.metrics.logger.info("Training started")
        self.metrics.logger.info(
            f"max_steps={self.max_steps} | "
            f"batch={self.config.batch_size} | "
            f"accum={self.config.accum_steps} | "
            f"effective_batch={self.config.effective_batch_size}"
        )
        self.metrics.logger.info("=" * 60)

        # Time tracking for throughput computation
        t_start     = time.time()
        tokens_seen = 0

        # Infinite iterator over the dataset
        # (needed since max_steps may span more than 1 epoch)
        def infinite_loader():
            while True:
                for batch in self.train_loader:
                    yield batch

        loader_iter      = infinite_loader()
        accumulated_loss = 0.0

        while self.step < self.max_steps:

            # ── Update learning rate ───────────────────────────────────────
            lr = self._set_lr(self.step)

            # ── Gradient Accumulation Loop ─────────────────────────────────
            # Accumulate gradients over `accum_steps` micro-batches
            # before applying the weight update
            self.optimizer.zero_grad(set_to_none=True)
            # set_to_none=True frees memory instead of zeroing — more efficient

            for _ in range(self.config.accum_steps):
                batch = next(loader_iter)

                # Prepare input and targets (shift of 1 token)
                input_ids = batch[:, :-1].to(self.device, non_blocking=True)
                targets   = batch[:, 1:].to(self.device, non_blocking=True)

                tokens_seen += input_ids.numel()

                # Forward with autocast (mixed precision)
                with torch.autocast(
                    device_type=self.device.type,
                    dtype=self.dtype,
                ):
                    _, loss = self.model(input_ids, targets)

                # Scale the loss by the number of micro-steps so that
                # the accumulated gradient is equivalent to the gradient
                # of a batch of size effective_batch
                loss = loss / self.config.accum_steps
                accumulated_loss += loss.item()

                # Backward: accumulate gradients (do not zero yet)
                loss.backward()

            # ── Weight update ──────────────────────────────────────────────

            # Gradient clipping: prevents gradient explosion
            # Returns the norm before clipping (useful for monitoring)
            grad_norm = nn.utils.clip_grad_norm_(
                self.model.parameters(),
                self.config.grad_clip,
            )

            # Optimization step
            self.optimizer.step()

            self.step += 1

            # ── Logging ────────────────────────────────────────────────────
            self.metrics.update(accumulated_loss)
            accumulated_loss = 0.0

            if self.step % self.config.log_interval == 0:
                elapsed     = time.time() - t_start
                tok_per_sec = tokens_seen / elapsed
                lr_now      = self.optimizer.param_groups[0]["lr"]

                self.metrics.log_step(
                    step=self.step,
                    lr=lr_now,
                    tokens_per_sec=tok_per_sec,
                    split="train",
                )

                # Reset throughput counters
                tokens_seen = 0
                t_start     = time.time()

            # ── Evaluation ─────────────────────────────────────────────────
            if self.step % self.config.eval_interval == 0:
                val_loss = evaluate(
                    model=self.model,
                    val_loader=self.val_loader,
                    device=self.device,
                    dtype=self.dtype,
                    eval_steps=self.config.eval_steps,
                )

                self.metrics._loss_accum  = val_loss
                self.metrics._accum_count = 1
                self.metrics.log_step(
                    step=self.step,
                    lr=self.optimizer.param_groups[0]["lr"],
                    tokens_per_sec=0,
                    split="val",
                )

                is_best = val_loss < self.best_loss
                if is_best:
                    self.best_loss = val_loss

                save_checkpoint(
                    model=self.model,
                    optimizer=self.optimizer,
                    step=self.step,
                    loss=val_loss,
                    config=self.config,
                    model_config=self.model_config,
                    is_best=is_best,
                )

            # ── Periodic checkpoint ────────────────────────────────────────
            elif self.step % self.config.save_interval == 0:
                save_checkpoint(
                    model=self.model,
                    optimizer=self.optimizer,
                    step=self.step,
                    loss=self.best_loss,
                    config=self.config,
                    model_config=self.model_config,
                    is_best=False,
                )

        # ── End of training ────────────────────────────────────────────────
        self.metrics.logger.info("=" * 60)
        self.metrics.logger.info(
            f"Training complete. "
            f"Best val loss: {self.best_loss:.4f} | "
            f"PPL: {math.exp(self.best_loss):.2f}"
        )
        self.metrics.logger.info("=" * 60)
        self.metrics.save_history()

        print(f"\nBest model saved to: "
              f"{os.path.join(self.config.checkpoint_dir, 'best_model.pt')}")


# ─────────────────────────────────────────────────────────────
# HuggingFace export
# ─────────────────────────────────────────────────────────────

def export_to_huggingface(
    checkpoint_path: str,
    output_dir: str,
    tokenizer_path: str,
) -> None:
    """
    Export the trained model to HuggingFace format.

    Saves the model in a format compatible with AutoModel.from_pretrained(),
    allowing anyone to load the model with:
        model = AutoModel.from_pretrained("your-username/your-model")

    The process:
        1. Load the trained checkpoint
        2. Save weights in safetensors (safe and efficient format)
        3. Create config.json in HuggingFace format
        4. Copy tokenizer files
        5. Create the model card (README.md)

    After this step, use the HuggingFace CLI to publish:
        huggingface-cli upload your-username/minilm ./hf_export

    Args:
        checkpoint_path: Path to best_model.pt.
        output_dir:      Output directory for HF files.
        tokenizer_path:  Directory with BPE tokenizer files.
    """
    import shutil

    os.makedirs(output_dir, exist_ok=True)
    print(f"Exporting to HuggingFace format in '{output_dir}'...")

    # Load checkpoint
    ckpt        = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
    model_cfg_dict = ckpt["model_config"]
    # d_head is derived automatically in ModelConfig.__post_init__
    # and must not be passed as a constructor argument
    model_cfg_dict.pop("d_head", None)
    model_config   = ModelConfig(**model_cfg_dict)

    # Instantiate and load weights
    model = MiniLM(model_config)

    # If the model was trained with torch.compile(), the state_dict keys
    # will have a '_orig_mod.' prefix — strip it before loading
    state_dict = ckpt["model_state"]
    if any(k.startswith("_orig_mod.") for k in state_dict):
        state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}

    model.load_state_dict(state_dict)
    model.eval()

    # Save weights in safetensors (safer than .bin)
    # Note: weight tying means lm_head.weight and token_emb.weight share
    # the same tensor in memory. safetensors does not allow shared tensors,
    # so we save lm_head.weight as an independent copy.
    try:
        from safetensors.torch import save_file
        tensors = {}
        for k, v in model.state_dict().items():
            # Skip complex tensors (e.g. freqs_complex from RoPE) —
            # safetensors does not support complex dtypes.
            # These buffers are recomputed automatically on model load.
            if v.is_complex():
                continue
            tensors[k] = v.clone()   # clone breaks shared memory references
        save_file(tensors, os.path.join(output_dir, "model.safetensors"))
        print("  Weights saved to model.safetensors")
    except ImportError:
        # Fallback to pytorch_model.bin — supports complex tensors
        state_dict = {k: v for k, v in model.state_dict().items()
                      if not v.is_complex()}
        torch.save(state_dict, os.path.join(output_dir, "pytorch_model.bin"))
        print("  Weights saved to pytorch_model.bin")
        print("  (install safetensors for the recommended format: pip install safetensors)")

    # Save config.json in HuggingFace format
    hf_config = {
        "model_type":              "minilm",
        "architectures":           ["MiniLM"],
        "vocab_size":              model_config.vocab_size,
        "hidden_size":             model_config.d_model,
        "num_hidden_layers":       model_config.n_layers,
        "num_attention_heads":     model_config.n_heads,
        "intermediate_size":       model_config.d_ff,
        "max_position_embeddings": model_config.seq_len,
        "hidden_dropout_prob":     model_config.dropout,
        "torch_dtype":             "bfloat16",
        "transformers_version":    "4.0.0",
    }
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(hf_config, f, indent=2)
    print("  config.json saved")

    # Copy tokenizer files
    for fname in ["tokenizer.json", "vocab.json"]:
        src = os.path.join(tokenizer_path, fname)
        if os.path.exists(src):
            shutil.copy(src, os.path.join(output_dir, fname))
    print("  Tokenizer files copied")

    # Create model card (README.md)
    params_m = model_config.n_params / 1e6
    readme = f"""---
language:
- pt
- en
license: mit
tags:
- language-model
- bilingual
- portuguese
- english
- from-scratch
---

# MiniLM — Bilingual PT+EN Language Model

A decoder-only Transformer language model trained from scratch,
including a BPE tokenizer and training loop implemented without
high-level frameworks.

## Specifications

| Attribute            | Value                  |
|----------------------|------------------------|
| Parameters           | {params_m:.0f}M               |
| Architecture         | Transformer Decoder-only |
| Normalization        | RMSNorm                |
| Positional Encoding  | RoPE                   |
| FFN Activation       | SwiGLU                 |
| Vocabulary           | {model_config.vocab_size:,} tokens (BPE) |
| Max context          | {model_config.seq_len} tokens          |
| Languages            | Brazilian Portuguese + English |

## Training corpus

- **TinyStories** (EN): short synthetic stories ~60%
- **CulturaX PT** (PT-BR): curated Portuguese web ~40%

## How to use

```python
from bpe_tokenizer import BPETokenizer
from transformer import MiniLM, ModelConfig
import torch, json

tokenizer = BPETokenizer.load("./")

with open("config.json") as f:
    cfg = json.load(f)

model_config = ModelConfig(
    vocab_size=cfg["vocab_size"],
    d_model=cfg["hidden_size"],
    n_layers=cfg["num_hidden_layers"],
    n_heads=cfg["num_attention_heads"],
    d_ff=cfg["intermediate_size"],
    seq_len=cfg["max_position_embeddings"],
)
model = MiniLM(model_config)
model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
model.eval()

prompt = "Once upon a time"
input_ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long)
output = model.generate(input_ids, max_new_tokens=100, temperature=0.8, top_k=50)
print(tokenizer.decode(output[0].tolist()))
```

## Development

All training code is available in the repository:
- `bpe_tokenizer.py` — BPE tokenizer from scratch
- `data_pipeline.py` — Corpus preparation pipeline
- `transformer.py`   — Model architecture
- `training_loop.py` — Custom training loop

## Citation

```
@misc{{minilm2025,
  title={{MiniLM: A bilingual PT+EN language model built from scratch}},
  author={{André Costa}},
  year={{2026}},
  url={{https://huggingface.co/AndreCosta/minilm}}
}}
```
"""
    with open(os.path.join(output_dir, "README.md"), "w", encoding="utf-8") as f:
        f.write(readme)
    print("  README.md (model card) created")

    print(f"\nExport complete!")
    print(f"To publish on HuggingFace:")
    print(f"  huggingface-cli login")
    print(f"  huggingface-cli upload [your-username]/minilm {output_dir}")


# ─────────────────────────────────────────────────────────────
# Entry point
# ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="MiniLM Training")
    parser.add_argument("--mode", choices=["train", "export"],
                        default="train", help="Execution mode")
    parser.add_argument("--resume", type=str, default=None,
                        help="Path to checkpoint to resume from")
    parser.add_argument("--checkpoint", type=str, default=None,
                        help="Checkpoint to export (export mode)")
    parser.add_argument("--output-dir", type=str, default="./hf_export",
                        help="Output directory for HF export")
    parser.add_argument("--tokenizer-path", type=str, default="./tokenizer",
                        help="Path to the BPE tokenizer")
    parser.add_argument("--small", action="store_true",
                        help="Use Tiny config (~15M params) for quick tests")
    args = parser.parse_args()

    if args.mode == "train":
        # Model configuration
        if args.small:
            print("Using Tiny configuration (~15M params) for quick test")
            model_config = ModelConfig(
                vocab_size=16384,
                seq_len=512,   # must match the seq_len used in data_pipeline.py
                d_model=256,
                n_heads=4,
                n_layers=4,
                d_ff=768,
                dropout=0.1,
            )
            train_config = TrainingConfig(
                batch_size=4,
                accum_steps=2,
                max_steps=100,
                log_interval=10,
                eval_interval=50,
                save_interval=50,
            )
        else:
            model_config = ModelConfig()    # Small (~85M) by default
            train_config = TrainingConfig()

        print("\nModel configuration:")
        print(f"  {model_config.n_params / 1e6:.1f}M parameters")

        trainer = Trainer(model_config, train_config)
        trainer.train(resume_from=args.resume)

    elif args.mode == "export":
        if args.checkpoint is None:
            args.checkpoint = "./checkpoints/best_model.pt"
        export_to_huggingface(
            checkpoint_path=args.checkpoint,
            output_dir=args.output_dir,
            tokenizer_path=args.tokenizer_path,
        )