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.