---
pipeline_tag: text-classification
library_name: transformers
language:
- en
license: mit
datasets:
- thesofakillers/jigsaw-toxic-comment-classification-challenge
base_model: bert-base-uncased
tags:
- toxic
- moderation
- safety
- content-moderation
- onnx
- quantized
---

# 🌟 Shuddhi v1: BERT-Base Toxicity Checker


[![Model Type: BERT-Base](https://img.shields.io/badge/Model_Type-BERT--Base-blue.svg?style=flat-square)](https://huggingface.co/bert-base-uncased)
[![Dataset: JIGSAW](https://img.shields.io/badge/Dataset-JIGSAW-orange.svg?style=flat-square)](https://www.kaggle.com/c/jigsaw-toxic-comment-classification-challenge)
[![License: MIT](https://img.shields.io/badge/License-MIT-green.svg?style=flat-square)](./LICENSE)
[![Framework: PyTorch / ONNX](https://img.shields.io/badge/Framework-PyTorch%20%2F%20ONNX-blueviolet.svg?style=flat-square)](#quick-start)
[![Task: Multi-label Classification](https://img.shields.io/badge/Task-Multi--label%20Classification-yellowgreen.svg?style=flat-square)](#model-card)

**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`](./thresholds.json):

| Category        | Description                                      | Optimal Threshold |
| :-------------- | :----------------------------------------------- | :---------------- |
| `toxic`         | General toxic, rude, or disrespectful comment    | `0.7800`          |
| `severe_toxic`  | Extremely aggressive or highly offensive comment | `0.8539`          |
| `obscene`       | Obscene, vulgar, or profane language             | `0.9070`          |
| `threat`        | Threats of violence, physical harm, or death     | `0.3861`          |
| `insult`        | Insults or derogatory remarks                    | `0.8832`          |
| `identity_hate` | Hate speech targeting identity groups            | `0.7942`          |

---

## ⚡ 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)

```python
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`):

```python
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.

| Runtime / Format   | Precision | Avg. Latency (CPU)      | Storage Size |
| :----------------- | :-------- | :---------------------- | :----------- |
| **PyTorch (Base)** | FP32      | ~120ms                  | ~438 MB      |
| **ONNX Quantized** | INT8      | **~25ms** (4.8x faster) | **105 MB**   |

_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`](https://huggingface.co/datasets/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](./LICENSE) file for details.