"""Stage 1: build a diverse task set with checkable answers and run every capable agent in the pool on it. Sources: GSM8K, MATH, MMLU-Pro, ARC (text reasoning), MBPP + HumanEval (code, graded by running the tests in a no-network Docker sandbox), ChartQA (images), LongBench-v2 (long context, up to ~260k tokens). Agents that cannot take a task (no image input, context too small) are recorded as failures without being called. Reuses results already in data/runs.jsonl and only runs missing (task, agent) pairs. Writes data/runs_v3.jsonl. """ import asyncio, base64, io, json, os, random, sys, time import httpx from datasets import load_dataset from tqdm import tqdm from common import (AGENTS, CODE_SYSTEM, POOL, SOLVER_SYSTEM, VISION, capable, endpoints, est_tokens, extract_code, extract_final, grade, max_out, truncate) OLD, OUT = "data/runs.jsonl", os.environ.get("OUT", "data/runs_v3.jsonl") SIZES = {"gsm8k": 1000, "math": 1000, "mmlu_pro": 1200, "arc": 800, "mbpp": 700, "chartqa": 300} LONG_TOKENS = 32768 # prompts above this get their own small concurrency budget random.seed(0) def load_tasks(): tasks = [] # --- text reasoning (identical to v1/v2 so existing results are reused by id) --- ds = load_dataset("openai/gsm8k", "main", split="train").shuffle(seed=0).select(range(SIZES["gsm8k"])) for i, r in enumerate(ds): tasks.append(dict(id=f"gsm8k-{i}", source="gsm8k", domain="math", kind="num", prompt=r["question"], gold=r["answer"].split("####")[-1].strip())) rows = [] for cfg in ["algebra", "counting_and_probability", "geometry", "intermediate_algebra", "number_theory", "prealgebra", "precalculus"]: rows += list(load_dataset("EleutherAI/hendrycks_math", cfg, split="train")) from common import last_boxed rows = [r for r in rows if last_boxed(r["solution"])] random.shuffle(rows) for i, r in enumerate(rows[:SIZES["math"]]): tasks.append(dict(id=f"math-{i}", source="math", domain="math", kind="math", level=r["level"], prompt=r["problem"], gold=last_boxed(r["solution"]))) ds = load_dataset("TIGER-Lab/MMLU-Pro", split="test").shuffle(seed=0).select(range(SIZES["mmlu_pro"])) for i, r in enumerate(ds): opts = "\n".join(f"{chr(65 + k)}. {o}" for k, o in enumerate(r["options"])) tasks.append(dict(id=f"mmlupro-{i}", source="mmlu_pro", domain=r["category"], kind="mc", prompt=f"{r['question']}\n\n{opts}", gold=r["answer"])) ds = load_dataset("allenai/ai2_arc", "ARC-Challenge", split="train").shuffle(seed=0) ds = ds.select(range(min(SIZES["arc"], len(ds)))) for i, r in enumerate(ds): labels = r["choices"]["label"] opts = "\n".join(f"{chr(65 + k)}. {t}" for k, t in enumerate(r["choices"]["text"])) gold = chr(65 + labels.index(r["answerKey"])) tasks.append(dict(id=f"arc-{i}", source="arc", domain="science", kind="mc", prompt=f"{r['question']}\n\n{opts}", gold=gold)) # --- code --- mbpp = [] for split in ["train", "test", "validation", "prompt"]: mbpp += list(load_dataset("google-research-datasets/mbpp", "full", split=split)) random.Random(2).shuffle(mbpp) for r in mbpp[:SIZES["mbpp"]]: tasks.append(dict(id=f"mbpp-{r['task_id']}", source="mbpp", domain="code", kind="code", prompt=f"{r['text']}\nYour code should pass these tests:\n" + "\n".join(r["test_list"]), gold=(r["test_setup_code"] or "") + "\n" + "\n".join(r["test_list"]))) for r in load_dataset("openai/openai_humaneval", split="test"): tasks.append(dict(id=f"humaneval-{r['task_id'].split('/')[-1]}", source="humaneval", domain="code", kind="code", prompt=f"Complete the following Python function:\n```python\n{r['prompt']}```", gold=f"{r['test']}\ncheck({r['entry_point']})")) # --- vision --- ds = load_dataset("HuggingFaceM4/ChartQA", split="test").shuffle(seed=0).select(range(SIZES["chartqa"])) for i, r in enumerate(ds): buf = io.BytesIO() r["image"].convert("RGB").save(buf, "PNG") tasks.append(dict(id=f"chartqa-{i}", source="chartqa", domain="vision", kind="relaxed", has_image=True, image="data:image/png;base64," + base64.b64encode(buf.getvalue()).decode(), prompt=r["query"], gold=str(r["label"][0]))) # --- long context --- for r in load_dataset("THUDM/LongBench-v2", split="train"): prompt = (f"Read the following text and answer the question.\n\n{r['context']}\n\nQuestion: {r['question']}\n" f"A. {r['choice_A']}\nB. {r['choice_B']}\nC. {r['choice_C']}\nD. {r['choice_D']}") if any(capable(a, prompt, False) for a in AGENTS): tasks.append(dict(id=f"longbench-{r['_id']}", source="longbench", domain=f"long:{r['domain']}", kind="mc", prompt=prompt, gold=r["answer"])) for t in tasks: t.setdefault("has_image", False) t["est_tokens"] = est_tokens(t["prompt"]) return tasks async def call(client, url, model, sem, messages, max_tokens=8192): """POST a chat completion. HTTP 400 (e.g. context overflow) is returned immediately as `rejected`.""" if isinstance(messages, str): messages = [{"role": "system", "content": SOLVER_SYSTEM}, {"role": "user", "content": messages}] async with sem: t0, err = time.time(), None for attempt in range(4): try: r = await client.post(url, json={"model": model, "max_tokens": max_tokens, "temperature": 0.6, "messages": messages}) if r.status_code == 400: return dict(text="", rejected=r.text[:200], prompt_tokens=0, completion_tokens=0, latency=time.time() - t0) r.raise_for_status() j = r.json() return dict(text=j["choices"][0]["message"].get("content") or "", finish=j["choices"][0].get("finish_reason"), prompt_tokens=j["usage"]["prompt_tokens"], completion_tokens=j["usage"]["completion_tokens"], latency=time.time() - t0) except Exception as e: # noqa: BLE001 err = repr(e)[:200] await asyncio.sleep(2 ** attempt) return dict(text="", error=err, prompt_tokens=0, completion_tokens=0, latency=time.time() - t0) SANDBOX = asyncio.Semaphore(24) async def run_tests(code: str, tests: str) -> bool: """Run model code + tests in a throwaway container: no network, 512 MB RAM, 1 CPU, 15 s.""" if not code.strip(): return False async with SANDBOX: proc = await asyncio.create_subprocess_exec( "docker", "run", "--rm", "-i", "--network", "none", "--memory", "512m", "--cpus", "1", "--pids-limit", "64", "--read-only", "--tmpfs", "/tmp:size=64m", "python:3.12-slim", "timeout", "15", "python", "-", stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL) try: await asyncio.wait_for(proc.communicate((code + "\n\n" + tests).encode()), timeout=60) except asyncio.TimeoutError: proc.kill() return False return proc.returncode == 0 def messages_for(task): if task["kind"] == "code": return [{"role": "system", "content": CODE_SYSTEM}, {"role": "user", "content": task["prompt"]}] user = task["prompt"] if task["has_image"]: user = [{"type": "image_url", "image_url": {"url": task["image"]}}, {"type": "text", "text": task["prompt"]}] return [{"role": "system", "content": SOLVER_SYSTEM}, {"role": "user", "content": user}] async def run_agent(client, urls, sems, task, agent): if not capable(agent, task["prompt"], task["has_image"]): reason = "no_vision" if task["has_image"] and agent not in VISION else "context" return dict(pred="", correct=False, prompt_tokens=0, completion_tokens=0, latency=0, finish=None, error=None, skipped=reason) sem = sems["long" if task["est_tokens"] > LONG_TOKENS else "short"][agent] o = await call(client, urls[agent], POOL[agent]["model"], sem, messages_for(task), max_out(agent, task["est_tokens"] + 1500 * task["has_image"])) if task["kind"] == "code": code = extract_code(o["text"]) pred, correct = code[:300], await run_tests(code, task["gold"]) else: pred = extract_final(o["text"]) correct = grade(task["kind"], pred, task["gold"]) return dict(pred=pred[:300], correct=correct, **{k: o[k] for k in ("prompt_tokens", "completion_tokens", "latency")}, finish=o.get("finish"), error=o.get("error"), skipped="rejected" if o.get("rejected") else None) async def run_task(client, urls, sems, task, prior): todo = [a for a in AGENTS if a not in prior] outs = await asyncio.gather(*[run_agent(client, urls, sems, task, a) for a in todo]) task = {k: v for k, v in task.items() if k != "image"} if task["source"] == "longbench": # keep the dataset small: store only what the router reads of huge contexts task["prompt"] = truncate(task["prompt"]) task["prompt_truncated"] = True task["results"] = {**prior, **dict(zip(todo, outs))} return task async def main(): os.makedirs("data", exist_ok=True) tasks = load_tasks() old = {} if os.path.exists(OLD): old = {r["id"]: r["results"] for r in map(json.loads, open(OLD)) if not any(v.get("error") for v in r["results"].values())} done = {json.loads(l)["id"] for l in open(OUT)} if os.path.exists(OUT) else set() todo = [t for t in tasks if t["id"] not in done] if os.environ.get("SOURCES"): todo = [t for t in todo if t["source"] in os.environ["SOURCES"].split(",")] if os.environ.get("LIMIT"): todo = random.Random(1).sample(todo, int(os.environ["LIMIT"])) print(f"{len(tasks)} tasks ({sum(t['source'] == 'longbench' for t in tasks)} long), {len(done)} done, " f"{len(todo)} to run, {sum(t['id'] in old for t in todo)} reuse prior results", flush=True) urls = {n: f"{base.rstrip('/')}/chat/completions" for n, base in endpoints().items()} sems = {"short": {n: asyncio.Semaphore(POOL[n]["concurrency"]) for n in POOL}, "long": {n: asyncio.Semaphore(4) for n in POOL}} limits = httpx.Limits(max_connections=400, max_keepalive_connections=400) async with httpx.AsyncClient(timeout=httpx.Timeout(1800), limits=limits) as client: with open(OUT, "a") as f: jobs = [run_task(client, urls, sems, t, old.get(t["id"], {})) for t in todo] for fut in tqdm(asyncio.as_completed(jobs), total=len(jobs), mininterval=30, file=sys.stdout): f.write(json.dumps(await fut) + "\n") f.flush() if __name__ == "__main__": asyncio.run(main())