Modelpublic

somukandula/prompt-router-distilbert

sourceHugging Facemitupdated 5mo agoView on Hugging Face

2likes237downloads

Model Card

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

python

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.

python

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 cheapsmalltext	Yes	Fixed
"Analyze this chart..." routed to cheapsmalltext	Yes	Fixed
"model"/"Hugging Face" in non-code contexts → code_model	Yes	Fixed

Dataset improvements

—+50 summarization templates → cheapsmalltext
—+35 rewriting/editing templates → cheapsmalltext
—+120 vision templates (screenshots, charts, images, diagrams, OCR) → vision_model
—+120 hard negatives for cheapsmalltext: "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
cheapsmalltext	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.	cheapsmalltext	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:

bash

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