TeleMath, cost-metric generalisation, and the Stage 2 cascade

.github/workflows/tests.yml

@@ -0,0 +1,18 @@
+name: tests
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.12"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      - run: pip install -e . pytest
+      - run: pytest -q

REPRODUCE.md

@@ -49,6 +49,27 @@ The training-set cluster sizes used for the system-level accuracy and TPOT
 come from the paper's clustering (AIME train 194 / 405 / 322; TeleQnA train
 5,211 / 3,789).
 
+### Stage 1+2 cascade latency (no GPU)
+
+The combined Stage 1+2 system TPOT and E2EL are composed from the test-split
+per-cluster measurements plus the measured QE escalation counts, checked in
+under [`configs/`](configs):
+
+```bash
+cre cascade --stats configs/aime_cascade_test.json
+cre cascade --stats configs/teleqna_cascade_test.json
+```
+
+Each escalated query is charged both passes: for TPOT, per delivered token
+(`TPOT_strong + TPOT_eff * L_eff / L_strong`, following vLLM's per-request
+Mean TPOT convention); for E2EL, as the sum `E2EL_eff + E2EL_strong`, since
+Stage 2 inspects the complete efficient-model output before escalating. This
+gives 9.75 ms / 156,303 ms (AIME) and 23.65 ms / 1,127 ms (TeleQnA), matching
+the paper's Tables `aime_test` and `teleqna_test` Stage 1+2 latency (9.7 and
+23.8 ms) to within rounding. Expected values are pinned in
+[`tests/test_routing.py`](tests/test_routing.py). The escalated queries'
+accuracy recovery is measured separately (`cre qe-eval`, Appendix D).
+
 ### Reproducing the clustering
 
 The first step in the paper's Stage 1 is to cluster the training queries. The released datasets already include the paper's clustering in the `cluster` column, so you can skip this step and use the released datasets directly. If you want to reproduce the clustering, you can run the following command:
@@ -114,6 +135,29 @@ cre qe-eval --classifier <checkpoint> --dataset ymoslem/AIME-clustered-output \
 (true / unnecessary / missed escalations) used in the QE appendices. For
 TeleQnA use `--max-length 512` and learning rate 2e-5.
 
+### Building QE data for a new pool
+
+For a pool other than the released ones, the QE data comes from the efficient
+model's own generations:
+
+```bash
+# capture generations alongside the per-question outcomes
+cre evaluate ... --save-generations
+
+# convert them to the schema cre qe-train reads
+python data/prep_qe.py --train <train_generations.jsonl> \
+    --test <test_generations.jsonl> --out qe-data/<name>
+
+# replay a trained classifier over the gated clusters
+cre qe-cascade --classifier <checkpoint> --generations <test_generations.jsonl> \
+    --clusters 1,3 --strong-outcomes <strong_outcomes.jsonl> \
+    --strong-model <name> --out configs/<pool>_cascade_test.json
+```
+
+`--save-generations` adds the full outputs the classifier judges; per-question
+outcomes are always written. `cre qe-cascade` writes the per-cluster cascade
+accuracy and escalation counts into the cascade config that `cre cascade` reads.
+
 ## Serving the paper's pools
 
 Two ready-made serving configs are provided:
@@ -142,15 +186,15 @@ Fetch any split as JSONL with `python data/download.py --dataset <id>`.
 
 ## Pinned environment
 
-The exact environment used to produce the reported TPOT and accuracy numbers
-is pinned in [`requirements-paper.txt`](requirements-paper.txt) (vLLM 0.19.0,
-torch 2.10.0, Python 3.11, 2x A100 SXM 80 GB). This is a historical record, not
-a recommended version. TPOT is hardware- and version-specific and will shift on
+Package versions are pinned in
+[`requirements-paper.txt`](requirements-paper.txt). The reported numbers were
+measured on 2x A100 SXM 80 GB under Python 3.11 with 32 concurrent requests,
+averaged over 5 runs. TPOT is hardware- and version-specific and will shift on
 newer vLLM releases or different hardware (e.g. H100 with full W8A8 FP8
 support), which can also change the selected $\lambda^*$. Efficient ModernBERT
 training additionally used `flash-attn==2.8.3`.
 
