Fix Routerbench data integrity and RouteFM/BERT reproduction bugs

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

configs/benchmarks/routerbench.json

@@ -8,6 +8,7 @@
    "dataset": {
     "name": "Routerbench",
     "dataset_dir": "./data/routerbench",
+    "protocol": "grouped-clean",
     "split": {
       "mode": "in-domain",
       "ratios": {

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

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

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

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.

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()

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)

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"] = {