presentation/src/finetune.py

"""Phase 2: fine-tune ``EncoderForClassification`` -> legitimate / phishing / spam.

The instructor's "experimentation dials" (number of layers, epochs, loss function, plus
learning rate, batch size, frozen layers, synthetic data) are first-class CLI flags, and
every run leaves ONE JSON line in reports/experiments.jsonl so the experimentation table
(``finetune.py list``) is generated from evidence rather than typed up afterwards.

Subcommands::

    train [--run-name R] [--init pretrained:<dir>|scratch] [--n-layers N] [--freeze-layers K]
          [--loss cross_entropy|weighted_cross_entropy|focal] [--focal-gamma G] [--lr LR]
          [--batch-size B] [--micro-batch M] [--epochs E] [--max-steps S] [--eval-every S]
          [--log-every S] [--synthetic on|off] [--limit-rows N] [--num-labels 2|3] [--cpu]
          [--resume] [--overwrite]
    plot  --run-name R [--out PNG]       reports/finetune_curves.png
    list                                  experiments.jsonl as a markdown table

ML decisions, in brief:

* **Three classes, partial labels (RUN_LOG entry 6).** ``finetune.num_labels`` (3) sizes the
  classification head; config ``model:`` is never touched because the live pretraining run
  refuses to resume if its model snapshot changes. Every row carries ``label3`` in
  {0 legitimate, 1 phishing, 2 spam, -1 unwanted-but-unsplit} and ``label_set`` in
  {legit, phish, spam, unwanted}; the binary ``label`` column (1 = phishing OR spam) stays.
  The TREC-05/07 and CEAS-08 positives never separated spam from phishing, so they train on
  the probability MASS of {phishing, spam}: per row an allowed-class mask ``[B, 3]`` built
  from label_set, and the loss is the partial-label negative log-likelihood
  ``logsumexp(all logits) - logsumexp(logits[allowed])`` = -log sum_{c in allowed} p_c
  (Jin & Ghahramani 2002; Cour et al. 2011). For a single allowed class this IS
  ``F.cross_entropy``; for an "unwanted" row it is ``-log(p_phish + p_spam)``, which still
  shapes the legitimate-vs-unwanted decision without asserting which kind it was. With
  ``num_labels: 2`` the old binary behaviour is reproduced exactly (label3 ignored; the mask
  is the one-hot label) so the binary baseline runs remain reproducible.
* **Two metric sets.** (a) BINARY roll-up on all val rows: P(unwanted) = p1 + p2 vs the
  binary label at threshold 0.5 -> accuracy/precision/recall/F1/ROC-AUC/PR-AUC with the same
  definitions as the binary runs (so results stay comparable). (b) THREE-WAY on rows with
  label3 in {0, 1, 2}: argmax over 3 classes -> 3x3 confusion, per-class P/R/F1, macro-F1
  over all three classes, plus a phishing-vs-spam confusion restricted to true unwanted rows
  (argmax over {phishing, spam} only: "given it is unwanted, which kind?"). Display names
  are {0: Legitimate, 1: Phishing, 2: Spam}; the roll-up is "Unwanted (phishing or spam)".
  Selection / early stopping follow ``finetune.early_stopping.metric``: ``val_f1_unwanted``
  (default; = today's binary F1) or ``val_macro_f1_3way``.

* **Data path.** Train = finetune_train.parquet (+ optional synthetic rows), selection =
  finetune_val.parquet. test_holdout.parquet (TREC-06) is refused by file name AND by the
  ``source`` column of anything loaded, so a misconfigured path cannot leak the holdout.
  Texts are encoded once with ``tokenizer.encode_batch`` (truncate to model.max_seq_len,
  pad to a multiple of finetune.dynamic_padding_multiple), stored trimmed, and re-padded per
  batch with the identical rule; the first batch is checked tensor-for-tensor against a
  direct ``encode_batch`` call so the app / evaluate.py path (which calls ``encode_batch``
  per request) sees exactly the inputs the classifier was trained on. No sequence packing:
  the bool key-padding mask takes SDPA's memory-efficient kernel (research_model-math.md sec. 5).
* **DataLoader.** Every batch goes through ``torch.utils.data.DataLoader`` (``make_loader``):
  a row-index ``Dataset``, the seeded length-bucketed batch order as ``batch_sampler`` and a
  collate function that pads per micro-batch. Shuffling, batching and preprocessing therefore
  live in the standard PyTorch data pipeline the course reference uses.
* **Length bucketing.** Each epoch: seeded permutation -> mega-batches of 50 x batch_size ->
  sort by token length inside each -> cut into batches -> shuffle the batch order. Padding
  waste drops from ~50% (random batches over a 20..512 length spread) to a few percent while
  the batch composition stays random at the mega-batch level. The order is a pure function
  of (seed, epoch) so ``--resume`` can skip exactly the batches already consumed.
* **Init.** ``pretrained:<dir>`` loads ``encoder.*`` from a Phase 1 checkpoint with strict
  key matching (``EncoderForClassification.from_pretrained_encoder``); ``scratch`` is the
  ablation that shows what MLM pretraining buys. ``--n-layers N`` keeps the FIRST N encoder
  layers (``encoder.layers.0..N-1``) of the pretrained stack and drops the rest; the final
  LayerNorm is kept although it was trained on layer-24 output (fine-tuning re-fits it).
  ``--freeze-layers K`` freezes embeddings + layers 0..K-1: cheaper steps (autograd stops at
  layer K) and less catastrophic forgetting on a small train set.
* **Loss dial.** All three variants are defined over the partial-label NLL above (``nll``):
  cross_entropy = mean(nll); weighted_cross_entropy = sum(w_r nll_r) / sum(w_r) with
  inverse-frequency class weights w_c = N / (C * n_c) from the fully-labelled rows of the
  train split actually used and w_r = mean of w over the row's allowed classes (so an
  "unwanted" row gets mean(w_phish, w_spam)) -- the sum-of-weights normalisation is
  ``F.cross_entropy(weight=)``'s, so the two agree on fully-labelled batches; focal loss
  (Lin et al. 2017) ``(1 - exp(-nll))^gamma * nll`` = ``-(1 - p_t)^gamma log p_t`` without
  alpha, where p_t is the allowed-mass. Loss is computed in fp32 on bf16-autocast logits.
* **Optimizer / schedule.** The same ``build_optimizer`` (AdamW, no weight decay on biases
  and LayerNorm, fused) and the same fwd/bwd -> clip 1.0 -> step -> zero_grad(set_to_none)
  pattern the smoke test measured. Linear warmup over finetune.warmup_fraction of the total
  optimizer steps, then linear decay to 0 (BERT fine-tuning schedule). Gradient
  accumulation: a logical batch of ``--batch-size`` rows is split into ``--micro-batch``
  chunks (default 16, the smoke test's 13.65 GiB point for 16 x 512 with the heavier MLM
  head); each chunk's loss is scaled by its row share so the step gradient is the mean over
  the logical batch even when the last batch of an epoch is short.
* **Selection / early stopping.** Full validation every ``--eval-every`` steps and at every
  epoch end; best/ is rewritten whenever the selection metric (val F1 of the unwanted
  roll-up at threshold 0.5 by default, or 3-way macro-F1) improves;
  finetune.early_stopping.patience evaluations without improvement stop the run. The
  F1-optimal threshold on val is recorded next to the 0.5 decision threshold so evaluate.py
  / the app can show both; selection itself uses 0.5 to stay comparable across runs.
  best/config.json carries num_labels, the label names, the selection metric, both metric
  sets and the thresholds.
* **--cpu (dev only).** Runs the identical code path on the CPU with whatever config it is
  given (tiny encoder + fixture parquet files) so the whole script can be exercised while
  the GPU is busy; autocast and the VRAM guard are disabled and torch.cuda is never initialised:
  device info, allocator stats and the CUDA RNG snapshot are gated on the chosen device, and
  CUDA_VISIBLE_DEVICES is cleared for the process before the runtime is first queried (torch
  2.14's ``AdamW.step`` health check would otherwise create a context on any visible GPU). If a
  device is still visible after that, the run refuses to start. The default path still requires CUDA.
* **Resume.** last/ holds model + trainer_state.pt (optimizer, scheduler, step, epoch,
  batch-in-epoch, Python/NumPy/torch CPU/torch CUDA RNG states, best score, patience counter,
  cumulative seconds, the resolved dials). ``--resume`` restores all of it and skips the
  already-consumed batches of the current epoch. Dials that shape the optimizer state or the
  batch order (init, n_layers, freeze_layers, batch_size, seed) are refused on resume; ``--lr``
  is re-applied to the restored optimizer/scheduler so the record never claims an lr that was
  not used; metrics.csv / epochs.csv rows logged after the checkpoint (a crash between
  checkpoints) are dropped before appending; a run already at its --max-steps / --epochs is
  reported as "nothing to resume" without training an extra step. SIGINT/SIGTERM save last/
  and exit 0.
* **VRAM guard.** ``PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`` before CUDA starts;
  after every optimizer step the reserved peak is compared with config.vram.hard_cap_gb
  (GiB, the smoke test's convention) and the run aborts (after saving last/) on a breach.
"""

from __future__ import annotations

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import argparse  # noqa: E402
import contextlib  # noqa: E402
import csv  # noqa: E402
import json  # noqa: E402
import logging  # noqa: E402
import math  # noqa: E402
import random  # noqa: E402
import signal  # noqa: E402
import sys  # noqa: E402
import time  # noqa: E402
import zlib  # noqa: E402
from dataclasses import asdict, dataclass, field, fields  # noqa: E402
from pathlib import Path  # noqa: E402
from typing import Any  # noqa: E402

sys.path.insert(0, str(Path(__file__).parent))

import numpy as np  # noqa: E402
import pyarrow.parquet as pq  # noqa: E402
import torch  # noqa: E402
import torch.nn.functional as F  # noqa: E402
from safetensors import safe_open  # noqa: E402
from safetensors.torch import load_file, save_file  # noqa: E402
from torch import Tensor  # noqa: E402
from torch.utils.data import DataLoader, Dataset  # noqa: E402

import common  # noqa: E402
from model import EncoderConfig, EncoderForClassification, count_parameters  # noqa: E402
from smoke_test import build_optimizer  # noqa: E402  (the measured optimizer code path)
from tokenizer import encode_batch, load_tokenizer  # noqa: E402

LOG = logging.getLogger("finetune")
GIB = float(2**30)

TRAIN_FILE = "finetune_train.parquet"
VAL_FILE = "finetune_val.parquet"
SYNTHETIC_FILE = "synthetic.parquet"
HOLDOUT_FILE = "test_holdout.parquet"
LOSSES = ("cross_entropy", "weighted_cross_entropy", "focal")
MEGA_BATCH_FACTOR = 50
ENCODE_CHUNK = 2048
DECISION_THRESHOLD = 0.5
METRIC_KEYS = ("accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc")
RUN_LOG_MIN_STEPS = 100
# Dials whose change on --resume would leave the restored optimizer state / batch order / head shape inconsistent.
RESUME_FIXED_DIALS = ("init", "n_layers", "freeze_layers", "batch_size", "seed", "num_labels")
# The only implemented values of the config-only dials; anything else in config.yaml is an error, not a no-op.
IMPLEMENTED = {"schedule": "linear_decay", "pooling": "cls"}
CONFIG_PATH = {"schedule": "finetune.schedule", "pooling": "finetune.pooling",
               "selection_metric": "finetune.early_stopping.metric", "num_labels": "finetune.num_labels"}

