code/mmjev_eval.py

"""Benchmarks for MM-Jev (run in the Colab kernel; `jev`, `DATA` in globals).

Accuracy:  typed-decisions (Laya notebook scoring), BTZSC AG News / DAIR Emotion / Banking77 (jev-benchmarks pilot-v1
protocol), MASSIVE scenario en/ko + intent en, XNLI en, JevBench public (easy / standard / hard), and our multimodal
held-out sets. Latency: batch-1 CUDA-graph path, p50 / p95, end-to-end (towers included for raw media).
"""
import json, math, time
import numpy as np, torch
from mmjev import FastConfig, Seg, options_of, set_ffn_width, QTYPES

LOGFILE = "/content/eval.log"


def say(*a):
    with open(LOGFILE, "a") as f:
        print(time.strftime("%H:%M:%S"), *a, file=f)


def ece(conf, corr, bins=10):
    conf, corr = np.asarray(conf), np.asarray(corr)
    e, edges = 0.0, np.linspace(0, 1, bins + 1)
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.any():
            e += m.mean() * abs(conf[m].mean() - corr[m].mean())
    return float(e)


def _T(temps, t, k):
    from mmjev import k_bucket
    temps = temps or {}
    return temps.get(f"{t}:{k_bucket(k)}", temps.get(t, 1.0))


@torch.no_grad()
def predict(jev, recs, fc, width, temps=None, bs=8):
    """-> list (per record) of list (per question) of calibrated probability vectors."""
    set_ffn_width(jev.lm, width)
    jev.eval()
    out = []
    order = sorted(range(len(recs)), key=lambda i: sum(len(options_of(q)[0]) for q in recs[i]["qs"]))
    res = [None] * len(recs)
    def _go(idx):
        batch = [recs[i] for i in idx]
        try:
            plans = [jev.plan(jev.encode_state(r["state"], fc), r["qs"], fc) for r in batch]
            lg = jev.run(jev.pack(plans), fc)
        except torch.OutOfMemoryError:
            torch.cuda.empty_cache()
            if len(idx) == 1:
                raise
            h = len(idx) // 2
            _go(idx[:h]); _go(idx[h:]); return
        for i, r, per_q in zip(idx, batch, lg):
            res[i] = [torch.softmax(l.float() / _T(temps, q["type"], len(l)), -1).cpu().numpy()
                      for q, l in zip(r["qs"], per_q)]

    for s in range(0, len(order), bs):
        _go(order[s:s + bs])
    set_ffn_width(jev.lm, None)
    return res


def basic_metrics(recs, probs):
    acc, conf, corr, brier, nll = [], [], [], [], []
    for r, pq in zip(recs, probs):
        for q, p, y in zip(r["qs"], pq, r["ys"]):
            c = float(int(p.argmax()) == y)
            acc.append(c); conf.append(float(p.max())); corr.append(c)
            oh = np.zeros_like(p); oh[y] = 1
            brier.append(float(((p - oh) ** 2).sum())); nll.append(float(-math.log(max(p[y], 1e-12))))
    return dict(n=len(acc), acc=float(np.mean(acc)), brier=float(np.mean(brier)), nll=float(np.mean(nll)),
                ece=ece(conf, corr))


def typed_metrics(recs, probs):
    """Exactly the Laya fine-tuning notebook scoring (accuracy, soft acc, Brier, ECE, score MAE in levels)."""
    accs, soft, brier, maes, confs, corrs = [], [], [], [], [], []
    by_type, by_wf = {}, {}
    for r, pq in zip(recs, probs):
        for qid, q, p in zip(r["qnames"], r["qs"], pq):
            g, t = r["gold"][qid], q["type"]
            if t == "choice":
                keys = list(q["criteria"])
                gp = np.array([g["probabilities"].get(k, 1e-6) for k in keys]); gp /= gp.sum()
                c = float(keys[int(p.argmax())] == str(g["label"]))
                soft.append(float((p * gp).sum())); brier.append(float(((p - gp) ** 2).sum()))
                confs.append(float(p.max()))
            elif t == "noul":
                pv = float(p[1]); gv = float(g.get("noul", g.get("probabilities", {}).get("true", 0.5)))
                c = float(("true" if pv >= 0.5 else "false") == str(g["label"]).lower())
                pd, gd = np.array([1 - pv, pv]), np.array([1 - gv, gv])
                soft.append(float((pd * gd).sum())); brier.append(float(((pd - gd) ** 2).sum()))
                confs.append(max(pv, 1 - pv))
            else:
                ev = float((p * np.arange(len(p))).sum())
                maes.append(abs(ev - float(g.get("score", 0.0))))
                c = float(int(p.argmax()) == int(g.get("label", round(g.get("score", 0)))))
                confs.append(float(p.max()))
            accs.append(c); corrs.append(c)
            by_type.setdefault(t, []).append(c); by_wf.setdefault(r["workflow"], []).append(c)
    return dict(n=len(accs), acc=float(np.mean(accs)), soft_acc=float(np.mean(soft)), brier=float(np.mean(brier)),
                ece=ece(confs, corrs, 15), score_mae=float(np.mean(maes)),
                by_type={k: round(float(np.mean(v)), 3) for k, v in by_type.items()},
                by_workflow={k: round(float(np.mean(v)), 3) for k, v in by_wf.items()})


