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