README.md

metadata

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

pip install catboost scikit-learn numpy pandas joblib scipy

Loading and Predicting

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.