Detoxify ONNX 🚀

This project provides an ONNX-exported and quantized version of the Detoxify multilingual model, optimized for runtime inference.
It enables faster and lighter toxicity detection using ONNX Runtime.

---

📦 Installation

pip install -r requirements.txt

---

🛠️ Features

• Export Detoxify multilingual model to ONNX
• Quantize the ONNX model for faster inference
• Update model config.json with correct label mappings
• Evaluate the model using standard metrics (accuracy, precision, recall, F1, AUC-ROC)

---

🧪 ONNX Evaluation Results

Original Model (using Detoxify lib and ONNX):

Time for 1 threshold evaluation =~ 3 min 30s

Quantized model:

Time for 1 threshold evaluation =~ 2 min 41s

---

📥 Dataset: Jigsaw Toxic Comment Classification

This project uses the validation set from the Jigsaw Toxic Comment Classification Challenge.

🔽 How to Download:

1. Create a Kaggle account and accept the competition rules.
2. Download the dataset from the Data tab of the competition page.
3. Extract or locate the validation.csv.
4. Place the validation file under a local folder:

mkdir -p jigsaw_data
cp path/to/validation.csv jigsaw_data/

Note: You may want to install the Kaggle CLI for automated downloading:

# Make sur to have kaggle.json loaded accordingly
pip install kaggle
kaggle competitions download -c jigsaw-toxic-comment-classification-challenge
unzip jigsaw-toxic-comment-classification-challenge.zip -d jigsaw_data

🧵 Usage

Convert and Quantize

python convert_to_onnx.py --model-dir detoxify_model \
  --save-hf-model \
  --export-onnx \
  --quantize-onnx \
  --update-config

Evaluate

python evaluate.py --validation-file ./jigsaw_data/validation.csv \
  --model-path ./detoxify_model/model.quant.onnx \
  --tokenizer-path ./detoxify_model \
  --thresholds 0.2 0.4 0.5 0.7 0.9 \
  --plot-roc

---

🤗 Hugging Face Model

You can find and use the quantized ONNX model from the Gravitee.io organization on Hugging Face:

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForSequenceClassification
import numpy as np

# Load model and tokenizer using optimum
model = ORTModelForSequenceClassification.from_pretrained("gravitee-io/detoxify-onnx", file_name="model.quant.onnx")
tokenizer = AutoTokenizer.from_pretrained("gravitee-io/detoxify-onnx")

# Tokenize input
text = "Your comment here"
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)

# Run inference
outputs = model(**inputs)
logits = outputs.logits

# Optional: convert to probabilities
probs = 1 / (1 + np.exp(-logits))
print(probs)