evaluate.py

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