From 90cf5bdd8dc09c70dd5925f252adcb5cf07fadc2 Mon Sep 17 00:00:00 2001 From: AIGNLAI Date: Fri, 16 Oct 2026 13:07:26 +0800 Subject: [PATCH] Fix Routerbench input consistency and RouteFM/BERT reproduction bugs --- README.md | 15 ++- configs/benchmarks/routerbench.json | 1 + configs/examples/routefm_qwen.json | 2 +- configs/routers/RouteFM.json | 2 +- docs/multimodal_adapters.md | 7 ++ docs/routerbench_protocol.md | 43 +++++++ methods/RouteFM.py | 57 +++++----- methods/RouteLLM.py | 35 +++--- tests/test_paper_adapters.py | 51 +++++++++ tests/test_routerbench_bugfixes.py | 168 ++++++++++++++++++++++++++++ utils/data.py | 61 ++++++++++ 11 files changed, 395 insertions(+), 47 deletions(-) create mode 100644 docs/routerbench_protocol.md create mode 100644 tests/test_routerbench_bugfixes.py diff --git a/README.md b/README.md index 030f005..482c18a 100644 --- a/README.md +++ b/README.md @@ -208,7 +208,20 @@ RouteFM uses the official immutable frozen-router release. Its default BGE adapt downloaded and checksum-verified by the official package, or can be supplied through `routefm.checkpoint` in `configs/routers/RouteFM.json`. SAVERouter uses the paper's per-benchmark profiles and acquires only the configured fixed-K feedback before fitting. - + +Routerbench now explicitly defaults to the **grouped-clean** data protocol: +conflicting-input query IDs are excluded with an audit log, and in-domain +train/test splits keep whitespace-normalized query groups disjoint. Original +prompts, candidate order, scores and monetary costs are preserved. This changes +the old data protocol; select `dataset.protocol="legacy"` to reproduce it. +See [Routerbench protocol and compatibility](docs/routerbench_protocol.md). + +RouteFM uses the **entire training split as aligned Context** by default +(`routefm.context_size="all"`), then predicts only on the held-out test split. +An integer `context_size` explicitly requests a fixed-K ablation, using the +same sampled query rows for every candidate. `target_batch_size` only batches +test queries; it never truncates Context or supplies held-out outcomes. + --- ## Benchmark Suite diff --git a/configs/benchmarks/routerbench.json b/configs/benchmarks/routerbench.json index 81f81b7..a0904e0 100644 --- a/configs/benchmarks/routerbench.json +++ b/configs/benchmarks/routerbench.json @@ -8,6 +8,7 @@ "dataset": { "name": "Routerbench", "dataset_dir": "./data/routerbench", + "protocol": "grouped-clean", "split": { "mode": "in-domain", "ratios": { diff --git a/configs/examples/routefm_qwen.json b/configs/examples/routefm_qwen.json index 8e4ed41..380899a 100644 --- a/configs/examples/routefm_qwen.json +++ b/configs/examples/routefm_qwen.json @@ -17,7 +17,7 @@ "encoder": "qwen", "checkpoint": null, "cache_dir": null, - "context_size": 8, + "context_size": "all", "target_batch_size": 128 } } diff --git a/configs/routers/RouteFM.json b/configs/routers/RouteFM.json index 8c207b1..4b26f97 100644 --- a/configs/routers/RouteFM.json +++ b/configs/routers/RouteFM.json @@ -14,7 +14,7 @@ "encoder": "bge", "checkpoint": null, "cache_dir": null, - "context_size": 8, + "context_size": "all", "target_batch_size": 128 } } diff --git a/docs/multimodal_adapters.md b/docs/multimodal_adapters.md index 9dafddd..25cfd85 100644 --- a/docs/multimodal_adapters.md +++ b/docs/multimodal_adapters.md @@ -15,6 +15,13 @@ python main.py --dataset VLRouterBench --method RouteFM \ --config configs/examples/routefm_qwen.json --device cuda:0 ``` +Both encoder variants default to `routefm.context_size="all"`: the full +training split is the query-aligned observed Context for every candidate, and +the test split is Target only. Integer sizes are explicit fixed-K ablations, +with one shared training-row sample across candidates. Test-query batching +does not reduce the training Context. No per-task oracle or test labels are +used to select Context. + `--config` applies **top-level replacements** after benchmark/router defaults; it does not recursively merge nested objects. `--device` takes final precedence. Prepare the benchmark's upstream exports as described in diff --git a/docs/routerbench_protocol.md b/docs/routerbench_protocol.md new file mode 100644 index 0000000..72540b4 --- /dev/null +++ b/docs/routerbench_protocol.md @@ -0,0 +1,43 @@ +# Routerbench input-consistent evaluation + +The standard Routerbench configuration now explicitly selects +`dataset.protocol="grouped-clean"`. This corrects two problems in the earlier +loader: taking the first prompt could silently combine candidate outcomes for +different inputs under the same query ID, and row-random splitting could put +identical query text in both train and test. + +The clean protocol: + +1. Checks original prompt equality across candidate records for each query ID. + IDs with conflicting inputs are excluded before pivoting; their IDs are + logged and exposed as `loader.input_conflict_sample_ids`. Exclusion depends + only on input consistency, never quality, response, or cost values. +2. For in-domain splitting, groups queries by whitespace-normalized prompt and + uses `GroupShuffleSplit` with the configured seed and train fraction. The + original prompt strings are retained for encoding. The query fraction is + approximate because whole groups must stay together. + +Candidate order, performance values, original monetary costs, and metric +definitions are unchanged. This is a corrected data protocol, not the original +row-random protocol; identify it explicitly when reporting results. It prevents +cross-split input overlap, not exposure during an external model's pretraining. +Other benchmarks and Routerbench's out-of-domain task split are unchanged. + +## Reproduce historical loading and splitting + +Set `dataset.protocol="legacy"` in the Routerbench benchmark configuration to +retain first-prompt aggregation and the original seeded row-random in-domain +split. Warnings identify conflicting-input merges and the potentially leaky +split. Use this only when reproducing historical results, not as a clean +evaluation. Other dataset options must be retained when replacing the +`dataset` object through a top-level `--config` override. + +## RouteFM context + +The default RouteFM adapter uses **every training row** as observed Context, +with query rows aligned across all candidates. Test queries are Target only; +their observed outcomes are used only by evaluation. An integer +`routefm.context_size` requests an explicit fixed-K ablation using the same +training query rows for each candidate. Missing candidate observations are +masked, not filled with fictitious observations. Test-query batch size changes +memory usage, not the amount of training Context. diff --git a/methods/RouteFM.py b/methods/RouteFM.py index f3465c3..ac01789 100644 --- a/methods/RouteFM.py +++ b/methods/RouteFM.py @@ -14,8 +14,9 @@ class RouteFMRouter(BaseRouter): """Build a RouteFM episode from ORBIT's training split and run it frozen. Candidate identities are deliberately absent from the model input. Each - candidate is represented only by sampled training-query embeddings and its - observed quality/cost, matching RouteFM's public custom-data contract. + candidate is represented by the complete, aligned training-query context + and its observed quality/cost. Fixed-size context is an explicit ablation, + never the default for full-training-set benchmark evaluation. """ def __init__(self, args): @@ -40,10 +41,11 @@ def __init__(self, args): super().__init__(args) self.route_config = route_config self.encoder = encoder - self.context_size = int(route_config.get("context_size", 8)) + requested_context = route_config.get("context_size", "all") + self.context_size = None if requested_context in (None, "all") else int(requested_context) self.target_batch_size = int(route_config.get("target_batch_size", 128)) - if self.context_size < 1 or self.target_batch_size < 1: - raise ValueError("RouteFM context_size and target_batch_size must be positive.") + if (self.context_size is not None and self.context_size < 1) or self.target_batch_size < 1: + raise ValueError("RouteFM context_size must be 'all' or positive; target_batch_size must be positive.") def _load_official_router(self): try: @@ -66,29 +68,25 @@ def _load_official_router(self): def _build_context(self, embeddings, performance, costs): n_queries, dimension = embeddings.shape n_models = performance.shape[1] - width = min(self.context_size, n_queries) - context_query = np.zeros((n_models, width, dimension), dtype=np.float32) - context_score = np.zeros((n_models, width), dtype=np.float32) - context_cost = np.zeros((n_models, width), dtype=np.float32) - context_mask = np.zeros((n_models, width), dtype=bool) - rng = np.random.default_rng(self.seed) - - for model_index in range(n_models): - valid = np.flatnonzero( - np.isfinite(performance[:, model_index]) - & np.isfinite(costs[:, model_index]) - & (costs[:, model_index] >= 0) - ) - if valid.size == 0: - raise ValueError( - f"RouteFM candidate {model_index} has no finite context observation." - ) - count = min(width, valid.size) - chosen = rng.choice(valid, size=count, replace=False) - context_query[model_index, :count] = embeddings[chosen] - context_score[model_index, :count] = performance[chosen, model_index] - context_cost[model_index, :count] = costs[chosen, model_index] - context_mask[model_index, :count] = True + if n_queries < 1 or n_models < 1: + raise ValueError("RouteFM requires nonempty training context and candidates.") + width = n_queries if self.context_size is None else min(self.context_size, n_queries) + # One shared row index: candidate outcomes remain query-aligned. The + # full-context path preserves every training row in its original order. + chosen = np.arange(n_queries) + if width < n_queries: + chosen = np.sort(np.random.default_rng(self.seed).choice(n_queries, width, replace=False)) + self.context_row_indices = chosen + context_score = np.asarray(performance[chosen].T, dtype=np.float32) + context_cost = np.asarray(costs[chosen].T, dtype=np.float32) + context_mask = np.isfinite(context_score) & np.isfinite(context_cost) & (context_cost >= 0) + if not context_mask.any(axis=1).all(): + missing = np.flatnonzero(~context_mask.any(axis=1)).tolist() + raise ValueError(f"RouteFM candidates {missing} have no finite context observation.") + context_query = np.broadcast_to(embeddings[chosen][None], (n_models, width, dimension)).copy() + context_query[~context_mask] = 0.0 + context_score = np.where(context_mask, context_score, 0.0) + context_cost = np.where(context_mask, context_cost, 0.0) valid_scores = context_score[context_mask] if np.any((valid_scores < 0) | (valid_scores > 1)): @@ -121,10 +119,11 @@ def train(self): f"but ORBIT produced {embeddings.shape[1]} dimensions." ) logging.info( - "[RouteFM] Loaded frozen %s checkpoint with %d candidates and K=%d context slots.", + "[RouteFM] Loaded frozen %s checkpoint with %d candidates and K=%d aligned context slots (%s).", self.encoder, len(self.model_list), self.context_query.shape[1], + "entire training split" if self.context_size is None else "explicit fixed-K ablation", ) @torch.inference_mode() diff --git a/methods/RouteLLM.py b/methods/RouteLLM.py index 700e7fc..3afeb23 100644 --- a/methods/RouteLLM.py +++ b/methods/RouteLLM.py @@ -486,11 +486,14 @@ def __init__(self, args: Dict[str, Any]) -> None: ) train_cfg = args.get("training", {}) - self.lr: float = float(train_cfg.get("lr", 2e-5)) - self.epochs: int = int(train_cfg.get("epochs", 2)) - self.batch_size: int = int(train_cfg.get("batch_size", 16)) + self.lr: float = float(train_cfg.get("lr", args.get("lr", 2e-5))) + self.epochs: int = int(train_cfg.get("epochs", args.get("epochs", 2))) + self.batch_size: int = int(train_cfg.get("batch_size", args.get("batch_size", 16))) self.weight_decay: float = float(train_cfg.get("weight_decay", 0.01)) self.max_length: int = int(args.get("max_length", 256)) + self.prediction_batch_size = int(args.get("prediction_batch_size", self.batch_size)) + if self.prediction_batch_size < 1: + raise ValueError("RouteLLM_BERT prediction_batch_size must be positive.") # Lazy import so that environments without transformers fail loudly only when used. try: @@ -577,18 +580,20 @@ def predict(self, test_texts: Any) -> Tuple[np.ndarray, np.ndarray]: if not isinstance(test_texts, (list, tuple)) or (len(test_texts) > 0 and not isinstance(test_texts[0], str)): raise ValueError("RouteLLM_BERT.predict expects List[str] test texts.") - enc = self.tokenizer( - list(test_texts), - padding=True, - truncation=True, - max_length=self.max_length, - return_tensors="pt", - ).to(self.device) - - out = self.encoder(**enc) - cls = out.last_hidden_state[:, 0, :] - logits = self.classifier(cls).squeeze(1) - p_strong = torch.sigmoid(logits).detach().cpu().numpy().astype(np.float32) + probabilities = [] + for start in range(0, len(test_texts), self.prediction_batch_size): + enc = self.tokenizer( + list(test_texts[start:start + self.prediction_batch_size]), + padding=True, + truncation=True, + max_length=self.max_length, + return_tensors="pt", + ).to(self.device) + out = self.encoder(**enc) + cls = out.last_hidden_state[:, 0, :] + logits = self.classifier(cls).squeeze(1) + probabilities.append(torch.sigmoid(logits).detach().cpu().numpy().astype(np.float32)) + p_strong = np.concatenate(probabilities) if probabilities else np.empty(0, dtype=np.float32) B = len(test_texts) K = len(self.model_list) diff --git a/tests/test_paper_adapters.py b/tests/test_paper_adapters.py index e27d10c..f3a3dde 100644 --- a/tests/test_paper_adapters.py +++ b/tests/test_paper_adapters.py @@ -270,6 +270,57 @@ def test_routefm_official_episode_adapter_smoke(self): self.assertTrue(np.isfinite(cost).all()) self.assertTrue((cost >= 0).all()) + def test_routefm_default_context_uses_every_training_row_aligned(self): + args = self._args("RouteFM") + args["routefm"] = {"encoder": "bge"} + dataset_patch, embedder_patch = self._patches() + with dataset_patch, embedder_patch: + router = RouteFMRouter(args) + self.assertIsNone(router.context_size) + embedding, quality, cost = router._prepare_training_data() + context, features, mask = router._build_context(embedding, quality, cost) + np.testing.assert_array_equal(router.context_row_indices, np.arange(len(self.train_df))) + self.assertEqual(context.shape, (3, len(self.train_df), 4)) + for model in range(3): + np.testing.assert_array_equal(context[model], embedding) + np.testing.assert_array_equal(features[model, :, 0], quality[:, model]) + self.assertTrue(mask.all()) + # Held-out inputs are not among the source rows used above. + self.assertTrue(set(self.train_df.prompt).isdisjoint(self.test_df.prompt)) + + def test_routefm_fixed_k_is_query_aligned_and_deterministic(self): + router = RouteFMRouter.__new__(RouteFMRouter) + router.context_size, router.seed = 3, 42 + embeddings = np.arange(40, dtype=np.float32).reshape(10, 4) + quality = np.arange(30, dtype=np.float32).reshape(10, 3) / 30 + cost = np.ones((10, 3), dtype=np.float32) + first, features, mask = router._build_context(embeddings, quality, cost) + indices = router.context_row_indices.copy() + second, _, _ = router._build_context(embeddings, quality, cost) + np.testing.assert_array_equal(first, second) + for model in range(3): + np.testing.assert_array_equal(first[model], embeddings[indices]) + np.testing.assert_array_equal(features[model, :, 0], quality[indices, model]) + self.assertTrue(mask.all()) + + def test_routefm_full_context_masks_missing_cells_without_dropping_rows(self): + router = RouteFMRouter.__new__(RouteFMRouter) + router.context_size, router.seed = None, 42 + embeddings = np.ones((4, 2), dtype=np.float32) + quality = np.array([[1, np.nan], [0, 1], [1, 0], [0, 1]], dtype=np.float32) + cost = np.ones((4, 2), dtype=np.float32) + cost[2, 0] = -1 + context, features, mask = router._build_context(embeddings, quality, cost) + self.assertEqual(mask.shape, (2, 4)) + self.assertFalse(mask[1, 0]) + self.assertFalse(mask[0, 2]) + self.assertTrue(np.isfinite(features).all()) + self.assertTrue((context[~mask] == 0).all()) + self.assertEqual(mask.sum(), 6) + quality[:, 0] = np.nan + with self.assertRaisesRegex(ValueError, "no finite context"): + router._build_context(embeddings, quality, cost) + def test_saverouter_sparse_supervision_adapter_smoke(self): args = self._args("SaveRouter") args["saverouter"] = { diff --git a/tests/test_routerbench_bugfixes.py b/tests/test_routerbench_bugfixes.py new file mode 100644 index 0000000..9921b5c --- /dev/null +++ b/tests/test_routerbench_bugfixes.py @@ -0,0 +1,168 @@ +"""Regressions found during full Routerbench reproduction; no experiment helpers.""" + +import json +from pathlib import Path +import pickle +import tempfile +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +import numpy as np +import pandas as pd +from sklearn.model_selection import GroupShuffleSplit +import torch + +from main import load_config +from methods.base import BaseRouter +from methods.RouteLLM import RouteLLM_BERT +from utils.data import RouterbenchLoader + + +class BertRuntimeRegressionTest(unittest.TestCase): + def _construct(self, options): + def initialize(router, args): + router.device = torch.device("cpu") + encoder = torch.nn.Linear(2, 2) + encoder.config = SimpleNamespace(hidden_size=2) + with patch.object(BaseRouter, "__init__", initialize), patch( + "transformers.AutoTokenizer.from_pretrained" + ), patch("transformers.AutoModel.from_pretrained", return_value=encoder): + return RouteLLM_BERT(options) + + def test_top_level_training_settings_are_not_silently_ignored(self): + args = json.loads(Path("configs/routers/RouteLLM_BERT.json").read_text()) + router = self._construct(args) + self.assertEqual(router.epochs, args["epochs"]) + self.assertEqual(router.batch_size, args["batch_size"]) + self.assertEqual(router.lr, args["lr"]) + self.assertEqual(router.prediction_batch_size, args["batch_size"]) + + def test_nested_training_settings_take_precedence_with_per_key_fallback(self): + router = self._construct({"epochs": 10, "batch_size": 128, "lr": 2e-5, + "training": {"epochs": 3, "lr": 1e-5}, + "prediction_batch_size": 4}) + self.assertEqual((router.epochs, router.batch_size, router.lr), (3, 128, 1e-5)) + self.assertEqual(router.prediction_batch_size, 4) + defaults = self._construct({}) + self.assertEqual((defaults.epochs, defaults.batch_size, defaults.lr), (2, 16, 2e-5)) + + def test_nonpositive_prediction_batch_is_rejected(self): + with self.assertRaisesRegex(ValueError, "prediction_batch_size"): + self._construct({"prediction_batch_size": 0}) + + def test_prediction_is_bounded_and_independent_of_batch_partition(self): + router = RouteLLM_BERT.__new__(RouteLLM_BERT) + router.device = torch.device("cpu") + router.max_length = 256 + router.model_list = ["weak", "strong", "disabled"] + router.cfg = SimpleNamespace(strong_model_idx=1, weak_model_idx=0) + router.prediction_batch_size = 2 + batch_sizes = [] + class Tokens(dict): + def to(self, device): + return self + def tokenize(texts, **kwargs): + batch_sizes.append(len(texts)) + return Tokens(input_ids=torch.tensor([[len(text), 1.] for text in texts])) + router.tokenizer = tokenize + router.encoder = lambda input_ids: SimpleNamespace(last_hidden_state=input_ids[:, None, :]) + router.classifier = torch.nn.Linear(2, 1) + router.embedder = SimpleNamespace(run_embed=lambda texts, images=None: torch.zeros((len(texts), 2))) + router._predict_shared_cost = lambda values: np.ones((len(values), 3), dtype=np.float32) + texts = ["x", "xx", "xxx", "xxxx", "xxxxx"] + batched_quality, batched_cost = router.predict(texts) + self.assertEqual(batch_sizes, [2, 2, 1]) + router.prediction_batch_size = 10 + whole_quality, whole_cost = router.predict(texts) + np.testing.assert_allclose(batched_quality[:, :2], whole_quality[:, :2], atol=1e-6) + np.testing.assert_array_equal(batched_cost, whole_cost) + self.assertTrue(np.isneginf(batched_quality[:, 2]).all()) + + +class RouterbenchProtocolRegressionTest(unittest.TestCase): + MODELS = ["WizardLM/WizardLM-13B-V1.2", "claude-instant-v1", "claude-v1", "claude-v2", + "gpt-3.5-turbo-1106", "gpt-4-1106-preview", "meta/code-llama-instruct-34b-chat", + "meta/llama-2-70b-chat", "mistralai/mistral-7b-chat", "mistralai/mixtral-8x7b-chat", + "zero-one-ai/Yi-34B-Chat"] + + def setUp(self): + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name) + self.args = {"seed": 42, "dataset": {"protocol": "grouped-clean", + "split": {"mode": "in-domain", "ratios": {"train": .2, "test": .8}}}} + self.raw = pd.DataFrame([ + {"sample_id": f"q{i:02d}", "model_name": model, "prompt": f" original query {i} \n", + "eval_name": f"task-{i % 2}", "performance": ((i + j) % 5) / 4, + "cost": .0001 * (i + 1) * (j + 1), "model_response": f"response {i} {j}"} + for i in range(20) for j, model in enumerate(self.MODELS) + ]) + + def _load(self): + path = self.root / "routerbench_raw.pkl" + with path.open("wb") as stream: + pickle.dump(self.raw, stream) + loader = RouterbenchLoader(str(self.root), self.args) + models, frame = loader.process(str(path)) + return loader, models, frame + + def test_input_conflict_is_excluded_without_score_based_filtering(self): + self.raw.loc[(self.raw.sample_id == "q05") & (self.raw.model_name == self.MODELS[-1]), "prompt"] = "different query" + loader, models, frame = self._load() + self.assertEqual(loader.input_conflict_sample_ids, ["q05"]) + self.assertEqual(models, self.MODELS) + self.assertEqual(len(frame), 19) + self.assertNotIn("q05", set(frame.sample_id)) + for j, model in enumerate(models): + expected = self.raw[(self.raw.model_name == model) & (self.raw.sample_id != "q05")].set_index("sample_id") + actual = frame.set_index("sample_id").loc[expected.index] + np.testing.assert_array_equal(actual[f"model_{j}_performance"], expected.performance) + np.testing.assert_array_equal(actual[f"model_{j}_cost"], expected.cost) + + def test_grouped_split_is_disjoint_deterministic_and_retains_original_prompts(self): + self.raw.loc[self.raw.sample_id == "q01", "prompt"] = " original query 0\t" + loader, _, frame = self._load() + train, test = loader.split(frame) + normalize = lambda values: values.str.replace(r"\s+", " ", regex=True).str.strip() + self.assertFalse(set(normalize(train.prompt)) & set(normalize(test.prompt))) + self.assertEqual(len(train) + len(test), len(frame)) + expected = GroupShuffleSplit(n_splits=1, train_size=.2, random_state=42) + train_idx, test_idx = next(expected.split(frame, groups=normalize(frame.prompt))) + pd.testing.assert_frame_equal(train, frame.iloc[train_idx].reset_index(drop=True)) + pd.testing.assert_frame_equal(test, frame.iloc[test_idx].reset_index(drop=True)) + original = frame.set_index("sample_id").prompt + combined = pd.concat([train, test]).set_index("sample_id").prompt + pd.testing.assert_series_equal(combined.loc[original.index], original) + + def test_legacy_reproduces_old_merge_and_row_random_split_with_warning(self): + self.args["dataset"]["protocol"] = "legacy" + self.raw.loc[(self.raw.sample_id == "q05") & (self.raw.model_name == self.MODELS[-1]), "prompt"] = "different query" + with self.assertLogs(level="WARNING") as recorded: + loader, _, frame = self._load() + train, test = loader.split(frame) + self.assertEqual(len(frame), 20) + self.assertEqual(loader.input_conflict_sample_ids, ["q05"]) + self.assertIn("legacy", " ".join(recorded.output)) + shuffled = frame.sample(frac=1, random_state=42).reset_index(drop=True) + pd.testing.assert_frame_equal(train, shuffled.iloc[:4].reset_index(drop=True)) + pd.testing.assert_frame_equal(test, shuffled.iloc[4:].reset_index(drop=True)) + + def test_missing_empty_input_and_unknown_protocol_fail_loudly(self): + for prompt in (None, " \t\n"): + with self.subTest(prompt=prompt): + self.raw.loc[self.raw.sample_id == "q00", "prompt"] = prompt + with self.assertRaisesRegex(ValueError, "prompt"): + self._load() + self.args["dataset"]["protocol"] = "typo" + with self.assertRaisesRegex(ValueError, "protocol"): + self._load() + + def test_default_configs_select_clean_protocol_and_full_context(self): + config = load_config("Routerbench", "RouteFM") + self.assertEqual(config["dataset"]["protocol"], "grouped-clean") + self.assertEqual(config["routefm"]["context_size"], "all") + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/data.py b/utils/data.py index a7005b1..55ca6d1 100644 --- a/utils/data.py +++ b/utils/data.py @@ -210,6 +210,48 @@ class RouterbenchLoader(BaseDatasetLoader): hf_id = "withmartian/routerbench" default_target_path = "./data/routerbench" + def _protocol(self) -> str: + protocol = self.args.get("dataset", {}).get("protocol", "grouped-clean") + if protocol not in {"grouped-clean", "legacy"}: + raise ValueError("Routerbench dataset.protocol must be 'grouped-clean' or 'legacy'.") + return protocol + + @staticmethod + def _input_groups(frame: pd.DataFrame) -> pd.Series: + # Normalization is only a grouping key. Original encoder input stays intact. + if "prompt" not in frame or frame["prompt"].isna().any(): + raise ValueError("Routerbench grouped input handling requires non-missing prompts.") + groups = frame["prompt"].astype(str).str.replace(r"\s+", " ", regex=True).str.strip() + if (groups == "").any(): + raise ValueError("Routerbench grouped input handling requires nonempty prompts.") + return groups + + def split_in(self, dataset: Any) -> Tuple[pd.DataFrame, pd.DataFrame]: + protocol = self._protocol() + if protocol == "legacy": + logging.warning("[Routerbench] Using legacy row-random split; duplicate inputs may cross train/test.") + return super().split_in(dataset) + from sklearn.model_selection import GroupShuffleSplit + + frame = dataset.copy() if isinstance(dataset, pd.DataFrame) else pd.DataFrame(dataset) + groups = self._input_groups(frame) + splitter = GroupShuffleSplit( + n_splits=1, + train_size=self.args["dataset"]["split"]["ratios"]["train"], + random_state=self.args["seed"], + ) + train_ids, test_ids = next(splitter.split(frame, groups=groups)) + train = frame.iloc[train_ids].reset_index(drop=True) + test = frame.iloc[test_ids].reset_index(drop=True) + shared = set(self._input_groups(train)) & set(self._input_groups(test)) + if shared: + raise RuntimeError("Routerbench grouped-clean split leaked normalized input groups.") + logging.info( + "[Routerbench] protocol=grouped-clean seed=%s train=%d test=%d cross_split_input_groups=%d", + self.args["seed"], len(train), len(test), len(shared), + ) + return train, test + def download(self) -> str: """Download routerbench_raw.pkl via snapshot_download and return its local path.""" from huggingface_hub import snapshot_download @@ -307,6 +349,25 @@ def resolve_column(name_candidates): df = df[selected_cols].copy() + # A query ID must identify one actual input across candidates. Never + # silently bind outcomes from different prompts via groupby.first(). + protocol = self._protocol() + if not prompt_col: + raise ValueError("Routerbench requires a prompt column to validate query identity.") + prompt_counts = df.groupby(sample_id_col)[prompt_col].nunique(dropna=False) + conflicting_ids = prompt_counts[prompt_counts > 1].index.tolist() + self.input_conflict_sample_ids = [str(value) for value in conflicting_ids] + if conflicting_ids: + logging.warning( + "[Routerbench] protocol=%s: %s %d query IDs with conflicting candidate inputs: %s", + protocol, "excluding" if protocol == "grouped-clean" else "legacy merge retains", + len(conflicting_ids), self.input_conflict_sample_ids, + ) + if protocol == "grouped-clean": + df = df[~df[sample_id_col].isin(conflicting_ids)].copy() + if protocol == "grouped-clean": + self._input_groups(df.rename(columns={prompt_col: "prompt"})) + perf_table = df.pivot_table(index=sample_id_col, columns=model_name_col, values=performance_col, aggfunc="first") cost_table = df.pivot_table(index=sample_id_col, columns=model_name_col, values=cost_col, aggfunc="first") resp_table = df.pivot_table(index=sample_id_col, columns=model_name_col, values=response_col, aggfunc="first")