One-command GPU run for the transformer grids

.gitignore

@@ -33,6 +33,8 @@ mlartifacts/
 .env.*
 !.env.example
 .venv/
+.venv-gpu/
+gpu-out/
 venv/
 
 # Python

README.md

@@ -241,14 +241,13 @@ uv run triggermap evaluate annotations/batches/batch-03.jsonl reports/agreement/
     --add-gold annotations/batches/batch-06.jsonl --split dev   # agreement with model drafts, not accuracy
 ```
 
-**Transformers** (once gold labels exist; DistilBERT, DeBERTa-v3-small; GPU recommended, see
+**Transformers** (DistilBERT, DeBERTa-v3-small; needs a CUDA GPU, see
 [`docs/training_on_gpu.md`](docs/training_on_gpu.md)):
 
 ```bash
-uv sync --extra neural
+bash scripts/gpu/run_neural.sh                       # on the GPU machine: both grids, packed results
+bash scripts/gpu/score_neural.sh triggermap-neural-results.tgz   # back here: reports/neural/
 uv run triggermap experiment --config configs/experiments/deberta.toml --dry-run      # planned runs
-uv run triggermap experiment --config configs/experiments/deberta.toml \
-    --gold annotations/pilot/pilot-1.jsonl annotations/batches/batch-0*.jsonl          # lr x seed grid
 uv run triggermap predict --model transformer --model-dir artifacts/experiments/<run> --out preds.jsonl
 ```
 

docs/training_on_gpu.md

@@ -1,39 +1,69 @@
 # Training the transformer models on a GPU
 
-Fine-tuning DistilBERT and DeBERTa-v3-small needs a GPU; the local Mac is only for smoke
-tests. These steps use Google Colab (a free T4 GPU, 16 GB, is enough). The same commands
-work on any Linux GPU machine. Training never runs in CI or on the serving host.
+DistilBERT and DeBERTa-v3-small are trained offline on one CUDA GPU, never in CI or on the
+serving host. `scripts/gpu/run_neural.sh` does the whole run in one command, and
+`scripts/gpu/score_neural.sh` scores the packed results on the development machine.
 
-## 1. Runtime
+## What the run does
 
-In Colab: **Runtime → Change runtime type → T4 GPU**. Then, in a cell:
+1. Creates a Python 3.12 environment with CUDA PyTorch (`.venv-gpu`).
+2. Rebuilds the corpus from the pinned snapshot (`1912cbd…`, a sparse clone of about 360 MB) and
+   stops unless `data/processed/advisories.jsonl` has the documented SHA-256.
+3. Runs `configs/experiments/distilbert.toml` and `configs/experiments/deberta.toml`: two
+   learning rates (2e-5, 5e-5) × three seeds each, 12 runs, with up to 6 epochs and early stopping.
+   - **Training:** the train split of `pilot-2`, `batch-01` and `batch-02` (239 documents).
+   - **Early stopping and run selection:** only the model-drafted dev batches `batch-03` and
+     `batch-06` (`--dev-gold`). The 10 hand-labelled dev documents and the test split play no
+     part in selection.
+4. Predicts with every run on the pilot, `batch-03` and `batch-06` documents.
+5. Packs two archives:
+   - `triggermap-neural-results.tgz`: reports, predictions, run metadata and MLflow runs. No
+     weights. This is the one to bring back.
+   - `triggermap-neural-bundles.tgz`: the best bundle of each grid (weights, tokenizer,
+     metadata), for later export and serving.
 
-```bash
-!nvidia-smi --query-gpu=name,memory.total --format=csv
-```
+Runs resume: if the session stops, run the script again and finished runs are reused.
+
+## On Kaggle
+
+1. Create a notebook. Under *Settings*, choose **GPU T4** as the accelerator and turn
+   **Internet** on. (The P100 option is too old for the CUDA 13 PyTorch wheels.)
+2. Run one cell:
 
-## 2. Code and environment
+   ```bash
+   !git clone https://github.com/professor3333/TriggerMap.git && cd TriggerMap && bash scripts/gpu/run_neural.sh
+   ```
+
+3. Download both `.tgz` files from `/kaggle/working` (the *Output* panel).
+
+If the script stops with `torch has no CUDA`, the driver is older than CUDA 13 needs. Re-run
+with the CUDA 12.6 wheels:
 
 ```bash