# ---- three-way label contract (RUN_LOG entry 6) -----------------------------------------------
LABEL_NAMES = {0: "Legitimate", 1: "Phishing", 2: "Spam"}
UNWANTED_NAME = "Unwanted (phishing or spam)"
LABEL_SETS: dict[str, tuple[int, ...]] = {"legit": (0,), "phish": (1,), "spam": (2,), "unwanted": (1, 2)}
LABEL3_OF_SET = {"legit": 0, "phish": 1, "spam": 2, "unwanted": -1}
NUM_LABELS_ALLOWED = (2, 3)
# Selection metrics. val_f1_unwanted is today's binary F1 (p_unwanted = p1 + p2 at 0.5) so results stay
# comparable with the binary runs; val_f1_phishing is the pre-entry-6 name of the same number and is
# accepted so the recorded binary baselines resolve. val_macro_f1_3way needs num_labels == 3.
SELECTION_METRICS = ("val_f1_unwanted", "val_macro_f1_3way", "val_f1_phishing")
SELECTION_KEY = {"val_f1_unwanted": "f1", "val_f1_phishing": "f1", "val_macro_f1_3way": "macro_f1_3way"}
CLASS_SLUGS = ("legitimate", "phishing", "spam")
CM3_KEYS = tuple(f"cm3_{a[0]}{b[0]}" for a in CLASS_SLUGS for b in CLASS_SLUGS)   # cm3_ll ... cm3_ss (true, pred)
PVS_KEYS = ("pvs_phish_as_phish", "pvs_phish_as_spam", "pvs_spam_as_phish", "pvs_spam_as_spam")
THREE_WAY_KEYS = ("n_3way", "accuracy_3way", "macro_f1_3way",
                  *(f"{m}_{s}" for s in CLASS_SLUGS for m in ("precision", "recall", "f1")),
                  *CM3_KEYS, "n_unwanted_3way", *PVS_KEYS)
# validated categorical palette (dataviz skill): blue, orange, aqua, yellow
PALETTE = ("#2a78d6", "#eb6834", "#1baf7a", "#eda100")


# ---------------------------------------------------------------------------------- dials
@dataclass
class TrainArgs:
    """Every dial of a run, resolved from CLI > saved state (on --resume) > config.yaml.

    Stored verbatim in trainer_state.pt and experiments.jsonl so a run is reproducible from
    its record alone.
    """

    run_name: str = "finetune"
    init: str = "pretrained"          # "pretrained" | "pretrained:<path>" | "scratch"
    n_layers: int = 24
    freeze_layers: int = 0
    loss: str = "cross_entropy"
    focal_gamma: float = 2.0
    lr: float = 3e-5
    batch_size: int = 32
    micro_batch: int = 16
    epochs: int = 3
    max_steps: int | None = None
    eval_every: int = 500
    log_every: int = 20
    eval_batch_size: int = 64
    warmup_fraction: float = 0.06
    weight_decay: float = 0.01
    betas: tuple[float, float] = (0.9, 0.98)
    eps: float = 1e-6
    grad_clip: float = 1.0
    patience: int = 2
    schedule: str = "linear_decay"            # config-only: finetune.schedule
    selection_metric: str = "val_f1_unwanted" # config: finetune.early_stopping.metric (val_f1_unwanted | val_macro_f1_3way)
    pooling: str = "cls"                      # config-only: finetune.pooling (EncoderForClassification reads token 0)
    num_labels: int = 3                       # finetune.num_labels (NOT model.num_labels): 3 = legit/phish/spam, 2 = binary baseline
    synthetic: str = "off"
    limit_rows: int | None = None
    seed: int = 42
    cpu: bool = False                         # dev only: run the whole script on the CPU

    @classmethod
    def resolve(cls, ns: argparse.Namespace, cfg: dict, saved: dict | None) -> "TrainArgs":
        fcfg, mcfg = cfg["finetune"], cfg["model"]
        pcfg = cfg.get("pretrain", {})
        base = cls(
            n_layers=int(mcfg["n_layers"]),
            # Absent key -> 2 so a config.yaml predating entry 6 still reproduces the binary baseline.
            num_labels=int(fcfg.get("num_labels", 2)),
            freeze_layers=int(fcfg.get("freeze_encoder_layers", 0)),
            loss=str(fcfg.get("loss", "cross_entropy")),
            focal_gamma=float(fcfg.get("focal_gamma", 2.0)),
            lr=float(fcfg["lr"]),
            batch_size=int(fcfg["batch_size"]),
            micro_batch=int(fcfg.get("micro_batch", 16)),
            epochs=int(fcfg["epochs"]),
            eval_every=int(fcfg.get("eval_every_steps", fcfg.get("checkpoint_every_steps", 500))),
            log_every=int(fcfg.get("log_every_steps", pcfg.get("log_every_steps", 20))),
            eval_batch_size=int(fcfg.get("eval_batch_size", 64)),
            warmup_fraction=float(fcfg["warmup_fraction"]),
            weight_decay=float(fcfg.get("weight_decay", pcfg.get("weight_decay", 0.01))),
            betas=tuple(fcfg.get("betas", pcfg.get("betas", (0.9, 0.98)))),
            eps=float(fcfg.get("eps", pcfg.get("eps", 1e-6))),
            grad_clip=float(fcfg.get("grad_clip", pcfg.get("grad_clip", 1.0))),
            patience=int(fcfg["early_stopping"]["patience"]),
            schedule=str(fcfg.get("schedule", IMPLEMENTED["schedule"])),
            selection_metric=str(fcfg["early_stopping"].get("metric", "val_f1_unwanted")),
            pooling=str(fcfg.get("pooling", IMPLEMENTED["pooling"])),
            synthetic="on" if cfg["data"].get("synthetic", {}).get("enabled") else "off",
            seed=int(cfg["seed"]),
        )
        if saved:
            for k, v in saved.items():
                if k in {f.name for f in fields(cls)}:
                    setattr(base, k, tuple(v) if k == "betas" else v)
        for k, v in vars(ns).items():
            if k in {f.name for f in fields(cls)} and v is not None and not (k == "cpu" and v is False):
                setattr(base, k, v)
        if getattr(ns, "micro_batch", None) is None:
            # A smaller --batch-size without an explicit --micro-batch means "one micro-batch per step".
            base.micro_batch = min(base.micro_batch, base.batch_size)
        for key, allowed in IMPLEMENTED.items():
            if getattr(base, key) != allowed:
                raise ValueError(f"config {CONFIG_PATH[key]}={getattr(base, key)!r} is not implemented; only {allowed!r} is "
                                 f"(editing it would otherwise silently change nothing)")
        if base.num_labels not in NUM_LABELS_ALLOWED:
            raise ValueError(f"{CONFIG_PATH['num_labels']} must be one of {NUM_LABELS_ALLOWED}, got {base.num_labels!r}")
        if base.selection_metric not in SELECTION_METRICS:
            raise ValueError(f"{CONFIG_PATH['selection_metric']}={base.selection_metric!r} is not implemented; "
                             f"one of {SELECTION_METRICS[:2]} (val_f1_phishing is the legacy alias of val_f1_unwanted)")
        if base.selection_metric == "val_macro_f1_3way" and base.num_labels != 3:
            raise ValueError("selection metric val_macro_f1_3way needs finetune.num_labels: 3")
        if base.loss not in LOSSES:
            raise ValueError(f"--loss must be one of {LOSSES}, got {base.loss!r}")
        if base.batch_size % base.micro_batch:
            raise ValueError(f"--batch-size {base.batch_size} must be a multiple of --micro-batch {base.micro_batch}")
        if not 0 <= base.freeze_layers <= base.n_layers:
            raise ValueError(f"--freeze-layers {base.freeze_layers} outside [0, {base.n_layers}]")
        return base

    def to_record(self) -> dict[str, Any]:
        d = asdict(self)
        d["betas"] = list(self.betas)
        return d


# ----------------------------------------------------------------------------------- data
@dataclass
class Labelled:
    """The label columns of one parquet split after validation against the three-way contract."""

    texts: list[str]
    labels: np.ndarray          # int64 [n] binary: 1 = phishing OR spam (unchanged from the binary build)
    label3: np.ndarray          # int64 [n] in {0, 1, 2, -1}; with num_labels == 2 this is a copy of ``labels``
    allowed: np.ndarray         # bool [n, C] allowed-class mask built from label_set (one-hot for full labels)
    sources: list[str]

    def __len__(self) -> int:
        return len(self.texts)


@dataclass
class EncodedSplit:
    """Token ids for one split, truncated to max_len and stored trimmed (flat int32 + offsets)."""

    name: str
    flat: np.ndarray            # int32 concatenation of every row's ids
    offsets: np.ndarray         # int64 [n + 1]
    labels: np.ndarray          # int64 [n] binary roll-up target (metrics set a)
    label3: np.ndarray          # int64 [n] three-way target, -1 = unwanted-but-unsplit (metrics set b uses >= 0)
    allowed: np.ndarray         # bool [n, C] partial-label mask (the loss target)
    sources: list[str]
    texts: list[str]            # kept for the encode_batch consistency check only

    def __len__(self) -> int:
        return int(self.labels.shape[0])

    @property
    def lengths(self) -> np.ndarray:
        return np.diff(self.offsets)

    @property
    def num_labels(self) -> int:
        return int(self.allowed.shape[1])

    def row(self, i: int) -> np.ndarray:
        return self.flat[self.offsets[i]:self.offsets[i + 1]]

    def collate(self, idx: np.ndarray, multiple: int, device: torch.device) -> tuple[Tensor, Tensor, Tensor, Tensor]:
        """Pad to the longest row in the batch rounded up to ``multiple`` (encode_batch's rule).

        Returns (ids int64 [B, L], key-padding mask bool [B, L], binary labels int64 [B],
        allowed-class mask bool [B, C]).
        """
        lens = self.lengths[idx]
        L = int(lens.max())
        if multiple:
            L = int(math.ceil(L / multiple) * multiple)
        ids = np.zeros((len(idx), L), dtype=np.int64)
        mask = np.zeros((len(idx), L), dtype=np.bool_)
        for r, i in enumerate(idx):
            n = int(lens[r])
            ids[r, :n] = self.row(int(i))
            mask[r, :n] = True
        return (torch.from_numpy(ids).to(device, non_blocking=True),
                torch.from_numpy(mask).to(device, non_blocking=True),
                torch.from_numpy(self.labels[idx]).to(device, non_blocking=True),
                torch.from_numpy(self.allowed[idx]).to(device, non_blocking=True))


def allowed_mask(label3: np.ndarray, label_set: list[str], num_labels: int) -> np.ndarray:
    """Per-row allowed-class mask ``[n, C]`` from label_set (3 classes) or the binary label (2 classes)."""
    out = np.zeros((len(label3), num_labels), dtype=np.bool_)
    if num_labels == 2:
        out[np.arange(len(label3)), label3] = True
        return out
    for r, s in enumerate(label_set):
        out[r, list(LABEL_SETS[s])] = True
    return out


def _assert_not_holdout(path: Path, holdout_source: str) -> None:
    """Hard rule: TREC-06 is touched only by evaluate.py. Refuse by file name and by source value."""
    if path.name == HOLDOUT_FILE:
        raise RuntimeError(f"refusing to load {path.name}: the test holdout is only read by evaluate.py")
    names = set(pq.read_schema(path).names)
    if "source" in names:
        values = {str(v) for v in pq.read_table(path, columns=["source"]).column("source").unique().to_pylist()}
        norm = holdout_source.lower().replace("_", "-")
        hits = {v for v in values if v.strip().lower().replace("_", "-") == norm}
        if hits:
            raise RuntimeError(f"{path.name} contains holdout source {hits}: test-set leakage")


