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

Downloads last month
7
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support