| """Choose the arm of the §4.4 screen (pure Python). Run as ``python -m training.screen RUN... --data BUNDLE``. |
| |
| Arms: B-last, I-last, B-mean, I-mean (checkpoint base/it x pooling last/mean), 1 epoch each. Each run's best val |
| predictions (``RUN/predictions/val.jsonl``) are scored on val ``eval`` posts: |
| - primary: mean val log-loss over relevance, direction, commercial (lower is better); |
| - secondary: direction macro-F1 (gated argmax, uncalibrated). |
| Rule (classifier-plan §4.4, made exact here): B-last is the default. Another arm beats it if its primary is more |
| than 2% lower than B-last's, or if its primary is within 2% of B-last's and its secondary is at least 1 point |
| higher. Among arms that beat B-last the lowest primary wins (then the higher secondary, then arm order). |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| from .calibrate import nll, outputs, pair_up |
| from .data import load_jsonl, load_split |
| from .taxonomy import Taxonomy |
|
|
| ARMS = ("B-last", "I-last", "B-mean", "I-mean") |
| DEFAULT = "B-last" |
| PRIMARY_FIELDS = ("relevance", "direction", "commercial") |
| REL_TOL = 0.02 |
| F1_MARGIN = 0.01 |
|
|
|
|
| def macro_f1(gold: list[str], pred: list[str], classes: list[str]) -> float | None: |
| scores = [] |
| for c in classes: |
| tp = sum(1 for g, p in zip(gold, pred) if g == c and p == c) |
| support = sum(1 for g in gold if g == c) |
| predicted = sum(1 for p in pred if p == c) |
| if support + predicted: |
| scores.append(2 * tp / (support + predicted)) |
| return sum(scores) / len(scores) if scores else None |
|
|
|
|
| def score(preds: list[dict], rows: list[dict], tax: Taxonomy) -> dict: |
| pairs = [(p, r) for p, r in pair_up(preds, rows) if r.get("eval")] |
| losses = [nll(pairs, tax, f, 1.0) for f in PRIMARY_FIELDS] |
| losses = [v for v in losses if v is not None] |
| gold, pred = [], [] |
| for p, r in pairs: |
| if "direction" in r.get("mask", ()) or r["targets"].get("direction") is None: |
| continue |
| gold.append(r["targets"]["direction"]) |
| pred.append(outputs(p["logits"], tax)["labels"]["direction"]) |
| return {"primary": sum(losses) / len(losses) if losses else None, |
| "secondary": macro_f1(gold, pred, tax.fields["direction"]), "eval_rows": len(pairs)} |
|
|
|
|
| def choose(scores: dict[str, dict]) -> dict: |
| """Apply the rule to {arm: {"primary", "secondary"}}; returns the winner and the reasons.""" |
| if DEFAULT not in scores or scores[DEFAULT]["primary"] is None: |
| raise ValueError(f"the {DEFAULT} arm is required (it is the default and the reference)") |
| ref = scores[DEFAULT] |
| beats = [] |
| for arm, s in scores.items(): |
| if arm == DEFAULT or s["primary"] is None: |
| continue |
| rel = s["primary"] / ref["primary"] - 1.0 |
| better_primary = rel < -REL_TOL |
| better_secondary = abs(rel) <= REL_TOL and s["secondary"] is not None and ref["secondary"] is not None \ |
| and s["secondary"] >= ref["secondary"] + F1_MARGIN |
| if better_primary or better_secondary: |
| beats.append(arm) |
| order = {a: i for i, a in enumerate(ARMS)} |
| if beats: |
| beats.sort(key=lambda a: (scores[a]["primary"], -(scores[a]["secondary"] or 0.0), order.get(a, 99))) |
| winner = beats[0] |
| else: |
| winner = DEFAULT |
| return {"winner": winner, "beats_default": beats, "scores": scores} |
|
|
|
|
| def main(argv=None) -> int: |
| p = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| p.add_argument("runs", nargs="+", help="Run directories (their config.json names the arm)") |
| p.add_argument("--data", required=True) |
| p.add_argument("--out") |
| args = p.parse_args(argv) |
| tax = Taxonomy.load(args.data) |
| rows = load_split(args.data, "val") |
| scores = {} |
| for run in map(Path, args.runs): |
| cfg = json.loads((run / "config.json").read_text()) |
| arm = cfg.get("arm") or run.name |
| preds = load_jsonl(run / "predictions" / "val.jsonl") |
| by_id = {r["id"] for r in preds} |
| scores[arm] = score(preds, [r for r in rows if r["id"] in by_id], tax) |
| scores[arm]["run"] = str(run) |
| result = choose(scores) |
| text = json.dumps(result, indent=1) |
| if args.out: |
| Path(args.out).write_text(text + "\n") |
| print(text) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