-!git clone https://github.com/professor3333/TriggerMap.git
-%cd TriggerMap
-!pip install -q -e ".[neural]"
+!cd TriggerMap && TORCH_INDEX=https://download.pytorch.org/whl/cu126 bash scripts/gpu/run_neural.sh
 ```
 
-Use plain `pip` here, not `uv sync`. The project's uv configuration pins **CPU** PyTorch on
-Linux (so CI stays small); `pip` ignores that setting and keeps Colab's CUDA build.
+On Colab or a rented machine the same commands work (without the leading `!` in a shell). The
+archives are then written to `./gpu-out`.
 
-## 3. Rebuild the corpus and check it matches
+## Scoring on the development machine
 
 ```bash
-!triggermap ingest fetch --commit 1912cbd6ce9d4cb141636adc4abafb4eb77a93e6
-!triggermap ingest build
-!git diff --exit-code data/manifests/snapshot.json && echo "corpus identical to the pinned manifest"
+bash scripts/gpu/score_neural.sh ~/Downloads/triggermap-neural-results.tgz
 ```
 
-If the manifest differs, stop: the gold annotations would not match the text.
+This writes `reports/neural/`:
+- `<run>_dev_agreement.md`: agreement with the 172 model-drafted dev records. This is not
+  accuracy.
+- `<run>_dev_manual.md`: accuracy on the 10 hand-labelled dev documents.
+- `compare_*`: paired comparisons of each grid's best run, chosen by its early-stopping score,
+  against `crf-batches-01-02` on both slices.
+
+The test split is not touched.
+
+## Manual steps and notes
 
-## 4. Persist results to Google Drive (survives disconnects)
+### Colab: keep results on Google Drive (survives disconnects)
 
 ```python
 from google.colab import drive
@@ -50,22 +80,22 @@ drive.mount("/content/drive")
 `triggermap experiment` skips any run whose bundle already exists, so after a disconnect
 you can rerun the same command and only unfinished runs train again.
 
-## 5. Run the experiments
+### Running the grids by hand
 
-Gold files are the committed annotation JSONL files: the pilot and every imported batch.
+The script runs these commands. Use them directly for ablations or a single config, with the
+environment the script created (`.venv-gpu`).
 
 ```bash
-!triggermap experiment --config configs/experiments/distilbert.toml --dry-run
-!triggermap experiment --config configs/experiments/distilbert.toml \
-    --gold annotations/pilot/pilot-1.jsonl annotations/batches/batch-0*.jsonl \
+!.venv-gpu/bin/triggermap experiment --config configs/experiments/distilbert.toml --dry-run
+!.venv-gpu/bin/triggermap experiment --config configs/experiments/distilbert.toml \
+    --gold annotations/pilot/pilot-2.jsonl annotations/batches/batch-01.jsonl annotations/batches/batch-02.jsonl \
+    --dev-gold annotations/batches/batch-03.jsonl annotations/batches/batch-06.jsonl \
     --report /content/drive/MyDrive/triggermap/reports/distilbert.md
-!triggermap experiment --config configs/experiments/deberta.toml \
-    --gold annotations/pilot/pilot-1.jsonl annotations/batches/batch-0*.jsonl \
-    --report /content/drive/MyDrive/triggermap/reports/deberta.md
 ```
 
 Each config runs learning rates {2e-5, 5e-5} × seeds {0, 1, 2}. Training uses only the train
