File size: 7,264 Bytes
f285c7f 74c0320 f285c7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 | ---
license: apache-2.0
language:
- en
tags:
- prompt-routing
- difficulty-classifier
- deberta-v3
- llm-router
datasets:
- RowRed/prompts-24000-en
base_model:
- microsoft/deberta-v3-xsmall
---
# DifficultyRouter: A Lightweight 3‑Tier Prompt Difficulty Router
DifficultyRouter is a successor to [ComplexityRouter](https://huggingface.co/RowRed/ComplexityRouter), specifically designed for **cost optimization** in LLM routing. It is finetuned from **microsoft/deberta-v3-xsmall** (≈22M params, ~8x smaller than the base model used previously) on a 20,000‑prompt dataset (2,000 L0 + 6,000 L1 → Tier 0; 6,000 L2 → Tier 1; 6,000 L3 → Tier 2).
It classifies prompts into **3 difficulty tiers** with a single classification head.
## Model Details
### Model Description
- **Model type:** Text Classification (3‑tier, single head)
- **Language:** English
- **License:** Apache‑2.0
- **Finetuned from model:** microsoft/deberta-v3-xsmall
- **Training data:** `RowRed/prompts-24000-en` (L0 downsampled to 2,000, L1/L2/L3 kept fully)
### Model Sources
- **Dataset repository:** `https://huggingface.co/datasets/RowRed/prompts-24000-en`
- **Old model:** `RowRed/ComplexityRouter`
## Uses
### Direct Use
Route prompts to appropriate LLM tiers based on predicted difficulty:
| Tier | Meaning | Original Levels | Suggested LLM Tier |
|------|---------|-----------------|--------------------|
| 0 (Easy) | Simple lookups, basic Q&A, light reasoning | L0 + L1 | Fast/cheap model |
| 1 (Moderate) | Complex reasoning, deep domain knowledge | L2 | Standard model |
| 2 (Complex) | Very complex reasoning, niche expertise, edge cases | L3 | Frontier model |
### Out‑of‑Scope Use
- Multi‑turn conversation routing (single prompts only).
- Non‑English prompts (training data is English‑only).
- Prompts requiring image or multimodal understanding.
- 4‑level classification (use the old ComplexityRouter for 4 classes).
## Bias, Risks, and Limitations
- Training data includes synthetic augmentation; distribution may not match all real‑world prompt patterns.
- Tier 0 (merged L0+L1) has inherent ambiguity—some "trivial" and "simple" prompts are hard to distinguish from "moderate".
- DeBERTa‑v3‑xsmall has a smaller representation capacity than the base model, so it may miss very subtle difficulty cues in niche technical domains.
### ⚠️ Not Production‑Ready
This model is a research prototype.
- Accuracy is ~63.4%, meaning ~4 out of 10 prompts will be misrouted.
- Adjacent accuracy of 87.1% means 1 in 8 prompts will be sent to a tier that is still off by one level, leading to noticeable latency/cost misses.
- The model has not been stress‑tested on real‑world, messy, multi‑domain prompts. It was trained on synthetic augmentations.
## Training Details
### Training Data
- **Source:** `RowRed/prompts-24000-en` (approximately 24,000 raw prompts)
- **Used:** 20,000 prompts (L0 downsampled to 2,000; L1, L2, L3 kept at 6,000 each)
- **Preprocessing:** Original difficulty levels (0–3) are mapped to 3 tiers: `0+1 -> 0`, `2 -> 1`, `3 -> 2`.
- **L0 downsampling:** To combat overfitting, exactly `2,000` L0 samples were randomly sampled (config flag `l0_sample_size=2000`, locked).
- **Split:** 70% train / 18% validation / 12% held-out test (stratified).
### Training Procedure
- **Hardware:** NVIDIA T4 (16 GB VRAM, Google Colab)
- **Framework:** PyTorch + Hugging Face Transformers
- **Optimizer:** AdamW (lr=3e-5, weight_decay=0.1)
- **Scheduler:** Linear warmup (6% steps) → linear decay
- **Loss:** Weighted Cross‑Entropy with label smoothing=0.1
- **Batch size:** 16 (effective 32 with gradient accumulation)
- **Max sequence length:** 256 tokens
- **Epochs:** 12 max (Early stopping patience = 4 on F1-macro; stopped at epoch 7, best checkpoint at epoch 3)
- **Class balancing:** WeightedRandomSampler only; class weights in loss removed to avoid double-counting
- **Head:** 256-dim, dropout=0.1
- **Precision:** FP16 mixed precision
## Evaluation Results
Reported on the **held‑out test set** (~12% of 20k prompts) using the **best model checkpoint** (epoch 3, selected via validation F1-macro):
| Metric | Value |
|---------------------|--------|
| Exact Match Accuracy | 63.36% |
| Adjacent (±1) Accuracy | 87.14% |
| F1 Macro | 0.6375 |
| F1 Weighted | 0.6314 |
Training loss continued to decrease but validation metrics peaked at epoch 3; early stopping correctly caught the onset of overfitting.
## How to Get Started with the Model
The model uses a single-head architecture and saves via **safetensors** (load with `strict=True`). Use the same class as in training:
```python
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn as nn
class DifficultyRouter(nn.Module):
def __init__(self, model_name="microsoft/deberta-v3-xsmall", num_labels=3):
super().__init__()
self.backbone = AutoModel.from_pretrained(model_name)
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_out = outputs.last_hidden_state[:, 0, :].to(torch.float32)
logits = self.classifier(cls_out)
return logits
# Load (safe, no pickle)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("RowRed/DifficultyRouter")
model = DifficultyRouter()
model.load_state_dict(torch.load("model.safetensors", map_location=device), strict=True)
model.to(device).eval()
# Predict
prompts = ["What is 2+2?", "Explain quantum entanglement in detail."]
encoded = tokenizer(prompts, padding=True, truncation=True, max_length=256, return_tensors="pt").to(device)
with torch.no_grad():
logits = model(encoded["input_ids"], encoded["attention_mask"])
probs = torch.softmax(logits, dim=-1)
tiers = torch.argmax(probs, dim=-1)
for prompt, tier in zip(prompts, tiers):
print(f"Tier {tier.item()}: {prompt}")
```
## Citation
If you use this model, please cite:
```bibtex
@software{DifficultyRouter,
author = {RowRed},
title = {DifficultyRouter},
year = {2026},
url = {https://huggingface.co/RowRed/DifficultyRouter}
}
```
Additionally, acknowledge the base model:
```bibtex
@misc{he2021debertav3,
title={DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing},
author={Pengcheng He and Jianfeng Gao and Weizhu Chen},
year={2021},
eprint={2111.09543},
archivePrefix={arXiv},
primaryClass={cs.CL}
}
```
```bibtex
@inproceedings{
he2021deberta,
title={DEBERTA: DECODING-ENHANCED BERT WITH DISENTANGLED ATTENTION},
author={Pengcheng He and Xiaodong Liu and Jianfeng Gao and Weizhu Chen},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=XPZIaotutsD}
}
```
## License
This model is released under Apache‑2.0.
The backbone (microsoft/deberta-v3-xsmall) is MIT‑licensed. |