def read_split(path: Path, holdout_source: str, limit: int | None, seed: int, num_labels: int,
               force_source: str | None = None) -> Labelled:
    """Read text + labels of one split and validate the three-way contract.

    With ``num_labels == 3`` the ``label3`` / ``label_set`` columns are REQUIRED (a split built
    before RUN_LOG entry 6 is an error, not a silent all-"unwanted" run) and checked for
    internal consistency (label_set <-> label3 <-> binary label). With ``num_labels == 2``
    they are ignored and the binary label is the target, reproducing the binary baseline.
    """
    _assert_not_holdout(path, holdout_source)
    names = set(pq.read_schema(path).names)
    three = num_labels == 3
    if three and not {"label3", "label_set"} <= names:
        raise ValueError(f"{path.name}: finetune.num_labels is 3 but the columns label3/label_set are missing "
                         f"(have {sorted(names)}); rebuild the split with the three-way data_prep")
    cols = ["text", "label"] + (["source"] if "source" in names else []) + (["label3", "label_set"] if three else [])
    table = pq.read_table(path, columns=cols)
    if limit is not None and table.num_rows > limit:
        # Seeded random subsample (not a head slice) so a dev run still sees every source.
        idx = np.sort(np.random.default_rng([seed, zlib.crc32(path.name.encode())]).choice(table.num_rows, size=limit, replace=False))
        table = table.take(idx)
    texts = [t or "" for t in table.column("text").to_pylist()]
    labels = np.asarray(table.column("label").to_pylist(), dtype=np.int64)
    if not set(np.unique(labels).tolist()) <= {0, 1}:
        raise ValueError(f"{path.name}: labels must be 0/1, got {np.unique(labels).tolist()}")
    if force_source is not None:
        sources = [force_source] * len(texts)
    else:
        sources = [str(s) for s in table.column("source").to_pylist()] if "source" in names else ["?"] * len(texts)
    if not three:
        return Labelled(texts, labels, labels.copy(), allowed_mask(labels, [], 2), sources)

    label3 = np.asarray(table.column("label3").to_pylist(), dtype=np.int64)
    label_set = [str(s) for s in table.column("label_set").to_pylist()]
    bad_sets = sorted(set(label_set) - set(LABEL_SETS))
    if bad_sets:
        raise ValueError(f"{path.name}: label_set values outside {sorted(LABEL_SETS)}: {bad_sets}")
    expected3 = np.fromiter((LABEL3_OF_SET[s] for s in label_set), dtype=np.int64, count=len(label_set))
    if not np.array_equal(expected3, label3):
        n_bad = int((expected3 != label3).sum())
        raise ValueError(f"{path.name}: {n_bad} row(s) where label3 disagrees with label_set "
                         f"(contract: legit->0, phish->1, spam->2, unwanted->-1)")
    expected_bin = (label3 != 0).astype(np.int64)
    if not np.array_equal(expected_bin, labels):
        raise ValueError(f"{path.name}: {int((expected_bin != labels).sum())} row(s) where the binary label disagrees "
                         f"with label3 (label must be 1 for phish/spam/unwanted, 0 for legit)")
    return Labelled(texts, labels, label3, allowed_mask(label3, label_set, 3), sources)


def concat_labelled(a: Labelled, b: Labelled) -> Labelled:
    return Labelled(a.texts + b.texts, np.concatenate([a.labels, b.labels]), np.concatenate([a.label3, b.label3]),
                    np.concatenate([a.allowed, b.allowed]), a.sources + b.sources)


def label_counts(label3: np.ndarray, num_labels: int) -> dict[str, int]:
    """Human-readable class counts for logs and records (keys: display names + 'unwanted')."""
    if num_labels == 2:
        return {LABEL_NAMES[0]: int((label3 == 0).sum()), UNWANTED_NAME: int((label3 == 1).sum())}
    out = {LABEL_NAMES[c]: int((label3 == c).sum()) for c in range(3)}
    out["unwanted (unsplit)"] = int((label3 == -1).sum())
    return out


def encode_split(name: str, tok, data: Labelled, max_len: int, multiple: int) -> EncodedSplit:
    """Encode once with the shared ``encode_batch`` contract, then store rows trimmed of padding."""
    t0 = common.timer()
    texts, labels = data.texts, data.labels
    rows: list[np.ndarray] = []
    for i in range(0, len(texts), ENCODE_CHUNK):
        ids, mask = encode_batch(tok, texts[i:i + ENCODE_CHUNK], max_len, pad_to_multiple_of=multiple)
        n = mask.sum(dim=1).numpy()
        ids_np = ids.numpy()
        rows.extend(ids_np[r, :n[r]].astype(np.int32) for r in range(ids_np.shape[0]))
    lengths = np.fromiter((len(r) for r in rows), dtype=np.int64, count=len(rows))
    offsets = np.zeros(len(rows) + 1, dtype=np.int64)
    np.cumsum(lengths, out=offsets[1:])
    flat = np.concatenate(rows) if rows else np.zeros(0, dtype=np.int32)
    split = EncodedSplit(name, flat, offsets, labels, data.label3, data.allowed, data.sources, texts)
    LOG.info("%s: %d rows encoded in %.1fs; tokens mean %.0f / median %.0f / max %d; %.1f%% truncated at %d; "
             "binary positives %d (%.1f%%); classes %s", name, len(rows), common.timer() - t0, lengths.mean(),
             np.median(lengths), lengths.max(), 100.0 * float((lengths >= max_len).mean()), max_len, int(labels.sum()),
             100.0 * float(labels.mean()), label_counts(data.label3, split.num_labels))
    return split


def epoch_batches(lengths: np.ndarray, batch_size: int, seed: int, epoch: int) -> list[np.ndarray]:
    """Length-bucketed batch order for one epoch; a pure function of (seed, epoch) for resume."""
    rng = np.random.default_rng([seed, epoch])
    perm = rng.permutation(len(lengths))
    mega = MEGA_BATCH_FACTOR * batch_size
    batches: list[np.ndarray] = []
    for start in range(0, len(perm), mega):
        chunk = perm[start:start + mega]
        chunk = chunk[np.argsort(-lengths[chunk], kind="stable")]
        batches.extend(chunk[j:j + batch_size] for j in range(0, len(chunk), batch_size))
    order = rng.permutation(len(batches))
    return [batches[k] for k in order]


def eval_batches(lengths: np.ndarray, batch_size: int) -> list[np.ndarray]:
    order = np.argsort(-lengths, kind="stable")
    return [order[i:i + batch_size] for i in range(0, len(order), batch_size)]


class RowIndexDataset(Dataset):
    """``torch.utils.data.Dataset`` over an ``EncodedSplit``: an item is a row index.

    Membership and order come from the batch sampler (``epoch_batches`` / ``eval_batches``:
    seeded shuffling + length bucketing), padding from the collate function, so every batch
    the trainer sees goes through ``torch.utils.data.DataLoader`` exactly as the course
    reference does, while the resume-exact batch order is kept.
    """

    def __init__(self, split: EncodedSplit) -> None:
        self.split = split

    def __len__(self) -> int:
        return len(self.split)

    def __getitem__(self, i: int) -> int:
        return int(i)


def make_loader(split: EncodedSplit, batches: list[np.ndarray], multiple: int,
                micro_batch: int | None = None) -> DataLoader:
    """DataLoader whose batch sampler is the precomputed batch list.

    Yields ``(idx, ids, mask, labels, allowed)`` per batch, or with ``micro_batch`` set
    ``(idx, [(ids, mask, labels, allowed), ...])`` where each micro-batch is padded to ITS
    OWN longest row (the gradient-accumulation geometry the smoke test measured). Tensors
    are produced on the CPU; the training loop moves them to the device. ``num_workers=0``
    keeps the order a pure function of (seed, epoch) and avoids forking next to a live
    CUDA context.
    """
    cpu = torch.device("cpu")

    def collate(rows: list[int]):
        idx = np.asarray(rows, dtype=np.int64)
        if micro_batch is None:
            return (idx, *split.collate(idx, multiple, cpu))
        return idx, [split.collate(idx[c:c + micro_batch], multiple, cpu) for c in range(0, len(idx), micro_batch)]

    # An explicit generator keeps the DataLoader from drawing its worker base seed from the
    # global torch RNG at iteration start, which would shift the dropout stream and break
    # bit-exact resume on the CPU dev path; the loader itself uses no randomness (num_workers=0,
    # explicit batch_sampler), so the seed value is irrelevant.
    return DataLoader(RowIndexDataset(split), batch_sampler=[[int(i) for i in b] for b in batches],
                      collate_fn=collate, num_workers=0, generator=torch.Generator().manual_seed(0))


def load_data(cfg: dict, args: TrainArgs, tok) -> tuple[EncodedSplit, EncodedSplit, int]:
    processed = Path(cfg["paths"]["data_processed"])
    holdout_source = str(cfg["data"]["holdout_test_source"])
    max_len = int(cfg["model"]["max_seq_len"])
    multiple = int(cfg["finetune"]["dynamic_padding_multiple"])

    data = read_split(processed / TRAIN_FILE, holdout_source, args.limit_rows, args.seed, args.num_labels)
    n_synth = 0
    if args.synthetic == "on":
        synth = processed / SYNTHETIC_FILE
        if synth.exists():
            # Synthetic rows carry their own label3/label_set (category "1:*" -> phishing, "0:*" -> legit).
            sd = read_split(synth, holdout_source, int(cfg["data"]["synthetic"].get("max_rows", 2000)),
                            args.seed, args.num_labels, force_source="synthetic")
            n_synth = len(sd)
            data = concat_labelled(data, sd)
            LOG.info("synthetic supplement: %d rows appended to TRAIN only; classes %s", n_synth,
                     label_counts(sd.label3, args.num_labels))
        else:
            LOG.warning("--synthetic on but %s is missing; training without synthetic rows", synth)
    train = encode_split("train", tok, data, max_len, multiple)
    vd = read_split(processed / VAL_FILE, holdout_source, args.limit_rows, args.seed, args.num_labels)
    if args.limit_rows is not None:
        LOG.warning("--limit-rows %d: validation is a seeded SUBSAMPLE too (dev only; metrics are not comparable)",
                    args.limit_rows)
    val = encode_split("val", tok, vd, max_len, multiple)
    assert "synthetic" not in set(val.sources), "synthetic rows must never reach validation"
    if args.num_labels == 3 and args.selection_metric == "val_macro_f1_3way" and not (val.label3 >= 0).any():
        raise ValueError("selection metric val_macro_f1_3way but finetune_val has no row with label3 in {0,1,2}")
    return train, val, n_synth


def check_collate_matches_encode_batch(split: EncodedSplit, tok, idx: np.ndarray, max_len: int, multiple: int) -> None:
    """The trained-on tensors must equal what ``encode_batch`` yields for the same texts."""
    ids, mask, _, _ = split.collate(idx, multiple, torch.device("cpu"))
    ref_ids, ref_mask = encode_batch(tok, [split.texts[int(i)] for i in idx], max_len, pad_to_multiple_of=multiple)
    if ids.shape != ref_ids.shape or not torch.equal(ids, ref_ids) or not torch.equal(mask, ref_mask):
        raise AssertionError("collated batch differs from tokenizer.encode_batch output")
    LOG.info("collate == encode_batch on the first batch (%s): OK", tuple(ids.shape))


# ---------------------------------------------------------------------------------- model
def newest_pretrain_final(cfg: dict) -> Path | None:
    root = Path(cfg["paths"]["checkpoints"]) / "pretrain"
    finals = [p for p in root.glob("*/final") if (p / "model.safetensors").exists()] if root.exists() else []
    return max(finals, key=lambda p: p.stat().st_mtime) if finals else None


def read_pretrained(path: Path) -> tuple[dict[str, Tensor], dict[str, Any]]:
    """state_dict + EncoderConfig dict from a Phase 1 final/ or step-NNNNNNN/ directory."""
    weights = path / "model.safetensors"
    if not weights.exists():
        raise FileNotFoundError(f"{weights} not found")
    cfg_json = path / "config.json"
    if cfg_json.exists():
        meta = common.read_json(cfg_json)
    else:
        with safe_open(str(weights), "pt") as fh:
            raw = (fh.metadata() or {}).get("config")
        if raw is None:
            raise RuntimeError(f"{path}: no config.json and no 'config' safetensors metadata")
        meta = json.loads(raw)
    model_cfg = meta["model"] if isinstance(meta, dict) and "model" in meta else meta
    return load_file(str(weights), device="cpu"), dict(model_cfg)