-pool, model selection and the report use only the dev pool, and the test pool is never read.
+split, early stopping and the report use only the dev records of `--dev-gold`, and the test
+split is never read.
 After choosing a backbone, set it in `configs/experiments/ablations.toml` and run the
 ablations (no context, no group objective) and learning curves (125 and 250 documents)
 the same way.
@@ -75,16 +105,16 @@ split hash, hyperparameters, seed, per-epoch dev metrics, training time and peak
 memory. When an MLflow experiment is configured (the default), the same values go to
 `mlruns/`.
 
-## 6. Profile before spending
+### Profile before spending
 
 Run one configuration with `--dry-run` first, then a single run, and read `train_seconds`
 and `peak_gpu_memory_mb` from its `metadata.json` before launching a whole grid. Paid GPU
 time is capped at about US$50 for the project. If memory runs out, reduce `batch_size` and
 raise `gradient_accumulation`, keeping their product the same.
 
-## 7. Back on your machine
+### MLflow runs
 
-Copy the reports and the chosen bundle from Drive. Browse runs with:
+The results archive includes `mlruns/`. Browse it with:
 
 ```bash
 uvx --from mlflow==3.17.1 mlflow ui --backend-store-uri ./mlruns

scripts/gpu/run_neural.sh

@@ -0,0 +1,89 @@
+#!/usr/bin/env bash
+# Train and score the DistilBERT and DeBERTa-v3-small grids on one CUDA GPU (Kaggle, Colab or a
+# rented box), then pack the results for scoring on the development machine.
+#
+#   git clone https://github.com/professor3333/TriggerMap.git && cd TriggerMap
+#   bash scripts/gpu/run_neural.sh                 # both grids (12 runs)
+#   CONFIGS="distilbert" bash scripts/gpu/run_neural.sh
+#   TORCH_INDEX=https://download.pytorch.org/whl/cu126 bash scripts/gpu/run_neural.sh
+#
+# PyPI's torch 2.14 wheels target CUDA 13: they need a recent NVIDIA driver and a Turing (T4) or
+# newer GPU. If the CUDA check below fails, set TORCH_INDEX to a CUDA 12 wheel index as shown.
+#
+# Writes $OUT_DIR/triggermap-neural-results.tgz (reports, predictions, run metadata, MLflow runs;
+# no weights) and $OUT_DIR/triggermap-neural-bundles.tgz (the best bundle of each grid, by dev
+# typed F1). OUT_DIR defaults to /kaggle/working on Kaggle, otherwise ./gpu-out.
+#
+# Selection never sees the test split or the hand-labelled dev documents: training uses the
+# train split of the pilot and batches 01-02, and early stopping uses only the model-drafted
+# dev batches 03 and 06 (--dev-gold). Re-running resumes: finished runs are not retrained.
+set -euo pipefail
+
+CONFIGS=${CONFIGS:-"distilbert deberta"}
+if [ -d /kaggle/working ]; then OUT_DIR=${OUT_DIR:-/kaggle/working}; else OUT_DIR=${OUT_DIR:-gpu-out}; fi
+SNAPSHOT=1912cbd6ce9d4cb141636adc4abafb4eb77a93e6
+CORPUS_SHA256=734043ff982a337f6cc593745160e997f20ee62609183ff857edce3d47d2b655
+UV_VERSION=0.12.5
+TRAIN_GOLD=(annotations/pilot/pilot-2.jsonl annotations/batches/batch-01.jsonl annotations/batches/batch-02.jsonl)
+DEV_GOLD=(annotations/batches/batch-03.jsonl annotations/batches/batch-06.jsonl)
+SAMPLES=(annotations/pilot/pilot_manifest.json annotations/batches/batch-03.json annotations/batches/batch-06.json)
+
+[ -f pyproject.toml ] && [ -d annotations ] || { echo "run from the repository root" >&2; exit 1; }
+nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || { echo "no CUDA GPU found" >&2; exit 1; }
+mkdir -p "$OUT_DIR"
+
+# 1. Environment: Python 3.12 with CUDA PyTorch from PyPI. --no-sources skips the CPU-only torch
+#    index that pyproject.toml pins for Linux development machines.
+command -v uv >/dev/null || python3 -m pip install --quiet "uv==$UV_VERSION"
+uv venv --quiet --python 3.12 .venv-gpu
+uv pip install --quiet --python .venv-gpu --no-sources -e ".[neural]"
+if [ -n "${TORCH_INDEX:-}" ]; then
+  uv pip install --quiet --python .venv-gpu --reinstall-package torch \
+    --index-url "$TORCH_INDEX" --extra-index-url https://pypi.org/simple "torch>=2.14,<2.15"
+fi
+TM=(.venv-gpu/bin/triggermap)
+.venv-gpu/bin/python -c "import torch; assert torch.cuda.is_available(), 'torch has no CUDA: see TORCH_INDEX at the top of this script'; print('torch', torch.__version__, torch.cuda.get_device_name(0))"
+
+# 2. Corpus from the pinned snapshot; stop if it is not byte-identical to the documented one.
+if [ ! -f data/processed/advisories.jsonl ]; then
+  "${TM[@]}" ingest fetch --commit "$SNAPSHOT"
+  "${TM[@]}" ingest build
+fi
+echo "$CORPUS_SHA256  data/processed/advisories.jsonl" | sha256sum -c -
+
+# 3. Grids, then predictions of every run on the pilot and dev-batch documents.
+mkdir -p reports/neural/predictions
+for config in $CONFIGS; do
+  "${TM[@]}" experiment --config "configs/experiments/$config.toml" \
+    --gold "${TRAIN_GOLD[@]}" --dev-gold "${DEV_GOLD[@]}" \
+    --out-root artifacts/experiments --report "reports/neural/experiment_$config.md"
+done
+for bundle in artifacts/experiments/*/; do
+  run=$(basename "$bundle")
+  out="reports/neural/predictions/$run.jsonl"
+  [ -s "$out" ] || "${TM[@]}" predict --model transformer --model-dir "$bundle" \
+    --sample "${SAMPLES[@]}" --out "$out"
+  cp "$bundle/metadata.json" "reports/neural/predictions/$run.metadata.json"
+done
+
+# 4. Pack. Results carry no weights; bundles hold the best run of each grid.
+pack=(reports/neural)
+[ -d mlruns ] && pack+=(mlruns)
+tar czf "$OUT_DIR/triggermap-neural-results.tgz" "${pack[@]}"
+best=$(.venv-gpu/bin/python - "$CONFIGS" <<'EOF'
+import json, sys
+from pathlib import Path
+for config in sys.argv[1].split():
+    prefix = {"distilbert": "distilbert-1", "deberta": "deberta-1"}.get(config, config)
+    runs = [
+        (json.loads((p / "metadata.json").read_text()).get("best_dev_typed_f1") or 0.0, p)
+        for p in Path("artifacts/experiments").glob(f"{prefix}-*") if (p / "metadata.json").exists()
+    ]
+    if runs:
+        print(max(runs)[1])
+EOF
+)
+# shellcheck disable=SC2086
+[ -n "$best" ] && tar czf "$OUT_DIR/triggermap-neural-bundles.tgz" $best
+ls -lh "$OUT_DIR"/triggermap-neural-*.tgz
+echo "Copy triggermap-neural-results.tgz back and run: bash scripts/gpu/score_neural.sh <path>"

scripts/gpu/score_neural.sh

@@ -0,0 +1,58 @@
+#!/usr/bin/env bash
+# Score the GPU runs packed by run_neural.sh on the development machine (no GPU needed).
+#
+#   bash scripts/gpu/score_neural.sh ~/Downloads/triggermap-neural-results.tgz
+#
+# Writes reports/neural/<run>_dev_agreement.md (agreement with the model-drafted dev batches 03
+# and 06, not accuracy), <run>_dev_manual.md (the 10 hand-labelled dev documents) and paired
+# comparisons of each grid's best run against crf-batches-01-02. The test split is not touched.
+set -euo pipefail
+
+archive=${1:?usage: score_neural.sh RESULTS.tgz}
+[ -f pyproject.toml ] && [ -d annotations ] || { echo "run from the repository root" >&2; exit 1; }
+
+# Unpack into a fresh directory and copy only the expected files.
+tmp=$(mktemp -d)
+trap 'rm -rf "$tmp"' EXIT
+tar -xzf "$archive" -C "$tmp" reports/neural
+mkdir -p reports/neural/predictions
+cp "$tmp"/reports/neural/experiment_*.md reports/neural/ 2>/dev/null || true
+cp "$tmp"/reports/neural/predictions/*.jsonl "$tmp"/reports/neural/predictions/*.metadata.json \
+  reports/neural/predictions/
+
+TM=(uv run --no-sync triggermap)
+AGREE=(annotations/batches/batch-03.jsonl)
+AGREE_ADD=(--add-gold annotations/batches/batch-06.jsonl)
+CRF=reports/agreement/predictions/crf-batches-01-02.jsonl
+
+for pred in reports/neural/predictions/*.jsonl; do
+  run=$(basename "$pred" .jsonl)
+  "${TM[@]}" evaluate "${AGREE[@]}" "$pred" "${AGREE_ADD[@]}" --split dev \
+    --title "Agreement with model drafts: $run" --out "reports/neural/${run}_dev_agreement.md" >/dev/null
+  "${TM[@]}" evaluate annotations/pilot/pilot-2.jsonl "$pred" --split dev --gold-origin manual \
+    --title "Hand-labelled dev: $run" --out "reports/neural/${run}_dev_manual.md" >/dev/null
+done
+
+# Best run per grid by early-stopping score (model-drafted dev), never by the manual slice.
+best=$(python3 -I - <<'EOF'
+import json
+from pathlib import Path
+best = {}
+for meta in Path("reports/neural/predictions").glob("*.metadata.json"):
+    data = json.loads(meta.read_text())
+    family = data["model_id"].split("-")[0]
+    score = data.get("best_dev_typed_f1") or 0.0
+    if family not in best or score > best[family][0]:
+        best[family] = (score, meta.name.removesuffix(".metadata.json"))
+print(" ".join(run for _, run in best.values()))
+EOF
+)
+for run in $best; do
+  pred=reports/neural/predictions/$run.jsonl
+  "${TM[@]}" compare "${AGREE[@]}" "$CRF" "$pred" "${AGREE_ADD[@]}" --split dev \
+    --title "Agreement: crf-batches-01-02 -> $run" --out "reports/neural/compare_agreement_$run.md" >/dev/null
+  "${TM[@]}" compare annotations/pilot/pilot-2.jsonl "$CRF" "$pred" --split dev --gold-origin manual \
+    --title "Hand-labelled dev: crf-batches-01-02 -> $run" --out "reports/neural/compare_manual_$run.md" >/dev/null
+done
+echo "best runs: $best"
+ls reports/neural

src/triggermap/cli.py

@@ -392,7 +392,7 @@ def _train(args: argparse.Namespace) -> int:
         settings["split_assignment_sha256"] = load_split(args.split_manifest)["assignment_sha256"]
         settings.setdefault("model_id", args.model_id)
         settings["seed"] = args.seed if args.seed is not None else settings.get("seed", 0)
-        dev = _gold_examples(args.gold, args.processed, args.split_manifest, "dev")
+        dev = _gold_examples(args.dev_gold or args.gold, args.processed, args.split_manifest, "dev")
         transformer = train(TrainConfig(**settings), examples, dev, args.out)
         best = transformer.metadata["best_dev_typed_f1"]
         print(
@@ -434,7 +434,7 @@ def _experiment(args: argparse.Namespace) -> int:
 
     split = load_split(args.split_manifest)
     examples = _gold_examples(args.gold, args.processed, args.split_manifest, origins=args.origin)
-    dev = _gold_examples(args.gold, args.processed, args.split_manifest, "dev")
+    dev = _gold_examples(args.dev_gold or args.gold, args.processed, args.split_manifest, "dev")
     if not examples or not dev:
         raise SystemExit("experiments need gold documents in both the train and dev splits")
     results = []
@@ -459,7 +459,9 @@ def _experiment(args: argparse.Namespace) -> int:
             results,
             {
                 "Config": f"`{args.config}`",
-                "Gold": f"{len(args.gold)} file(s): {len(examples)} train, {len(dev)} dev docs",
+                "Gold": f"{len(args.gold)} file(s): {len(examples)} train docs",
+                "Dev gold": ", ".join(f"`{p}`" for p in (args.dev_gold or args.gold))
+                + f": {len(dev)} dev docs (early stopping and this table)",
                 "Split": f"`{args.split_manifest}` ({split['assignment_sha256'][:12]})",
             },
         ),
@@ -944,6 +946,16 @@ def _gold_names(args: argparse.Namespace) -> str:
     return ", ".join(f"`{path}`" for path in _gold_paths(args))
 
 
+def _dev_gold_argument(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument(
+        "--dev-gold",
+        type=Path,
+        nargs="+",
+        help="gold file(s) whose dev records drive early stopping (default: --gold); keeps an "
+        "evaluation slice out of model selection",
+    )
+
+
 def _scoring_arguments(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "--add-gold",
@@ -1089,6 +1101,7 @@ def build_parser() -> argparse.ArgumentParser:
     train.add_argument("--model", choices=["crf", "transformer"], default="crf")
     train.add_argument("--config", type=Path, help="TOML experiment config (transformer)")
     train.add_argument("--gold", type=Path, nargs="+", required=True, help="gold JSONL file(s)")
+    _dev_gold_argument(train)
     train.add_argument("--processed", type=Path, default=DEFAULT_PROCESSED)
     train.add_argument("--out", type=Path, required=True, help="bundle directory to write")
     train.add_argument("--split-manifest", type=Path, default=DEFAULT_SPLIT)
@@ -1246,6 +1259,7 @@ def build_parser() -> argparse.ArgumentParser:
     experiment = commands.add_parser("experiment", help="run a grid of transformer trainings")
     experiment.add_argument("--config", type=Path, required=True)
     experiment.add_argument("--gold", type=Path, nargs="+", help="gold JSONL file(s)")
+    _dev_gold_argument(experiment)
     experiment.add_argument("--processed", type=Path, default=DEFAULT_PROCESSED)
     experiment.add_argument("--split-manifest", type=Path, default=DEFAULT_SPLIT)
     experiment.add_argument("--out-root", type=Path, default=Path("artifacts/experiments"))

tests/test_transformer.py

@@ -234,6 +234,12 @@ def test_cli_train_predict_and_experiment(base: Path, tmp_path: Path) -> None:
     # Resume: a second invocation reuses both finished bundles instead of retraining.
     assert main([*experiment_args, "--samples", "20", *common]) == 0
     assert report.read_text().count("tiny-base-seed") == text.count("tiny-base-seed")
+    # --dev-gold: early stopping and the report use only the given files' dev records.
+    dev_gold = tmp_path / "dev_gold.jsonl"
+    dev_gold.write_text("".join(r.model_dump_json() + "\n" for _, _, r in data[12:14]))
+    dev_args = [*experiment_args[:-1], str(tmp_path / "dev_report.md"), "--dev-gold", str(dev_gold)]
+    assert main([*dev_args, "--samples", "20", *common]) == 0
+    assert ": 2 dev docs" in (tmp_path / "dev_report.md").read_text()
 
 
 @pytest.mark.slow