shuddhi-base-onnx-int8

🌟 Shuddhi v1: BERT-Base Toxicity Checker

Shuddhi is a high-performance, production-ready moderation model based on the bert-base-uncased architecture. It is fine-tuned on the JIGSAW Toxic Comment Classification dataset to detect and classify toxic text into six distinct labels.

The repository includes both the standard PyTorch model configuration and a quantized ONNX version (model_quantized.onnx) optimized for low-latency CPU and edge deployments.

---

🚀 Model Details

• Developed by: Shuddhi Project Authors
• Model Type: Transformer (bert)
• Base Model: bert-base-uncased
• Language(s) (NLP): English
• License: MIT
• Task: Multi-Label Text Classification (Toxicity Moderation)
• Input Limit: 512 tokens

Detected Categories & Optimal Thresholds

The model classifies text across the 6 JIGSAW standard categories. To optimize moderation accuracy and balance precision/recall, use the pre-calculated classification thresholds from thresholds.json:

---

⚡ Quick Start

You can load and perform inference with this model using either Python's transformers library or using the optimized ONNX runtime.

Option 1: Standard PyTorch Inference (via Hugging Face Transformers)

import torch
import json
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load model, tokenizer, and thresholds
model_path = "./"  # Path to the shuddhi_v1 directory
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)

with open(f"{model_path}/thresholds.json") as f:
    thresholds = json.load(f)

# Prepare inputs
text = "Go play in traffic!"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)

# Run prediction
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits
    # Apply sigmoid since it is multi-label classification
    probabilities = torch.sigmoid(logits).cpu().numpy()[0]

# Class mapping and threshold application
results = {}
for i in range(len(probabilities)):
    label = model.config.id2label[i]
    score = float(probabilities[i])
    results[label] = {
        "score": score,
        "flagged": score >= thresholds.get(label, 0.5)
    }

print("Moderation Scores:")
for label, res in results.items():
    status = "🚨 FLAGGED" if res["flagged"] else "✅ CLEAN"
    print(f" - {label:<15}: {res['score']:.4f} [{status}]")

Option 2: High-Performance ONNX Runtime Inference

For low-latency production applications, load the pre-quantized ONNX model (model_quantized.onnx):

import numpy as np
import json
from transformers import AutoTokenizer
import onnxruntime as ort

# Load tokenizer, ONNX session, and thresholds
model_path = "./"
tokenizer = AutoTokenizer.from_pretrained(model_path)
ort_session = ort.InferenceSession(f"{model_path}/model_quantized.onnx")

with open(f"{model_path}/thresholds.json") as f:
    thresholds = json.load(f)

# Prepare inputs
text = "This is a clean, helpful, and respectful comment."
inputs = tokenizer(text, return_tensors="np", truncation=True, max_length=512)

# Cast token inputs to INT64 for ONNX compatibility
onnx_inputs = {
    "input_ids": inputs["input_ids"].astype(np.int64),
    "attention_mask": inputs["attention_mask"].astype(np.int64),
}
if "token_type_ids" in inputs:
    onnx_inputs["token_type_ids"] = inputs["token_type_ids"].astype(np.int64)

# Run ONNX inference
logits = ort_session.run(None, onnx_inputs)[0]

# Compute probabilities (Sigmoid)
probabilities = 1 / (1 + np.exp(-logits))[0]

# Output results using thresholds
labels = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
results = {}
for label, score in zip(labels, probabilities):
    results[label] = {
        "score": float(score),
        "flagged": float(score) >= thresholds.get(label, 0.5)
    }

print("ONNX Moderation Scores:")
for label, res in results.items():
    status = "🚨 FLAGGED" if res["flagged"] else "✅ CLEAN"
    print(f" - {label:<15}: {res['score']:.4f} [{status}]")

---

📈 Performance & Benchmark

The quantization of Shuddhi to ONNX format yields significant latency reductions with minimal loss in classification accuracy.

Note: Benchmarks conducted on a typical AMD Ryzen 5 5500U CPU with sequence lengths of 128 tokens.

---

📊 Dataset: JIGSAW Toxicity

The model was trained on the dataset from the JIGSAW Toxic Comment Classification Challenge on Kaggle. The dataset contains comments from Wikipedia talk pages labeled by human raters for toxic behavior.

• Total Samples: 465,899 comments (source: thesofakillers/jigsaw-toxic-comment-classification-challenge)
• Toxicity Rate: ~10% of the comments in the training set are labeled as toxic or hostile.

---

⚠️ Intended Use & Limitations

Intended Use

• Moderation engines for chat applications, comment threads, and online communities.
• Real-time safety filters for collaborative platforms.
• Analysis tools for historical community sentiment or behavior metrics.

Limitations & Biases

• Nuance and Context: The model is trained at the comment/sentence level and may struggle with subtle sarcasm, irony, or highly contextual toxicity.
• Bias in Training Data: Because the model is trained on JIGSAW data sourced from Wikipedia talk pages, it may reflect historical biases present in the labeling process (e.g., higher false-positive rates for text containing certain demographic keywords). We advise monitoring predictions and using a confidence threshold suited to your application needs.

---

📄 License

This model card and the Shuddhi model are distributed under the MIT License. See the accompanying LICENSE file for details.