gen.py

"""Stage 1: collect graded (task, agent answer) pairs - the raw material for Jedi.

Tasks: the text and code sources Shaktimaan was built on (same ids: GSM8K, MATH, MMLU-Pro, ARC, MBPP, HumanEval) plus
new domains for coverage (AQuA-RAT, CommonsenseQA) and BIG-Bench Hard, which is kept out of training entirely as an
out-of-domain test. Every agent in $AGENTS answers every task it can take, $SAMPLES times (temperature 0.6), and the
full visible answer is stored with an automatic grade (exact / numeric / math-normalised match; code is graded by
running its tests in a no-network Docker sandbox).

Resumable: (task, agent, sample) triples already in data/answers.jsonl without an API error are skipped.
BORDERLINE=1 restricts the run to tasks whose existing answers disagree (some graded correct, some wrong) - used for
a second pass with more samples there, where the hard-to-judge answers are.
"""
import asyncio, json, os, random, sys, time
import httpx
from datasets import get_dataset_config_names, load_dataset
from tqdm import tqdm
from common import (CODE_SYSTEM, POOL, SOLVER_SYSTEM, capable, endpoints, est_tokens, extract_code, extract_final,
                    grade, last_boxed, max_out, strip_think, truncate)

OUT = "data/answers.jsonl"
AGENTS = os.environ.get("AGENTS", "glm-5.3-flash,qwen3.8-27b,kimi-vl").split(",")
SAMPLES = {a: int(n) for a, n in (x.split("=") for x in os.environ.get(
    "SAMPLES", "glm-5.3-flash=2,qwen3.8-27b=1,kimi-vl=1").split(","))}
SIZES = {"gsm8k": 1000, "math": 1000, "mmlu_pro": 1200, "arc": 800, "mbpp": 700, "aqua": 800, "csqa": 600,
         "bbh": 30}  # bbh: per sub-task (27 sub-tasks)
random.seed(0)


def load_tasks():
    tasks = []
    # --- sources shared with Shaktimaan (identical loading, so ids line up with the Shaktimaan-routing dataset) ---
    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"))
    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",
                          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)
    for i, r in enumerate(ds.select(range(min(SIZES["arc"], len(ds))))):
        labels = r["choices"]["label"]
        opts = "\n".join(f"{chr(65 + k)}. {t}" for k, t in enumerate(r["choices"]["text"]))
        tasks.append(dict(id=f"arc-{i}", source="arc", domain="science", kind="mc",
                          prompt=f"{r['question']}\n\n{opts}", gold=chr(65 + labels.index(r["answerKey"]))))
    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']})"))
    # --- new domains ---
    ds = load_dataset("deepmind/aqua_rat", "raw", split="train").shuffle(seed=0).select(range(SIZES["aqua"]))
    for i, r in enumerate(ds):
        opts = "\n".join(o.replace(")", ". ", 1) for o in r["options"])
        tasks.append(dict(id=f"aqua-{i}", source="aqua", domain="math", kind="mc",
                          prompt=f"{r['question']}\n\n{opts}", gold=r["correct"]))
    ds = load_dataset("tau/commonsense_qa", split="train").shuffle(seed=0).select(range(SIZES["csqa"]))
    for r in ds:
        opts = "\n".join(f"{l}. {t}" for l, t in zip(r["choices"]["label"], r["choices"]["text"]))
        tasks.append(dict(id=f"csqa-{r['id']}", source="csqa", domain="commonsense", kind="mc",
                          prompt=f"{r['question']}\n\n{opts}", gold=r["answerKey"]))
    # --- out-of-domain test only ---
    for cfg in get_dataset_config_names("lukaemon/bbh"):
        ds = load_dataset("lukaemon/bbh", cfg, split="test").shuffle(seed=0)
        for i, r in enumerate(ds.select(range(min(SIZES["bbh"], len(ds))))):
            tasks.append(dict(id=f"bbh-{cfg}-{i}", source="bbh", domain=f"bbh:{cfg}", kind="exact",
                              prompt=r["input"], gold=r["target"]))
    for t in tasks:
        t["est_tokens"] = est_tokens(t["prompt"])
    return tasks


async def call(client, url, model, sem, messages, max_tokens):
    """POST a chat completion. HTTP 400 (e.g. context overflow) is returned immediately as `rejected`."""
    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


async def answer(client, urls, sems, task, agent, sample):
    key = dict(task_id=task["id"], agent=agent, sample=sample)
    if not capable(agent, task["prompt"], False):
        return dict(key, skipped="context")
    system = CODE_SYSTEM if task["kind"] == "code" else SOLVER_SYSTEM
    o = await call(client, urls[agent], POOL[agent]["model"], sems[agent],
                   [{"role": "system", "content": system}, {"role": "user", "content": task["prompt"]}],
                   max_out(agent, task["est_tokens"]))
    if task["kind"] == "code":
        pred = extract_code(o["text"])
        correct = await run_tests(pred, task["gold"])
    else:
        pred = extract_final(o["text"])
        correct = grade(task["kind"], pred, task["gold"])
    return dict(key, text=truncate(strip_think(o["text"]), 6000, 6000), pred=pred[:300], correct=correct,
                finish=o.get("finish"), error=o.get("error"), skipped="rejected" if o.get("rejected") else None,
                **{k: o[k] for k in ("prompt_tokens", "completion_tokens", "latency")})


async def alive(client, urls, agent):
    try:
        r = await client.post(urls[agent], timeout=120, json={
            "model": POOL[agent]["model"], "max_tokens": 4, "messages": [{"role": "user", "content": "hi"}]})
        return r.status_code == 200
    except Exception:  # noqa: BLE001
        return False


async def main():
    os.makedirs("data", exist_ok=True)
    tasks = load_tasks()
    with open("data/tasks.jsonl", "w") as f:
        for t in tasks:
            f.write(json.dumps(t) + "\n")
    done, grades = set(), {}
    if os.path.exists(OUT):
        for r in map(json.loads, open(OUT)):
            if not r.get("error"):
                done.add((r["task_id"], r["agent"], r["sample"]))
                if not r.get("skipped"):
                    grades.setdefault(r["task_id"], set()).add(bool(r["correct"]))
    if os.environ.get("BORDERLINE"):
        tasks = [t for t in tasks if len(grades.get(t["id"], ())) == 2 and t["source"] != "bbh"]
    urls = {n: f"{base.rstrip('/')}/chat/completions" for n, base in endpoints().items()}
    sems = {n: asyncio.Semaphore(int(os.environ.get(f"CONC_{n}", POOL[n]["concurrency"]))) 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:
        up = [a for a, ok in zip(AGENTS, await asyncio.gather(*[alive(client, urls, a) for a in AGENTS])) if ok]
        print(f"{len(tasks)} tasks; agents up: {up}; down: {sorted(set(AGENTS) - set(up))}", flush=True)
        jobs = [(t, a, s) for t in tasks for a in up for s in range(SAMPLES.get(a, 1))
                if (t["id"], a, s) not in done]
        print(f"{len(done)} answers done, {len(jobs)} to run", flush=True)
        with open(OUT, "a") as f:
            futs = [answer(client, urls, sems, t, a, s) for t, a, s in jobs]
            for fut in tqdm(asyncio.as_completed(futs), total=len(futs), mininterval=60, file=sys.stdout):
                f.write(json.dumps(await fut) + "\n")
                f.flush()


if __name__ == "__main__":
    asyncio.run(main())