def slice_layers(state: dict[str, Tensor], n_layers: int) -> dict[str, Tensor]:
    """Keep encoder.layers.0..n_layers-1 (the instructor's 'number of layers' dial on a pretrained stack)."""
    out: dict[str, Tensor] = {}
    dropped = set()
    for k, v in state.items():
        if k.startswith("encoder.layers."):
            i = int(k.split(".")[2])
            if i >= n_layers:
                dropped.add(i)
                continue
        out[k] = v
    if dropped:
        LOG.info("--n-layers %d: dropped pretrained layers %d..%d", n_layers, min(dropped), max(dropped))
    return out


def build_model(cfg: dict, args: TrainArgs, device: torch.device) -> tuple[EncoderForClassification, str]:
    """Returns (model, init_descriptor) where init_descriptor is 'scratch' or 'pretrained:<path>'."""
    # The head size comes from finetune.num_labels, never from config model: (the pretraining snapshot).
    mcfg_dict = {**cfg["model"], "num_labels": args.num_labels}
    if args.init == "scratch":
        mcfg_dict["n_layers"] = args.n_layers
        ecfg = EncoderConfig.from_dict(mcfg_dict)
        common.set_seed(args.seed)
        with torch.device(device):
            model = EncoderForClassification(ecfg)
        return model, "scratch"

    path = Path(args.init.split(":", 1)[1]) if ":" in args.init else newest_pretrain_final(cfg)
    if path is None:
        raise FileNotFoundError("no checkpoints/pretrain/*/final found; pass --init pretrained:<dir> or --init scratch")
    if not path.is_absolute():
        path = common.project_root() / path
    state, pre_cfg = read_pretrained(path)
    if int(pre_cfg["n_layers"]) < args.n_layers:
        raise ValueError(f"--n-layers {args.n_layers} > pretrained n_layers {pre_cfg['n_layers']}")
    # Architecture from the checkpoint; regularisation from config.yaml; depth from the dial.
    merged = {**pre_cfg, "n_layers": args.n_layers, "num_labels": args.num_labels,
              "dropout": mcfg_dict["dropout"], "attention_dropout": mcfg_dict["attention_dropout"],
              "grad_checkpointing": mcfg_dict.get("grad_checkpointing", False), "tie_mlm_head": pre_cfg.get("tie_mlm_head", True)}
    ecfg = EncoderConfig.from_dict(merged)
    common.set_seed(args.seed)
    with torch.device(device):
        model = EncoderForClassification.from_pretrained_encoder(slice_layers(state, args.n_layers), ecfg)
    LOG.info("loaded pretrained encoder from %s (%d of %d layers)", path, args.n_layers, pre_cfg["n_layers"])
    return model, f"pretrained:{path}"


def freeze(model: EncoderForClassification, k: int) -> None:
    if k <= 0:
        return
    for p in model.encoder.embeddings.parameters():
        p.requires_grad_(False)
    for layer in list(model.encoder.layers)[:k]:
        for p in layer.parameters():
            p.requires_grad_(False)
    LOG.info("frozen: embeddings + encoder layers 0..%d", k - 1)


# ----------------------------------------------------------------------------------- loss
def class_weights(label3: np.ndarray, num_labels: int, device: torch.device) -> Tensor:
    """Inverse-frequency class weights ``w_c = N / (C * n_c)`` over the FULLY labelled rows
    (label3 >= 0) of the train split actually used. Unwanted rows have no class of their own;
    they receive mean(w_phish, w_spam) through ``row_weights`` at loss time."""
    full = label3[label3 >= 0]
    counts = np.bincount(full, minlength=num_labels).astype(np.float64)
    w = full.shape[0] / (float(num_labels) * np.maximum(counts, 1))
    return torch.tensor(w, dtype=torch.float32, device=device)


def partial_label_nll(logits: Tensor, allowed: Tensor) -> Tensor:
    """Per-row ``-log sum_{c in allowed} softmax(logits)_c`` = logsumexp(all) - logsumexp(allowed).

    One allowed class -> exactly ``F.cross_entropy(reduction="none")``; {phish, spam} ->
    ``-log(p1 + p2)``. Computed in fp32 from bf16-autocast logits.
    """
    logits = logits.float()
    if allowed.shape != logits.shape:
        raise ValueError(f"allowed mask {tuple(allowed.shape)} does not match logits {tuple(logits.shape)}")
    lse_all = torch.logsumexp(logits, dim=-1)
    lse_allowed = torch.logsumexp(logits.masked_fill(~allowed, float("-inf")), dim=-1)
    return lse_all - lse_allowed


def row_weights(allowed: Tensor, weights: Tensor) -> Tensor:
    """Mean of the class weights over each row's allowed classes: w[c] for a full label,
    mean(w1, w2) for an "unwanted" row."""
    a = allowed.to(weights.dtype)
    return (a @ weights) / a.sum(dim=-1)


def compute_loss(logits: Tensor, allowed: Tensor, kind: str, weights: Tensor | None, gamma: float) -> Tensor:
    """The three loss dials, all defined over the partial-label NLL (see module docstring).

    cross_entropy          mean(nll)
    weighted_cross_entropy sum(w_r nll_r) / sum(w_r)   (F.cross_entropy(weight=)'s normalisation)
    focal                  mean((1 - exp(-nll))^gamma * nll)   (= -(1 - p_t)^gamma log p_t; gamma 0 = plain)
    """
    nll = partial_label_nll(logits, allowed)
    if kind == "cross_entropy":
        return nll.mean()
    if kind == "weighted_cross_entropy":
        if weights is None:
            raise ValueError("weighted_cross_entropy needs class weights")
        w = row_weights(allowed, weights)
        return (w * nll).sum() / w.sum()
    if kind == "focal":
        return ((1.0 - torch.exp(-nll)).pow(gamma) * nll).mean()
    raise ValueError(f"unknown loss {kind!r}")


# -------------------------------------------------------------------------------- metrics
def classification_metrics(labels: np.ndarray, probs: np.ndarray, threshold: float = DECISION_THRESHOLD) -> dict[str, Any]:
    from sklearn.metrics import (accuracy_score, average_precision_score, confusion_matrix,
                                 precision_recall_curve, precision_recall_fscore_support, roc_auc_score)

    pred = (probs >= threshold).astype(np.int64)
    p, r, f1, _ = precision_recall_fscore_support(labels, pred, average="binary", pos_label=1, zero_division=0)
    cm = confusion_matrix(labels, pred, labels=[0, 1])
    try:
        roc = float(roc_auc_score(labels, probs))
        pr_auc = float(average_precision_score(labels, probs))
        pc, rc, th = precision_recall_curve(labels, probs)
        f1s = 2 * pc[:-1] * rc[:-1] / np.maximum(pc[:-1] + rc[:-1], 1e-12)
        j = int(np.argmax(f1s))
        opt_t, opt_f1 = float(th[j]), float(f1s[j])
    except ValueError:  # a single class present (tiny dev subsets)
        roc = pr_auc = opt_t = opt_f1 = float("nan")
    return {
        "accuracy": float(accuracy_score(labels, pred)), "precision": float(p), "recall": float(r), "f1": float(f1),
        "roc_auc": roc, "pr_auc": pr_auc,
        "tn": int(cm[0, 0]), "fp": int(cm[0, 1]), "fn": int(cm[1, 0]), "tp": int(cm[1, 1]),
        "threshold": threshold, "f1_optimal_threshold": opt_t, "f1_at_optimal_threshold": opt_f1,
    }


def three_way_metrics(label3: np.ndarray, probs: np.ndarray) -> dict[str, Any]:
    """Metric set (b) on the rows with a full three-way label (label3 in {0, 1, 2}).

    ``probs`` is ``[n, 3]``. Prediction = argmax over the three classes. Keys (flat, so they
    drop straight into epochs.csv): n_3way, accuracy_3way, macro_f1_3way (macro over ALL three
    classes, a class absent from the slice scores 0, so tiny dev subsets read low),
    {precision,recall,f1}_{legitimate,phishing,spam}, cm3_<true><pred> with l/p/s initials
    (cm3_ps = true phishing predicted spam), and the phishing-vs-spam confusion restricted to
    true unwanted rows where the prediction is the argmax over {phishing, spam} only
    ("given it is unwanted, which kind?"): n_unwanted_3way, pvs_phish_as_phish,
    pvs_phish_as_spam, pvs_spam_as_phish, pvs_spam_as_spam.
    Returns None/0 placeholders (same keys) when the slice is empty or the model is binary: None,
    not NaN, so best/config.json and experiments.jsonl of a binary run stay strict JSON.
    """
    from sklearn.metrics import confusion_matrix, precision_recall_fscore_support

    out: dict[str, Any] = {k: (0 if k.startswith(("n_", "cm3_", "pvs_")) else None) for k in THREE_WAY_KEYS}
    sel = label3 >= 0
    if probs.ndim != 2 or probs.shape[1] != 3 or not sel.any():
        return out
    y, p = label3[sel], probs[sel]
    pred = p.argmax(axis=1)
    cm = confusion_matrix(y, pred, labels=[0, 1, 2])
    prec, rec, f1, _ = precision_recall_fscore_support(y, pred, labels=[0, 1, 2], zero_division=0)
    out.update(n_3way=int(sel.sum()), accuracy_3way=float((pred == y).mean()), macro_f1_3way=float(f1.mean()))
    for c, slug in enumerate(CLASS_SLUGS):
        out[f"precision_{slug}"], out[f"recall_{slug}"], out[f"f1_{slug}"] = float(prec[c]), float(rec[c]), float(f1[c])
    for i in range(3):
        for j in range(3):
            out[f"cm3_{CLASS_SLUGS[i][0]}{CLASS_SLUGS[j][0]}"] = int(cm[i, j])
    unw = y >= 1
    if unw.any():
        kind = 1 + p[unw][:, 1:3].argmax(axis=1)
        cm2 = confusion_matrix(y[unw], kind, labels=[1, 2])
        out.update(n_unwanted_3way=int(unw.sum()), pvs_phish_as_phish=int(cm2[0, 0]), pvs_phish_as_spam=int(cm2[0, 1]),
                   pvs_spam_as_phish=int(cm2[1, 0]), pvs_spam_as_spam=int(cm2[1, 1]))
    return out


def metric_sets(res: dict[str, Any]) -> dict[str, Any]:
    """Structured view of a flat evaluation result for best/config.json and experiments.jsonl:
    ``binary_unwanted`` (metric set a) and ``three_way`` (metric set b, None for a binary model
    or an empty slice)."""
    binary = {k: res.get(k) for k in (*METRIC_KEYS, "tn", "fp", "fn", "tp", "threshold", "f1_optimal_threshold",
                                      "f1_at_optimal_threshold", "n")}
    binary["positive_class"] = UNWANTED_NAME
    three = None
    if res.get("n_3way"):
        three = {
            "n": res["n_3way"], "accuracy": res["accuracy_3way"], "macro_f1": res["macro_f1_3way"],
            "per_class": {LABEL_NAMES[c]: {m: res[f"{m}_{slug}"] for m in ("precision", "recall", "f1")}
                          for c, slug in enumerate(CLASS_SLUGS)},
            "confusion": {"order": [LABEL_NAMES[c] for c in range(3)], "rows_true_cols_pred":
                          [[res[f"cm3_{a[0]}{b[0]}"] for b in CLASS_SLUGS] for a in CLASS_SLUGS]},
            "phish_vs_spam_on_true_unwanted": {
                "n": res["n_unwanted_3way"], "order": [LABEL_NAMES[1], LABEL_NAMES[2]],
                "rows_true_cols_pred": [[res["pvs_phish_as_phish"], res["pvs_phish_as_spam"]],
                                        [res["pvs_spam_as_phish"], res["pvs_spam_as_spam"]]]},
        }
    return {"binary_unwanted": binary, "three_way": three}


def autocast_ctx(device: torch.device):
    """bf16 autocast on CUDA; a no-op on the CPU dev path (fp32 end to end)."""
    if device.type == "cuda":
        return torch.autocast(device_type="cuda", dtype=torch.bfloat16)
    return contextlib.nullcontext()


