vLLM-based route_text(), uv install guide, complete runnable examples

README.md

@@ -33,17 +33,22 @@ Official leaderboard results on 8,400 queries:
 
 ### Installation
 
+We recommend using [uv](https://docs.astral.sh/uv/) for fast, reliable environment setup:
+
 ```bash
-pip install scikit-learn numpy joblib huggingface_hub torch transformers>=4.51
+# Install uv (if not already installed)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Create environment and install dependencies
+uv venv .venv && source .venv/bin/activate
+uv pip install scikit-learn numpy joblib huggingface_hub vllm
 ```
 
-### Complete Example
+### Complete Example (GPU)
 
 ```python
 from huggingface_hub import snapshot_download
-import sys, torch
-import numpy as np
-from transformers import AutoModel, AutoTokenizer
+import sys
 
 # 1. Download router
 path = snapshot_download("JiaqiXue/r2-router")
@@ -54,48 +59,54 @@ from router import R2Router
 # 2. Load pre-trained KNN checkpoints
 router = R2Router.from_pretrained(path)
 
-# 3. Embed your query with Qwen3-0.6B (1024-dim)
-tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
-model = AutoModel.from_pretrained("Qwen/Qwen3-0.6B")
-
-query = "What is the capital of France?"
-inputs = tokenizer(query, return_tensors="pt", padding=True, truncation=True)
-with torch.no_grad():
-    output = model(**inputs)
-embedding = output.last_hidden_state.mean(dim=1).squeeze().numpy()
-
-# 4. Route!
-result = router.route(embedding)
+# 3. Route a query (auto-embeds with Qwen3-0.6B via vLLM)
+result = router.route_text("What is the capital of France?")
 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
+`route_text()` automatically loads Qwen3-0.6B via vLLM on first call and caches it. Batch routing is also supported:
 
 ```python
-from huggingface_hub import snapshot_download
-import sys
+queries = [
+    "What is the capital of France?",
+    "Write a Python function to sort a list",
+    "Translate 'hello' to Japanese",
+]
+results = router.route_text(queries)
+for q, r in zip(queries, results):
+    print(f"{q[:40]:40s} -> {r['model']} (budget={r['token_limit']})")
+```
 
-path = snapshot_download("JiaqiXue/r2-router")
-sys.path.insert(0, path)
+### CPU-Only (No GPU)
+
+If you don't have a GPU, provide pre-computed embeddings directly:
 
+```python
+import numpy as np
 from router import R2Router
 
-# Train KNN from the provided sub_10 training data (custom hyperparameters)
-router = R2Router.from_training_data(path, k=80)
-```
+router = R2Router.from_pretrained(path)
 
-### Alternative: vLLM Embeddings (Faster for Batches)
+# Your own 1024-dim embedding (e.g., from an API or pre-computed)
+embedding = np.random.randn(1024)  # replace with real embedding
+result = router.route(embedding)
+```
 
-For GPU-accelerated batch embedding:
+### Train from Scratch
 
 ```python
-from vllm import LLM
+from huggingface_hub import snapshot_download
+import sys
+
+path = snapshot_download("JiaqiXue/r2-router")
+sys.path.insert(0, path)
+
+from router import R2Router
 
-llm = LLM(model="Qwen/Qwen3-0.6B", runner="pooling")
-outputs = llm.embed(["What is the capital of France?"])
-embedding = outputs[0].outputs.embedding
+# Train KNN with custom hyperparameters
+router = R2Router.from_training_data(path, k=80, lambda_val=0.999)
 ```
 
 ## Architecture
@@ -155,7 +166,7 @@ Output: (model_name, token_budget)
 
 ```
 config.json             # Router configuration (models, budgets, prices, hyperparams)
-router.py               # Self-contained inference code
+router.py               # Self-contained inference code (embed + route)
 training_data/
   embeddings.npy        # Sub_10 training embeddings (809 x 1024)
   labels.json           # Per-(model, budget) accuracy & token labels
@@ -164,10 +175,13 @@ checkpoints/
   token_knn_*.joblib    # Pre-fitted KNN token predictors (6 total)
 ```
 
-### Two Ways to Use
+### Three 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).
+| Method | GPU? | Description |
+|--------|------|-------------|
+| `router.route_text(query)` | Yes | End-to-end: auto-embeds with vLLM, then routes |
+| `router.route(embedding)` | No | Route from pre-computed 1024-dim embedding |
+| `R2Router.from_training_data(path)` | No | Train your own KNN with custom hyperparameters |
 
 ## Training Details
 

router.py

@@ -6,11 +6,17 @@ 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")
+    router = R2Router.from_pretrained(path)
+
+    # Option 1: Route from text (auto-embeds with vLLM)
+    result = router.route_text("What is the capital of France?")
+
+    # Option 2: Route from pre-computed embedding
+    result = router.route(embedding)  # np.ndarray (1024,)
+
+    # Option 3: Train from scratch
+    router = R2Router.from_training_data(path, k=80)
 """
 
 import os
@@ -45,6 +51,7 @@ class R2Router:
         self.model_names = model_names    # {short_name: full_name}
         self.budgets = budgets            # {budget_name: token_limit}
         self.lambda_val = lambda_val
+        self._embedder = None
 
     @classmethod
     def from_pretrained(cls, path: str, lambda_val: float = 0.999) -> "R2Router":
@@ -52,10 +59,9 @@ class R2Router:
         Load pre-trained KNN checkpoints.
 
         Args:
-            path: Local directory or HuggingFace repo ID (e.g., "jqxue1999/r2-router")
+            path: Local directory or HuggingFace repo ID (e.g., "JiaqiXue/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)
 
@@ -138,7 +144,6 @@ class R2Router:
                 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)
@@ -181,6 +186,57 @@ class R2Router:
                 "Install with: pip install huggingface_hub"
             )
 
+    def embed(self, queries: Union[str, List[str]]) -> np.ndarray:
+        """
+        Embed queries using Qwen3-0.6B via vLLM (loaded on first call).
+
+        Args:
+            queries: Single query string or list of queries
+
+        Returns:
+            numpy array of shape (N, 1024)
+        """
+        if self._embedder is None:
+            try:
+                from vllm import LLM
+            except ImportError:
+                raise ImportError(
+                    "vLLM is required for text embedding. "
+                    "Install with: uv pip install vllm"
+                )
+            self._embedder = LLM(
+                model="Qwen/Qwen3-0.6B",
+                runner="pooling",
+                trust_remote_code=True,
+                dtype="half",
+            )
+
+        if isinstance(queries, str):
+            queries = [queries]
+
+        outputs = self._embedder.embed(queries)
+        return np.array([o.outputs.embedding for o in outputs])
+
+    def route_text(
+        self,
+        query: Union[str, List[str]],
+        lambda_val: Optional[float] = None,
+    ) -> Union[Dict, List[Dict]]:
+        """
+        Route text query(ies) end-to-end: embed with Qwen3-0.6B, then route.
+
+        Args:
+            query: Single query string or list of queries
+            lambda_val: Override default lambda
+
+        Returns:
+            Routing decision dict (single) or list of dicts (batch)
+        """
+        embeddings = self.embed(query)
+        if isinstance(query, str):
+            return self.route(embeddings[0], lambda_val)
+        return self.route_batch(embeddings, lambda_val)
+
     def route(
         self,
         embedding: np.ndarray,
@@ -206,7 +262,6 @@ class R2Router:
         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: