--- license: apache-2.0 base_model: answerdotai/ModernBERT-large tags: - onnx - sequence-classification - text-classification - query-classifier - modernbert - llm-routing pipeline_tag: text-classification library_name: onnxruntime --- # 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`](https://huggingface.co/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/`](tests) folder: - [`tests/test_dataset.csv`](tests/test_dataset.csv): 301 held-out test queries (100 `low`, 100 `medium`, 101 `hard`). - [`tests/evaluate.py`](tests/evaluate.py): Standalone evaluation script for computing accuracy, F1, confusion matrix, and latency profile. To run the benchmark suite locally: ```bash 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 ```bash pip install onnxruntime tokenizers numpy huggingface-hub ``` ### 2. Python Code Example ```python 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](https://huggingface.co/answerdotai/ModernBERT-large) developed by **Answer.AI** and **LightOn**. - Frameworks: [ONNX Runtime](https://onnxruntime.ai/) and [Hugging Face Tokenizers](https://github.com/huggingface/tokenizers).