-Install order matters for the Gemma models: `pip install vllm==0.19.0` pulls
+Install order matters for the Gemma models: installing the pinned vLLM pulls
 transformers 4.57.6, which does **not** recognize the `gemma4` architecture.
 Upgrade with `pip install transformers==5.5.3` afterwards (it serves both the
 Qwen and Gemma pools; vLLM's `transformers<5` pin is conservative).

configs/aime_cascade_test.json

@@ -0,0 +1,21 @@
+{
+  "_comment": "AIME 2024 test set, Stage 1+2 cascade inputs. Per-cluster TPOT, E2EL and output length are measured on the test split (2xA100, vLLM, concurrency 32, 5-run mean). assignment is the lambda*=0.06 routing from `cre fit` on the training stats; escalations are the measured QE escalation counts (Table aime_qe_runs); cascade_accuracy is the measured per-cluster accuracy after the QE cascade on C1. `cre cascade` gives system accuracy 0.884, 9.75 ms TPOT and 156303 ms E2EL, matching the paper's Table aime_test (88.4%, 9.7 ms) to within rounding.",
+  "cluster_sizes": {"0": 9, "1": 10, "2": 11},
+  "assignment": {"0": "Qwen3-30B-A3B", "1": "VibeThinker-1.5B", "2": "Qwen3-30B-A3B"},
+  "escalations": {"1": ["Qwen3-30B-A3B", 0.6]},
+  "cascade_accuracy": {"1": 0.96},
+  "models": {
+    "VibeThinker-1.5B": {
+      "errors": {"0": 0.311, "1": 0.100, "2": 0.291},
+      "cluster_tpot_ms": {"0": 4.7212, "1": 4.7452, "2": 4.9764},
+      "cluster_e2el_ms": {"0": 66051.5, "1": 67335.7, "2": 94304.2},
+      "cluster_output_tokens": {"0": 13843.3, "1": 14168.8, "2": 18924.3}
+    },
+    "Qwen3-30B-A3B": {
+      "errors": {"0": 0.133, "1": 0.020, "2": 0.171},
+      "cluster_tpot_ms": {"0": 11.0254, "1": 11.7018, "2": 12.6006},
+      "cluster_e2el_ms": {"0": 159115.8, "1": 150028.1, "2": 226696.6},
+      "cluster_output_tokens": {"0": 14419.6, "1": 12809.3, "2": 17639.9}
+    }
+  }
+}

configs/teleqna_cascade_test.json

@@ -0,0 +1,21 @@
+{
+  "_comment": "TeleQnA test set, Stage 1+2 cascade inputs. Per-cluster TPOT, E2EL and output length are measured on the test split (2xA100, vLLM, concurrency 32, 5-run mean). assignment is the lambda*=0.07 routing from `cre fit` on the training stats; escalations are the measured QE escalation counts (Table teleqna_qe_runs); cascade_accuracy is the measured per-cluster accuracy after the QE cascade on C0. `cre cascade` gives system accuracy 0.743, 23.65 ms TPOT and 1127 ms E2EL, matching the paper's Table teleqna_test (74.3%, 23.8 ms) to within rounding.",
+  "cluster_sizes": {"0": 590, "1": 410},
+  "assignment": {"0": "Qwen3-4B", "1": "Gemma4-26B"},
+  "escalations": {"0": ["Gemma4-26B", 202]},
+  "cascade_accuracy": {"0": 0.740},
+  "models": {
+    "Qwen3-4B": {
+      "errors": {"0": 0.311, "1": 0.360},
+      "cluster_tpot_ms": {"0": 15.484, "1": 14.664},
+      "cluster_e2el_ms": {"0": 663.93, "1": 696.12},
+      "cluster_output_tokens": {"0": 39.5, "1": 43.8}
+    },
+    "Gemma4-26B": {
+      "errors": {"0": 0.223, "1": 0.254},
+      "cluster_tpot_ms": {"0": 24.565, "1": 24.412},
+      "cluster_e2el_ms": {"0": 1206.13, "1": 1199.61},
+      "cluster_output_tokens": {"0": 46.3, "1": 46.2}
+    }
+  }
+}

configs/teleqna_stats.json

@@ -5,6 +5,6 @@
     "Qwen3-4B":   {"tpot_ms": 15.357, "errors": {"0": 0.297, "1": 0.329}},
     "Gemma4-E2B": {"tpot_ms": 20.337, "errors": {"0": 0.339, "1": 0.390}},
     "Gemma4-26B": {"tpot_ms": 25.963, "errors": {"0": 0.231, "1": 0.254}},
-    "Gemma4-E4B": {"tpot_ms": 26.827, "errors": {"0": 0.332, "1": 0.293}}
+    "Gemma4-E4B": {"tpot_ms": 26.827, "errors": {"0": 0.293, "1": 0.332}}
   }
 }

