README.md

---
license: apache-2.0
language:
- en
tags:
- prompt-routing
- category-classifier
- deberta-v3
- llm-router
- cost-optimization
datasets:
- RowRed/ComplexityRouter
- OpenAssistant/oasst2
base_model:
- microsoft/deberta-v3-base
---

# CategoryRouter: A Category‑based LLM Router

**CategoryRouter** is a lightweight prompt category classifier finetuned from **microsoft/deberta-v3-base**. It assigns prompts to one of **6 categories**, making it useful for routing queries to the appropriate specialist LLM or service.

## Model Details

### Model Description

- **Model type:** Text Classification (multi‑class, 6 classes)
- **Language:** English
- **Backbone:** microsoft/deberta-v3-base
- **License:** Apache‑2.0
- **Finetuned from model:** microsoft/deberta-v3-base
- **Training data:** OASST2 + synthetic augmentations + manually created prompts  
  Labels generated by **Qwen3.5‑4B** (non‑thinking mode).

### Model Sources

- **Dataset repository:** https://huggingface.co/datasets/RowRed/ComplexityRouter

## Categories

| Label | Category | Example Prompt |
|-------|----------|----------------|
| 0 | Coding | “Write a Python function to merge two sorted lists.” |
| 1 | Information | “What is the normal force? Explain with an example.” |
| 2 | Guidance | “How do I change a car tire?” |
| 3 | Media Generation | “Generate an image prompt for a sunset.” |
| 4 | Writing | “Write a short story about a robot learning to paint.” |
| 5 | Other | “Invent a new holiday.” |

> **Note:** The original ComplexityRouter dataset had 7 categories; “Math” prompts were merged into “Coding” to retain only 6 distinct categories.

## Uses

### Direct Use

Route prompts to the appropriate LLM backend:

| Category | Suggested Backend |
|----------|-------------------|
| Coding | Code‑specialized model |
| Information | General‑knowledge model |
| Guidance | Instruction‑tuned model |
| Media Generation | Multimodal / image‑generation model |
| Writing | Creative writing model |
| Other | Catch‑all (generic model) |

### Out‑of‑Scope Use

- Multi‑turn conversations (single prompts only).
- Non‑English prompts.
- Prompts requiring image or multimodal understanding.

## Bias, Risks, and Limitations

- Training data is synthetic and may not represent all real‑world prompt distributions.
- Categories with low support (Media Generation, Writing) have lower per‑class F1 scores – boundary cases are inherently ambiguous.
- The model may struggle with very domain‑specific technical jargon.
- The “Other” category is a catch‑all; many “Information” prompts leak into it and vice‑versa.
- Performance may degrade on prompts very different from the training distribution.

## Notice
This is my first attempt making a widespread finetune ( just the part 2 version of it :) ). There are probably lots of issues, but thought the idea was sound. I might make a second (hopefully better) version eventually, but am not sure where to get lots of high-quality open source data.

## Training Details

### Training Data

| Split       | Samples | Source File          | Notes |
|-------------|---------|----------------------|-------|
| Training    | 2,800   | TRAINING.jsonl       | Used for model training |
| Validation  | 600     | TRAINING.jsonl       | Used for early stopping / hyperparameter tuning |
| Test (internal) | 600 | TRAINING.jsonl      | Used for in‑distribution evaluation |
| Test (held‑out) | 400 | TEST.jsonl   | Fully independent test set (reported results) |

**Total unique prompts:** 4,400

Class distribution (training):
Coding: 508 (18.1%) • Information: 937 (33.5%) • Guidance: 609 (21.8%) • Media Generation: 111 (4.0%) • Writing: 88 (3.1%) • Other: 547 (19.6%)

### Training Procedure

- Hardware: NVIDIA T4 (16 GB VRAM, Google Colab)
- Framework: PyTorch 2.11 + Hugging Face Transformers
- Optimizer: AdamW (lr=2e-5, weight_decay=0.01)
- Scheduler: Linear warmup (10% of steps) → linear decay
- Loss: Weighted Cross‑Entropy (sqrt‑scaled class weights) + label smoothing (0.15)
- Batch size: 16 (effective 32 with gradient accumulation)
- Epochs: 7 (early stopping patience = 3, best epoch = 7)
- Training time: ~16 minutes
- Class balancing: sqrt‑scaled class weights + weighted random sampler

