Jedi 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,357 @@
+---
+base_model: Qwen/Qwen2.5-0.5B-Instruct
+license: apache-2.0
+language: [en]
+library_name: transformers
+pipeline_tag: text-generation
+tags: [verifier, answer-verification, llm-cascade, model-cascade, escalation, llm-router, orchestration, multi-agent]
+model-index:
+- name: Jedi v1.0
+  results:
+  - task: {type: text-classification, name: answer verification (accept / escalate)}
+    dataset: {type: custom, name: held-out test answers (2,532 answers to 633 unseen tasks)}
+    metrics:
+    - {type: roc_auc, value: 0.9195, name: AUROC (correct vs wrong answers)}
+    - {type: ece, value: 0.0170, name: expected calibration error}
+  - task: {type: text-generation, name: LLM escalation ladder}
+    dataset: {type: custom, name: Shaktimaan held-out text and code tasks (491 tasks)}
+    metrics:
+    - {type: accuracy, value: 0.9185, name: end-to-end task accuracy (threshold 0.7)}
+    - {type: relative_cost, value: 0.4538, name: cost vs always-strongest agent}
+---
+# Jedi v1.0 (0.5B)
+
+**Jedi** is a 0.5B-parameter **answer verifier** that runs an LLM escalation ladder. A cheap model answers first.
+Jedi reads the task *and that answer* and returns a calibrated probability that the answer is correct. If the
+probability is high enough, the answer is accepted. Otherwise the task goes to the next, more expensive model,
+whose answer Jedi checks in turn. Jedi never answers tasks itself; it emits a single token.
+
+Fine-tuned from [Qwen2.5-0.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-0.5B-Instruct) on 38,558 graded
+answers from three LLMs. It is the successor to [Shaktimaan](https://huggingface.co/podstack/Shaktimaan), which
+picks a model *before* any answer exists (see [How Jedi differs from Shaktimaan](#how-jedi-differs-from-shaktimaan)).
+
+## At a glance
+- **91.9%** end-to-end accuracy on 491 held-out text and code tasks at **0.45x** the cost of always using the
+  strongest model. Always using the strongest model gets 84.7%; Shaktimaan v1.0 gets 84.9% at 0.32x on the same tasks.
+- Beats every single model in the pool. At threshold 0.3 it also beats both standard cascade checks (a format
+  check and self-consistency) on accuracy, cost and silent failures at once.
+- Within 1.4 points of a *perfect* verifier (93.3%) on the same ladder.
+- **AUROC 0.92** at telling correct answers from wrong ones on unseen tasks; calibrated confidence (ECE 0.017).
+- One forward pass and **one output token** per check: ~56 ms at batch size 1, ~450 input tokens.
+- **Weak outside its training domains**: AUROC 0.72 on BIG-Bench Hard and 0.59 on CommonsenseQA. See *Limitations*.
+
+## How it works
+
+```
+             task
+               |
+               v
+   +-----------------------+      answer      +------------------+
+   | rung 1: glm-5.3-flash | ---------------> |       Jedi       |  P(correct) >= threshold ?
+   |      (cost 0.2)       |                  |  0.5B, 1 token   | --- yes --> return answer
+   +-----------------------+                  +------------------+
+               | no (escalate)
+               v
+   +-----------------------+      answer      +------------------+
+   | rung 2: qwen3.8-27b   | ---------------> |       Jedi       | --- yes --> return answer
+   |      (cost 0.5)       |                  +------------------+
+   +-----------------------+
+               | no (escalate)
+               v
+   +-----------------------+
+   | rung 3: nemotron-ultra| -----------------------------------------------> return answer (always accepted)
+   |      (cost 1.0)       |
+   +-----------------------+
+```
+
+1. **What Jedi reads.** `common.judge_view()` builds a short header naming the agent, the task type, the answer
+   length and whether the answer was cut off at the token limit. It then appends the task (first 2,500 + last 1,000
+   characters) and the candidate's visible answer (first 1,500 + last 2,500 characters, with `<think>` blocks
+   removed). The system prompt is `common.JEDI_SYSTEM`.
+2. **What Jedi outputs.** One word, `accept` or `escalate`. Jedi is never sampled. The runtime reads the logits of
+   those two tokens at the first output position and takes the gap `z = logit(accept) - logit(escalate)`.
+3. **Calibration.** `P(correct) = sigmoid(a·z + b)`, with `a`, `b` fitted by Platt scaling on a held-out validation
+   split (`calibration.json`). This drops calibration error from 0.061 to 0.017 on the test split, so a threshold
+   of 0.7 means roughly "accept when at least 70% likely to be right".
+4. **The single dial.** The threshold trades cost for accuracy. Raise it to escalate more and catch more wrong
+   answers; lower it to save money. The top rung is always accepted, so every task gets an answer.
+
+## Quick start
+
+### Score one answer
+```python
+import importlib.util, json, math, torch
+from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "podstack/Jedi"
+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)
+cal = json.load(open(hf_hub_download(repo, "calibration.json")))
+tok = AutoTokenizer.from_pretrained(repo)
+model = AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16).eval()
+
+def p_correct(task, answer, agent="glm-5.3-flash", kind="num", finish="stop"):
+    msgs = [{"role": "system", "content": common.JEDI_SYSTEM},
+            {"role": "user", "content": common.judge_view(task, answer, agent, kind, finish)}]
+    enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", return_dict=True)
+    with torch.no_grad():
+        logits = model(**enc).logits[0, -1].float()
+    gap = (logits[cal["accept_token"]] - logits[cal["escalate_token"]]).item()
+    return 1 / (1 + math.exp(-(cal["a"] * gap + cal["b"])))
+
+task = "A shop sells pencils at 3 for $1.20. How much do 15 pencils cost, in dollars?"
+print(p_correct(task, "Each pencil costs 1.20 / 3 = 0.40. 15 x 0.40 = 6.00.\nFinal answer: 6"))
+print(p_correct(task, "15 pencils is 15 / 3 = 5 groups. 5 x 1.20 = 5.00.\nFinal answer: 5"))
+```
+
+Always build the input with `common.judge_view(...)` and use `common.JEDI_SYSTEM` as the system prompt; Jedi was
+trained on exactly that format. `kind` is the task type it was trained with: `num`, `math`, `mc`
+(multiple choice), `code`, `exact`. `finish` is the agent's `finish_reason` (`"length"` means the answer was truncated).
+
+### Run the full ladder
+`jedi.py` runs the ladder end to end against OpenAI-compatible endpoints:
+
+```bash
+export JEDI_ENDPOINTS='{"glm-5.3-flash": "http://host-a:8000/v1", "qwen3.8-27b": "http://host-b:8000/v1",
+                       "nemotron-ultra": "http://host-c:8000/v1"}'
+python jedi.py "If a train travels 120 km in 1.5 hours, what is its average speed in km/h?"
+python jedi.py --threshold 0.9 --code "Write a function that checks whether a string is a palindrome."
+```
+
+It prints the final answer to stdout and a per-rung trace to stderr, e.g.
+`{"agent": "glm-5.3-flash", "verdict": "accept", "confidence": 0.97, "escalate_to": null}`. Rungs missing from
+`JEDI_ENDPOINTS` are skipped, so a two-rung ladder works too.
+
+## Agent pool (the ladder)
+| rung | agent | model | relative cost / token |
+|---|---|---|---|
+| 1 | `glm-5.3-flash` | `zai-org/GLM-5.3-Flash` | 0.2 |
+| 2 | `qwen3.8-27b` | `Qwen/Qwen3.8-27B` | 0.5 |
+| 3 | `nemotron-ultra` | `nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4` | 1.0 |
+
+`kimi-vl` (`Kimi-VL-A3B-Thinking`) is not on the ladder, but its answers are in the training data. It is the weakest
+of the three answering models, so it supplied most of the wrong answers Jedi learned from. "Relative cost" is a
+per-token weight for roughly relative serving cost, not a vendor price.
+
+## How Jedi differs from Shaktimaan
+Both models are 0.5B fine-tunes of Qwen2.5-0.5B-Instruct built for the same agent pool. They solve the same
+problem, getting good answers without paying for the biggest model every time, at different moments.
+
+| | **Shaktimaan v1.0** (router) | **Jedi v1.0** (verifier) |
+|---|---|---|
+| When it decides | **before** any model answers | **after** each answer |
+| What it reads | the task only | the task **and** a specific candidate answer |
+| What it outputs | route JSON: domain, difficulty, strategy, agent list (~30 tokens) | one token -> calibrated P(correct) |
+| Question it learns | "which agent is likely to solve this kind of task?" | "is *this* answer right?" |
+| Model calls per task | exactly the routed agent(s), all up front | 1 call, plus more only on escalation (1.15 on average at threshold 0.7) |
+| Cost/quality control | fixed by training labels | one threshold, tunable at run time |
+| Inputs it handles | text, code, images (routes to the vision agent), up to ~260k tokens (routes to long-context agents) | text and code; long inputs are truncated for judging |
+| Extra latency | one router call before the agent | ~56 ms per check, plus the next rung's latency when it escalates (~14 s on average at 0.7) |
+
+**Why the switch.** Shaktimaan's main finding was that, on ordinary text and code tasks, *which* mid-size model
+solves a given task is largely sampling noise that cannot be predicted from the prompt. A router has to guess
+blind. Jedi waits until a cheap answer exists and judges that answer: unsupported reasoning, arithmetic that doesn't
+add up, truncation, hedging and code that won't pass its tests. Its errors are mostly answers that look right but
+aren't, which a router cannot see at all.
+
+**Head to head**, on the same 491 held-out tasks with the same fresh agent samples:
+
+| policy | accuracy | relative cost (always-Ultra = 1.00) |
+|---|---|---|
+| Shaktimaan v1.0 | 84.9% | 0.323 |
+| Jedi ladder, threshold 0.1 | 91.5% | 0.263 |
+| Jedi ladder, threshold 0.7 | **91.9%** | 0.454 |
+
+At threshold 0.1, Jedi is both more accurate (+6.5 points) and cheaper than Shaktimaan. (Shaktimaan's ensemble
+routes are scored as the aggregating agent's own answer here, which is slightly conservative for Shaktimaan.)
+
+**They are complementary.** Jedi's ladder is text-only with a 32k-context first rung; it has no notion of images or
+256k-token inputs. A sensible deployment routes by capability first (images -> the vision agent, very long inputs
+-> long-context agents, as Shaktimaan or simple rules do) and sends everything else through Jedi's ladder.
+
+## Training data
+Every task was answered by `glm-5.3-flash` (2 samples), `qwen3.8-27b` (1) and `kimi-vl` (1) at temperature 0.6
+with a "be concise, end with `Final answer:`" prompt. Each answer was graded automatically: exact, numeric or
+LaTeX-normalised matching, and execution of unit tests in a network-less Docker sandbox for code. That gave 28,296
+graded answers on 7,074 tasks. A second **borderline pass** collected 8,580 more GLM and Qwen answers on the 2,145
+tasks where the models disagreed. This is where the hard-to-judge answers are; those extra samples are used for
+training only.
+
+| source | category | tasks | grading | license |
+|---|---|---|---|---|
+| [GSM8K](https://huggingface.co/datasets/openai/gsm8k) | math word problems | 1,000 | numeric match | MIT |
+| [MATH](https://huggingface.co/datasets/EleutherAI/hendrycks_math) | competition math | 1,000 | normalised LaTeX / numeric | MIT |
+| [MMLU-Pro](https://huggingface.co/datasets/TIGER-Lab/MMLU-Pro) | knowledge MCQ (10 options) | 1,200 | option letter | MIT |
+| [ARC-Challenge](https://huggingface.co/datasets/allenai/ai2_arc) | science MCQ | 800 | option letter | CC-BY-SA-4.0 |
+| [MBPP](https://huggingface.co/datasets/google-research-datasets/mbpp) | Python coding | 700 | unit tests in sandbox | CC-BY-4.0 |
+| [HumanEval](https://huggingface.co/datasets/openai/openai_humaneval) | Python coding | 164 | unit tests in sandbox | MIT |
+| [AQuA-RAT](https://huggingface.co/datasets/deepmind/aqua_rat) | algebra word problems (MCQ) | 800 | option letter | Apache-2.0 |
+| [CommonsenseQA](https://huggingface.co/datasets/tau/commonsense_qa) | commonsense MCQ | 600 | option letter | MIT |
+| [BIG-Bench Hard](https://huggingface.co/datasets/lukaemon/bbh) | 27 reasoning tasks | 810 | exact match | MIT |
+
+The first six sources use the same task ids as
+[Shaktimaan-routing](https://huggingface.co/datasets/podstack/Shaktimaan-routing), so the two models can be
+compared on identical tasks. Charts and long documents were left out.
+
+**Solve rate per model** (share of answers graded correct):
+
+| source | glm-5.3-flash | qwen3.8-27b | kimi-vl |
+|---|---|---|---|
+| GSM8K | 0.96 | 0.95 | 0.81 |
+| MATH | 0.75 | 0.75 | 0.42 |
+| MMLU-Pro | 0.81 | 0.80 | 0.49 |
+| ARC-Challenge | 0.97 | 0.97 | 0.92 |
+| MBPP | 0.93 | 0.82 | 0.65 |
+| HumanEval | 0.94 | 0.91 | 0.81 |
+| AQuA-RAT | 0.89 | 0.92 | 0.64 |
+| CommonsenseQA | 0.80 | 0.81 | 0.75 |
+| BIG-Bench Hard | 0.90 | 0.91 | 0.54 |
+
+**Splits are by task, never by answer**, so no task appears in two splits:
+
+| split | tasks | answers | wrong answers | used for |
+|---|---|---|---|---|
+| train | 5,261 | 38,558 | 40.0% (oversampled from ~18%) | fine-tuning |
+| val | 370 | 1,480 | 17.6% | Platt calibration only |
+| test | 633 | 2,532 | 17.6% | Shaktimaan's held-out test tasks + 10% of AQuA / CommonsenseQA |
+| ood | 810 | 3,240 | 18.8% | all of BIG-Bench Hard, never trained on |
+
+Each example's label is `accept` if the answer was graded correct and `escalate` otherwise. Held-out splits use
+only the base samples, without the borderline extras, so every task is sampled the same way.
+
+## Training
+| setting | value |
+|---|---|
+| base model | Qwen2.5-0.5B-Instruct (494M parameters) |
+| method | full fine-tune, supervised (TRL `SFTTrainer`); loss on the one-word verdict only |
+| epochs | 2 (2,380 steps) |
+| learning rate | 2e-5, cosine schedule, 30 warm-up steps |
+| effective batch size | 32 |
+| max sequence length | 2,048 tokens |
+| precision | bf16, gradient checkpointing |
+| hardware / time | one MIG 1g.67gb slice of an NVIDIA B300, 1 h 54 min |
+| eval loss (val) | 0.095 after epoch 1, 0.142 after epoch 2 (the published weights are from epoch 2) |
+
+## Results
+
+### As a verifier
+Calibrated P(correct) on answers to tasks Jedi never saw. *False accept* = share of wrong answers scored >= 0.5;
+*false escalate* = share of correct answers scored < 0.5.
+
+| split | answers | wrong | AUROC | false accept | false escalate | ECE | Brier |
+|---|---|---|---|---|---|---|---|
+| test (in-domain) | 2,532 | 17.6% | **0.920** | 35.9% | 1.9% | 0.017 | 0.063 |
+| val | 1,480 | 17.6% | 0.884 | 42.3% | 1.3% | 0.017 | 0.072 |
+| ood (BIG-Bench Hard) | 3,240 | 18.8% | 0.720 | 70.2% | 6.7% | 0.060 | 0.140 |
+
+**Test split by answering model and by source:**
+
+| slice | answers | wrong | AUROC | false accept |
+|---|---|---|---|---|
+| glm-5.3-flash | 1,266 | 11.8% | 0.910 | 41.6% |
+| qwen3.8-27b | 633 | 12.3% | 0.915 | 50.0% |
+| kimi-vl | 633 | 34.6% | 0.896 | 26.9% |
+| HumanEval | 72 | 9.7% | 0.991 | 14.3% |
+| MATH | 424 | 31.8% | 0.963 | 14.1% |
+| MBPP | 268 | 9.7% | 0.955 | 38.5% |
+| MMLU-Pro | 480 | 29.4% | 0.928 | 34.0% |
+| AQuA-RAT | 316 | 16.8% | 0.896 | 30.2% |
+| GSM8K | 420 | 6.2% | 0.854 | 53.8% |
+| ARC-Challenge | 300 | 4.0% | 0.842 | 91.7% |
+| CommonsenseQA | 252 | 18.3% | 0.593 | 89.1% |
+
+Jedi is strongest where correctness can be checked from the answer itself: code, multi-step math and
+knowledge questions with worked reasoning. It is weakest where a wrong answer reads exactly like a right one,
+such as single-hop commonsense and easy science multiple choice.
+
+### As the gate of the ladder
+491 held-out tasks: Shaktimaan's test tasks from GSM8K, MATH, MMLU-Pro, ARC, MBPP and HumanEval. GLM and Qwen
+answers are fresh samples; Ultra's outcome and token count come from its graded run in Shaktimaan-routing (same task
+ids). Cost = sum over called agents of (relative cost x output tokens), normalised so always-Ultra = 1.00. *Silent
+failure* = a wrong answer accepted by a lower rung, i.e. a mistake the ladder had a chance to escalate and didn't.
+
+| policy | accuracy | relative cost | silent failures | calls / task |
+|---|---|---|---|---|
+| Always `glm-5.3-flash` | 88.8% | 0.115 | - | 1.00 |
+| Always `qwen3.8-27b` | 87.4% | 0.421 | - | 1.00 |
+| Always `nemotron-ultra` | 84.7% | 1.000 | - | 1.00 |
+| Shaktimaan v1.0 routing | 84.9% | 0.323 | - | - |
+| Ladder + format check (escalate truncated / unparseable answers) | 89.8% | 0.406 | 7.3% | 1.07 |
+| Ladder + self-consistency (escalate if two samples disagree) | 90.4% | 0.554 | 6.5% | 1.14 |
+| **Ladder + Jedi, threshold 0.7** | **91.9%** | **0.454** | **2.9%** | 1.15 |
+| Ladder + perfect verifier (oracle, upper bound) | 93.3% | 0.473 | 0.0% | 1.19 |
+
+**Threshold sweep** (same 491 tasks):
+
+| threshold | accuracy | relative cost | silent failures | calls / task | extra latency / task |
+|---|---|---|---|---|---|
+| 0.05 | 91.0% | 0.244 | 7.5% | 1.09 | 8.9 s |
+| 0.1 | 91.5% | 0.263 | 6.5% | 1.10 | 9.5 s |
+| 0.3 | 91.5% | 0.323 | 5.5% | 1.11 | 10.7 s |
+| 0.5 | 91.7% | 0.387 | 4.5% | 1.12 | 12.1 s |
+| **0.7** | **91.9%** | **0.454** | **2.9%** | 1.15 | 14.3 s |
+| 0.8 | 91.0% | 0.506 | 2.9% | 1.20 | 16.4 s |
+| 0.9 | 90.2% | 0.609 | 1.6% | 1.32 | 20.8 s |
+| 0.95 | 87.0% | 0.891 | 1.4% | 1.65 | 34.5 s |
+
+Above 0.7, accuracy *falls* even as silent failures keep dropping. Escalating a correct answer is not free: the
+next rung can get it wrong, and in this pool Ultra's recorded solve rate (84.7%) is below GLM's. Use **0.7** for
+the best accuracy, or **0.1-0.3** for most of the gain at a quarter of always-Ultra's cost.
+
+### Outside the training domains (two-rung ladder: GLM -> Qwen)
+Cost is relative to always using Qwen.
+
+| tasks | policy | accuracy | cost vs always-Qwen | silent failures |
+|---|---|---|---|---|
+| AQuA + CommonsenseQA test (142) | always GLM / always Qwen | 85.2% / 88.7% | 0.30 / 1.00 | - |
+| | self-consistency | 88.7% | 1.05 | 7.0% |
+| | Jedi @ 0.5 | 88.7% | 0.68 | 8.5% |
+| | Jedi @ 0.7 | 87.3% | 0.72 | 8.5% |
+| | oracle | 92.3% | 0.71 | 0.0% |
+| BIG-Bench Hard, unseen (810) | always GLM / always Qwen | 90.1% / 90.6% | 0.14 / 1.00 | - |
+| | self-consistency | 91.4% | 0.53 | 6.1% |
+| | Jedi @ 0.7 | 91.1% | 0.27 | 7.8% |
+| | Jedi @ 0.9 | 92.0% | 0.69 | 2.2% |
+| | oracle | 94.6% | 0.37 | 0.0% |
+
+On unseen task families the ladder still improves on both single models, because the cheap model is usually right
+and Jedi rarely escalates correct answers. It catches far fewer wrong ones than it does in-domain, though. On the
+CommonsenseQA-heavy new-domain set, Jedi at 0.7 does not beat always using Qwen.
+
+### Overhead
+~56 ms per check (batch size 1, one MIG 1g.67gb slice of a B300), ~451 input tokens, 1 output token. Checks are a
+single forward pass and batch well.
+
+## Limitations
+- **Domain-bound.** Verification skill transfers poorly to unseen task families (BIG-Bench Hard AUROC 0.72) and is
+  near chance on CommonsenseQA (0.59). Measure it on your own traffic before relying on a threshold.
+- **A third of wrong in-domain answers get through at 0.5.** False-accept rates are highest for the strongest
+  models, whose mistakes are rarer and more fluent. Jedi is a cost-saving filter, not a correctness guarantee.
+- **Pool-specific.** Trained on answers from GLM-5.3-Flash, Qwen3.8-27B and Kimi-VL; the agent name is part of its
+  input. Other models' answers are out of distribution, and a new pool needs new data and retraining.
+- **Text and code only**, English only, single-turn tasks with a checkable final answer. Not trained on open-ended
+  writing, agentic / tool use, multi-turn or image inputs. Long tasks and answers are truncated (head + tail) for judging.
+- **Labels inherit grading noise** from automatic answer extraction and matching.
+- **Ladder numbers use recorded Ultra outcomes**, a single sample per task from the Shaktimaan-routing runs, not
+  fresh calls.
+- **Validation loss rose in epoch 2** (0.095 -> 0.142) while AUROC and calibration stayed good; a one-epoch model
+  may be marginally better.
+
+## Files
+| file | purpose |
+|---|---|
+| `model.safetensors`, `config.json`, tokenizer files | the verifier |
+| `calibration.json` | Platt scaling `a`, `b` and the `accept` / `escalate` token ids |
+| `common.py` | agent pool, ladder, prompts, `judge_view`, answer extraction and graders |
+| `jedi.py` | end-to-end runtime (answer -> verify -> escalate) |
+| `gen.py` | collects and grades agent answers (stage 1) |
+| `build.py` | builds the accept / escalate dataset and task-level splits (stage 2) |
+| `train.py` | fine-tuning script (stage 3) |
+| `evaluate.py` | verifier metrics, ladder, baselines and threshold sweep (stage 4) |
+| `eval_results.json` | raw evaluation output behind every number on this card |
+
+The graded answer data is not published; `gen.py` and `build.py` regenerate it given endpoints for the three models.
+
+## License
+Model weights: Apache-2.0 (as the base model). Task sources keep their own licenses (table above).

build.py

@@ -0,0 +1,67 @@
+"""Stage 2: turn graded answers into Jedi's verifier dataset.
+
+One example per (task, agent answer): Jedi reads `judge_view(...)` and must reply `accept` (the answer was graded
+correct) or `escalate` (graded wrong). Splits are by task, never by answer, so no task leaks across splits:
+  * test - the in-domain Shaktimaan-routing test tasks (text + code) plus 10% of the new-domain tasks;
+  * ood  - every BIG-Bench Hard task (never trained on: measures generalisation to unseen task families);
+  * val  - 6% of the remaining tasks, used only to calibrate confidence (Platt scaling);
+  * train - the rest. Wrong answers are the minority, so they are oversampled in train to ~40% of examples.
+Extra samples from the borderline pass (gen.py BORDERLINE=1) are used for train only, so held-out splits keep the
+same, unselected sampling for every task.
+"""
+import json, random
+from collections import Counter
+from datasets import Dataset, DatasetDict, load_dataset
+from common import JEDI_SYSTEM, judge_view
+
+random.seed(0)
+WRONG_SHARE = 0.40
+BASE_SAMPLES = {"glm-5.3-flash": 2, "qwen3.8-27b": 1, "kimi-vl": 1}
+
+
+def main():
+    tasks = {t["id"]: t for t in map(json.loads, open("data/tasks.jsonl"))}
+    answers = [a for a in map(json.loads, open("data/answers.jsonl")) if not a.get("error") and not a.get("skipped")]
+    shak_test = set(load_dataset("podstack/Shaktimaan-routing", split="test")["id"])
+    split = {}
+    for tid, t in tasks.items():
+        if t["source"] == "bbh":
+            split[tid] = "ood"
+        elif tid in shak_test or (t["source"] in ("aqua", "csqa") and random.random() < 0.10):
+            split[tid] = "test"
+        else:
+            split[tid] = "val" if random.random() < 0.06 else "train"
+    rows = {k: [] for k in ("train", "val", "test", "ood")}
+    for a in answers:
+        t = tasks[a["task_id"]]
+        if split[t["id"]] != "train" and a["sample"] >= BASE_SAMPLES.get(a["agent"], 1):
+            continue
+        verdict = "accept" if a["correct"] else "escalate"
+        view = judge_view(t["prompt"], a["text"], a["agent"], t["kind"], a["finish"])
+        rows[split[t["id"]]].append(dict(
+            task_id=t["id"], source=t["source"], kind=t["kind"], agent=a["agent"], sample=a["sample"],
+            correct=bool(a["correct"]), finish=a["finish"] or "", completion_tokens=a["completion_tokens"],
+            latency=a["latency"], pred=a["pred"], gold=t["gold"], judge_input=view, verdict=verdict,
+            prompt=[{"role": "system", "content": JEDI_SYSTEM}, {"role": "user", "content": view}],
+            completion=[{"role": "assistant", "content": verdict}]))
+    train = rows["train"]
+    wrong = [r for r in train if not r["correct"]]
+    need = int(WRONG_SHARE * (len(train) - len(wrong)) / (1 - WRONG_SHARE)) - len(wrong)
+    train += random.choices(wrong, k=max(0, need))
+    random.shuffle(train)
+    dd = DatasetDict({k: Dataset.from_list(v) for k, v in rows.items()})
+    dd.save_to_disk("data/jedi_judge")
+    print(dd)
+    for k, v in rows.items():
+        print(f"{k:5s} tasks={len({r['task_id'] for r in v}):5d} answers={len(v):6d} "
+              f"wrong={sum(not r['correct'] for r in v) / len(v):.1%} "
+              f"agents={dict(Counter(r['agent'] for r in v))}")
+    print("solve rate by source x agent (all answers):")
+    by = Counter((tasks[a["task_id"]]["source"], a["agent"]) for a in answers)
+    ok = Counter((tasks[a["task_id"]]["source"], a["agent"]) for a in answers if a["correct"])
+    for src in sorted({s for s, _ in by}):
+        print(f"  {src:10s} " + " ".join(f"{ag}={ok[src, ag] / by[src, ag]:.2f}" for s, ag in sorted(by) if s == src))
+
+
+if __name__ == "__main__":
+    main()

calibration.json

@@ -0,0 +1,6 @@
+{
+  "a": 0.32936498522758484,
+  "b": 0.13781841099262238,
+  "accept_token": 10330,
+  "escalate_token": 81082
+}

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,165 @@
+"""Shared config for Jedi: the agent pool, the escalation ladder, prompts, capabilities, graders and the judge view."""
+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
+
+
+LADDER = ["glm-5.3-flash", "qwen3.8-27b", "nemotron-ultra"]  # cheap -> mid -> strongest; Jedi checks each rung
+
+
+def endpoints() -> dict:
+    """Agent -> OpenAI-compatible base URL (".../v1"), from $JEDI_ENDPOINTS: inline JSON or a path to a JSON file."""
+    raw = os.environ["JEDI_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.")
+
+JEDI_SYSTEM = (
+    "You are Jedi, a verifier. You are shown a task and a candidate answer written by an AI agent. Do not solve "
+    "the task yourself. Judge whether the candidate's final answer is correct: look for reasoning that does not hold "
+    "up, arithmetic that does not add up, code that would fail its tests, hedging, truncation or a wrong answer "
+    "format. Reply with one word: accept if the answer is correct, escalate if it is likely wrong."
+)
+VERDICTS = ("accept", "escalate")
+
+
+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 judge_view(prompt: str, answer: str, agent: str, kind: str, finish: str | None) -> str:
+    """What Jedi reads: who answered and how the answer ended, then the task and the visible answer (head + tail)."""
+    answer = strip_think(answer)
+    head = f"[agent: {agent}; task type: {kind}; answer ~{est_tokens(answer)} tokens" + \
+           ("; answer was cut off at the token limit" if finish == "length" else "") + "]"
+    return f"{head}\n### Task\n{truncate(prompt, 2500, 1000)}\n\n### Candidate answer\n{truncate(answer, 1500, 2500)}"
+
+
+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
+    if kind == "exact":  # BIG-Bench Hard targets: an option like "(B)", or a short string / number / list
+        norm = lambda x: re.sub(r"\s+", " ", x.strip().strip("*").strip().rstrip(".").lower())
+        p, g = norm(pred), norm(gold)
+        if re.fullmatch(r"\([a-r]\)", g):
+            m = re.match(r"^\(?([a-r])\)?(?:\W|$)", p)
+            return bool(m) and m.group(1) == g[1]
+        return p == g
+    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,709 @@
+{
+  "calibration": {
+    "a": 0.32936498522758484,
+    "b": 0.13781841099262238,
+    "accept_token": 10330,
+    "escalate_token": 81082
+  },
+  "verifier": {
+    "val": {
+      "all": {
+        "n": 1480,
+        "answers_wrong": 0.1757,
+        "auroc": 0.8839,
+        "false_accept_rate": 0.4231,
+        "false_escalate_rate": 0.0131,
+        "ece": 0.0173,
+        "brier": 0.0723
+      },
+      "uncalibrated_ece": 0.0724,
+      "glm-5.3-flash": {
+        "n": 740,
+        "answers_wrong": 0.123,
+        "auroc": 0.8637,
+        "false_accept_rate": 0.5604,
+        "false_escalate_rate": 0.0046,
+        "ece": 0.0223,
+        "brier": 0.0629
+      },
+      "kimi-vl": {
+        "n": 370,
+        "answers_wrong": 0.3243,
+        "auroc": 0.9005,
+        "false_accept_rate": 0.2583,
+        "false_escalate_rate": 0.04,
+        "ece": 0.0632,
+        "brier": 0.0928
+      },
+      "qwen3.8-27b": {
+        "n": 370,
+        "answers_wrong": 0.1324,
+        "auroc": 0.8666,
+        "false_accept_rate": 0.5714,
+        "false_escalate_rate": 0.0093,
+        "ece": 0.0299,
+        "brier": 0.0704
+      },
+      "source:aqua": {
+        "n": 212,
+        "answers_wrong": 0.184,
+        "auroc": 0.8519,
+        "false_accept_rate": 0.4359,
+        "false_escalate_rate": 0.0058,
+        "ece": 0.0353,
+        "brier": 0.0735
+      },
+      "source:arc": {
+        "n": 160,
+        "answers_wrong": 0.0625,
+        "auroc": 0.742,
+        "false_accept_rate": 1.0,
+        "false_escalate_rate": 0.0067,
+        "ece": 0.0419,
+        "brier": 0.0556
+      },
+      "source:csqa": {
+        "n": 116,
+        "answers_wrong": 0.25,
+        "auroc": 0.4935,
+        "false_accept_rate": 0.931,
+        "false_escalate_rate": 0.023,
+        "ece": 0.1432,
+        "brier": 0.203
+      },
+      "source:gsm8k": {
+        "n": 264,
+        "answers_wrong": 0.0417,
+        "auroc": 0.9677,
+        "false_accept_rate": 0.2727,
+        "false_escalate_rate": 0.004,
+        "ece": 0.0499,
+        "brier": 0.0147
+      },
+      "source:humaneval": {
+        "n": 36,
+        "answers_wrong": 0.0278,
+        "auroc": 1.0,
+        "false_accept_rate": 0.0,
+        "false_escalate_rate": 0.0,
+        "ece": 0.0512,
+        "brier": 0.0055
+      },
+      "source:math": {
+        "n": 248,
+        "answers_wrong": 0.3024,
+        "auroc": 0.9272,
+        "false_accept_rate": 0.1467,
+        "false_escalate_rate": 0.0173,
+        "ece": 0.0534,
+        "brier": 0.0539
+      },
+      "source:mbpp": {
+        "n": 168,
+        "answers_wrong": 0.131,
+        "auroc": 0.9692,
+        "false_accept_rate": 0.2273,
+        "false_escalate_rate": 0.0274,
+        "ece": 0.0483,
+        "brier": 0.0449
+      },
+      "source:mmlu_pro": {
+        "n": 276,
+        "answers_wrong": 0.2645,
+        "auroc": 0.8385,
+        "false_accept_rate": 0.5068,
+        "false_escalate_rate": 0.0197,
+        "ece": 0.0517,
+        "brier": 0.1229
+      },
+      "reliability": [
+        {
+          "bin": "0.0-0.1",
+          "n": 104,
+          "confidence": 0.029,
+          "accuracy": 0.0
+        },
+        {
+          "bin": "0.1-0.2",
+          "n": 20,
+          "confidence": 0.13,
+          "accuracy": 0.1
+        },
+        {
+          "bin": "0.2-0.3",
+          "n": 17,
+          "confidence": 0.263,
+          "accuracy": 0.176
+        },
+        {
+          "bin": "0.3-0.4",
+          "n": 11,
+          "confidence": 0.355,
+          "accuracy": 0.273
+        },
+        {
+          "bin": "0.4-0.5",
+          "n": 14,
+          "confidence": 0.464,
+          "accuracy": 0.571
+        },
+        {
+          "bin": "0.5-0.6",
+          "n": 23,
+          "confidence": 0.563,
+          "accuracy": 0.696
+        },
+        {
+          "bin": "0.6-0.7",
+          "n": 36,
+          "confidence": 0.66,
+          "accuracy": 0.778
+        },
+        {
+          "bin": "0.7-0.8",
+          "n": 84,
+          "confidence": 0.762,
+          "accuracy": 0.81
+        },
+        {
+          "bin": "0.8-0.9",
+          "n": 252,
+          "confidence": 0.858,
+          "accuracy": 0.857
+        },
+        {
+          "bin": "0.9-1.0",
+          "n": 919,
+          "confidence": 0.96,
+          "accuracy": 0.953
+        }
+      ]
+    },
+    "test": {
+      "all": {
+        "n": 2532,
+        "answers_wrong": 0.1761,
+        "auroc": 0.9195,
+        "false_accept_rate": 0.3587,
+        "false_escalate_rate": 0.0187,
+        "ece": 0.017,
+        "brier": 0.0631
+      },
+      "uncalibrated_ece": 0.0614,
+      "glm-5.3-flash": {
+        "n": 1266,
+        "answers_wrong": 0.1177,
+        "auroc": 0.9102,
+        "false_accept_rate": 0.4161,
+        "false_escalate_rate": 0.0063,
+        "ece": 0.0106,
+        "brier": 0.0458
+      },
+      "kimi-vl": {
+        "n": 633,
+        "answers_wrong": 0.346,
+        "auroc": 0.896,
+        "false_accept_rate": 0.2694,
+        "false_escalate_rate": 0.0507,
+        "ece": 0.0418,
+        "brier": 0.1013
+      },
+      "qwen3.8-27b": {
+        "n": 633,
+        "answers_wrong": 0.1232,
+        "auroc": 0.9152,
+        "false_accept_rate": 0.5,
+        "false_escalate_rate": 0.0198,
+        "ece": 0.0248,
+        "brier": 0.0595
+      },
+      "source:aqua": {
+        "n": 316,
+        "answers_wrong": 0.1677,
+        "auroc": 0.8961,
+        "false_accept_rate": 0.3019,
+        "false_escalate_rate": 0.0304,
+        "ece": 0.0288,
+        "brier": 0.0627
+      },
+      "source:arc": {
+        "n": 300,
+        "answers_wrong": 0.04,
+        "auroc": 0.8423,
+        "false_accept_rate": 0.9167,
+        "false_escalate_rate": 0.0069,
+        "ece": 0.0286,
+        "brier": 0.0392
+      },
+      "source:csqa": {
+        "n": 252,
+        "answers_wrong": 0.1825,
+        "auroc": 0.5933,
+        "false_accept_rate": 0.8913,
+        "false_escalate_rate": 0.0583,
+        "ece": 0.087,
+        "brier": 0.1548
+      },
+      "source:gsm8k": {
+        "n": 420,
+        "answers_wrong": 0.0619,
+        "auroc": 0.8544,
+        "false_accept_rate": 0.5385,
+        "false_escalate_rate": 0.0025,
+        "ece": 0.0175,
+        "brier": 0.0335
+      },
+      "source:humaneval": {
+        "n": 72,
+        "answers_wrong": 0.0972,
+        "auroc": 0.9912,
+        "false_accept_rate": 0.1429,
+        "false_escalate_rate": 0.0,
+        "ece": 0.0366,
+        "brier": 0.0148
+      },
+      "source:math": {
+        "n": 424,
+        "answers_wrong": 0.3184,
+        "auroc": 0.9625,
+        "false_accept_rate": 0.1407,
+        "false_escalate_rate": 0.0035,
+        "ece": 0.033,
+        "brier": 0.0427
+      },
+      "source:mbpp": {
+        "n": 268,
+        "answers_wrong": 0.097,
+        "auroc": 0.9545,
+        "false_accept_rate": 0.3846,
+        "false_escalate_rate": 0.0124,
+        "ece": 0.041,
+        "brier": 0.0411
+      },
+      "source:mmlu_pro": {
+        "n": 480,
+        "answers_wrong": 0.2937,
+        "auroc": 0.9275,
+        "false_accept_rate": 0.3404,
+        "false_escalate_rate": 0.0354,
+        "ece": 0.0448,
+        "brier": 0.0935
+      },
+      "reliability": [
+        {
+          "bin": "0.0-0.1",
+          "n": 204,
+          "confidence": 0.03,
+          "accuracy": 0.005
+        },
+        {
+          "bin": "0.1-0.2",
+          "n": 32,
+          "confidence": 0.143,
+          "accuracy": 0.156
+        },
+        {
+          "bin": "0.2-0.3",
+          "n": 40,
+          "confidence": 0.248,
+          "accuracy": 0.175
+        },
+        {
+          "bin": "0.3-0.4",
+          "n": 23,
+          "confidence": 0.359,
+          "accuracy": 0.478
+        },
+        {
+          "bin": "0.4-0.5",
+          "n": 26,
+          "confidence": 0.454,
+          "accuracy": 0.577
+        },
+        {
+          "bin": "0.5-0.6",
+          "n": 33,
+          "confidence": 0.553,
+          "accuracy": 0.636
+        },
+        {
+          "bin": "0.6-0.7",
+          "n": 71,
+          "confidence": 0.657,
+          "accuracy": 0.746
+        },
+        {
+          "bin": "0.7-0.8",
+          "n": 153,
+          "confidence": 0.758,
+          "accuracy": 0.771
+        },
+        {
+          "bin": "0.8-0.9",
+          "n": 360,
+          "confidence": 0.86,
+          "accuracy": 0.878
+        },
+        {
+          "bin": "0.9-1.0",
+          "n": 1590,
+          "confidence": 0.961,
+          "accuracy": 0.968
+        }
+      ]
+    },
+    "ood": {
+      "all": {
+        "n": 3240,
+        "answers_wrong": 0.1877,
+        "auroc": 0.72,
+        "false_accept_rate": 0.7023,
+        "false_escalate_rate": 0.0672,
+        "ece": 0.0602,
+        "brier": 0.1398
+      },
+      "uncalibrated_ece": 0.1358,
+      "glm-5.3-flash": {
+        "n": 1620,
+        "answers_wrong": 0.0963,
+        "auroc": 0.6469,
+        "false_accept_rate": 0.9487,
+        "false_escalate_rate": 0.0342,
+        "ece": 0.0679,
+        "brier": 0.1001
+      },
+      "kimi-vl": {
+        "n": 810,
+        "answers_wrong": 0.4642,
+        "auroc": 0.6174,
+        "false_accept_rate": 0.5851,
+        "false_escalate_rate": 0.1797,
+        "ece": 0.119,
+        "brier": 0.2476
+      },
+      "qwen3.8-27b": {
+        "n": 810,
+        "answers_wrong": 0.0938,
+        "auroc": 0.6742,
+        "false_accept_rate": 0.7763,
+        "false_escalate_rate": 0.0668,
+        "ece": 0.1386,
+        "brier": 0.1115
+      },
+      "source:bbh": {
+        "n": 3240,
+        "answers_wrong": 0.1877,
+        "auroc": 0.72,
+        "false_accept_rate": 0.7023,
+        "false_escalate_rate": 0.0672,
+        "ece": 0.0602,
+        "brier": 0.1398
+      },
+      "reliability": [
+        {
+          "bin": "0.0-0.1",
+          "n": 30,
+          "confidence": 0.066,
+          "accuracy": 0.033
+        },
+        {
+          "bin": "0.1-0.2",
+          "n": 56,
+          "confidence": 0.149,
+          "accuracy": 0.125
+        },
+        {
+          "bin": "0.2-0.3",
+          "n": 73,
+          "confidence": 0.254,
+          "accuracy": 0.589
+        },
+        {
+          "bin": "0.3-0.4",
+          "n": 82,
+          "confidence": 0.357,
+          "accuracy": 0.598
+        },
+        {
+          "bin": "0.4-0.5",
+          "n": 117,
+          "confidence": 0.449,
+          "accuracy": 0.658
+        },
+        {
+          "bin": "0.5-0.6",
+          "n": 199,
+          "confidence": 0.553,
+          "accuracy": 0.734
+        },
+        {
+          "bin": "0.6-0.7",
+          "n": 375,
+          "confidence": 0.654,
+          "accuracy": 0.773
+        },
+        {
+          "bin": "0.7-0.8",
+          "n": 525,
+          "confidence": 0.755,
+          "accuracy": 0.794
+        },
+        {
+          "bin": "0.8-0.9",
+          "n": 906,
+          "confidence": 0.854,
+          "accuracy": 0.87
+        },
+        {
+          "bin": "0.9-1.0",
+          "n": 877,
+          "confidence": 0.939,
+          "accuracy": 0.928
+        }
+      ]
+    }
+  },
+  "ladder": [
+    "glm-5.3-flash",
+    "qwen3.8-27b",
+    "nemotron-ultra"
+  ],
+  "n_ladder_tasks": 491,
+  "policies": {
+    "always:glm-5.3-flash": {
+      "accuracy": 0.888,
+      "relative_cost": 0.1146,
+      "silent_failure_rate": 0.0,
+      "calls_per_task": 1.0,
+      "extra_agent_latency_s": 0.0
+    },
+    "always:qwen3.8-27b": {
+      "accuracy": 0.8737,
+      "relative_cost": 0.4208,
+      "silent_failure_rate": 0.0,
+      "calls_per_task": 1.0,
+      "extra_agent_latency_s": 0.0
+    },
+    "always:nemotron-ultra": {
+      "accuracy": 0.8473,
+      "relative_cost": 1.0,
+      "silent_failure_rate": 0.0,
+      "calls_per_task": 1.0,
+      "extra_agent_latency_s": 0.0
+    },
+    "ladder:format-check": {
+      "accuracy": 0.8982,
+      "relative_cost": 0.4058,
+      "silent_failure_rate": 0.0733,
+      "calls_per_task": 1.073,
+      "extra_agent_latency_s": 10.95
+    },
+    "ladder:self-consistency": {
+      "accuracy": 0.9043,
+      "relative_cost": 0.5535,
+      "silent_failure_rate": 0.0652,
+      "calls_per_task": 1.141,
+      "extra_agent_latency_s": 15.29
+    },
+    "ladder:oracle": {
+      "accuracy": 0.9328,
+      "relative_cost": 0.473,
+      "silent_failure_rate": 0.0,
+      "calls_per_task": 1.185,
+      "extra_agent_latency_s": 15.89
+    },
+    "shaktimaan-v1.0": {
+      "accuracy": 0.8493,
+      "relative_cost": 0.3226
+    }
+  },
+  "jedi_curve": [
+    {
+      "threshold": 0.05,
+      "accuracy": 0.9104,
+      "relative_cost": 0.2441,
+      "silent_failure_rate": 0.0754,
+      "calls_per_task": 1.086,
+      "extra_agent_latency_s": 8.94
+    },
+    {
+      "threshold": 0.1,
+      "accuracy": 0.9145,
+      "relative_cost": 0.2634,
+      "silent_failure_rate": 0.0652,
+      "calls_per_task": 1.096,
+      "extra_agent_latency_s": 9.49
+    },
+    {
+      "threshold": 0.2,
+      "accuracy": 0.9124,
+      "relative_cost": 0.2746,
+      "silent_failure_rate": 0.0652,
+      "calls_per_task": 1.098,
+      "extra_agent_latency_s": 9.69
+    },
+    {
+      "threshold": 0.3,
+      "accuracy": 0.9145,
+      "relative_cost": 0.3229,
+      "silent_failure_rate": 0.055,
+      "calls_per_task": 1.108,
+      "extra_agent_latency_s": 10.73
+    },
+    {
+      "threshold": 0.4,
+      "accuracy": 0.9145,
+      "relative_cost": 0.3551,
+      "silent_failure_rate": 0.0509,
+      "calls_per_task": 1.114,
+      "extra_agent_latency_s": 11.42
+    },
+    {
+      "threshold": 0.5,
+      "accuracy": 0.9165,
+      "relative_cost": 0.3865,
+      "silent_failure_rate": 0.0448,
+      "calls_per_task": 1.122,
+      "extra_agent_latency_s": 12.08
+    },
+    {
+      "threshold": 0.6,
+      "accuracy": 0.9165,
+      "relative_cost": 0.4313,
+      "silent_failure_rate": 0.0387,
+      "calls_per_task": 1.13,
+      "extra_agent_latency_s": 13.14
+    },
+    {
+      "threshold": 0.7,
+      "accuracy": 0.9185,
+      "relative_cost": 0.4538,
+      "silent_failure_rate": 0.0285,
+      "calls_per_task": 1.147,
+      "extra_agent_latency_s": 14.32
+    },
+    {
+      "threshold": 0.8,
+      "accuracy": 0.9104,
+      "relative_cost": 0.5055,
+      "silent_failure_rate": 0.0285,
+      "calls_per_task": 1.196,
+      "extra_agent_latency_s": 16.37
+    },
+    {
+      "threshold": 0.85,
+      "accuracy": 0.9084,
+      "relative_cost": 0.5408,
+      "silent_failure_rate": 0.0224,
+      "calls_per_task": 1.238,
+      "extra_agent_latency_s": 17.9
+    },
+    {
+      "threshold": 0.9,
+      "accuracy": 0.9022,
+      "relative_cost": 0.6093,
+      "silent_failure_rate": 0.0163,
+      "calls_per_task": 1.316,
+      "extra_agent_latency_s": 20.84
+    },
+    {
+      "threshold": 0.95,
+      "accuracy": 0.8697,
+      "relative_cost": 0.8913,
+      "silent_failure_rate": 0.0143,
+      "calls_per_task": 1.654,
+      "extra_agent_latency_s": 34.46
+    },
+    {
+      "threshold": 0.98,
+      "accuracy": 0.8615,
+      "relative_cost": 1.2413,
+      "silent_failure_rate": 0.0,
+      "calls_per_task": 2.393,
+      "extra_agent_latency_s": 55.17
+    }
+  ],
+  "two_rung": {
+    "new-domains (aqua, csqa)": {
+      "n_tasks": 142,
+      "always:glm-5.3-flash": {
+        "accuracy": 0.8521,
+        "cost_vs_always_mid": 0.2999,
+        "silent_failure_rate": 0.0
+      },
+      "always:qwen3.8-27b": {
+        "accuracy": 0.8873,
+        "cost_vs_always_mid": 1.0,
+        "silent_failure_rate": 0.0
+      },
+      "oracle": {
+        "accuracy": 0.9225,
+        "cost_vs_always_mid": 0.7111,
+        "silent_failure_rate": 0.0
+      },
+      "self-consistency": {
+        "accuracy": 0.8873,
+        "cost_vs_always_mid": 1.0475,
+        "silent_failure_rate": 0.0704
+      },
+      "jedi@0.5": {
+        "accuracy": 0.8873,
+        "cost_vs_always_mid": 0.6761,
+        "silent_failure_rate": 0.0845
+      },
+      "jedi@0.7": {
+        "accuracy": 0.8732,
+        "cost_vs_always_mid": 0.7191,
+        "silent_failure_rate": 0.0845
+      },
+      "jedi@0.9": {
+        "accuracy": 0.8873,
+        "cost_vs_always_mid": 0.8953,
+        "silent_failure_rate": 0.0493
+      }
+    },
+    "ood: bbh": {
+      "n_tasks": 810,
+      "always:glm-5.3-flash": {
+        "accuracy": 0.9012,
+        "cost_vs_always_mid": 0.1436,
+        "silent_failure_rate": 0.0
+      },
+      "always:qwen3.8-27b": {
+        "accuracy": 0.9062,
+        "cost_vs_always_mid": 1.0,
+        "silent_failure_rate": 0.0
+      },
+      "oracle": {
+        "accuracy": 0.9457,
+        "cost_vs_always_mid": 0.3743,
+        "silent_failure_rate": 0.0
+      },
+      "self-consistency": {
+        "accuracy": 0.9136,
+        "cost_vs_always_mid": 0.5322,
+        "silent_failure_rate": 0.0605
+      },
+      "jedi@0.5": {
+        "accuracy": 0.9049,
+        "cost_vs_always_mid": 0.178,
+        "silent_failure_rate": 0.0938
+      },
+      "jedi@0.7": {
+        "accuracy": 0.9111,
+        "cost_vs_always_mid": 0.2665,
+        "silent_failure_rate": 0.0778
+      },
+      "jedi@0.9": {
+        "accuracy": 0.9198,
+        "cost_vs_always_mid": 0.6854,
+        "silent_failure_rate": 0.0222
+      }
+    }
+  },
+  "overhead": {
+    "jedi_ms_per_check_batch1": 55.6,
+    "jedi_input_tokens_mean": 451,
+    "jedi_output_tokens": 1
+  }
+}

evaluate.py

@@ -0,0 +1,279 @@
+"""Stage 4: evaluate Jedi as a verifier and as the gate of an escalation ladder.
+
+Jedi's confidence is P(accept) read from its first output token (accept vs escalate), Platt-calibrated on the val
+split. Reports, on held-out data:
+  * verifier quality per split / agent: AUROC, false-accept rate (wrong answers it lets through), ECE, Brier;
+  * the ladder glm-5.3-flash -> qwen3.8-27b -> nemotron-ultra, with Jedi checking rungs 1 and 2 and the top rung's
+    answer always accepted, swept over the confidence threshold (accuracy, relative cost, silent-failure rate);
+  * baselines on exactly the same tasks: every single agent, Shaktimaan v1.0 routing, a format check (escalate only
+    truncated / unparseable answers), self-consistency (escalate when two cheap samples disagree) and the oracle
+    ladder (a perfect verifier);
+  * overhead: Jedi's own latency per check and the extra agent latency on escalated tasks.
+Answers for glm / qwen are fresh samples (sample 0). Nemotron-ultra's outcome and token count come from the graded
+record in the Shaktimaan-routing dataset (same task ids), as do agents Shaktimaan routes to that were not re-run.
+Relative cost = sum over called agents of (agent price x completion tokens), divided by always-nemotron-ultra.
+"""
+import json, sys, time
+from collections import defaultdict
+import numpy as np
+import torch
+from datasets import load_dataset, load_from_disk
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from common import LADDER, POOL, JEDI_SYSTEM, STRONGEST
+
+path = sys.argv[1] if len(sys.argv) > 1 else "jedi-0.5b"
+tok = AutoTokenizer.from_pretrained(path)
+tok.padding_side = "left"
+model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).cuda().eval()
+ds = load_from_disk("data/jedi_judge")
+ACC, ESC = (tok.encode(w, add_special_tokens=False)[0] for w in ("accept", "escalate"))
+assert ACC != ESC
+CHEAP, MID = LADDER[0], LADDER[1]
+
+
+def logit_gap(views, bs=16):
+    """accept-vs-escalate logit gap at Jedi's first output token, for each judge input."""
+    out = []
+    order = np.argsort([len(v) for v in views])  # length-sorted batches: far less padding
+    for i in range(0, len(views), bs):
+        idx = order[i:i + bs]
+        msgs = [[{"role": "system", "content": JEDI_SYSTEM}, {"role": "user", "content": views[j]}] for j in idx]
+        enc = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True,
+                                      return_dict=True).to("cuda")
+        with torch.no_grad():
+            last = model(**enc).logits[:, -1, :].float()
+        out += list(zip(idx, (last[:, ACC] - last[:, ESC]).tolist()))
+    gap = np.zeros(len(views))
+    for j, g in out:
+        gap[j] = g
+    return gap
+
+
+def fit_platt(z, y):
+    a, b = torch.ones(1, requires_grad=True), torch.zeros(1, requires_grad=True)
+    zt, yt = torch.tensor(z, dtype=torch.float32), torch.tensor(y, dtype=torch.float32)
+    opt = torch.optim.LBFGS([a, b], max_iter=200)
+
+    def closure():
+        opt.zero_grad()
+        loss = torch.nn.functional.binary_cross_entropy_with_logits(a * zt + b, yt)
+        loss.backward()
+        return loss
+    opt.step(closure)
+    return a.item(), b.item()
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+def auroc(p, y):
+    p, y = np.asarray(p), np.asarray(y, bool)
+    if y.all() or not y.any():
+        return None
+    ranks = np.argsort(np.argsort(p)) + 1
+    return float((ranks[y].sum() - y.sum() * (y.sum() + 1) / 2) / (y.sum() * (~y).sum()))
+
+
+def ece(p, y, bins=10):
+    p, y = np.asarray(p), np.asarray(y, float)
+    edges = np.linspace(0, 1, bins + 1)
+    total, table = 0.0, []
+    for lo, hi in zip(edges[:-1], edges[1:]):
+        m = (p >= lo) & ((p < hi) | (hi == 1))
+        if m.any():
+            total += m.mean() * abs(p[m].mean() - y[m].mean())
+            table.append(dict(bin=f"{lo:.1f}-{hi:.1f}", n=int(m.sum()), confidence=round(float(p[m].mean()), 3),
+                              accuracy=round(float(y[m].mean()), 3)))
+    return float(total), table
+
+
+def verifier_metrics(p, y, tau=0.5):
+    p, y = np.asarray(p), np.asarray(y, bool)
+    e, _ = ece(p, y)
+    return dict(n=int(len(y)), answers_wrong=round(float(1 - y.mean()), 4), auroc=round(auroc(p, y) or 0, 4),
+                false_accept_rate=round(float(((p >= tau) & ~y).sum() / max(1, (~y).sum())), 4),
+                false_escalate_rate=round(float(((p < tau) & y).sum() / max(1, y.sum())), 4),
+                ece=round(e, 4), brier=round(float(((p - y) ** 2).mean()), 4))
+
+
+# ---------- score every held-out answer ----------
+z_val = logit_gap(ds["val"]["judge_input"])
+a, b = fit_platt(z_val, ds["val"]["correct"])
+cal = dict(a=a, b=b, accept_token=ACC, escalate_token=ESC)
+P = {}
+report = {"calibration": cal, "verifier": {}}
+for split in ("val", "test", "ood"):
+    d = ds[split]
+    z = z_val if split == "val" else logit_gap(d["judge_input"])
+    p_raw, p = sigmoid(z), sigmoid(a * z + b)
+    y = np.array(d["correct"])
+    rep = {"all": verifier_metrics(p, y), "uncalibrated_ece": round(ece(p_raw, y)[0], 4)}
+    for ag in sorted(set(d["agent"])):
+        m = np.array(d["agent"]) == ag
+        rep[ag] = verifier_metrics(p[m], y[m])
+    for src in sorted(set(d["source"])):
+        m = np.array(d["source"]) == src
+        rep[f"source:{src}"] = verifier_metrics(p[m], y[m])
+    rep["reliability"] = ece(p, y)[1]
+    report["verifier"][split] = rep
+    for r, pi in zip(d.select_columns(["task_id", "agent", "sample", "correct", "completion_tokens", "latency",
+                                       "finish", "pred", "kind"]), p):
+        P[r["task_id"], r["agent"], r["sample"]] = dict(r, p=float(pi), split=split)
+    print(split, json.dumps(rep["all"]), flush=True)
+
+# ---------- ladder ----------
+shak = {r["id"]: r for sp in ("train", "test") for r in load_dataset("podstack/Shaktimaan-routing", split=sp)}
+
+
+def norm_pred(s):
+    return (s or "").strip().strip(".").lower()
+
+
+def format_ok(ans):
+    return ans["finish"] != "length" and bool(ans["pred"].strip())
+
+
+def run_ladder(tasks, rungs, accept, final_record=None):
+    """tasks: task ids. rungs: agent names; accept(task_id, agent) -> (bool, extra_cost). The last rung is always
+    accepted; if final_record is given it supplies the last rung's (correct, tokens)."""
+    correct, cost, silent, calls, extra_latency = 0, 0.0, 0, 0, 0.0
+    for t in tasks:
+        for i, ag in enumerate(rungs):
+            last = i == len(rungs) - 1
+            if last and final_record:
+                c, toks, lat = final_record(t)
+            else:
+                ans = P[t, ag, 0]
+                c, toks, lat = ans["correct"], ans["completion_tokens"], ans["latency"]
+            cost += POOL[ag]["cost"] * toks
+            calls += 1
+            if i:
+                extra_latency += lat
+            if last:
+                correct += c
+                break
+            ok, extra = accept(t, ag)
+            cost += extra
+            if ok:
+                correct += c
+                silent += not c
+                break
+    n = len(tasks)
+    return dict(accuracy=correct / n, cost=cost, silent_failure_rate=silent / n, calls_per_task=calls / n,
+                extra_agent_latency_s=extra_latency / n)
+
+
+def ultra_record(t):
+    r = json.loads(shak[t]["agent_results"])[STRONGEST]
+    return r["correct"], r["completion_tokens"], r["latency"]
+
+
+def jedi(tau):
+    return lambda t, ag: (P[t, ag, 0]["p"] >= tau, 0.0)
+
+
+def fmt(t, ag):
+    return format_ok(P[t, ag, 0]), 0.0
+
+
+def self_consistency(t, ag):
+    s0, s1 = P[t, ag, 0], P.get((t, ag, 1))
+    if s1 is None or s0["kind"] == "code":
+        return format_ok(s0), 0.0
+    return format_ok(s0) and norm_pred(s0["pred"]) == norm_pred(s1["pred"]), POOL[ag]["cost"] * s1["completion_tokens"]
+
+
+def oracle(t, ag):
+    return P[t, ag, 0]["correct"], 0.0
+
+
+def ladder_tasks(split, sources):
+    return sorted({t for (t, _, _), v in P.items() if v["split"] == split and t.split("-")[0] in sources
+                   and (t, CHEAP, 0) in P and (t, MID, 0) in P})
+
+
+test_tasks = [t for t in ladder_tasks("test", {"gsm8k", "math", "mmlupro", "arc", "mbpp", "humaneval"}) if t in shak]
+ultra_cost = sum(POOL[STRONGEST]["cost"] * ultra_record(t)[1] for t in test_tasks)
+
+
+def rel(res):
+    return {"accuracy": round(res["accuracy"], 4), "relative_cost": round(res["cost"] / ultra_cost, 4),
+            "silent_failure_rate": round(res["silent_failure_rate"], 4),
+            "calls_per_task": round(res["calls_per_task"], 3),
+            "extra_agent_latency_s": round(res["extra_agent_latency_s"], 2)}
+
+
+policies = {
+    f"always:{CHEAP}": rel(run_ladder(test_tasks, [CHEAP], None)),
+    f"always:{MID}": rel(run_ladder(test_tasks, [MID], None)),
+    f"always:{STRONGEST}": rel(run_ladder(test_tasks, [STRONGEST], None, ultra_record)),
+    "ladder:format-check": rel(run_ladder(test_tasks, LADDER, fmt, ultra_record)),
+    "ladder:self-consistency": rel(run_ladder(test_tasks, LADDER, self_consistency, ultra_record)),
+    "ladder:oracle": rel(run_ladder(test_tasks, LADDER, oracle, ultra_record)),
+}
+
+# Shaktimaan v1.0 on the same tasks: route with the published router, then score its choice.
+rtok = AutoTokenizer.from_pretrained("podstack/Shaktimaan")
+rtok.padding_side = "left"
+router = AutoModelForCausalLM.from_pretrained("podstack/Shaktimaan", dtype=torch.bfloat16).cuda().eval()
+routes = []
+for i in range(0, len(test_tasks), 32):
+    msgs = [shak[t]["messages"][:2] for t in test_tasks[i:i + 32]]
+    enc = rtok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt", padding=True,
+                                   return_dict=True).to("cuda")
+    with torch.no_grad():
+        out = router.generate(**enc, max_new_tokens=96, do_sample=False)
+    routes += rtok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
+del router
+s_correct, s_cost = 0, 0.0
+for t, rt in zip(test_tasks, routes):
+    try:
+        agents = [x for x in json.loads(rt)["agents"] if x in POOL] or [STRONGEST]
+    except Exception:  # noqa: BLE001
+        agents = [STRONGEST]
+    rec = json.loads(shak[t]["agent_results"])
+    for ag in agents:  # fresh sample where one exists, otherwise the recorded run
+        fresh = P.get((t, ag, 0))
+        s_cost += POOL[ag]["cost"] * (fresh["completion_tokens"] if fresh else rec[ag]["completion_tokens"])
+    fin = P.get((t, agents[-1], 0))  # ensembles are scored as the aggregating agent's own answer (conservative)
+    s_correct += fin["correct"] if fin else rec[agents[-1]]["correct"]
+policies["shaktimaan-v1.0"] = {"accuracy": round(s_correct / len(test_tasks), 4),
+                               "relative_cost": round(s_cost / ultra_cost, 4)}
+
+curve = []
+for tau in [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.85, 0.9, 0.95, 0.98]:
+    curve.append(dict(threshold=tau, **rel(run_ladder(test_tasks, LADDER, jedi(tau), ultra_record))))
+
+# two-rung ladder (no Ultra record) on the new domains and on unseen BIG-Bench Hard
+two_rung = {}
+for name, split, srcs in [("new-domains (aqua, csqa)", "test", {"aqua", "csqa"}), ("ood: bbh", "ood", {"bbh"})]:
+    ids = ladder_tasks(split, srcs)
+    base = sum(POOL[MID]["cost"] * P[t, MID, 0]["completion_tokens"] for t in ids)
+    r = lambda res: {"accuracy": round(res["accuracy"], 4), "cost_vs_always_mid": round(res["cost"] / base, 4),
+                     "silent_failure_rate": round(res["silent_failure_rate"], 4)}
+    two_rung[name] = {"n_tasks": len(ids),
+                      f"always:{CHEAP}": r(run_ladder(ids, [CHEAP], None)),
+                      f"always:{MID}": r(run_ladder(ids, [MID], None)),
+                      "oracle": r(run_ladder(ids, [CHEAP, MID], oracle)),
+                      "self-consistency": r(run_ladder(ids, [CHEAP, MID], self_consistency)),
+                      **{f"jedi@{tau}": r(run_ladder(ids, [CHEAP, MID], jedi(tau))) for tau in (0.5, 0.7, 0.9)}}
+
+# overhead: one check at a time, as a live ladder would issue it
+views = ds["test"]["judge_input"][:200]
+torch.cuda.synchronize()
+t0 = time.time()
+for v in views:
+    logit_gap([v], bs=1)
+torch.cuda.synchronize()
+n_tok = np.mean([len(tok(tok.apply_chat_template([{"role": "system", "content": JEDI_SYSTEM},
+                                                  {"role": "user", "content": v}], add_generation_prompt=True,
+                                                 tokenize=False)).input_ids) for v in views])
+
+report.update({"ladder": LADDER, "n_ladder_tasks": len(test_tasks), "policies": policies, "jedi_curve": curve,
+               "two_rung": two_rung,
+               "overhead": {"jedi_ms_per_check_batch1": round((time.time() - t0) / len(views) * 1000, 1),
+                            "jedi_input_tokens_mean": round(float(n_tok)), "jedi_output_tokens": 1}})
+print(json.dumps({k: v for k, v in report.items() if k != "verifier"}, indent=2))
+json.dump(report, open(f"{path}/eval_results.json", "w"), indent=2)
+json.dump(cal, open(f"{path}/calibration.json", "w"), indent=2)

gen.py

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

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"
+}

jedi.py

@@ -0,0 +1,99 @@
+"""Jedi runtime: answer with the cheapest agent, let the 0.5B verifier check it, escalate up the ladder if needed.
+
+Usage:
+  JEDI_ENDPOINTS='{"glm-5.3-flash": "http://host:8000/v1", ...}' python jedi.py [--threshold 0.7] [--code] "task"
+
+The threshold is the single cost-vs-quality dial: raise it for accuracy (more escalations), lower it for savings.
+Prints the final answer; the per-rung trace goes to stderr as JSON lines like
+  {"agent": "glm-5.3-flash", "verdict": "escalate", "confidence": 0.87, "escalate_to": "qwen3.8-27b"}
+"""
+import argparse, json, math, os, re, sys
+import httpx
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from common import CODE_SYSTEM, LADDER, POOL, JEDI_SYSTEM, SOLVER_SYSTEM, endpoints, est_tokens, judge_view, max_out
+
+MODEL = os.environ.get("JEDI_MODEL", "podstack/Jedi")
+
+
+class Jedi:
+    def __init__(self, path=MODEL):
+        self.tok = AutoTokenizer.from_pretrained(path)
+        self.model = AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).eval()
+        if torch.cuda.is_available():
+            self.model.cuda()
+        try:
+            from huggingface_hub import hf_hub_download
+            cal_file = os.path.join(path, "calibration.json") if os.path.isdir(path) else \
+                hf_hub_download(path, "calibration.json")
+            cal = json.load(open(cal_file))
+        except Exception:  # noqa: BLE001
+            cal = {"a": 1.0, "b": 0.0}
+        self.a, self.b = cal["a"], cal["b"]
+        self.ids = [self.tok.encode(w, add_special_tokens=False)[0] for w in ("accept", "escalate")]
+
+    def p_correct(self, task, answer, agent, kind, finish=None):
+        """Calibrated probability that `answer` (written by `agent`) is correct."""
+        enc = self.tok.apply_chat_template(
+            [{"role": "system", "content": JEDI_SYSTEM},
+             {"role": "user", "content": judge_view(task, answer, agent, kind, finish)}],
+            add_generation_prompt=True, return_tensors="pt", return_dict=True).to(self.model.device)
+        with torch.no_grad():
+            logits = self.model(**enc).logits[0, -1].float()
+        gap = (logits[self.ids[0]] - logits[self.ids[1]]).item()
+        return 1 / (1 + math.exp(-(self.a * gap + self.b)))
+
+
+def ask(agent, task, code):
+    base = endpoints()[agent].rstrip("/")
+    msgs = [{"role": "system", "content": CODE_SYSTEM if code else SOLVER_SYSTEM}, {"role": "user", "content": task}]
+    r = httpx.post(f"{base}/chat/completions", timeout=1800, json={
+        "model": POOL[agent]["model"], "messages": msgs, "max_tokens": max_out(agent, est_tokens(task))})
+    r.raise_for_status()
+    c = r.json()["choices"][0]
+    return c["message"].get("content") or "", c.get("finish_reason")
+
+
+def infer_kind(task, code):
+    """Task type shown to Jedi (its training data used: code, mc, num, math, exact)."""
+    if code:
+        return "code"
+    return "mc" if re.search(r"^\(?[A-J][.)]\s", task, flags=re.M) else "math"
+
+
+def solve(task, threshold=0.7, code=False, kind=None, judge=None):
+    """Run the ladder. Returns (final answer text, trace)."""
+    judge = judge or Jedi()
+    kind = kind or infer_kind(task, code)
+    rungs = [a for a in LADDER if a in endpoints()]
+    trace, text = [], ""
+    for i, agent in enumerate(rungs):
+        try:
+            text, finish = ask(agent, task, code)
+        except httpx.HTTPError as e:
+            if i == 0:
+                raise
+            trace.append({"agent": agent, "verdict": "unavailable", "error": repr(e)[:120]})
+            return text, trace  # keep the best answer so far
+        if i == len(rungs) - 1:
+            trace.append({"agent": agent, "verdict": "final", "confidence": None, "escalate_to": None})
+            return text, trace
+        p = judge.p_correct(task, text, agent, kind, finish)
+        verdict = "accept" if p >= threshold else "escalate"
+        trace.append({"agent": agent, "verdict": verdict, "confidence": round(p if verdict == "accept" else 1 - p, 3),
+                      "escalate_to": None if verdict == "accept" else rungs[i + 1]})
+        if verdict == "accept":
+            return text, trace
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("task", nargs="+")
+    ap.add_argument("--threshold", type=float, default=0.7, help="min P(correct) to accept an answer")
+    ap.add_argument("--code", action="store_true", help="ask agents for a Python code solution")
+    ap.add_argument("--kind", choices=["math", "num", "mc", "code", "exact"], help="task type (default: inferred)")
+    args = ap.parse_args()
+    answer, trace = solve(" ".join(args.task), args.threshold, args.code, args.kind)
+    for step in trace:
+        print(json.dumps(step), file=sys.stderr)
+    print(answer)

model.safetensors

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

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,18 @@
+"""Stage 3: full fine-tune of Qwen2.5-0.5B-Instruct into Jedi (loss on the one-word verdict 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/jedi_judge").select_columns(["prompt", "completion"])
+
+cfg = SFTConfig(
+    output_dir="ckpt", num_train_epochs=2, per_device_train_batch_size=8, gradient_accumulation_steps=4,
+    learning_rate=2e-5, lr_scheduler_type="cosine", warmup_steps=30, weight_decay=0.0, bf16=True,
+    max_length=2048, completion_only_loss=True, logging_steps=20, 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["val"])
+trainer.train()
+trainer.save_model("jedi-0.5b")
+trainer.processing_class.save_pretrained("jedi-0.5b")