Upload folder using huggingface_hub

README.md

@@ -16,7 +16,7 @@ metrics:
 - f1
 ---
 
-# Prompt Router — DistilBERT Classifier
+# 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.
 
@@ -42,29 +42,27 @@ router = pipeline("text-classification", model="somukandula/prompt-router-distil
 router("What is the capital of France?")
 # [{'label': 'cheap_small_text', 'score': 0.993}]
 
-router("Summarize this paragraph: The solar system consists of...")
-# [{'label': 'cheap_small_text', 'score': 0.997}]
+# 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 linked list")
-# [{'label': 'code_model', 'score': 0.991}]
+router("Write a Python function to reverse a string")
+# [{'label': 'code_model', 'score': 0.994}]
 
-router("Debug this TypeScript error: Type 'string' is not assignable...")
-# [{'label': 'code_model', 'score': 0.989}]
+router("Fix this React error: Cannot read properties of undefined")
+# [{'label': 'code_model', 'score': 0.993}]
 
 # Images/vision → 3B vision model
-router("Describe what is shown in this screenshot")
-# [{'label': 'vision_model', 'score': 0.995}]
+router("What is in this screenshot?")
+# [{'label': 'vision_model', 'score': 0.993}]
 
-router("Extract the text from this chart image")
-# [{'label': 'vision_model', 'score': 0.992}]
+router("Analyze this chart and tell me the trend")
+# [{'label': 'vision_model', 'score': 0.986}]
 
 # Hard reasoning → 24B strong model
-router("Develop a 5-year strategic plan for a startup entering the EV market")
-# [{'label': 'strong_general', 'score': 0.992}]
-
-router("Analyze the geopolitical consequences of rare earth mineral scarcity")
-# [{'label': 'strong_general', 'score': 0.994}]
+router("Plan a research project comparing small language models on math reasoning.")
+# [{'label': 'strong_general', 'score': 0.995}]
 ```
 
 ## Confidence Fallback: The Safety Net
@@ -110,49 +108,67 @@ MODEL_MAP = {
     "vision_model": "Qwen/Qwen2.5-VL-3B-Instruct",
     "strong_general": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
 }
-
-# Example
-print(route("How do I fix this numpy import error?"))
-# {
-#   'route': 'code_model',
-#   'model': 'Qwen/Qwen2.5-Coder-1.5B-Instruct',
-#   'confidence': 0.9876,
-#   'reason': 'confidence 0.99',
-#   'all_probs': {
-#     'cheap_small_text': 0.0035,
-#     'code_model': 0.9908,
-#     'vision_model': 0.0029,
-#     'strong_general': 0.0028
-#   }
-# }
 ```
 
-## Why Use This?
+## 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** |
 
-- **Cost**: Running everything through a 24B model is ~8× more expensive than the 1.7B model.
-- **Latency**: Smaller models are faster.
-- **Safety**: The 0.60 confidence threshold means uncertain prompts still go to the big model.
+### 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 140 held-out test prompts:
+Evaluated on 388 held-out test prompts:
 
 | Metric | Value |
 |--------|-------|
-| Accuracy | **1.0000** |
-| Macro F1 | **1.0000** |
+| Accuracy | **0.9768** |
+| Macro F1 | **0.9763** |
+
+### Confusion Matrix
 
-Per-class F1: 1.0000 for all 4 categories.
+```
+                    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
+```
 
-### Cost Savings vs Always Using the 24B Model
+### Per-class F1
 
-| Router | Accuracy | Macro F1 | Effective Cost Savings |
-|--------|----------|----------|------------------------|
-| Rule-Based (keywords) | 0.8786 | 0.8712 | 54.6% |
-| Embeddings + LogReg | 0.9857 | 0.9855 | 60.0% |
-| **DistilBERT (this model)** | **1.0000** | **1.0000** | **60.3%** |
+| Class | F1 Score |
+|-------|----------|
+| cheap_small_text | 0.9648 |
+| code_model | 0.9855 |
+| vision_model | 0.9697 |
+| strong_general | 0.9854 |
 
-> "Effective savings" includes a quality penalty: if a prompt is wrongly sent to a cheap model, we count that as a partial loss. This model had **zero misroutes** on the test set, so the penalty is 0.
+### 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
 
@@ -162,18 +178,29 @@ Per-class F1: 1.0000 for all 4 categories.
 | Architecture | DistilBERTForSequenceClassification |
 | Classes | 4 |
 | Max sequence length | 128 tokens |
-| Training data | 560 synthetic prompts |
-| Training epochs | 10 (early stopping patience=3) |
-| Learning rate | 2e-5 |
-| Batch size | 16 |
+| 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](https://huggingface.co/datasets/somukandula/prompt-router-dataset)
-- 560 training prompts, 140 test prompts
-- 4 balanced classes
-- Synthetic but diverse — covers simple chat, code, vision-language, and complex reasoning
+- 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
 
@@ -183,27 +210,11 @@ A ready-to-use script is included in this repo:
 python router_inference.py "Your prompt here"
 ```
 
-Output:
-```
-Prompt: Your prompt here
-------------------------------------------------------------
-Routed to: code_model
-Model:     Qwen/Qwen2.5-Coder-1.5B-Instruct
-Reason:    confidence 0.991
-
-All class probabilities:
-  cheap_small_text    : 0.0035
-  code_model          : 0.9908
-  vision_model        : 0.0029
-  strong_general      : 0.0028
-```
-
 ## 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.
-- The 1.0 test accuracy is on a held-out synthetic test set; expect some degradation on real-world ambiguous prompts.
 
 ## License
 

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f1aaf37f51a1de3a2d7b4090ff02a13d98a4df96b9087b0da4039dd462def7d3
+oid sha256:67f0c2e4d35255562a324f134d584faedb5a35bbc4a15f8c20e21f3e209413a2
 size 267838720