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
- Downloads last month
- 7