| """MNLI validation-matched accuracy: CPU HF reference vs the served TT model (/classify). |
| |
| python scripts/eval_mnli.py --n 500 [--url http://localhost:8000] [--skip-cpu] [--cpu-only] |
| |
| Writes results/mnli_<n>.json with per-example labels/probs for both sides. |
| """ |
| import argparse |
| import json |
| import os |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
| import torch |
|
|
| TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}" |
| |
| MNLI_TO_OPENJEV = {0: 1, 1: 2, 2: 0} |
| CKPT = os.environ.get("HF_MODEL", "/home/ttuser/experiments/openjev/weights/openjev-qwen3.5-4b-nli") |
|
|
|
|
| def load_mnli(n, seed=0): |
| from datasets import load_dataset |
|
|
| ds = load_dataset("nyu-mll/multi_nli", split="validation_matched") |
| ds = ds.shuffle(seed=seed).select(range(n)) |
| texts = [TEMPLATE.format(premise=p.strip(), hypothesis=h.strip()) for p, h in zip(ds["premise"], ds["hypothesis"])] |
| gold = [MNLI_TO_OPENJEV[int(l)] for l in ds["label"]] |
| return texts, gold |
|
|
|
|
| def post(url, payload, timeout=1800): |
| req = urllib.request.Request(url, data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}) |
| with urllib.request.urlopen(req, timeout=timeout) as r: |
| return json.loads(r.read()) |
|
|
|
|
| def run_cpu(texts, bs=16): |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
| tok = AutoTokenizer.from_pretrained(CKPT) |
| tok.padding_side = "right" |
| model = AutoModelForSequenceClassification.from_pretrained(CKPT, dtype=torch.bfloat16).eval() |
| model.config.get_text_config().pad_token_id = tok.pad_token_id |
| out = [] |
| t0 = time.time() |
| with torch.no_grad(): |
| for i in range(0, len(texts), bs): |
| enc = tok(texts[i : i + bs], truncation=True, max_length=4096, padding=True, return_tensors="pt") |
| out.append(torch.softmax(model(**enc).logits.float(), -1)) |
| print(f" cpu {min(i + bs, len(texts))}/{len(texts)} ({time.time() - t0:.0f}s)", flush=True) |
| return torch.cat(out, 0) |
|
|
|
|
| def run_server(texts, url, model_name, bs=32): |
| out = [] |
| t0 = time.time() |
| for i in range(0, len(texts), bs): |
| res = post(f"{url}/classify", {"model": model_name, "input": texts[i : i + bs]}) |
| out.extend([d["probs"] for d in res["data"]]) |
| print(f" server {min(i + bs, len(texts))}/{len(texts)} ({time.time() - t0:.1f}s)", flush=True) |
| return torch.tensor(out, dtype=torch.float32), time.time() - t0 |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--n", type=int, default=500) |
| ap.add_argument("--url", default="http://localhost:8000") |
| ap.add_argument("--model", default="AlexWortega/openjev") |
| ap.add_argument("--skip-cpu", action="store_true") |
| ap.add_argument("--cpu-only", action="store_true") |
| ap.add_argument("--batch", type=int, default=32) |
| args = ap.parse_args() |
|
|
| texts, gold = load_mnli(args.n) |
| gold_t = torch.tensor(gold) |
| res_dir = Path(__file__).resolve().parent.parent / "results" |
| res_dir.mkdir(exist_ok=True) |
| out_path = res_dir / f"mnli_{args.n}.json" |
| result = json.load(open(out_path)) if out_path.exists() else {"n": args.n, "texts": texts, "gold": gold} |
|
|
| if not args.skip_cpu: |
| cpu = run_cpu(texts) |
| result["cpu_probs"] = cpu.tolist() |
| cpu = torch.tensor(result["cpu_probs"]) if "cpu_probs" in result else None |
| if cpu is not None: |
| print(f"CPU bf16 accuracy: {(cpu.argmax(-1) == gold_t).float().mean():.4f}") |
|
|
| if not args.cpu_only: |
| tt, dt = run_server(texts, args.url, args.model, bs=args.batch) |
| result["tt_probs"] = tt.tolist() |
| result["tt_seconds"] = dt |
| print(f"TT accuracy: {(tt.argmax(-1) == gold_t).float().mean():.4f} ({dt:.1f}s, {dt / args.n * 1e3:.0f} ms/example)") |
| if cpu is not None: |
| print(f"TT vs CPU label agreement: {(tt.argmax(-1) == cpu.argmax(-1)).float().mean():.4f}; " |
| f"max|dprob| {(tt - cpu).abs().max():.4f} mean|dprob| {(tt - cpu).abs().mean():.4f}") |
| json.dump(result, open(out_path, "w")) |
| print("saved", out_path) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
