File size: 4,583 Bytes
b1cf75e 0052e99 b1cf75e 0052e99 b1cf75e 0052e99 b1cf75e | 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 | """Record one model answering one prompt through llama-server, with a timestamp for every streamed chunk.
python3 scripts/race_record.py --server BIN --model GGUF --name "Dory-Beam-4B" --loops 2 \
--problem math500-123 --seed 0 --out race/dory-beam.json [--server-args "..."]
Starts llama-server, sends the problem as one user message (temperature 1.0, top_p 0.95, fixed seed), streams the
reply with raw <think> tags kept in the content, and writes {name, loops, problem, prompt, gold, events: [[t, text]],
completion_tokens, seconds, timings}. t is seconds since the request was sent. Stops the server afterwards.
"""
import argparse
import json
import shlex
import subprocess
import time
import urllib.request
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
PORT = 8099
def wait_ready(proc, timeout=600):
t0 = time.time()
while time.time() - t0 < timeout:
if proc.poll() is not None:
raise RuntimeError("llama-server exited during startup")
try:
with urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=2) as r:
if r.status == 200:
return
except Exception:
pass
time.sleep(1)
raise TimeoutError("llama-server did not become ready")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--server", required=True)
ap.add_argument("--model", required=True)
ap.add_argument("--name", required=True)
ap.add_argument("--loops", type=int, default=2)
ap.add_argument("--problem", required=True)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--max-tokens", type=int, default=16000)
ap.add_argument("--server-args", default="")
ap.add_argument("--effort", default="", help="reasoning_effort chat template argument (Dory-Beam-4B v2)")
ap.add_argument("--out", required=True)
args = ap.parse_args()
prob = next(json.loads(l) for l in open(ROOT / "data" / "evalsets.jsonl") if json.loads(l)["id"] == args.problem)
cmd = [args.server, "-m", args.model, "--host", "127.0.0.1", "--port", str(PORT), "-ngl", "999",
"--ctx-size", "20480", "--parallel", "1", "--no-context-shift", "--jinja", "--reasoning-format", "none",
"--override-kv", f"dory.num_loops=int:{args.loops},dory.recurrent_kv_cache_mode=str:per_loop",
*shlex.split(args.server_args)]
log = open(Path(args.out).with_suffix(".server.log"), "w")
proc = subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)
try:
wait_ready(proc)
body = {"messages": [{"role": "user", "content": prob["prompt"]}], "stream": True, "temperature": 1.0,
"top_p": 0.95, "top_k": 0, "min_p": 0, "seed": args.seed, "max_tokens": args.max_tokens,
"cache_prompt": False, "timings_per_token": False}
if args.effort:
body["chat_template_kwargs"] = {"reasoning_effort": args.effort}
req = urllib.request.Request(f"http://127.0.0.1:{PORT}/v1/chat/completions", json.dumps(body).encode(),
{"Content-Type": "application/json"})
events, timings, t0 = [], None, time.time()
with urllib.request.urlopen(req, timeout=7200) as resp:
for raw in resp:
line = raw.decode("utf-8", "replace").strip()
if not line.startswith("data: ") or line == "data: [DONE]":
continue
chunk = json.loads(line[6:])
timings = chunk.get("timings", timings)
for ch in chunk.get("choices", []):
text = (ch.get("delta") or {}).get("content")
if text:
events.append([round(time.time() - t0, 3), text])
seconds = time.time() - t0
finally:
proc.terminate()
proc.wait(timeout=60)
text = "".join(e[1] for e in events)
out = {"name": args.name, "model_file": Path(args.model).name, "loops": args.loops, "effort": args.effort,
"problem": args.problem,
"prompt": prob["prompt"], "gold": prob["answer"], "seed": args.seed, "events": events,
"completion_tokens": (timings or {}).get("predicted_n"), "seconds": round(seconds, 2), "timings": timings,
"text": text}
Path(args.out).parent.mkdir(parents=True, exist_ok=True)
Path(args.out).write_text(json.dumps(out, ensure_ascii=False))
print(f"{args.name}: {out['completion_tokens']} tokens in {seconds:.1f}s -> {args.out}")
if __name__ == "__main__":
main()
|