| --- |
| 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). |
