gen.py

"""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())