From db5158c584c824e71a6c5fec3ff02b0692b08a0f Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Thu, 23 Jul 2026 22:50:42 +0100 Subject: [PATCH 1/9] Add E2EL cost metric and Stage 1+2 cascade system metrics Measure end-to-end latency (E2EL), request throughput, and truncation rate alongside TPOT in `cre evaluate`, and allow routing to be fit on either metric with `cre fit --cost-metric {tpot,e2el}` (default tpot, which reproduces the published results). E2EL is the informative cost when pool members differ in output length rather than decode speed, where TPOT cannot tell them apart. Add `cascade_system_metrics` and a `cre cascade` command that compose Stage 1+2 system TPOT (query-weighted, charging an escalated query's discarded efficient pass per delivered token) and E2EL (the per-request sum of both passes) from per-cluster measurements and the escalation rate. Add the two per-cluster test-split stat files (configs/*_cascade_test.json) so `cre cascade` reproduces the paper's Stage 1+2 latency with no GPU, and document it in REPRODUCE. --- REPRODUCE.md | 21 ++++ configs/aime_cascade_test.json | 20 ++++ configs/teleqna_cascade_test.json | 20 ++++ src/cre_router/cli.py | 66 +++++++++-- src/cre_router/evaluate.py | 123 +++++++++++++++++--- src/cre_router/routing.py | 157 +++++++++++++++++++++++--- tests/test_cost_metric.py | 182 ++++++++++++++++++++++++++++++ tests/test_routing.py | 37 ++++++ 8 files changed, 580 insertions(+), 46 deletions(-) create mode 100644 configs/aime_cascade_test.json create mode 100644 configs/teleqna_cascade_test.json create mode 100644 tests/test_cost_metric.py diff --git a/REPRODUCE.md b/REPRODUCE.md index 6505716..a3ed6e2 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: diff --git a/configs/aime_cascade_test.json b/configs/aime_cascade_test.json new file mode 100644 index 0000000..48989de --- /dev/null +++ b/configs/aime_cascade_test.json @@ -0,0 +1,20 @@ +{ + "_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). `cre cascade` gives 9.75 ms TPOT and 156303 ms E2EL, matching the paper's Table aime_test (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]}, + "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..f0c3756 --- /dev/null +++ b/configs/teleqna_cascade_test.json @@ -0,0 +1,20 @@ +{ + "_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). `cre cascade` gives 23.65 ms TPOT and 1127 ms E2EL, matching the paper's Table teleqna_test (23.8 ms) to within rounding.", + "cluster_sizes": {"0": 590, "1": 410}, + "assignment": {"0": "Qwen3-4B", "1": "Gemma4-26B"}, + "escalations": {"0": ["Gemma4-26B", 202]}, + "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/src/cre_router/cli.py b/src/cre_router/cli.py index 59596a2..650f4b1 100644 --- a/src/cre_router/cli.py +++ b/src/cre_router/cli.py @@ -21,6 +21,7 @@ from cre_router.artifacts import RouterArtifacts from cre_router.clustering import DEFAULT_EMBEDDING_MODEL from cre_router.routing import ( + cascade_system_metrics, eta, models_from_stats, pareto_prune, @@ -83,13 +84,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 +103,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 +136,27 @@ 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) + + 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}" + print(line) + print(f"\n 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, @@ -258,10 +287,25 @@ def main(argv: list[str] | None = None) -> None: 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)", diff --git a/src/cre_router/evaluate.py b/src/cre_router/evaluate.py index 0ff6b40..37c235b 100644 --- a/src/cre_router/evaluate.py +++ b/src/cre_router/evaluate.py @@ -163,6 +163,19 @@ 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 = ( + "e2el_ms", + "request_throughput", + "mean_output_tokens", + "truncated_frac", +) + + @dataclass class RunMeasurement: """One (cluster, run) benchmark outcome, kept for provenance.""" @@ -172,27 +185,95 @@ class RunMeasurement: error: float tpot_ms: float num_prompts: int + 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 + + 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 +297,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") # --------------------------------------------------------------------------- @@ -286,6 +367,11 @@ def run_vllm_benchmark( 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. @@ -350,6 +436,7 @@ def evaluate_model( error=error, tpot_ms=float(tpot_ms), num_prompts=len(items), + **optional_metrics(result, len(items), task.max_tokens), ) ) return measurements diff --git a/src/cre_router/routing.py b/src/cre_router/routing.py index 0caa61e..6ec4d0d 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,62 @@ 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( + 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 the full 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. + + Both metrics are query-weighted, matching Stage 1 and vLLM's per-request Mean + TPOT. They differ only in how an escalated query is charged: + + - **TPOT** charges both passes per delivered token, + ``TPOT_strong + TPOT_eff * (L_eff / L_strong)``: the discarded efficient + generation is amortised over the strong model's delivered answer. + - **E2EL** is a plain sum, ``E2EL_eff + E2EL_strong``: the user waits for the + efficient generation to complete, since Stage 2 inspects the whole output + before escalating, and then for the strong generation. + + Returns ``(tpot_ms, e2el_ms)``. Reproduces the paper's Stage 1+2 latency, + 9.7 ms on AIME and 23.8 ms on TeleQnA. + """ + by_name = {m.name: m for m in models} + total = sum(cluster_sizes.values()) + tpot_sum = 0.0 + e2el_sum = 0.0 + for cluster, name in assignment.items(): + efficient = by_name[name] + size = cluster_sizes[cluster] + if cluster in escalations: + strong_name, count = escalations[cluster] + strong = by_name[strong_name] + direct = size - count + ratio = efficient.output_tokens_for(cluster) / strong.output_tokens_for(cluster) + escalated_tpot = strong.tpot_for(cluster) + efficient.tpot_for(cluster) * ratio + tpot_sum += direct * efficient.tpot_for(cluster) + count * escalated_tpot + e2el_sum += direct * efficient.e2el_for(cluster) + count * ( + efficient.e2el_for(cluster) + strong.e2el_for(cluster) + ) + else: + tpot_sum += size * efficient.tpot_for(cluster) + e2el_sum += size * efficient.e2el_for(cluster) + return tpot_sum / total, e2el_sum / total + + def eta( models: list[ModelStats], assignment: dict[str, str], @@ -229,9 +329,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 +346,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 +359,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 +376,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/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_routing.py b/tests/test_routing.py index ffa27a1..0b6b98f 100644 --- a/tests/test_routing.py +++ b/tests/test_routing.py @@ -13,6 +13,7 @@ from cre_router.routing import ( assign, + cascade_system_metrics, crossover_candidates, eta, models_from_stats, @@ -133,3 +134,39 @@ 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) From ffd4de7dd2b20001b04a12d7c709d62d27a2bef4 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:52:32 +0100 Subject: [PATCH 2/9] Add TeleMath task with per-question outcome capture Add TeleMath (netop/TeleMath) as an evaluation task. Its answers are numerical, so they are matched by relative tolerance rather than exact equality. Separate arms cover thinking and non-thinking models, plus a pre-rendered Gemma 4 variant whose thinking switch is a chat-template kwarg baked into the prompt text at prep time, since vLLM's benchmark loader never forwards template kwargs. Record per-question outcomes (qid, cluster, run, correct, output_len) and, behind a flag, full generations with token counts. Both are irrecoverable once a run finishes, and are needed for confidence intervals and for building Stage 2 QE training data. A stable question_id lets outcomes join across models and runs. Bound the TeleMath answer search to the output tail and use a number pattern with no adjacent variable-length digit runs, keeping parsing linear in the window. Move split_thinking into cre_router.textutils, now shared by the answer parsers and the QE input formatter. --- data/prep_gemma4_thinking.py | 89 ++++++++++++++ data/prep_telemath.py | 98 +++++++++++++++ src/cre_router/evaluate.py | 230 +++++++++++++++++++++++++++++++++-- src/cre_router/textutils.py | 23 ++++ tests/test_evaluate.py | 148 ++++++++++++++++++++++ tests/test_prep_gemma4.py | 29 +++++ tests/test_telemath.py | 115 ++++++++++++++++++ 7 files changed, 721 insertions(+), 11 deletions(-) create mode 100644 data/prep_gemma4_thinking.py create mode 100644 data/prep_telemath.py create mode 100644 src/cre_router/textutils.py create mode 100644 tests/test_prep_gemma4.py create mode 100644 tests/test_telemath.py 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_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/evaluate.py b/src/cre_router/evaluate.py index 37c235b..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 @@ -169,6 +282,7 @@ def cluster_sizes(dataset: list[dict]) -> dict[str, int]: # 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", @@ -185,6 +299,7 @@ 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 @@ -212,6 +327,10 @@ def optional_metrics( 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") @@ -363,6 +482,19 @@ 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") @@ -382,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, @@ -395,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: @@ -425,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'") @@ -439,4 +635,16 @@ def evaluate_model( **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/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_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_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) From a67eeb4440d3e335b228a49e84ccaccd6a7bd551 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:52:43 +0100 Subject: [PATCH 3/9] Generalise cascade cost to N tiers, add Stage 1+2 system accuracy Cascade cost now takes a full escalation chain per cluster, ordered from the Stage 1 model up to the strongest tier, with a reach count per tier. A cluster can therefore route direct, gate once, or gate repeatedly. The two-tier cascade_system_metrics becomes a thin wrapper over it, so the published 9.7 ms AIME and 23.8 ms TeleQnA figures serve as the N=2 regression test. Accuracy composes per query at a single gate: an accepted output keeps the efficient model's correctness, an escalated one takes the strong model's. It remains two-tier. Extending it to walk a longer chain is not done here. Record the measured per-cluster cascade accuracy in the AIME and TeleQnA cascade configs, so cre cascade reproduces the reported system accuracy of 88.4% and 74.3% rather than cost alone. --- configs/aime_cascade_test.json | 3 +- configs/teleqna_cascade_test.json | 3 +- src/cre_router/routing.py | 142 ++++++++++++++++++++++++------ tests/test_routing.py | 129 +++++++++++++++++++++++++++ 4 files changed, 248 insertions(+), 29 deletions(-) diff --git a/configs/aime_cascade_test.json b/configs/aime_cascade_test.json index 48989de..7c6c069 100644 --- a/configs/aime_cascade_test.json +++ b/configs/aime_cascade_test.json @@ -1,8 +1,9 @@ { - "_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). `cre cascade` gives 9.75 ms TPOT and 156303 ms E2EL, matching the paper's Table aime_test (9.7 ms) to within rounding.", + "_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}, diff --git a/configs/teleqna_cascade_test.json b/configs/teleqna_cascade_test.json index f0c3756..cd349ea 100644 --- a/configs/teleqna_cascade_test.json +++ b/configs/teleqna_cascade_test.json @@ -1,8 +1,9 @@ { - "_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). `cre cascade` gives 23.65 ms TPOT and 1127 ms E2EL, matching the paper's Table teleqna_test (23.8 ms) to within rounding.", + "_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}, diff --git a/src/cre_router/routing.py b/src/cre_router/routing.py index 6ec4d0d..67fcc1c 100644 --- a/src/cre_router/routing.py +++ b/src/cre_router/routing.py @@ -233,13 +233,69 @@ def system_metrics( 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 the full Stage 1+2 cascade. + """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 @@ -248,40 +304,72 @@ def cascade_system_metrics( route entirely to their Stage 1 model, so ``system_metrics`` is the special case with no escalations. - Both metrics are query-weighted, matching Stage 1 and vLLM's per-request Mean - TPOT. They differ only in how an escalated query is charged: + 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 - - **TPOT** charges both passes per delivered token, - ``TPOT_strong + TPOT_eff * (L_eff / L_strong)``: the discarded efficient - generation is amortised over the strong model's delivered answer. - - **E2EL** is a plain sum, ``E2EL_eff + E2EL_strong``: the user waits for the - efficient generation to complete, since Stage 2 inspects the whole output - before escalating, and then for the strong generation. - Returns ``(tpot_ms, e2el_ms)``. Reproduces the paper's Stage 1+2 latency, - 9.7 ms on AIME and 23.8 ms on TeleQnA. +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()) - tpot_sum = 0.0 - e2el_sum = 0.0 + acc = 0.0 for cluster, name in assignment.items(): - efficient = by_name[name] size = cluster_sizes[cluster] - if cluster in escalations: - strong_name, count = escalations[cluster] - strong = by_name[strong_name] - direct = size - count - ratio = efficient.output_tokens_for(cluster) / strong.output_tokens_for(cluster) - escalated_tpot = strong.tpot_for(cluster) + efficient.tpot_for(cluster) * ratio - tpot_sum += direct * efficient.tpot_for(cluster) + count * escalated_tpot - e2el_sum += direct * efficient.e2el_for(cluster) + count * ( - efficient.e2el_for(cluster) + strong.e2el_for(cluster) - ) + if cluster in cascade_accuracy: + acc += size * cascade_accuracy[cluster] else: - tpot_sum += size * efficient.tpot_for(cluster) - e2el_sum += size * efficient.e2el_for(cluster) - return tpot_sum / total, e2el_sum / total + acc += size * (1.0 - by_name[name].errors[cluster]) + return acc / total def eta( diff --git a/tests/test_routing.py b/tests/test_routing.py index 0b6b98f..c27e95f 100644 --- a/tests/test_routing.py +++ b/tests/test_routing.py @@ -12,8 +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, @@ -170,3 +174,128 @@ def test_no_escalation_matches_stage1(self): _, 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) From de81abad113ade324d7bf27c295f726a519f9d55 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:52:56 +0100 Subject: [PATCH 4/9] Add Stage 2 QE cascade evaluation Add a qe-cascade command that replays a trained quality-estimation classifier over saved generations, decides accept or escalate per query, and reports the resulting per-cluster cascade accuracy and escalation counts. Those feed the Stage 1+2 cost and accuracy functions. Add prep_qe.py to build QE training data from the generations JSONL that cre evaluate now writes, and wire --save-generations through the CLI. Restrict the QE dataset loader to train* and test* files. It previously globbed every JSONL in the directory, which pulled a cascade-test sidecar into the test split. --- data/prep_qe.py | 115 +++++++++++++++++++++++ src/cre_router/cli.py | 104 ++++++++++++++++++-- src/cre_router/qe/cascade.py | 121 ++++++++++++++++++++++++ src/cre_router/qe/classifier.py | 17 +++- src/cre_router/qe/train.py | 19 +++- tests/test_prep_qe.py | 83 ++++++++++++++++ tests/test_qe_cascade.py | 162 ++++++++++++++++++++++++++++++++ tests/test_qe_classifier.py | 47 +++++++++ 8 files changed, 659 insertions(+), 9 deletions(-) create mode 100644 data/prep_qe.py create mode 100644 src/cre_router/qe/cascade.py create mode 100644 tests/test_prep_qe.py create mode 100644 tests/test_qe_cascade.py create mode 100644 tests/test_qe_classifier.py 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/src/cre_router/cli.py b/src/cre_router/cli.py index 650f4b1..b4cfab7 100644 --- a/src/cre_router/cli.py +++ b/src/cre_router/cli.py @@ -21,6 +21,7 @@ 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, @@ -144,6 +145,10 @@ def cmd_cascade(args: argparse.Namespace) -> None: 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}:") @@ -152,9 +157,15 @@ def cmd_cascade(args: argparse.Namespace) -> None: 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) - print(f"\n system TPOT: {tpot:.2f} ms") - print(f" system E2EL: {e2el:.0f} ms") + + 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: @@ -188,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, @@ -199,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) @@ -211,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}") @@ -241,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 @@ -283,6 +351,11 @@ 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") @@ -318,6 +391,25 @@ 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="answerdotai/ModernBERT-base") + 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/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/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/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]") From 8346125cab8ad33d285bbc0880c43f3cefecf5f5 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 16:54:01 +0100 Subject: [PATCH 5/9] Fix swapped Gemma4-E4B cluster errors in the TeleQnA stats The C0 and C1 error rates were transposed relative to the measurement sheet, which records 0.293 on C0 and 0.332 on C1. Routing is unaffected. Gemma4-E4B lies inside the convex hull under either ordering, with margins of 0.039 and 0.052 before and 0.078 and 0.013 after, so it is never the argmin at any lambda. lambda*, the assignment and every reported accuracy are unchanged. (cherry picked from commit 80a7b68f3fb1851bfcc3a13a710de31b0c5c1613) --- configs/teleqna_stats.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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}} } } From f580770b08438eb34ea16d9280e4b933a4f9ce2c Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:08:03 +0100 Subject: [PATCH 6/9] Restate the pinned environment section Give the measurement conditions (2x A100 SXM 80 GB, Python 3.11, 32 concurrent requests, 5-run mean) alongside the pinned versions, and refer to the pinned file rather than repeating version numbers in prose, so the two cannot drift apart. --- REPRODUCE.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/REPRODUCE.md b/REPRODUCE.md index a3ed6e2..23f8ff3 100644 --- a/REPRODUCE.md +++ b/REPRODUCE.md @@ -163,15 +163,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). From cc28f02893697187aaf99ea7f17bcd74a5d66cbe Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:01:06 +0100 Subject: [PATCH 7/9] Default the QE tokenizer to the checkpoint itself Both qe-cascade and the qe evaluate script defaulted --base-tokenizer to a Hub model id, which overrode QEClassifier's own fallback and made either command fail on a machine without internet, even though every checkpoint we train ships its tokenizer alongside the weights. Default to None so the fallback applies; a base model id is now only needed for a checkpoint saved without a tokenizer. --- src/cre_router/cli.py | 4 +++- src/cre_router/qe/evaluate.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/cre_router/cli.py b/src/cre_router/cli.py index b4cfab7..d6a8518 100644 --- a/src/cre_router/cli.py +++ b/src/cre_router/cli.py @@ -404,7 +404,9 @@ def main(argv: list[str] | None = None) -> None: 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="answerdotai/ModernBERT-base") + 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) 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) From dde48b8152a53e44f670ace5799be30367f9fa35 Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:27:41 +0100 Subject: [PATCH 8/9] Document building QE data for a new pool Add the command chain from a run's own generations to a cascade config: capture generations during evaluation, convert them to the schema qe-train reads, then replay the classifier over the gated clusters. --- REPRODUCE.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/REPRODUCE.md b/REPRODUCE.md index 23f8ff3..2b5a380 100644 --- a/REPRODUCE.md +++ b/REPRODUCE.md @@ -135,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: From 4a870db31b3e6145e4857ec54291c2874796555c Mon Sep 17 00:00:00 2001 From: Yasmin Moslem <48152713+ymoslem@users.noreply.github.com> Date: Wed, 5 Aug 2026 22:38:00 +0100 Subject: [PATCH 9/9] Run the test suite on push and pull request Add a GitHub Actions workflow that installs the package and runs pytest on Python 3.10 and 3.12, the ends of the supported range. The suite needs no GPU and no optional extras. --- .github/workflows/tests.yml | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 .github/workflows/tests.yml 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