| |
| """Element-head evaluation on held-out sites: does each image get a rendering Haiku judged acceptable? |
| |
| Compares the v2 element head with what v1 shipped for the same images (its 5-class image head for the first 32 |
| images of a page, the hand-written runtime rules for everything else), per kind of image. |
| """ |
| import argparse, collections, gzip, importlib.util, json, pickle, sys |
| from pathlib import Path |
| import numpy as np |
| import torch |
|
|
| from model import NadirNet, TREAT |
| from train import collate, prep, element_logits |
| import safety |
|
|
| V1_MAP = {"keep": 0, "dim": 0, "invert": 1, "light": 1, "plate": 1, "clear": 3} |
|
|
|
|
| def v1_decision(t, im, surf_L): |
| """v1 runtime: model treatment (or 'keep' for images it never saw) through the hand-written rules.""" |
| s = safety.safe_treatment(t, im, surf_L) |
| if s.startswith("adapt:"): |
| return 2 |
| return V1_MAP.get(s, 0) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--crawl", action="append", required=True, help="TAG=CRAWL_DIR") |
| ap.add_argument("--v1-ckpt", default="") |
| ap.add_argument("--v1-src", default="", help="v1 model.py (its architecture)") |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--rows", default="", help="also write every element (probabilities, labels) as jsonl") |
| args = ap.parse_args() |
| crawls = dict(c.split("=", 1) for c in args.crawl) |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| ck = torch.load(args.ckpt, map_location="cpu") |
| T = float(ck.get("temperature", 1.0)) |
| m = NadirNet(**ck["cfg"]); m.load_state_dict(ck["model"]); m.eval().to(dev) |
| v1 = None |
| if args.v1_ckpt and args.v1_src: |
| spec = importlib.util.spec_from_file_location("model_v1", args.v1_src) |
| mv1 = importlib.util.module_from_spec(spec); spec.loader.exec_module(mv1) |
| c1 = torch.load(args.v1_ckpt, map_location="cpu") |
| v1 = mv1.NadirNet(**c1["cfg"]); v1.load_state_dict(c1["model"]); v1.eval().to(dev) |
| data = pickle.load(open(args.data, "rb")) |
| val = [prep(s) for s in data["val"] if len(s.get("E_ut", []))] |
| rows = [] |
| with torch.no_grad(): |
| for s in val: |
| b = collate([s], dev) |
| lab, _, ctx = m(b["tok_f"], b["role"], b["lab"], b["img_f"], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"]) |
| tl, rl = element_logits(m, b, lab.float(), ctx) |
| k = int(b["e_mask"][0].sum()) |
| p = torch.softmax(tl[0, :k] / T, -1).cpu().numpy() |
| role = rl[0, :k].argmax(-1).cpu().numpy() |
| rprob = torch.softmax(rl[0, :k].float() / T, -1).max(-1).values.cpu().numpy() |
| v1t = None |
| if v1 is not None and len(s["IF"]): |
| _, _, il = v1(b["tok_f"], b["role"], b["lab"], b["img_f"][..., :39], b["thumbs"], b["page"], b["src"], b["rel"], b["mask"]) |
| v1t = il[0].float().argmax(-1).cpu().numpy() |
| L = json.load(gzip.open(Path(crawls[s["crawl"]]) / "pages" / s["id"] / "light.json.gz")) |
| n = len(s["role"]) |
| outL = lab[0, :n, 0].float().cpu().numpy() |
| for q in range(k): |
| j = int(s["E_idx"][q]) |
| im = L["images"][j] |
| u = im.get("underTok", -1) |
| surf_L = float(outL[u]) if 0 <= u < n else None |
| ud = [round(float(x), 4) for x in lab[0, u].float().cpu().numpy()] if 0 <= u < n else None |
| old = ["keep", "dim", "invert", "light", "plate"][int(v1t[j])] if (v1t is not None and j < 32 and j < len(v1t)) else "keep" |
| soft = s["E_soft"][q] |
| rows.append({"bucket": s["E_bucket"][q], "kind": s["E_kind"][q], "soft": soft.tolist(), "best": int(np.argmax(soft)), |
| "v2": int(np.argmax(p[q])), "conf": float(p[q].max()), "p": [round(float(x), 4) for x in p[q]], "role_pred": int(role[q]), "role": int(s["E_role"][q]), |
| "v1": v1_decision(old, im, surf_L), "v1_seen": bool(v1t is not None and j < 32), "nlab": int(s["E_nlab"][q]), |
| **({"rp": round(float(rprob[q]), 4), "f": [round(float(x), 4) for x in s["E_f"][q]], "ud": ud, "ul": [round(float(x), 4) for x in s["E_ul"][q]]} if args.rows else {})}) |
| def score(rs, key): |
| if not rs: |
| return {} |
| acc = np.mean([r[key] == r["best"] for r in rs]) |
| ok = np.mean([r["soft"][r[key]] > 0 for r in rs]) |
| return {"n": len(rs), "acc": round(float(acc), 4), "acceptable": round(float(ok), 4)} |
| rep = {"temperature": T, "overall": {"v2": score(rows, "v2"), "v1": score(rows, "v1")}, "by_bucket": {}, "by_kind": {}} |
| for key in ("bucket", "kind"): |
| groups = collections.defaultdict(list) |
| for r in rows: |
| groups[r[key]].append(r) |
| rep["by_" + key] = {g: {"v2": score(rs, "v2"), "v1": score(rs, "v1")} for g, rs in sorted(groups.items())} |
| |
| bins = collections.defaultdict(list) |
| for r in rows: |
| bins[min(9, int(r["conf"] * 10))].append(r["v2"] == r["best"]) |
| rep["calibration"] = {f"{b / 10:.1f}": {"n": len(v), "acc": round(float(np.mean(v)), 3)} for b, v in sorted(bins.items())} |
| rr = [r for r in rows if r["role"] >= 0] |
| rep["role_acc"] = round(float(np.mean([r["role_pred"] == r["role"] for r in rr])), 4) if rr else None |
| |
| rep["n_double_labelled"] = sum(r["nlab"] > 1 for r in rows) |
| json.dump(rep, open(args.out, "w"), indent=1) |
| if args.rows: |
| with open(args.rows, "w") as f: |
| for r in rows: |
| f.write(json.dumps(r) + "\n") |
| print(json.dumps({k: rep[k] for k in ("overall", "role_acc", "temperature")}, indent=1)) |
| print(json.dumps(rep["by_bucket"], indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