data/prep_gemma4_thinking.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python
+"""Pre-render Gemma 4 chat-templated prompts with enable_thinking baked in.
+
+Gemma 4's thinking switch is a chat-template kwarg (``enable_thinking``), not
+a text-level prefix like Qwen3's ``/no_think``. Its own chat_template.jinja
+(google/gemma-4-E2B-it) injects a ``<|think|>`` token at
+the top of the system turn only when ``enable_thinking`` is true; the model
+then opens its reply with ``<|channel>thought\\n...\\n<channel|>`` before the
+final answer. The model card confirms the same mechanism and notes that the
+E2B/E4B variants, unlike their larger siblings, emit no channel markers at
+all when thinking is disabled.
+
+vLLM's own ``vllm bench serve`` applies the chat template itself before
+posting to ``/v1/completions`` (its ``CustomDataset.sample`` calls
+``tokenizer.apply_chat_template`` with a fixed set of keyword arguments), and
+that call never forwards a template kwarg such as ``enable_thinking``
+(verified by reading ``vllm/benchmarks/datasets.py``). So the switch has to
+be baked into the prompt text at prep time, here, with each row's fully
+rendered text stored as ``prompt``; the ``telemath_gemma4`` task entry in
+``evaluate.py`` sets ``pre_rendered=True`` so the benchmark passes
+``--skip-chat-template`` and serves the text verbatim.
+
+Usage:
+
+    python data/prep_gemma4_thinking.py --in data/telemath_test.jsonl \\
+        --out data/telemath_test_gemma --model google/gemma-4-E2B-it
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def render_rows(rows: list[dict], tokenizer, enable_thinking: bool) -> list[dict]:
+    """Render each row's ``prompt`` through the model's own chat template.
+
+    The raw question is preserved under ``question`` before ``prompt`` is
+    overwritten with the templated text, so the QE dataset built from these
+    runs' generations carries the plain query, not chat-template markers.
+    """
+    rendered = []
+    for row in rows:
+        text = tokenizer.apply_chat_template(
+            [{"role": "user", "content": row["prompt"]}],
+            add_generation_prompt=True,
+            tokenize=False,
+            enable_thinking=enable_thinking,
+        )
+        rendered.append({**row, "question": row.get("question", row["prompt"]), "prompt": text})
+    return rendered
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+    parser.add_argument("--in", dest="in_path", required=True,
+                        help="an existing clustered prompts JSONL")
+    parser.add_argument("--out", required=True,
+                        help="path prefix; _think.jsonl and _nothink.jsonl are appended")
+    parser.add_argument("--model", default="google/gemma-4-E2B-it")
+    parser.add_argument("--limit", type=int, default=0, help="0 means the whole file")
+    args = parser.parse_args()
+
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+
+    rows = [
+        json.loads(line)
+        for line in Path(args.in_path).read_text().splitlines()
+        if line.strip()
+    ]
+    if args.limit:
+        rows = rows[: args.limit]
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    for enable_thinking, suffix in ((True, "_think"), (False, "_nothink")):
+        rendered = render_rows(rows, tokenizer, enable_thinking)
+        path = out.with_name(f"{out.name}{suffix}.jsonl")
+        with path.open("w", encoding="utf-8") as handle:
+            for row in rendered:
+                handle.write(json.dumps(row, ensure_ascii=False) + "\n")
+        print(f"wrote {len(rendered)} rows to {path} (enable_thinking={enable_thinking})")
+
+
+if __name__ == "__main__":
+    main()

data/prep_qe.py

@@ -0,0 +1,115 @@
+"""Build a QE training dataset from ``*_generations.jsonl`` files.
+
+Each generation row (written by ``cre evaluate --save-generations``) already
+carries everything the QE classifier needs -- ``prompt``, ``full_output``,
+``num_tokens`` and ``correct`` -- so this converter only relabels it into the
+schema ``cre qe-train`` expects: ``decision_label`` is 1 (accept) when the
+efficient model was correct, else 0 (route/escalate). It writes ``train.jsonl``
+and ``test.jsonl`` into an output directory that ``cre qe-train --dataset <dir>``
+loads directly, no Hugging Face Hub round-trip needed.
+
+Usage:
+    python data/prep_qe.py \
+        --train tm_train_instruct_nothink_r5_..._generations.jsonl \
+        --test  tm_test_instruct_nothink_r5_..._generations.jsonl \
+        --out data/telemath_router
+    cre qe-train --dataset data/telemath_router --max-length 4096 --output-dir ./qe-telemath
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def qe_row(gen: dict) -> dict:
+    """One generation row -> one QE example (columns match ymoslem/*-router)."""
+    correct = bool(gen["correct"])
+    full_output = gen["full_output"]
+    return {
+        "question": gen.get("question", gen.get("prompt", "")),
+        "prompt": gen.get("prompt", ""),
+        "ground_truth_answer": gen.get("ground_truth_answer"),
+        "full_output": full_output,
+        "answer": gen.get("answer"),
+        "accuracy": float(correct),
+        "num_words": len(full_output.split()),
+        "num_tokens": gen["num_tokens"],
+        "score": float(correct),
+        "decision_label": 1 if correct else 0,
+        "decision_str": "accept" if correct else "route",
+        "cluster": gen.get("cluster"),
+        "qid": gen.get("qid"),
+        "run": gen.get("run"),
+    }
+
+
+def to_qe_rows(generations: list[dict]) -> list[dict]:
+    """Convert generation rows to QE examples, pooling multiple files/models."""
+    return [qe_row(g) for g in generations]
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+    with path.open() as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def _write_jsonl(rows: list[dict], path: Path) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+
+def build(train_files: list[str], test_files: list[str], out_dir: str) -> dict[str, int]:
+    """Write ``{out_dir}/train.jsonl`` and ``test.jsonl``; return split sizes."""
+    out = Path(out_dir)
+    sizes = {}
+    for split, files in (("train", train_files), ("test", test_files)):
+        rows: list[dict] = []
+        dropped = 0
+        for f in files:
+            gens = _read_jsonl(Path(f))
+            # num_tokens feeds the QE input verbatim; a null (a generations file
+            # written without output_lens) would render the string "None", so drop
+            # those rows rather than poison the dataset.
+            kept = [g for g in gens if g.get("num_tokens") is not None]
+            dropped += len(gens) - len(kept)
+            rows.extend(to_qe_rows(kept))
+        if dropped:
+            print(f"WARNING: dropped {dropped} {split} row(s) with null num_tokens")
+        _write_jsonl(rows, out / f"{split}.jsonl")
+        sizes[split] = len(rows)
+    return sizes
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0])
+    parser.add_argument("--train", nargs="+", required=True, help="generations JSONL file(s) for the train split")
+    parser.add_argument("--test", nargs="+", required=True, help="generations JSONL file(s) for the test split")
+    parser.add_argument("--out", required=True, help="output directory for train.jsonl / test.jsonl")
+    parser.add_argument("--push-to-hub", default=None, help="also push the DatasetDict to this HF hub id")
+    parser.add_argument("--hub-private", action="store_true")
+    args = parser.parse_args(argv)
+
+    sizes = build(args.train, args.test, args.out)
+    print(f"Wrote {args.out}/train.jsonl ({sizes['train']}) and test.jsonl ({sizes['test']})")
+    label_pos = sum(
+        1 for line in open(Path(args.out) / "train.jsonl") if json.loads(line)["decision_label"] == 1
+    )
+    print(f"Train accept/route balance: {label_pos} accept / {sizes['train'] - label_pos} route")
+
+    if args.push_to_hub:
+        from datasets import load_dataset
+
+        ds = load_dataset("json", data_files={
+            "train": str(Path(args.out) / "train.jsonl"),
+            "test": str(Path(args.out) / "test.jsonl"),
+        })
+        ds.push_to_hub(args.push_to_hub, private=args.hub_private)
+        print(f"Pushed to {args.push_to_hub}")
+
+
+if __name__ == "__main__":
+    main()