Initial release: R2-Router (#1 on RouterArena)

README.md

@@ -0,0 +1,178 @@
+---
+license: apache-2.0
+tags:
+  - llm-routing
+  - model-selection
+  - budget-optimization
+  - knn
+language:
+  - en
+library_name: sklearn
+pipeline_tag: text-classification
+---
+
+# R2-Router: LLM Router with Joint Model-Budget Optimization
+
+**R2-Router** intelligently routes each query to the optimal (LLM, token budget) pair, jointly optimizing accuracy and inference cost. Ranked **#1** on the [RouterArena](https://routerarena.github.io/) leaderboard.
+
+**Paper**: [R2-Router (arxiv)](https://arxiv.org/abs/TODO)
+
+## RouterArena Performance
+
+Official leaderboard results on 8,400 queries:
+
+| Metric | Value |
+|--------|-------|
+| Accuracy | 71.23% |
+| Cost per 1K Queries | $0.061 |
+| Arena Score (beta=0.1) | **71.60** |
+| Robustness Score | 45.71% |
+| Rank | **#1** |
+
+## Quick Start
+
+### Installation
+
+```bash
+pip install scikit-learn numpy joblib huggingface_hub
+```
+
+### Load Pre-trained Checkpoints
+
+```python
+from router import R2Router
+
+# Load pre-trained KNN checkpoints (no training needed)
+router = R2Router.from_pretrained("jqxue1999/r2-router")
+
+# Route a query (requires 1024-dim embedding from Qwen3-0.6B)
+result = router.route(embedding)
+print(f"Model: {result['model_full_name']}")
+print(f"Token Budget: {result['token_limit']}")
+print(f"Predicted Quality: {result['predicted_quality']:.3f}")
+```
+
+### Train from Scratch
+
+```python
+from router import R2Router
+
+# Train KNN from the provided sub_10 training data
+router = R2Router.from_training_data("jqxue1999/r2-router", k=80)
+
+# Route a query
+result = router.route(embedding)
+```
+
+### Get Query Embeddings
+
+R2-Router uses [Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B) embeddings (1024-dim). You can generate them with:
+
+```python
+from sentence_transformers import SentenceTransformer
+
+model = SentenceTransformer("Qwen/Qwen3-0.6B")
+embedding = model.encode("What is the capital of France?")
+```
+
+Or with vLLM for faster batch inference:
+
+```python
+from vllm import LLM
+llm = LLM(model="Qwen/Qwen3-0.6B", runner="pooling")
+outputs = llm.embed(["What is the capital of France?"])
+embedding = outputs[0].outputs.embedding
+```
+
+## Architecture
+
+R2-Router jointly optimizes **which model** to use and **how many tokens** to allocate per query.
+
+### Routing Formula
+
+```
+risk(M, b) = (1 - lambda) * predicted_quality(query, M, b) - lambda * predicted_tokens(query, M) * price_M / 1e6
+(M*, b*) = argmax risk
+```
+
+### Pipeline
+
+```
+Input Query
+    |
+[1] Embed with Qwen3-0.6B -> 1024-dim vector
+    |
+[2] For each (model, budget) pair:
+      - KNN predicts quality (accuracy)
+      - KNN predicts output token count
+      - Compute risk = (1-lambda) * quality - lambda * cost
+    |
+[3] Select (model, budget) with highest risk
+    |
+Output: (model_name, token_budget)
+```
+
+### Model Pool (6 LLMs)
+
+| Model | Output $/M tokens |
+|-------|------------------|
+| Qwen3-235B-A22B | $0.463 |
+| Qwen3-Next-80B-A3B | $1.10 |
+| Qwen3-30B-A3B | $0.33 |
+| Qwen3-Coder-Next | $0.30 |
+| Gemini 2.5 Flash | $2.50 |
+| Claude 3 Haiku | $1.25 |
+
+### Token Budgets
+
+4 output token limits: **100, 200, 400, 800** tokens.
+
+### Key Parameters
+
+| Parameter | Value |
+|-----------|-------|
+| KNN K | 80 |
+| Lambda | 0.999 |
+| Distance Metric | Cosine |
+| KNN Weights | Distance-weighted |
+| Embedding Dim | 1024 |
+
+## Repository Contents
+
+```
+config.json             # Router configuration (models, budgets, prices, hyperparams)
+router.py               # Self-contained inference code
+training_data/
+  embeddings.npy        # Sub_10 training embeddings (809 x 1024)
+  labels.json           # Per-(model, budget) accuracy & token labels
+checkpoints/
+  quality_knn_*.joblib  # Pre-fitted KNN quality predictors (18 total)
+  token_knn_*.joblib    # Pre-fitted KNN token predictors (6 total)
+```
+
+### Two Ways to Use
+
+1. **Load checkpoints** (`from_pretrained`): Directly load pre-fitted KNN models. No training needed.
+2. **Train from data** (`from_training_data`): Use the provided training embeddings and labels to fit your own KNN with custom hyperparameters (e.g., different K, distance metric).
+
+## Training Details
+
+- **Training Data**: RouterArena sub_10 split (809 queries, 10% of full 8,400)
+- **Method**: KNeighborsRegressor with cosine distance, distance-weighted
+- **Evaluation**: Full 8,400 RouterArena queries (no data leakage)
+- **Training Time**: < 1 second (KNN fitting)
+
+## Citation
+
+```bibtex
+@article{r2router2026,
+  title={R2-Router: A New Paradigm for LLM Routing with Reasoning},
+  author={TODO},
+  year={2026},
+  url={https://arxiv.org/abs/TODO}
+}
+```
+
+## License
+
+Apache 2.0

checkpoints/quality_knn_235b_budget_200.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f66789acf7675f2b4113c9ac173aaaa8199fdd270dfa6869609ed17b2feb7dbe
+size 3317588

checkpoints/quality_knn_235b_budget_400.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c55dbaa6b9cda97db793b63471a95ff1b1389bc1e46acf8ee724f4deda1280b7
+size 3317588

checkpoints/quality_knn_235b_budget_800.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:51ef60d08cbb7e2c309af411054aaa1c2deee661f8dc0c1953228823fca56066
+size 3317588

checkpoints/quality_knn_235b_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:76a9a2f93c64e1eb8243a35b7b94b1dcb66deb7e0d49e095866f4e3057bd7ea5
+size 3317588

checkpoints/quality_knn_30b_budget_200.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:101348a979019d849d4dc56b1629844bd6c5d77adf1746b6cd28dbec6702aecb
+size 3317588

checkpoints/quality_knn_30b_budget_400.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb2d64cca2cf69d56d72e9d1c700fbc3247d04831df2f1a154212f3cc98202fd
+size 3317588

checkpoints/quality_knn_30b_budget_800.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:136bea24e03a67f7580e21e0c54306abeb883aab4ef804bdf9b1be238ad3bdbd
+size 3317588

checkpoints/quality_knn_30b_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7e69c75c0aeb15ae58298e34c42a834c1e69cce9b844ab25c38787384490550
+size 3317588

checkpoints/quality_knn_80b_budget_200.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6f86c6981e31a9057ee7a2112db13eaa34c459471dfffb52f8aa9a3bdcce29e
+size 3317588

checkpoints/quality_knn_80b_budget_400.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38cc8af776d03877dd0a513437e6637d9dd5a46b3a26ce47918548744de0f2e9
+size 3317588

checkpoints/quality_knn_80b_budget_800.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6702cd8020a07189ab5943c20061776941bc3eae3611920a0d07500cf57daea
+size 3317588

checkpoints/quality_knn_80b_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:002f5021ea92a5d846244ba1ec76a03d773c93174281c789284bd63aecd54348
+size 3317588

checkpoints/quality_knn_coder-next_budget_200.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3cea7b7f12e959e401319642d4aac688fb1a96f76f591466e5cfe4ad09b30ad
+size 3317588

checkpoints/quality_knn_coder-next_budget_400.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:528b2e6e7ab9bde25d4d6639dd2b359d9e750a73ddb8a1695c496756e79cd8d3
+size 3317588

checkpoints/quality_knn_coder-next_budget_800.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2a4e599bf2490595a222792037329e97ef7630d917db67f05b4d3bfae17e60be
+size 3317588

checkpoints/quality_knn_coder-next_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4c2df7aabc3fee5ae12a05b3296eafd6f54053c6b1bf6a3cf5e6cec4cf4574c
+size 3317588

checkpoints/quality_knn_gemini-flash_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ffefebf08d5de8c2784ee3074a38e470fe307085bc2f8a62f463c855ab0e9c6
+size 3317588

checkpoints/quality_knn_haiku_concise.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22633da4eecf1d0b1b3da5fc1bc2e20889ec1271816b949909863ce6c5a624b0
+size 3317588

checkpoints/token_knn_235b.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d8e62844ac6f97e8d4cdf7d85af9a801c3850032600634a41378bd5ae436900
+size 3317588

checkpoints/token_knn_30b.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f1a7972693a4b5727324c401dbe8beefa187be05e07fefdc0ff2f490ee256b5
+size 3317588

checkpoints/token_knn_80b.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aea91620048b766624696d0dfe2ad0998fceb854183aeb011a7c41d723b2e217
+size 3317588

checkpoints/token_knn_coder-next.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8374d3e3159cf9583c3a4dbb71fd7c3c2f69664bd49cf40e05851922d3e65f4a
+size 3317588

checkpoints/token_knn_gemini-flash.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3bd50eb44333e5e7e8ff5cbcec995b673ac54224a3a86fe8dffdf8d4e1d129e
+size 3317588

checkpoints/token_knn_haiku.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89578d06d0aa2640b720f0599bb91d66e4a52b39f475475db1aa33a833f53178
+size 3317588

config.json

@@ -0,0 +1,54 @@
+{
+  "router_name": "R2-Router",
+  "method": "Global KNN (cosine, distance-weighted)",
+  "embedding_model": "Qwen/Qwen3-0.6B",
+  "embedding_dim": 1024,
+  "knn_k": 80,
+  "lambda": 0.999,
+  "models": {
+    "235b": {
+      "full_name": "qwen/qwen3-235b-a22b-2507",
+      "output_price_per_million": 0.463,
+      "input_price_per_million": 0.071,
+      "type": "vllm"
+    },
+    "80b": {
+      "full_name": "qwen/qwen3-next-80b-a3b-instruct",
+      "output_price_per_million": 1.1,
+      "input_price_per_million": 0.09,
+      "type": "vllm"
+    },
+    "30b": {
+      "full_name": "qwen/qwen3-30b-a3b-instruct-2507",
+      "output_price_per_million": 0.33,
+      "input_price_per_million": 0.08,
+      "type": "vllm"
+    },
+    "coder-next": {
+      "full_name": "Qwen/Qwen3-Coder-Next",
+      "output_price_per_million": 0.3,
+      "input_price_per_million": 0.07,
+      "type": "vllm"
+    },
+    "gemini-flash": {
+      "full_name": "gemini-2.5-flash",
+      "output_price_per_million": 2.5,
+      "input_price_per_million": 0.3,
+      "type": "api"
+    },
+    "haiku": {
+      "full_name": "claude-3-haiku-20240307",
+      "output_price_per_million": 1.25,
+      "input_price_per_million": 0.25,
+      "type": "api"
+    }
+  },
+  "budgets": {
+    "concise": 100,
+    "budget_200": 200,
+    "budget_400": 400,
+    "budget_800": 800
+  },
+  "training_size": 809,
+  "training_source": "RouterArena sub_10 (10% official split)"
+}

router.py

@@ -0,0 +1,252 @@
+"""
+R2-Router: LLM Router with Joint Model-Budget Optimization
+
+Self-contained inference module. Routes queries to the optimal (model, token_budget)
+pair by predicting per-query quality and cost using KNN.
+
+Usage:
+    from router import R2Router
+    router = R2Router.from_pretrained("jqxue1999/r2-router")
+    result = router.route(embedding)  # embedding: np.ndarray (1024,)
+
+    # Or train from scratch:
+    router = R2Router.from_training_data("jqxue1999/r2-router")
+"""
+
+import os
+import json
+import numpy as np
+import joblib
+from typing import Dict, List, Optional, Union
+from sklearn.neighbors import KNeighborsRegressor
+
+
+class R2Router:
+    """
+    R2-Router: Routes queries to optimal (LLM, token_budget) pair.
+
+    Uses KNN to predict quality for each (model, budget) combination,
+    then selects the pair that maximizes:
+        risk = (1 - lambda) * quality - lambda * tokens * price / 1e6
+    """
+
+    def __init__(
+        self,
+        quality_knns: Dict[str, Dict[str, KNeighborsRegressor]],
+        token_knns: Dict[str, KNeighborsRegressor],
+        model_prices: Dict[str, float],
+        model_names: Dict[str, str],
+        budgets: Dict[str, int],
+        lambda_val: float = 0.999,
+    ):
+        self.quality_knns = quality_knns  # {model: {budget: KNN}}
+        self.token_knns = token_knns      # {model: KNN}
+        self.model_prices = model_prices  # {model: price_per_million_output_tokens}
+        self.model_names = model_names    # {short_name: full_name}
+        self.budgets = budgets            # {budget_name: token_limit}
+        self.lambda_val = lambda_val
+
+    @classmethod
+    def from_pretrained(cls, path: str, lambda_val: float = 0.999) -> "R2Router":
+        """
+        Load pre-trained KNN checkpoints.
+
+        Args:
+            path: Local directory or HuggingFace repo ID (e.g., "jqxue1999/r2-router")
+            lambda_val: Cost-accuracy tradeoff (higher = more cost-sensitive)
+        """
+        # If HF repo ID, download first
+        if not os.path.isdir(path):
+            path = cls._download_from_hf(path)
+
+        with open(os.path.join(path, "config.json")) as f:
+            config = json.load(f)
+
+        ckpt_dir = os.path.join(path, "checkpoints")
+        quality_knns = {}
+        token_knns = {}
+
+        for model_name in config["models"]:
+            quality_knns[model_name] = {}
+            for budget_name in config["budgets"]:
+                ckpt_path = os.path.join(ckpt_dir, f"quality_knn_{model_name}_{budget_name}.joblib")
+                if os.path.exists(ckpt_path):
+                    quality_knns[model_name][budget_name] = joblib.load(ckpt_path)
+
+            tok_path = os.path.join(ckpt_dir, f"token_knn_{model_name}.joblib")
+            if os.path.exists(tok_path):
+                token_knns[model_name] = joblib.load(tok_path)
+
+        model_prices = {
+            mn: cfg["output_price_per_million"]
+            for mn, cfg in config["models"].items()
+        }
+        model_names = {
+            mn: cfg["full_name"]
+            for mn, cfg in config["models"].items()
+        }
+
+        return cls(
+            quality_knns=quality_knns,
+            token_knns=token_knns,
+            model_prices=model_prices,
+            model_names=model_names,
+            budgets=config["budgets"],
+            lambda_val=lambda_val,
+        )
+
+    @classmethod
+    def from_training_data(
+        cls,
+        path: str,
+        k: int = 80,
+        lambda_val: float = 0.999,
+    ) -> "R2Router":
+        """
+        Train KNN from scratch using the provided training data.
+
+        Args:
+            path: Local directory or HuggingFace repo ID
+            k: Number of KNN neighbors
+            lambda_val: Cost-accuracy tradeoff
+        """
+        if not os.path.isdir(path):
+            path = cls._download_from_hf(path)
+
+        with open(os.path.join(path, "config.json")) as f:
+            config = json.load(f)
+
+        X_train = np.load(os.path.join(path, "training_data", "embeddings.npy"))
+        with open(os.path.join(path, "training_data", "labels.json")) as f:
+            labels = json.load(f)
+
+        quality_knns = {}
+        token_knns = {}
+
+        for model_name, model_labels in labels.items():
+            quality_knns[model_name] = {}
+            for budget_name, bdata in model_labels.items():
+                acc = np.array([x if x is not None else np.nan for x in bdata["accuracy"]])
+                valid = ~np.isnan(acc)
+                if valid.sum() < 3:
+                    continue
+                knn = KNeighborsRegressor(
+                    n_neighbors=min(k, int(valid.sum()) - 1),
+                    metric="cosine",
+                    weights="distance",
+                )
+                knn.fit(X_train[valid], acc[valid])
+                quality_knns[model_name][budget_name] = knn
+
+            # Token predictor (use concise budget's output_tokens)
+            if "concise" in model_labels and "output_tokens" in model_labels["concise"]:
+                tok = np.array([x if x is not None else np.nan for x in model_labels["concise"]["output_tokens"]])
+                valid = ~np.isnan(tok)
+                if valid.sum() >= 3:
+                    tknn = KNeighborsRegressor(
+                        n_neighbors=min(k, int(valid.sum()) - 1),
+                        metric="cosine",
+                        weights="distance",
+                    )
+                    tknn.fit(X_train[valid], tok[valid])
+                    token_knns[model_name] = tknn
+
+        model_prices = {
+            mn: cfg["output_price_per_million"]
+            for mn, cfg in config["models"].items()
+        }
+        model_names = {
+            mn: cfg["full_name"]
+            for mn, cfg in config["models"].items()
+        }
+
+        return cls(
+            quality_knns=quality_knns,
+            token_knns=token_knns,
+            model_prices=model_prices,
+            model_names=model_names,
+            budgets=config["budgets"],
+            lambda_val=lambda_val,
+        )
+
+    @staticmethod
+    def _download_from_hf(repo_id: str) -> str:
+        """Download model from Hugging Face Hub."""
+        try:
+            from huggingface_hub import snapshot_download
+            return snapshot_download(repo_id)
+        except ImportError:
+            raise ImportError(
+                "huggingface_hub is required to download from HF. "
+                "Install with: pip install huggingface_hub"
+            )
+
+    def route(
+        self,
+        embedding: np.ndarray,
+        lambda_val: Optional[float] = None,
+    ) -> Dict:
+        """
+        Route a query to the optimal (model, token_budget) pair.
+
+        Args:
+            embedding: Query embedding vector, shape (1024,) or (1, 1024)
+            lambda_val: Override default lambda (higher = more cost-sensitive)
+
+        Returns:
+            Dict with keys: model, model_full_name, budget, token_limit,
+                           predicted_quality, predicted_cost, risk, all_options
+        """
+        if embedding.ndim == 1:
+            embedding = embedding.reshape(1, -1)
+
+        lam = lambda_val if lambda_val is not None else self.lambda_val
+        all_options = []
+
+        for mn in self.quality_knns:
+            price = self.model_prices.get(mn, 0)
+
+            # Predict output tokens
+            if mn in self.token_knns:
+                tok = max(1.0, float(self.token_knns[mn].predict(embedding)[0]))
+            else:
+                tok = 50.0
+
+            for budget_name, knn in self.quality_knns[mn].items():
+                q = float(knn.predict(embedding)[0])
+                risk = (1 - lam) * q - lam * tok * price / 1e6
+
+                all_options.append({
+                    "model": mn,
+                    "model_full_name": self.model_names.get(mn, mn),
+                    "budget": budget_name,
+                    "token_limit": self.budgets.get(budget_name, budget_name),
+                    "predicted_quality": q,
+                    "predicted_tokens": tok,
+                    "predicted_cost": tok * price / 1e6,
+                    "risk": risk,
+                })
+
+        if not all_options:
+            raise RuntimeError("No valid routing options")
+
+        best = max(all_options, key=lambda x: x["risk"])
+        best["all_options"] = all_options
+        return best
+
+    def route_batch(
+        self,
+        embeddings: np.ndarray,
+        lambda_val: Optional[float] = None,
+    ) -> List[Dict]:
+        """
+        Route a batch of queries.
+
+        Args:
+            embeddings: Query embeddings, shape (N, 1024)
+            lambda_val: Override default lambda
+
+        Returns:
+            List of routing decisions
+        """
+        return [self.route(embeddings[i], lambda_val) for i in range(len(embeddings))]

training_data/embeddings.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24f0f5116dc1d153739d6a30604dc4a6553bbd9d7d4708a1d14fa5b00041bd6c
+size 3313792