train.py

"""
Stage 3 — Train classifier.

Trains a regression head on Snowflake/snowflake-arctic-embed-m-v2.0 against
Gemma-annotated marketing scores from Stage 2.  Run twice to produce both
classifier heads:

  Head A (natural distribution, primary):
    python -m classification.train \\
        --shard-dir data/annotations \\
        --output-dir models/head_a \\
        --head natural \\
        --embed-cache-dir data/embed_cache

  Head B (balanced distribution, secondary):
    python -m classification.train \\
        --shard-dir data/annotations \\
        --output-dir models/head_b \\
        --head balanced \\
        --embed-cache-dir data/embed_cache

--embed-cache-dir (recommended): precomputes the 445k [CLS] embeddings once on
the first call (~25 min on A100) and caches them to disk.  Subsequent calls
(including Head B) load from cache in seconds, then train only the tiny MLP head
for 20 epochs — total runtime drops from ~10-20h to ~2-3h for both heads.

Both heads are evaluated after every epoch on the same 50k held-out set.
The checkpoint with the best held-out Spearman correlation is saved as
best_model.pt.  The final epoch checkpoint is saved as final_model.pt.

Artifacts written to --output-dir:
  best_model.pt         — state_dict at peak Spearman
  final_model.pt        — state_dict at end of training
  training_history.json — per-epoch loss, F1@3, Spearman
  run_config.json       — all hyperparameters for reproducibility
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, SequentialSampler
from tqdm import tqdm
from transformers import AutoTokenizer

from .dataset import (
    MAX_LEN,
    RANDOM_SEED,
    AnnotationDataset,
    EmbeddingDataset,
    compute_lds_weights,
    load_annotation_shards,
    make_balanced_sampler,
    make_splits,
)
from .metrics import compute_metrics
from .model import BACKBONE, MarketingClassifier

# ── Hyperparameters (pinned; do not change without spec update) ───────────────
EPOCHS = 20
LR = 3e-4
BATCH_SIZE = 32
EVAL_BATCH_SIZE = 64
EMBED_BATCH_SIZE = 256  # larger batch for one-shot embedding computation (no grad)


def weighted_mse_loss(
    pred: torch.Tensor,
    target: torch.Tensor,
    weight: torch.Tensor,
) -> torch.Tensor:
    return (weight * (pred - target) ** 2).mean()


def balanced_mse_loss(
    pred: torch.Tensor,
    target: torch.Tensor,
    noise_var: float = 1.0,
) -> torch.Tensor:
    """
    Balanced MSE — Ren et al. CVPR 2022 (BMC batch variant).

    Equivalent to cross-entropy over a Gaussian mixture where each
    training label in the batch is a mixture component with shared
    variance noise_var.  The "correct" component for sample i is i
    itself (diagonal of the logit matrix).

    Gradient: MSE gradient minus the weighted mean toward the batch
    label density, so gradient updates are smaller where labels are
    dense (naturally over-represented) and larger where sparse.
    """
    pred_col = pred.view(-1, 1)    # (B, 1)
    target_row = target.view(1, -1)  # (1, B)
    logits = -(pred_col - target_row) ** 2 / (2 * noise_var)  # (B, B)
    labels = torch.arange(pred.shape[0], device=pred.device)
    loss = F.cross_entropy(logits, labels)
    return loss * (2 * noise_var)


def _worker_init_fn(worker_id: int) -> None:  # noqa: ARG001
    # Seed each DataLoader worker from the worker's torch-assigned seed so
    # training data order is deterministic given the same --seed.
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)


def compute_and_cache_embeddings(
    model: MarketingClassifier,
    texts: list[str],
    scores: list[float],
    tokenizer: AutoTokenizer,
    emb_path: Path,
    scores_path: Path,
    max_len: int,
    batch_size: int,
    device: torch.device,
) -> EmbeddingDataset:
    """
    Run the frozen encoder once over all texts and cache [CLS] embeddings.

    If the cache files already exist (same path), loads them directly —
    so Head B reuses Head A's cache at no extra cost.

    Cache layout:
      {emb_path}        — float32 ndarray (N, hidden_size), ~0.75 GB for 445k docs
      {scores_path}     — float32 ndarray (N,)
    """
    if emb_path.exists() and scores_path.exists():
        print(f"  Loading cached embeddings from {emb_path}")
        return EmbeddingDataset(np.load(emb_path), np.load(scores_path).tolist())

    print(f"  Computing {len(texts):,} embeddings (batch_size={batch_size}) …")
    model.eval()
    all_embs: list[np.ndarray] = []
    use_autocast = device.type == "cuda"
    with torch.no_grad(), torch.autocast(device_type=device.type, enabled=use_autocast):
        for i in tqdm(range(0, len(texts), batch_size), desc="  embed", unit="batch"):
            batch_texts = texts[i : i + batch_size]
            enc = tokenizer(
                batch_texts,
                max_length=max_len,
                truncation=True,
                padding=True,          # pad to longest in batch (not full max_len)
                return_tensors="pt",
            )
            cls_emb = model.encode(
                enc["input_ids"].to(device),
                enc["attention_mask"].to(device),
            )
            all_embs.append(cls_emb.cpu().float().numpy())

    embeddings = np.concatenate(all_embs, axis=0)
    np.save(emb_path, embeddings)
    np.save(scores_path, np.array(scores, dtype=np.float32))
    print(f"  Cached {embeddings.shape} embeddings ({embeddings.nbytes / 1e9:.2f} GB) → {emb_path}")
    return EmbeddingDataset(embeddings, scores)


def train(
    shard_dir: Path,
    output_dir: Path,
    head: str,
    epochs: int = EPOCHS,
    lr: float = LR,
    batch_size: int = BATCH_SIZE,
    eval_batch_size: int = EVAL_BATCH_SIZE,
    max_len: int = MAX_LEN,
    seed: int = RANDOM_SEED,
    embed_cache_dir: Path | None = None,
    embed_batch_size: int = EMBED_BATCH_SIZE,
    lds_sigma: float = 1.0,
    noise_var: float = 1.0,
) -> None:
    if head not in {"natural", "balanced", "lds", "bmse"}:
        raise ValueError(f"--head must be 'natural', 'balanced', 'lds', or 'bmse', got {head!r}")

    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if torch.cuda.is_available():
        device = torch.device("cuda")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
    else:
        device = torch.device("cpu")
    print(f"Device: {device}")
    print(f"Head: {head}")

    # ── Data ──────────────────────────────────────────────────────────────────
    print("Loading annotation shards …")
    df = load_annotation_shards(shard_dir)
    print(f"  Total annotated docs: {len(df):,}")

    train_df, eval_df = make_splits(df, seed=seed)
    print(f"  Train: {len(train_df):,}  |  Eval: {len(eval_df):,}")

    tokenizer = AutoTokenizer.from_pretrained(BACKBONE, trust_remote_code=True)

    # ── Model ─────────────────────────────────────────────────────────────────
    # Loaded before dataset construction so the encoder can be used for
    # embedding precomputation when --embed-cache-dir is set.
    print(f"Loading backbone: {BACKBONE}")
    model = MarketingClassifier(backbone=BACKBONE).to(device)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(f"  Trainable params: {trainable:,} / {total:,}")

    # ── LDS weights (computed before dataset construction) ────────────────────
    train_weights: np.ndarray | None = None
    if head == "lds":
        print(f"Computing LDS weights (sigma={lds_sigma}) …")
        train_weights = compute_lds_weights(train_df["score"].tolist(), sigma=lds_sigma)
        print(f"  weight mean={train_weights.mean():.3f}  max={train_weights.max():.1f}  "
              f"min={train_weights.min():.3f}")
        if train_weights.max() > 1000:
            print(f"  WARNING: max weight {train_weights.max():.0f} > 1000 — "
                  "score-5 samples will dominate gradient when they appear. "
                  "Consider increasing sigma.")

    # ── Datasets ───────────────────────────────────────────────────────────────
    if embed_cache_dir is not None:
        embed_cache_dir.mkdir(parents=True, exist_ok=True)
        print("Embedding cache enabled — encoder runs once, head trains on cached vectors.")
        print("Train embeddings:")
        train_dataset = compute_and_cache_embeddings(
            model=model,
            texts=train_df["text"].tolist(),
            scores=train_df["score"].tolist(),
            tokenizer=tokenizer,
            emb_path=embed_cache_dir / "train_embeddings.npy",
            scores_path=embed_cache_dir / "train_scores.npy",
            max_len=max_len,
            batch_size=embed_batch_size,
            device=device,
        )
        if train_weights is not None:
            train_dataset = EmbeddingDataset(
                train_dataset.embeddings, train_dataset.scores, weights=train_weights
            )
        print("Eval embeddings:")
        eval_dataset = compute_and_cache_embeddings(
            model=model,
            texts=eval_df["text"].tolist(),
            scores=eval_df["score"].tolist(),
            tokenizer=tokenizer,
            emb_path=embed_cache_dir / "eval_embeddings.npy",
            scores_path=embed_cache_dir / "eval_scores.npy",
            max_len=max_len,
            batch_size=embed_batch_size,
            device=device,
        )

        def forward_batch(batch: dict) -> torch.Tensor:
            return model.score_from_embedding(batch["embedding"].to(device))

    else:
        train_dataset = AnnotationDataset(
            texts=train_df["text"].tolist(),
            scores=train_df["score"].tolist(),
            tokenizer=tokenizer,
            max_len=max_len,
            weights=train_weights,
        )
        eval_dataset = AnnotationDataset(
            texts=eval_df["text"].tolist(),
            scores=eval_df["score"].tolist(),
            tokenizer=tokenizer,
            max_len=max_len,
        )

        def forward_batch(batch: dict) -> torch.Tensor:
            return model(batch["input_ids"].to(device), batch["attention_mask"].to(device))

    # ── DataLoaders ────────────────────────────────────────────────────────────
    if head == "balanced":
        sampler = make_balanced_sampler(train_df["score"].tolist())
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            sampler=sampler,
            num_workers=4,
            pin_memory=True,
            worker_init_fn=_worker_init_fn,
        )
    else:  # "natural", "lds", "bmse" — all use random shuffle
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=4,
            pin_memory=True,
            worker_init_fn=_worker_init_fn,
        )

    eval_loader = DataLoader(
        eval_dataset,
        batch_size=eval_batch_size,
        sampler=SequentialSampler(eval_dataset),
        num_workers=4,
        pin_memory=True,
        worker_init_fn=_worker_init_fn,
    )

    optimizer = torch.optim.Adam(
        [p for p in model.parameters() if p.requires_grad],
        lr=lr,
    )

    # ── Output directory ──────────────────────────────────────────────────────
    output_dir.mkdir(parents=True, exist_ok=True)

    run_config = {
        "head": head,
        "lds": head == "lds",
        "lds_n_bins": 51,
        "lds_sigma": lds_sigma if head == "lds" else None,
        "bmse": head == "bmse",
        "noise_var": noise_var if head == "bmse" else None,
        "backbone": BACKBONE,
        "epochs": epochs,
        "lr": lr,
        "batch_size": batch_size,
        "max_len": max_len,
        "seed": seed,
        "embed_cache_dir": str(embed_cache_dir) if embed_cache_dir else None,
        "embed_batch_size": embed_batch_size if embed_cache_dir else None,
        "train_size": len(train_df),
        "eval_size": len(eval_df),
    }
    with open(output_dir / "run_config.json", "w") as f:
        json.dump(run_config, f, indent=2)

    # ── Training loop ─────────────────────────────────────────────────────────
    EARLY_STOP_PATIENCE = 3

    best_spearman = -1.0
    best_epoch = -1
    epochs_without_improvement = 0
    history: list[dict] = []

    for epoch in range(1, epochs + 1):
        model.train()
        train_loss = 0.0
        n_train = 0

        nan_abort = False
        for batch_idx, batch in enumerate(tqdm(train_loader, desc=f"Epoch {epoch}/{epochs} [train]", leave=False)):
            targets = batch["score"].to(device)
            preds = forward_batch(batch)

            if head == "bmse":
                loss = balanced_mse_loss(preds, targets, noise_var=noise_var)
            else:
                weights = batch["weight"].to(device)
                loss = weighted_mse_loss(preds, targets, weights)
                if epoch == 1 and batch_idx == 0:
                    print(f"  weight sanity — mean={weights.mean():.3f}  max={weights.max():.1f}")

            if torch.isnan(loss):
                print(f"\n  ABORT: loss=NaN at epoch {epoch} batch {batch_idx}. "
                      "Check noise_var or label distribution.")
                nan_abort = True
                break

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * len(targets)
            n_train += len(targets)

        if nan_abort:
            break

        train_loss /= n_train

        # ── Evaluation ────────────────────────────────────────────────────────
        model.eval()
        all_preds: list[float] = []
        all_targets: list[float] = []

        with torch.no_grad():
            for batch in tqdm(eval_loader, desc=f"Epoch {epoch}/{epochs} [eval]", leave=False):
                preds = forward_batch(batch)
                all_preds.extend(preds.cpu().tolist())
                all_targets.extend(batch["score"].tolist())

        metrics = compute_metrics(all_preds, all_targets)
        row = {"epoch": epoch, "train_loss": train_loss, **metrics}
        history.append(row)

        print(
            f"Epoch {epoch:02d} | loss={train_loss:.4f}"
            f" | F1@3={metrics['f1_at_3']:.4f}"
            f" | Spearman={metrics['spearman']:.4f}"
        )

        if metrics["spearman"] > best_spearman:
            best_spearman = metrics["spearman"]
            best_epoch = epoch
            epochs_without_improvement = 0
            torch.save(model.state_dict(), output_dir / "best_model.pt")
            print(f"  ✓ New best Spearman={best_spearman:.4f} (epoch {best_epoch})")
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= EARLY_STOP_PATIENCE:
                print(
                    f"  Early stop: no improvement for {EARLY_STOP_PATIENCE} epochs "
                    f"(best epoch {best_epoch}, Spearman={best_spearman:.4f})"
                )
                break

    # ── Save final artifacts ───────────────────────────────────────────────────
    torch.save(model.state_dict(), output_dir / "final_model.pt")
    with open(output_dir / "training_history.json", "w") as f:
        json.dump(history, f, indent=2)

    if best_epoch < 1:
        raise RuntimeError(
            "No checkpoint was saved — Spearman was NaN every epoch. "
            "Check for model collapse (constant predictions)."
        )
    best_metrics = history[best_epoch - 1]
    print(
        f"\nBest checkpoint: epoch {best_epoch}"
        f" | Spearman={best_metrics['spearman']:.4f}"
        f" | F1@3={best_metrics['f1_at_3']:.4f}"
    )
    print(f"Artifacts saved to {output_dir}")

    if best_metrics["spearman"] < 0.80:
        print(
            f"\nWARNING: best Spearman {best_metrics['spearman']:.4f} < 0.80 target."
            " Consider more epochs or a higher learning rate."
        )
    if best_metrics["f1_at_3"] < 0.82:
        print(
            f"WARNING: F1@3 {best_metrics['f1_at_3']:.4f} < 0.82 target."
        )


if __name__ == "__main__":
    ap = argparse.ArgumentParser(
        description="Train a marketing quality regression classifier."
    )
    ap.add_argument(
        "--shard-dir", type=Path, required=True,
        help="Directory with Stage 2 annotation shards (shard_*.parquet)",
    )
    ap.add_argument(
        "--output-dir", type=Path, required=True,
        help="Directory to save model checkpoints and training history",
    )
    ap.add_argument(
        "--head", choices=["natural", "balanced", "lds", "bmse"], required=True,
        help=(
            "'natural' = Head A (plain MSE, shuffle sampler); "
            "'balanced' = Head B (WeightedRandomSampler); "
            "'lds' = Label Distribution Smoothing (weighted MSE); "
            "'bmse' = Balanced MSE (Ren et al. CVPR 2022)"
        ),
    )
    ap.add_argument(
        "--lds-sigma", type=float, default=1.0,
        help="Gaussian kernel sigma for LDS weight smoothing (default 1.0, in bin units)",
    )
    ap.add_argument(
        "--noise-var", type=float, default=1.0,
        help="Gaussian noise variance for Balanced MSE (default 1.0)",
    )
    ap.add_argument(
        "--embed-cache-dir", type=Path, default=None,
        help=(
            "Directory to cache pre-computed [CLS] embeddings. "
            "Encoder runs once on first call; subsequent calls (e.g. Head B) load from disk. "
            "Strongly recommended for GPU runs — reduces total training time from ~20h to ~3h."
        ),
    )
    ap.add_argument("--epochs", type=int, default=EPOCHS)
    ap.add_argument("--lr", type=float, default=LR)
    ap.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    ap.add_argument("--eval-batch-size", type=int, default=EVAL_BATCH_SIZE)
    ap.add_argument("--embed-batch-size", type=int, default=EMBED_BATCH_SIZE,
                    help="Batch size for embedding precomputation (default 256, larger = faster)")
    ap.add_argument("--max-len", type=int, default=MAX_LEN)
    ap.add_argument("--seed", type=int, default=RANDOM_SEED)
    args = ap.parse_args()

    train(
        shard_dir=args.shard_dir,
        output_dir=args.output_dir,
        head=args.head,
        epochs=args.epochs,
        lr=args.lr,
        batch_size=args.batch_size,
        eval_batch_size=args.eval_batch_size,
        max_len=args.max_len,
        seed=args.seed,
        embed_cache_dir=args.embed_cache_dir,
        embed_batch_size=args.embed_batch_size,
        lds_sigma=args.lds_sigma,
        noise_var=args.noise_var,
    )