File size: 4,838 Bytes
888eee4 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """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: # network hiccup
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()
|