prvn-ramesh/query-classifier-onnx

Query Classifier ONNX (ModernBERT-large Fine-Tuned)

A CPU-optimized ONNX sequence classification model fine-tuned from ModernBERT-large to classify user queries into low, medium, and hard complexity levels. It is designed for LLM routing, enabling applications to dispatch simple requests to smaller, lower-cost models while reserving larger reasoning models for complex, multi-step tasks.

---

📌 Model Details

• Base Architecture: answerdotai/ModernBERT-large (ModernBertForSequenceClassification)
• Training Data: Fine-tuned on a manually curated dataset of 1,800 user queries spanning factual QA, programming, reasoning, mathematics, and technical problem solving.
• Model Format: Int8 quantized ONNX (model.onnx, ~397 MB).
• Inference Runtime: Pure ONNX Runtime + Hugging Face Rust tokenizers — zero PyTorch or transformers runtime dependencies are required for inference.

---

🏷️ Class Labels & Definitions

Class ID	Label	Description & Examples
0	low	Simple factual questions, basic definitions, direct short lookup, unit conversions. Example: "What is the capital of France?", "How many days in a leap year?"
1	medium	Explanations, code snippets, summarizing text, multi-step instructions. Example: "Explain how key-value storage works in Redis", "Write a python script to parse CSV files."
2	hard	Complex algorithms, multi-file code synthesis, lock-free concurrency, advanced mathematics. Example: "Implement a lock-free SPMC queue in C++ using atomics", "Calculate the integral of x^2 * sin(x) dx."

---

📊 Benchmark & Evaluation Results

Evaluated on an independent held-out test set of 301 unseen queries (balanced across low, medium, and hard buckets):

Accuracy & F1 Breakdown

• Overall Accuracy: 88.04%
• Macro F1-Score: 0.8797

Class	Precision	Recall	F1-Score	Support
low	100.00%	81.00%	0.8950	100
medium	82.18%	83.00%	0.8259	100
hard	84.87%	100.00%	0.9182	101
Overall	88.04% Acc	—	0.8797 Macro-F1	301 total

🎯 Confusion Matrix (Actual vs. Predicted)

Actual \ Predicted	low	medium	hard
low	81	18	1
medium	0	83	17
hard	0	0	101

Latency & Throughput Profile (CPU Execution)

Metric	Measurement
Median Latency (p50)	67.68 ms
Mean Latency	81.44 ms
p90 Latency	134.69 ms
p99 Latency	161.74 ms
Throughput	12.3 queries/sec

🧪 Reproducing Evaluation & Running Local Tests

The held-out test dataset and evaluation script are provided directly in the repository tests/ folder:

• tests/test_dataset.csv: 301 held-out test queries (100 low, 100 medium, 101 hard).
• tests/evaluate.py: Standalone evaluation script for computing accuracy, F1, confusion matrix, and latency profile.

To run the benchmark suite locally:

python tests/evaluate.py

---

🚀 Quickstart: Pure ONNX Runtime Usage

The model can be used directly with ONNX Runtime without requiring PyTorch or Transformers for inference.

1. Installation

pip install onnxruntime tokenizers numpy huggingface-hub

2. Python Code Example

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer
from huggingface_hub import hf_hub_download

# Repository details
REPO_ID = "prvn-ramesh/query-classifier-onnx"

# 1. Download model artifacts from Hugging Face Hub
model_path = hf_hub_download(repo_id=REPO_ID, filename="model.onnx")
tok_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")

# 2. Load tokenizer and initialize ONNX Runtime session
tokenizer = Tokenizer.from_file(tok_path)
session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
label_map = {0: "low", 1: "medium", 2: "hard"}

def classify_query(text: str):
    # Tokenize input string
    encoded = tokenizer.encode(text)
    input_ids = np.array([encoded.ids], dtype=np.int64)
    attention_mask = np.array([encoded.attention_mask], dtype=np.int64)

    # Run inference session
    outputs = session.run(None, {
        "input_ids": input_ids,
        "attention_mask": attention_mask
    })
    
    # Softmax over logits
    logits = outputs[0][0]
    exp_logits = np.exp(logits - np.max(logits))
    probs = exp_logits / np.sum(exp_logits)
    
    pred_idx = int(np.argmax(probs))
    
    return {
        "label": label_map[pred_idx],
        "confidence": float(probs[pred_idx]),
        "scores": {label_map[i]: float(probs[i]) for i in range(len(probs))}
    }

# Example usage
query = "Write a lock-free multi-threaded SPMC queue in C++ using atomics"
result = classify_query(query)

print(f"Query: {query}")
print(f"Predicted Class: {result['label'].upper()} (Confidence: {result['confidence']:.2%})")
print(f"All Scores: {result['scores']}")

---

🙌 Citation & Acknowledgements

• Base model: ModernBERT developed by Answer.AI and LightOn.
• Frameworks: ONNX Runtime and Hugging Face Tokenizers.