| --- |
| language: |
| - en |
| license: mit |
| library_name: catboost |
| pipeline_tag: text-classification |
| tags: |
| - prompt-routing |
| - llm-routing |
| - multi-label-classification |
| - prompt-complexity |
| - catboost |
| - scikit-learn |
| - model-router |
| - token-budgeting |
| pretty_name: Prompt Router (CatBoost Multi-Label Classifier) |
| datasets: |
| - Nasim435/Multi-label-Prompt-Dataset |
| --- |
| |
| # Multi-Label Prompt Classifier |
|
|
| A fast, lightweight multi-label machine learning model designed for prompt complexity estimation, task intent classification, output token length forecasting, and dynamic LLM routing. The model executes inference in **< 10ms on CPU** with **zero GPU dependencies**. |
|
|
| ## Model Summary |
|
|
| - **Architecture**: Scikit-Learn `OneVsRestClassifier` ensemble of 23 binary `CatBoostClassifier` estimators |
| - **Feature Pipeline**: 5,000 TF-IDF features (unigram + bigram) combined with 19 handcrafted structural/semantic text features |
| - **Number of Target Classes**: 23 multi-label categories across 4 semantic dimensions |
| - **Inference Latency**: < 10ms per prompt on standard CPU |
| - **Memory Footprint**: ~13 MB model weights |
| - **Primary Use Case**: Classifying raw user prompts to route them to the most cost-effective LLM tier and enforce pre-inference token budgets without calling an auxiliary LLM. |
|
|
| ## Model Files & Artifacts |
|
|
| The repository contains four serialized artifacts: |
|
|
| | File | Size | Description | |
| |:---|:---:|:---| |
| | **`feature_extractor.pkl`** | 211 KB | Scikit-Learn transformer pipeline combining 5,000 TF-IDF n-gram features with 19 structural heuristics (sentence count, code blocks, math symbols, domain keywords). | |
| | **`prompt_router.pkl`** | 13.0 MB | Trained `OneVsRestClassifier` wrapping 23 individual `CatBoostClassifier` models (iterations=300, depth=6, learning_rate=0.1). | |
| | **`label_binarizer.pkl`** | 826 B | Fitted Scikit-Learn `MultiLabelBinarizer` mapping categorical label names to binary arrays. | |
| | **`thresholds.npy`** | 312 B | Optimal decision threshold matrix ($t_{\text{opt}}$) tuned per class to maximize individual F1 scores. | |
|
|
| ## Target Classes (23 Multi-Label Tags) |
|
|
| The model predicts across 23 categorical dimensions simultaneously: |
|
|
| 1. **Complexity Tier**: `easy`, `moderate`, `hard` |
| 2. **Reasoning Depth**: `reasoning-light`, `reasoning-moderate`, `reasoning-intensive` |
| 3. **Expected Output Token Length**: `short-output` ($\le 200$), `medium-output` ($\approx 500$), `long-output` ($\ge 1,200$) |
| 4. **Execution Priority & Compute Tier**: `cheap`, `balanced`, `premium`, `realtime`, `interactive`, `background` |
| 5. **Task & Domain Intent**: `coding`, `debugging`, `infrastructure`, `architecture`, `architecture-heavy`, `mlops`, `analysis`, `research` |
|
|
| ## Evaluation & Benchmark Performance |
|
|
| Evaluated on an independent 20% holdout test set (372 samples): |
|
|
| | Metric | Baseline ($t=0.50$) | Tuned Thresholds ($t=t_{\text{opt}}$) | Relative Change | |
| |:---|:---:|:---:|:---:| |
| | **Macro F1 Score** | **0.8094** | **0.8320** | **+2.79%** | |
| | **Micro F1 Score** | **0.8282** | **0.8419** | **+1.65%** | |
| | **Weighted F1 Score** | **0.8300** | **0.8447** | **+1.77%** | |
| | **Hamming Loss** | **0.0907** | **0.0840** | **-7.39% (Lower is better)** | |
| | **Inference Latency** | **< 10ms** | **< 10ms** | **CPU Real-Time** | |
| |
| ### Per-Class Evaluation Breakdown |
| |
| | Label | Precision | Recall | F1-Score | Optimal Threshold ($t_{\text{opt}}$) | Test Support | |
| |:---|:---:|:---:|:---:|:---:|:---:| |
| | `architecture-heavy` | 1.00 | 0.90 | **0.95** | 0.40 | 29 | |
| | `interactive` | 0.91 | 0.98 | **0.94** | 0.35 | 230 | |
| | `mlops` | 1.00 | 0.85 | **0.92** | 0.45 | 27 | |
| | `hard` | 0.92 | 0.90 | **0.91** | 0.50 | 136 | |
| | `reasoning-intensive` | 0.92 | 0.90 | **0.91** | 0.50 | 136 | |
| | `realtime` | 0.90 | 0.92 | **0.91** | 0.40 | 48 | |
| | `background` | 0.93 | 0.85 | **0.89** | 0.55 | 91 | |
| | `long-output` | 0.94 | 0.86 | **0.89** | 0.55 | 104 | |
| | `premium` | 0.84 | 0.93 | **0.88** | 0.40 | 114 | |
| | `medium-output` | 0.85 | 0.92 | **0.88** | 0.40 | 177 | |
| | `debugging` | 0.90 | 0.80 | **0.85** | 0.50 | 46 | |
| | `short-output` | 0.86 | 0.81 | **0.84** | 0.50 | 91 | |
| | `coding` | 0.77 | 0.90 | **0.83** | 0.40 | 105 | |
| | `easy` | 0.88 | 0.77 | **0.82** | 0.55 | 96 | |
| | `reasoning-light` | 0.88 | 0.76 | **0.82** | 0.55 | 96 | |
| | `cheap` | 0.82 | 0.79 | **0.80** | 0.50 | 90 | |
| | `balanced` | 0.75 | 0.83 | **0.79** | 0.45 | 122 | |
| | `moderate` | 0.67 | 0.89 | **0.77** | 0.35 | 140 | |
| | `reasoning-moderate` | 0.65 | 0.92 | **0.76** | 0.35 | 140 | |
| | `research` | 0.71 | 0.77 | **0.74** | 0.45 | 22 | |
| | `infrastructure` | 0.62 | 0.83 | **0.71** | 0.35 | 77 | |
| | `analysis` | 0.56 | 0.85 | **0.68** | 0.35 | 41 | |
| | `architecture` | 0.80 | 0.56 | **0.66** | 0.55 | 43 | |
|
|
| ## Quick Start & Inference |
|
|
| ### Installation |
|
|
| ```bash |
| pip install catboost scikit-learn numpy pandas joblib scipy |
| ``` |
|
|
| ### Loading and Predicting |
|
|
| ```python |
| import joblib |
| import numpy as np |
| import pandas as pd |
| |
| # 1. Load serialized artifacts |
| feature_extractor = joblib.load("feature_extractor.pkl") |
| classifier = joblib.load("prompt_router.pkl") |
| mlb = joblib.load("label_binarizer.pkl") |
| thresholds = np.load("thresholds.npy") |
| |
| def predict_prompt_labels(prompt: str, return_scores: bool = False): |
| # Transform input text into combined TF-IDF + structural feature matrix |
| X = feature_extractor.transform(pd.Series([prompt])) |
| |
| # Predict probabilities for each binary classifier in the ensemble |
| probs = np.array(classifier.predict_proba(X)) |
| scores = np.array([p[0][1] if np.ndim(p) == 2 else p[1] for p in probs]) |
| |
| # Apply calibrated decision thresholds |
| predictions = (scores >= thresholds).astype(int) |
| |
| # Fallback to top-scoring class if no threshold is met |
| if predictions.sum() == 0: |
| predictions[np.argmax(scores)] = 1 |
| |
| labels = list(mlb.inverse_transform(predictions.reshape(1, -1))[0]) |
| |
| if return_scores: |
| score_dict = {label: round(float(score), 4) for label, score in zip(mlb.classes_, scores)} |
| return labels, score_dict |
| |
| return labels |
| |
| # Example usage |
| query = "Design a distributed real-time fraud detection pipeline with Apache Flink and Kafka." |
| labels, scores = predict_prompt_labels(query, return_scores=True) |
| print("Predicted labels:", labels) |
| # Output: ['architecture-heavy', 'hard', 'infrastructure', 'interactive', 'long-output', 'premium', 'realtime', 'reasoning-intensive'] |
| ``` |
|
|
| ## Intended Use & Integration |
|
|
| - **LLM Routing Middleware**: Classify incoming prompts to route between small/nano (e.g. 8B–9B), medium (e.g. 30B–70B), and large/frontier (e.g. 120B–405B) models. |
| - **Pre-Inference Token Budgeting**: Forecast expected output token lengths (`short-output`, `medium-output`, `long-output`) before generation to prevent token overspend. |
| - **Domain Specialization**: Direct code queries to coding models, debugging queries to specialized debug agents, and theoretical research questions to reasoning models. |
|
|
| ## Limitations |
|
|
| - **Domain Scope**: The training dataset focuses on technical engineering prompts (software engineering, cloud infrastructure, mathematics, algorithms). Predictions on general casual conversation or creative fiction may be less accurate. |
| - **Language**: English prompts only (`language: en`). |
|
|
| ## License |
|
|
| This model is distributed under the **MIT License**. |
