CoolFace
Modelpublic

addyo07/distilbert-query-classifier

sourceHugging Facemitupdated 3mo agoView on Hugging Face
0likes
Model Card

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)

python
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)

toml
[dependencies]
query-sieve = { git = "https://github.com/addy-47/query-sieve-rs" }
rust
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

bash
# 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

SplitAccuracy
Test (15% holdout)98.39%

Latency

ModeP50P99
Multi-thread CPU8.39 ms11.81 ms
Single-thread CPU (intraopthreads=1)14.81 ms16.87 ms

Training Data

Dataset: addyo07/query-classification-dataset

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

CategoryEnglishHindi
GENERIC3,0033,019
SEMANTIC3,0173,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. 1.python scripts/generate_dataset.py --category en_semantic — generate synthetic data via Ollama
  2. 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