---
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**.