addyo07/distilbert-query-classifier

Query Sieve Classifier

DistilBERT multilingual fine-tuned to classify user queries as GENERIC or SEMANTIC — filtering chit-chat from durable knowledge worth storing.

Model Description

• Architecture: distilbert-base-multilingual-cased (134M params)
• Quantization: INT8 dynamic (ONNX Runtime)
• Input: Short text queries in English or Hindi (≤10 words; longer sentences bypass to SEMANTIC)
• Output: Binary — GENERIC (0) or SEMANTIC (1)
• Inference: ONNX Runtime CPU, single-thread P99 = 16.87ms

Repository Structure

distilbert-query-classifier/
├── README.md                    # Model card (this file)
├── model/
│   ├── pytorch/                 # PyTorch safetensors (for transformers library)
│   │   ├── model.safetensors    # Full precision PyTorch weights
│   │   ├── config.json          # Model configuration
│   │   ├── tokenizer.json       # HuggingFace tokenizer
│   │   └── tokenizer_config.json
│   └── onnx/                    # ONNX INT8 quantized (for CPU inference)
│       └── model_quantized.onnx # INT8 quantized model (~130 MB)
├── scripts/                     # Python training pipeline
│   ├── config.py                # Constants and paths
│   ├── generate_dataset.py      # Synthetic data generation via Ollama
│   ├── train.py                 # Fine-tuning + ONNX export + latency benchmark
│   └── ...

Usage

Python (with transformers + ONNX Runtime)

from transformers import AutoTokenizer
import onnxruntime as ort

# Load tokenizer from the pytorch folder
tokenizer = AutoTokenizer.from_pretrained(
    "addyo07/distilbert-query-classifier",
    subfolder="model/pytorch",
)

# Load ONNX model
session = ort.InferenceSession("model/onnx/model_quantized.onnx")

def classify(text: str) -> str:
    inputs = tokenizer(text, return_tensors="np", max_length=64, truncation=True, padding="max_length")
    logits = session.run(None, {
        "input_ids": inputs["input_ids"].astype(np.int64),
        "attention_mask": inputs["attention_mask"].astype(np.int64),
    })[0]
    return "SEMANTIC" if logits[0][1] > logits[0][0] else "GENERIC"

Rust (with query-sieve crate)

[dependencies]
query-sieve = { git = "https://github.com/addy-47/query-sieve-rs" }

use query_sieve::GenericSemanticClassifier;

let classifier = GenericSemanticClassifier::load(
    "model/onnx/model_quantized.onnx",
    "model/pytorch/tokenizer.json",
)?;
let result = classifier.classify("my name is John")?;

Download individual files

# ONNX model (for CPU inference)
wget https://huggingface.co/addyo07/distilbert-query-classifier/resolve/main/model/onnx/model_quantized.onnx

# PyTorch weights (for fine-tuning)
wget https://huggingface.co/addyo07/distilbert-query-classifier/resolve/main/model/pytorch/model.safetensors

# Tokenizer
wget https://huggingface.co/addyo07/distilbert-query-classifier/resolve/main/model/pytorch/tokenizer.json

Performance

Split	Accuracy
Test (15% holdout)	98.39%

Latency

Mode	P50	P99
Multi-thread CPU	8.39 ms	11.81 ms
Single-thread CPU (intra_op_threads=1)	14.81 ms	16.87 ms

Training Data

Dataset: addyo07/query-classification-dataset

12,044 synthetic examples generated by llama3.1:8b:

Category	English	Hindi
GENERIC	3,003	3,019
SEMANTIC	3,017	3,005

The SEMANTIC category contains ~40% short standalone statements (3-7 words) to prevent the model from learning "semantic = long sentence."

Training Scripts

The scripts/ directory contains the full training pipeline:

1. python scripts/generate_dataset.py --category en_semantic — generate synthetic data via Ollama
2. python scripts/train.py — fine-tune DistilBERT + export ONNX INT8 + benchmark

Rust Crate

The query-sieve Rust crate provides the inference runtime:

• GitHub: addy-47/query-sieve-rs
• ONNX Runtime (ort) with HuggingFace tokenizer
• Configurable single/multi-thread CPU

• 10 word bypass (long sentences auto-classify as SEMANTIC)

License

MIT