"""Stage 4: score the router on held-out tasks. Single-agent routes are scored from the recorded agent outcomes (an agent that cannot take a task - no image input, context too small - counts as a failure at zero cost). Shaktimaan's ensemble routes on text/code tasks are executed live (agents + aggregator via their OpenAI-compatible endpoints) and graded; oracle ensemble labels are scored conservatively as the aggregator's recorded outcome. Baselines: every always-one-agent policy, plus capability rules ("rule:X" = image -> kimi-vl, too long for X -> the cheapest agent with enough context, otherwise X). A useful router has to beat the rules, not just fixed agents. """ import asyncio, json, sys from collections import Counter, defaultdict import httpx import torch from datasets import load_from_disk from transformers import AutoModelForCausalLM, AutoTokenizer from common import (AGENTS, CODE_SYSTEM, POOL, ROUTER_SYSTEM, SOLVER_SYSTEM, STRONGEST, capable, endpoints, extract_code, extract_final, grade, strip_think) from gen import call, run_tests path = sys.argv[1] if len(sys.argv) > 1 else "shaktimaan-0.5b" tok = AutoTokenizer.from_pretrained(path) tok.padding_side = "left" model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).cuda().eval() test = load_from_disk("data/shaktimaan_routing")["test"] BY_CONTEXT = ["kimi-vl", "glm-5.3-flash", "nemotron-lightning", STRONGEST] # cheapest first, growing context def parse(text): try: r = json.loads(text.strip()) agents = [a for a in r.get("agents", []) if a in POOL] return r, agents or [STRONGEST], True except Exception: # noqa: BLE001 return {}, [STRONGEST], False def score(row, agents): res = json.loads(row["agent_results"]) return bool(res[agents[-1]]["correct"]), sum(POOL[a]["cost"] * res[a]["completion_tokens"] for a in agents) def rule(row, default): res = json.loads(row["agent_results"]) if row["has_image"]: return ["kimi-vl"] if not res[default]["skipped"]: return [default] return [next((a for a in BY_CONTEXT if not res[a]["skipped"]), STRONGEST)] async def run_ensembles(jobs): """jobs: list of (row, agents). Returns [(correct, cost)] from real agent + aggregator calls.""" urls = {a: f"{base.rstrip('/')}/chat/completions" for a, base in endpoints().items()} sems = {a: asyncio.Semaphore(POOL[a]["concurrency"]) for a in POOL} async def one(client, row, agents): code = row["kind"] == "code" system = CODE_SYSTEM if code else SOLVER_SYSTEM msgs = [{"role": "system", "content": system}, {"role": "user", "content": row["prompt"]}] drafts = await asyncio.gather(*[call(client, urls[a], POOL[a]["model"], sems[a], msgs) for a in agents]) joined = "\n\n".join(f"### Candidate {i + 1}\n{strip_think(d['text'])}" for i, d in enumerate(drafts)) agg = await call(client, urls[agents[-1]], POOL[agents[-1]]["model"], sems[agents[-1]], [ {"role": "system", "content": system}, {"role": "user", "content": f"{row['prompt']}\n\nIndependent candidate solutions:\n{joined}\n\n" "Check them, resolve disagreements, and give the best final answer."}]) cost = sum(POOL[a]["cost"] * d["completion_tokens"] for a, d in zip(agents, drafts)) cost += POOL[agents[-1]]["cost"] * agg["completion_tokens"] ok = await run_tests(extract_code(agg["text"]), row["gold"]) if code else \ grade(row["kind"], extract_final(agg["text"]), row["gold"]) return ok, cost async with httpx.AsyncClient(timeout=httpx.Timeout(1800)) as client: return await asyncio.gather(*[one(client, r, a) for r, a in jobs]) preds = [] for i in range(0, len(test), 32): batch = test[i:i + 32] msgs = [[{"role": "system", "content": ROUTER_SYSTEM}, {"role": "user", "content": v}] for v in batch["router_input"]] enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True, return_dict=True).to("cuda") with torch.no_grad(): out = model.generate(**enc, max_new_tokens=96, do_sample=False) preds += tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True) parsed = [parse(p) for p in preds] ens = [(row, [a for a in ag if capable(a, row["prompt"], False)]) for row, (r, ag, ok) in zip(test, parsed) if r.get("strategy") == "ensemble" and len(ag) > 1 and not row["has_image"] and not row["prompt_truncated"]] ens = [(row, ag) for row, ag in ens if len(ag) > 1] ens_results = dict(zip([row["id"] for row, _ in ens], asyncio.run(run_ensembles(ens)))) print(f"ran {len(ens)} live ensembles, solved {sum(c for c, _ in ens_results.values())}", flush=True) policies = ["shaktimaan", "oracle"] + [f"always:{a}" for a in AGENTS] + \ ["rule:glm-5.3-flash", "rule:qwen3.8-27b", f"rule:{STRONGEST}"] stats = {k: [0, 0.0] for k in policies} by_src = defaultdict(lambda: defaultdict(list)) valid = exact = 0 for row, (r, agents, ok) in zip(test, parsed): valid += ok exact += ok and agents == json.loads(row["route"])["agents"] choices = {"shaktimaan": None, "oracle": json.loads(row["route"])["agents"], **{f"always:{a}": [a] for a in AGENTS}, **{f"rule:{a}": rule(row, a) for a in ["glm-5.3-flash", "qwen3.8-27b", STRONGEST]}} for key, ag in choices.items(): if key == "shaktimaan": c, cost = ens_results[row["id"]] if row["id"] in ens_results else score(row, agents) else: c, cost = score(row, ag) stats[key][0] += c stats[key][1] += cost by_src[row["source"]][key].append(c) n = len(test) base_cost = stats[f"always:{STRONGEST}"][1] report = {"n_test": n, "valid_json": valid / n, "route_match": exact / n, "live_ensembles": len(ens), "live_ensembles_solved": sum(c for c, _ in ens_results.values()), "route_counts": dict(Counter(",".join(ag) for _, ag, _ in parsed).most_common()), "policies": {k: {"accuracy": round(v[0] / n, 4), "relative_cost": round(v[1] / base_cost, 4)} for k, v in stats.items()}, "accuracy_by_source": {s: {"n": len(d["shaktimaan"]), **{k: round(sum(v) / len(v), 3) for k, v in d.items() if k in ("shaktimaan", "oracle", "rule:qwen3.8-27b", "rule:glm-5.3-flash", f"always:{STRONGEST}")}} for s, d in sorted(by_src.items())}} print(json.dumps(report, indent=2)) json.dump(report, open(f"{path}/eval_results.json", "w"), indent=2)