| """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 |
| import contextlib |
| import csv |
| import json |
| import logging |
| import math |
| import random |
| import signal |
| import sys |
| import time |
| import zlib |
| from dataclasses import asdict, dataclass, field, fields |
| from pathlib import Path |
| from typing import Any |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
| from safetensors.torch import load_file, save_file |
| from torch import Tensor |
| from torch.utils.data import DataLoader, Dataset |
|
|
| import common |
| from model import EncoderConfig, EncoderForClassification, count_parameters |
| from smoke_test import build_optimizer |
| from tokenizer import encode_batch, load_tokenizer |
|
|
| 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 |
| |
| RESUME_FIXED_DIALS = ("init", "n_layers", "freeze_layers", "batch_size", "seed", "num_labels") |
| |
| 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"} |
|
|
| |
| 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", "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) |
| 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) |
| |
| PALETTE = ("#2a78d6", "#eb6834", "#1baf7a", "#eda100") |
|
|
|
|
| |
| @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" |
| 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" |
| selection_metric: str = "val_f1_unwanted" |
| pooling: str = "cls" |
| num_labels: int = 3 |
| synthetic: str = "off" |
| limit_rows: int | None = None |
| seed: int = 42 |
| cpu: bool = False |
|
|
| @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"]), |
| |
| 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: |
| |
| 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 |
|
|
|
|
| |
| @dataclass |
| class Labelled: |
| """The label columns of one parquet split after validation against the three-way contract.""" |
|
|
| texts: list[str] |
| labels: np.ndarray |
| label3: np.ndarray |
| allowed: np.ndarray |
| 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 |
| offsets: np.ndarray |
| labels: np.ndarray |
| label3: np.ndarray |
| allowed: np.ndarray |
| sources: list[str] |
| texts: list[str] |
|
|
| 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: |
| |
| 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)] |
|
|
| |
| |
| |
| |
| 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(): |
| |
| 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)) |
|
|
|
|
| |
| 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>'.""" |
| |
| 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']}") |
| |
| 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) |
|
|
|
|
| |
| 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}") |
|
|
|
|
| |
| 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: |
| 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()) |
| |
| |
| 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) |
|
|
|
|
| |
| @dataclass |
| class RunState: |
| step: int = 0 |
| epoch: int = 0 |
| batch_in_epoch: int = 0 |
| best_score: float = -1.0 |
| 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"] |
| |
| 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"] |
|
|
|
|
| |
| 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, |
| |
| "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) |
| 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.", |
| ]) |
|
|
|
|
| |
| 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() |
| |
| |
| |
| |
| if ns.cpu: |
| |
| |
| |
| |
| 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 |
| |
| 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) |
| 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: |
| 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: |
| 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"]): |
| |
| |
| 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 |
| 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 |
|
|
| |
| 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 |
|
|
| |
| |
| 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() |
| 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(): |
| |
| 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}" |
|
|
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
|
|
| |
| 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)) |
|
|
| |
| |
| 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 |
|
|
|
|
| |
| 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()) |
|
|
