somukandula/prompt-router-distilbert

Prompt Router — DistilBERT Classifier (v2)

This model reads a user prompt and decides which LLM should answer it, so you don't waste money running simple queries through massive models.

What It Actually Does

You have a pool of 4 models with very different sizes and costs. Instead of sending every prompt to the biggest one, this classifier looks at the prompt and picks the cheapest model that can handle it. If it's unsure, it safely falls back to the big model.

Route	Target Model	Size	When to use
cheap_small_text	Qwen/Qwen3-1.7B	1.7B params	Simple QA, rewriting, summarization, casual chat
code_model	Qwen/Qwen2.5-Coder-1.5B-Instruct	1.5B params	Coding, debugging, code explanation, devops
vision_model	Qwen/Qwen2.5-VL-3B-Instruct	3B params	Anything with images, screenshots, charts, diagrams
strong_general	mistralai/Mistral-Small-3.2-24B-Instruct-2506	24B params	Complex reasoning, research, planning, uncertain cases

Live Examples

from transformers import pipeline

router = pipeline("text-classification", model="somukandula/prompt-router-distilbert")

# Simple questions → tiny 1.7B model
router("What is the capital of France?")
# [{'label': 'cheap_small_text', 'score': 0.993}]

# Summarization (hard negative: contains "model" and "Hugging Face")
router("Summarize this paragraph: Hugging Face hosts models and datasets.")
# [{'label': 'cheap_small_text', 'score': 0.991}]

# Coding → 1.5B code-specialized model
router("Write a Python function to reverse a string")
# [{'label': 'code_model', 'score': 0.994}]

router("Fix this React error: Cannot read properties of undefined")
# [{'label': 'code_model', 'score': 0.993}]

# Images/vision → 3B vision model
router("What is in this screenshot?")
# [{'label': 'vision_model', 'score': 0.993}]

router("Analyze this chart and tell me the trend")
# [{'label': 'vision_model', 'score': 0.986}]

# Hard reasoning → 24B strong model
router("Plan a research project comparing small language models on math reasoning.")
# [{'label': 'strong_general', 'score': 0.995}]

Confidence Fallback: The Safety Net

If the model is unsure (confidence < 0.60), it routes to strong_general rather than risking a bad answer from a cheap model.

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("somukandula/prompt-router-distilbert")
model = AutoModelForSequenceClassification.from_pretrained("somukandula/prompt-router-distilbert")

def route(prompt: str, threshold: float = 0.60):
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128)
    with torch.no_grad():
        probs = torch.softmax(model(**inputs).logits, dim=-1)[0]

    pred_id = probs.argmax().item()
    confidence = probs.max().item()

    id2label = {0: "cheap_small_text", 1: "code_model", 2: "vision_model", 3: "strong_general"}
    label = id2label[pred_id]

    if confidence < threshold:
        label = "strong_general"
        reason = f"low confidence ({confidence:.2f} < {threshold})"
    else:
        reason = f"confidence {confidence:.2f}"

    return {
        "prompt": prompt,
        "route": label,
        "model": MODEL_MAP[label],
        "confidence": round(confidence, 4),
        "reason": reason,
        "all_probs": {id2label[i]: round(float(probs[i]), 4) for i in range(4)}
    }

MODEL_MAP = {
    "cheap_small_text": "Qwen/Qwen3-1.7B",
    "code_model": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
    "vision_model": "Qwen/Qwen2.5-VL-3B-Instruct",
    "strong_general": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
}

v2 Improvements

This v2 model fixes known routing mistakes from v1:

Issue	v1	v2
Summarization routed to code_model	Yes	Fixed
"What is in this screenshot?" routed to cheap_small_text	Yes	Fixed
"Analyze this chart..." routed to cheap_small_text	Yes	Fixed
"model"/"Hugging Face" in non-code contexts → code_model	Yes	Fixed

Dataset improvements

• +50 summarization templates → cheap_small_text
• +35 rewriting/editing templates → cheap_small_text
• +120 vision templates (screenshots, charts, images, diagrams, OCR) → vision_model
• +120 hard negatives for cheap_small_text: "What is a language model?", "Summarize the BERT model", etc.
• +30 hard negatives for code_model: "What is the dress code?", "Morse code for SOS", etc.
• +80 hard negatives for vision_model: text-only discussions of charts/diagrams

Performance

Evaluated on 388 held-out test prompts:

Metric	Value
Accuracy	0.9768
Macro F1	0.9763

Confusion Matrix

                    pred→
              cheap  code  vision  strong
true cheap     96     1      1       3
true code       0   102      0       0
true vision     2     2     80       0
true strong     0     0      0     101

Per-class F1

Class	F1 Score
cheap_small_text	0.9648
code_model	0.9855
vision_model	0.9697
strong_general	0.9854

Required Test Prompts

Prompt	Expected	Predicted	Confidence
Summarize this paragraph: Hugging Face hosts models and datasets.	cheap_small_text	cheap_small_text	0.9907
What is in this screenshot?	vision_model	vision_model	0.9926
Analyze this chart and tell me the trend	vision_model	vision_model	0.9856
Write a Python function to reverse a string	code_model	code_model	0.9937
Fix this React error: Cannot read properties of undefined	code_model	code_model	0.9932
Plan a research project comparing small language models on math reasoning.	strong_general	strong_general	0.9946

Model Details

Property	Value
Base model	distilbert-base-uncased
Architecture	DistilBERTForSequenceClassification
Classes	4
Max sequence length	128 tokens
Training data	600 prompts (v2, balanced with hard negatives)
Test data	388 prompts
Training epochs	4
Learning rate	3e-5
Batch size	8
Model size	~255 MB

Training Data

• Dataset: somukandula/prompt-router-dataset
• 600 training prompts, 388 test prompts
• 4 classes with hard negatives for each

Comparison with Baselines

Method	Accuracy	Macro F1
Rule-Based (keywords)	0.8786	0.8712
Embeddings + LogReg	0.9857	0.9855
DistilBERT v1	1.0000	1.0000
DistilBERT v2 (this model)	0.9768	0.9763

v1 achieved 1.0 on a smaller synthetic test set. v2 uses a larger, more challenging test set with hard negatives and still achieves >97% accuracy.

Inference Script

A ready-to-use script is included in this repo:

python router_inference.py "Your prompt here"

Limitations

• Trained on synthetic prompts, not real user traffic. Performance may vary on out-of-distribution production prompts.
• Does not inspect actual images — it only routes based on the text prompt (e.g. "describe this screenshot").
• Cost savings depend on your actual pricing; the relative cost model used here is illustrative.

License

MIT