diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..5da9d21 --- /dev/null +++ b/.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 diff --git a/REPRODUCE.md b/REPRODUCE.md index 6505716..2b5a380 100644 --- a/REPRODUCE.md +++ b/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 --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 \ + --test --out qe-data/ + +# replay a trained classifier over the gated clusters +cre qe-cascade --classifier --generations \ + --clusters 1,3 --strong-outcomes \ + --strong-model --out configs/_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 `. ## 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). diff --git a/configs/aime_cascade_test.json b/configs/aime_cascade_test.json new file mode 100644 index 0000000..7c6c069 --- /dev/null +++ b/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} + } + } +} diff --git a/configs/teleqna_cascade_test.json b/configs/teleqna_cascade_test.json new file mode 100644 index 0000000..cd349ea --- /dev/null +++ b/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} + } + } +} diff --git a/configs/teleqna_stats.json b/configs/teleqna_stats.json index d9afc47..d7a1d3a 100644 --- a/configs/teleqna_stats.json +++ b/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}} } } diff --git a/data/prep_gemma4_thinking.py b/data/prep_gemma4_thinking.py new file mode 100644 index 0000000..c7af013 --- /dev/null +++ b/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`` 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() diff --git a/data/prep_qe.py b/data/prep_qe.py new file mode 100644 index 0000000..fc0d2df --- /dev/null +++ b/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 `` +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() diff --git a/data/prep_telemath.py b/data/prep_telemath.py new file mode 100644 index 0000000..2fb0520 --- /dev/null +++ b/data/prep_telemath.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python +"""Prepare TeleMath (netop/TeleMath) for the CRE pipeline. + +Downloads the 500-item TeleMath benchmark, telecom mathematical problems with +numerical answers, appends a numerical-answer prompt suffix, and writes a +stratified-by-category train/test split (approx 300/200). TeleMath ships a +single ``test`` split with no train partition, so the split is created here and +pinned by seed. + +TeleMath is a gated dataset; set ``HF_TOKEN`` in the environment before running. + +Usage: + HF_TOKEN=... python data/prep_telemath.py --out data/telemath +""" + +from __future__ import annotations + +import argparse +import json +import random +from collections import Counter, defaultdict +from pathlib import Path + +# The questions already name the quantity and its unit, so the suffix only fixes +# the output format: a single numerical answer, in the stated unit, in a box. +PROMPT_SUFFIX = ( + "\n\nSolve the problem and give the final numerical answer, in the unit " + "stated in the question, inside \\boxed{}." +) + + +def build_prompt(row: dict) -> str: + return row["question"].strip() + PROMPT_SUFFIX + + +def stratified_split(rows: list[dict], train_size: int, seed: int) -> tuple[list, list]: + """Split rows into train/test, stratified by category so every topic is + represented on both sides in proportion to its frequency.""" + frac = train_size / len(rows) + by_cat: dict[str, list] = defaultdict(list) + for r in rows: + by_cat[r["category"]].append(r) + rng = random.Random(seed) + train, test = [], [] + for _, items in sorted(by_cat.items()): + items = items[:] + rng.shuffle(items) + n_train = round(len(items) * frac) + train.extend(items[:n_train]) + test.extend(items[n_train:]) + return train, test + + +def write_split(rows: list[dict], path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("w", encoding="utf-8") as handle: + for i, row in enumerate(rows): + record = { + "id": row["id"], + "prompt": build_prompt(row), + "answer": row["answer"], + "category": row["category"], + "difficulty": row["difficulty"], + } + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + + +def describe(rows: list[dict], label: str) -> None: + cats = Counter(r["category"] for r in rows) + print(f"{label}: {len(rows)} rows | {dict(cats.most_common())}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument("--out", default="data/telemath", help="output path prefix") + parser.add_argument("--train-size", type=int, default=300) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--split", default="test", + help="TeleMath ships a single 'test' split") + args = parser.parse_args() + + from datasets import load_dataset + + dataset = list(load_dataset("netop/TeleMath")[args.split]) + rows = [dict(row, id=i) for i, row in enumerate(dataset)] + + train, test = stratified_split(rows, args.train_size, args.seed) + out = Path(args.out) + write_split(train, out.with_name(f"{out.name}_train.jsonl")) + write_split(test, out.with_name(f"{out.name}_test.jsonl")) + + describe(train, "train") + describe(test, "test") + print(f"-> {out}_train.jsonl / {out}_test.jsonl") + + +if __name__ == "__main__": + main() diff --git a/src/cre_router/cli.py b/src/cre_router/cli.py index 59596a2..d6a8518 100644 --- a/src/cre_router/cli.py +++ b/src/cre_router/cli.py @@ -21,6 +21,8 @@ from cre_router.artifacts import RouterArtifacts from cre_router.clustering import DEFAULT_EMBEDDING_MODEL from cre_router.routing import ( + cascade_system_accuracy, + cascade_system_metrics, eta, models_from_stats, pareto_prune, @@ -83,13 +85,18 @@ def cmd_cluster(args: argparse.Namespace) -> None: def cmd_fit(args: argparse.Namespace) -> None: stats = json.loads(Path(args.stats).read_text()) - models, cluster_sizes = models_from_stats(stats) + models, cluster_sizes = models_from_stats(stats, args.cost_metric) efficient, dominated = pareto_prune(models) - print("Pareto analysis:") - for m in sorted(models, key=lambda m: m.tpot_ms): + label = args.cost_metric.upper() + # eta is accuracy points per millisecond of cost. That reads well for TPOT + # (single-digit ms) but collapses to 0.00 for E2EL, whose costs run to + # hundreds of thousands of ms, so E2EL is reported per second instead. + eta_scale, eta_unit = (1000.0, "pp/s") if args.cost_metric == "e2el" else (1.0, "pp/ms") + print(f"Pareto analysis (cost = {label}):") + for m in sorted(models, key=lambda m: m.cost_ms): status = "dominated" if m in dominated else "efficient" - print(f" {m.name:<24} TPOT {m.tpot_ms:7.3f} ms {status}") + print(f" {m.name:<24} {label} {m.cost_ms:10.3f} ms {status}") if dominated: print(f"Pruned {len(dominated)} dominated model(s); " f"routing over: {[m.name for m in efficient]}") @@ -97,21 +104,23 @@ def cmd_fit(args: argparse.Namespace) -> None: clusters = sorted(cluster_sizes) print(f"\nRouting regions (lambda sweep) over clusters {clusters}:") header = f" {'lambda range':>16} " + " ".join(f"{('C' + c):>24}" for c in clusters) \ - + f" {'Acc':>7} {'TPOT':>8} {'eta':>6}" + + f" {'Acc':>7} {label:>11} {('eta ' + eta_unit):>10}" print(header) for region in routing_regions(efficient): - acc, tpot = system_metrics(efficient, region.assignment, cluster_sizes) + acc, cost = system_metrics(efficient, region.assignment, cluster_sizes) e = eta(efficient, region.assignment, cluster_sizes) row = f" {region.interval_str:>16} " + " ".join( f"{region.assignment[c]:>24}" for c in clusters - ) + f" {acc:>6.1%} {tpot:>6.1f}ms " + (f"{e:>6.2f}" if e is not None else " ---") + ) + f" {acc:>6.1%} {cost:>9.1f}ms " + ( + f"{e * eta_scale:>10.3f}" if e is not None else f"{'---':>10}") print(row) selection = select_lambda(efficient, cluster_sizes, args.budget) - print(f"\nBudget B = {args.budget} ms -> lambda* = {selection.lambda_star}") + print(f"\nBudget B = {args.budget} ms {label} -> lambda* = {selection.lambda_star}") print(f" assignment: {selection.region.assignment}") - print(f" training accuracy {selection.accuracy:.1%} at {selection.tpot_ms:.1f} ms TPOT" - + (f", eta {selection.eta:.2f} pp/ms" if selection.eta is not None else "")) + print(f" training accuracy {selection.accuracy:.1%} at {selection.tpot_ms:.1f} ms {label}" + + (f", eta {selection.eta * eta_scale:.3f} {eta_unit}" + if selection.eta is not None else "")) if args.output: out_dir = Path(args.output) @@ -128,6 +137,37 @@ def cmd_fit(args: argparse.Namespace) -> None: print(f"Saved routing table to {out_dir}/router.json") +def cmd_cascade(args: argparse.Namespace) -> None: + stats = json.loads(Path(args.stats).read_text()) + models, cluster_sizes = models_from_stats(stats) + assignment = {str(k): str(v) for k, v in stats["assignment"].items()} + escalations = { + str(k): (str(v[0]), float(v[1])) for k, v in stats.get("escalations", {}).items() + } + tpot, e2el = cascade_system_metrics(models, assignment, cluster_sizes, escalations) + # Per-cluster cascade accuracy (efficient outputs gated by the QE classifier, + # rejects escalated to the strong model), produced by the QE cascade step. + # Clusters absent route entirely to their Stage 1 model. + cascade_accuracy = {str(k): float(v) for k, v in stats.get("cascade_accuracy", {}).items()} + + clusters = sorted(cluster_sizes) + print(f"Stage 1+2 cascade over clusters {clusters}:") + for c in clusters: + line = f" C{c} ({int(cluster_sizes[c])} queries) -> {assignment[c]}" + if c in escalations: + model, count = escalations[c] + line += f", escalate {count:g} -> {model}" + if c in cascade_accuracy: + line += f" (cascade acc {cascade_accuracy[c]:.3f})" + print(line) + + stage1_acc, _ = system_metrics(models, assignment, cluster_sizes) + system_acc = cascade_system_accuracy(models, assignment, cluster_sizes, cascade_accuracy) + print(f"\n system accuracy: {system_acc:.3f} (Stage 1 alone: {stage1_acc:.3f})") + print(f" system TPOT: {tpot:.2f} ms") + print(f" system E2EL: {e2el:.0f} ms") + + def cmd_evaluate(args: argparse.Namespace) -> None: from cre_router.evaluate import ( TASKS, @@ -159,6 +199,15 @@ def cmd_evaluate(args: argparse.Namespace) -> None: f"benchmark at {args.host}:{args.port}" ) + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + safe_model = args.model.replace("/", "_") + stem = f"{args.task}_{safe_model}_{timestamp}" + outcomes_path = Path(args.results_dir) / f"{stem}_outcomes.jsonl" + generations_path = ( + Path(args.results_dir) / f"{stem}_generations.jsonl" + if getattr(args, "save_generations", False) + else None + ) measurements = evaluate_model( dataset, model=args.model, @@ -170,11 +219,11 @@ def cmd_evaluate(args: argparse.Namespace) -> None: base_seed=args.seed, workdir=Path(args.results_dir) / "splits", download_dir=args.download_dir, + outcomes_out=outcomes_path, + generations_out=generations_path, ) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - safe_model = args.model.replace("/", "_") - raw_path = Path(args.results_dir) / f"{args.task}_{safe_model}_{timestamp}.jsonl" + raw_path = Path(args.results_dir) / f"{stem}.jsonl" save_raw_measurements(measurements, raw_path) entry = model_entry(measurements) @@ -182,12 +231,22 @@ def cmd_evaluate(args: argparse.Namespace) -> None: merge_model_into_stats(args.stats_out, args.model, entry, sizes) print(f"\nPer-cluster results for {args.model}:") + e2el = entry.get("cluster_e2el_ms", {}) + trunc = entry.get("cluster_truncated_frac", {}) for c in sorted(entry["errors"]): - print( + line = ( f" C{c}: error {entry['errors'][c]:.3f} " f"TPOT {entry['cluster_tpot_ms'][c]:.3f} ms" ) + if c in e2el: + line += f" E2EL {e2el[c]:.0f} ms" + if c in trunc: + line += f" trunc {trunc[c]:.2f}" + print(line) print(f"\nRaw measurements: {raw_path}") + print(f"Per-question: {outcomes_path}") + if generations_path is not None: + print(f"Generations: {generations_path}") print(f"Updated stats: {args.stats_out}") @@ -212,6 +271,44 @@ def cmd_qe_train(args: argparse.Namespace, extra: list[str]) -> None: qe_train_main(extra) +def cmd_qe_cascade(args: argparse.Namespace) -> None: + from cre_router.qe import QEClassifier + from cre_router.qe.cascade import ( + compose_cascade, + run_qe, + strong_correct_by_qid, + write_cascade_stats, + ) + + generations = _read_jsonl(Path(args.generations)) + if args.clusters: + keep = set(args.clusters.split(",")) + generations = [g for g in generations if str(g["cluster"]) in keep] + if not generations: + raise SystemExit("no generations to score (check --generations and --clusters)") + strong = strong_correct_by_qid(_read_jsonl(Path(args.strong_outcomes))) + + classifier = QEClassifier( + model_name=args.classifier, + base_tokenizer=args.base_tokenizer, + accept_threshold=args.accept_threshold, + max_length=args.max_length, + ) + escalate = run_qe(classifier, generations, batch_size=args.batch_size) + report = compose_cascade(generations, escalate, strong) + + print(f"QE cascade over {len(generations)} generations, escalating to {args.strong_model}:") + for cluster in sorted(report): + r = report[cluster] + print( + f" C{cluster} (n={r['n']}): cascade acc {r['cascade_accuracy']:.3f}, " + f"escalate {r['escalations']:g}/run" + ) + if args.out: + write_cascade_stats(args.out, args.strong_model, report) + print(f"\nUpdated {args.out} (escalations + cascade_accuracy); run `cre cascade --stats {args.out}`") + + def cmd_serve(args: argparse.Namespace) -> None: from cre_router.server.app import serve @@ -254,14 +351,34 @@ def main(argv: list[str] | None = None) -> None: p.add_argument("--seed", type=int, default=0, help="base seed; run r uses seed+r") p.add_argument("--download-dir", default=None, help="vLLM model download/cache directory") p.add_argument("--results-dir", default="results", help="where to write raw measurements and cluster splits") + p.add_argument( + "--save-generations", + action="store_true", + help="also write per-question full_output + num_tokens (QE training data); off by default", + ) p.set_defaults(func=cmd_evaluate) p = sub.add_parser("fit", help="compute the routing table from model stats") p.add_argument("--stats", required=True, help="JSON stats file (see configs/)") - p.add_argument("--budget", type=float, required=True, help="TPOT budget B in ms") + p.add_argument("--budget", type=float, required=True, + help="cost budget B in ms, in the units of --cost-metric") + p.add_argument("--cost-metric", choices=("tpot", "e2el"), default="tpot", + help="measurement used as Cost: 'tpot' (default, reproduces the " + "published results) or 'e2el' end-to-end request latency, " + "needed when pool members differ in output length rather " + "than decode speed, such as a thinking/non-thinking pair") p.add_argument("--output", default=None, help="artifacts directory to update") p.set_defaults(func=cmd_fit) + p = sub.add_parser( + "cascade", + help="Stage 1+2 system latency (TPOT and E2EL) from measured stats", + ) + p.add_argument("--stats", required=True, + help="cascade stats JSON with assignment and escalations " + "(see configs/*_cascade_test.json)") + p.set_defaults(func=cmd_cascade) + p = sub.add_parser( "qe-train", help="fine-tune the QE classifier (all flags forwarded, see --help)", @@ -274,6 +391,27 @@ def main(argv: list[str] | None = None) -> None: add_help=False, ) + p = sub.add_parser( + "qe-cascade", + help="run the QE classifier over an efficient model's generations and " + "compose per-cluster cascade accuracy + escalation counts (requires [qe])", + ) + p.add_argument("--classifier", required=True, help="trained QE checkpoint") + p.add_argument("--generations", required=True, + help="efficient model's *_generations.jsonl (qid, cluster, run, correct, full_output, num_tokens)") + p.add_argument("--strong-outcomes", required=True, + help="strong model's *_outcomes.jsonl or *_generations.jsonl (qid, correct)") + p.add_argument("--strong-model", required=True, help="strong model name recorded in the escalations") + p.add_argument("--clusters", default=None, help="comma-separated clusters to cascade (default: all present)") + p.add_argument("--out", default=None, help="cascade stats JSON to update in place with the Stage 2 fields") + p.add_argument("--base-tokenizer", default=None, + help="defaults to the checkpoint itself, which ships its own tokenizer; " + "give a base model id only for a checkpoint saved without one") + p.add_argument("--max-length", type=int, default=4096, help="4096 for long reasoning, 512 for short MCQ") + p.add_argument("--accept-threshold", type=float, default=0.5) + p.add_argument("--batch-size", type=int, default=32) + p.set_defaults(func=cmd_qe_cascade) + p = sub.add_parser("serve", help="run the cascade router") p.add_argument("--config", required=True, help="YAML serving config") p.add_argument("--host", default="0.0.0.0") diff --git a/src/cre_router/evaluate.py b/src/cre_router/evaluate.py index 0ff6b40..7058d0c 100644 --- a/src/cre_router/evaluate.py +++ b/src/cre_router/evaluate.py @@ -19,7 +19,9 @@ from __future__ import annotations +import hashlib import json +import math import re from collections import Counter, defaultdict from dataclasses import dataclass @@ -27,18 +29,13 @@ from statistics import mean from typing import Any, Callable +from cre_router.textutils import split_thinking + # --------------------------------------------------------------------------- # Answer parsing # --------------------------------------------------------------------------- -def split_thinking(text: str) -> str: - """Return the post-reasoning content, dropping a leading ... - block when present.""" - end = text.find("") - return text[end + len("") :].strip() if end != -1 else text.strip() - - def parse_aime_answer(text: str) -> int | None: """Extract an AIME answer (integer 0-999) from a completion. @@ -70,6 +67,49 @@ def parse_teleqna_answer(text: str) -> int | None: return int(numbers[-1]) if numbers else None +# TeleMath gold answers are short numbers (at most 17 characters across the +# 500-question dataset) and a well-formed completion states the answer at the +# very end. Degenerate or truncated generations, however, can run to hundreds of +# kilobytes; searching all of it made the number pattern below backtrack +# quadratically and stall for hours. Restricting the search to a tail window +# keeps every pattern linear in the window, and a real answer (a few characters, +# at the end) is never clipped. +_TELEMATH_TAIL_CHARS = 4000 +# One integer run with an optional fraction, or a leading-dot decimal, plus an +# optional exponent. Unlike ``\d*\.?\d+`` this has no two adjacent +# variable-length digit runs, so it cannot backtrack catastrophically. +_TELEMATH_NUMBER = r"[-+]?(?:\d+(?:\.\d+)?|\.\d+)(?:[eE][-+]?\d+)?" + + +def parse_telemath_answer(text: str) -> float | None: + """Extract a TeleMath numerical answer (a float) from a completion. + + TeleMath answers are numerical quantities, often long decimals or in + scientific notation (e.g. 233.333333333333, 7.2e-05, -62.0854). The final + value is taken from a ``\\boxed{}`` when present, then an explicit + ``Answer:``, then the last number in the text. LaTeX scientific notation + (``7.2 \\times 10^{-5}``) is normalised to ``7.2e-5`` before matching. + """ + content = split_thinking(text)[-_TELEMATH_TAIL_CHARS:] + content = re.sub( + r"([-+]?(?:\d+(?:\.\d+)?|\.\d+))\s*\\times\s*10\^\{?(-?\d+)\}?", + r"\1e\2", + content, + ) + for pattern in ( + rf"\\boxed\{{\s*({_TELEMATH_NUMBER})\s*\}}", + rf"\*{{0,2}}Answer\*{{0,2}}\s*[:：]\s*({_TELEMATH_NUMBER})", + rf"({_TELEMATH_NUMBER})", + ): + matches = re.findall(pattern, content) + if matches: + try: + return float(matches[-1]) + except ValueError: + continue + return None + + def answers_match(predicted: int | None, gold: Any) -> bool: if predicted is None: return False @@ -79,6 +119,21 @@ def answers_match(predicted: int | None, gold: Any) -> bool: return str(predicted).strip() == str(gold).strip() +def numeric_match(predicted: float | None, gold: Any, rel_tol: float = 1e-2) -> bool: + """Correct if the predicted value is within a relative tolerance of the gold. + + Uses ``math.isclose`` with a 1% relative tolerance, which accepts the same + quantity reported at different rounding (233.33 vs 233.333333) and rejects + genuinely different values; the small absolute floor covers near-zero golds. + """ + if predicted is None: + return False + try: + return math.isclose(float(predicted), float(gold), rel_tol=rel_tol, abs_tol=1e-9) + except (TypeError, ValueError): + return False + + # --------------------------------------------------------------------------- # Tasks # --------------------------------------------------------------------------- @@ -100,6 +155,19 @@ class Task: top_k: int min_p: float max_tokens: int + # Gemma 4's thinking switch (`enable_thinking`) is a chat-template kwarg, + # not a text-level prefix like Qwen's /no_think, and vLLM's own bench-serve + # dataset loader never forwards template kwargs (confirmed by reading + # vllm/benchmarks/datasets.py CustomDataset.sample: it calls + # apply_chat_template with a fixed argument list). So for this family the + # dataset prep step renders the template itself and stores the finished + # text as the prompt; this flag tells the benchmark not to template it + # again on top. + pre_rendered: bool = False + # How a parsed answer is compared to the gold label. Defaults to exact + # integer/string equality (`answers_match`); numerical-answer tasks such as + # TeleMath set this to a tolerance-based comparison instead. + match: Callable[[Any, Any], bool] | None = None TASKS: dict[str, Task] = { @@ -121,6 +189,46 @@ class Task: min_p=0.0, max_tokens=1024, ), + # TeleMath: telecom mathematical problems with numerical (float) answers, + # scored by relative tolerance rather than exact match. Two arms, since the + # pool mixes thinking and non-thinking models and each has its own + # recommended sampling: `telemath` for thinking models, `telemath_nothink` + # for instruct models. + "telemath": Task( + name="telemath", + parse=parse_telemath_answer, + match=numeric_match, + temperature=0.6, + top_p=0.95, + top_k=20, + min_p=0.0, + max_tokens=40960, + ), + "telemath_nothink": Task( + name="telemath_nothink", + parse=parse_telemath_answer, + match=numeric_match, + temperature=0.7, + top_p=0.8, + top_k=20, + min_p=0.0, + max_tokens=16384, + ), + # Gemma 4 with thinking enabled. Same reasoning sampling as `telemath`, but + # the thinking switch is a chat-template kwarg baked into the prompt text by + # data/prep_gemma4_thinking.py, so the prompts are pre_rendered and served + # verbatim rather than templated again by the benchmark. + "telemath_gemma4": Task( + name="telemath_gemma4", + parse=parse_telemath_answer, + match=numeric_match, + temperature=0.6, + top_p=0.95, + top_k=20, + min_p=0.0, + max_tokens=40960, + pre_rendered=True, + ), } @@ -130,18 +238,23 @@ class Task: def score_generations( - generated_texts: list[str], gold_answers: list[Any], parse: Callable[[str], int | None] + generated_texts: list[str], + gold_answers: list[Any], + parse: Callable[[str], Any], + match: Callable[[Any, Any], bool] | None = None, ) -> tuple[float, list[bool]]: """Return (error_rate, per-item correctness) for one benchmark run. ``generated_texts`` are assumed aligned with ``gold_answers`` (the vLLM - benchmark preserves dataset order with shuffling disabled). + benchmark preserves dataset order with shuffling disabled). ``match`` + defaults to exact equality; numerical tasks pass a tolerance comparison. """ if len(generated_texts) != len(gold_answers): raise ValueError( f"{len(generated_texts)} generations vs {len(gold_answers)} gold answers" ) - correct = [answers_match(parse(t), g) for t, g in zip(generated_texts, gold_answers)] + matcher = match or answers_match + correct = [matcher(parse(t), g) for t, g in zip(generated_texts, gold_answers)] error = 1.0 - sum(correct) / len(correct) if correct else 1.0 return error, correct @@ -163,6 +276,20 @@ def cluster_sizes(dataset: list[dict]) -> dict[str, int]: return dict(Counter(str(row["cluster"]) for row in dataset)) +# Optional per-run metrics. TPOT alone cannot price a thinking/non-thinking pool, +# because both modes decode at the same speed and differ only in how many tokens +# they emit; E2EL (= TTFT + TPOT x output length) captures that, and unlike +# request throughput it is per-request, so it adds up correctly across cascade +# rungs. All are optional so an injected benchmark may supply only TPOT. +OPTIONAL_METRICS = ( + "ttft_ms", + "e2el_ms", + "request_throughput", + "mean_output_tokens", + "truncated_frac", +) + + @dataclass class RunMeasurement: """One (cluster, run) benchmark outcome, kept for provenance.""" @@ -172,27 +299,100 @@ class RunMeasurement: error: float tpot_ms: float num_prompts: int + ttft_ms: float | None = None + e2el_ms: float | None = None + request_throughput: float | None = None + mean_output_tokens: float | None = None + truncated_frac: float | None = None + + +def optional_metrics( + result: dict, num_prompts: int, max_output_len: int | None = None +) -> dict[str, float]: + """Pull the optional cost metrics out of a benchmark result, tolerating any + that this vLLM version (or an injected fake) does not report. + + ``mean_e2el_ms`` is preferred when present; otherwise it is reconstructed as + ``TTFT + TPOT x (output length - 1)``, which is how it is recoverable from + the summary block of runs that predate this capture. + + ``max_output_len`` is the task's token cap; when given alongside the + per-request ``output_lens``, the fraction of requests that hit the cap is + reported as ``truncated_frac``. + """ + metrics: dict[str, float] = {} + + output_tokens = result.get("total_output_tokens") + mean_output = float(output_tokens) / num_prompts if output_tokens and num_prompts else None + if mean_output is not None: + metrics["mean_output_tokens"] = mean_output + + ttft = result.get("mean_ttft_ms") + if ttft is not None: + metrics["ttft_ms"] = float(ttft) + + e2el = result.get("mean_e2el_ms") + if e2el is None: + ttft, tpot = result.get("mean_ttft_ms"), result.get("mean_tpot_ms") + if ttft is not None and tpot is not None and mean_output: + e2el = float(ttft) + float(tpot) * max(mean_output - 1.0, 0.0) + if e2el is not None: + metrics["e2el_ms"] = float(e2el) + + throughput = result.get("request_throughput") + if throughput is not None: + metrics["request_throughput"] = float(throughput) + + # A capped generation corrupts accuracy and understates cost at the same + # time, so the truncation rate has to travel with the numbers. vLLM's serve + # benchmark does not report a finish reason, but it does return per-request + # output_lens; a request that reached the cap was cut off. EOS-terminated + # requests stop below the cap, so equality with it is the truncation test. + output_lens = result.get("output_lens") + if output_lens and max_output_len: + truncated = sum(1 for length in output_lens if length >= max_output_len) + metrics["truncated_frac"] = truncated / len(output_lens) + + return metrics def aggregate_runs(measurements: list[RunMeasurement]) -> dict[str, dict[str, float]]: - """Average error and TPOT across runs, per cluster.""" + """Average error, TPOT and any available optional metric across runs, per + cluster. An optional metric is reported only when every run supplied it.""" errors: dict[str, list[float]] = defaultdict(list) tpots: dict[str, list[float]] = defaultdict(list) + extra: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list)) for m in measurements: errors[m.cluster].append(m.error) tpots[m.cluster].append(m.tpot_ms) - return { - c: {"error": mean(errors[c]), "tpot_ms": mean(tpots[c])} for c in sorted(errors) - } + for metric in OPTIONAL_METRICS: + value = getattr(m, metric) + if value is not None: + extra[m.cluster][metric].append(float(value)) + + agg = {} + for c in sorted(errors): + entry = {"error": mean(errors[c]), "tpot_ms": mean(tpots[c])} + for metric, values in extra[c].items(): + if len(values) == len(errors[c]): + entry[metric] = mean(values) + agg[c] = entry + return agg def model_entry(measurements: list[RunMeasurement]) -> dict: - """The per-model block for a stats file: per-cluster error and TPOT.""" + """The per-model block for a stats file: per-cluster error and TPOT, plus any + optional metric that every run reported.""" agg = aggregate_runs(measurements) - return { + entry = { "errors": {c: round(agg[c]["error"], 6) for c in agg}, "cluster_tpot_ms": {c: round(agg[c]["tpot_ms"], 6) for c in agg}, } + for metric in OPTIONAL_METRICS: + present = {c: agg[c][metric] for c in agg if metric in agg[c]} + if len(present) == len(agg): + entry[f"cluster_{metric}"] = {c: round(v, 6) for c, v in present.items()} + return entry def merge_model_into_stats( @@ -216,18 +416,18 @@ def save_raw_measurements(measurements: list[RunMeasurement], path: str | Path) path.parent.mkdir(parents=True, exist_ok=True) with path.open("w") as f: for m in sorted(measurements, key=lambda m: (m.cluster, m.run)): - f.write( - json.dumps( - { - "cluster": m.cluster, - "run": m.run, - "error": round(m.error, 6), - "tpot_ms": round(m.tpot_ms, 6), - "num_prompts": m.num_prompts, - } - ) - + "\n" - ) + record = { + "cluster": m.cluster, + "run": m.run, + "error": round(m.error, 6), + "tpot_ms": round(m.tpot_ms, 6), + "num_prompts": m.num_prompts, + } + for metric in OPTIONAL_METRICS: + value = getattr(m, metric) + if value is not None: + record[metric] = round(float(value), 6) + f.write(json.dumps(record) + "\n") # --------------------------------------------------------------------------- @@ -282,10 +482,28 @@ def run_vllm_benchmark( args.top_k = task.top_k args.min_p = task.min_p args.custom_output_len = task.max_tokens + # Prompts for a pre_rendered task already carry the fully-templated text + # (chat-template kwargs like Gemma 4's enable_thinking baked in), so the + # loader must serve them verbatim rather than templating a second time. + args.skip_chat_template = task.pre_rendered + if task.pre_rendered: + # vLLM's completions endpoint decodes with skip_special_tokens=True by + # default, which silently removes Gemma 4's <|channel> marker (a + # registered special token) from the returned text while leaving the + # ordinary word "thought" that follows it untouched -- confirmed by a + # direct A/B request against a live google/gemma-4-E2B-it server. + # Qwen's / are not special tokens and are + # unaffected either way, so this is only needed for pre_rendered tasks. + args.extra_body = {"skip_special_tokens": False} args.disable_shuffle = True args.no_oversample = True args.request_rate = float("inf") args.burstiness = 1.0 + # vLLM reports only the metrics named here; its generative default is + # "ttft,tpot,itl", which drops e2el entirely (see benchmarks/serve.py, + # process_one_metric). Ask for it explicitly so `--cost-metric e2el` reads a + # measured value instead of falling back to reconstructing it from means. + args.percentile_metrics = "ttft,tpot,itl,e2el" args.save_result = False # Keep the per-request fields (generated_texts, errors) in the returned # dict; without this vLLM strips them for a summary-only result. @@ -296,6 +514,19 @@ def run_vllm_benchmark( return benchmark_main(args) +def question_id(item: dict) -> str: + """A stable per-question key for joining outcomes across models/runs. + + Uses the dataset's ``id`` when present (preserved through clustering and + Gemma pre-rendering), otherwise a hash of the raw prompt. It must identify + the same question identically across every model, so downstream analysis can + pair per-question correctness (confidence intervals, McNemar significance).""" + qid = item.get("id") + if qid is not None: + return str(qid) + return hashlib.md5(item["prompt"].encode("utf-8")).hexdigest()[:12] + + def evaluate_model( dataset: list[dict], model: str, @@ -309,18 +540,33 @@ def evaluate_model( workdir: str | Path = "results/splits", download_dir: str | None = None, benchmark: Callable[..., dict] | None = None, + outcomes_out: str | Path | None = None, + generations_out: str | Path | None = None, ) -> list[RunMeasurement]: """Benchmark ``model`` on each cluster for ``runs`` repetitions. Each dataset row needs ``prompt``, ``answer``, and ``cluster``. Returns the per-(cluster, run) measurements; aggregate them with ``model_entry``. ``benchmark`` defaults to ``run_vllm_benchmark`` and is injected in tests. + + When ``outcomes_out`` is given, writes one JSONL row per (question, run) with + ``qid``, ``cluster``, ``run``, ``correct`` and ``output_len``. This is the + per-question record needed for confidence intervals and paired significance + tests; it is irrecoverable once discarded, so capture it during the run. + + When ``generations_out`` is given, writes one JSONL row per (question, run) + additionally carrying the model's ``full_output`` text and ``num_tokens``. + This is the training data for the Stage 2 accept/escalate QE classifier + (``question [SEP] full_output [SEP] num_tokens`` -> correct), and like the + generations themselves it is irrecoverable once the run ends. """ run_benchmark = benchmark or run_vllm_benchmark workdir = Path(workdir) workdir.mkdir(parents=True, exist_ok=True) measurements: list[RunMeasurement] = [] + outcomes: list[dict] | None = [] if outcomes_out is not None else None + generations: list[dict] | None = [] if generations_out is not None else None for cluster, items in split_by_cluster(dataset).items(): split_path = workdir / f"{task.name}_cluster_{cluster}.jsonl" with split_path.open("w") as f: @@ -339,7 +585,43 @@ def evaluate_model( seed=base_seed + run, download_dir=download_dir, ) - error, _ = score_generations(result["generated_texts"], gold, task.parse) + error, correct = score_generations( + result["generated_texts"], gold, task.parse, task.match + ) + output_lens = result.get("output_lens") or [None] * len(items) + if outcomes is not None: + for item, ok, olen in zip(items, correct, output_lens): + outcomes.append( + { + "qid": question_id(item), + "cluster": cluster, + "run": run, + "correct": bool(ok), + "output_len": olen, + } + ) + if generations is not None: + texts = result["generated_texts"] + for item, ok, olen, text in zip(items, correct, output_lens, texts): + generations.append( + { + "qid": question_id(item), + "cluster": cluster, + "run": run, + # ``question`` is the raw query the QE classifier and + # humans read; ``prompt`` is the exact served input + # (chat-templated for pre_rendered tasks). They differ + # only for pre_rendered models, where prep stores the + # original question under ``question``. + "question": item.get("question", item.get("prompt", "")), + "prompt": item.get("prompt", ""), + "ground_truth_answer": item["answer"], + "answer": task.parse(text), + "full_output": text, + "num_tokens": olen, + "correct": bool(ok), + } + ) tpot_ms = result.get("mean_tpot_ms") if tpot_ms is None: raise ValueError("benchmark result is missing 'mean_tpot_ms'") @@ -350,6 +632,19 @@ def evaluate_model( error=error, tpot_ms=float(tpot_ms), num_prompts=len(items), + **optional_metrics(result, len(items), task.max_tokens), ) ) + if outcomes is not None: + out_path = Path(outcomes_out) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w") as f: + for o in outcomes: + f.write(json.dumps(o) + "\n") + if generations is not None: + gen_path = Path(generations_out) + gen_path.parent.mkdir(parents=True, exist_ok=True) + with gen_path.open("w") as f: + for g in generations: + f.write(json.dumps(g, ensure_ascii=False) + "\n") return measurements diff --git a/src/cre_router/qe/cascade.py b/src/cre_router/qe/cascade.py new file mode 100644 index 0000000..0adfe8e --- /dev/null +++ b/src/cre_router/qe/cascade.py @@ -0,0 +1,121 @@ +"""Stage 1+2 cascade evaluation. + +Runs the QE classifier over an efficient model's per-cluster generations, +escalates the rejected outputs to the strong model, and composes the per-cluster +cascade accuracy and escalation counts that ``cre cascade`` consumes +(``routing.cascade_system_accuracy`` / ``cascade_system_metrics``). + +The composition is split so the arithmetic is testable without a GPU: +``compose_cascade`` is a pure function over already-made accept/route decisions, +and ``run_qe`` is the thin wrapper that produces those decisions from a trained +classifier. Requires the ``qe`` extra only for ``run_qe``. +""" + +from __future__ import annotations + +import json +from collections import defaultdict +from pathlib import Path + + +def strong_correct_by_qid(outcomes: list[dict]) -> dict[str, float]: + """Mean correctness per question id from the strong model's records. + + Accepts either the strong model's ``*_outcomes.jsonl`` or its + ``*_generations.jsonl`` (both carry ``qid`` and ``correct``). When the strong + model was run several times, the per-qid mean is its expected correctness on + that query, which is what an escalation to it earns. + """ + total: dict[str, float] = defaultdict(float) + count: dict[str, int] = defaultdict(int) + for row in outcomes: + qid = str(row["qid"]) + total[qid] += 1.0 if row["correct"] else 0.0 + count[qid] += 1 + return {qid: total[qid] / count[qid] for qid in total} + + +def compose_cascade( + generations: list[dict], + escalate: list[bool], + strong_correct: dict[str, float], +) -> dict[str, dict]: + """Per-cluster cascade accuracy and average escalations-per-run. + + ``generations`` are the efficient model's per-question rows (``qid``, + ``cluster``, ``run``, ``correct``); ``escalate[i]`` is the QE decision to + escalate row ``i`` (True = route to the strong model). For each row the + system answer is the strong model's (paired by ``qid``) when escalated, else + the efficient model's own correctness. + + Returns ``{cluster: {"cascade_accuracy", "escalations", "n"}}`` where + ``escalations`` is the mean number of escalated queries per run, the ``count`` + that ``cascade_system_metrics`` charges (``direct = size - count``). + """ + if len(escalate) != len(generations): + raise ValueError( + f"escalate ({len(escalate)}) must align with generations ({len(generations)})" + ) + correct_sum: dict[str, float] = defaultdict(float) + esc_count: dict[str, int] = defaultdict(int) + n_rows: dict[str, int] = defaultdict(int) + runs: dict[str, set] = defaultdict(set) + for row, esc in zip(generations, escalate): + cluster = str(row["cluster"]) + n_rows[cluster] += 1 + runs[cluster].add(row["run"]) + if esc: + qid = str(row["qid"]) + if qid not in strong_correct: + raise KeyError( + f"no strong-model outcome for escalated qid {qid!r}; the strong " + f"model must be evaluated on every query it can be escalated -- run " + f"it on cluster {cluster!r} (its outcomes/generations feed strong_correct)" + ) + esc_count[cluster] += 1 + correct_sum[cluster] += strong_correct[qid] + else: + correct_sum[cluster] += 1.0 if row["correct"] else 0.0 + report: dict[str, dict] = {} + for cluster in n_rows: + n_runs = len(runs[cluster]) or 1 + report[cluster] = { + "cascade_accuracy": correct_sum[cluster] / n_rows[cluster], + "escalations": esc_count[cluster] / n_runs, + "n": n_rows[cluster], + } + return report + + +def run_qe(classifier, generations: list[dict], batch_size: int = 32) -> list[bool]: + """Return the escalate decision (True = route) per generation. + + ``classifier`` needs a ``predict_batch(list[(question, output, num_tokens)])`` + returning objects with an ``.accept`` flag (``QEClassifier`` or a test stub). + """ + escalate: list[bool] = [] + for start in range(0, len(generations), batch_size): + chunk = generations[start : start + batch_size] + items = [ + (row.get("question", row.get("prompt", "")), row["full_output"], row["num_tokens"]) + for row in chunk + ] + escalate.extend(not d.accept for d in classifier.predict_batch(items)) + return escalate + + +def write_cascade_stats(out_path: str | Path, strong_model: str, report: dict[str, dict]) -> None: + """Merge ``escalations`` and ``cascade_accuracy`` into a cascade stats JSON. + + ``out_path`` must already hold the routing side (``assignment``, + ``cluster_sizes``, ``models``) as produced for ``cre cascade``; this adds the + Stage 2 fields per cluster in ``report`` and leaves the rest untouched. + """ + path = Path(out_path) + stats = json.loads(path.read_text()) + stats.setdefault("escalations", {}) + stats.setdefault("cascade_accuracy", {}) + for cluster, r in report.items(): + stats["escalations"][cluster] = [strong_model, r["escalations"]] + stats["cascade_accuracy"][cluster] = r["cascade_accuracy"] + path.write_text(json.dumps(stats, indent=2) + "\n") diff --git a/src/cre_router/qe/classifier.py b/src/cre_router/qe/classifier.py index e63e2eb..6a1560a 100644 --- a/src/cre_router/qe/classifier.py +++ b/src/cre_router/qe/classifier.py @@ -11,6 +11,8 @@ from dataclasses import dataclass +from cre_router.textutils import split_thinking + ROUTE, ACCEPT = 0, 1 CLASS_NAMES = ("Route", "Accept") @@ -24,9 +26,20 @@ def format_qe_input( ) -> str: """Training/inference input format: ``query [SEP] output [SEP] num_tokens``. - Including the output length gives the classifier direct access to - chain-of-thought length, a proxy for model confidence. + The reasoning block is dropped first (``split_thinking``, all model families) + so the classifier sees the delivered answer, not the chain of thought -- the + thinking is not part of the output. ``num_tokens`` still reflects the *full* + generation length (the real serving cost), giving the classifier direct + access to output length as a proxy for model confidence. The output is then + truncated to its last ``max_output_words`` words. """ + if num_tokens is None: + raise ValueError( + "num_tokens is None: the generation length is required for the QE input " + "(a generations file written without output_lens). Build the dataset via " + "prep_qe, which drops such rows, or re-run the benchmark so output_lens is present." + ) + output = split_thinking(output) words = output.split() if len(words) > max_output_words: output = " ".join(words[-max_output_words:]) diff --git a/src/cre_router/qe/evaluate.py b/src/cre_router/qe/evaluate.py index 039ac65..c67713c 100644 --- a/src/cre_router/qe/evaluate.py +++ b/src/cre_router/qe/evaluate.py @@ -88,7 +88,9 @@ def main(argv: list[str] | None = None) -> None: parser.add_argument("--classifier", required=True, help="trained QE checkpoint") parser.add_argument("--dataset", required=True, help="HF dataset of model outputs") parser.add_argument("--split", required=True, help="dataset split to evaluate") - parser.add_argument("--base-tokenizer", default="answerdotai/ModernBERT-base") + parser.add_argument("--base-tokenizer", default=None, + help="defaults to the checkpoint itself, which ships its own tokenizer; " + "give a base model id only for a checkpoint saved without one") parser.add_argument("--max-length", type=int, default=4096, help="4096 AIME, 512 TeleQnA") parser.add_argument("--accept-threshold", type=float, default=0.5) parser.add_argument("--batch-size", type=int, default=32) diff --git a/src/cre_router/qe/train.py b/src/cre_router/qe/train.py index 1fe68af..f944bb6 100644 --- a/src/cre_router/qe/train.py +++ b/src/cre_router/qe/train.py @@ -74,7 +74,24 @@ def main(argv: list[str] | None = None) -> None: set_seed(args.seed) tokenizer = AutoTokenizer.from_pretrained(args.base_model) - dataset = load_dataset(args.dataset, cache_dir=args.cache_dir) + # A local directory of {train,test}.jsonl (from `python -m cre_router.data.prep_qe`) + # loads via the json builder; anything else is a Hub dataset id. + from pathlib import Path as _Path + + if _Path(args.dataset).is_dir(): + # Only the split files themselves; a dir may also hold sidecar jsonl (e.g. + # cascade_test_gens.jsonl) with a different schema that would break the + # single-schema json builder if globbed in. + data_files = { + p.stem: str(p) + for p in sorted(_Path(args.dataset).glob("*.jsonl")) + if p.stem.startswith(("train", "test")) + } + if not data_files: + raise SystemExit(f"{args.dataset} is a directory but has no train*/test* .jsonl splits") + dataset = load_dataset("json", data_files=data_files, cache_dir=args.cache_dir) + else: + dataset = load_dataset(args.dataset, cache_dir=args.cache_dir) train_split = args.train_split or next(s for s in dataset if s.startswith("train")) eval_split = args.eval_split or next(s for s in dataset if s.startswith("test")) diff --git a/src/cre_router/routing.py b/src/cre_router/routing.py index 0caa61e..67fcc1c 100644 --- a/src/cre_router/routing.py +++ b/src/cre_router/routing.py @@ -18,20 +18,69 @@ class ModelStats: """Statistics for one candidate model, measured on the training corpus. - ``tpot_ms`` is the pool-level average Time Per Output Token used in the - normalised cost term. ``cluster_tpot_ms`` optionally refines system-TPOT - estimates with per-cluster measurements; when absent, ``tpot_ms`` is used - for every cluster. + ``tpot_ms`` is the pool-level average Time Per Output Token. ``cluster_tpot_ms`` + optionally refines system estimates with per-cluster measurements; when + absent, the pool-level value is used for every cluster. + + ``cost_metric`` selects which measurement the routing arithmetic treats as + Cost. ``"tpot"`` is the default and reproduces the published results exactly. + ``"e2el"`` uses end-to-end request latency, which is required whenever pool + members differ in output *length* rather than decode speed (a thinking and a + non-thinking mode of one model have near-identical TPOT). E2EL generalises + TPOT: since ``E2EL = TTFT + TPOT x L``, uniform L and TTFT make it affine in + TPOT, and min-max normalisation is affine-invariant, so the two agree. """ name: str tpot_ms: float errors: dict[str, float] cluster_tpot_ms: dict[str, float] = field(default_factory=dict) + e2el_ms: float | None = None + cluster_e2el_ms: dict[str, float] = field(default_factory=dict) + cluster_output_tokens: dict[str, float] = field(default_factory=dict) + cost_metric: str = "tpot" + + def __post_init__(self) -> None: + if self.cost_metric not in ("tpot", "e2el"): + raise ValueError(f"unknown cost_metric {self.cost_metric!r}") + if self.cost_metric == "e2el" and self.e2el_ms is None: + raise ValueError( + f"{self.name}: cost_metric 'e2el' needs e2el_ms, which this stats " + "entry does not have; re-measure, or fit with --cost-metric tpot" + ) + + @property + def cost_ms(self) -> float: + """The pool-level cost under the selected metric.""" + return self.tpot_ms if self.cost_metric == "tpot" else float(self.e2el_ms) + + def cost_for(self, cluster: str) -> float: + """The per-cluster cost under the selected metric.""" + if self.cost_metric == "tpot": + return self.cluster_tpot_ms.get(cluster, self.tpot_ms) + return self.cluster_e2el_ms.get(cluster, float(self.e2el_ms)) def tpot_for(self, cluster: str) -> float: return self.cluster_tpot_ms.get(cluster, self.tpot_ms) + def e2el_for(self, cluster: str) -> float: + """Per-cluster E2EL, independent of the selected cost metric.""" + if cluster in self.cluster_e2el_ms: + return self.cluster_e2el_ms[cluster] + if self.e2el_ms is None: + raise ValueError(f"{self.name}: E2EL requested but none measured") + return float(self.e2el_ms) + + def output_tokens_for(self, cluster: str) -> float: + """Mean output length on a cluster; needed to charge an escalated + query's discarded efficient pass per delivered token.""" + length = self.cluster_output_tokens.get(cluster) + if length is None: + raise ValueError( + f"{self.name}: cascade TPOT needs cluster_output_tokens[{cluster!r}]" + ) + return length + @dataclass class Region: @@ -89,13 +138,13 @@ def clusters_of(models: list[ModelStats]) -> list[str]: def normalized_costs(models: list[ModelStats]) -> dict[str, float]: - """Eq. 2: min-max normalise pool TPOT so the fastest model costs 0 and - the slowest costs 1.""" - lo = min(m.tpot_ms for m in models) - hi = max(m.tpot_ms for m in models) + """Eq. 2: min-max normalise the pool's cost metric so the cheapest model + costs 0 and the most expensive costs 1.""" + lo = min(m.cost_ms for m in models) + hi = max(m.cost_ms for m in models) if hi == lo: return {m.name: 0.0 for m in models} - return {m.name: (m.tpot_ms - lo) / (hi - lo) for m in models} + return {m.name: (m.cost_ms - lo) / (hi - lo) for m in models} def assign(models: list[ModelStats], lam: float) -> dict[str, str]: @@ -104,19 +153,19 @@ def assign(models: list[ModelStats], lam: float) -> dict[str, str]: costs = normalized_costs(models) table = {} for c in clusters_of(models): - best = min(models, key=lambda m: (m.errors[c] + lam * costs[m.name], m.tpot_ms)) + best = min(models, key=lambda m: (m.errors[c] + lam * costs[m.name], m.cost_ms)) table[c] = best.name return table def dominates(a: ModelStats, b: ModelStats) -> bool: - """True if ``a`` Pareto-dominates ``b``: no worse on TPOT and on every + """True if ``a`` Pareto-dominates ``b``: no worse on cost and on every cluster's error, and strictly better on at least one of those.""" - if a.tpot_ms > b.tpot_ms: + if a.cost_ms > b.cost_ms: return False if any(a.errors[c] > b.errors[c] for c in b.errors): return False - return a.tpot_ms < b.tpot_ms or any(a.errors[c] < b.errors[c] for c in b.errors) + return a.cost_ms < b.cost_ms or any(a.errors[c] < b.errors[c] for c in b.errors) def pareto_prune(models: list[ModelStats]) -> tuple[list[ModelStats], list[ModelStats]]: @@ -179,11 +228,150 @@ def system_metrics( cluster_sizes[c] * (1.0 - by_name[name].errors[c]) for c, name in assignment.items() ) / total tpot = sum( - cluster_sizes[c] * by_name[name].tpot_for(c) for c, name in assignment.items() + cluster_sizes[c] * by_name[name].cost_for(c) for c, name in assignment.items() ) / total return acc, tpot +def cascade_system_metrics_ntier( + models: list[ModelStats], + cascades: dict[str, list[tuple[str, int | float]]], + cluster_sizes: dict[str, int | float], +) -> tuple[float, float]: + """System TPOT and E2EL for an N-tier cascade (>= 1 tier per cluster). + + ``cascades`` maps each cluster to its escalation chain, ordered from the + Stage 1 (base) model up to the strongest tier: + ``[(model_0, reach_0), (model_1, reach_1), ...]``. ``reach_i`` is the number + of the cluster's queries that *execute* tier ``i`` -- ``reach_0`` equals the + cluster size (every query runs the base model), and the sequence is + non-increasing as the QE gate at each tier accepts some outputs and escalates + the rest. A single-element chain is a direct (un-gated) assignment. Two-tier + ``cascade_system_metrics`` is the special case ``[(eff, size), (strong, count)]``. + + Both metrics are query-weighted, matching Stage 1 and vLLM's per-request Mean + TPOT. A query delivered at tier ``k`` has executed tiers ``0..k``: + + - **E2EL** is the plain sum of every executed pass, so each query reaching + tier ``i`` pays ``E2EL_i`` in full: ``E2EL = sum_i reach_i * E2EL_i``. + - **TPOT** charges all executed passes to the delivered tokens ``L_k``: + ``(sum_{i<=k} TPOT_i * L_i) / L_k``, generalising the two-tier rule that + amortises the discarded efficient generation over the delivered answer. + + Returns ``(tpot_ms, e2el_ms)``. + """ + by_name = {m.name: m for m in models} + total = sum(cluster_sizes.values()) + tpot_sum = 0.0 + e2el_sum = 0.0 + for cluster, chain in cascades.items(): + if not chain: + raise ValueError(f"cluster {cluster!r} has an empty cascade chain") + size = cluster_sizes[cluster] + names = [t[0] for t in chain] + reach = [t[1] for t in chain] + if not math.isclose(reach[0], size): + raise ValueError( + f"cluster {cluster!r}: base tier reach {reach[0]} must equal cluster size {size}" + ) + if any(reach[i + 1] > reach[i] for i in range(len(reach) - 1)): + raise ValueError(f"cluster {cluster!r}: tier reach must be non-increasing, got {reach}") + tpots = [by_name[n].tpot_for(cluster) for n in names] + e2els = [by_name[n].e2el_for(cluster) for n in names] + lengths = [by_name[n].output_tokens_for(cluster) for n in names] + cumulative_decode = 0.0 + for k in range(len(names)): + e2el_sum += reach[k] * e2els[k] + cumulative_decode += tpots[k] * lengths[k] + reach_next = reach[k + 1] if k + 1 < len(names) else 0 + delivered_k = reach[k] - reach_next + tpot_sum += delivered_k * (cumulative_decode / lengths[k]) + return tpot_sum / total, e2el_sum / total + + +def cascade_system_metrics( + models: list[ModelStats], + assignment: dict[str, str], + cluster_sizes: dict[str, int | float], + escalations: dict[str, tuple[str, float]], +) -> tuple[float, float]: + """System TPOT and E2EL for a two-tier Stage 1+2 cascade. + + ``assignment`` is the Stage 1 cluster-to-model map. ``escalations`` maps a + cluster to ``(strong_model, count)``: that many of the cluster's queries run + the Stage 1 (efficient) model, are judged low-quality by the QE classifier, + and are then re-run on ``strong_model``. Clusters absent from ``escalations`` + route entirely to their Stage 1 model, so ``system_metrics`` is the special + case with no escalations. + + A thin wrapper over :func:`cascade_system_metrics_ntier` that builds a + one- or two-tier chain per cluster. Reproduces the paper's Stage 1+2 latency, + 9.7 ms on AIME and 23.8 ms on TeleQnA. + """ + cascades: dict[str, list[tuple[str, int | float]]] = {} + for cluster, name in assignment.items(): + size = cluster_sizes[cluster] + if cluster in escalations: + strong_name, count = escalations[cluster] + cascades[cluster] = [(name, size), (strong_name, count)] + else: + cascades[cluster] = [(name, size)] + return cascade_system_metrics_ntier(models, cascades, cluster_sizes) + + +def cluster_cascade_accuracy( + weak_correct: list[bool], + strong_correct: list[bool], + escalate: list[bool], +) -> float: + """Per-cluster accuracy of the Stage 2 cascade, composed per query. + + For each query the QE classifier either accepts the efficient model's output + (keep ``weak_correct``) or escalates it to the strong model (take + ``strong_correct``). The three lists are aligned per query on the same + cluster. This is the exact composition the paper reports: a wrongly escalated + correct answer (false positive) still gets whatever the strong model returns, + and a wrongly accepted wrong answer (false negative) stays wrong. + """ + n = len(weak_correct) + if not (len(strong_correct) == len(escalate) == n): + raise ValueError("weak_correct, strong_correct, escalate must align per query") + if n == 0: + raise ValueError("cannot compute cascade accuracy over an empty cluster") + correct = sum( + (s if esc else w) for w, s, esc in zip(weak_correct, strong_correct, escalate) + ) + return correct / n + + +def cascade_system_accuracy( + models: list[ModelStats], + assignment: dict[str, str], + cluster_sizes: dict[str, int | float], + cascade_accuracy: dict[str, float], +) -> float: + """System accuracy for the full Stage 1+2 cascade, query-weighted. + + A cluster listed in ``cascade_accuracy`` contributes that measured cascade + accuracy (efficient outputs gated by the QE classifier, rejects escalated to + the strong model -- see ``cluster_cascade_accuracy``). A cluster absent from + it routes entirely to its Stage 1 model and contributes ``1 - error`` for the + assigned model, so ``system_metrics``' accuracy is the no-escalation special + case. Reproduces the paper's Stage 1+2 accuracy, 88.4% on AIME and 74.3% on + TeleQnA. + """ + by_name = {m.name: m for m in models} + total = sum(cluster_sizes.values()) + acc = 0.0 + for cluster, name in assignment.items(): + size = cluster_sizes[cluster] + if cluster in cascade_accuracy: + acc += size * cascade_accuracy[cluster] + else: + acc += size * (1.0 - by_name[name].errors[cluster]) + return acc / total + + def eta( models: list[ModelStats], assignment: dict[str, str], @@ -229,9 +417,12 @@ def select_lambda( system_metrics(models, r.assignment, cluster_sizes)[1] for r in routing_regions(models) ) + metric = models[0].cost_metric.upper() if models else "cost" raise ValueError( f"No routing strategy satisfies budget {budget_ms} ms " - f"(fastest achievable system TPOT is {fastest:.1f} ms)." + f"(fastest achievable system {metric} is {fastest:.1f} ms). " + f"Note the budget is in {metric} units: a per-token TPOT budget is " + f"orders of magnitude smaller than a per-request E2EL one." ) acc, tpot, region = best return Selection( @@ -243,7 +434,9 @@ def select_lambda( ) -def models_from_stats(stats: dict) -> tuple[list[ModelStats], dict[str, float]]: +def models_from_stats( + stats: dict, cost_metric: str = "tpot" +) -> tuple[list[ModelStats], dict[str, float]]: """Build ``ModelStats`` from a stats dict (see configs/*_stats.json). Expected shape:: @@ -254,10 +447,14 @@ def models_from_stats(stats: dict) -> tuple[list[ModelStats], dict[str, float]]: "name": { "tpot_ms": 9.15, # optional if cluster_tpot_ms given "errors": {"0": 0.130, ...}, - "cluster_tpot_ms": {"0": 9.282, ...} # optional + "cluster_tpot_ms": {"0": 9.282, ...}, # optional + "cluster_e2el_ms": {"0": 8123.4, ...} # optional, needed for e2el } } } + + ``cost_metric`` defaults to ``"tpot"``, under which this reads exactly the + fields it always has and reproduces published results unchanged. """ models = [] for name, spec in stats["models"].items(): @@ -267,12 +464,26 @@ def models_from_stats(stats: dict) -> tuple[list[ModelStats], dict[str, float]]: if not cluster_tpot: raise ValueError(f"Model {name!r} needs tpot_ms or cluster_tpot_ms") tpot = sum(cluster_tpot.values()) / len(cluster_tpot) + + cluster_e2el = {str(k): float(v) for k, v in spec.get("cluster_e2el_ms", {}).items()} + e2el = spec.get("e2el_ms") + if e2el is None and cluster_e2el: + e2el = sum(cluster_e2el.values()) / len(cluster_e2el) + + cluster_tokens = { + str(k): float(v) for k, v in spec.get("cluster_output_tokens", {}).items() + } + models.append( ModelStats( name=name, tpot_ms=float(tpot), errors={str(k): float(v) for k, v in spec["errors"].items()}, cluster_tpot_ms=cluster_tpot, + e2el_ms=float(e2el) if e2el is not None else None, + cluster_e2el_ms=cluster_e2el, + cluster_output_tokens=cluster_tokens, + cost_metric=cost_metric, ) ) cluster_sizes = {str(k): float(v) for k, v in stats["cluster_sizes"].items()} diff --git a/src/cre_router/textutils.py b/src/cre_router/textutils.py new file mode 100644 index 0000000..85ec542 --- /dev/null +++ b/src/cre_router/textutils.py @@ -0,0 +1,23 @@ +"""Small, dependency-free text helpers shared across the pipeline.""" + +from __future__ import annotations + + +def split_thinking(text: str) -> str: + """Return the post-reasoning content, dropping a leading reasoning block. + + Two marker styles are in use across the pool: Qwen's ``...`` + and Gemma 4's ``<|channel>thought\\n...\\n`` (confirmed against + google/gemma-4-E2B-it's own chat_template.jinja and model card). + Gemma's E2B/E4B variants emit no channel markers at all when thinking is + disabled, unlike its larger siblings and unlike Qwen, which always emits an + empty ; checking for both end markers handles every case + without needing to know which family produced the text. Idempotent: text with + no reasoning marker is returned unchanged (stripped), so it is safe to apply + to outputs that were already de-thought (e.g. the released router datasets). + """ + for end_marker in ("", ""): + end = text.find(end_marker) + if end != -1: + return text[end + len(end_marker) :].strip() + return text.strip() diff --git a/tests/test_cost_metric.py b/tests/test_cost_metric.py new file mode 100644 index 0000000..21b7b76 --- /dev/null +++ b/tests/test_cost_metric.py @@ -0,0 +1,182 @@ +"""The selectable cost metric and the optional E2EL/throughput capture. + +The load-bearing property is that `cost_metric="tpot"` (the default) leaves the +published behaviour untouched, so every test here that exercises e2el has a tpot +counterpart asserting the old path still answers the same way. +""" + +import pytest + +from cre_router.evaluate import ( + RunMeasurement, + aggregate_runs, + model_entry, + optional_metrics, +) +from cre_router.routing import ( + ModelStats, + assign, + models_from_stats, + normalized_costs, + pareto_prune, +) + + +class TestOptionalMetrics: + def test_prefers_reported_e2el(self): + m = optional_metrics( + {"mean_e2el_ms": 8000.0, "mean_ttft_ms": 100.0, "mean_tpot_ms": 10.0, + "total_output_tokens": 400}, num_prompts=4 + ) + assert m["e2el_ms"] == 8000.0 + assert m["mean_output_tokens"] == 100.0 + + def test_reconstructs_e2el_when_absent(self): + # TTFT + TPOT x (L - 1) = 100 + 10 x 99 + m = optional_metrics( + {"mean_ttft_ms": 100.0, "mean_tpot_ms": 10.0, "total_output_tokens": 400}, + num_prompts=4, + ) + assert m["e2el_ms"] == pytest.approx(1090.0) + + def test_truncation_rate(self): + # Two of four requests reached the 30k cap; the others stopped at EOS. + m = optional_metrics( + {"output_lens": [512, 30000, 1024, 30000]}, num_prompts=4, max_output_len=30000 + ) + assert m["truncated_frac"] == 0.5 + + def test_no_truncation_rate_without_the_cap(self): + # output_lens alone cannot say what counts as truncated. + m = optional_metrics({"output_lens": [512, 30000]}, num_prompts=2) + assert "truncated_frac" not in m + + def test_tolerates_a_bare_result(self): + # An injected fake benchmark reporting only TPOT must not break capture. + assert optional_metrics({"mean_tpot_ms": 10.0}, num_prompts=4) == {} + + +class TestAggregation: + def _runs(self, **extra): + return [ + RunMeasurement("0", 0, error=0.2, tpot_ms=10.0, num_prompts=5, **extra), + RunMeasurement("0", 1, error=0.4, tpot_ms=12.0, num_prompts=5, **extra), + ] + + def test_optional_metric_averaged(self): + runs = [ + RunMeasurement("0", 0, error=0.2, tpot_ms=10.0, num_prompts=5, e2el_ms=1000.0), + RunMeasurement("0", 1, error=0.4, tpot_ms=12.0, num_prompts=5, e2el_ms=2000.0), + ] + assert aggregate_runs(runs)["0"]["e2el_ms"] == pytest.approx(1500.0) + + def test_partial_metric_is_dropped(self): + """A metric only some runs reported would be an average over a different + denominator than the rest, so it is omitted rather than half-reported.""" + runs = [ + RunMeasurement("0", 0, error=0.2, tpot_ms=10.0, num_prompts=5, e2el_ms=1000.0), + RunMeasurement("0", 1, error=0.4, tpot_ms=12.0, num_prompts=5), + ] + assert "e2el_ms" not in aggregate_runs(runs)["0"] + assert "cluster_e2el_ms" not in model_entry(runs) + + def test_entry_shape_unchanged_without_optional_metrics(self): + assert set(model_entry(self._runs())) == {"errors", "cluster_tpot_ms"} + + + +class TestCostMetricSelection: + def _pair(self, metric): + """A thinking/non-thinking pair: near-identical TPOT, very different E2EL. + + This is the case TPOT cannot price, since both modes run the same weights + at the same decode speed and differ only in how many tokens they emit. + """ + return [ + ModelStats("nothink", tpot_ms=10.0, errors={"0": 0.40}, e2el_ms=2_000.0, + cost_metric=metric), + ModelStats("think", tpot_ms=10.2, errors={"0": 0.20}, e2el_ms=40_000.0, + cost_metric=metric), + ] + + def test_tpot_is_the_default(self): + assert ModelStats("m", tpot_ms=10.0, errors={"0": 0.1}).cost_ms == 10.0 + + def test_cost_ms_follows_the_metric(self): + nothink, think = self._pair("e2el") + assert (nothink.cost_ms, think.cost_ms) == (2_000.0, 40_000.0) + + def test_per_cluster_cost_falls_back_to_pool_level(self): + m = ModelStats("m", tpot_ms=10.0, errors={"0": 0.1}, e2el_ms=500.0, + cluster_e2el_ms={"0": 900.0}, cost_metric="e2el") + assert m.cost_for("0") == 900.0 + assert m.cost_for("absent") == 500.0 + + def test_e2el_without_measurement_is_refused(self): + with pytest.raises(ValueError, match="needs e2el_ms"): + ModelStats("m", tpot_ms=10.0, errors={"0": 0.1}, cost_metric="e2el") + + def test_unknown_metric_is_refused(self): + with pytest.raises(ValueError, match="unknown cost_metric"): + ModelStats("m", tpot_ms=10.0, errors={"0": 0.1}, cost_metric="dollars") + + def test_normalisation_ranks_by_the_selected_metric(self): + # Under TPOT the thinking mode is (misleadingly) the marginally dearer + # one by 0.2 ms; under E2EL it is dearer by 20x. Min-max maps both to + # {0, 1} for a two-model pool, so the ordering is what differs. + assert normalized_costs(self._pair("tpot"))["nothink"] == 0.0 + assert normalized_costs(self._pair("e2el"))["nothink"] == 0.0 + + def test_pareto_pruning_uses_the_selected_metric(self): + """With three rungs the spacing matters, not just the ordering: a cheap + thinking mode must not dominate a genuinely faster non-thinking one.""" + models = [ + ModelStats("nothink", tpot_ms=10.0, errors={"0": 0.40}, e2el_ms=2_000.0, + cost_metric="e2el"), + ModelStats("think", tpot_ms=10.2, errors={"0": 0.20}, e2el_ms=40_000.0, + cost_metric="e2el"), + ModelStats("big", tpot_ms=30.0, errors={"0": 0.35}, e2el_ms=60_000.0, + cost_metric="e2el"), + ] + efficient, dominated = pareto_prune(models) + # "big" is both dearer and worse than "think", so it cannot ever be chosen. + assert [m.name for m in dominated] == ["big"] + assert {m.name for m in efficient} == {"nothink", "think"} + + def test_lambda_sweep_still_trades_error_against_cost(self): + models = self._pair("e2el") + assert assign(models, lam=0.0)["0"] == "think" # accuracy at any cost + assert assign(models, lam=10.0)["0"] == "nothink" # cost dominates + + +class TestModelsFromStats: + def _stats(self): + return { + "cluster_sizes": {"0": 10}, + "models": { + "fast": {"errors": {"0": 0.4}, "cluster_tpot_ms": {"0": 10.0}, + "cluster_e2el_ms": {"0": 2000.0}}, + "slow": {"errors": {"0": 0.2}, "cluster_tpot_ms": {"0": 10.2}, + "cluster_e2el_ms": {"0": 40000.0}}, + }, + } + + def test_defaults_to_tpot(self): + models, _ = models_from_stats(self._stats()) + assert [m.cost_ms for m in models] == [10.0, 10.2] + + def test_e2el_selected(self): + models, _ = models_from_stats(self._stats(), "e2el") + assert [m.cost_ms for m in models] == [2000.0, 40000.0] + + def test_pool_level_derived_from_per_cluster(self): + models, _ = models_from_stats(self._stats(), "e2el") + assert models[0].e2el_ms == 2000.0 + + def test_legacy_stats_still_load(self): + """A stats file written before this change has no e2el and must keep + working under the default metric.""" + legacy = {"cluster_sizes": {"0": 10}, + "models": {"m": {"tpot_ms": 9.15, "errors": {"0": 0.13}}}} + models, sizes = models_from_stats(legacy) + assert models[0].cost_ms == 9.15 and sizes == {"0": 10.0} diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py index fc8d66b..e07bc54 100644 --- a/tests/test_evaluate.py +++ b/tests/test_evaluate.py @@ -41,6 +41,17 @@ def test_none_when_no_number(self): def test_strips_thinking_block(self): assert split_thinking("99 is wrong Answer: 3") == "Answer: 3" + def test_strips_gemma4_channel_block(self): + # Gemma 4's marker pair, distinct from Qwen's .... + text = "<|channel>thought\nmaybe 99\n Answer: 3" + assert split_thinking(text) == "Answer: 3" + + def test_no_marker_passes_through(self): + # Gemma 4's E2B/E4B variants emit no channel markers at all when + # thinking is disabled, unlike Qwen (always ) and + # unlike Gemma 4's own larger siblings (empty channel block). + assert split_thinking("Answer: 3") == "Answer: 3" + def test_aime_ignores_reasoning_numbers(self): assert parse_aime_answer("try 500 then 600\\boxed{7}") == 7 @@ -48,6 +59,17 @@ def test_teleqna_answer_label(self): assert parse_teleqna_answer("Explanation: foo\nAnswer: 2") == 2 +class TestGemma4TaskWiring: + def test_gemma4_tasks_are_pre_rendered(self): + # Only the Gemma 4 arms bake enable_thinking into the prompt text; + # every other task lets the benchmark apply the chat template. + assert TASKS["telemath_gemma4"].pre_rendered + pre_rendered = {"telemath_gemma4"} + for name, task in TASKS.items(): + if name not in pre_rendered: + assert not task.pre_rendered, name + + class TestAnswersMatch: def test_int_string_equivalence(self): assert answers_match(42, "42") @@ -150,6 +172,132 @@ def fake_benchmark(dataset_path, model, task, *, host, port, max_concurrency, se # seed advances per run assert sorted({seed for _, seed, _ in calls}) == [0, 1] + def test_per_question_outcomes_saved(self, tmp_path): + dataset = [ + {"id": "a", "prompt": "q0", "answer": 1, "cluster": 0}, + {"id": "b", "prompt": "q1", "answer": 2, "cluster": 0}, + {"id": "c", "prompt": "q2", "answer": 9, "cluster": 1}, + ] + + def fake_benchmark(dataset_path, model, task, *, host, port, max_concurrency, seed, download_dir): + rows = [json.loads(line) for line in open(dataset_path)] + replies = {"q0": "Answer: 1", "q1": "Answer: 3", "q2": "Answer: 9"} + return { + "generated_texts": [replies[r["prompt"]] for r in rows], + "mean_tpot_ms": 10.0, + "output_lens": [5] * len(rows), + } + + out = tmp_path / "outcomes.jsonl" + evaluate_model( + dataset, model="m", task=TASKS["teleqna"], runs=2, + workdir=tmp_path / "splits", benchmark=fake_benchmark, outcomes_out=out, + ) + recs = [json.loads(line) for line in open(out)] + assert len(recs) == 3 * 2 # questions x runs + for r in recs: + assert set(r) == {"qid", "cluster", "run", "correct", "output_len"} + assert r["output_len"] == 5 + by_qid: dict = {} + for r in recs: + by_qid.setdefault(r["qid"], []).append(r["correct"]) + assert by_qid["a"] == [True, True] # q0 -> 1 correct + assert by_qid["b"] == [False, False] # q1 -> 3 wrong + assert by_qid["c"] == [True, True] # q2 -> 9 correct + assert sorted({r["run"] for r in recs}) == [0, 1] + + def test_generations_saved(self, tmp_path): + dataset = [ + {"id": "a", "prompt": "q0", "answer": 1, "cluster": 0}, + {"id": "b", "prompt": "q1", "answer": 2, "cluster": 0}, + {"id": "c", "prompt": "q2", "answer": 9, "cluster": 1}, + ] + + def fake_benchmark(dataset_path, model, task, *, host, port, max_concurrency, seed, download_dir): + rows = [json.loads(line) for line in open(dataset_path)] + replies = {"q0": "Answer: 1", "q1": "Answer: 3", "q2": "Answer: 9"} + return { + "generated_texts": [replies[r["prompt"]] for r in rows], + "mean_tpot_ms": 10.0, + "output_lens": [7] * len(rows), + } + + gens = tmp_path / "generations.jsonl" + evaluate_model( + dataset, model="m", task=TASKS["teleqna"], runs=1, + workdir=tmp_path / "splits", benchmark=fake_benchmark, generations_out=gens, + ) + recs = [json.loads(line) for line in open(gens)] + assert len(recs) == 3 # questions x 1 run + for r in recs: + assert set(r) == { + "qid", "cluster", "run", "question", "prompt", "ground_truth_answer", + "answer", "full_output", "num_tokens", "correct", + } + assert r["num_tokens"] == 7 + by_qid = {r["qid"]: r for r in recs} + # no explicit question field -> falls back to the (raw) prompt + assert by_qid["a"]["question"] == "q0" and by_qid["a"]["prompt"] == "q0" + # a: gold 1, model says "Answer: 1" -> parsed 1, correct + assert by_qid["a"]["full_output"] == "Answer: 1" and by_qid["a"]["correct"] is True + assert by_qid["a"]["ground_truth_answer"] == 1 and by_qid["a"]["answer"] == 1 + # b: gold 2, model says "Answer: 3" -> parsed 3, wrong + assert by_qid["b"]["correct"] is False + assert by_qid["b"]["ground_truth_answer"] == 2 and by_qid["b"]["answer"] == 3 + + def test_generations_question_is_raw_for_prerendered(self, tmp_path): + # a pre_rendered item carries the templated text in `prompt` and the raw + # query in `question`; generations must store the raw question for the QE. + dataset = [ + {"id": "a", "question": "what is 2+2?", + "prompt": "what is 2+2?", "answer": 4, "cluster": 0}, + ] + + def fake_benchmark(dataset_path, model, task, **kw): + rows = [json.loads(line) for line in open(dataset_path)] + return {"generated_texts": ["Answer: 4"] * len(rows), + "mean_tpot_ms": 10.0, "output_lens": [5] * len(rows)} + + gens = tmp_path / "g.jsonl" + evaluate_model( + dataset, model="m", task=TASKS["teleqna"], runs=1, + workdir=tmp_path / "s", benchmark=fake_benchmark, generations_out=gens, + ) + r = json.loads(open(gens).readline()) + assert r["question"] == "what is 2+2?" # raw, QE-facing + assert r["prompt"] == "what is 2+2?" # exact served input + + def test_generations_off_by_default(self, tmp_path): + def fake_benchmark(dataset_path, model, task, **kw): + rows = [json.loads(line) for line in open(dataset_path)] + return {"generated_texts": ["Answer: 1"] * len(rows), "mean_tpot_ms": 10.0} + + # default generations_out=None: no file, no error + evaluate_model( + [{"prompt": "q", "answer": 1, "cluster": 0}], model="m", + task=TASKS["teleqna"], runs=1, workdir=tmp_path / "splits", + benchmark=fake_benchmark, + ) + + def test_question_id_uses_id_then_hash(self): + from cre_router.evaluate import question_id + assert question_id({"id": 7, "prompt": "x"}) == "7" + h = question_id({"prompt": "x"}) + assert len(h) == 12 and question_id({"prompt": "x"}) == h # stable + assert question_id({"prompt": "y"}) != h # different prompt -> different id + + def test_outcomes_off_by_default(self, tmp_path): + def fake_benchmark(dataset_path, model, task, **kw): + rows = [json.loads(line) for line in open(dataset_path)] + return {"generated_texts": ["Answer: 1"] * len(rows), "mean_tpot_ms": 10.0} + + # default outcomes_out=None: no file, no error + evaluate_model( + [{"prompt": "q", "answer": 1, "cluster": 0}], model="m", + task=TASKS["teleqna"], runs=1, workdir=tmp_path / "splits", + benchmark=fake_benchmark, + ) + def test_missing_tpot_raises(self, tmp_path): def bad_benchmark(dataset_path, model, task, **kwargs): rows = [json.loads(line) for line in open(dataset_path)] diff --git a/tests/test_prep_gemma4.py b/tests/test_prep_gemma4.py new file mode 100644 index 0000000..7ca3cea --- /dev/null +++ b/tests/test_prep_gemma4.py @@ -0,0 +1,29 @@ +"""render_rows preserves the raw question when baking the chat template.""" + +import importlib.util +from pathlib import Path + +_spec = importlib.util.spec_from_file_location( + "prep_gemma4", Path(__file__).parent.parent / "data" / "prep_gemma4_thinking.py" +) +prep_gemma4 = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(prep_gemma4) + + +class _FakeTokenizer: + def apply_chat_template(self, messages, add_generation_prompt, tokenize, enable_thinking): + return f"{messages[0]['content']}" + + +def test_render_preserves_raw_question_and_other_fields(): + rows = [{"id": "a", "prompt": "what is 2+2?", "answer": 4, "category": "math"}] + out = prep_gemma4.render_rows(rows, _FakeTokenizer(), enable_thinking=True) + assert out[0]["question"] == "what is 2+2?" # raw preserved + assert out[0]["prompt"] == "what is 2+2?" # templated + assert out[0]["answer"] == 4 and out[0]["category"] == "math" # untouched + + +def test_render_keeps_explicit_question_if_present(): + rows = [{"prompt": "already raw", "question": "the real question", "answer": 1}] + out = prep_gemma4.render_rows(rows, _FakeTokenizer(), enable_thinking=False) + assert out[0]["question"] == "the real question" diff --git a/tests/test_prep_qe.py b/tests/test_prep_qe.py new file mode 100644 index 0000000..18188e1 --- /dev/null +++ b/tests/test_prep_qe.py @@ -0,0 +1,83 @@ +"""QE-dataset builder: generations JSONL -> qe-train {train,test}.jsonl.""" + +import importlib.util +import json +from pathlib import Path + +import pytest + +_spec = importlib.util.spec_from_file_location( + "prep_qe", Path(__file__).parent.parent / "data" / "prep_qe.py" +) +prep_qe = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(prep_qe) + + +def _gen(qid, cluster, correct, out="the answer is 4", ntok=12): + return { + "qid": qid, "cluster": cluster, "run": 0, "prompt": f"question {qid}?", + "ground_truth_answer": 4, "answer": 4 if correct else 3, + "full_output": out, "num_tokens": ntok, "correct": correct, + } + + +class TestQeRow: + def test_correct_maps_to_accept(self): + row = prep_qe.qe_row(_gen("a", 0, True)) + assert row["decision_label"] == 1 and row["decision_str"] == "accept" + assert row["score"] == 1.0 and row["accuracy"] == 1.0 + + def test_wrong_maps_to_route(self): + row = prep_qe.qe_row(_gen("b", 1, False)) + assert row["decision_label"] == 0 and row["decision_str"] == "route" + assert row["score"] == 0.0 + + def test_schema_and_derived_fields(self): + row = prep_qe.qe_row(_gen("a", 0, True, out="one two three", ntok=7)) + # columns cre qe-train needs plus the router-parity extras + assert {"question", "full_output", "num_tokens", "decision_label"} <= set(row) + assert row["question"] == "question a?" + assert row["num_words"] == 3 and row["num_tokens"] == 7 + assert row["cluster"] == 0 and row["qid"] == "a" + + def test_question_prefers_explicit_field(self): + row = prep_qe.qe_row({**_gen("a", 0, True), "question": "raw?"}) + assert row["question"] == "raw?" + + +class TestBuild: + def test_writes_pooled_train_test_splits(self, tmp_path): + trainf = tmp_path / "train_gen.jsonl" + testf = tmp_path / "test_gen.jsonl" + trainf.write_text("\n".join(json.dumps(_gen(f"t{i}", i % 2, i % 2 == 0)) for i in range(4))) + testf.write_text("\n".join(json.dumps(_gen(f"e{i}", 0, True)) for i in range(2))) + + out = tmp_path / "router" + sizes = prep_qe.build([str(trainf)], [str(testf)], str(out)) + assert sizes == {"train": 4, "test": 2} + + train_rows = [json.loads(l) for l in open(out / "train.jsonl")] + test_rows = [json.loads(l) for l in open(out / "test.jsonl")] + assert len(train_rows) == 4 and len(test_rows) == 2 + # labels present and binary + assert all(r["decision_label"] in (0, 1) for r in train_rows) + # i even -> correct -> accept(1); i odd -> route(0) + assert [r["decision_label"] for r in train_rows] == [1, 0, 1, 0] + + def test_pools_multiple_files(self, tmp_path): + f1 = tmp_path / "a.jsonl"; f1.write_text(json.dumps(_gen("a", 0, True))) + f2 = tmp_path / "b.jsonl"; f2.write_text(json.dumps(_gen("b", 0, False))) + out = tmp_path / "router" + sizes = prep_qe.build([str(f1), str(f2)], [str(f1)], str(out)) + assert sizes["train"] == 2 + + def test_drops_rows_with_null_num_tokens(self, tmp_path): + good = _gen("a", 0, True) + bad = {**_gen("b", 0, False), "num_tokens": None} + f = tmp_path / "g.jsonl" + f.write_text("\n".join(json.dumps(r) for r in [good, bad])) + out = tmp_path / "router" + sizes = prep_qe.build([str(f)], [str(f)], str(out)) + assert sizes["train"] == 1 # the null-num_tokens row is dropped + rows = [json.loads(l) for l in open(out / "train.jsonl")] + assert [r["qid"] for r in rows] == ["a"] diff --git a/tests/test_qe_cascade.py b/tests/test_qe_cascade.py new file mode 100644 index 0000000..fcdee45 --- /dev/null +++ b/tests/test_qe_cascade.py @@ -0,0 +1,162 @@ +"""Stage 1+2 cascade evaluation: QE decisions -> per-cluster cascade accuracy +and escalation counts, with no GPU (the classifier is stubbed).""" + +import json +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from cre_router.qe.cascade import ( + compose_cascade, + run_qe, + strong_correct_by_qid, + write_cascade_stats, +) +from cre_router.routing import cluster_cascade_accuracy + +CONFIGS = Path(__file__).parent.parent / "configs" + + +def _gen(qid, cluster, run, correct): + return {"qid": qid, "cluster": cluster, "run": run, "correct": correct, + "full_output": f"out-{qid}", "num_tokens": 10, "prompt": f"q-{qid}"} + + +class TestStrongCorrectByQid: + def test_mean_over_runs(self): + outcomes = [ + {"qid": "a", "correct": True}, {"qid": "a", "correct": False}, + {"qid": "b", "correct": True}, {"qid": "b", "correct": True}, + ] + assert strong_correct_by_qid(outcomes) == {"a": 0.5, "b": 1.0} + + +class TestComposeCascade: + def test_per_query_composition_and_escalation_count(self): + # cluster 0, 2 queries x 2 runs. a: weak-correct + accepted; b: weak-wrong + escalated. + gens = [ + _gen("a", 0, 0, True), _gen("b", 0, 0, False), + _gen("a", 0, 1, True), _gen("b", 0, 1, False), + ] + escalate = [False, True, False, True] + strong = {"b": 0.5} # strong right on b half its runs + report = compose_cascade(gens, escalate, strong) + # correct: a,a -> 1+1 ; b,b escalated -> 0.5+0.5 ; /4 = 0.75 + assert report["0"]["cascade_accuracy"] == pytest.approx(0.75) + # 2 escalated rows over 2 runs -> 1 escalation/run + assert report["0"]["escalations"] == pytest.approx(1.0) + assert report["0"]["n"] == 4 + + def test_multiple_clusters(self): + gens = [_gen("a", 0, 0, True), _gen("b", 1, 0, False)] + report = compose_cascade(gens, [False, True], {"b": 1.0}) + assert report["0"]["cascade_accuracy"] == pytest.approx(1.0) + assert report["1"]["cascade_accuracy"] == pytest.approx(1.0) + assert report["1"]["escalations"] == pytest.approx(1.0) + + def test_escalation_absent_in_some_runs_divides_by_all_runs(self): + # AIME C1 shape: one query over 5 runs; escalated in runs 0,1,2 only. + # escalations must divide by the 5 runs present, not the 3 escalated ones. + gens = [_gen("x", 1, r, correct=(r >= 3)) for r in range(5)] + escalate = [True, True, True, False, False] + report = compose_cascade(gens, escalate, {"x": 1.0}) + assert report["1"]["escalations"] == pytest.approx(0.6) # 3/5, not 3/3 + assert report["1"]["n"] == 5 + # 3 escalated -> strong 1.0 ; runs 3,4 accepted + weak-correct -> 5/5 + assert report["1"]["cascade_accuracy"] == pytest.approx(1.0) + + def test_misaligned_raises(self): + with pytest.raises(ValueError, match="align"): + compose_cascade([_gen("a", 0, 0, True)], [False, True], {}) + + def test_missing_strong_outcome_raises(self): + with pytest.raises(KeyError, match="qid"): + compose_cascade([_gen("a", 0, 0, False)], [True], strong_correct={}) + + +class TestRunQe: + def test_stub_classifier_produces_escalate(self): + @dataclass + class Decision: + accept: bool + + class Stub: + # accept when the output ends in "keep", else route + def predict_batch(self, items): + return [Decision(accept=o.endswith("keep")) for _, o, _ in items] + + gens = [ + {"prompt": "q0", "full_output": "... keep", "num_tokens": 5, "cluster": 0, "run": 0, "qid": "0", "correct": True}, + {"prompt": "q1", "full_output": "... drop", "num_tokens": 5, "cluster": 0, "run": 0, "qid": "1", "correct": False}, + ] + assert run_qe(Stub(), gens, batch_size=1) == [False, True] # keep->accept->no escalate; drop->route + + +class TestWriteCascadeStats: + def test_merges_into_routing_json_and_feeds_cre_cascade(self, tmp_path): + from cre_router.routing import ( + cascade_system_accuracy, + models_from_stats, + ) + + # a routing-side cascade config, no Stage 2 fields yet + cfg = tmp_path / "cascade.json" + cfg.write_text(json.dumps({ + "cluster_sizes": {"0": 2, "1": 2}, + "assignment": {"0": "weak", "1": "strong"}, + "models": { + "weak": {"errors": {"0": 0.5, "1": 0.5}, "cluster_tpot_ms": {"0": 10.0, "1": 10.0}}, + "strong": {"errors": {"0": 0.1, "1": 0.2}, "cluster_tpot_ms": {"0": 20.0, "1": 20.0}}, + }, + })) + report = {"0": {"cascade_accuracy": 0.9, "escalations": 1.0, "n": 4}} + write_cascade_stats(cfg, "strong", report) + + stats = json.loads(cfg.read_text()) + assert stats["escalations"]["0"] == ["strong", 1.0] + assert stats["cascade_accuracy"]["0"] == 0.9 + # cre cascade composition: C0 cascade 0.9, C1 direct = strong 1-0.2=0.8 + models, sizes = models_from_stats(stats) + assignment = {str(k): str(v) for k, v in stats["assignment"].items()} + acc = cascade_system_accuracy(models, assignment, sizes, + {str(k): float(v) for k, v in stats["cascade_accuracy"].items()}) + assert acc == pytest.approx((2 * 0.9 + 2 * 0.8) / 4) + + +class TestPaperCascadeReconstruction: + """Locks the session's validation that the composition reproduces the paper's + per-cluster cascade accuracy from the real per-run data: + TeleQnA C0 from route.zip + `QE-Route TeleQnA.numbers` + inference_teleqna_log_50; + AIME C1 from `AIME-Clusters-Test.numbers` + route.zip. See memory + cre-stage1plus2-metrics-composition.""" + + def test_teleqna_c0_reconstructs_0_742(self): + # per efficient/QE run (route.zip cluster_0_run_0..4): routed count, the + # accepted-and-efficient-correct count (log_50 "Accept: c/d"), and Gemma-26B + # accuracy on that routed set (mean of its 5 repeats in QE-Route TeleQnA). + routed = [199, 205, 206, 200, 202] + accepted_correct = [305, 295, 301, 310, 304] + s = [0.671, 0.683, 0.670, 0.654, 0.647] + N = 590 + per_run = [(a + r * si) / N for r, a, si in zip(routed, accepted_correct, s)] + mean = sum(per_run) / len(per_run) + assert mean == pytest.approx(0.742, abs=0.002) # paper Table teleqna_test 0.743 + cfg = json.loads((CONFIGS / "teleqna_cascade_test.json").read_text()) + assert cfg["cascade_accuracy"]["0"] == pytest.approx(mean, abs=0.005) + + def test_aime_c1_reconstructs_0_96(self): + # cluster 1: size 10 x 5 runs = 50 instances. VibeThinker acc 0.9 every run + # = exactly one miss/run (the same hard query); the QE escalates it in 3 of + # 5 runs (route.zip counts 1,1,1,0,0); Qwen3-30B is correct on those 3. + weak, strong, escalate = [], [], [] + for run in range(5): + for q in range(10): + is_hard = q == 5 + weak.append(not is_hard) # 9 correct, the hard one wrong + escalate.append(is_hard and run in {0, 1, 2}) + strong.append(True) # strong right on the escalated hard query + acc = cluster_cascade_accuracy(weak, strong, escalate) + assert acc == pytest.approx(0.96) # paper Table aime_test 0.96 + cfg = json.loads((CONFIGS / "aime_cascade_test.json").read_text()) + assert cfg["cascade_accuracy"]["1"] == pytest.approx(acc) diff --git a/tests/test_qe_classifier.py b/tests/test_qe_classifier.py new file mode 100644 index 0000000..a05844b --- /dev/null +++ b/tests/test_qe_classifier.py @@ -0,0 +1,47 @@ +"""format_qe_input: think-stripping (all model families) + last-N-words +truncation. Pure string formatting, no torch/transformers needed.""" + +from cre_router.qe.classifier import format_qe_input + + +def test_strips_qwen_thinking(): + out = format_qe_input("q", "long private reasoningFinal: 42", 123, "[SEP]") + assert out == "q [SEP] Final: 42 [SEP] 123" + assert "reasoning" not in out + + +def test_strips_gemma_channel(): + out = format_qe_input("q", "<|channel>thought\nscratch workAnswer 7", 5, "[SEP]") + assert "thought" not in out and "scratch" not in out + assert out == "q [SEP] Answer 7 [SEP] 5" + + +def test_idempotent_on_plain_output(): + # no reasoning marker -> unchanged (safe on already-de-thought data) + out = format_qe_input("q", "just the answer", 9, "[SEP]") + assert out == "q [SEP] just the answer [SEP] 9" + + +def test_num_tokens_is_the_full_length_not_the_stripped_text(): + # num_tokens reflects the whole generation (cost), even though the thinking + # text is dropped from the classifier's view. + out = format_qe_input("q", "" + "x " * 50 + "done", 999, "[SEP]") + assert out.endswith("[SEP] 999") + assert "x" not in out.split(" [SEP] ")[1] + + +def test_truncation_applies_after_stripping(): + answer = " ".join(str(i) for i in range(2000)) # 2000 words after the think block + out = format_qe_input( + "q", "" + ("noise " * 5000) + "" + answer, 42, "[SEP]", max_output_words=1000 + ) + body = out.split(" [SEP] ")[1] + assert "noise" not in body # reasoning gone + assert len(body.split()) == 1000 # kept the last 1000 words + assert body.split()[-1] == "1999" # of the answer, not the thinking + + +def test_none_num_tokens_raises(): + import pytest + with pytest.raises(ValueError, match="num_tokens is None"): + format_qe_input("q", "answer", None, "[SEP]") diff --git a/tests/test_routing.py b/tests/test_routing.py index ffa27a1..c27e95f 100644 --- a/tests/test_routing.py +++ b/tests/test_routing.py @@ -12,7 +12,12 @@ import pytest from cre_router.routing import ( + ModelStats, assign, + cascade_system_accuracy, + cascade_system_metrics, + cascade_system_metrics_ntier, + cluster_cascade_accuracy, crossover_candidates, eta, models_from_stats, @@ -133,3 +138,164 @@ def test_dominated_models_never_selected(self, teleqna): chosen = set(region.assignment.values()) assert "Gemma4-E2B" not in chosen assert "Gemma4-E4B" not in chosen + + +def _load_cascade(name: str): + """Load a checked-in cascade config the way ``cre cascade`` does.""" + stats = json.loads((CONFIGS / name).read_text()) + models, cluster_sizes = models_from_stats(stats) + assignment = {str(k): str(v) for k, v in stats["assignment"].items()} + escalations = { + str(k): (str(v[0]), float(v[1])) for k, v in stats.get("escalations", {}).items() + } + return models, assignment, cluster_sizes, escalations + + +class TestCascadeSystemMetrics: + """Stage 1+2 system latency from the checked-in test-split cascade configs, + verified against the paper's Tables `aime_test` and `teleqna_test`. The + composed values are 9.75 / 23.65 ms; the paper reports 9.7 / 23.8 ms.""" + + def test_aime_stage1plus2_latency(self): + models, assignment, sizes, escalations = _load_cascade("aime_cascade_test.json") + tpot, e2el = cascade_system_metrics(models, assignment, sizes, escalations) + assert tpot == pytest.approx(9.7, abs=0.1) # paper Table aime_test + assert e2el == pytest.approx(156300, rel=1e-3) + + def test_teleqna_stage1plus2_latency(self): + models, assignment, sizes, escalations = _load_cascade("teleqna_cascade_test.json") + tpot, e2el = cascade_system_metrics(models, assignment, sizes, escalations) + assert tpot == pytest.approx(23.8, abs=0.2) # paper Table teleqna_test + assert e2el == pytest.approx(1127, rel=1e-3) + + def test_no_escalation_matches_stage1(self): + """With no escalations the cascade collapses to Stage 1 TPOT exactly.""" + models, assignment, sizes, _ = _load_cascade("teleqna_cascade_test.json") + _, stage1_tpot = system_metrics(models, assignment, sizes) + tpot, _ = cascade_system_metrics(models, assignment, sizes, escalations={}) + assert tpot == pytest.approx(stage1_tpot) + + +class TestCascadeSystemMetricsNTier: + """The N-tier generalisation; 2-tier ``cascade_system_metrics`` delegates to + it, so the AIME/TeleQnA tests above are the N=2 regression.""" + + def _models(self): + # per-cluster tpot / e2el / output length for a single cluster "0" + eff = ModelStats(name="eff", tpot_ms=10.0, errors={"0": 0.5}, + cluster_tpot_ms={"0": 10.0}, e2el_ms=100.0, + cluster_e2el_ms={"0": 100.0}, cluster_output_tokens={"0": 50.0}) + mid = ModelStats(name="mid", tpot_ms=20.0, errors={"0": 0.3}, + cluster_tpot_ms={"0": 20.0}, e2el_ms=300.0, + cluster_e2el_ms={"0": 300.0}, cluster_output_tokens={"0": 100.0}) + strong = ModelStats(name="strong", tpot_ms=30.0, errors={"0": 0.1}, + cluster_tpot_ms={"0": 30.0}, e2el_ms=600.0, + cluster_e2el_ms={"0": 600.0}, cluster_output_tokens={"0": 200.0}) + return [eff, mid, strong] + + def test_three_tier_hand_computed(self): + # reach [10,4,2]: 10 run eff, 4 escalate to mid, 2 further to strong. + # E2EL = 10*100 + 4*300 + 2*600 = 3400 -> /10 = 340 + # TPOT: t0 6*(500/50)=60 ; t1 2*(2500/100)=50 ; t2 2*(8500/200)=85 -> 195/10 = 19.5 + cascades = {"0": [("eff", 10), ("mid", 4), ("strong", 2)]} + tpot, e2el = cascade_system_metrics_ntier(self._models(), cascades, {"0": 10}) + assert e2el == pytest.approx(340.0) + assert tpot == pytest.approx(19.5) + + def test_single_tier_is_direct_assignment(self): + cascades = {"0": [("mid", 10)]} + tpot, e2el = cascade_system_metrics_ntier(self._models(), cascades, {"0": 10}) + assert (tpot, e2el) == pytest.approx((20.0, 300.0)) + + def test_two_tier_wrapper_equals_ntier(self): + models = self._models() + sizes = {"0": 10} + direct = cascade_system_metrics_ntier( + models, {"0": [("eff", 10), ("strong", 4)]}, sizes + ) + wrapped = cascade_system_metrics( + models, {"0": "eff"}, sizes, {"0": ("strong", 4)} + ) + assert wrapped == pytest.approx(direct) + + def test_rejects_increasing_reach(self): + with pytest.raises(ValueError, match="non-increasing"): + cascade_system_metrics_ntier( + self._models(), {"0": [("eff", 10), ("mid", 12)]}, {"0": 10} + ) + + def test_rejects_base_reach_mismatch(self): + with pytest.raises(ValueError, match="cluster size"): + cascade_system_metrics_ntier( + self._models(), {"0": [("eff", 8), ("mid", 4)]}, {"0": 10} + ) + + +class TestClusterCascadeAccuracy: + def test_per_query_composition(self): + # accept -> keep weak; escalate -> take strong. FP (escalated-correct) and + # FN (accepted-wrong) both handled by taking the actual per-query outcome. + weak = [True, True, False, False] + strong = [False, False, True, False] + escalate = [False, False, True, True] + # q0,q1 accepted+weak-correct; q2 escalated+strong-correct; q3 escalated+strong-wrong + assert cluster_cascade_accuracy(weak, strong, escalate) == pytest.approx(3 / 4) + + def test_no_escalation_equals_weak(self): + weak = [True, False, True] + assert cluster_cascade_accuracy(weak, [False, False, False], [False, False, False]) == pytest.approx(2 / 3) + + def test_misaligned_lengths_raise(self): + with pytest.raises(ValueError, match="align"): + cluster_cascade_accuracy([True], [True, False], [False, False]) + + def test_empty_raises(self): + with pytest.raises(ValueError, match="empty"): + cluster_cascade_accuracy([], [], []) + + +class TestCascadeSystemAccuracy: + """Stage 1+2 system accuracy composition, verified against the paper's + combined-cascade slides (AIME 88.4%, TeleQnA 74.3%).""" + + def test_aime_stage1plus2_accuracy(self): + # AIME test: C0->Q3, C1->V (cascade to Q3), C2->Q3. Q3 test errors give + # 0.867/0.980/0.829; C1's cascade accuracy is 0.96. + v = ModelStats(name=V, tpot_ms=4.8, errors={"0": 0.311, "1": 0.100, "2": 0.291}) + q = ModelStats(name=Q, tpot_ms=11.8, errors={"0": 0.133, "1": 0.020, "2": 0.171}) + assignment = {"0": Q, "1": V, "2": Q} + sizes = {"0": 9, "1": 10, "2": 11} + acc = cascade_system_accuracy([v, q], assignment, sizes, cascade_accuracy={"1": 0.96}) + assert acc == pytest.approx(0.884, abs=0.001) + + def test_teleqna_stage1plus2_accuracy(self): + # TeleQnA test: C0->Q-4B (cascade to G-26B, acc 0.740), C1->G-26B direct. + q4b = ModelStats(name="Q-4B", tpot_ms=15.1, errors={"0": 0.311, "1": 0.360}) + g26 = ModelStats(name="G-26B", tpot_ms=24.5, errors={"0": 0.223, "1": 0.254}) + assignment = {"0": "Q-4B", "1": "G-26B"} + sizes = {"0": 590, "1": 410} + acc = cascade_system_accuracy([q4b, g26], assignment, sizes, cascade_accuracy={"0": 0.740}) + assert acc == pytest.approx(0.743, abs=0.001) + + def test_no_cascade_matches_stage1_accuracy(self): + # With no escalated clusters, system accuracy == Stage 1 accuracy. + q4b = ModelStats(name="Q-4B", tpot_ms=15.1, errors={"0": 0.311, "1": 0.360}) + g26 = ModelStats(name="G-26B", tpot_ms=24.5, errors={"0": 0.223, "1": 0.254}) + assignment = {"0": "Q-4B", "1": "G-26B"} + sizes = {"0": 590, "1": 410} + stage1_acc, _ = system_metrics([q4b, g26], assignment, sizes) + acc = cascade_system_accuracy([q4b, g26], assignment, sizes, cascade_accuracy={}) + assert acc == pytest.approx(stage1_acc) + + def _cascade_acc_from_config(self, name: str) -> float: + stats = json.loads((CONFIGS / name).read_text()) + models, sizes = models_from_stats(stats) + assignment = {str(k): str(v) for k, v in stats["assignment"].items()} + cascade_accuracy = {str(k): float(v) for k, v in stats.get("cascade_accuracy", {}).items()} + return cascade_system_accuracy(models, assignment, sizes, cascade_accuracy) + + def test_aime_config_reproduces_884(self): + assert self._cascade_acc_from_config("aime_cascade_test.json") == pytest.approx(0.884, abs=0.001) + + def test_teleqna_config_reproduces_743(self): + assert self._cascade_acc_from_config("teleqna_cascade_test.json") == pytest.approx(0.743, abs=0.001) diff --git a/tests/test_telemath.py b/tests/test_telemath.py new file mode 100644 index 0000000..eae0d5c --- /dev/null +++ b/tests/test_telemath.py @@ -0,0 +1,115 @@ +"""TeleMath numeric answer parsing and tolerance matching. + +TeleMath answers are numerical quantities (floats, scientific notation, and +negatives), scored by a relative tolerance rather than exact match. +""" + +import pytest + +from cre_router.evaluate import ( + TASKS, + numeric_match, + parse_telemath_answer, + score_generations, +) + + +class TestParseTelemathAnswer: + @pytest.mark.parametrize( + "text, expected", + [ + (r"The result is \boxed{233.333333333333}.", 233.333333333333), + ("Answer: 6.0", 6.0), + (r"so \boxed{7.2e-05}", 7.2e-05), + (r"final \boxed{7.2 \times 10^{-5}}", 7.2e-5), + (r"\boxed{7.2 \times 10^-5}", 7.2e-5), + ("Answer: -62.085424660791375", -62.085424660791375), + ("the current is 0.25 A", 0.25), + ("Answer: 50000.0", 50000.0), + ], + ) + def test_extracts_value(self, text, expected): + assert parse_telemath_answer(text) == pytest.approx(expected) + + def test_boxed_wins_over_earlier_numbers(self): + assert parse_telemath_answer(r"tried 12 and 3, so \boxed{204}") == 204.0 + + def test_answer_label_wins_over_reasoning(self): + assert parse_telemath_answer("we get 3.14 then 2.71.\nAnswer: 42.0") == 42.0 + + def test_reasoning_is_stripped(self): + assert parse_telemath_answer(r"maybe 5\boxed{9.0}") == 9.0 + + def test_no_number_returns_none(self): + assert parse_telemath_answer("no numeric answer here") is None + + def test_leading_dot_decimal(self): + assert parse_telemath_answer(r"\boxed{.5}") == pytest.approx(0.5) + + def test_answer_at_end_of_long_reasoning(self): + # A long completion whose answer sits at the very end still parses; the + # tail window keeps parsing bounded without clipping a real answer. + text = "reasoning that goes on and on. " * 4000 + r"\boxed{233.5}" + assert parse_telemath_answer(text) == pytest.approx(233.5) + + @pytest.mark.parametrize( + "text", + [ + "9" * 200_000 + " done", # long digit run, no closing token + r"\boxed{" + "1" * 200_000, # truncated boxed, unclosed brace + "3." + "3" * 200_000, # runaway decimal expansion + ], + ) + def test_pathological_input_is_fast(self, text): + # Degenerate/truncated generations used to send the number regex into + # O(n^2) backtracking and stall for hours. Parsing must stay well under a + # second regardless of output length. + import time + + start = time.perf_counter() + parse_telemath_answer(text) + assert time.perf_counter() - start < 1.0 + + +class TestNumericMatch: + def test_exact(self): + assert numeric_match(233.333333, 233.333333) + + def test_within_one_percent(self): + # rounding of the same quantity is accepted + assert numeric_match(233.33, 233.333333) + + def test_outside_tolerance_rejected(self): + assert not numeric_match(200.0, 233.333) + + def test_scientific_within_tolerance(self): + assert numeric_match(7.2e-05, 7.21e-05) + + def test_negative(self): + assert numeric_match(-62.09, -62.085424660791375) + + def test_near_zero_gold(self): + assert numeric_match(0.0, 0.0) + assert numeric_match(1e-10, 0.0) + assert not numeric_match(0.5, 0.0) + + def test_gold_as_string(self): + assert numeric_match(6.0, "6.0") + + def test_none_prediction(self): + assert not numeric_match(None, 6.0) + + +class TestTelemathTaskWiring: + def test_task_uses_numeric_match(self): + assert TASKS["telemath"].match is numeric_match + assert TASKS["telemath"].parse is parse_telemath_answer + + def test_score_generations_uses_task_matcher(self): + texts = [r"\boxed{233.33}", r"\boxed{6.0}", r"\boxed{999.0}"] + gold = [233.333333, 6.0, 6.0] # third is wrong + error, correct = score_generations( + texts, gold, TASKS["telemath"].parse, TASKS["telemath"].match + ) + assert correct == [True, True, False] + assert error == pytest.approx(1 / 3)