File size: 5,963 Bytes
864f427 | 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 104 105 106 107 108 109 110 111 | #!/usr/bin/env python3
"""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 # the logo adaptation path
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())}
# calibration: confidence vs accuracy
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
# Haiku's own consistency where a second, independent label exists
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()
|