def unwanted_prob(probs: Tensor) -> Tensor:
    """P(unwanted) = p_phish + p_spam (3 classes) or p_1 (binary)."""
    return probs[:, 1:].sum(dim=-1)


@torch.inference_mode()
def evaluate(model: EncoderForClassification, split: EncodedSplit, args: TrainArgs, multiple: int,
             device: torch.device) -> dict[str, Any]:
    """Both metric sets on ``split``. ``val_loss`` is the unweighted partial-label NLL (the
    plain cross_entropy dial) so it is comparable across loss variants."""
    model.eval()
    t0 = common.timer()
    probs, labels, label3s, loss_sum = [], [], [], 0.0
    loader = make_loader(split, eval_batches(split.lengths, args.eval_batch_size), multiple)
    for idx, ids, mask, y, allowed in loader:
        ids, mask, y, allowed = (t.to(device, non_blocking=True) for t in (ids, mask, y, allowed))
        with autocast_ctx(device):
            logits = model(ids, mask)
        logits = logits.float()
        loss_sum += float(partial_label_nll(logits, allowed).sum())
        probs.append(logits.softmax(-1).cpu().numpy())
        labels.append(y.cpu().numpy())
        # eval_batches() visits rows in length order, so EVERY target must be gathered with the same idx
        # as the predictions; split.label3 in storage order would score a permutation of the model.
        label3s.append(split.label3[idx])
    model.train()
    probs_np, labels_np, label3_np = np.concatenate(probs), np.concatenate(labels), np.concatenate(label3s)
    out = classification_metrics(labels_np, unwanted_prob(torch.from_numpy(probs_np)).numpy())
    out["f1_unwanted"] = out["f1"]
    out.update(three_way_metrics(label3_np, probs_np))
    out["val_loss"] = loss_sum / max(1, len(labels_np))
    out["n"] = int(len(labels_np))
    out["eval_seconds"] = round(common.timer() - t0, 1)
    return out


def selection_value(res: dict[str, Any], metric: str) -> float:
    v = res.get(SELECTION_KEY[metric])
    return float("-inf") if v is None or (isinstance(v, float) and math.isnan(v)) else float(v)


# ------------------------------------------------------------------------------- run state
@dataclass
class RunState:
    step: int = 0
    epoch: int = 0
    batch_in_epoch: int = 0
    best_score: float = -1.0        # value of the selection metric at best/ (named best_f1 before entry 6)
    best_step: int = 0
    best_epoch: int = 0
    evals_without_improvement: int = 0
    train_seconds: float = 0.0
    eval_seconds: float = 0.0
    peak_reserved_gb: float = 0.0
    peak_allocated_gb: float = 0.0
    best_val: dict[str, Any] | None = None
    last_val: dict[str, Any] | None = None
    n_train: int = 0
    n_val: int = 0
    n_synthetic: int = 0
    init: str = ""
    stopped_by: str = ""
    resumed_from_step: int | None = None
    metrics_tail: list[dict[str, Any]] = field(default_factory=list)


class StopFlag:
    """SIGINT/SIGTERM -> finish the current micro-step, save last/, exit 0. Second signal -> default."""

    def __init__(self) -> None:
        self.requested: str | None = None
        self._old = {s: signal.signal(s, self._handle) for s in (signal.SIGINT, signal.SIGTERM)}

    def _handle(self, signum: int, frame: Any) -> None:
        if self.requested:
            signal.signal(signum, self._old[signum])
            raise KeyboardInterrupt
        self.requested = signal.Signals(signum).name
        LOG.warning("%s received: finishing the current step, saving last/ and exiting", self.requested)


def rng_states(device: torch.device) -> dict[str, Any]:
    """RNG snapshot for resume. The CUDA generator is read only when training ON CUDA:
    ``torch.cuda.get_rng_state_all`` goes through ``_lazy_init`` and would create a CUDA context
    on a GPU the --cpu dev path must leave alone (the production pretraining run lives there)."""
    return {"python": random.getstate(), "numpy": np.random.get_state(), "torch_cpu": torch.get_rng_state(),
            "torch_cuda": torch.cuda.get_rng_state_all() if device.type == "cuda" else []}


def restore_rng(states: dict[str, Any], device: torch.device) -> None:
    random.setstate(states["python"])
    np.random.set_state(states["numpy"])
    torch.set_rng_state(states["torch_cpu"])
    if device.type == "cuda" and states.get("torch_cuda"):
        torch.cuda.set_rng_state_all(states["torch_cuda"])


