Shaktimaan v1.0

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,206 @@
+---
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+license: apache-2.0
+language: [en]
+datasets: [podstack/Shaktimaan-routing]
+library_name: transformers
+pipeline_tag: text-generation
+tags: [router, llm-router, orchestration, multi-agent, model-selection, fugu]
+model-index:
+- name: Shaktimaan v1.0
+  results:
+  - task: {type: text-classification, name: LLM routing}
+    dataset: {type: podstack/Shaktimaan-routing, name: Shaktimaan routing (held-out test, 554 tasks)}
+    metrics:
+    - {type: accuracy, value: 0.8339, name: end-to-end task accuracy}
+    - {type: relative_cost, value: 0.4240, name: cost vs always-strongest agent}
+---
+# Shaktimaan v1.0 (0.5B)
+
+**Shaktimaan** is a 0.5B-parameter LLM orchestrator in the spirit of Sakana Fugu. It reads a task and decides which
+agent(s) from a pool of six LLMs should handle it - one agent (`single`), or several that answer independently with
+the last one aggregating (`ensemble`) - trading off cost, capability (image input, context length) and accuracy. It
+does not answer tasks itself. Fine-tuned from [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct)
+on [podstack/Shaktimaan-routing](https://huggingface.co/datasets/podstack/Shaktimaan-routing).
+
+## At a glance
+- **83.4%** end-to-end accuracy on 554 held-out tasks (math, knowledge,
+  science, code, long context, charts) at **0.42x** the cost of always using the
+  strongest agent.
+- Beats every single agent used alone (best: Always `glm-5.3-flash` at 79.2%) and the
+  strongest agent alone (78.5% at 1.00x).
+- Learns capability routing from data: charts -> the only vision agent, very long inputs -> the long-context agents.
+- **Does not beat hand-written capability rules overall**: Capability rules -> `glm-5.3-flash` reaches
+  85.7% at 0.13x. See *Results* and *Findings*.
+- 100% of outputs are valid route JSON (100.0% on the test set).
+
+## Output format
+```json
+{"domain": "math", "difficulty": "easy|medium|hard", "strategy": "single|ensemble", "agents": ["glm-5.3-flash"]}
+```
+
+## Agent pool
+| agent | model | context | image input | relative cost / token |
+|---|---|---|---|---|
+| `gemma-3-1b` | `google/gemma-3-1b-it` | 32k | no | 0.02 |
+| `kimi-vl` | `Kimi-VL-A3B-Thinking` | 8k | yes | 0.05 |
+| `nemotron-lightning` | `nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4` | 128k | no | 0.1 |
+| `glm-5.3-flash` | `zai-org/GLM-5.3-Flash` | 32k | no | 0.2 |
+| `qwen3.8-27b` | `Qwen/Qwen3.8-27B` | 32k | no | 0.5 |
+| `nemotron-ultra` | `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4` | 256k | no | 1.0 |
+
+Agents are reached through OpenAI-compatible chat-completions endpoints. "Relative cost" is a per-token weight used
+to compute routing cost (roughly relative serving cost), not a vendor price.
+
+## Quick start
+```python
+import json
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import importlib.util, sys
+
+repo = "podstack/Shaktimaan"
+spec = importlib.util.spec_from_file_location("common", hf_hub_download(repo, "common.py"))
+common = importlib.util.module_from_spec(spec); spec.loader.exec_module(common)
+
+tok = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForCausalLM.from_pretrained(repo)
+task = "Write a Python function that returns the n-th Fibonacci number."
+msgs = [{"role": "system", "content": common.ROUTER_SYSTEM},
+        {"role": "user", "content": common.router_view(task, has_image=False)}]
+enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)
+out = model.generate(**enc, max_new_tokens=96, do_sample=False)
+print(json.loads(tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)))
+```
+
+Always build the input with `common.router_view(task, has_image)` - the router was trained on that exact format (a
+`[input ~Nk tokens; 1 image attached]` header plus the first 5,000 and last 1,500 characters of the task) with
+`common.ROUTER_SYSTEM` as the system prompt.
+
+### End-to-end orchestration
+`orchestrate.py` routes with Shaktimaan, drops agents that cannot take the input, then calls the chosen agent(s) and
+aggregates ensembles:
+
+```bash
+export SHAKTIMAAN_ENDPOINTS='{"kimi-vl": "http://host-a:8000/v1", "nemotron-lightning": "http://host-b:8000/v1", ...}'
+python orchestrate.py "If a train travels 120 km in 1.5 hours, what is its average speed?"
+python orchestrate.py --code "Write a function that checks whether a string is a palindrome."
+python orchestrate.py --image chart.png "What was the revenue in Q2 2021?"
+```
+
+## Training data
+4989 train / 554 test tasks across eight sources. Every capable agent attempted every
+task once and was graded automatically; the label is the cheapest *reliable* agent that solved the task (or an
+ensemble when none did). Full details, solve rates per agent and the label policy are in the
+[dataset card](https://huggingface.co/datasets/podstack/Shaktimaan-routing).
+
+| source | category | train | test | grading | license |
+|---|---|---|---|---|---|
+| [GSM8K](https://huggingface.co/datasets/openai/gsm8k) | math word problems | 895 | 105 | numeric match | MIT |
+| [MATH](https://huggingface.co/datasets/EleutherAI/hendrycks_math) | competition math | 894 | 106 | normalised LaTeX / numeric match | MIT |
+| [MMLU-Pro](https://huggingface.co/datasets/TIGER-Lab/MMLU-Pro) | knowledge MCQ (10 options) | 1080 | 120 | option letter | MIT |
+| [ARC-Challenge](https://huggingface.co/datasets/allenai/ai2_arc) | science MCQ | 725 | 75 | option letter | CC-BY-SA-4.0 |
+| [MBPP](https://huggingface.co/datasets/google-research-datasets/mbpp) | Python coding | 633 | 67 | unit tests executed in a network-less sandbox | CC-BY-4.0 |
+| [HumanEval](https://huggingface.co/datasets/openai/openai_humaneval) | Python coding | 146 | 18 | unit tests executed in a network-less sandbox | MIT |
+| [LongBench-v2](https://huggingface.co/datasets/THUDM/LongBench-v2) | long-context MCQ (up to ~260k tokens) | 347 | 32 | option letter | Apache-2.0 |
+| [ChartQA](https://huggingface.co/datasets/HuggingFaceM4/ChartQA) | chart images | 269 | 31 | relaxed match (numbers within 5%) | GPL-3.0 |
+
+## Training
+| setting | value |
+|---|---|
+| base model | Qwen2.5-0.5B-Instruct (494M parameters) |
+| method | full fine-tune, supervised; loss on the route JSON only |
+| epochs | 3 |
+| learning rate | 2e-5, cosine schedule, 20 warm-up steps |
+| effective batch size | 32 |
+| max sequence length | 3,072 tokens |
+| precision | bf16 |
+| final eval loss | 0.0484 |
+
+## Results
+Held-out test set: 554 tasks the router never saw. A policy's accuracy is whether the agent it picks
+actually solved the task. Shaktimaan's 31 ensemble decisions were **executed live** (all agents
+plus aggregator called and graded); single-agent decisions use the recorded outcomes. Cost is the sum of
+`relative cost x output tokens`, normalised so that always using `nemotron-ultra` = 1.00.
+
+| policy | accuracy | relative cost (always-Ultra = 1.00) |
+|---|---|---|
+| **Shaktimaan v1.0** | 83.4% | 0.424 |
+| Capability rules -> `glm-5.3-flash` | 85.7% | 0.133 |
+| Capability rules -> `qwen3.8-27b` | 85.4% | 0.403 |
+| Capability rules -> `nemotron-ultra` | 82.7% | 1.001 |
+| Always `gemma-3-1b` | 25.6% | 0.002 |
+| Always `kimi-vl` | 61.4% | 0.074 |
+| Always `nemotron-lightning` | 75.6% | 0.109 |
+| Always `glm-5.3-flash` | 79.2% | 0.094 |
+| Always `qwen3.8-27b` | 78.9% | 0.364 |
+| Always `nemotron-ultra` | 78.5% | 1.000 |
+| Oracle (hindsight-best label) | 89.7% | 0.490 |
+
+**Capability rules** are a hand-written baseline: image -> `kimi-vl`; input too long for the default agent -> the
+cheapest agent with a large enough context window; otherwise the default agent. **Oracle** is the hindsight-best
+training label (an upper bound; its ensembles are scored as the aggregator's solo answer).
+
+### Accuracy by source
+| source | n | Shaktimaan | rules -> glm | rules -> qwen | always Ultra | oracle |
+|---|---|---|---|---|---|---|
+| ARC-Challenge | 75 | 98.7% | 98.7% | 98.7% | 98.7% | 100.0% |
+| ChartQA | 31 | 74.2% | 74.2% | 74.2% | 0.0% | 74.2% |
+| GSM8K | 105 | 89.5% | 96.2% | 98.1% | 97.1% | 98.1% |
+| HumanEval | 18 | 88.9% | 94.4% | 94.4% | 94.4% | 100.0% |
+| LongBench-v2 | 32 | 59.4% | 46.9% | 50.0% | 59.4% | 59.4% |
+| MATH | 106 | 75.5% | 74.5% | 75.5% | 65.1% | 82.1% |
+| MBPP | 67 | 94.0% | 98.5% | 92.5% | 97.0% | 100.0% |
+| MMLU-Pro | 120 | 77.5% | 83.3% | 81.7% | 74.2% | 87.5% |
+
+### Routes chosen on the test set
+| route chosen by Shaktimaan | test tasks |
+|---|---|
+| `nemotron-lightning` | 218 |
+| `glm-5.3-flash` | 192 |
+| `kimi-vl` | 80 |
+| `glm-5.3-flash,qwen3.8-27b,nemotron-ultra` | 32 |
+| `nemotron-ultra` | 31 |
+| `qwen3.8-27b` | 1 |
+
+## Findings
+- **Routing beats any single model.** Shaktimaan is 4.1
+  points above the best single agent and 4.9
+  points above always using the strongest (and most expensive) one, at
+  0.42x its cost.
+- **Capability routing is learned from data.** On long-context tasks Shaktimaan matches the oracle
+  (59.4%) and beats the capability rules (46.9%),
+  sending inputs to the agent most likely to solve them rather than merely one that fits. On charts it routes to the
+  only vision agent, as the rules do.
+- **Simple capability rules are stronger overall** (85.7% at
+  0.13x). On ordinary text and code tasks the mid-size agents are close to
+  interchangeable, and which one solves a particular task is largely sampling noise that a router cannot predict from
+  the prompt. Shaktimaan sometimes picks the cheaper `nemotron-lightning` or `kimi-vl` where the rule's fixed
+  `glm-5.3-flash` default is slightly more reliable, and its ensembles (chosen for hard tasks) solved
+  11 of the 31 hard tasks they were sent, at a high cost.
+- **Practical recommendation:** for this pool, apply capability rules first and use Shaktimaan where its learned
+  judgement helps most - choosing among long-context agents and flagging hard tasks.
+
+## Limitations
+- Routes only among the six agents above, with their context windows and image support; a different pool needs new
+  data and retraining.
+- Labels come from a single sample per agent and inherit grading noise (e.g. answer extraction from free text).
+- Not trained on agentic / tool-use, multi-turn or open-ended writing tasks; English only.
+- `est_tokens` uses a conservative characters/3 estimate, so borderline inputs may be routed to a larger-context
+  agent than strictly needed.
+
+## Files
+| file | purpose |
+|---|---|
+| `model.safetensors`, `config.json`, tokenizer files | the router |
+| `common.py` | agent pool, prompts, `router_view`, capability checks, graders |
+| `orchestrate.py` | end-to-end runtime (route -> call agents -> aggregate) |
+| `gen.py` | builds tasks and collects graded agent outcomes |
+| `build.py` | labels outcomes and builds the train / test split |
+| `train.py` | fine-tuning script |
+| `evaluate.py` | held-out evaluation incl. live ensembles and baselines |
+| `eval_results.json` | raw evaluation output |
+
+## License
+Model weights: Apache-2.0 (as the base model). Training data: see the
+[dataset card](https://huggingface.co/datasets/podstack/Shaktimaan-routing) (CC-BY-SA-4.0; ChartQA rows GPL-3.0, not redistributed).

build.py

@@ -0,0 +1,93 @@
+"""Stage 2: turn graded runs into routing labels (chat-format SFT data) with a held-out test split.
+
+Label policy (cost-aware oracle):
+  * cheapest routable agent that solved the task, counting a win only if the agent is reliable on that source
+    (solve rate >= 70%). Wins by unreliable agents are mostly lucky single samples a router cannot predict: v1 sent
+    gemma-3-1b 68 held-out tasks and only 53% were solved, and kimi-vl solves just 63% of text tasks that >= 4 other
+    agents solve (lightning 95%, glm 98%). gemma-3-1b is excluded as a target outright;
+  * if no routable agent solved it -> ensemble of the capable agents among glm / qwen / ultra (aggregated by the
+    last one) when at least two can take the task, otherwise the strongest capable agent alone.
+The router sees `router_view(...)`: a size / image header plus the start and end of the task.
+"""
+import json, random
+from collections import Counter
+from datasets import Dataset, DatasetDict
+from common import AGENTS, POOL, ROUTER_SYSTEM, STRONGEST, capable, est_tokens, router_view
+
+random.seed(0)
+ROUTABLE = [a for a in AGENTS if a != "gemma-3-1b"]
+ENSEMBLE_POOL = ["glm-5.3-flash", "qwen3.8-27b", STRONGEST]
+RELIABLE = 0.70
+SOLVE_RATE = {}  # (source, agent) -> solve rate, filled in main()
+
+
+def cost(res, name):
+    return POOL[name]["cost"] * (res[name]["prompt_tokens"] + res[name]["completion_tokens"])
+
+
+def can(task, agent):
+    """Capability check; long prompts are stored truncated, so trust the recorded skip instead of re-estimating."""
+    return not task["results"][agent].get("skipped") and capable(agent, "" if task.get("prompt_truncated")
+                                                                 else task["prompt"], task["has_image"])
+
+
+def label(task):
+    res = task["results"]
+    n_correct = sum(res[a]["correct"] for a in AGENTS)
+    correct = [a for a in ROUTABLE if res[a]["correct"] and SOLVE_RATE[task["source"], a] >= RELIABLE]
+    difficulty = "easy" if n_correct >= 4 else "medium" if n_correct >= 2 else "hard"
+    if correct:
+        return dict(strategy="single", agents=[min(correct, key=lambda a: cost(res, a))], difficulty=difficulty)
+    ens = [a for a in ENSEMBLE_POOL if can(task, a)]
+    if len(ens) >= 2 and not task["has_image"]:
+        return dict(strategy="ensemble", agents=ens, difficulty="hard")
+    fallback = [a for a in ROUTABLE if can(task, a)] or [STRONGEST]
+    return dict(strategy="single", agents=[fallback[-1]], difficulty="hard")
+
+
+def to_example(task):
+    lab = label(task)
+    target = json.dumps({"domain": task["domain"], "difficulty": lab["difficulty"],
+                         "strategy": lab["strategy"], "agents": lab["agents"]})
+    view = router_view(task["prompt"], task["has_image"], task["est_tokens"])
+    return dict(id=task["id"], source=task["source"], domain=task["domain"], kind=task["kind"],
+                gold=task["gold"], prompt=task["prompt"], has_image=task["has_image"],
+                est_tokens=task["est_tokens"], prompt_truncated=bool(task.get("prompt_truncated")),
+                router_input=view, route=target,
+                agent_results=json.dumps({a: {k: task["results"][a].get(k) for k in
+                                              ("correct", "completion_tokens", "latency", "skipped")}
+                                          for a in AGENTS}),
+                messages=[{"role": "system", "content": ROUTER_SYSTEM},
+                          {"role": "user", "content": view},
+                          {"role": "assistant", "content": target}])
+
+
+def main():
+    rows = [json.loads(l) for l in open("data/runs_v3.jsonl")]
+    n_all = len(rows)
+    rows = [r for r in rows if not any(v.get("error") for v in r["results"].values())]
+    print(f"{n_all} runs, {n_all - len(rows)} dropped for transient API errors")
+    for r in rows:
+        r.setdefault("has_image", False)
+        r.setdefault("est_tokens", est_tokens(r["prompt"]))
+    for src in {r["source"] for r in rows}:
+        sub = [r for r in rows if r["source"] == src]
+        for a in AGENTS:
+            SOLVE_RATE[src, a] = sum(r["results"][a]["correct"] for r in sub) / len(sub)
+    random.shuffle(rows)
+    ex = [to_example(r) for r in rows]
+    n_test = min(max(200, len(ex) // 10), len(ex) // 5)
+    dd = DatasetDict(train=Dataset.from_list(ex[n_test:]), test=Dataset.from_list(ex[:n_test]))
+    dd.save_to_disk("data/shaktimaan_routing")
+    print(dd)
+    routes = Counter(json.loads(e["route"])["strategy"] + ":" + ",".join(json.loads(e["route"])["agents"])
+                     for e in ex)
+    print("label distribution:", *routes.most_common(), sep="\n  ")
+    for src in sorted({r["source"] for r in rows}):
+        sub = [r for r in rows if r["source"] == src]
+        print(f"{src:10s} n={len(sub):4d} " + " ".join(
+            f"{a}={sum(r['results'][a]['correct'] for r in sub) / len(sub):.2f}" for a in AGENTS))
+
+
+if __name__ == "__main__":
+    main()

chat_template.jinja

@@ -0,0 +1,54 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\n<tool_call>\n{"name": "' }}
+            {{- tool_call.name }}
+            {{- '", "arguments": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+{%- endif %}

common.py

@@ -0,0 +1,157 @@
+"""Shared config for Shaktimaan: the agent pool, prompts, capabilities and graders."""
+import json, os, re
+
+# agent -> served model id, max concurrent requests we send, relative cost per token
+POOL = {
+    "gemma-3-1b":         dict(model="google/gemma-3-1b-it", concurrency=32, cost=0.02),
+    "kimi-vl":            dict(model="Kimi-VL-A3B-Thinking", concurrency=64, cost=0.05),
+    "nemotron-lightning": dict(model="nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", concurrency=64, cost=0.10),
+    "glm-5.3-flash":      dict(model="zai-org/GLM-5.3-Flash", concurrency=32, cost=0.20),
+    "qwen3.8-27b":        dict(model="Qwen/Qwen3.8-27B", concurrency=12, cost=0.50),
+    "nemotron-ultra":     dict(model="nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", concurrency=64, cost=1.00),
+}
+AGENTS = list(POOL)  # ordered cheapest -> most expensive
+STRONGEST = "nemotron-ultra"
+CONTEXT = {"gemma-3-1b": 32768, "kimi-vl": 8192, "nemotron-lightning": 131072, "glm-5.3-flash": 32768,
+           "qwen3.8-27b": 32768, "nemotron-ultra": 262144}
+VISION = {"kimi-vl"}  # only agent served with image input enabled
+MAX_OUT = 8192
+
+
+def endpoints() -> dict:
+    """Agent -> OpenAI-compatible base URL (".../v1"), from $SHAKTIMAAN_ENDPOINTS: inline JSON or a path to a JSON file."""
+    raw = os.environ["SHAKTIMAAN_ENDPOINTS"]
+    return json.load(open(raw)) if os.path.exists(raw) else json.loads(raw)
+
+SOLVER_SYSTEM = ("Solve the task. Be concise. End your response with a final line exactly of the form:\n"
+                 "Final answer: <answer>\nFor multiple-choice questions, the answer is the option letter only.")
+CODE_SYSTEM = ("Write a correct, self-contained Python solution. Put the complete code (imports, required function "
+               "names and signatures) in a single ```python code block. Do not include tests or example usage.")
+
+ROUTER_SYSTEM = (
+    "You are Shaktimaan, an orchestrator. You do not answer the task yourself. You read the task and decide which "
+    "agent(s) from the pool should handle it, minimising cost while maximising the chance of a correct answer.\n"
+    "Agents (cheapest first): " + ", ".join(AGENTS) + ".\n"
+    "Capabilities: only kimi-vl can see images; context limits: nemotron-ultra 262k tokens, nemotron-lightning 131k, "
+    "kimi-vl 8k, all others 32k.\n"
+    "Strategies: single = one agent answers; ensemble = several agents answer independently and the last one "
+    "aggregates.\n"
+    "Reply with JSON only: {\"domain\": ..., \"difficulty\": \"easy|medium|hard\", \"strategy\": \"single|ensemble\", "
+    "\"agents\": [...]}"
+)
+
+
+def est_tokens(text: str) -> int:
+    """Conservative token estimate (real tokenizers average ~3.5-4 chars/token on English)."""
+    return len(text) // 3 + 1
+
+
+def max_out(agent: str, prompt_tokens: int) -> int:
+    return min(MAX_OUT, CONTEXT[agent] - prompt_tokens - 256)
+
+
+def capable(agent: str, prompt: str, has_image: bool) -> bool:
+    if has_image and agent not in VISION:
+        return False
+    return max_out(agent, est_tokens(prompt) + 1500 * has_image) >= 1024
+
+
+def truncate(prompt: str, head: int = 5000, tail: int = 1500) -> str:
+    return prompt if len(prompt) <= head + tail + 7 else prompt[:head] + "\n[...]\n" + prompt[-tail:]
+
+
+def router_view(prompt: str, has_image: bool = False, n_tokens: int | None = None) -> str:
+    """What the router reads: a size/attachment header plus the start and end of the task."""
+    n = est_tokens(prompt) if n_tokens is None else n_tokens
+    return f"[input ~{n // 1000}k tokens{'; 1 image attached' if has_image else ''}]\n" + truncate(prompt)
+
+
+def strip_think(text: str) -> str:
+    text = text or ""
+    if "</think>" in text:
+        text = text.rsplit("</think>", 1)[1]
+    text = re.sub(r"◁think▷.*?◁/think▷", "", text, flags=re.S)  # Kimi-VL thinking markers
+    return text.strip()
+
+
+def extract_final(text: str) -> str:
+    text = strip_think(text)
+    m = re.findall(r"final answer\s*[:：]\s*(.+)", text, flags=re.I)
+    if m:
+        return m[-1].strip()
+    m = re.findall(r"^\s*answer\s*[:：]\s*(.+)", text, flags=re.I | re.M)
+    if m:
+        return m[-1].strip()
+    boxed = last_boxed(text)
+    if boxed is not None:
+        return boxed
+    lines = [l for l in text.splitlines() if l.strip()]
+    return lines[-1].strip() if lines else ""
+
+
+def extract_code(text: str) -> str:
+    text = strip_think(text)
+    blocks = re.findall(r"```(?:python|py)?\s*\n(.*?)```", text, flags=re.S)
+    return blocks[-1] if blocks else ""
+
+
+def last_boxed(s: str):
+    i = s.rfind("\\boxed{")
+    if i < 0:
+        return None
+    i += len("\\boxed{")
+    depth, j = 1, i
+    while j < len(s) and depth:
+        depth += {"{": 1, "}": -1}.get(s[j], 0)
+        j += 1
+    return s[i:j - 1] if depth == 0 else None
+
+
+def _norm_math(s: str) -> str:
+    s = s.strip().strip("$").strip()
+    b = last_boxed(s)
+    if b is not None:
+        s = b
+    for a, b2 in [("\\left", ""), ("\\right", ""), ("\\!", ""), ("\\,", ""), ("\\;", ""), ("dfrac", "frac"),
+                  ("tfrac", "frac"), ("^\\circ", ""), ("^{\\circ}", ""), ("\\%", ""), ("%", ""), ("\\text{", "{"),
+                  ("\\mbox{", "{"), (" ", ""), ("\\$", "")]:
+        s = s.replace(a, b2)
+    s = re.sub(r"\\frac(\d)(\d)", r"\\frac{\1}{\2}", s)
+    s = s.rstrip(".")
+    if s.startswith("{") and s.endswith("}") and s.count("{") == 1:
+        s = s[1:-1]
+    return s
+
+
+def _num(s: str):
+    s = s.replace(",", "").replace("$", "").strip()
+    m = re.findall(r"-?\d+(?:\.\d+)?", s)
+    return float(m[-1]) if m else None
+
+
+def grade(kind: str, pred: str, gold: str) -> bool:
+    """Grades extracted answers. Code is graded by execution (see gen.run_tests), not here."""
+    if not pred:
+        return False
+    if kind == "mc":
+        m = re.search(r"\b([A-J])\b", pred.upper().replace("(", " ").replace(")", " "))
+        return bool(m) and m.group(1) == gold.upper()
+    if kind == "num":
+        p, g = _num(pred), _num(gold)
+        return p is not None and g is not None and abs(p - g) < 1e-6 * max(1, abs(g))
+    if kind == "relaxed":  # ChartQA convention: numbers within 5%, otherwise case-insensitive exact match
+        pred = last_boxed(pred) or pred
+        pred = re.sub(r"\\(?:mathbf|textbf|text|mathrm)\{([^}]*)\}", r"\1", pred).replace("\\(", "").replace("\\)", "")
+        p, g = _num(pred), _num(gold)
+        if g is not None and p is not None and re.fullmatch(r"[-\d.,%$ ]+", gold.strip()):
+            return abs(p - g) <= 0.05 * abs(g) if g else p == 0
+        return pred.strip().strip(".").lower() == gold.strip().lower()
+    if kind == "math":
+        p, g = _norm_math(pred), _norm_math(gold)
+        if p == g:
+            return True
+        try:
+            return abs(float(p) - float(g)) < 1e-6
+        except ValueError:
+            return False
+    raise ValueError(kind)

config.json

@@ -0,0 +1,57 @@
+{
+  "architectures": [
+    "Qwen2ForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "dtype": "float32",
+  "eos_token_id": 151645,
+  "hidden_act": "silu",
+  "hidden_size": 896,
+  "initializer_range": 0.02,
+  "intermediate_size": 4864,
+  "layer_types": [
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention",
+    "full_attention"
+  ],
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 14,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 2,
+  "pad_token_id": 151643,
+  "rms_norm_eps": 1e-06,
+  "rope_parameters": {
+    "rope_theta": 1000000.0,
+    "rope_type": "default"
+  },
+  "sliding_window": null,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.19.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

eval_results.json

@@ -0,0 +1,127 @@
+{
+  "n_test": 554,
+  "valid_json": 1.0,
+  "route_match": 0.5306859205776173,
+  "live_ensembles": 31,
+  "route_counts": {
+    "nemotron-lightning": 218,
+    "glm-5.3-flash": 192,
+    "kimi-vl": 80,
+    "glm-5.3-flash,qwen3.8-27b,nemotron-ultra": 32,
+    "nemotron-ultra": 31,
+    "qwen3.8-27b": 1
+  },
+  "policies": {
+    "shaktimaan": {
+      "accuracy": 0.8339,
+      "relative_cost": 0.424
+    },
+    "oracle": {
+      "accuracy": 0.8971,
+      "relative_cost": 0.49
+    },
+    "always:gemma-3-1b": {
+      "accuracy": 0.2563,
+      "relative_cost": 0.0024
+    },
+    "always:kimi-vl": {
+      "accuracy": 0.6137,
+      "relative_cost": 0.0736
+    },
+    "always:nemotron-lightning": {
+      "accuracy": 0.7563,
+      "relative_cost": 0.1089
+    },
+    "always:glm-5.3-flash": {
+      "accuracy": 0.7924,
+      "relative_cost": 0.0944
+    },
+    "always:qwen3.8-27b": {
+      "accuracy": 0.7888,
+      "relative_cost": 0.3641
+    },
+    "always:nemotron-ultra": {
+      "accuracy": 0.7852,
+      "relative_cost": 1.0
+    },
+    "rule:glm-5.3-flash": {
+      "accuracy": 0.8574,
+      "relative_cost": 0.1331
+    },
+    "rule:qwen3.8-27b": {
+      "accuracy": 0.8538,
+      "relative_cost": 0.4028
+    },
+    "rule:nemotron-ultra": {
+      "accuracy": 0.8267,
+      "relative_cost": 1.0009
+    }
+  },
+  "accuracy_by_source": {
+    "arc": {
+      "n": 75,
+      "shaktimaan": 0.987,
+      "oracle": 1.0,
+      "always:nemotron-ultra": 0.987,
+      "rule:glm-5.3-flash": 0.987,
+      "rule:qwen3.8-27b": 0.987
+    },
+    "chartqa": {
+      "n": 31,
+      "shaktimaan": 0.742,
+      "oracle": 0.742,
+      "always:nemotron-ultra": 0.0,
+      "rule:glm-5.3-flash": 0.742,
+      "rule:qwen3.8-27b": 0.742
+    },
+    "gsm8k": {
+      "n": 105,
+      "shaktimaan": 0.895,
+      "oracle": 0.981,
+      "always:nemotron-ultra": 0.971,
+      "rule:glm-5.3-flash": 0.962,
+      "rule:qwen3.8-27b": 0.981
+    },
+    "humaneval": {
+      "n": 18,
+      "shaktimaan": 0.889,
+      "oracle": 1.0,
+      "always:nemotron-ultra": 0.944,
+      "rule:glm-5.3-flash": 0.944,
+      "rule:qwen3.8-27b": 0.944
+    },
+    "longbench": {
+      "n": 32,
+      "shaktimaan": 0.594,
+      "oracle": 0.594,
+      "always:nemotron-ultra": 0.594,
+      "rule:glm-5.3-flash": 0.469,
+      "rule:qwen3.8-27b": 0.5
+    },
+    "math": {
+      "n": 106,
+      "shaktimaan": 0.755,
+      "oracle": 0.821,
+      "always:nemotron-ultra": 0.651,
+      "rule:glm-5.3-flash": 0.745,
+      "rule:qwen3.8-27b": 0.755
+    },
+    "mbpp": {
+      "n": 67,
+      "shaktimaan": 0.94,
+      "oracle": 1.0,
+      "always:nemotron-ultra": 0.97,
+      "rule:glm-5.3-flash": 0.985,
+      "rule:qwen3.8-27b": 0.925
+    },
+    "mmlu_pro": {
+      "n": 120,
+      "shaktimaan": 0.775,
+      "oracle": 0.875,
+      "always:nemotron-ultra": 0.742,
+      "rule:glm-5.3-flash": 0.833,
+      "rule:qwen3.8-27b": 0.817
+    }
+  },
+  "live_ensembles_solved": 11
+}

evaluate.py

@@ -0,0 +1,129 @@
+"""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)

gen.py

@@ -0,0 +1,198 @@
+"""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())

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 151643,
+  "do_sample": true,
+  "eos_token_id": [
+    151645,
+    151643
+  ],
+  "pad_token_id": 151643,
+  "repetition_penalty": 1.1,
+  "temperature": 0.7,
+  "top_k": 20,
+  "top_p": 0.8,
+  "transformers_version": "5.19.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e9af94dfa0109e0d3fc82dc339210890b3eeaa2c9261a2adf84c202a71d7805
+size 1976163472

orchestrate.py

@@ -0,0 +1,72 @@
+"""Shaktimaan runtime: route a task with the 0.5B router, then call the chosen agent(s) via OpenAI-compatible APIs.
+
+Usage:
+  SHAKTIMAAN_ENDPOINTS='{"glm-5.3-flash": "http://host:8000/v1", ...}' \
+      python orchestrate.py [--image chart.png] [--code] "task"
+"""
+import argparse, asyncio, base64, json, mimetypes, os, sys
+import httpx
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from common import (CODE_SYSTEM, POOL, ROUTER_SYSTEM, SOLVER_SYSTEM, STRONGEST, capable, endpoints, est_tokens,
+                    max_out, router_view, strip_think)
+
+ROUTER = os.environ.get("SHAKTIMAAN_ROUTER", "podstack/Shaktimaan")
+ENDPOINTS = endpoints()
+
+
+def route(task, has_image):
+    tok = AutoTokenizer.from_pretrained(ROUTER)
+    model = AutoModelForCausalLM.from_pretrained(ROUTER, dtype=torch.bfloat16)
+    enc = tok.apply_chat_template([{"role": "system", "content": ROUTER_SYSTEM},
+                                   {"role": "user", "content": router_view(task, has_image)}],
+                                  add_generation_prompt=True, return_tensors="pt", return_dict=True)
+    out = model.generate(**enc, max_new_tokens=96, do_sample=False)
+    try:
+        plan = json.loads(tok.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True))
+    except Exception:  # noqa: BLE001
+        plan = {"strategy": "single", "agents": []}
+    # keep only agents that are reachable and can actually take this input
+    ok = [a for a in plan.get("agents", []) if a in ENDPOINTS and capable(a, task, has_image)]
+    plan["agents"] = ok or [a for a in [STRONGEST, "kimi-vl"] if a in ENDPOINTS and capable(a, task, has_image)][:1]
+    return plan
+
+
+async def ask(client, agent, messages, n_tokens):
+    r = await client.post(f"{ENDPOINTS[agent]}/chat/completions",
+                          json={"model": POOL[agent]["model"], "messages": messages,
+                                "max_tokens": max_out(agent, n_tokens)})
+    r.raise_for_status()
+    return strip_think(r.json()["choices"][0]["message"]["content"])
+
+
+async def run(task, image=None, code=False):
+    has_image = image is not None
+    plan = route(task, has_image)
+    print("route:", json.dumps(plan), file=sys.stderr)
+    system = CODE_SYSTEM if code else SOLVER_SYSTEM
+    user = [{"type": "image_url", "image_url": {"url": image}}, {"type": "text", "text": task}] if has_image else task
+    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
+    n = est_tokens(task) + 1500 * has_image
+    async with httpx.AsyncClient(timeout=1800) as client:
+        if plan.get("strategy") != "ensemble" or len(plan["agents"]) == 1:
+            return await ask(client, plan["agents"][0], msgs, n)
+        drafts = await asyncio.gather(*[ask(client, a, msgs, n) for a in plan["agents"]])
+        joined = "\n\n".join(f"### Candidate {i + 1}\n{d}" for i, d in enumerate(drafts))
+        agg = f"{task}\n\nIndependent candidate solutions:\n{joined}\n\n" \
+              "Check them, resolve disagreements, and give the best final answer."
+        return await ask(client, plan["agents"][-1], [{"role": "system", "content": system},
+                                                      {"role": "user", "content": agg}], est_tokens(agg))
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("task", nargs="+")
+    p.add_argument("--image", help="path to an image to attach")
+    p.add_argument("--code", action="store_true", help="ask agents for a Python code solution")
+    args = p.parse_args()
+    uri = None
+    if args.image:
+        mime = mimetypes.guess_type(args.image)[0] or "image/png"
+        uri = f"data:{mime};base64," + base64.b64encode(open(args.image, "rb").read()).decode()
+    print(asyncio.run(run(" ".join(args.task), uri, args.code)))

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fd169731d2cbde95e10bf356d66d5997fd885dd8dbb6fb4684da3f23b2585d8
+size 11421892

tokenizer_config.json

@@ -0,0 +1,30 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

train.py

@@ -0,0 +1,19 @@
+"""Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into the Shaktimaan router (loss on the route only)."""
+import os
+from datasets import load_from_disk
+from trl import SFTConfig, SFTTrainer
+
+BASE = os.environ.get("BASE", "Qwen/Qwen2.5-0.5B-Instruct")
+ds = load_from_disk("data/shaktimaan_routing").map(
+    lambda r: {"prompt": r["messages"][:2], "completion": r["messages"][2:]}).select_columns(["prompt", "completion"])
+
+cfg = SFTConfig(
+    output_dir="ckpt", num_train_epochs=3, per_device_train_batch_size=8, gradient_accumulation_steps=4,
+    learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=20, weight_decay=0.0, bf16=True,
+    max_length=3072, completion_only_loss=True, logging_steps=10, eval_strategy="epoch", save_strategy="epoch",
+    save_total_limit=1, report_to=[], gradient_checkpointing=True, dataloader_num_workers=4,
+)
+trainer = SFTTrainer(model=BASE, args=cfg, train_dataset=ds["train"], eval_dataset=ds["test"])
+trainer.train()
+trainer.save_model("shaktimaan-0.5b")
+trainer.processing_class.save_pretrained("shaktimaan-0.5b")