## Evaluation Results

### Internal Test (600 held‑out samples from training split)

| Metric | Value |
|--------|-------|
| Exact Match Accuracy | 67.0% |
| Macro F1 | 0.617 |
| Weighted F1 | 0.665 |

**Per‑Class Performance (internal test, 600 samples)**

| Category | Precision | Recall | F1 | Support |
|----------|-----------|--------|----|---------|
| Coding | 0.714 | 0.688 | 0.701 | 109 |
| Information | 0.740 | 0.766 | 0.753 | 201 |
| Guidance | 0.707 | 0.800 | 0.751 | 130 |
| Media Generation | 0.619 | 0.542 | 0.578 | 24 |
| Writing | 0.571 | 0.421 | 0.485 | 19 |
| Other | 0.457 | 0.410 | 0.432 | 117 |

**Confusion Matrix (internal test)**

```
           Pred  Pred  Pred  Pred  Pred  Pred
           Cod   Inf   Gui   Med   Wri   Oth
True Cod    75     8    12     1     1    12
True Inf     5   154    13     0     1    28
True Gui     9     6   104     0     1    10
True Med     3     2     0    13     2     4
True Wri     0     0     3     5     8     3
True Oth    13    38    15     2     1    48
```

### Held‑Out Test (400 independent samples)

| Metric | Value |
|--------|-------|
| Exact Match Accuracy | 76.5% |
| Macro F1 | – (per‑category below) |

**Per‑Category Breakdown**

| Category | Exact Count | Exact Accuracy |
|----------|-------------|----------------|
| Coding | 35/44 | 79.5% |
| Information | 184/229 | 80.3% |
| Guidance | 55/63 | 87.3% |
| Media Generation | 8/14 | 57.1% |
| Writing | 9/20 | 45.0% |
| Other | 15/30 | 50.0% |

## How to Get Started with the Model

```python
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn as nn

class CategoryRouter(nn.Module):
    def __init__(self, backbone="microsoft/deberta-v3-base", num_labels=6):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(backbone)
        hidden_size = self.backbone.config.hidden_size
        self.classifier = nn.Sequential(
            nn.Dropout(0.1),
            nn.Linear(hidden_size, 256),
            nn.GELU(),
            nn.Dropout(0.1),
            nn.Linear(256, num_labels),
        )

    def forward(self, input_ids, attention_mask):
        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]
        return self.classifier(cls_output)

# Load
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("RowRed/CategoryRouter")
model = CategoryRouter()
model.load_state_dict(
    torch.load("pytorch_model.bin", map_location=device),
    strict=False
)
model.to(device)
model.eval()

# Predict
prompts = ["Write a Python function", "What is the normal force?"]
encoded = tokenizer(prompts, padding=True, truncation=True, return_tensors="pt").to(device)
with torch.no_grad():
    logits = model(encoded["input_ids"], encoded["attention_mask"])
    probs = torch.softmax(logits, dim=-1)
    predictions = torch.argmax(probs, dim=-1)

category_map = ["Coding", "Information", "Guidance", "Media Generation", "Writing", "Other"]
for prompt, idx in zip(prompts, predictions):
    print(f"Category: {category_map[idx.item()]} – {prompt}")
```

## Citation
If you use this model, please cite:

```bibtex
@software{CategoryRouter,
  author = {RowRed},
  title = {CategoryRouter},
  year = {2026},
  url = {https://huggingface.co/RowRed/CategoryRouter}
}
```

Additionally, acknowledge the base dataset and labeling model:

```bibtex
@dataset{oasst2,
  author = {OpenAssistant Contributors},
  title = {Open Assistant Conversations Dataset Release 2},
  year = {2023},
  url = {https://huggingface.co/datasets/OpenAssistant/oasst2}
}

@software{qwen3.5-4b,
  author = {Qwen Team},
  title = {Qwen3.5-4B},
  year = {2026},
  url = {https://huggingface.co/Qwen/Qwen3.5-4B}
}
```

## License
This model is released under Apache‑2.0.
The backbone (microsoft/deberta-v3-base) is MIT‑licensed.
The training dataset is derived from OASST2 (Apache‑2.0) and Qwen3.5‑4B outputs (Apache‑2.0).