def json_clean(obj: Any) -> Any:
    """Recursively replace non-finite floats (NaN/inf from a degenerate metric) with None so every
    JSON artifact (best/config.json, last/config.json, experiments.jsonl) is strict JSON; the stdlib
    encoder would otherwise emit the bare token ``NaN`` that non-Python parsers reject."""
    if isinstance(obj, float):
        return obj if math.isfinite(obj) else None
    if isinstance(obj, np.floating):
        return float(obj) if np.isfinite(obj) else None
    if isinstance(obj, dict):
        return {k: json_clean(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [json_clean(v) for v in obj]
    return obj


def save_model_dir(model: EncoderForClassification, out: Path, step: int, extra: dict[str, Any]) -> None:
    out.mkdir(parents=True, exist_ok=True)
    sd = {k: v.detach().to("cpu").contiguous() for k, v in model.state_dict().items()}
    cfg_dict = model.config.to_dict()
    save_file(sd, str(out / "model.safetensors"),
              metadata={"format": "pt", "step": str(step), "config": json.dumps(cfg_dict)})
    common.write_json(json_clean({"model": cfg_dict, "tokenizer": "tokenizer.json", **extra}), out / "config.json")


def save_last(run_dir: Path, model: EncoderForClassification, opt: torch.optim.Optimizer,
              sched: torch.optim.lr_scheduler.LRScheduler, state: RunState, args: TrainArgs,
              schedule: dict[str, int], device: torch.device) -> None:
    t0 = common.timer()
    last = run_dir / "last"
    save_model_dir(model, last, state.step, {"finetune": args.to_record(), "state": _state_record(state)})
    torch.save({
        "optimizer": opt.state_dict(), "scheduler": sched.state_dict(), "state": asdict(state),
        "args": args.to_record(), "schedule": schedule, "rng": rng_states(device),
    }, last / "trainer_state.pt")
    LOG.info("saved last/ at step %d (%.1fs)", state.step, common.timer() - t0)


def _state_record(state: RunState) -> dict[str, Any]:
    d = asdict(state)
    d.pop("metrics_tail", None)
    return d


def write_csv_row(path: Path, row: dict[str, Any], header: list[str]) -> None:
    new = not path.exists()
    with path.open("a", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=header, extrasaction="ignore")
        if new:
            w.writeheader()
        w.writerow({k: ("" if row.get(k) is None else row.get(k)) for k in header})


def trim_csv(path: Path, max_step: int) -> None:
    """Drop rows with step > max_step: a crash between checkpoints leaves rows past the last/ being
    resumed, and the resumed steps would otherwise be appended a second time."""
    if not path.exists():
        return
    with path.open(newline="", encoding="utf-8") as fh:
        reader = csv.DictReader(fh)
        header, rows = list(reader.fieldnames or []), list(reader)
    keep = [r for r in rows if r.get("step") and int(float(r["step"])) <= max_step]
    if len(keep) == len(rows):
        return
    tmp = path.with_suffix(path.suffix + ".tmp")
    with tmp.open("w", newline="", encoding="utf-8") as fh:
        w = csv.DictWriter(fh, fieldnames=header)
        w.writeheader()
        w.writerows(keep)
    os.replace(tmp, path)
    LOG.info("%s: dropped %d row(s) logged after step %d (crash between checkpoints)", path.name, len(rows) - len(keep), max_step)


METRICS_HEADER = ["step", "epoch", "loss", "lr", "examples_per_s", "max_alloc_gb", "max_reserved_gb", "elapsed_s",
                  "eta_s", "val_loss", "val_f1", "val_roc_auc", "val_macro_f1_3way"]
# Binary roll-up columns first (unchanged names, so old epochs.csv readers keep working), then the three-way set.
EPOCHS_HEADER = ["epoch", "step", "is_epoch_end", "n", "val_loss", "accuracy", "precision", "recall", "f1", "roc_auc",
                 "pr_auc", "tn", "fp", "fn", "tp", "threshold", "f1_optimal_threshold", "f1_at_optimal_threshold",
                 "f1_unwanted", *THREE_WAY_KEYS, "eval_seconds", "is_best"]


# ----------------------------------------------------------------------------- experiments
def experiment_record(args: TrainArgs, state: RunState, params: dict[str, Any], steps_per_epoch: int) -> dict[str, Any]:
    val = state.best_val or {}
    return {
        "run_name": args.run_name,
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "init": state.init,
        "dials": {
            "n_layers": args.n_layers, "frozen_layers": args.freeze_layers,
            "epochs_planned": args.epochs, "epochs_run": round(state.step / max(1, steps_per_epoch), 3),
            "loss": args.loss, "focal_gamma": args.focal_gamma if args.loss == "focal" else None,
            "lr": args.lr, "batch_size": args.batch_size, "micro_batch": args.micro_batch,
            "class_weights": "inverse_frequency" if args.loss == "weighted_cross_entropy" else None,
            "synthetic": args.synthetic, "synthetic_rows": state.n_synthetic,
            "warmup_fraction": args.warmup_fraction, "max_steps": args.max_steps, "limit_rows": args.limit_rows,
            "seed": args.seed, "schedule": args.schedule, "selection_metric": args.selection_metric,
            "pooling": args.pooling, "num_labels": args.num_labels, "device": "cpu" if args.cpu else "cuda",
        },
        "params": params,
        "steps_run": state.step,
        "n_train": state.n_train, "n_val": state.n_val,
        "train_seconds": round(state.train_seconds, 1),
        "peak_vram_gb": round(state.peak_reserved_gb, 3),
        "best_epoch": state.best_epoch, "best_step": state.best_step,
        "best_score": state.best_score if val else None,
        # Old binary keys kept verbatim (finetune.py list / evaluate.py read them); the three-way set is added.
        "val": ({**{k: val.get(k) for k in METRIC_KEYS}, "f1_unwanted": val.get("f1"),
                 "macro_f1_3way": val.get("macro_f1_3way"),
                 "per_class": {LABEL_NAMES[c]: {m: val.get(f"{m}_{slug}") for m in ("precision", "recall", "f1")}
                               for c, slug in enumerate(CLASS_SLUGS)} if val.get("n_3way") else None,
                 "three_way": metric_sets(val)["three_way"]} if val else None),
        "val_threshold": DECISION_THRESHOLD,
        "label_names": [LABEL_NAMES[c] for c in range(args.num_labels)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME],
        "stopped_by": state.stopped_by,
        "resumed_from_step": state.resumed_from_step,
        "test": None,
    }


def upsert_experiment(path: Path, rec: dict[str, Any]) -> str:
    """ONE train line per run_name: a resumed invocation replaces its earlier train line instead
    of appending a duplicate; evaluate.py's lines (test != null) are never touched."""
    if rec.get("test") is not None:
        raise ValueError("upsert_experiment writes train lines only (test must be None); a test record would "
                         "overwrite the run's train line")
    lines = path.read_text(encoding="utf-8").splitlines() if path.exists() else []
    out, replaced = [], False
    rec = json_clean(rec)   # strict JSON: no bare NaN tokens for binary runs' three-way placeholders
    for line in lines:
        if not line.strip():
            continue
        try:
            old = json.loads(line)
        except json.JSONDecodeError:
            out.append(line)
            continue
        if not replaced and old.get("run_name") == rec["run_name"] and old.get("test") is None:
            out.append(json.dumps(rec, allow_nan=False))
            replaced = True
        else:
            out.append(line)
    if not replaced:
        out.append(json.dumps(rec, allow_nan=False))
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(".jsonl.tmp")
    tmp.write_text("\n".join(out) + "\n", encoding="utf-8")
    os.replace(tmp, path)
    return "replaced" if replaced else "appended"


def run_log_block(args: TrainArgs, state: RunState, run_dir: Path, params: dict[str, Any]) -> str:
    v = state.best_val or {}
    nan = float("nan")
    three = ""
    if v.get("n_3way"):
        three = (f" Three-way on {v['n_3way']:,} fully-labelled val rows: macro-F1 {v['macro_f1_3way']:.4f}, accuracy "
                 f"{v['accuracy_3way']:.4f}; per-class F1 " +
                 ", ".join(f"{LABEL_NAMES[c]} {v[f'f1_{s}']:.4f}" for c, s in enumerate(CLASS_SLUGS)) +
                 f"; phishing-vs-spam on true unwanted (n={v['n_unwanted_3way']:,}): phish->phish {v['pvs_phish_as_phish']}, "
                 f"phish->spam {v['pvs_phish_as_spam']}, spam->phish {v['pvs_spam_as_phish']}, spam->spam {v['pvs_spam_as_spam']}.")
    return "\n".join([
        f"## Fine-tune run `{args.run_name}` — {time.strftime('%Y-%m-%d %H:%M')}",
        f"**Dials.** init {state.init}; {args.num_labels} classes; {args.n_layers} layers, {args.freeze_layers} frozen; loss {args.loss}"
        f"{f' (gamma {args.focal_gamma})' if args.loss == 'focal' else ''}; lr {args.lr:g}; batch {args.batch_size} "
        f"(micro {args.micro_batch}); epochs {args.epochs} planned; synthetic {args.synthetic} ({state.n_synthetic} rows); "
        f"selection {args.selection_metric}; train rows {state.n_train:,}, val rows {state.n_val:,}; params {params['total']:,} "
        f"({params['trainable']:,} trainable).",
        f"**Result.** {state.step} steps, stopped by {state.stopped_by}; training time {state.train_seconds / 60:.1f} min "
        f"(eval {state.eval_seconds / 60:.1f} min of it); peak reserved {state.peak_reserved_gb:.2f} GiB. Best val at step "
        f"{state.best_step} (epoch {state.best_epoch}), {UNWANTED_NAME} roll-up: F1 {v.get('f1', nan):.4f}, precision "
        f"{v.get('precision', nan):.4f}, recall {v.get('recall', nan):.4f}, accuracy {v.get('accuracy', nan):.4f}, ROC-AUC "
        f"{v.get('roc_auc', nan):.4f}, PR-AUC {v.get('pr_auc', nan):.4f}; F1-optimal threshold {v.get('f1_optimal_threshold', nan):.3f}."
        f"{three} Checkpoints: {run_dir}/best, {run_dir}/last.",
    ])


# ------------------------------------------------------------------------------- training
CPU_VRAM = {"allocated_gb": 0.0, "reserved_gb": 0.0, "max_allocated_gb": 0.0, "max_reserved_gb": 0.0}


def backfill_legacy_args(saved: dict[str, Any], run_dir: Path) -> None:
    """A trainer_state.pt written before RUN_LOG entry 6 has no ``num_labels`` / ``selection_metric``
    in its args. Fill them from what that run actually was (the head size recorded in
    last/config.json model.num_labels; the only selection metric that existed, val_f1_phishing) so the
    saved-overrides-config rule resolves the run as the binary run it is, instead of the current config's
    3-class head crashing in load_state_dict with a size mismatch."""
    args = saved.get("args") or {}
    if "num_labels" not in args:
        try:
            n = int(common.read_json(run_dir / "last" / "config.json")["model"]["num_labels"])
        except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError):
            n = 2
        args["num_labels"] = n
        LOG.info("--resume: legacy trainer_state.pt without num_labels; using %d from last/config.json", n)
    if "selection_metric" not in args:
        args["selection_metric"] = "val_f1_phishing"
    saved["args"] = args


def lr_lambda_factory(total_steps: int, warmup_steps: int):
    def lr_lambda(completed: int) -> float:
        if completed < warmup_steps:
            return (completed + 1) / warmup_steps
        return max(0.0, (total_steps - completed) / max(1, total_steps - warmup_steps))
    return lr_lambda


def train(ns: argparse.Namespace, cfg: dict) -> int:
    t_invocation = common.timer()
    # Decide the device FIRST and touch torch.cuda only on the CUDA branch. Under --cpu the GPU may be
    # (and today is) occupied by the production pretraining run: torch.cuda.current_device() /
    # mem_get_info() / get_rng_state_all() all run _lazy_init(), which creates a CUDA context there.
    # Hiding the devices from this process is the belt to that braces: nothing has called CUDA yet.
    if ns.cpu:
        # Hide the devices before the CUDA runtime is first queried (torch import does not query it): then
        # torch.cuda.is_available() is False and torch's own AdamW.step() health check, which otherwise calls
        # torch.accelerator.current_stream() -> _lazy_init() on any visible GPU, is skipped. If a GPU is still
        # visible the contract cannot be honoured from inside this process, so refuse rather than risk it.
        os.environ["CUDA_VISIBLE_DEVICES"] = ""
        if torch.cuda.is_initialized() or torch.cuda.is_available():
            LOG.error("--cpu: a CUDA device is still visible to this process (CUDA was queried before the devices could "
                      "be hidden); run with CUDA_VISIBLE_DEVICES=\"\" in the environment so the occupied GPU is never touched")
            return 2
        device = torch.device("cpu")
        device_info = {"device": "cpu", "torch": torch.__version__}
        LOG.warning("--cpu: dev path on the CPU with the config as given (no autocast, no VRAM guard, CUDA never "
                    "initialised; CUDA_VISIBLE_DEVICES cleared for this process); not for real runs")
    elif not torch.cuda.is_available():
        LOG.error("CUDA is required for fine-tuning (--cpu is the dev-only escape hatch)")
        return 2
    else:
        device = torch.device("cuda")
        device_info = common.print_device_info()
        LOG.info("PYTORCH_CUDA_ALLOC_CONF=%s", os.environ.get("PYTORCH_CUDA_ALLOC_CONF"))

    run_dir = Path(cfg["paths"]["checkpoints"]) / "finetune" / ns.run_name
    last_state_path = run_dir / "last" / "trainer_state.pt"
    saved: dict[str, Any] | None = None
    if ns.resume:
        if not last_state_path.exists():
            LOG.error("--resume: %s not found", last_state_path)
            return 2
        # Own file (optimizer tensors + Python/NumPy RNG tuples), hence weights_only=False.
        saved = torch.load(last_state_path, map_location="cpu", weights_only=False)
        backfill_legacy_args(saved, run_dir)
    elif run_dir.exists() and any(run_dir.iterdir()):
        if not ns.overwrite:
            LOG.error("%s exists; use --resume to continue it, --overwrite to discard it, or a new --run-name", run_dir)
            return 2
        import shutil
        shutil.rmtree(run_dir)
        LOG.warning("--overwrite: removed %s", run_dir)

    try:
        args = TrainArgs.resolve(ns, cfg, saved["args"] if saved else None)
    except ValueError as e:
        LOG.error("dial error: %s", e)
        return 2
    args.run_name = ns.run_name
    args.cpu = bool(ns.cpu)   # the device is this invocation's choice, never restored from last/
    if saved:
        changed = {k: (saved["args"].get(k), getattr(ns, k)) for k in RESUME_FIXED_DIALS
                   if getattr(ns, k, None) is not None and getattr(ns, k) != saved["args"].get(k)}
        if changed:
            LOG.error("--resume: %s cannot change on resume (saved -> requested); the restored optimizer state and "
                      "batch order would not match. Start a new --run-name instead", changed)
            return 2
    LOG.info("resolved dials: %s", args.to_record())
    common.set_seed(args.seed)
    vr = cfg["vram"]
    max_len = int(cfg["model"]["max_seq_len"])
    multiple = int(cfg["finetune"]["dynamic_padding_multiple"])

    tok = load_tokenizer(cfg["paths"]["tokenizer"])
    train_split, val_split, n_synth = load_data(cfg, args, tok)
    check_collate_matches_encode_batch(train_split, tok, epoch_batches(train_split.lengths, args.batch_size, args.seed, 0)[0],
                                       max_len, multiple)
    if saved and (saved["state"]["n_train"] != len(train_split) or saved["state"]["n_val"] != len(val_split)):
        LOG.error("--resume: data changed (train %d -> %d, val %d -> %d); the batch order would not line up",
                  saved["state"]["n_train"], len(train_split), saved["state"]["n_val"], len(val_split))
        return 2

    try:
        model, init_desc = build_model(cfg, args, device)
    except (FileNotFoundError, ValueError, RuntimeError) as e:
        LOG.error("model init failed: %s", e)
        return 2
    if saved:
        sd = load_file(str(run_dir / "last" / "model.safetensors"), device="cpu")
        try:
            model.load_state_dict(sd, strict=True)
        except RuntimeError as e:   # head / depth mismatch that slipped past the dial checks
            LOG.error("--resume: last/model.safetensors does not fit the resolved dials (num_labels %d, n_layers %d): %s",
                      args.num_labels, args.n_layers, str(e).splitlines()[0])
            return 2
    freeze(model, args.freeze_layers)
    model.train()
    params = count_parameters(model)
    LOG.info("EncoderForClassification params: %s", params)
    if tok.get_vocab_size() > model.config.vocab_size:
        raise AssertionError("tokenizer vocab larger than model vocab")

    steps_per_epoch = math.ceil(len(train_split) / args.batch_size)
    total_steps = steps_per_epoch * args.epochs
    if args.max_steps is not None:
        total_steps = min(total_steps, args.max_steps)
    warmup_steps = max(1, math.ceil(args.warmup_fraction * total_steps))
    schedule = {"steps_per_epoch": steps_per_epoch, "total_steps": total_steps, "warmup_steps": warmup_steps}
    LOG.info("schedule: %d steps/epoch x %d epochs -> %d optimizer steps (max_steps=%s), warmup %d, "
             "micro-batch %d x %d accumulation", steps_per_epoch, args.epochs, total_steps, args.max_steps, warmup_steps,
             args.micro_batch, args.batch_size // args.micro_batch)

    opt_cfg = {"weight_decay": args.weight_decay, "betas": list(args.betas), "eps": args.eps}
    opt = build_optimizer(model, args.lr, opt_cfg)
    lr_lambda = lr_lambda_factory(total_steps, warmup_steps)
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda)
    weights = class_weights(train_split.label3, args.num_labels, device) if args.loss == "weighted_cross_entropy" else None
    if weights is not None:
        names = [LABEL_NAMES[c] for c in range(3)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME]
        LOG.info("inverse-frequency class weights %s from fully-labelled counts %s%s", dict(zip(names, [round(float(w), 4) for w in weights])),
                 np.bincount(train_split.label3[train_split.label3 >= 0], minlength=args.num_labels).tolist(),
                 f"; {int((train_split.label3 < 0).sum())} unwanted rows weighted mean(w_phish, w_spam)={float(weights[1:].mean()):.4f}"
                 if args.num_labels == 3 else "")
    if args.num_labels == 3:
        LOG.info("label contract: classes %s; binary roll-up '%s' = P(%s) + P(%s); selection metric %s",
                 LABEL_NAMES, UNWANTED_NAME, LABEL_NAMES[1], LABEL_NAMES[2], args.selection_metric)

    state = RunState(n_train=len(train_split), n_val=len(val_split), n_synthetic=n_synth, init=init_desc)
    if saved:
        st = dict(saved["state"])
        if "best_f1" in st:   # trainer_state.pt written before entry 6
            st["best_score"] = st.pop("best_f1")
        known = {f.name for f in fields(RunState)}
        state = RunState(**{**{k: v for k, v in st.items() if k in known}, "stopped_by": "", "init": st["init"]})
        state.resumed_from_step = state.step
        opt.load_state_dict(saved["optimizer"])
        sched.load_state_dict(saved["scheduler"])
        if args.lr != float(saved["args"]["lr"]):
            # load_state_dict restored the saved base lr; re-apply the override so the next optimizer step
            # (which runs before the next sched.step()) and every later one use the lr the record claims.
            scale = lr_lambda(sched.last_epoch)
            for g in opt.param_groups:
                g["initial_lr"] = args.lr
                g["lr"] = args.lr * scale
            sched.base_lrs = [args.lr] * len(opt.param_groups)
            LOG.warning("--lr %g overrides the saved %g from step %d onward (steps 1..%d used the old lr); current lr %.3e",
                        args.lr, float(saved["args"]["lr"]), state.step, state.step, args.lr * scale)
        if saved["schedule"] != schedule:
            LOG.warning("schedule changed on resume %s -> %s; the LR lambda now follows the new totals from step %d",
                        saved["schedule"], schedule, state.step)
        restore_rng(saved["rng"], device)
        LOG.info("resumed %s at step %d (epoch %d, batch %d/%d); best %s %.4f at step %d; %.1f s trained so far",
                 ns.run_name, state.step, state.epoch, state.batch_in_epoch, steps_per_epoch, args.selection_metric,
                 state.best_score, state.best_step, state.train_seconds)
        if args.max_steps is not None and state.step >= args.max_steps:
            LOG.warning("nothing to resume: step %d already reached --max-steps %d; raise --max-steps or --epochs to continue",
                        state.step, args.max_steps)
            return 0
        if state.epoch >= args.epochs:
            LOG.warning("nothing to resume: epoch %d already reached --epochs %d; raise --epochs to continue",
                        state.epoch, args.epochs)
            return 0
        if state.evals_without_improvement >= args.patience:
            LOG.warning("resuming after early stopping (%d/%d evaluations without improvement): the run stops again at the "
                        "next evaluation unless %s improves", state.evals_without_improvement, args.patience, args.selection_metric)

    run_dir.mkdir(parents=True, exist_ok=True)
    metrics_csv, epochs_csv = run_dir / "metrics.csv", run_dir / "epochs.csv"
    if saved:
        trim_csv(metrics_csv, state.step)
        trim_csv(epochs_csv, state.step)
    stop = StopFlag()
    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats()
    band_warned = False
    label_names = [LABEL_NAMES[c] for c in range(3)] if args.num_labels == 3 else [LABEL_NAMES[0], UNWANTED_NAME]

    def record_eval(is_epoch_end: bool) -> None:
        nonlocal state
        res = evaluate(model, val_split, args, multiple, device)
        state.eval_seconds += res["eval_seconds"]
        state.last_val = res
        score = selection_value(res, args.selection_metric)
        improved = score > state.best_score
        if improved:
            state.best_score, state.best_step, state.best_epoch, state.best_val = score, state.step, state.epoch, res
            state.evals_without_improvement = 0
            save_model_dir(model, run_dir / "best", state.step, {
                "finetune": args.to_record(), "best_epoch": state.epoch, "best_step": state.step,
                "num_labels": args.num_labels, "label_names": label_names, "binary_rollup_name": UNWANTED_NAME,
                "selection_metric": args.selection_metric, "selection_value": score,
                "val_metrics": res, "val_metric_sets": metric_sets(res),
                "thresholds": {"decision": DECISION_THRESHOLD, "decision_applies_to": "P(unwanted) = sum of non-legitimate probabilities",
                               "val_f1_optimal": res["f1_optimal_threshold"], "val_f1_at_optimal": res["f1_at_optimal_threshold"]},
                "init": state.init, "params": params,
            })
        else:
            state.evals_without_improvement += 1
        three = ""
        if res.get("n_3way"):
            three = (f" | 3-way n={res['n_3way']} macro-F1 {res['macro_f1_3way']:.4f} acc {res['accuracy_3way']:.4f} F1[L P S]=[%s] "
                     f"cm3[rows true L,P,S; cols pred]=[%s] phish/spam on unwanted [pp ps sp ss]=[%d %d %d %d]" % (
                         " ".join(f"{res[f'f1_{s}']:.3f}" for s in CLASS_SLUGS),
                         " ".join(f"{res[k]}" for k in CM3_KEYS), *(res[k] for k in PVS_KEYS)))
        LOG.info("eval step %d epoch %d%s: n=%d loss %.4f | unwanted roll-up acc %.4f P %.4f R %.4f F1 %.4f ROC-AUC %.4f "
                 "PR-AUC %.4f cm[tn fp fn tp]=[%d %d %d %d] F1-opt thr %.3f (F1 %.4f)%s | %s=%.4f %s (%.1fs)",
                 state.step, state.epoch, " [epoch end]" if is_epoch_end else "", res["n"], res["val_loss"], res["accuracy"],
                 res["precision"], res["recall"], res["f1"], res["roc_auc"], res["pr_auc"], res["tn"], res["fp"], res["fn"],
                 res["tp"], res["f1_optimal_threshold"], res["f1_at_optimal_threshold"], three, args.selection_metric, score,
                 "NEW BEST -> best/" if improved else f"no improvement ({state.evals_without_improvement}/{args.patience})",
                 res["eval_seconds"])
        write_csv_row(epochs_csv, {"epoch": state.epoch, "step": state.step, "is_epoch_end": int(is_epoch_end),
                                   "is_best": int(improved), **res}, EPOCHS_HEADER)

    t_loop = t_inv_start = common.timer()
    steps_this_invocation = 0
    overhead_inv = 0.0   # eval + checkpoint seconds this invocation, excluded from the ETA's per-step cost
    window_loss, window_rows, window_steps, t_window = 0.0, 0, 0, common.timer()
    last_eval_step = state.step if saved else -1
    exit_code = 0

    def flush_time() -> None:
        nonlocal t_loop
        now = common.timer()
        state.train_seconds += now - t_loop
        t_loop = now

    try:
        while state.epoch < args.epochs and not state.stopped_by:
            batches = epoch_batches(train_split.lengths, args.batch_size, args.seed, state.epoch)
            loader = make_loader(train_split, batches[state.batch_in_epoch:], multiple, args.micro_batch)
            for b, (idx, micro_batches) in enumerate(loader, start=state.batch_in_epoch):
                n_rows = len(idx)
                loss_acc = 0.0
                for ids, mask, _y_bin, allowed in micro_batches:
                    ids, mask, allowed = (t.to(device, non_blocking=True) for t in (ids, mask, allowed))
                    n_sub = int(ids.shape[0])
                    with autocast_ctx(device):
                        logits = model(ids, mask)
                    loss = compute_loss(logits, allowed, args.loss, weights, args.focal_gamma)
                    (loss * (n_sub / n_rows)).backward()
                    loss_acc += float(loss.detach()) * n_sub / n_rows
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
                opt.step()
                opt.zero_grad(set_to_none=True)
                sched.step()
                state.step += 1
                state.batch_in_epoch = b + 1
                steps_this_invocation += 1
                window_loss += loss_acc
                window_rows += n_rows
                window_steps += 1
                if not math.isfinite(loss_acc):
                    LOG.error("non-finite loss at step %d; stopping", state.step)
                    state.stopped_by = "non_finite_loss"
                    exit_code = 3

                # Read the allocator only on CUDA; the CPU dev path reports zeros without touching torch.cuda.
                vram = common.vram_stats() if device.type == "cuda" else CPU_VRAM
                state.peak_reserved_gb = max(state.peak_reserved_gb, vram["max_reserved_gb"])
                state.peak_allocated_gb = max(state.peak_allocated_gb, vram["max_allocated_gb"])
                if device.type == "cuda" and vram["max_reserved_gb"] > float(vr["hard_cap_gb"]):
                    LOG.error("VRAM hard cap breached: %.2f GiB reserved > %s GiB; lower --micro-batch. Saving last/ and aborting",
                              vram["max_reserved_gb"], vr["hard_cap_gb"])
                    state.stopped_by = "vram_cap"
                    exit_code = 3
                elif not band_warned and vram["max_reserved_gb"] > float(vr["target_gb_high"]):
                    LOG.warning("reserved peak %.2f GiB is above the target band [%s, %s] GiB", vram["max_reserved_gb"],
                                vr["target_gb_low"], vr["target_gb_high"])
                    band_warned = True

                if stop.requested:
                    state.stopped_by = f"signal:{stop.requested}"
                is_eval = args.eval_every > 0 and state.step % args.eval_every == 0
                epoch_end = b == len(batches) - 1
                hit_max = args.max_steps is not None and state.step >= args.max_steps
                if hit_max and not state.stopped_by:
                    state.stopped_by = "max_steps"
                is_log = state.step % args.log_every == 0 or is_eval or epoch_end or bool(state.stopped_by)
                row: dict[str, Any] | None = None
                if is_log:
                    now = common.timer()
                    elapsed = state.train_seconds + (now - t_loop)
                    per_step = (now - t_inv_start - overhead_inv) / max(1, steps_this_invocation)
                    eta = per_step * max(0, total_steps - state.step)
                    row = {"step": state.step, "epoch": state.epoch, "loss": round(window_loss / max(1, window_steps), 5),
                           "lr": sched.get_last_lr()[0], "examples_per_s": round(window_rows / max(1e-9, now - t_window), 1),
                           "max_alloc_gb": round(vram["max_allocated_gb"], 3), "max_reserved_gb": round(vram["max_reserved_gb"], 3),
                           "elapsed_s": round(elapsed, 1), "eta_s": round(eta)}
                    LOG.info("step %d/%d epoch %d loss %.4f lr %.2e %.1f ex/s alloc %.2f GiB reserved %.2f GiB elapsed %s ETA %s",
                             state.step, total_steps, state.epoch, row["loss"], row["lr"], row["examples_per_s"],
                             row["max_alloc_gb"], row["max_reserved_gb"], _hms(elapsed), _hms(eta))
                    window_loss, window_rows, window_steps, t_window = 0.0, 0, 0, now

                # Evaluate at the cadence, at every epoch end and when --max-steps ends the run; never
                # after a signal (fast exit), a VRAM breach or a NaN model (it must not become best/).
                do_eval = (is_eval or epoch_end or hit_max) and state.step != last_eval_step \
                    and state.stopped_by in ("", "max_steps")
                if do_eval:
                    t_ev = common.timer()
                    record_eval(epoch_end)
                    overhead_inv += common.timer() - t_ev
                    last_eval_step = state.step
                    if row is not None:
                        m3 = state.last_val.get("macro_f1_3way")
                        row.update(val_loss=round(state.last_val["val_loss"], 5), val_f1=round(state.last_val["f1"], 5),
                                   val_roc_auc=round(state.last_val["roc_auc"], 5),
                                   val_macro_f1_3way=None if m3 is None or math.isnan(m3) else round(m3, 5))
                    if state.evals_without_improvement >= args.patience and not state.stopped_by:
                        state.stopped_by = "early_stopping"
                        LOG.info("early stopping: %d evaluations without %s improvement (best %.4f at step %d)",
                                 state.evals_without_improvement, args.selection_metric, state.best_score, state.best_step)
                if row is not None:
                    write_csv_row(metrics_csv, row, METRICS_HEADER)
                    state.metrics_tail = (state.metrics_tail + [row])[-20:]
                if state.stopped_by:
                    break
                if do_eval:
                    flush_time()
                    t_sv = common.timer()
                    save_last(run_dir, model, opt, sched, state, args, schedule, device)
                    overhead_inv += common.timer() - t_sv
                    t_window = common.timer()   # examples/s must not absorb eval + checkpoint time
            if not state.stopped_by:
                state.batch_in_epoch = 0
                state.epoch += 1
                if state.epoch >= args.epochs:
                    state.stopped_by = "epochs_done"
    except KeyboardInterrupt:
        LOG.error("second interrupt: hard stop without a checkpoint")
        return 130

    flush_time()
    save_last(run_dir, model, opt, sched, state, args, schedule, device)
    if args.cpu and torch.cuda.is_initialized():
        # The --cpu contract is "the GPU is never touched"; a regression here would OOM the pretraining run.
        LOG.error("--cpu run initialised CUDA; a code path touched torch.cuda on the CPU dev path (bug)")
        exit_code = exit_code or 3

    rec = experiment_record(args, state, params, steps_per_epoch)
    exp_path = Path(cfg["paths"]["reports"]) / "experiments.jsonl"
    how = upsert_experiment(exp_path, rec)
    LOG.info("%s experiment record for %s in %s", how, args.run_name, exp_path)
    if state.step >= RUN_LOG_MIN_STEPS:
        common.append_run_log(run_log_block(args, state, run_dir, params))
        LOG.info("appended a RUN_LOG.md entry")
    v = state.best_val or {}
    print(f"\nTotal training time: {state.train_seconds:.1f} s ({_hms(state.train_seconds)}) over {state.step} optimizer steps"
          f" [{state.stopped_by}]; this invocation {common.timer() - t_invocation:.1f} s wall")
    print(f"Peak VRAM reserved {state.peak_reserved_gb:.2f} GiB (allocated {state.peak_allocated_gb:.2f} GiB) on {device_info['device']}")
    if v:
        print(f"Best val by {args.selection_metric} (step {state.best_step}, epoch {state.best_epoch}), {UNWANTED_NAME} roll-up: "
              + " ".join(f"{k}={v[k]:.4f}" for k in METRIC_KEYS) + f"  -> {run_dir / 'best'}")
        if v.get("n_3way"):
            print(f"Three-way (n={v['n_3way']}): macro-F1={v['macro_f1_3way']:.4f} accuracy={v['accuracy_3way']:.4f} "
                  + " ".join(f"F1[{LABEL_NAMES[c]}]={v[f'f1_{s}']:.4f}" for c, s in enumerate(CLASS_SLUGS))
                  + f"; phishing-vs-spam on true unwanted [pp ps sp ss]=[{v['pvs_phish_as_phish']} {v['pvs_phish_as_spam']} "
                    f"{v['pvs_spam_as_phish']} {v['pvs_spam_as_spam']}]")
    return exit_code


def _hms(seconds: float) -> str:
    s = int(round(seconds))
    return f"{s // 3600:d}:{(s % 3600) // 60:02d}:{s % 60:02d}"


# ------------------------------------------------------------------------------------ plot
def plot(ns: argparse.Namespace, cfg: dict) -> int:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    run_dir = Path(cfg["paths"]["checkpoints"]) / "finetune" / ns.run_name
    metrics_csv, epochs_csv = run_dir / "metrics.csv", run_dir / "epochs.csv"
    if not metrics_csv.exists():
        LOG.error("%s not found", metrics_csv)
        return 2
    m = _read_csv(metrics_csv)
    e = _read_csv(epochs_csv) if epochs_csv.exists() else []
    out = Path(ns.out) if ns.out else Path(cfg["paths"]["reports"]) / "finetune_curves.png"
    out.parent.mkdir(parents=True, exist_ok=True)

    # Two panels, one y-axis each (loss and scores live on different scales; no dual axis).
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.2), dpi=150)
    steps = [float(r["step"]) for r in m]
    ax1.plot(steps, [float(r["loss"]) for r in m], color=PALETTE[0], lw=2, label="train loss")
    if e:
        ax1.plot([float(r["step"]) for r in e], [float(r["val_loss"]) for r in e], color=PALETTE[1], lw=2,
                 marker="o", ms=4, label="val loss (CE)")
    ax1.set_xlabel("optimizer step")
    ax1.set_ylabel("loss")
    ax1.set_title(f"{ns.run_name}: loss")
    ax1.legend(frameon=False)
    if e:
        es = [float(r["step"]) for r in e]
        series = [("f1", "val F1 (unwanted roll-up)", PALETTE[0]), ("roc_auc", "val ROC-AUC", PALETTE[1]),
                  ("pr_auc", "val PR-AUC", PALETTE[2])]
        if any(_num(r.get("macro_f1_3way")) is not None for r in e):
            series.append(("macro_f1_3way", "val macro-F1 (3-way)", PALETTE[3]))
        lows = []
        for key, label, color in series:
            pts = [(s, _num(r.get(key))) for s, r in zip(es, e)]
            pts = [(s, y) for s, y in pts if y is not None]
            if not pts:
                continue
            ax2.plot([p[0] for p in pts], [p[1] for p in pts], color=color, lw=2, marker="o", ms=4, label=label)
            lows.append(min(p[1] for p in pts))
        best = max(e, key=lambda r: float(r["f1"]))
        ax2.annotate(f"best F1 {float(best['f1']):.4f}\nstep {int(float(best['step']))}", (float(best["step"]), float(best["f1"])),
                     textcoords="offset points", xytext=(8, -28), fontsize=8, color="#444444")
        ax2.set_ylim(max(0.0, min(0.5, min(lows) - 0.05)) if lows else 0.0, 1.0)
        ax2.legend(frameon=False, loc="lower right")
    else:
        ax2.text(0.5, 0.5, "no evaluations yet", ha="center", va="center", transform=ax2.transAxes)
    ax2.set_xlabel("optimizer step")
    ax2.set_ylabel("score")
    ax2.set_title(f"{ns.run_name}: validation")
    for ax in (ax1, ax2):
        ax.grid(True, color="#e6e6e6", lw=0.6)
        for side in ("top", "right"):
            ax.spines[side].set_visible(False)
    fig.tight_layout()
    fig.savefig(out)
    LOG.info("wrote %s (%d train rows, %d eval rows)", out, len(m), len(e))
    print(out)
    return 0


