# ML-Based Model Selection for Semantic Router

> **Issue**: [#986 - Implement ML-based model selection techniques](https://github.com/vllm-project/semantic-router/issues/986)  
> **Milestone**: v0.2 - Athena  
> **Reference**: [vLLM SR Iris Blog](https://blog.vllm.ai/2026/01/05/vllm-sr-iris.html)
>
> **📝 Note on Training Data:**  
> All training data files (`training_data_with_category.jsonl`, `benchmark_training_data.jsonl`, `trained_models/*.json`) are **deployment-specific** and should be generated by each user based on their own LLMs, queries, and benchmarks. These files are NOT included in this repository. See [Training Data Format](#training-data-format) for required file formats.
>
> **📥 Download Pretrained Models from HuggingFace:**
>
> ```bash
> pip install huggingface-hub
> cd src/training/model_selection/ml_model_selection
> python download_model.py --output-dir ../../../.cache/ml-models
> ```
>
> See: https://huggingface.co/abdallah1008/semantic-router-ml-models
>
> **Note:** The E2E test (`ml-model-selection` profile) automatically downloads pretrained models
> from HuggingFace during setup. No local training or Python venv setup is required for E2E testing.

## Overview

This module implements machine learning-based model selection algorithms that intelligently route queries to the most appropriate LLM based on query characteristics, historical performance, and learned patterns.

### Implemented Algorithms

| Algorithm | Implementation | Backend | Best For | Key Parameters |
|-----------|----------------|---------|----------|----------------|
| **KNN** (K-Nearest Neighbors) | Rust (`ml-binding/`) | Linfa (`linfa-nn`) | Quality-weighted voting among similar queries | `k` (neighbors) |
| **KMeans** | Rust (`ml-binding/`) | Linfa (`linfa-clustering`) | Cluster-based routing with quality-weighted assignment | `num_clusters` |
| **SVM** (Support Vector Machine) | Rust (`ml-binding/`) | Linfa (`linfa-svm`) | Decision boundaries with RBF kernel (gamma=1.0) | `kernel`, `gamma` |
| **MLP** (Multi-Layer Perceptron) | Rust (`candle-binding/`) | Candle (GPU) | GPU-accelerated neural network routing | `hidden_sizes`, `device` |

> **Note:** KNN, KMeans, and SVM use [Linfa](https://github.com/rust-ml/linfa) via `ml-binding/`. MLP uses [Candle](https://github.com/huggingface/candle) via `candle-binding/` for GPU acceleration (CUDA/Metal).
>
> **Reference:** Implementation aligned with [FusionFactory (arXiv:2507.10540)](https://arxiv.org/abs/2507.10540) query-level fusion approach.

---

## Architecture

```
┌─────────────────────────────────────────────────────────────────────┐
│                    OFFLINE TRAINING (Python)                        │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  Training is done in Python: src/training/model_selection/ml_model_selection/       │
│                                                                     │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────────────┐  │
│  │ Benchmark    │───▶│  Python      │───▶│ Trained Model Files  │  │
│  │ Data (JSONL) │    │  train.py    │    │ knn/kmeans/svm/mlp   │  │
│  └──────────────┘    └──────────────┘    └──────────────────────┘  │
│                             │                       │               │
│                             ▼                       ▼               │
│                      ┌─────────────┐        ┌─────────────┐        │
│                      │ sentence-   │        │ HuggingFace │        │
│                      │ transformers│        │ Hub (models)│        │
│                      └─────────────┘        └─────────────┘        │
│                                                                     │
│  Training Process:                                                  │
│  1. Load benchmark data (query, model, quality, latency)           │
│  2. Generate embeddings (1024-dim via Qwen3)                       │
│  3. Train KNN, KMeans, SVM using scikit-learn                      │
│  4. Save models to JSON format                                      │
│  5. Upload to HuggingFace Hub (optional)                           │
│                                                                     │
└─────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────┐
│                         ONLINE INFERENCE                            │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  ┌──────────────┐    ┌──────────────┐    ┌──────────────────────┐  │
│  │  User Query  │───▶│  Embedding   │───▶│  Feature Vector      │  │
│  │              │    │   Model      │    │  (1024 + 14 = 1038)  │  │
│  └──────────────┘    └──────────────┘    └──────────────────────┘  │
│                                                    │                │
│                                                    ▼                │
│  ┌─────────────────────────────────────────────────────────────┐   │
│  │                    SELECTOR                                  │   │
│  │  • Loads trained model from JSON                            │   │
│  │  • Applies algorithm-specific inference                     │   │
│  │  • Returns best model for query                             │   │
│  └─────────────────────────────────────────────────────────────┘   │
│                               │                                     │
│                               ▼                                     │
│                      ┌─────────────┐                               │
│                      │  Selected   │                               │
│                      │    LLM      │                               │
│                      └─────────────┘                               │
└─────────────────────────────────────────────────────────────────────┘
```

---

## Router Integration

The ML-based selectors (KNN, KMeans, SVM) are integrated into the main semantic router via the `pkg/selection` package.

### How It Works

```
Request (model="auto")
    ↓
Router (pkg/extproc)
    ↓
Decision matched → algorithm.type: "knn" (from config)
    ↓
getSelectionMethod() → returns MethodKNN
    ↓
selection.Registry.Get(MethodKNN) → MLSelectorAdapter
    ↓
MLSelectorAdapter.Select() → modelselection.KNNSelector.Select()
    ↓
Returns selected model
```

### Configuration

Enable ML-based selection in your config YAML:

```yaml
# Per-decision algorithm configuration
decisions:
  - name: "math_decision"
    algorithm:
      type: "knn"  # Options: knn, kmeans, svm, mlp
      knn:
        k: 5
    modelRefs:
      - model: "llama-3.2-1b"
      - model: "mistral-7b"

  - name: "coding_decision"
    algorithm:
      type: "svm"
      svm:
        kernel: "rbf"
    modelRefs:
      - model: "codellama-7b"
      - model: "mistral-7b"

# Global ML selector settings (optional)
model_selection:
  ml:
    models_path: ".cache/ml-models"  # Default location for downloaded models
    knn:
      k: 5
      pretrained_path: "path/to/knn_model.json"
    svm:
      kernel: "rbf"
      gamma: 1.0
      pretrained_path: "path/to/svm_model.json"
    mlp:
      device: "cuda"  # Options: cpu, cuda, metal
      pretrained_path: "path/to/mlp_model.json"
```

### Key Files

| File | Description |
|------|-------------|
| `pkg/selection/ml_adapter.go` | Bridges `modelselection.Selector` → `selection.Selector` |
| `pkg/selection/factory.go` | Creates and registers ML selectors |
| `pkg/extproc/req_filter_classification.go` | Routes to ML selectors based on config |

---

## File Structure

```
src/semantic-router/pkg/modelselection/
├── README.md                    # This file
├── selector.go                  # Inference implementations (loads trained models)
├── persistence.go               # JSON serialization/deserialization
└── selector_test.go             # Unit tests

src/training/model_selection/ml_model_selection/         # Training and validation scripts
├── train.py                     # Main training script (Python)
├── validate.go                  # Production validation (Go/Rust)
├── go.mod                       # Go module for validation
├── benchmark.py                 # Benchmark LLMs for training data
├── download_model.py            # Download from HuggingFace
├── upload_model.py              # Upload to HuggingFace
├── models.py                    # KNN, KMeans, SVM implementations
├── embeddings.py                # Embedding generation (Qwen3)
├── data_loader.py               # Data loading utilities
└── requirements.txt             # Python dependencies

HuggingFace Repositories (downloaded at runtime, NOT in repo):
├── abdallah1008/semantic-router-ml-models     # Trained models
│   ├── knn_model.json
│   ├── kmeans_model.json
│   ├── svm_model.json
│   └── mlp_model.json
└── abdallah1008/ml-selection-benchmark-data   # Benchmark datasets
    └── validation_benchmark_with_gt.jsonl
```

> **Note:** The `data/` directory and its contents (`trained_models/`, `*.jsonl` files) are
> **downloaded from HuggingFace at runtime** by `validate.go` or the E2E tests.
> These files are NOT part of the repository.

---

## Algorithm Details

### 1. KNN (K-Nearest Neighbors)

**How it works:**

- Stores all training embeddings with their associated best-performing model and quality scores
- For new queries, finds K nearest neighbors using cosine similarity
- Uses **quality-weighted voting**: sums quality scores for each model among neighbors
- Returns the model with highest total quality (not just most votes)

**Quality-Weighted Voting Example:**

```text
Query: "Explain quantum entanglement"
5 Nearest Neighbors:
  - llama-3b (quality: 0.92)
  - llama-3b (quality: 0.88)  
  - mistral (quality: 0.95)
  - llama-1b (quality: 0.75)
  - mistral (quality: 0.91)

Quality Sums:
  - llama-3b: 0.92 + 0.88 = 1.80
  - mistral:  0.95 + 0.91 = 1.86  ← Winner (higher quality sum)
  - llama-1b: 0.75

Selected: mistral-7b (despite llama-3b having same vote count)
```

**JSON Structure (`knn_model.json`):**

```json
{
  "algorithm": "knn",
  "trained": true,
  "k": 5,
  "training_records": [
    {
      "embedding": [0.123, -0.456, ...],  // 1038-dim vector (1024 + 14)
      "best_model": "llama-3.2-3b",
      "quality": 0.92,
      "latency_ns": 1500000000,
      "category": "math"
    }
  ]
}
```

### 2. KMeans (Clustering)

**How it works:**

- Clusters training embeddings into K groups
- Assigns the best-performing model to each cluster centroid
- For new queries, finds nearest centroid and returns its model

**JSON Structure (`kmeans_model.json`):**

```json
{
  "algorithm": "kmeans",
  "trained": true,
  "num_clusters": 8,
  "centroids": [
    [0.123, -0.456, ...],  // 1038-dim centroid vectors (1024 + 14)
    ...
  ],
  "cluster_models": ["llama-3.2-1b", "llama-3.2-3b", ...]
}
```

### 3. SVM (Support Vector Machine)

**How it works:**

- Binary classifiers for each model (one-vs-all)
- Uses **RBF kernel by default** with gamma=1.0 (optimal for 782-dim normalized embeddings)
- **Quality+Latency weighted training**: Records are weighted by `0.9*quality + 0.1*speed`
- Stores support vectors and alpha coefficients
- Model with highest decision score wins

**Weighted Training (Oversampling):**

- High-quality, fast responses are duplicated more in training (1-5 copies)
- Low-quality, slow responses have less influence
- Formula: `weight = 0.9 * quality + 0.1 * speed_factor`
- Matches global `--quality-weight 0.9` hyperparameter (90% quality, 10% speed)

**Kernel Options:**

| Kernel | Gamma | Best For |
|--------|-------|----------|
| **RBF** (default) | 1.0 | High-dimensional embeddings, non-linear boundaries |
| Linear | N/A | Simple linear separable data |

**JSON Structure (`svm_model.json`):**

```json
{
  "algorithm": "svm",
  "trained": true,
  "kernel": "rbf",
  "gamma": 1.0,
  "support_vectors": {
    "0": [[...], [...], ...],  // Support vectors for model 0
    "1": [[...], [...], ...]   // Support vectors for model 1
  },
  "alphas": {
    "0": [...],  // Alpha coefficients for model 0
    "1": [...]   // Alpha coefficients for model 1
  },
  "biases": [0.5, -0.3, ...]
}
```

### 4. MLP (Multi-Layer Perceptron)

**How it works:**

- Feedforward neural network with configurable hidden layers (default: 256→128)
- Uses ReLU activation, BatchNorm, and Dropout for regularization
- **GPU-accelerated inference** via Candle (CUDA/Metal support)
- Trained with quality+latency weighted CrossEntropyLoss
- Early stopping based on validation loss

**Network Architecture:**

```text
Input (1038-dim)
    ↓
Linear(1038, 256) → ReLU → BatchNorm → Dropout(0.1)
    ↓
Linear(256, 128) → ReLU → BatchNorm → Dropout(0.1)
    ↓
Linear(128, num_models)
    ↓
Softmax → Selected Model
```

**GPU Device Options:**

| Device | Flag | Description |
|--------|------|-------------|
| **CPU** | `cpu` | Default, works everywhere |
| **CUDA** | `cuda` | NVIDIA GPU acceleration |
| **Metal** | `metal` | Apple Silicon GPU acceleration |

**JSON Structure (`mlp_model.json`):**

```json
{
  "algorithm": "mlp",
  "trained": true,
  "model_names": ["llama-3.2-1b", "llama-3.2-3b", "mistral-7b"],
  "input_dim": 1038,
  "hidden_sizes": [256, 128],
  "layers": [
    {"type": "Linear", "weight": [[...]], "bias": [...]},
    {"type": "ReLU"},
    {"type": "BatchNorm", "weight": [...], "bias": [...], "running_mean": [...], "running_var": [...], "eps": 1e-5},
    {"type": "Dropout", "p": 0.1},
    {"type": "Linear", "weight": [[...]], "bias": [...]}
  ]
}
```

> **Reference:** MLP architecture inspired by [FusionFactory (arXiv:2507.10540)](https://arxiv.org/abs/2507.10540).

---

## Quality-Based Training Approach

All algorithms now use **response quality** and **latency** from benchmark data to make better model selections:

### How Quality is Used

| Algorithm | Quality Usage |
|-----------|--------------|
| **KNN** | Quality-weighted voting: neighbors with higher quality scores have more influence |
| **KMeans** | Quality-weighted cluster assignment: models with best quality×efficiency assigned to clusters |
| **SVM** | Quality filtering: only trains on records with quality ≥ 0.5 (high-quality responses) |
| **MLP** | Quality+latency weighted loss: CrossEntropyLoss with quality×efficiency sample weights |

### Training Record Structure

Each training record now includes:

```json
{
  "query_embedding": [0.123, -0.456, ...],  // 1038-dim feature vector (1024 + 14)
  "selected_model": "llama-3.2-3b",
  "response_quality": 0.92,                  // 0-1 quality score
  "response_latency_ns": 1500000000          // Latency in nanoseconds
}
```

### Why This Matters

**Before:** Algorithms learned to predict which model was *selected* (potentially arbitrary)

**After:** Algorithms learn to predict which model produces *high-quality responses*

This improves selection accuracy because:

1. KNN prefers neighbors that had good outcomes (quality-weighted)
2. SVM only learns from successful model selections (quality ≥ 0.5)
3. KMeans assigns clusters based on quality×efficiency, not just selection count

---

## Embedding & Feature Vector Generation

### How Query Embedding Works

When a query arrives, it goes through a **two-stage embedding process**:

```
┌─────────────────────────────────────────────────────────────────────┐
│  Query: "What is the derivative of sin(x)?"                        │
└─────────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────────┐
│  STAGE 1: Text Embedding (Candle + Qwen3)                          │
│  ─────────────────────────────────────────                          │
│  • Uses Qwen3-Embedding model via Candle (Rust ML framework)       │
│  • Converts query text to dense semantic vector                    │
│  • Output: 1024-dimensional float vector                           │
│                                                                     │
│  embedding = [0.0234, -0.1456, 0.0821, ..., 0.0512]  ← 1024 dims   │
└─────────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────────┐
│  STAGE 2: Category One-Hot Encoding                                 │
│  ──────────────────────────────────                                 │
│  • Category determined by domain classifier                        │
│  • Converted to binary vector (1 for matching category, 0 else)    │
│  • Output: 14-dimensional binary vector                            │
│                                                                     │
│  category = "math"                                                  │
│  one_hot  = [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]  ← 14 dims  │
│              ↑                                                      │
│            math                                                     │
└─────────────────────────────────────────────────────────────────────┘
                              │
                              ▼
┌─────────────────────────────────────────────────────────────────────┐
│  STAGE 3: Concatenation → Feature Vector                           │
│  ────────────────────────────────────────                           │
│                                                                     │
│  feature_vector = embedding ⊕ category_one_hot                     │
│                                                                     │
│  [0.0234, -0.1456, ..., 0.0512, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
│  └────────── 1024 dims ─────────┘└────────── 14 dims ──────────────┘
│                                                                     │
│  Total: 1038 dimensions                                             │
└─────────────────────────────────────────────────────────────────────┘
```

### Why Combine Query + Category?

| Component | Purpose |
|-----------|---------|
| **Query Embedding (1024d)** | Captures semantic meaning - "what the query is about" |
| **Category One-Hot (14d)** | Explicit domain signal - helps algorithms learn domain-specific patterns |

**Example:** Two queries might have similar embeddings:

- "Calculate the integral of x²" → math
- "Calculate the trajectory of a projectile" → physics

The category one-hot helps the model distinguish that **different LLMs** may excel in different domains, even for semantically similar queries.

### Embedding Models

**Training (Python)** uses sentence-transformers:

```python
# embeddings.py
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Alibaba-NLP/gte-Qwen2-1.5B-instruct")
embedding = model.encode(query)  # Qwen3/GTE family embedding
```

| Model | Library | Dim | Use Case |
|-------|---------|-----|----------|
| **qwen3** (default) | sentence-transformers | 1024 | Qwen3-Embedding-0.6B, matches router |
| **gte** | sentence-transformers | 768 | Original GTE |
| **mpnet** | sentence-transformers | 768 | High quality alternative |
| **e5** | sentence-transformers | 768 | Multilingual |
| **bge** | sentence-transformers | 768 | Chinese/English |

**Inference (Go/Rust)** uses Candle with Qwen3:

```go
// At runtime, router uses candle_binding
import candle_binding "github.com/vllm-project/semantic-router/candle-binding"

embedding, err := candle_binding.GetEmbeddingBatched(text, "qwen3", 1024)
```

> **Note:** Training and inference use the same Qwen3-Embedding-0.6B model (1024-dim).
> Both Python (sentence-transformers) and Rust (Candle) load the same model weights.

### Category Mapping

```go
// trainer.go - VSRCategories (MUST match Python training)
var VSRCategories = []string{
    "math",             // 0
    "physics",          // 1
    "chemistry",        // 2
    "biology",          // 3
    "computer science", // 4  (note: space, not underscore)
    "history",          // 5
    "economics",        // 6
    "business",         // 7
    "law",              // 8
    "health",           // 9
    "psychology",       // 10
    "philosophy",       // 11
    "other",            // 12
    "unknown",          // 13
}
```

### Complete Feature Vector Structure

| Index | Content | Description |
|-------|---------|-------------|
| 0-1023 | `embedding[0:1024]` | Qwen3 semantic embedding |
| 1024 | `category == "math"` | 1.0 if math, else 0.0 |
| 1025 | `category == "physics"` | 1.0 if physics, else 0.0 |
| 1026 | `category == "chemistry"` | 1.0 if chemistry, else 0.0 |
| 1027 | `category == "biology"` | 1.0 if biology, else 0.0 |
| 1028 | `category == "computer science"` | 1.0 if CS, else 0.0 |
| 1029 | `category == "history"` | 1.0 if history, else 0.0 |
| 1030 | `category == "economics"` | 1.0 if economics, else 0.0 |
| 1031 | `category == "business"` | 1.0 if business, else 0.0 |
| 1032 | `category == "law"` | 1.0 if law, else 0.0 |
| 1033 | `category == "health"` | 1.0 if health, else 0.0 |
| 1034 | `category == "psychology"` | 1.0 if psychology, else 0.0 |
| 1035 | `category == "philosophy"` | 1.0 if philosophy, else 0.0 |
| 1036 | `category == "other"` | 1.0 if other, else 0.0 |
| 1037 | `category == "unknown"` | 1.0 if unknown, else 0.0 |

---

## Training Data Format

> **⚠️ IMPORTANT: Data Files Are Deployment-Specific**
>
> Training data files are **NOT included in this repository**. Each deployment should generate its own training data based on:
>
> - Your specific LLM models (Ollama, OpenAI, vLLM, etc.)
> - Your query patterns and use cases
> - Your performance benchmarks
>
> The files referenced below (`training_data_with_category.jsonl`, `benchmark_training_data.jsonl`, `trained_models/*.json`) are generated by YOU during the training process.

### What You Provide vs. What System Generates

| Field | Source | File |
|-------|--------|------|
| `query` | ✍️ **You provide** | Your input file |
| `category` | ✍️ **You provide** (or use script) | Your input file |
| `ground_truth` | ✍️ **You provide** (recommended) | Your input file |
| `model_name` | 🤖 **System generates** | `benchmark_training_data.jsonl` |
| `response` | 🤖 **System generates** (LLM output) | `benchmark_training_data.jsonl` |
| `response_time` | 🤖 **System generates** (latency) | `benchmark_training_data.jsonl` |
| `performance` | 🤖 **System generates** (quality) | `benchmark_training_data.jsonl` |

**Data Flow:**

```
┌─────────────────────────────────────┐
│  YOUR INPUT FILE                    │
│  (training_data_with_category.jsonl)│
├─────────────────────────────────────┤
│  query: "What is 2+2?"              │
│  category: "math"                   │
│  ground_truth: "4"                  │
└──────────────┬──────────────────────┘
               │
               ▼  vsr train --benchmark
┌─────────────────────────────────────┐
│  BENCHMARK RUNNER                   │
│  • Sends query to each LLM          │
│  • Measures response time (latency) │
│  • Captures model response          │
│  • Calculates quality score         │
└──────────────┬──────────────────────┘
               │
               ▼
┌─────────────────────────────────────┐
│  GENERATED: benchmark_training_data │
├─────────────────────────────────────┤
│  query: "What is 2+2?"              │  ← from your input
│  category: "math"                   │  ← from your input
│  ground_truth: "4"                  │  ← from your input
│  model_name: "llama-3.2-1b"         │  ← GENERATED
│  response: "The answer is 4"        │  ← GENERATED (LLM output)
│  response_time: 0.523               │  ← GENERATED (latency)
│  performance: 1.0                   │  ← GENERATED (quality)
└─────────────────────────────────────┘
```

### Step 1: Prepare Your Training Data

Your training data file should be a JSONL file with one JSON object per line. This is YOUR dataset - name it anything (e.g., `my_training_data.jsonl`, `queries.jsonl`).

**Required Fields:**

| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `query` | string | Yes | The input query/prompt text |
| `category` | string | Yes | Query category (see categories below) |
| `ground_truth` | string | ⚡ Recommended | Expected correct answer (for quality scoring) |
| `task_name` | string | Optional | Task identifier (e.g., "mmlu", "gsm8k") |

**Supported Categories (VSRCategories - 14 domains):**

> **⚠️ Important:** Use exact names with **spaces** (not underscores).

```
biology, business, chemistry, computer science, economics, engineering,
health, history, law, math, other, philosophy, physics, psychology
```

| Category | Example Queries |
|----------|-----------------|
| `math` | Calculus, algebra, geometry, statistics |
| `computer science` | Programming, algorithms, data structures |
| `physics` | Mechanics, relativity, quantum physics |
| `chemistry` | Chemical reactions, molecular structures |
| `biology` | Genetics, ecology, cell biology |
| `health` | Medical conditions, anatomy, treatments |
| `engineering` | Civil, mechanical, electrical engineering |
| `economics` | Markets, finance, economic theory |
| `business` | Management, marketing, strategy |
| `history` | Historical events, periods, figures |
| `law` | Legal matters, regulations, cases |
| `philosophy` | Ethics, logic, metaphysics |
| `psychology` | Behavior, cognition, mental health |
| `other` | General knowledge (catch-all) |

**Minimal Format (query + category only):**

```jsonl
{"query": "What is the integral of x^2?", "category": "math"}
{"query": "Explain quantum entanglement", "category": "physics"}
{"query": "Write a Python function to sort a list", "category": "computer science"}
```

**Recommended Format (with ground_truth for quality scoring):**

```jsonl
{"query": "What is 2 + 2?", "category": "math", "ground_truth": "4"}
{"query": "What is the capital of France?", "category": "history", "ground_truth": "Paris"}
{"query": "Write a function to reverse a string in Python", "category": "computer science", "ground_truth": "def reverse(s): return s[::-1]"}
```

**Full Format (MMLU-style with multiple choice):**

```jsonl
{
  "task_name": "mmlu",
  "query": "What is the derivative of x^3?\nA) x^2\nB) 3x^2\nC) 3x\nD) x^3",
  "category": "math",
  "ground_truth": "B"
}
```

### Step 2: Add Categories (if needed)

If your data doesn't have categories, use the provided script:

```bash
cd src/training/model_selection/ml_model_selection
python add_category_to_training_data.py \
  --input your_training_data.jsonl \
  --output training_data_with_category.jsonl
```

### Step 3: Generated Benchmark Data

After running `vsr train --benchmark`, the system generates `benchmark_training_data.jsonl` with performance data for each model:

```jsonl
{
  "task_name": "mmlu",
  "query": "What is the integral of x^2?",
  "category": "math",
  "model_name": "llama-3.2-1b",
  "response": "The integral of x^2 is x^3/3 + C",
  "ground_truth": "x^3/3 + C",
  "response_time": 1.234,
  "performance": 0.85,
  "embedding_id": 1
}
{
  "task_name": "mmlu",
  "query": "What is the integral of x^2?",
  "category": "math",
  "model_name": "llama-3.2-3b",
  "response": "To find the integral of x^2, we use the power rule...",
  "ground_truth": "x^3/3 + C",
  "response_time": 2.567,
  "performance": 0.92,
  "embedding_id": 1
}
```

**Benchmark Data Fields:**

| Field | Type | Description |
|-------|------|-------------|
| `query` | string | Original query text |
| `category` | string | Query category |
| `model_name` | string | LLM that processed this query |
| `response` | string | Model's response |
| `ground_truth` | string | Expected answer (if provided) |
| `response_time` | float | Latency in seconds |
| `performance` | float | Quality score 0-1 |
| `embedding_id` | int | Groups responses for same query |

### Quality Scoring

The system calculates quality scores based on `ground_truth`:

| Query Type | Scoring Method |
|------------|----------------|
| Multiple Choice (A/B/C/D) | Exact match of selected option |
| Numeric/Math | Parse and compare numbers |
| Text/Code | F1 score between response and ground_truth |
| No ground_truth | Uses `performance` field if available, else 0.5 |

---

## Configuration

### Supported LLM Providers

> **Note:** This documentation uses **Ollama** as an example for local development and testing. 
> However, the model selection system supports **any LLM provider** with proper authentication.

| Provider | Type | Authentication | Use Case |
|----------|------|----------------|----------|
| **Ollama** | `ollama` | None (local) | Local development, testing |
| **vLLM** | `vllm` | None or API key | Self-hosted production |
| **OpenAI** | `openai` | API key | GPT-4, GPT-3.5 |
| **HuggingFace** | `huggingface` | HF Token | Inference API |
| **NVIDIA NIM** | `nvidia` | API key | Enterprise deployment |
| **OpenRouter** | `openrouter` | API key | Multi-provider gateway |

### Example 1: Ollama (Local - No Auth)

```yaml
providers:
  models:
    - name: "llama-3.2-1b"
      provider_model_id: "llama3.2:1b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
          weight: 1
    - name: "llama-3.2-3b"
      provider_model_id: "llama3.2:3b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
          weight: 1

routing:
  modelCards:
    - name: "llama-3.2-1b"
    - name: "llama-3.2-3b"
```

### Example 2: HuggingFace (Token Auth)

```yaml
providers:
  models:
    - name: "mistral-7b-hf"
      provider_model_id: "mistralai/Mistral-7B-Instruct-v0.2"
      external_model_ids:
        huggingface: "mistralai/Mistral-7B-Instruct-v0.2"
      backend_refs:
        - name: "huggingface"
          base_url: "https://api-inference.huggingface.co"
          provider: "huggingface"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "HF_TOKEN"
          weight: 1

routing:
  modelCards:
    - name: "mistral-7b-hf"
```

### Example 3: OpenAI (API Key)

```yaml
providers:
  models:
    - name: "gpt-4"
      provider_model_id: "gpt-4-turbo"
      external_model_ids:
        openai: "gpt-4-turbo"
      backend_refs:
        - name: "openai"
          base_url: "https://api.openai.com/v1"
          provider: "openai"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "OPENAI_API_KEY"
          weight: 1
    - name: "gpt-3.5"
      provider_model_id: "gpt-3.5-turbo"
      external_model_ids:
        openai: "gpt-3.5-turbo"
      backend_refs:
        - name: "openai"
          base_url: "https://api.openai.com/v1"
          provider: "openai"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "OPENAI_API_KEY"
          weight: 1

routing:
  modelCards:
    - name: "gpt-4"
    - name: "gpt-3.5"
```

### Example 4: NVIDIA NIM (Enterprise)

```yaml
providers:
  models:
    - name: "llama-3.1-70b-nim"
      provider_model_id: "meta/llama-3.1-70b-instruct"
      external_model_ids:
        nvidia: "meta/llama-3.1-70b-instruct"
      backend_refs:
        - name: "nvidia-nim"
          base_url: "https://integrate.api.nvidia.com/v1"
          provider: "nvidia"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "NVIDIA_API_KEY"
          weight: 1

routing:
  modelCards:
    - name: "llama-3.1-70b-nim"
```

### Example 5: Mixed Providers (Production)

You can mix multiple providers in a single configuration:

```yaml
providers:
  models:
    - name: "llama-3.2-1b"
      provider_model_id: "llama3.2:1b"
      external_model_ids:
        ollama: "llama3.2:1b"
      backend_refs:
        - name: "ollama-local"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
          weight: 1
    - name: "gpt-4"
      provider_model_id: "gpt-4-turbo"
      external_model_ids:
        openai: "gpt-4-turbo"
      backend_refs:
        - name: "openai"
          base_url: "https://api.openai.com/v1"
          provider: "openai"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "OPENAI_API_KEY"
          weight: 1
    - name: "codellama-34b"
      provider_model_id: "codellama/CodeLlama-34b-Instruct-hf"
      external_model_ids:
        huggingface: "codellama/CodeLlama-34b-Instruct-hf"
      backend_refs:
        - name: "huggingface"
          base_url: "https://api-inference.huggingface.co"
          provider: "huggingface"
          auth_header: "Authorization"
          auth_prefix: "Bearer"
          api_key_env: "HF_TOKEN"
          weight: 1

routing:
  modelCards:
    - name: "llama-3.2-1b"
    - name: "gpt-4"
    - name: "codellama-34b"
```

### Authentication Options

| Method | Config Field | Example |
|--------|--------------|---------|
| **Backend credential by env** | `api_key_env` in `backend_refs[]` | `api_key_env: "OPENAI_API_KEY"` |
| **Inline backend credential** | `api_key` in `backend_refs[]` | `api_key: "sk-..."` |
| **Provider auth header** | `auth_header` + `auth_prefix` | `Authorization` + `Bearer` |
| **Environment variable** | `${VAR_NAME}` syntax | Auto-resolved at runtime |
| **Direct value** | Plain string | `api_key: "sk-..."` (not recommended) |

### Full Example: 4 Models with Ollama (This Documentation)

```yaml
global:
  router:
    model_selection:
      enabled: true
      ml:
        models_path: ".cache/ml-models"
        embedding_dim: 1024

providers:
  models:
    - name: "llama-3.2-1b"
      provider_model_id: "llama3.2:1b"
      external_model_ids:
        ollama: "llama3.2:1b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
    - name: "llama-3.2-3b"
      provider_model_id: "llama3.2:3b"
      external_model_ids:
        ollama: "llama3.2:3b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
    - name: "codellama-7b"
      provider_model_id: "codellama:7b"
      external_model_ids:
        ollama: "codellama:7b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"
    - name: "mistral-7b"
      provider_model_id: "mistral:7b"
      external_model_ids:
        ollama: "mistral:7b"
      backend_refs:
        - name: "ollama"
          endpoint: "localhost:11434"
          protocol: "http"
          type: "ollama"

routing:
  modelCards:
    - name: "llama-3.2-1b"
    - name: "llama-3.2-3b"
    - name: "codellama-7b"
    - name: "mistral-7b"
```

### Per-Decision Algorithm Configuration

```yaml
routing:
  decisions:
    - name: "math_decision"
      description: "Mathematics and quantitative reasoning"
      priority: 100
      rules:
        operator: "AND"
        conditions:
          - type: "domain"
            name: "math"
      algorithm:
        type: "knn"
        knn:
          k: 5
      modelRefs:
        - model: "llama-3.2-1b"
          use_reasoning: false
        - model: "llama-3.2-3b"
          use_reasoning: false
        - model: "codellama-7b"
          use_reasoning: false
        - model: "mistral-7b"
          use_reasoning: false
```

---

## Training Commands (Python)

> **Training is now done in Python** using scikit-learn for consistency with other training scripts.
> See `src/training/model_selection/ml_model_selection/` for the full training pipeline.

### Option 1: Download Pretrained Models from HuggingFace (Recommended)

```bash
pip install huggingface-hub
cd src/training/model_selection/ml_model_selection

# Download pretrained models
python download_model.py \
  --output-dir ../../../.cache/ml-models \
  --repo-id abdallah1008/semantic-router-ml-models
```

### Option 2: Train Models Locally

> **Note:** This requires your own benchmark data generated with `benchmark.py`.
> The output directory is local and NOT committed to the repository.

```bash
cd src/training/model_selection/ml_model_selection
pip install -r requirements.txt

# Train from your own benchmark data
python train.py \
  --data-file your_benchmark_data.jsonl \
  --output-dir ./trained_models \
  --embedding-model qwen3
```

### Training Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `--data-file` | (required) | Path to JSONL benchmark data |
| `--output-dir` | `models/` | Where to save trained model JSON files |
| `--embedding-model` | `qwen3` | Embedding model: qwen3, gte, mpnet, e5, bge |
| `--algorithm` | `all` | Train specific algorithm: all, knn, kmeans, svm, mlp |
| `--knn-k` | 5 | KNN: Number of neighbors |
| `--kmeans-clusters` | 8 | KMeans: Number of clusters |
| `--svm-kernel` | `rbf` | SVM kernel: rbf, linear |
| `--svm-gamma` | 1.0 | SVM gamma for RBF kernel |
| `--mlp-hidden-sizes` | `256,128` | MLP: Hidden layer sizes |
| `--mlp-epochs` | 100 | MLP: Training epochs |
| `--mlp-learning-rate` | 0.001 | MLP: Learning rate |
| `--mlp-dropout` | 0.1 | MLP: Dropout rate |
| `--quality-weight` | 0.9 | Quality vs speed weight (0=speed, 1=quality) |
| `--batch-size` | 32 | Batch size for embedding generation |
| `--device` | `cpu` | Device: cpu, cuda, mps |

### Upload Models to HuggingFace

After training, share models:

```bash
export HF_TOKEN=your_token
python upload_model.py \
  --model-dir ./trained_models \
  --repo-id abdallah1008/semantic-router-ml-models
```

### Preparing Benchmark Data

Benchmark data should be JSONL format with fields:

```json
{
  "query": "What is the derivative of sin(x)?",
  "category": "math",
  "model_name": "llama-3.2-3b",
  "performance": 0.85,
  "response_time": 1.234
}
```

---

## Inference Flow

When a new query arrives:

```
1. Query: "What is the derivative of sin(x)?"
                    │
                    ▼
2. Category Classification
   └── category: "math"
                    │
                    ▼
3. Generate Embedding (Candle/Qwen3)
   └── embedding: [0.123, -0.456, ..., 0.789]  (1024 dims)
                    │
                    ▼
4. Build Feature Vector
   └── embedding + one_hot(category) = 1038 dims
                    │
                    ▼
5. Load Algorithm (based on decision config)
   └── KNN selector from knn_model.json
                    │
                    ▼
6. Run Inference
   └── Find 5 nearest neighbors
   └── Vote: 3× llama-3.2-3b, 2× llama-3.2-1b
                    │
                    ▼
7. Return Best Model
   └── "llama-3.2-3b"
```

---

## Testing

### Test Files

| File | Purpose |
|------|---------|
| `selector_test.go` | Unit tests for algorithm implementations |

### Run Tests

```bash
# On WSL (required for real Qwen3 embeddings)
export LD_LIBRARY_PATH=/mnt/c/vllmproject/semantic-router/candle-binding/target/release:$LD_LIBRARY_PATH
cd /mnt/c/vllmproject/semantic-router/src/semantic-router/pkg/modelselection

# Run all tests
go test -v . -timeout 30m
```

### Run Unit Tests

```bash
cd src/semantic-router
go test -v ./pkg/modelselection/... -run TestLoadPretrainedSelectors
go test -v ./pkg/modelselection/... -run TestModelSelectionWithVariedQueries
```

### Test Methodology

The tests in `selector_test.go` follow this flow:

```
1. Load pre-trained models from HuggingFace
   └── knn_model.json, kmeans_model.json, svm_model.json
                    │
                    ▼
2. Generate embeddings for test queries using Qwen3
                    │
                    ▼
3. Run inference with each algorithm
   └── Track which model is selected per query
                    │
                    ▼
4. Analyze model distribution
   └── Verify algorithms select appropriate models
```

### Test Results

All 3 algorithms successfully:

1. Load from JSON files
2. Mark as `trained: true`
3. Select **multiple different models** based on query characteristics
4. Use the quality × efficiency formula correctly

**Example test output (WSL with real Qwen3):**

```
=== RUN   TestGeneralizationAllCategories/knn
  [biology] 'Explain the structure of a eukaryotic cell' -> llama-3.2-3b
  [chemistry] 'Describe the structure of an atom' -> mistral-7b
  [computer science] 'Write a breadth-first search algorithm' -> llama-3.2-1b
  [math] 'Find the derivative of sin(x) * e^x' -> llama-3.2-1b
  [physics] 'Explain Maxwell's equations for electromagnetism' -> mistral-7b
  Overall model distribution: map[codellama-7b:1 llama-3.2-1b:11 llama-3.2-3b:8 mistral-7b:10]
  knn selected 4 different models across categories
```

---

## Model Loading (Code Reference)

### Loading a Selector

```go
// persistence.go
func LoadSelectorFromJSON(algorithm, modelsPath string, refs []config.ModelRef) (ModelSelector, error) {
    filePath := filepath.Join(modelsPath, algorithm+"_model.json")
    data, err := os.ReadFile(filePath)
    if err != nil {
        return nil, err
    }
    
    var modelData ModelData
    json.Unmarshal(data, &modelData)
    
    selector := createSelector(algorithm)
    selector.LoadFromJSON(modelData)
    
    return selector, nil
}
```

### Selector Interface

```go
// selector.go
type ModelSelector interface {
    Select(ctx context.Context, query string, refs []config.ModelRef, 
           embedding []float64, category string) (config.ModelRef, error)
    Train(training []TrainingRecord) error
    LoadFromJSON(data ModelData) error
    SaveToJSON() (ModelData, error)
    IsTrained() bool
}
```

---

## Performance Considerations

| Algorithm | Training Speed | Inference Speed | Memory Usage |
|-----------|---------------|-----------------|--------------|
| KNN | Fast | O(n) - slower with more data | High (stores all embeddings) |
| KMeans | Medium | O(k) - fast | Low (only centroids) |
| SVM | Slow | O(sv) - depends on support vectors | Medium |

---

## Troubleshooting

### Model not trained (always selects first model)

- Check `trained: true` in JSON files
- Ensure enough training records (minimum ~20-50)
- Verify embeddings are being generated (check for Candle errors)

### Dimension mismatch errors

- Ensure embedding model outputs 1024 dimensions (Qwen3-Embedding-0.6B)
- Check category one-hot encoding produces 14 dimensions
- Total feature vector should be 1038 dimensions

### Benchmark timeout errors

- Increase Ollama timeout in config
- Ensure models are loaded (first request is slow)
- Check available RAM for larger models

### Windows: Tests fail with dimension mismatch (384 vs 1024)

On Windows, the Candle library uses a **mock implementation** that generates 384-dim hash-based 
embeddings instead of real 1024-dim Qwen3 embeddings.

**Solution:** Run tests on WSL (Windows Subsystem for Linux):

```bash
# Set library path for Candle
export LD_LIBRARY_PATH=/mnt/c/vllmproject/semantic-router/candle-binding/target/release:$LD_LIBRARY_PATH

# Run tests
cd /mnt/c/vllmproject/semantic-router/src/semantic-router/pkg/modelselection
go test -v -run 'TestGeneralizationAllCategories' . -timeout 30m
```

The mock is triggered by this build constraint in `candle-binding/semantic-router_mock.go`:

```go
//go:build windows || !cgo
```

---

## Algorithm Accuracy & Benchmark Results

### Test Methodology

The tests follow this methodology:

1. **Load pre-trained models from HuggingFace** (not re-train)
2. **Generate embeddings** for test queries using Qwen3
3. **Test model selection** with different algorithms

This ensures we're testing the model's ability to generalize to unseen queries.

### Test Results Summary

#### TestGeneralizationAllCategories (30 NEW queries, 14 categories)

| Algorithm | Implementation | Linfa Crate | Models Selected | Model Distribution | Success |
|-----------|----------------|-------------|-----------------|-------------------|---------|
| **KNN** | Rust | `linfa-nn` | **4** | codellama:4, llama-1b:13, llama-3b:4, mistral:9 | 30/30 |
| **KMeans** | Rust | `linfa-clustering` | **1** | mistral:30 (cluster-based) | 30/30 |
| **SVM** | Rust | `linfa-svm` | **4** | codellama:4, llama-1b:13, llama-3b:3, mistral:10 | 30/30 |

#### Model Selection Diversity Chart

```
KNN     █████░░░░░░░░░░░░░░░░░░░░░░░ codellama (13%)
        █████████████████░░░░░░░░░░░ llama-1b (43%)
        █████░░░░░░░░░░░░░░░░░░░░░░░ llama-3b (13%)
        ██████████░░░░░░░░░░░░░░░░░░ mistral (30%)

KMeans  ░░░░░░░░░░░░░░░░░░░░░░░░░░░░ codellama (0%)
        ░░░░░░░░░░░░░░░░░░░░░░░░░░░░ llama-1b (0%)
        ░░░░░░░░░░░░░░░░░░░░░░░░░░░░ llama-3b (0%)
        ████████████████████████████ mistral (100%)

SVM     █████░░░░░░░░░░░░░░░░░░░░░░░ codellama (13%)
        █████████████████░░░░░░░░░░░ llama-1b (43%)
        ████░░░░░░░░░░░░░░░░░░░░░░░░ llama-3b (10%)
        ███████████░░░░░░░░░░░░░░░░░ mistral (33%)
```

#### Key Insights

| Metric | Best Algorithm |
|--------|----------------|
| **Most Diverse** | KNN, SVM (use all 4 models) |
| **Fastest Inference** | KMeans (< 1s for 30 queries) |
| **Best for Quality** | KNN (Ball Tree nearest neighbor) |
| **Cluster-based** | KMeans (routes to dominant cluster) |
| **Decision Boundaries** | SVM (RBF kernel) |

### Benchmark Results (WSL with Real Qwen3)

Validation using 109 test queries with 4 models (codellama-7b, llama-3.2-1b, llama-3.2-3b, mistral-7b):

**Validation Results:**

| Strategy | Avg Quality | Improvement over Random |
|----------|-------------|------------------------|
| **Oracle (best)** | 0.495 | - |
| **MLP Selection** | 0.286 | **+47.1%** |
| **KMEANS Selection** | 0.252 | **+29.9%** |
| **SVM Selection** | 0.233 | **+20.0%** |
| **KNN Selection** | 0.196 | **+1.0%** |
| Random Selection | 0.194 | baseline |

**Algorithm Performance Summary:**

| Algorithm | Quality | Best Model % | Notes |
|-----------|---------|--------------|-------|
| **MLP** | 0.286 | 20.2% | Best overall, 1.5x better than random (GPU-accelerated) |
| **KMEANS** | 0.252 | 23.9% | Strong performer, 1.7x better than random |
| **SVM** | 0.233 | 14.7% | Good quality, 1.1x better than random |
| **KNN** | 0.196 | 13.8% | Moderate improvement, 1.0x vs random |

All algorithms correctly select **multiple different models** based on query characteristics.

### RouteLLM Alignment

The implementation is aligned with [RouteLLM (arXiv:2406.18665)](https://arxiv.org/abs/2406.18665):

| RouteLLM Concept | Our Implementation |
|------------------|-------------------|
| Routes between stronger/weaker LLMs | Selects from multiple LLMs |
| Optimizes cost vs quality | Formula: `quality × efficiency` |
| Transfer learning capability | Pre-trained models generalize to new queries |
