distilbert-toxic-onnx-int8

DistilBERT Toxic Comment Classifier — ONNX INT8 Quantized

Fine-tuned distilbert-base-uncased for binary toxicity classification, exported to ONNX and INT8-quantized for fast CPU inference.

• Base model: distilbert-base-uncased
• Task: Binary text classification (toxic / clean)
• Dataset: Jigsaw Toxic Comment Classification
• Optimization: ONNX Runtime + INT8 dynamic quantization via HuggingFace Optimum
• AUC-ROC: 0.9217
• Size: 64 MB (vs 255 MB PyTorch — 75% reduction)
• p50 inference latency: ~29ms on CPU

Usage

from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import DistilBertTokenizerFast
import torch

model = ORTModelForSequenceClassification.from_pretrained("Trakerlack/distilbert-toxic-onnx-int8")
tokenizer = DistilBertTokenizerFast.from_pretrained("Trakerlack/distilbert-toxic-onnx-int8")

inputs = tokenizer("You are an idiot!", return_tensors="pt", truncation=True, max_length=216)
outputs = model(**inputs)
probs = torch.softmax(outputs.logits, dim=1)
toxic_prob = probs[0, 1].item()
print(f"Toxic probability: {toxic_prob:.3f}")

Project

GitHub | Live Demo