| """Generate Beam reasoning traces for a prompt pool, verify answers, append to JSONL. |
| |
| Resumable: (id, effort) pairs already in --out are skipped. Stops launching new |
| requests when the daily output-token or request budget (from Reflection's rate |
| limit headers) drops below the reserve. Reads the key from REFLECTION_API_KEY. |
| """ |
|
|
| import argparse |
| import asyncio |
| import json |
| import os |
| import random |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import httpx |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| from verify import check |
|
|
| BASE = "https://api.reflection.ai/openai/v1" |
| MODEL = "Beam-1.0-501B-A23B" |
|
|
|
|
| class Budget: |
| def __init__(self, reserve_out, reserve_req): |
| self.reserve_out = reserve_out |
| self.reserve_req = reserve_req |
| self.out_left = None |
| self.req_left = None |
| self.reset_s = None |
|
|
| def update(self, headers): |
| def num(name): |
| v = headers.get(name) |
| return int(float(v)) if v is not None else None |
|
|
| out_left = num("x-ratelimit-remaining-output-tokens-day") |
| req_left = num("x-ratelimit-remaining-requests-day") |
| if out_left is not None: |
| self.out_left = out_left |
| if req_left is not None: |
| self.req_left = req_left |
| r = headers.get("x-ratelimit-reset-output-tokens-day") |
| if r: |
| self.reset_s = float(r.rstrip("s")) |
|
|
| def exhausted(self): |
| return (self.out_left is not None and self.out_left < self.reserve_out) or ( |
| self.req_left is not None and self.req_left < self.reserve_req |
| ) |
|
|
|
|
| class Pacer: |
| """Keep under the per-minute request limit.""" |
|
|
| def __init__(self, per_minute): |
| self.interval = 60.0 / per_minute |
| self.next = time.monotonic() |
| self.lock = asyncio.Lock() |
|
|
| async def wait(self): |
| async with self.lock: |
| now = time.monotonic() |
| if self.next > now: |
| await asyncio.sleep(self.next - now) |
| self.next = max(now, self.next) + self.interval |
|
|
|
|
| async def call(client, item, args, budget, pacer): |
| body = { |
| "model": MODEL, |
| "messages": [{"role": "user", "content": item["prompt"]}], |
| "max_tokens": args.max_tokens, |
| "reasoning_effort": args.effort, |
| "temperature": args.temperature, |
| "top_p": args.top_p, |
| } |
| for attempt in range(6): |
| await pacer.wait() |
| t0 = time.time() |
| try: |
| r = await client.post(f"{BASE}/chat/completions", json=body) |
| except httpx.HTTPError as e: |
| err = f"transport: {e!r}" |
| await asyncio.sleep(min(60, 2**attempt + random.random())) |
| continue |
| budget.update(r.headers) |
| if r.status_code == 200: |
| return r.json(), time.time() - t0, None |
| err = f"HTTP {r.status_code}: {r.text[:300]}" |
| if r.status_code in (429, 500, 502, 503, 504): |
| if budget.exhausted(): |
| return None, 0, err |
| wait = r.headers.get("retry-after") |
| await asyncio.sleep(float(wait) if wait else min(90, 2 ** (attempt + 1))) |
| continue |
| return None, 0, err |
| return None, 0, err |
|
|
|
|
| async def worker(name, queue, client, args, budget, pacer, out_f, stats): |
| while True: |
| item = await queue.get() |
| try: |
| if budget.exhausted() or stats["stop"]: |
| continue |
| resp, latency, err = await call(client, item, args, budget, pacer) |
| rec = { |
| "id": item["id"], |
| "effort": args.effort, |
| "domain": item["domain"], |
| "source": item["source"], |
| "prompt": item["prompt"], |
| "gold": item["answer"], |
| "answer_type": item["answer_type"], |
| "difficulty": item.get("difficulty"), |
| "topic": item.get("topic"), |
| "time": time.strftime("%Y-%m-%dT%H:%M:%S"), |
| } |
| if resp is None: |
| rec.update({"error": err}) |
| stats["errors"] += 1 |
| if "HTTP 4" in (err or "") and "429" not in err: |
| print(f"[{name}] non-retryable error: {err}", flush=True) |
| else: |
| choice = resp["choices"][0] |
| msg = choice["message"] |
| usage = resp.get("usage") or {} |
| rec.update( |
| { |
| "reasoning": msg.get("reasoning_content") or "", |
| "content": msg.get("content") or "", |
| "finish_reason": choice.get("finish_reason"), |
| "usage": usage, |
| "latency_s": round(latency, 2), |
| "response_id": resp.get("id"), |
| } |
| ) |
| ok, pred = check(rec["content"], item["answer"], item["answer_type"]) |
| rec["pred"] = pred |
| rec["correct"] = bool(ok) and rec["finish_reason"] == "stop" |
| stats["done"] += 1 |
| stats["correct"] += rec["correct"] |
| stats["out_tokens"] += usage.get("completion_tokens", 0) |
| out_f.write(json.dumps(rec, ensure_ascii=False) + "\n") |
| out_f.flush() |
| n = stats["done"] + stats["errors"] |
| if n % args.log_every == 0: |
| print(progress(stats, budget), flush=True) |
| if args.target_correct and stats["correct"] >= args.target_correct: |
| stats["stop"] = True |
| finally: |
| queue.task_done() |
|
|
|
|
| def progress(stats, budget): |
| done = max(1, stats["done"]) |
| return ( |
| f"{time.strftime('%H:%M:%S')} done={stats['done']} err={stats['errors']} " |
| f"correct={stats['correct']} ({stats['correct'] / done:.0%}) " |
| f"mean_out={stats['out_tokens'] / done:.0f} " |
| f"out_left_day={budget.out_left} req_left_day={budget.req_left}" |
| ) |
|
|
|
|
| async def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--pool", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--effort", default="medium") |
| ap.add_argument("--max-tokens", type=int, default=16384) |
| ap.add_argument("--temperature", type=float, default=0.7) |
| ap.add_argument("--top-p", type=float, default=0.99) |
| ap.add_argument("--concurrency", type=int, default=12) |
| ap.add_argument("--per-minute", type=float, default=50) |
| ap.add_argument("--limit", type=int, default=0, help="max new prompts this run") |
| ap.add_argument("--ids", help="file with ids to run (one per line)") |
| ap.add_argument("--target-correct", type=int, default=0) |
| ap.add_argument("--reserve-out", type=int, default=40000) |
| ap.add_argument("--reserve-req", type=int, default=200) |
| ap.add_argument("--log-every", type=int, default=20) |
| args = ap.parse_args() |
|
|
| key = os.environ.get("REFLECTION_API_KEY") |
| if not key: |
| sys.exit("REFLECTION_API_KEY not set") |
|
|
| pool = [json.loads(line) for line in open(args.pool)] |
| if args.ids: |
| wanted = [line.strip() for line in open(args.ids) if line.strip()] |
| by_id = {p["id"]: p for p in pool} |
| pool = [by_id[i] for i in wanted] |
| done = set() |
| if Path(args.out).exists(): |
| for line in open(args.out): |
| r = json.loads(line) |
| if "error" not in r: |
| done.add((r["id"], r["effort"])) |
| todo = [p for p in pool if (p["id"], args.effort) not in done] |
| if args.limit: |
| todo = todo[: args.limit] |
| print(f"{len(todo)} prompts to run at effort={args.effort} ({len(done)} already done)", flush=True) |
|
|
| budget = Budget(args.reserve_out, args.reserve_req) |
| pacer = Pacer(args.per_minute) |
| stats = {"done": 0, "errors": 0, "correct": 0, "out_tokens": 0, "stop": False} |
| queue = asyncio.Queue() |
| for p in todo: |
| queue.put_nowait(p) |
|
|
| headers = {"Authorization": f"Bearer {key}"} |
| timeout = httpx.Timeout(connect=30, read=900, write=60, pool=60) |
| limits = httpx.Limits(max_connections=args.concurrency) |
| with open(args.out, "a") as out_f: |
| async with httpx.AsyncClient(headers=headers, timeout=timeout, limits=limits) as client: |
| workers = [ |
| asyncio.create_task(worker(f"w{i}", queue, client, args, budget, pacer, out_f, stats)) |
| for i in range(args.concurrency) |
| ] |
| await queue.join() |
| for w in workers: |
| w.cancel() |
| print("FINISHED " + progress(stats, budget), flush=True) |
| if budget.exhausted(): |
| print(f"STOPPED: daily budget reserve reached; resets in ~{(budget.reset_s or 0) / 3600:.1f}h", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| asyncio.run(main()) |
|
|
