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:

Complexity Tier: easy, moderate, hard
Reasoning Depth: reasoning-light, reasoning-moderate, reasoning-intensive
Expected Output Token Length: short-output ($\le 200$), medium-output ($\approx 500$), long-output ($\ge 1,200$)
Execution Priority & Compute Tier: cheap, balanced, premium, realtime, interactive, background
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.

Nasim435/Multi-Label-Prompt-Classifier