File size: 4,352 Bytes
210f9af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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())
|