"""Document-level K-fold cross-validation with error bars. The headline numbers in this repo come from a single 70/15/15 document split. Two models 0.009 macro-F1 apart on one split is inside the noise of a 151-document corpus, so no ranking claim survives without a variance estimate. This harness answers the only question that licenses a "beats X" claim: across independent, leak-free document folds, is the gap larger than its own spread? Design ------ * ``GroupKFold`` over ``doc_title`` assigns each of the 151 source reports to exactly one test fold — no report is ever split across train and test. * Inside each fold's training documents, a document-disjoint dev set is carved off for threshold and blend-weight tuning. Test is never touched during tuning. * macro-F1 is averaged over the techniques that actually have test support in that fold, so a technique that happens to land entirely in train does not drag every model's score toward zero. The same label set is used for every model within a fold, so the comparison stays fair. python scripts/06_cv.py --folds 5 --model modernbert ModernBERT is retrained from scratch inside every fold. On the 4060 this is a few minutes per fold; the whole run is well under an hour. """ import argparse import gc import json import shutil import sys import tempfile import time from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src")) import numpy as np # noqa: E402 import torch # noqa: E402 from sklearn.model_selection import GroupKFold # noqa: E402 from torch.utils.data import DataLoader # noqa: E402 from cti_attack import baselines, config, data, evaluate, modeling # noqa: E402 def macro_f1_supported(Yte: np.ndarray, Ypred: np.ndarray) -> tuple[float, int]: """Mean F1 over the labels that have >=1 positive in this fold's test set.""" supported = np.where(Yte.sum(axis=0) > 0)[0] fs = [] for j in supported: yt, yp = Yte[:, j], Ypred[:, j] tp = int(((yt == 1) & (yp == 1)).sum()) fp = int(((yt == 0) & (yp == 1)).sum()) fn = int(((yt == 1) & (yp == 0)).sum()) p = tp / (tp + fp) if tp + fp else 0.0 r = tp / (tp + fn) if tp + fn else 0.0 fs.append(2 * p * r / (p + r) if p + r else 0.0) return float(np.mean(fs)) if fs else 0.0, len(supported) def carve_dev(records, seed, dev_doc_frac=0.2): """Split records into (train, dev) by whole documents, dev-disjoint.""" docs = sorted({r["doc_title"] for r in records}) rng = np.random.RandomState(seed) rng.shuffle(docs) n_dev = max(1, int(len(docs) * dev_doc_frac)) dev_docs = set(docs[:n_dev]) tr = [r for r in records if r["doc_title"] not in dev_docs] dv = [r for r in records if r["doc_title"] in dev_docs] return tr, dv def eval_regimes(Ydv, s_dv, Yte, s_te): """Tune global + per-class thresholds on dev, report both on test.""" gt, _ = evaluate.tune_global_threshold(Ydv, s_dv) pct = evaluate.tune_per_class_thresholds(Ydv, s_dv) g, _ = macro_f1_supported(Yte, (s_te >= gt).astype(np.int8)) p, _ = macro_f1_supported(Yte, evaluate.apply_thresholds(s_te, pct)) return g, p def pick_alpha(Ydv, t_dv, b_dv): """Choose blend weight on dev: scores = alpha*tfidf + (1-alpha)*bert.""" best_a, best_f = 0.5, -1.0 for i in range(21): a = round(0.05 * i, 2) _, f = evaluate.tune_global_threshold(Ydv, a * t_dv + (1 - a) * b_dv) if f > best_f: best_a, best_f = a, f return best_a def run_fold(fold, tr, dv, te, labels, model_key, epochs): Ytr = evaluate.to_matrix(tr, labels) Ydv = evaluate.to_matrix(dv, labels).astype("int8") Yte = evaluate.to_matrix(te, labels).astype("int8") txt = lambda recs: [r["sentence"] for r in recs] # ---- TF-IDF ------------------------------------------------------------ s = baselines.tfidf_lr_scores(txt(tr), Ytr, {"dev": txt(dv), "test": txt(te)}) t_dv, t_te = s["dev"], s["test"] tf_g, tf_p = eval_regimes(Ydv, t_dv, Yte, t_te) # ---- encoder ----------------------------------------------------------- tmp = Path(tempfile.mkdtemp(prefix=f"cvfold{fold}_")) try: modeling.train(model_key, "cv", tr, dv, labels, tmp, epochs=epochs) model, tok = modeling.load_for_inference(tmp) dev_ = modeling.device() model.to(dev_) amp = model_key not in config.FP32_ONLY_MODELS ds_dv = modeling.SentenceDataset(dv, labels, tok, config.MAX_LENGTH) ds_te = modeling.SentenceDataset(te, labels, tok, config.MAX_LENGTH) b_dv = modeling.predict_scores(model, DataLoader(ds_dv, batch_size=32), dev_, amp=amp) b_te = modeling.predict_scores(model, DataLoader(ds_te, batch_size=32), dev_, amp=amp) del model gc.collect() torch.cuda.empty_cache() finally: shutil.rmtree(tmp, ignore_errors=True) bt_g, bt_p = eval_regimes(Ydv, b_dv, Yte, b_te) # ---- ensemble ---------------------------------------------------------- alpha = pick_alpha(Ydv, t_dv, b_dv) en_dv = alpha * t_dv + (1 - alpha) * b_dv en_te = alpha * t_te + (1 - alpha) * b_te en_g, en_p = eval_regimes(Ydv, en_dv, Yte, en_te) return { "fold": fold, "n_train": len(tr), "n_dev": len(dv), "n_test": len(te), "alpha": alpha, "tfidf_global": tf_g, "tfidf_perclass": tf_p, f"{model_key}_global": bt_g, f"{model_key}_perclass": bt_p, "ensemble_global": en_g, "ensemble_perclass": en_p, } def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--folds", type=int, default=5) ap.add_argument("--model", default="modernbert", choices=list(config.BASE_MODELS)) ap.add_argument("--epochs", type=int, default=config.EPOCHS) ap.add_argument("--include-single", action="store_true", help="merge single_label.json before splitting (Task 6 data lever)") args = ap.parse_args() records, labels, _ = data.build(verbose=False, include_single=args.include_single) groups = [r["doc_title"] for r in records] gkf = GroupKFold(n_splits=args.folds) print(f"{'=' * 62}\n document-level {args.folds}-fold CV " f"(tfidf_lr / {args.model} / ensemble)\n{'=' * 62}") print(f" records={len(records)} documents={len(set(groups))} labels={len(labels)}\n") fold_rows = [] t0 = time.time() for fold, (trval_idx, test_idx) in enumerate(gkf.split(records, groups=groups), 1): trval = [records[i] for i in trval_idx] te = [records[i] for i in test_idx] tr, dv = carve_dev(trval, seed=config.SPLIT_SEED + fold) print(f" --- fold {fold}/{args.folds} " f"train={len(tr)} dev={len(dv)} test={len(te)} " f"({len(set(g['doc_title'] for g in te))} test docs) ---") row = run_fold(fold, tr, dv, te, labels, args.model, args.epochs) fold_rows.append(row) print(f" tfidf={row['tfidf_perclass']:.4f} " f"{args.model}={row[args.model + '_perclass']:.4f} " f"ensemble={row['ensemble_perclass']:.4f} (alpha={row['alpha']:.2f}) " f"[{(time.time() - t0) / 60:.1f} min elapsed]\n") # ---- aggregate --------------------------------------------------------- keys = ["tfidf_global", "tfidf_perclass", f"{args.model}_global", f"{args.model}_perclass", "ensemble_global", "ensemble_perclass"] summary = {} print(f"{'=' * 62}\n {args.folds}-fold CV summary — macro-F1 mean +/- std\n{'=' * 62}") for k in keys: vals = np.array([r[k] for r in fold_rows]) summary[k] = {"mean": round(float(vals.mean()), 4), "std": round(float(vals.std(ddof=1)), 4), "folds": [round(float(v), 4) for v in vals]} print(f" {k:26} {vals.mean():.4f} +/- {vals.std(ddof=1):.4f}") # paired ensemble-vs-tfidf gap and its spread ens = np.array([r["ensemble_perclass"] for r in fold_rows]) tf = np.array([max(r["tfidf_perclass"], r["tfidf_global"]) for r in fold_rows]) gap = ens - tf print(f"\n paired gap (ensemble_perclass - best_tfidf), per fold: " f"{[round(float(g), 4) for g in gap]}") print(f" mean gap = {gap.mean():+.4f} +/- {gap.std(ddof=1):.4f} " f"(min {gap.min():+.4f})") tag = f"{args.model}_single" if args.include_single else args.model out = config.RESULTS_DIR / f"cv_{args.folds}fold__{tag}.json" out.write_text(json.dumps({ "folds": args.folds, "model": args.model, "include_single_label": args.include_single, "n_records": len(records), "n_documents": len(set(groups)), "per_fold": fold_rows, "summary": summary, "ensemble_vs_tfidf_gap": { "per_fold": [round(float(g), 4) for g in gap], "mean": round(float(gap.mean()), 4), "std": round(float(gap.std(ddof=1)), 4), "min": round(float(gap.min()), 4), }, }, indent=2), encoding="utf-8") print(f"\n -> {out.relative_to(config.REPO_ROOT)} " f"[total {(time.time() - t0) / 60:.1f} min]") if __name__ == "__main__": main()