def _read_csv(path: Path) -> list[dict[str, str]]:
    with path.open(newline="", encoding="utf-8") as fh:
        return [r for r in csv.DictReader(fh) if r.get("step")]


def _num(s: str | None) -> float | None:
    """CSV cell -> float, None for blank / NaN (a binary run has no three-way columns)."""
    if s is None or s == "":
        return None
    try:
        v = float(s)
    except ValueError:
        return None
    return None if math.isnan(v) else v


# ------------------------------------------------------------------------------------ list
def list_experiments(ns: argparse.Namespace, cfg: dict) -> int:
    path = Path(cfg["paths"]["reports"]) / "experiments.jsonl"
    if not path.exists():
        print(f"no experiments yet ({path})")
        return 0
    runs: dict[str, dict[str, Any]] = {}
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        rec = json.loads(line)
        slot = runs.setdefault(rec["run_name"], {})
        slot["test" if rec.get("test") is not None else "train"] = rec

    def f(x: Any, nd: int = 4) -> str:
        return "-" if x is None or (isinstance(x, float) and math.isnan(x)) else (f"{x:.{nd}f}" if isinstance(x, float) else str(x))

    # "val acc/P/R/F1" are the Unwanted (phishing or spam) roll-up at 0.5, identical to the binary runs' numbers;
    # the three-way columns are blank for binary (num_labels 2) runs.
    head = ["run", "init", "classes", "layers", "frozen", "epochs run/planned", "loss", "lr", "batch", "synthetic", "params (M)",
            "train min", "peak GiB", "best ep", "val acc", "val P", "val R", "val F1", "val ROC-AUC", "val PR-AUC",
            "val macro-F1 3w", "val P phish", "val R phish", "test F1", "test ROC-AUC", "stopped by"]
    rows = []
    for name, slot in runs.items():
        tr, te = slot.get("train"), slot.get("test")
        base = tr or te
        d = base.get("dials", {})
        v = (tr or {}).get("val") or {}
        t = (te or {}).get("test") or {}
        phish = (v.get("per_class") or {}).get(LABEL_NAMES[1]) or {}
        init = base.get("init", "-")
        init = "scratch" if init == "scratch" else "pretrained" if init.startswith("pretrained") else init
        rows.append([name, init, f(d.get("num_labels", 2)), f(d.get("n_layers")), f(d.get("frozen_layers")),
                     f"{f(d.get('epochs_run'), 2)}/{f(d.get('epochs_planned'))}", f(d.get("loss")),
                     f"{d['lr']:.0e}" if d.get("lr") else "-", f(d.get("batch_size")),
                     f"{d.get('synthetic', '-')}" + (f" ({d['synthetic_rows']})" if d.get("synthetic_rows") else ""),
                     f((base.get("params") or {}).get("total", 0) / 1e6, 1) if base.get("params") else "-",
                     f((tr or {}).get("train_seconds", 0) / 60.0, 1) if tr else "-", f((tr or {}).get("peak_vram_gb"), 2),
                     f((tr or {}).get("best_epoch")), f(v.get("accuracy")), f(v.get("precision")), f(v.get("recall")),
                     f(v.get("f1")), f(v.get("roc_auc")), f(v.get("pr_auc")), f(v.get("macro_f1_3way")),
                     f(phish.get("precision")), f(phish.get("recall")), f(t.get("f1")), f(t.get("roc_auc")),
                     f((tr or {}).get("stopped_by"))])
    print("| " + " | ".join(head) + " |")
    print("|" + "|".join("---" for _ in head) + "|")
    for r in rows:
        print("| " + " | ".join(r) + " |")
    print(f"\n{len(rows)} run(s) from {path}")
    return 0


