| """Needle-in-a-haystack: Groot vs Jev on the same requests, accuracy and latency per state length. |
| |
| python scripts/needle_eval.py --rows needle_eval.jsonl --model runs/x/final [--jev] [--limit-per-length 50] --out needle.json |
| """ |
| import argparse |
| import json |
| import os |
| import statistics |
| import time |
| from collections import defaultdict |
| from concurrent.futures import ThreadPoolExecutor |
|
|
| from groot1.data import read_jsonl |
|
|
|
|
| def to_question(r): |
| if r["type"] == "choice": |
| crit = {} |
| for o in r["options"]: |
| name, _, desc = o.partition(": ") |
| crit[name] = desc or None |
| return {"type": "choice", "instructions": r["question"], "criteria": crit}, list(crit) |
| return {"type": "noul", "instructions": r["question"]}, ["false", "true"] |
|
|
|
|
| def gold_of(r, names): |
| return names[max(range(len(r["target"])), key=r["target"].__getitem__)] |
|
|
|
|
| def pred_of(ans, kind): |
| if kind == "choice": |
| return ans["choice"] |
| return "true" if ans["noul"] >= 0.5 else "false" |
|
|
|
|
| def run_groot(rows, model, device, overrides=None): |
| from groot1.serve import GrootEngine |
| from groot1.schema import SystemOneRequest |
|
|
| eng = GrootEngine(model, device=device, overrides=overrides) |
| out = [] |
| for r in rows: |
| q, names = to_question(r) |
| req = SystemOneRequest(state=r["state"], questions={"q": q}) |
| t0 = time.perf_counter() |
| resp = eng.answer(req.state, req.questions) |
| dt = (time.perf_counter() - t0) * 1000 |
| a = resp.answers["q"].model_dump() |
| out.append({"ok": pred_of(a, r["type"]) == gold_of(r, names), "ms": dt, "tokens": resp.usage.input_tokens}) |
| return out |
|
|
|
|
| def run_jev(rows, workers=8): |
| import httpx |
|
|
| key = os.environ["TYPESAFE_API_KEY"] |
|
|
| def one(r): |
| q, names = to_question(r) |
| long = r["needle"]["length"] > 30000 |
| headers = {"Authorization": f"Bearer {key}", **({"TypeSafe-Danger-Zone-Long-Context": "true"} if long else {})} |
| body = {"model": "jev-latest", "state": r["state"], "questions": {"q": q}} |
| for attempt in range(4): |
| try: |
| t0 = time.perf_counter() |
| resp = httpx.post("https://api.typesafe.ai/v1/systemone", json=body, headers=headers, timeout=120) |
| dt = (time.perf_counter() - t0) * 1000 |
| if resp.status_code == 429: |
| time.sleep(2 ** attempt) |
| continue |
| if resp.status_code != 200: |
| return {"ok": False, "ms": dt, "error": f"{resp.status_code} {resp.text[:120]}"} |
| a = resp.json()["answers"]["q"] |
| return {"ok": pred_of(a, r["type"]) == gold_of(r, names), "ms": dt, "tokens": resp.json()["usage"]["input_tokens"]} |
| except Exception as e: |
| err = str(e) |
| time.sleep(2 ** attempt) |
| return {"ok": False, "ms": None, "error": err} |
|
|
| with ThreadPoolExecutor(workers) as ex: |
| return list(ex.map(one, rows)) |
|
|
|
|
| def summarize(rows, res): |
| by = defaultdict(list) |
| for r, x in zip(rows, res): |
| by[r["needle"]["length"]].append(x) |
| out = {} |
| for L in sorted(by): |
| xs = by[L] |
| ms = [x["ms"] for x in xs if x.get("ms") is not None] |
| out[L] = {"n": len(xs), "acc": round(sum(x["ok"] for x in xs) / len(xs), 3), "p50_ms": round(statistics.median(ms), 1) if ms else None, |
| "errors": sum(1 for x in xs if "error" in x)} |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--rows", required=True) |
| ap.add_argument("--model", default=None) |
| ap.add_argument("--device", default="cuda") |
| ap.add_argument("--jev", action="store_true") |
| ap.add_argument("--limit-per-length", type=int, default=None) |
| ap.add_argument("--out", default=None) |
| ap.add_argument("--set", action="append", default=[], help="config override key=value") |
| args = ap.parse_args() |
| rows, seen = [], defaultdict(int) |
| for r in read_jsonl(args.rows): |
| L = r["needle"]["length"] |
| if args.limit_per_length and seen[L] >= args.limit_per_length: |
| continue |
| seen[L] += 1 |
| rows.append(r) |
| report = {} |
| if args.model: |
| from groot1.config import parse_overrides |
|
|
| report["groot"] = summarize(rows, run_groot(rows, args.model, args.device, parse_overrides(args.set))) |
| print("groot", json.dumps(report["groot"]), flush=True) |
| if args.jev: |
| res = run_jev(rows) |
| report["jev"] = summarize(rows, res) |
| errs = [x["error"] for x in res if "error" in x][:3] |
| print("jev", json.dumps(report["jev"]), ("errors: " + str(errs)) if errs else "", flush=True) |
| if args.out: |
| open(args.out, "w").write(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