def accuracy_suite(jev, DATA, fc, width, temps=None, tag="", tasks=None):
    ev = [r for v in list(DATA.values()) for r in v if r["split"] == "eval"]
    tasks = tasks or sorted(set(r["task"] for r in ev))
    out = {}
    for t in tasks:
        rs = [r for r in ev if r["task"] == t]
        t0 = time.time()
        pr = predict(jev, rs, fc, width, temps)
        m = typed_metrics(rs, pr) if t == "typed_decisions" else basic_metrics(rs, pr)
        m["modality"] = rs[0]["modality"]
        out[t] = m
        say(f"[{tag}] {t:22s} n={m['n']:5d} acc {m['acc']:.3f} ece {m['ece']:.3f} ({time.time() - t0:.0f}s)")
    return out


# ------------------------------------------------------------------------------------------ latency
def _time(f, n):
    ts = []
    for _ in range(n):
        torch.cuda.synchronize(); t0 = time.perf_counter(); f(); torch.cuda.synchronize()
        ts.append((time.perf_counter() - t0) * 1000)
    return ts


@torch.no_grad()
def latency_suite(jev, DATA, fc, width, n_rep=3, max_items=40):
    """End-to-end decide() latency, batch 1, CUDA-graph path. Raw media -> towers are included."""
    set_ffn_width(jev.lm, width); jev._graphs = {}
    ev = [r for v in DATA.values() for r in v if r["split"] == "eval"]
    groups = {}
    for r in ev:
        if r.get("raw") is not None:
            groups.setdefault(f"{r['modality']}:{r['task']}", []).append((r["raw"], r["qs"]))
    for t in ("typed_decisions", "btzsc_agnews", "btzsc_banking77", "jevbench_hard"):
        rs = [r for r in ev if r["task"] == t][:max_items]
        groups[f"text:{t}"] = [(r["state"], r["qs"]) for r in rs]
    res = {}
    for g, items in groups.items():
        items = items[:max_items]
        for st, qs in items:                    # warm-up: capture every bucket graph first
            jev.decide(st, qs, fc=fc, graph=True)
        ts = []
        for st, qs in items:
            ts += _time(lambda: jev.decide(st, qs, fc=fc, graph=True), n_rep)
        nq = float(np.mean([len(qs) for _, qs in items]))
        res[g] = dict(n=len(items), questions_per_call=nq, p50_ms=float(np.percentile(ts, 50)),
                      p95_ms=float(np.percentile(ts, 95)), per_question_ms=float(np.percentile(ts, 50) / nq))
        say(f"  {g:32s} p50 {res[g]['p50_ms']:7.1f} ms  p95 {res[g]['p95_ms']:7.1f} ms  ({nq:.1f} q/call)")
    set_ffn_width(jev.lm, None)
    return res


@torch.no_grad()
def questions_scaling(jev, fc, width, text=None):
    """Laya-style: 1, 5, 10, 50 questions on one state in one call."""
    set_ffn_width(jev.lm, width)
    text = text or ("Hi, we were billed twice for March. Please refund the duplicate today or we will cancel our plan.")
    base_qs = [{"type": "choice", "instructions": "Which department should handle this?",
                "criteria": {"billing": "invoices, payments, refunds", "technical": "bugs, outages",
                             "other": "everything else"}},
               {"type": "score", "instructions": "How urgent is this?", "criteria": ["not urgent", "soon", "blocking"]},
               {"type": "noul", "instructions": "Does the user threaten to cancel or leave?"},
               {"type": "noul", "instructions": "Does the user explicitly request a refund?"},
               {"type": "noul", "instructions": "Is the user angry?"}]
    out = {}
    for n in (1, 5, 10, 50):
        qs = [dict(base_qs[i % 5], instructions=base_qs[i % 5]["instructions"] + (f" (#{i})" if i >= 5 else ""))
              for i in range(n)]
        jev.decide(text, qs, fc=fc, graph=True)
        ts = _time(lambda: jev.decide(text, qs, fc=fc, graph=True), 10)
        out[n] = float(np.median(ts))
        say(f"  {n:3d} questions/call: {out[n]:6.1f} ms ({out[n] / n:5.1f} ms/q)")
    set_ffn_width(jev.lm, None)
    return out


@torch.no_grad()
def tree_vs_naive(jev, DATA, fc, width, n=30):
    """openjev-style one pass per option vs the packed option tree (same model, eager path)."""
    set_ffn_width(jev.lm, width)
    rs = [r for r in DATA["typed"] if r["split"] == "eval"][:n]
    rows = {}
    for name, fn in (("tree (1 pass)", lambda p: jev.run(jev.pack([p]), fc)),
                     ("naive (1 pass / option)", lambda p: [jev.run(jev.pack([q]), fc) for q in jev.isolate(p)])):
        ts = []
        for r in rs:
            p = jev.plan(jev.encode_state(r["state"], fc), r["qs"], fc)
            fn(p)
            ts += _time(lambda: fn(p), 2)
        rows[name] = float(np.median(ts))
        say(f"  {name:26s} {rows[name]:7.1f} ms per typed-decisions case")
    set_ffn_width(jev.lm, None)
    return rows