# ------------------------------------------------------------------------------------- CLI
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--config", default="config.yaml")
    sub = ap.add_subparsers(dest="cmd", required=True)

    tr = sub.add_parser("train", help="fine-tune the classifier (dials default to config.yaml; --resume restores the saved ones)")
    tr.add_argument("--run-name", default="finetune")
    tr.add_argument("--init", default=None, help="pretrained | pretrained:<dir> | scratch (default: newest checkpoints/pretrain/*/final)")
    tr.add_argument("--n-layers", type=int, default=None, help="use only the first N encoder layers")
    tr.add_argument("--freeze-layers", type=int, default=None, help="freeze embeddings + first K layers")
    tr.add_argument("--loss", choices=LOSSES, default=None)
    tr.add_argument("--focal-gamma", type=float, default=None)
    tr.add_argument("--lr", type=float, default=None)
    tr.add_argument("--batch-size", type=int, default=None, help="logical batch (rows per optimizer step)")
    tr.add_argument("--micro-batch", type=int, default=None, help="rows per fwd/bwd; batch-size must be a multiple (default 16)")
    tr.add_argument("--epochs", type=int, default=None)
    tr.add_argument("--max-steps", type=int, default=None, help="dev: stop after this many optimizer steps")
    tr.add_argument("--eval-every", type=int, default=None, help="optimizer steps between full val evaluations")
    tr.add_argument("--log-every", type=int, default=None)
    tr.add_argument("--eval-batch-size", type=int, default=None)
    tr.add_argument("--synthetic", choices=("on", "off"), default=None, help="append data/processed/synthetic.parquet to TRAIN")
    tr.add_argument("--limit-rows", type=int, default=None, help="dev: seeded subsample of train AND val")
    tr.add_argument("--seed", type=int, default=None)
    tr.add_argument("--num-labels", type=int, choices=NUM_LABELS_ALLOWED, default=None,
                    help="override finetune.num_labels: 3 = legitimate/phishing/spam (partial-label loss), 2 = binary baseline")
    tr.add_argument("--cpu", action="store_true", help="dev only: run on the CPU with the config as given (tiny encoder + fixtures)")
    tr.add_argument("--resume", action="store_true",
                    help="continue <run>/last (dials restored; CLI flags override, except --init/--n-layers/--freeze-layers/"
                         "--batch-size/--seed/--num-labels which are refused; --lr is re-applied to the restored optimizer)")
    tr.add_argument("--overwrite", action="store_true", help="discard an existing run dir of the same name")

    pl = sub.add_parser("plot", help="train loss and val F1/AUC vs step -> reports/finetune_curves.png")
    pl.add_argument("--run-name", required=True)
    pl.add_argument("--out", default=None)

    sub.add_parser("list", help="reports/experiments.jsonl as a markdown table")
    return ap.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    ns = parse_args(argv)
    cfg = common.load_config(ns.config)
    global LOG
    LOG = common.setup_logging("finetune", cfg["paths"]["logs"])
    if ns.cmd == "train":
        return train(ns, cfg)
    if ns.cmd == "plot":
        return plot(ns, cfg)
    return list_experiments(ns, cfg)


if __name__ == "__main__":
    sys.exit(main())