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.

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:

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:

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

Required Test Prompts

Model Details

Training Data

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

Comparison with Baselines

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