| """Test-set metrics + latency for NCM generator, AR baseline, intent classifier.""" |
| import os, sys, json, time, argparse |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from paths import WORK |
| os.environ.setdefault("MFMT", "C") |
| import numpy as np, mlx.core as mx |
| from transformers import AutoTokenizer |
| from gen_mlx import * |
| from ar_mlx import * |
| from massive_eval import score |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ncm", default=WORK+"/ncm_pilot"); ap.add_argument("--ar", default=WORK+"/ar_pilot"); ap.add_argument("--cls", default=WORK+"/cls_pilot") |
| ap.add_argument("--split", default="test"); ap.add_argument("--out", default=WORK+"/massive_de/results.json"); ap.add_argument("--P", type=int, default=40) |
| ap.add_argument("--nreq", type=int, default=200) |
| a = ap.parse_args() |
| EOS, PAD = 151645, 151643 |
| tok = AutoTokenizer.from_pretrained(WORK+"/tok/Qwen3-0.6B"); fc = FrozenCodec() |
| D = [json.loads(l) for l in open(f"{WORK}/massive_de/{a.split}.jsonl")] |
| def pad(ids, n): ids = ids[:n]; return ids + [PAD] * (n - len(ids)) |
| PT = np.array([pad(tok.encode(d["prefix"])[:a.P - 1] + [EOS], a.P) for d in D], np.int32) |
| def cut(r): |
| r = list(r); r = r[:r.index(EOS)] if EOS in r else r |
| return tok.decode([t for t in r if t != PAD], skip_special_tokens=True) |
| def lat(fn, bs, n): |
| ts = [] |
| rng = np.random.RandomState(0) |
| for i in range(n + 3): |
| idx = rng.randint(0, len(D), bs); t = time.perf_counter(); fn(mx.array(PT[idx])); ts.append((time.perf_counter() - t) * 1000) |
| ts = np.array(ts[3:]); return dict(p50=float(np.percentile(ts, 50)), p95=float(np.percentile(ts, 95)), mean=float(ts.mean())) |
| res = {} |
| |
| cfg = json.load(open(f"{a.ncm}/ckpt.json"))["args"] |
| master = Gen(fc, cfg["P"], cfg["D"], cfg["H"], cfg["nl"], PAD, drop=0.0); master.load_weights(f"{a.ncm}/best.safetensors") |
| g = Gen(fc, cfg["P"], cfg["D"], cfg["H"], cfg["nl"], PAD, drop=0.0).set_compute(mx.float16) |
| from mlx.utils import tree_map |
| g.update(tree_map(lambda t: t.astype(mx.float16), master.parameters())); g.eval(); mx.eval(g.parameters()) |
| def ncm_gen(p): t, _ = generate(g, fc, p); mx.eval(t); return t |
| preds = [] |
| for i in range(0, len(D), 64): preds += [cut(r) for r in ncm_gen(mx.array(PT[i:i + 64])).tolist()] |
| r = score(preds, D); r["NFE"] = g.K; r["params_M"] = count_params(master) / 1e6 |
| r["lat_b1_ms"] = lat(ncm_gen, 1, a.nreq); r["lat_b16_ms"] = lat(ncm_gen, 16, max(a.nreq // 5, 20)); res["ncm"] = r |
| res["ncm_codec_ceiling_exact"] = None |
| json.dump({"preds_ncm": preds[:50]}, open(a.out.replace(".json", "_samples_ncm.json"), "w"), ensure_ascii=False, indent=1) |
| print("NCM", r, flush=True) |
| |
| if os.path.exists(f"{a.ar}/model.safetensors"): |
| voc = json.load(open(f"{a.ar}/vocab.json")); ck = json.load(open(f"{a.ar}/ckpt.json")); inv_v = {i + 3: t for i, t in enumerate(voc["vocab"])} |
| m = ARGen(fc.tok, ck["V"], ck["args"]["P"], ck["args"]["T"], pad_id=PAD, drop=0.0); m.load_weights(f"{a.ar}/best.safetensors") |
| m.set_dtype(mx.float16); m.eval(); mx.eval(m.parameters()) |
| nfe = [] |
| def ar_gen(p): o, n = m.generate(p, 1); mx.eval(o); nfe.append(n); return o |
| def dec(row): |
| row = row[:row.index(1)] if 1 in row else row |
| return tok.decode([inv_v[c] for c in row if c in inv_v], skip_special_tokens=True) |
| preds = [] |
| for i in range(0, len(D), 64): preds += [dec(r) for r in ar_gen(mx.array(PT[i:i + 64])).tolist()] |
| r = score(preds, D); r["NFE_mean_b64"] = float(np.mean(nfe)); r["params_M"] = count_params(m) / 1e6 |
| nfe.clear(); r["lat_b1_ms"] = lat(ar_gen, 1, a.nreq); r["NFE_mean_b1"] = float(np.mean(nfe)); nfe.clear() |
| r["lat_b16_ms"] = lat(ar_gen, 16, max(a.nreq // 5, 20)); r["NFE_mean_b16"] = float(np.mean(nfe)); res["ar"] = r |
| json.dump({"preds_ar": preds[:50]}, open(a.out.replace(".json", "_samples_ar.json"), "w"), ensure_ascii=False, indent=1) |
| print("AR", r, flush=True) |
| |
| if os.path.exists(f"{a.cls}/model.safetensors"): |
| ints = json.load(open(f"{a.cls}/vocab.json"))["ints"] |
| m = IntentCls(fc.tok, len(ints), a.P, pad_id=PAD, drop=0.0); m.load_weights(f"{a.cls}/best.safetensors"); m.set_dtype(mx.float16); m.eval() |
| def cl(p): o = mx.argmax(m(p), -1); mx.eval(o); return o |
| hit = 0 |
| for i in range(0, len(D), 256): hit += sum(ints[j] == d["intent"] for j, d in zip(cl(mx.array(PT[i:i + 256])).tolist(), D[i:i + 256])) |
| res["cls"] = dict(intent_acc=hit / len(D), params_M=count_params(m) / 1e6, NFE=1, lat_b1_ms=lat(cl, 1, a.nreq), lat_b16_ms=lat(cl, 16, max(a.nreq // 5, 20))) |
| print("CLS", res["cls"], flush=True) |
| json.dump(res, open(a.out, "w"), indent=1) |
|
|
