From 0a7573653a8907ee88942184783fd9f824fe6a96 Mon Sep 17 00:00:00 2001 From: AIGNLAI Date: Mon, 5 Oct 2026 10:53:01 +0800 Subject: [PATCH] Add RouteFM and SAVERouter adapters --- README.md | 22 +++- configs/routers/RouteFM.json | 20 ++++ configs/routers/SaveRouter.json | 13 ++ methods/RouteFM.py | 170 +++++++++++++++++++++++++++ methods/SaveRouter.py | 148 +++++++++++++++++++++++ methods/__init__.py | 2 + requirement.txt | 2 + requirements-ci.txt | 2 + tests/test_paper_adapters.py | 95 +++++++++++++++ tests/test_router_implementations.py | 4 +- train.py | 4 +- utils/embedding.py | 38 ++++++ 12 files changed, 513 insertions(+), 7 deletions(-) create mode 100644 configs/routers/RouteFM.json create mode 100644 configs/routers/SaveRouter.json create mode 100644 methods/RouteFM.py create mode 100644 methods/SaveRouter.py diff --git a/README.md b/README.md index 844837e..80309a1 100644 --- a/README.md +++ b/README.md @@ -56,9 +56,10 @@ Useful links: --- -## What's New - -- **2026-08**: ORBIT has been accepted by the *Frontiers of Computer Science* (FCS) special column **[Code & Data](https://journal.hep.com.cn/fcs/EN/subject/showCollection.do?subjectId=1710741206314)**. [[Paper]](https://doi.org/10.1007/s11704-026-61310-5) +## What's New + +- **2026-10**: Added **RouteFM**, the frozen in-context foundation router that transfers to new candidate pools through behavioral context, and **SAVERouter**, which adaptively acquires fixed-K sparse feedback and accounts for supervision expenditure. [[RouteFM Paper]](https://arxiv.org/abs/2609.37362) [[RouteFM Code]](https://github.com/LAMDA-Model-Reuse/RouteFM) [[SAVERouter Paper]](https://arxiv.org/abs/2609.37402) [[SAVERouter Code]](https://github.com/LAMDA-Model-Reuse/SaveRouter) +- **2026-08**: ORBIT has been accepted by the *Frontiers of Computer Science* (FCS) special column **[Code & Data](https://journal.hep.com.cn/fcs/EN/subject/showCollection.do?subjectId=1710741206314)**. [[Paper]](https://doi.org/10.1007/s11704-026-61310-5) - **2026-08**: Added InferenceDynamics. - **2026-08**: Added the LLMRouterBench Performance-Cost benchmark with 12 models across 10 tasks, using ORBIT's standardized split and globally scaled per-query costs. - **2026-06**: Initial release of ORBIT v1.0 with a unified routing pipeline, standardized budget-aware evaluation, unimodal and multimodal benchmarks, and reproduced routing methods. @@ -67,7 +68,7 @@ Useful links: ## Methods Reproduced -ORBIT reproduces **28 representative LLM routing methods** across training-free, retrieval-based, and learned routers under a unified pipeline and standardized budgeted evaluation. +ORBIT reproduces **30 representative LLM routing methods** across training-free, retrieval-based, and learned routers under a unified pipeline and standardized budgeted evaluation. - **Avengers**: A training-free recipe that clusters queries and routes by cluster-wise capability profiles with sampling/voting. [[Paper]](https://arxiv.org/abs/2505.19797) - **Avengers-Pro**: A test-time routing framework that traces a Pareto frontier via clustering and a tunable performance-efficiency objective. [[Paper]](https://arxiv.org/abs/2508.12631) @@ -97,6 +98,8 @@ ORBIT reproduces **28 representative LLM routing methods** across training-free, - **ProfileRouter**: Implements RouteProfile's training-free Emb-GNN profile construction and cosine-similarity SimRouter using public model descriptions, family metadata, and optional reported benchmark scores. [[Paper]](https://arxiv.org/abs/2605.00180) - **CarrotRouter**: Uses CARROT's KNN plug-in estimators for query-level performance and cost, then applies its original cost-aware rate-optimal scalarization. [[Paper]](https://arxiv.org/abs/2502.03261) - **EARAMRouter**: Trains independent provider-side success predictors and allocates queries through EA-RAM's largest-positive-surplus reverse auction with runner-up externality payments. [[Paper]](https://arxiv.org/abs/2608.12719) +- **RouteFM**: Uses the official frozen foundation router and anonymous per-candidate behavioral context to predict target-query quality and relative cost without target-domain parameter updates. ORBIT defaults to the released BGE text checkpoint. [[Paper]](https://arxiv.org/abs/2609.37362) [[Code]](https://github.com/LAMDA-Model-Reuse/RouteFM) +- **SAVERouter**: Adaptively reveals exactly K model outcomes per training query, then combines hierarchical group-model shrinkage with query-level residual refinement and native sparse cost estimates. [[Paper]](https://arxiv.org/abs/2609.37402) [[Code]](https://github.com/LAMDA-Model-Reuse/SaveRouter) --- @@ -119,7 +122,8 @@ ORBIT evaluates LLM routing under **budgeted inference** by sweeping budgets to Every score-based learned router supplies a query-by-model cost matrix without using test labels. Methods with a native cost mechanism reuse it: for example, cluster statistics for Avengers, neighbor retrieval for Eagle, cluster features for UniRoute, and profile indexes for -InferenceDynamics. Methods whose native score cannot represent cost (HybridLLM, RouterDC, +InferenceDynamics. RouteFM uses its released relative-cost head, while SAVERouter estimates +costs from its sparse group/model observations. Methods whose native score cannot represent cost (HybridLLM, RouterDC, RouteLLM, ModelSAT, ProfileRouter, and EA-RAM's offline execution-cost bids) use the same multi-output MLP trained only on training-query embeddings and costs. Its optional settings live under `cost_prediction` (`hidden_sizes`, `dropout`, @@ -181,10 +185,18 @@ python main.py --dataset Routerbench --method AvengersPro python main.py --dataset MMRBench --method EquiRouter python main.py --dataset Mixinstruct --method oracle python main.py --dataset LLMRouterBench --method knn +python main.py --dataset Routerbench --method RouteFM +python main.py --dataset Routerbench --method SaveRouter ``` Dataset and method names are matched case-insensitively. Device selection defaults to `auto`; use `--device cpu`, `--device cuda`, or `--device cuda:N` to override it. + +RouteFM uses the official immutable frozen-router release. Its default BGE adapter expects +`BAAI/bge-base-en-v1.5` at `./bge-base-en-v1.5`; the RouteFM router checkpoint itself is +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. --- diff --git a/configs/routers/RouteFM.json b/configs/routers/RouteFM.json new file mode 100644 index 0000000..8c207b1 --- /dev/null +++ b/configs/routers/RouteFM.json @@ -0,0 +1,20 @@ +{ + "method": "RouteFM", + "modality": "text", + "embeddings": { + "text_model": "BAAI/bge-base-en-v1.5", + "image_model": null, + "text_model_dir": "./bge-base-en-v1.5", + "batch_size": 32, + "normalize": true, + "training": false, + "out_dim": 768 + }, + "routefm": { + "encoder": "bge", + "checkpoint": null, + "cache_dir": null, + "context_size": 8, + "target_batch_size": 128 + } +} diff --git a/configs/routers/SaveRouter.json b/configs/routers/SaveRouter.json new file mode 100644 index 0000000..e16a26e --- /dev/null +++ b/configs/routers/SaveRouter.json @@ -0,0 +1,13 @@ +{ + "method": "SaveRouter", + "saverouter": { + "k": 4, + "acquisition_beta": 0.35, + "acquisition_prior_strength": 10.0, + "prior_alpha": 1.0, + "prior_beta": 1.0, + "seed_stride": 1009, + "prior_ridge_alpha": 10.0, + "min_model_observations": 8 + } +} diff --git a/methods/RouteFM.py b/methods/RouteFM.py new file mode 100644 index 0000000..b69f239 --- /dev/null +++ b/methods/RouteFM.py @@ -0,0 +1,170 @@ +"""ORBIT adapter for the released RouteFM frozen in-context router.""" + +from __future__ import annotations + +import logging + +import numpy as np +import torch + +from methods.base import BaseRouter + + +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. + """ + + def __init__(self, args): + route_config = dict(args.get("routefm", {})) + encoder = str(route_config.get("encoder", "bge")).lower() + if encoder != "bge": + raise ValueError( + "ORBIT's built-in RouteFM adapter currently supports the official " + "BGE text checkpoint. The Qwen checkpoint requires joint Qwen3-VL " + "embeddings from a compatible service." + ) + # The released BGE checkpoint is text-only. In a multimodal benchmark, + # this is an explicit text-only evaluation rather than concatenating an + # incompatible image representation. + args = dict(args) + args["modality"] = "text" + super().__init__(args) + self.route_config = route_config + self.encoder = encoder + self.context_size = int(route_config.get("context_size", 8)) + 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.") + + def _load_official_router(self): + try: + from routefm.checkpoints import resolve_checkpoint + from routefm.predict import load_router + except ImportError as exc: + raise RuntimeError( + "RouteFM requires the official routefm-router package. " + "Install ORBIT's requirement.txt or install RouteFM from its " + "official repository." + ) from exc + + checkpoint = resolve_checkpoint( + self.encoder, + self.route_config.get("checkpoint"), + cache_dir=self.route_config.get("cache_dir"), + ) + return load_router(checkpoint, str(self.device), self.encoder) + + 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 + + valid_scores = context_score[context_mask] + if np.any((valid_scores < 0) | (valid_scores > 1)): + raise ValueError("RouteFM context performance must be in [0, 1].") + log_cost = np.log1p(context_cost[context_mask]) + self._log_cost_low = float(log_cost.min()) + self._log_cost_high = float(log_cost.max()) + span = max(self._log_cost_high - self._log_cost_low, 1e-8) + normalized_cost = np.zeros_like(context_cost) + normalized_cost[context_mask] = ( + np.log1p(context_cost[context_mask]) - self._log_cost_low + ) / span + features = np.stack((context_score, normalized_cost), axis=-1) + return context_query, features.astype(np.float32), context_mask + + def train(self): + embeddings, performance, costs = self._prepare_training_data() + if not np.isfinite(embeddings).all(): + raise ValueError("RouteFM training embeddings must be finite.") + ( + self.context_query, + self.context_features, + self.context_mask, + ) = self._build_context(embeddings, performance, costs) + self.model = self._load_official_router() + query_dim = int(self.model.config.query_dim) + if embeddings.shape[1] != query_dim: + raise ValueError( + f"RouteFM {self.encoder} checkpoint expects {query_dim}-D embeddings, " + f"but ORBIT produced {embeddings.shape[1]} dimensions." + ) + logging.info( + "[RouteFM] Loaded frozen %s checkpoint with %d candidates and K=%d context slots.", + self.encoder, + len(self.model_list), + self.context_query.shape[1], + ) + + @torch.inference_mode() + def predict(self, test_embedding): + if self.model is None: + raise RuntimeError("RouteFM must be trained (contextualized) before predict().") + if isinstance(test_embedding, torch.Tensor): + targets = test_embedding.detach().cpu().numpy().astype(np.float32) + else: + targets = np.asarray(test_embedding, dtype=np.float32) + if targets.ndim != 2 or targets.shape[1] != int(self.model.config.query_dim): + raise ValueError( + f"RouteFM target embeddings must have shape (N, {self.model.config.query_dim})." + ) + if not np.isfinite(targets).all(): + raise ValueError("RouteFM target embeddings must be finite.") + + device = self.device + models = len(self.model_list) + shared = { + "context_query": torch.from_numpy(self.context_query[None]).to(device), + "context_features": torch.from_numpy(self.context_features[None]).to(device), + "context_mask": torch.from_numpy(self.context_mask[None]).to(device), + "candidate_mask": torch.ones((1, models), dtype=torch.bool, device=device), + } + score_parts, relative_cost_parts = [], [] + for offset in range(0, len(targets), self.target_batch_size): + local_targets = torch.from_numpy( + targets[offset : offset + self.target_batch_size][None] + ).to(device) + batch = dict(shared) + batch["target_query"] = local_targets + batch["target_mask"] = torch.ones( + (1, local_targets.shape[1], models), dtype=torch.bool, device=device + ) + output = self.model(batch) + score_parts.append(output["score_mean"][0].float().cpu()) + relative_cost_parts.append(output["cost_mean"][0].float().cpu()) + + scores = torch.cat(score_parts).numpy() + relative_costs = torch.cat(relative_cost_parts).numpy() + span = self._log_cost_high - self._log_cost_low + predicted_costs = np.expm1( + self._log_cost_low + np.clip(relative_costs, 0.0, 1.0) * span + ).astype(np.float32) + if not np.isfinite(scores).all() or not np.isfinite(predicted_costs).all(): + raise ValueError("RouteFM returned non-finite predictions.") + return scores.astype(np.float32), predicted_costs diff --git a/methods/SaveRouter.py b/methods/SaveRouter.py new file mode 100644 index 0000000..a42a547 --- /dev/null +++ b/methods/SaveRouter.py @@ -0,0 +1,148 @@ +"""ORBIT adapter for the official SAVERouter sparse-supervision pipeline.""" + +from __future__ import annotations + +import logging + +import numpy as np +import torch + +from methods.base import BaseRouter + + +_PAPER_DEFAULTS = { + "llmrouterbench": {"group_strategy": "auto", "n_groups": None, "prior_strength": 40.0, + "residual_ridge_alpha": 200.0, "residual_gamma": 2.0, + "include_dense_context": False}, + "routerbench": {"group_strategy": "auto", "n_groups": None, "prior_strength": 40.0, + "residual_ridge_alpha": 200.0, "residual_gamma": 2.0, + "include_dense_context": False}, + "mixinstruct": {"group_strategy": "latent", "n_groups": 1, "prior_strength": 40.0, + "residual_ridge_alpha": 100.0, "residual_gamma": 2.0, + "include_dense_context": False}, + "mmrbench": {"group_strategy": "auto", "n_groups": None, "prior_strength": 120.0, + "residual_ridge_alpha": 100.0, "residual_gamma": 0.25, + "include_dense_context": True}, +} + + +class SaveRouter(BaseRouter): + """Acquire fixed-K feedback, then fit SAVERouter's published estimator.""" + + def __init__(self, args): + super().__init__(args) + dataset_key = str(self.args["dataset"]["name"]).lower().replace("-", "") + config = dict(_PAPER_DEFAULTS.get(dataset_key, _PAPER_DEFAULTS["routerbench"])) + config.update(dict(self.args.get("saverouter", {}))) + config.setdefault("k", 4) + config.setdefault("acquisition_beta", 0.35) + config.setdefault("acquisition_prior_strength", 10.0) + config.setdefault("prior_alpha", 1.0) + config.setdefault("prior_beta", 1.0) + config.setdefault("seed_stride", 1009) + config.setdefault("prior_ridge_alpha", 10.0) + config.setdefault("min_model_observations", 8) + self.save_config = config + + @staticmethod + def _official_components(): + try: + from saverouter import ( + FixedKSparseRouter, + GroupAssigner, + TextFeatureEncoder, + simulate_fixed_k_supervision, + ) + except ImportError as exc: + raise RuntimeError( + "SaveRouter requires the official saverouter package. " + "Install ORBIT's requirement.txt or install SaveRouter from its " + "official repository." + ) from exc + return FixedKSparseRouter, GroupAssigner, TextFeatureEncoder, simulate_fixed_k_supervision + + @staticmethod + def _numpy(values): + if isinstance(values, torch.Tensor): + return values.detach().cpu().numpy().astype(np.float32) + return np.asarray(values, dtype=np.float32) + + def _group_features(self, embeddings, *, training): + dataset_key = str(self.args["dataset"]["name"]).lower().replace("-", "") + if dataset_key == "mixinstruct": + return np.zeros((len(embeddings), 1), dtype=np.float32) + return self._numpy(embeddings) + + def train(self): + FixedKSparseRouter, GroupAssigner, TextFeatureEncoder, simulate = ( + self._official_components() + ) + embeddings, performance, costs = self._prepare_training_data() + group_train = self._group_features(embeddings, training=True) + labels = ( + self.train_df["eval_name"].astype(str).to_numpy() + if "eval_name" in self.train_df + else None + ) + self.group_assigner = GroupAssigner( + strategy=self.save_config["group_strategy"], + n_groups=self.save_config["n_groups"], + seed=self.seed, + ) + train_groups = self.group_assigner.fit_predict(group_train, labels) + + self.feature_encoder = TextFeatureEncoder() + include_dense = bool(self.save_config["include_dense_context"]) + self.include_dense_context = include_dense + context_train = self.feature_encoder.fit_transform( + texts=self.train_df["prompt"].astype(str).tolist(), + dense=group_train if include_dense else None, + ) + k = int(self.save_config["k"]) + supervision = simulate( + performance, + train_groups, + costs=costs, + k=k, + seed=self.seed, + beta=float(self.save_config["acquisition_beta"]), + prior_strength=float(self.save_config["acquisition_prior_strength"]), + prior_alpha=float(self.save_config["prior_alpha"]), + prior_beta=float(self.save_config["prior_beta"]), + seed_stride=int(self.save_config["seed_stride"]), + reset_each_pass=True, + query_ids=[str(index) for index in self.train_df.index], + model_ids=self.model_list, + ) + self.model = FixedKSparseRouter( + prior_strength=float(self.save_config["prior_strength"]), + prior_ridge_alpha=float(self.save_config["prior_ridge_alpha"]), + residual_ridge_alpha=float(self.save_config["residual_ridge_alpha"]), + residual_gamma=float(self.save_config["residual_gamma"]), + min_model_observations=int(self.save_config["min_model_observations"]), + ).fit(context_train, supervision) + self.supervision = supervision + logging.info( + "[SaveRouter] Acquired %d/%d feedback pairs (density %.4f, K=%d).", + supervision.n_observations, + supervision.metadata.get("available_pairs", len(performance) * performance.shape[1]), + supervision.density, + k, + ) + + def predict(self, test_embedding): + if self.model is None: + raise RuntimeError("SaveRouter must be trained before predict().") + group_test = self._group_features(test_embedding, training=False) + if len(group_test) != len(self.test_df): + raise ValueError("SaveRouter predictions must align with ORBIT's test split.") + test_groups = self.group_assigner.predict(group_test) + context_test = self.feature_encoder.transform( + texts=self.test_df["prompt"].astype(str).tolist(), + dense=group_test if self.include_dense_context else None, + ) + scores = self.model.predict_scores(context_test, group_ids=test_groups) + costs = self.model.predict_costs(context_test, group_ids=test_groups) + if not np.isfinite(scores).all() or not np.isfinite(costs).all(): + raise ValueError("SaveRouter returned non-finite predictions.") + return scores.astype(np.float32), costs.astype(np.float32) diff --git a/methods/__init__.py b/methods/__init__.py index 49e8dd6..9c1ad9b 100644 --- a/methods/__init__.py +++ b/methods/__init__.py @@ -23,3 +23,5 @@ from methods.ProfileRouter import ProfileRouter from methods.CarrotRouter import CarrotRouter from methods.EARAMRouter import EARAMRouter +from methods.RouteFM import RouteFMRouter +from methods.SaveRouter import SaveRouter diff --git a/requirement.txt b/requirement.txt index 91d17b8..21cb397 100644 --- a/requirement.txt +++ b/requirement.txt @@ -16,6 +16,8 @@ accelerate>=0.30,<2.0 huggingface-hub>=0.24,<0.36 sentence-transformers>=2.6,<6.0 git+https://github.com/openai/CLIP.git +routefm-router @ git+https://github.com/LAMDA-Model-Reuse/RouteFM.git@3e23c7d903f7d7e1b943d0c867fadfdb3eb23eb1 +saverouter @ git+https://github.com/LAMDA-Model-Reuse/SaveRouter.git@95532ccce659377d2aa1a651c9908ec631c385b5 # ---------------------------- # Data processing diff --git a/requirements-ci.txt b/requirements-ci.txt index bcaf901..2f526b5 100644 --- a/requirements-ci.txt +++ b/requirements-ci.txt @@ -6,6 +6,8 @@ transformers>=4.45,<4.58 datasets>=2.20,<4.5 huggingface-hub>=0.24,<0.36 git+https://github.com/openai/CLIP.git +routefm-router @ git+https://github.com/LAMDA-Model-Reuse/RouteFM.git@3e23c7d903f7d7e1b943d0c867fadfdb3eb23eb1 +saverouter @ git+https://github.com/LAMDA-Model-Reuse/SaveRouter.git@95532ccce659377d2aa1a651c9908ec631c385b5 numpy>=2.0,<2.4 pandas>=2.2,<2.4 scipy>=1.13,<1.17 diff --git a/tests/test_paper_adapters.py b/tests/test_paper_adapters.py index 9e5f32f..b0de6ec 100644 --- a/tests/test_paper_adapters.py +++ b/tests/test_paper_adapters.py @@ -11,6 +11,8 @@ from methods.CarrotRouter import CarrotRouter from methods.EARAMRouter import EARAMRouter from methods.ProfileRouter import ProfileRouter +from methods.RouteFM import RouteFMRouter +from methods.SaveRouter import SaveRouter class _FakeEmbedder: @@ -224,6 +226,99 @@ def test_earam_train_predict_and_auction_smoke(self): ) np.testing.assert_array_equal(null_winner, [-1]) + def test_routefm_official_episode_adapter_smoke(self): + from routefm.models import RouteFM, RouteFMConfig + + official_model = RouteFM( + RouteFMConfig( + query_dim=4, + hidden_dim=8, + projection_dim=8, + ffn_dim=16, + heads=2, + profile_layers=1, + readout_layers=1, + pool_layers=1, + capability_tokens=2, + quality_bins=5, + observation_features=2, + observation_schema="score_cost", + dropout=0.0, + use_pool_transformer=False, + architecture="profile_context", + local_context_layers=1, + ) + ).eval() + args = self._args("RouteFM") + args["routefm"] = { + "encoder": "bge", + "context_size": 3, + "target_batch_size": 2, + } + dataset_patch, embedder_patch = self._patches() + with dataset_patch, embedder_patch, patch.object( + RouteFMRouter, + "_load_official_router", + return_value=official_model, + ): + router = RouteFMRouter(args) + router.train() + features = router.embedder.run_embed( + texts=self.test_df["prompt"].tolist(), images=None + ) + performance, cost = router.predict(features) + with patch.object(router, "cal_rci") as rci, patch.object( + router, "cal_metrics" + ) as metrics: + router.evaluate() + rci.assert_called_once() + metrics.assert_called_once() + self.assertEqual(router.context_query.shape, (3, 3, 4)) + self.assertEqual(router.context_features.shape, (3, 3, 2)) + self.assertTrue(router.context_mask.all()) + self.assertEqual(performance.shape, (4, 3)) + self.assertEqual(cost.shape, (4, 3)) + self.assertTrue(np.isfinite(performance).all()) + self.assertTrue(np.isfinite(cost).all()) + self.assertTrue((cost >= 0).all()) + + def test_saverouter_sparse_supervision_adapter_smoke(self): + args = self._args("SaveRouter") + args["saverouter"] = { + "k": 2, + "group_strategy": "auto", + "n_groups": None, + "include_dense_context": False, + "prior_strength": 40.0, + "prior_ridge_alpha": 10.0, + "residual_ridge_alpha": 20.0, + "residual_gamma": 2.0, + "min_model_observations": 2, + } + dataset_patch, embedder_patch = self._patches() + with dataset_patch, embedder_patch: + router = SaveRouter(args) + router.train() + features = router.embedder.run_embed( + texts=self.test_df["prompt"].tolist(), images=None + ) + performance, cost = router.predict(features) + with patch.object(router, "cal_rci") as rci, patch.object( + router, "cal_metrics" + ) as metrics: + router.evaluate() + rci.assert_called_once() + metrics.assert_called_once() + self.assertEqual(router.supervision.n_observations, len(self.train_df) * 2) + np.testing.assert_array_equal( + router.supervision.observation_counts, + np.full(len(self.train_df), 2), + ) + self.assertEqual(performance.shape, (4, 3)) + self.assertEqual(cost.shape, (4, 3)) + self.assertTrue(np.isfinite(performance).all()) + self.assertTrue(np.isfinite(cost).all()) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_router_implementations.py b/tests/test_router_implementations.py index a6fd238..c5af62d 100644 --- a/tests/test_router_implementations.py +++ b/tests/test_router_implementations.py @@ -162,10 +162,12 @@ def test_modelsat_selects_last_non_padding_position(self): torch.testing.assert_close(selected[1], logits[1, 3]) def test_registry_contains_only_importable_routers(self): - self.assertEqual(len(ROUTER_REGISTRY), 28) + self.assertEqual(len(ROUTER_REGISTRY), 30) self.assertIn("ProfileRouter", ROUTER_REGISTRY) self.assertIn("CarrotRouter", ROUTER_REGISTRY) self.assertIn("EARAMRouter", ROUTER_REGISTRY) + self.assertIn("RouteFM", ROUTER_REGISTRY) + self.assertIn("SaveRouter", ROUTER_REGISTRY) if __name__ == "__main__": diff --git a/train.py b/train.py index 7f301fd..efd8696 100644 --- a/train.py +++ b/train.py @@ -9,7 +9,7 @@ Eagle, EmbedLLM, kNNRouter, SVMRouter, RouteLLM_SWRanking, RouteLLM_MF, RouteLLM_BERT, OmniRouter, EquiRouter, RMClassification, RMSoftmax, RMInterval, OracleRouter, TRouter, UniRoute, InferenceDynamics, ModelSAT, ProfileRouter, CarrotRouter, - EARAMRouter + EARAMRouter, RouteFMRouter, SaveRouter ) # ---------------------------- @@ -44,6 +44,8 @@ "ProfileRouter": ProfileRouter, "CarrotRouter": CarrotRouter, "EARAMRouter": EARAMRouter, + "RouteFM": RouteFMRouter, + "SaveRouter": SaveRouter, } diff --git a/utils/embedding.py b/utils/embedding.py index b131db2..4fdaeae 100644 --- a/utils/embedding.py +++ b/utils/embedding.py @@ -312,6 +312,37 @@ def embed_fn(texts): return embed_fn, dim +def CreateBGEBaseEnV15( + model_dir: str = "./bge-base-en-v1.5", device: str = "cpu" +): + """Create RouteFM's official normalized BGE CLS encoder.""" + if not os.path.exists(model_dir): + raise FileNotFoundError( + f"[ERROR] Local BGE directory not found: {model_dir}\n" + "Download BAAI/bge-base-en-v1.5 and set embeddings.text_model_dir." + ) + tokenizer = AutoTokenizer.from_pretrained(model_dir) + model = AutoModel.from_pretrained(model_dir).to(device).eval() + dim = 768 + + @torch.no_grad() + def embed_fn(texts): + if isinstance(texts, str): + texts = [texts] + tokens = tokenizer( + list(texts), padding=True, truncation=True, max_length=512, + return_tensors="pt", + ) + tokens = {key: value.to(device) for key, value in tokens.items()} + cls = model(**tokens).last_hidden_state[:, 0].float() + values = torch.nn.functional.normalize(cls, dim=-1) + if values.shape[1] != dim or not torch.isfinite(values).all(): + raise ValueError("BGE encoder must produce finite normalized 768-D vectors.") + return values + + return embed_fn, dim + + def CreateQwen3_Embedding_8B( model_dir: str = "./Qwen3-Embedding-8B", device: str = "cpu" ): @@ -573,6 +604,13 @@ def _build_longformer(args: Dict[str, Any]): return CreateLongformerBase4096(model_dir=model_dir, device=device) +@TEXT_EMBEDDERS.register("BAAI/bge-base-en-v1.5") +def _build_bge_base_en_v15(args: Dict[str, Any]): + model_dir = args["embeddings"].get("text_model_dir", "./bge-base-en-v1.5") + device = args["embeddings"].get("device", "cpu") + return CreateBGEBaseEnV15(model_dir=model_dir, device=device) + + @TEXT_EMBEDDERS.register("ViT-B/16") def _build_clip_text(args: Dict[str, Any]): device = args["embeddings"].get("device", "cpu")