Add PyTorch training pipeline with ONNX/INT8 export and inference script

hatespeech.py

@@ -0,0 +1,267 @@
+#!/usr/bin/env python3
+"""
+hatespeech.py - ONNX Runtime Inference & Examples
+Performs fast inference on text using exported ONNX (Raw or INT8 Quantized) models.
+
+Key Features:
+- Runs via ONNX Runtime (CPU or CUDA).
+- Supports both raw `hatespeech.onnx` and quantized `hatespeech_int8.onnx`.
+- Preprocesses input text and tokenizes inputs.
+- Outputs predicted class, confidence scores across all classes, and latency in milliseconds.
+- Provides interactive mode, CLI text argument, and built-in benchmark examples.
+"""
+
+import os
+import sys
+import re
+import html
+import time
+import argparse
+from typing import List, Dict, Union, Any
+
+import numpy as np
+import onnxruntime as ort
+from transformers import AutoTokenizer
+
+LABEL_NAMES = {
+    0: "Hate Speech",
+    1: "Offensive Language",
+    2: "Neither",
+}
+
+
+def clean_text(text: str) -> str:
+    """Preprocesses input text matching the training pipeline."""
+    if not isinstance(text, str):
+        return ""
+    text = html.unescape(text)
+    text = re.sub(r"https?://\S+|www\.\S+", "", text)
+    text = re.sub(r"@\w+", "", text)
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+
+def softmax(x: np.ndarray) -> np.ndarray:
+    """Computes softmax over logits."""
+    e_x = np.exp(x - np.max(x, axis=-1, keepdims=True))
+    return e_x / e_x.sum(axis=-1, keepdims=True)
+
+
+class HateSpeechDetector:
+    """
+    Hate Speech & Offensive Language Detector powered by ONNX Runtime.
+    """
+
+    def __init__(
+        self,
+        model_path: str = "./model/hatespeech_int8.onnx",
+        tokenizer_name_or_path: str = "./model",
+        fallback_model_name: str = "distilbert-base-uncased",
+        max_length: int = 128,
+        use_gpu: bool = False,
+    ):
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(
+                f"Model file not found at: '{model_path}'.\n"
+                f"Please train the model first by running `python train.py`."
+            )
+
+        print(f"Loading ONNX model from: {model_path}")
+        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
+        
+        # Configure ONNX runtime session
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        self.session = ort.InferenceSession(model_path, sess_options=sess_options, providers=providers)
+
+        # Load tokenizer (first check local model dir, then fallback to HF hub)
+        if os.path.exists(os.path.join(tokenizer_name_or_path, "tokenizer_config.json")):
+            print(f"Loading tokenizer from local path: {tokenizer_name_or_path}")
+            self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
+        else:
+            print(f"Local tokenizer not found. Loading fallback: {fallback_model_name}")
+            self.tokenizer = AutoTokenizer.from_pretrained(fallback_model_name)
+
+        self.max_length = max_length
+        self.model_path = model_path
+        self.file_size_mb = os.path.getsize(model_path) / (1024 * 1024)
+
+    def predict(self, text: Union[str, List[str]]) -> List[Dict[str, Any]]:
+        """
+        Runs inference on a single string or a list of strings.
+        Returns prediction details including class label, confidence, and latency.
+        """
+        is_single = isinstance(text, str)
+        raw_texts = [text] if is_single else text
+        cleaned_texts = [clean_text(t) for t in raw_texts]
+
+        start_time = time.perf_counter()
+
+        # Tokenize
+        encoded = self.tokenizer(
+            cleaned_texts,
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="np",
+        )
+
+        ort_inputs = {
+            "input_ids": encoded["input_ids"].astype(np.int64),
+            "attention_mask": encoded["attention_mask"].astype(np.int64),
+        }
+
+        # Run ONNX inference
+        ort_outputs = self.session.run(None, ort_inputs)
+        logits = ort_outputs[0]
+        probs = softmax(logits)
+
+        elapsed_ms = (time.perf_counter() - start_time) * 1000
+        latency_per_sample = elapsed_ms / max(len(raw_texts), 1)
+
+        results = []
+        for i, original_text in enumerate(raw_texts):
+            pred_class_id = int(np.argmax(probs[i]))
+            class_probs = {
+                LABEL_NAMES[c]: float(probs[i][c]) for c in range(3)
+            }
+            results.append({
+                "text": original_text,
+                "cleaned_text": cleaned_texts[i],
+                "label_id": pred_class_id,
+                "label_name": LABEL_NAMES[pred_class_id],
+                "confidence": float(probs[i][pred_class_id]),
+                "probabilities": class_probs,
+                "latency_ms": latency_per_sample,
+            })
+
+        return results
+
+
+def run_examples(detector: HateSpeechDetector):
+    """Runs a series of representative test cases and prints formatted output."""
+    examples = [
+        "I really love this community, everyone is so supportive and kind!",
+        "Can you please send me the report by tomorrow morning?",
+        "Shut the fuck up, you are being so damn annoying tonight.",
+        "That bitch is crazy, get the hell out of here.",
+        "Kill all members of that group, they don't deserve to exist in our country.",
+    ]
+
+    print("\n" + "=" * 75)
+    print("                    RUNNING BENCHMARK EXAMPLES")
+    print(f" Model: {detector.model_path} ({detector.file_size_mb:.2f} MB)")
+    print("=" * 75)
+
+    results = detector.predict(examples)
+    for i, res in enumerate(results, 1):
+        print(f"\n[Example {i}]")
+        print(f"  Input       : \"{res['text']}\"")
+        print(f"  Prediction  : {res['label_name']} (Class {res['label_id']})")
+        print(f"  Confidence  : {res['confidence'] * 100:.2f}%")
+        print(f"  Latency     : {res['latency_ms']:.2f} ms")
+        print("  Breakdown   :")
+        for label, prob in res["probabilities"].items():
+            bar = "#" * int(prob * 20)
+            print(f"    - {label:<20}: {prob * 100:>5.1f}% | {bar}")
+
+    print("\n" + "=" * 75)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Hate Speech Detection with ONNX Runtime.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help="Path to ONNX model file. If not set, automatically prefers INT8 or raw ONNX in ./model",
+    )
+    parser.add_argument(
+        "--raw",
+        action="store_true",
+        help="Use raw ONNX model (./model/hatespeech.onnx) instead of INT8 quantized model",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default=None,
+        help="Specific text to classify via CLI",
+    )
+    parser.add_argument(
+        "--interactive",
+        action="store_true",
+        help="Start interactive prompt loop to test custom sentences",
+    )
+    parser.add_argument(
+        "--gpu",
+        action="store_true",
+        help="Use CUDA Execution Provider if available",
+    )
+    args = parser.parse_args()
+
+    # Determine model path
+    if args.model:
+        model_path = args.model
+    elif args.raw:
+        model_path = "./model/hatespeech.onnx"
+    else:
+        # Default to INT8 if present, else fallback to raw
+        if os.path.exists("./model/hatespeech_int8.onnx"):
+            model_path = "./model/hatespeech_int8.onnx"
+        else:
+            model_path = "./model/hatespeech.onnx"
+
+    try:
+        detector = HateSpeechDetector(
+            model_path=model_path,
+            tokenizer_name_or_path="./model",
+            use_gpu=args.gpu,
+        )
+    except FileNotFoundError as e:
+        print(f"\nError: {e}")
+        sys.exit(1)
+
+    # 1. Single text prediction via CLI
+    if args.text:
+        res = detector.predict(args.text)[0]
+        print("\n" + "=" * 50)
+        print(f"Text       : \"{res['text']}\"")
+        print(f"Prediction : {res['label_name']} (Class {res['label_id']})")
+        print(f"Confidence : {res['confidence'] * 100:.2f}%")
+        print(f"Latency    : {res['latency_ms']:.2f} ms")
+        print("Probabilities:")
+        for label, prob in res["probabilities"].items():
+            print(f"  {label:<20}: {prob * 100:.1f}%")
+        print("=" * 50)
+        return
+
+    # 2. Interactive prompt loop
+    if args.interactive:
+        print("\n" + "=" * 60)
+        print("  INTERACTIVE HATE SPEECH DETECTION (Type 'exit' to quit)")
+        print(f"  Model: {model_path}")
+        print("=" * 60)
+        while True:
+            try:
+                user_input = input("\nEnter text: ").strip()
+                if not user_input:
+                    continue
+                if user_input.lower() in ("exit", "quit", "q"):
+                    print("Exiting.")
+                    break
+
+                res = detector.predict(user_input)[0]
+                print(f" -> Result     : {res['label_name']} ({res['confidence'] * 100:.1f}%)")
+                print(f" -> Latency    : {res['latency_ms']:.2f} ms")
+                print(f" -> Breakdown  : " + ", ".join([f"{k}: {v*100:.1f}%" for k, v in res['probabilities'].items()]))
+            except (KeyboardInterrupt, EOFError):
+                print("\nExiting.")
+                break
+        return
+
+    # 3. Default: Run benchmark examples
+    run_examples(detector)
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,15 @@
+[project]
+name = "hatespeech"
+version = "0.1.0"
+description = "Hate speech and offensive language classifier with ONNX & 8q export"
+readme = "README.md"
+requires-python = ">=3.9"
+dependencies = [
+    "torch>=2.0.0",
+    "transformers>=4.30.0",
+    "scikit-learn>=1.2.0",
+    "pandas>=2.0.0",
+    "numpy>=1.24.0",
+    "onnx>=1.14.0",
+    "onnxruntime>=1.15.0",
+]

requirements.txt

@@ -0,0 +1,335 @@
+# This file was autogenerated by uv via the following command:
+#    uv export --no-hashes --output-file requirements.txt --no-emit-project
+annotated-doc==0.0.5 ; python_full_version >= '3.10'
+    # via typer
+anyio==4.15.0 ; python_full_version >= '3.10'
+    # via httpx
+certifi==2026.7.22
+    # via
+    #   httpcore
+    #   httpx
+    #   requests
+charset-normalizer==3.5.1 ; python_full_version < '3.10'
+    # via requests
+click==8.5.0 ; python_full_version >= '3.10'
+    # via huggingface-hub
+cloudpickle==3.1.2 ; python_full_version >= '3.10'
+    # via joblib
+colorama==0.4.6 ; sys_platform == 'win32'
+    # via
+    #   tqdm
+    #   typer
+coloredlogs==15.0.1 ; python_full_version < '3.10'
+    # via onnxruntime
+cuda-bindings==13.3.1 ; python_full_version >= '3.10' and python_full_version < '3.15' and sys_platform == 'linux'
+    # via torch
+cuda-pathfinder==1.8.1 ; python_full_version >= '3.10' and python_full_version < '3.15' and sys_platform == 'linux'
+    # via cuda-bindings
+cuda-toolkit==13.0.3.0 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via torch
+exceptiongroup==1.3.1 ; python_full_version == '3.10.*'
+    # via anyio
+filelock==3.19.1 ; python_full_version < '3.10'
+    # via
+    #   huggingface-hub
+    #   torch
+    #   transformers
+filelock==3.32.5 ; python_full_version >= '3.10'
+    # via
+    #   huggingface-hub
+    #   torch
+flatbuffers==25.12.19
+    # via onnxruntime
+fsspec==2025.10.0 ; python_full_version < '3.10'
+    # via
+    #   huggingface-hub
+    #   torch
+fsspec==2026.7.0 ; python_full_version >= '3.10'
+    # via
+    #   huggingface-hub
+    #   torch
+h11==0.16.0 ; python_full_version >= '3.10'
+    # via httpcore
+hf-xet==1.6.0 ; (python_full_version >= '3.10' and platform_machine == 'AMD64') or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'
+    # via huggingface-hub
+httpcore==1.0.9 ; python_full_version >= '3.10'
+    # via httpx
+httpx==0.28.1 ; python_full_version >= '3.10'
+    # via huggingface-hub
+huggingface-hub==0.36.2 ; python_full_version < '3.10'
+    # via
+    #   tokenizers
+    #   transformers
+huggingface-hub==1.30.0 ; python_full_version >= '3.10'
+    # via
+    #   tokenizers
+    #   transformers
+humanfriendly==10.0 ; python_full_version < '3.10'
+    # via coloredlogs
+idna==3.19
+    # via
+    #   anyio
+    #   httpx
+    #   requests
+importlib-metadata==8.7.1 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via triton
+jinja2==3.1.6
+    # via torch
+joblib==1.5.3 ; python_full_version < '3.10'
+    # via scikit-learn
+joblib==1.6.0 ; python_full_version >= '3.10'
+    # via scikit-learn
+markdown-it-py==4.2.0 ; python_full_version >= '3.10'
+    # via rich
+markupsafe==3.0.3
+    # via jinja2
+mdurl==0.1.2 ; python_full_version >= '3.10'
+    # via markdown-it-py
+ml-dtypes==0.5.4 ; python_full_version < '3.10'
+    # via onnx
+ml-dtypes==0.6.0 ; python_full_version >= '3.10'
+    # via onnx
+mpmath==1.3.0
+    # via sympy
+narwhals==2.25.0 ; python_full_version >= '3.11'
+    # via scikit-learn
+networkx==3.2.1 ; python_full_version < '3.10'
+    # via torch
+networkx==3.4.2 ; python_full_version == '3.10.*'
+    # via torch
+networkx==3.6.1 ; python_full_version >= '3.11'
+    # via torch
+numpy==2.0.2 ; python_full_version < '3.10'
+    # via
+    #   hatespeech
+    #   ml-dtypes
+    #   onnx
+    #   onnxruntime
+    #   pandas
+    #   scikit-learn
+    #   scipy
+    #   transformers
+numpy==2.2.6 ; python_full_version == '3.10.*'
+    # via
+    #   hatespeech
+    #   ml-dtypes
+    #   onnx
+    #   onnxruntime
+    #   pandas
+    #   scikit-learn
+    #   scipy
+    #   transformers
+numpy==2.4.6 ; python_full_version == '3.11.*'
+    # via
+    #   hatespeech
+    #   ml-dtypes
+    #   onnx
+    #   onnxruntime
+    #   pandas
+    #   scikit-learn
+    #   scipy
+    #   transformers
+numpy==2.5.2 ; python_full_version >= '3.12'
+    # via
+    #   hatespeech
+    #   ml-dtypes
+    #   onnx
+    #   onnxruntime
+    #   pandas
+    #   scikit-learn
+    #   scipy
+    #   transformers
+nvidia-cublas==13.1.1.3 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via
+    #   cuda-toolkit
+    #   nvidia-cudnn-cu13
+    #   nvidia-cusolver
+nvidia-cublas-cu12==12.8.4.1 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via
+    #   nvidia-cudnn-cu12
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cuda-cupti==13.0.85 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-cuda-cupti-cu12==12.8.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cuda-nvrtc==13.0.88 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via
+    #   cuda-toolkit
+    #   nvidia-cublas
+nvidia-cuda-nvrtc-cu12==12.8.93 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cuda-runtime==13.0.96 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-cuda-runtime-cu12==12.8.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cudnn-cu12==9.10.2.21 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cudnn-cu13==9.24.0.43 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via torch
+nvidia-cufft==12.0.0.61 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-cufft-cu12==11.3.3.83 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cufile==1.15.1.6 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-cufile-cu12==1.13.1.3 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-curand==10.4.0.35 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-curand-cu12==10.3.9.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cusolver==12.0.4.66 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-cusolver-cu12==11.7.3.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cusparse==12.6.3.3 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via
+    #   cuda-toolkit
+    #   nvidia-cusolver
+nvidia-cusparse-cu12==12.5.8.93 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via
+    #   nvidia-cusolver-cu12
+    #   torch
+nvidia-cusparselt-cu12==0.7.1 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-cusparselt-cu13==0.8.1 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via torch
+nvidia-nccl-cu12==2.27.3 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+nvidia-nccl-cu13==2.30.7 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via torch
+nvidia-nvjitlink==13.3.33 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via
+    #   cuda-toolkit
+    #   nvidia-cufft
+    #   nvidia-cusolver
+    #   nvidia-cusparse
+nvidia-nvjitlink-cu12==12.8.93 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via
+    #   nvidia-cufft-cu12
+    #   nvidia-cusolver-cu12
+    #   nvidia-cusparse-cu12
+    #   torch
+nvidia-nvshmem-cu13==3.4.5 ; python_full_version >= '3.10' and sys_platform == 'linux'
+    # via torch
+nvidia-nvtx==13.0.85 ; (python_full_version >= '3.10' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version >= '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux')
+    # via cuda-toolkit
+nvidia-nvtx-cu12==12.8.90 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+onnx==1.19.1 ; python_full_version < '3.10'
+    # via hatespeech
+onnx==1.22.0 ; python_full_version >= '3.10'
+    # via hatespeech
+onnxruntime==1.20.1 ; python_full_version < '3.10'
+    # via hatespeech
+onnxruntime==1.24.3 ; python_full_version == '3.10.*'
+    # via hatespeech
+onnxruntime==1.29.0 ; python_full_version >= '3.11'
+    # via hatespeech
+packaging==26.3
+    # via
+    #   huggingface-hub
+    #   onnxruntime
+    #   transformers
+pandas==2.3.3 ; python_full_version < '3.11'
+    # via hatespeech
+pandas==3.0.5 ; python_full_version >= '3.11'
+    # via hatespeech
+protobuf==6.33.6 ; python_full_version < '3.10'
+    # via
+    #   onnx
+    #   onnxruntime
+protobuf==7.36.1 ; python_full_version >= '3.10'
+    # via
+    #   onnx
+    #   onnxruntime
+pygments==2.21.0 ; python_full_version >= '3.10'
+    # via rich
+pyreadline3==3.5.6 ; python_full_version < '3.10' and sys_platform == 'win32'
+    # via humanfriendly
+python-dateutil==2.9.0.post0
+    # via pandas
+pytz==2026.3.post1 ; python_full_version < '3.11'
+    # via pandas
+pyyaml==6.0.3
+    # via
+    #   huggingface-hub
+    #   transformers
+regex==2026.1.15 ; python_full_version < '3.10'
+    # via transformers
+regex==2026.9.3 ; python_full_version >= '3.10'
+    # via transformers
+requests==2.32.5 ; python_full_version < '3.10'
+    # via
+    #   huggingface-hub
+    #   transformers
+rich==15.0.0 ; python_full_version >= '3.10'
+    # via typer
+safetensors==0.7.0 ; python_full_version < '3.10'
+    # via transformers
+safetensors==0.8.0 ; python_full_version >= '3.10'
+    # via transformers
+scikit-learn==1.6.1 ; python_full_version < '3.10'
+    # via hatespeech
+scikit-learn==1.7.2 ; python_full_version == '3.10.*'
+    # via hatespeech
+scikit-learn==1.9.0 ; python_full_version >= '3.11'
+    # via hatespeech
+scipy==1.13.1 ; python_full_version < '3.10'
+    # via scikit-learn
+scipy==1.15.3 ; python_full_version == '3.10.*'
+    # via scikit-learn
+scipy==1.17.1 ; python_full_version == '3.11.*'
+    # via scikit-learn
+scipy==1.18.1 ; python_full_version >= '3.12'
+    # via scikit-learn
+setuptools==82.0.1 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via triton
+setuptools==84.0.0 ; python_full_version >= '3.10'
+    # via torch
+shellingham==1.5.4 ; python_full_version >= '3.10'
+    # via typer
+six==1.17.0
+    # via python-dateutil
+sympy==1.14.0
+    # via
+    #   onnxruntime
+    #   torch
+threadpoolctl==3.6.0
+    # via scikit-learn
+tokenizers==0.22.2 ; python_full_version < '3.10'
+    # via transformers
+tokenizers==0.23.2 ; python_full_version >= '3.10'
+    # via transformers
+torch==2.8.0 ; python_full_version < '3.10'
+    # via hatespeech
+torch==2.14.0 ; python_full_version >= '3.10'
+    # via hatespeech
+tqdm==4.70.0
+    # via
+    #   huggingface-hub
+    #   transformers
+transformers==4.57.6 ; python_full_version < '3.10'
+    # via hatespeech
+transformers==5.16.1 ; python_full_version >= '3.10'
+    # via hatespeech
+triton==3.4.0 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via torch
+triton==3.8.0 ; python_full_version >= '3.10' and python_full_version < '3.15' and sys_platform == 'linux'
+    # via torch
+typer==0.27.2 ; python_full_version >= '3.10'
+    # via transformers
+typing-extensions==4.16.0
+    # via
+    #   anyio
+    #   exceptiongroup
+    #   huggingface-hub
+    #   onnx
+    #   torch
+tzdata==2026.3 ; python_full_version < '3.11' or sys_platform == 'emscripten' or sys_platform == 'win32'
+    # via pandas
+urllib3==2.6.3 ; python_full_version < '3.10'
+    # via requests
+zipp==3.23.1 ; python_full_version < '3.10' and platform_machine == 'x86_64' and sys_platform == 'linux'
+    # via importlib-metadata

train.py

@@ -0,0 +1,457 @@
+#!/usr/bin/env python3
+"""
+train.py - Hate Speech & Offensive Language Classifier
+Trains a PyTorch transformer model on Davidson et al. labeled_data.csv.
+Exports both raw ONNX and INT8 quantized (8q) ONNX models.
+
+Key Features:
+- Downloads dataset to a temporary directory and cleans it up after loading.
+- Preprocessing: unescapes HTML, strips URLs and Twitter mentions.
+- Prevents Overfitting & Underfitting:
+  * Stratified 80/10/10 train/val/test split.
+  * Balanced class weights for CrossEntropyLoss (combats ~5% hate speech imbalance).
+  * Weight decay (AdamW) + Linear learning rate warmup & decay.
+  * Early stopping tracking validation Macro F1 score.
+- Epoch Optimization:
+  * Sequence length truncated to 128 (optimal for tweets, fast epochs).
+  * PyTorch Automatic Mixed Precision (AMP / FP16) on CUDA.
+- Model Exports:
+  * Raw ONNX (FP32)
+  * Quantized INT8 ONNX (dynamic quantization for ~75% size reduction & fast inference)
+  * Tokenizer files saved alongside ONNX models for standalone portability.
+"""
+
+import os
+import sys
+import re
+import html
+import time
+import shutil
+import tempfile
+import urllib.request
+import argparse
+from typing import Tuple, Dict, Any
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torch.optim import AdamW
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    get_linear_schedule_with_warmup,
+)
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import classification_report, f1_score, accuracy_score
+import onnx
+import onnxruntime
+from onnxruntime.quantization import quantize_dynamic, QuantType
+
+# Label mappings
+LABEL_NAMES = {
+    0: "Hate Speech",
+    1: "Offensive Language",
+    2: "Neither",
+}
+
+DATASET_URL = (
+    "https://raw.githubusercontent.com/t-davidson/hate-speech-and-offensive-language/"
+    "master/data/labeled_data.csv"
+)
+
+
+def clean_text(text: str) -> str:
+    """Cleans tweet text by unescaping HTML, normalizing handles, and stripping URLs."""
+    if not isinstance(text, str):
+        return ""
+    # Decode HTML entities (e.g., &amp; -> &, &lt; -> <)
+    text = html.unescape(text)
+    # Remove URLs
+    text = re.sub(r"https?://\S+|www\.\S+", "", text)
+    # Remove user mentions
+    text = re.sub(r"@\w+", "", text)
+    # Normalize excessive whitespaces
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+
+def download_dataset_temp(url: str = DATASET_URL) -> pd.DataFrame:
+    """
+    Downloads labeled_data.csv into a temporary directory, loads it into pandas,
+    and cleans up the temporary directory immediately after.
+    """
+    print(f"\n[1/6] Fetching dataset from: {url}")
+    with tempfile.TemporaryDirectory() as temp_dir:
+        temp_file_path = os.path.join(temp_dir, "labeled_data.csv")
+        print(f" -> Temporary download location: {temp_file_path}")
+
+        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
+        with urllib.request.urlopen(req) as response, open(temp_file_path, "wb") as out_file:
+            shutil.copyfileobj(response, out_file)
+
+        file_size_mb = os.path.getsize(temp_file_path) / (1024 * 1024)
+        print(f" -> Download completed successfully ({file_size_mb:.2f} MB).")
+
+        df = pd.read_csv(temp_file_path)
+        # Required columns: 'class' (0, 1, 2) and 'tweet'
+        df = df[["class", "tweet"]].dropna()
+        df["tweet"] = df["tweet"].apply(clean_text)
+        df = df[df["tweet"].str.strip() != ""].reset_index(drop=True)
+
+        print(f" -> Loaded {len(df):,} valid samples.")
+        class_counts = df["class"].value_counts().sort_index()
+        for cls_id, count in class_counts.items():
+            pct = count / len(df) * 100
+            print(f"    Class {cls_id} ({LABEL_NAMES.get(cls_id, 'Unknown')}): {count:,} ({pct:.1f}%)")
+
+        print(" -> Temporary directory and file automatically cleaned up.")
+    return df
+
+
+class HateSpeechDataset(Dataset):
+    """PyTorch Dataset for tokenized hate speech data."""
+
+    def __init__(self, texts, labels, tokenizer, max_length: int = 128):
+        self.encodings = tokenizer(
+            texts,
+            truncation=True,
+            padding=True,
+            max_length=max_length,
+            return_tensors="pt",
+        )
+        self.labels = torch.tensor(labels, dtype=torch.long)
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        item = {key: val[idx] for key, val in self.encodings.items()}
+        item["labels"] = self.labels[idx]
+        return item
+
+
+def compute_class_weights(labels: np.ndarray, num_classes: int = 3) -> torch.Tensor:
+    """
+    Computes balanced class weights inversely proportional to class frequencies.
+    Crucial to prevent under-learning the minority Hate Speech class (~5.7%).
+    """
+    counts = np.bincount(labels, minlength=num_classes)
+    total_samples = len(labels)
+    # Balanced weighting: total / (num_classes * count)
+    weights = total_samples / (num_classes * counts.astype(np.float32))
+    # Normalize so mean weight is 1.0 to preserve learning rate stability
+    weights = weights / np.mean(weights)
+    return torch.tensor(weights, dtype=torch.float)
+
+
+def train_epoch(
+    model: nn.Module,
+    dataloader: DataLoader,
+    optimizer: optim.Optimizer,
+    scheduler: Any,
+    scaler: torch.amp.GradScaler,
+    criterion: nn.Module,
+    device: torch.device,
+) -> float:
+    """Trains the model for 1 epoch with optional AMP mixed-precision."""
+    model.train()
+    total_loss = 0.0
+    use_amp = (device.type == "cuda")
+
+    for batch in dataloader:
+        optimizer.zero_grad()
+        input_ids = batch["input_ids"].to(device)
+        attention_mask = batch["attention_mask"].to(device)
+        labels = batch["labels"].to(device)
+
+        if use_amp:
+            with torch.amp.autocast(device_type="cuda", dtype=torch.float16):
+                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+                loss = criterion(outputs.logits, labels)
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, labels)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+            optimizer.step()
+
+        scheduler.step()
+        total_loss += loss.item()
+
+    return total_loss / len(dataloader)
+
+
+def evaluate(
+    model: nn.Module,
+    dataloader: DataLoader,
+    criterion: nn.Module,
+    device: torch.device,
+) -> Tuple[float, float, float, np.ndarray, np.ndarray]:
+    """Evaluates the model and computes validation loss, accuracy, and Macro F1."""
+    model.eval()
+    total_loss = 0.0
+    all_preds = []
+    all_labels = []
+
+    with torch.no_grad():
+        for batch in dataloader:
+            input_ids = batch["input_ids"].to(device)
+            attention_mask = batch["attention_mask"].to(device)
+            labels = batch["labels"].to(device)
+
+            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
+            loss = criterion(outputs.logits, labels)
+            total_loss += loss.item()
+
+            preds = torch.argmax(outputs.logits, dim=1).cpu().numpy()
+            all_preds.extend(preds)
+            all_labels.extend(labels.cpu().numpy())
+
+    avg_loss = total_loss / len(dataloader)
+    all_preds = np.array(all_preds)
+    all_labels = np.array(all_labels)
+    acc = accuracy_score(all_labels, all_preds)
+    macro_f1 = f1_score(all_labels, all_preds, average="macro")
+
+    return avg_loss, acc, macro_f1, all_preds, all_labels
+
+
+def export_to_onnx(
+    model: nn.Module,
+    tokenizer: Any,
+    output_dir: str,
+    device: torch.device,
+    max_length: int = 128,
+) -> Tuple[str, str]:
+    """
+    Exports PyTorch model to:
+    1. Raw ONNX format (`hatespeech.onnx`) with dynamic batch and sequence length axes.
+    2. Quantized INT8 ONNX (`hatespeech_int8.onnx`) for CPU/edge speed and small file size.
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    raw_onnx_path = os.path.join(output_dir, "hatespeech.onnx")
+    int8_onnx_path = os.path.join(output_dir, "hatespeech_int8.onnx")
+
+    model.eval()
+    model.to("cpu")  # Export on CPU for universal ONNX portability
+
+    print("\n[5/6] Exporting to Raw ONNX...")
+    dummy_text = "Sample sentence for hate speech detection ONNX export."
+    dummy_encodings = tokenizer(
+        dummy_text,
+        max_length=max_length,
+        padding="max_length",
+        truncation=True,
+        return_tensors="pt",
+    )
+    dummy_inputs = (
+        dummy_encodings["input_ids"],
+        dummy_encodings["attention_mask"],
+    )
+
+    torch.onnx.export(
+        model,
+        dummy_inputs,
+        raw_onnx_path,
+        input_names=["input_ids", "attention_mask"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input_ids": {0: "batch_size", 1: "sequence_length"},
+            "attention_mask": {0: "batch_size", 1: "sequence_length"},
+            "logits": {0: "batch_size"},
+        },
+        opset_version=14,
+        do_constant_folding=True,
+    )
+
+    # Verify raw ONNX model validity
+    onnx_model = onnx.load(raw_onnx_path)
+    onnx.checker.check_model(onnx_model)
+    raw_size_mb = os.path.getsize(raw_onnx_path) / (1024 * 1024)
+    print(f" -> Raw ONNX saved to: {raw_onnx_path} ({raw_size_mb:.2f} MB)")
+
+    print("\n[6/6] Quantizing ONNX model to INT8 (8q dynamic)...")
+    quantize_dynamic(
+        model_input=raw_onnx_path,
+        model_output=int8_onnx_path,
+        weight_type=QuantType.QInt8,
+    )
+    int8_size_mb = os.path.getsize(int8_onnx_path) / (1024 * 1024)
+    compression = (1.0 - (int8_size_mb / raw_size_mb)) * 100
+    print(f" -> INT8 ONNX saved to: {int8_onnx_path} ({int8_size_mb:.2f} MB, {compression:.1f}% smaller)")
+
+    # Save tokenizer files alongside models for offline independence
+    tokenizer.save_pretrained(output_dir)
+    print(f" -> Tokenizer assets saved in: {output_dir}")
+
+    return raw_onnx_path, int8_onnx_path
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Hate Speech Classifier and Export to ONNX & INT8.")
+    parser.add_argument("--model_name", type=str, default="distilbert-base-uncased",
+                        help="Pretrained Hugging Face model architecture (default: distilbert-base-uncased)")
+    parser.add_argument("--epochs", type=int, default=3,
+                        help="Maximum training epochs (default: 3, prevents overfitting)")
+    parser.add_argument("--batch_size", type=int, default=32,
+                        help="Training batch size (default: 32)")
+    parser.add_argument("--lr", type=float, default=2e-5,
+                        help="Learning rate for AdamW (default: 2e-5)")
+    parser.add_argument("--max_length", type=int, default=128,
+                        help="Max sequence token length for tweets (default: 128 for fast epoch time)")
+    parser.add_argument("--patience", type=int, default=2,
+                        help="Early stopping patience epochs based on Macro F1 (default: 2)")
+    parser.add_argument("--output_dir", type=str, default="./model",
+                        help="Directory to save ONNX models and tokenizer (default: ./model)")
+    args = parser.parse_args()
+
+    # Hardware detection
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print("=" * 65)
+    print("      HATE SPEECH & OFFENSIVE LANGUAGE CLASSIFIER TRAINING")
+    print("=" * 65)
+    print(f"Using Compute Device   : {device}")
+    if device.type == "cuda":
+        print(f"GPU Model              : {torch.cuda.get_device_name(0)}")
+        print(f"Automatic Mixed Prec   : Enabled (FP16)")
+    print(f"Base Architecture      : {args.model_name}")
+    print(f"Max Sequence Length    : {args.max_length} tokens")
+    print(f"Batch Size / Max Epochs: {args.batch_size} / {args.epochs}")
+    print("=" * 65)
+
+    # 1. Download & Clean Dataset
+    df = download_dataset_temp(DATASET_URL)
+    texts = df["tweet"].tolist()
+    labels = df["class"].to_numpy()
+
+    # 2. Stratified Data Split (80% Train, 10% Validation, 10% Test)
+    print("\n[2/6] Preparing stratified data splits (80% Train, 10% Val, 10% Test)...")
+    train_texts, temp_texts, train_labels, temp_labels = train_test_split(
+        texts, labels, test_size=0.20, random_state=42, stratify=labels
+    )
+    val_texts, test_texts, val_labels, test_labels = train_test_split(
+        temp_texts, temp_labels, test_size=0.50, random_state=42, stratify=temp_labels
+    )
+    print(f" -> Train split : {len(train_labels):,} samples")
+    print(f" -> Val split   : {len(val_labels):,} samples")
+    print(f" -> Test split  : {len(test_labels):,} samples")
+
+    # 3. Tokenizer & DataLoaders
+    print(f"\n[3/6] Loading tokenizer '{args.model_name}' and tokenizing...")
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+
+    train_dataset = HateSpeechDataset(train_texts, train_labels, tokenizer, max_length=args.max_length)
+    val_dataset = HateSpeechDataset(val_texts, val_labels, tokenizer, max_length=args.max_length)
+    test_dataset = HateSpeechDataset(test_texts, test_labels, tokenizer, max_length=args.max_length)
+
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
+    val_loader = DataLoader(val_dataset, batch_size=args.batch_size * 2, shuffle=False)
+    test_loader = DataLoader(test_dataset, batch_size=args.batch_size * 2, shuffle=False)
+
+    # 4. Model Setup & Anti-Overfitting Controls
+    print("\n[4/6] Initializing model & regularization controls...")
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_name,
+        num_labels=3,
+        classifier_dropout=0.2,
+    )
+    model.to(device)
+
+    # Balanced class weights to counteract the severe minority of class 0
+    class_weights = compute_class_weights(train_labels, num_classes=3).to(device)
+    print(f" -> Computed balanced class weights: {class_weights.cpu().numpy().round(3)}")
+    criterion = nn.CrossEntropyLoss(weight=class_weights)
+
+    # Optimizer with weight decay to block overfitting
+    optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=0.01)
+
+    # Linear warmup scheduler
+    total_steps = len(train_loader) * args.epochs
+    warmup_steps = int(0.1 * total_steps)
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=warmup_steps, num_training_steps=total_steps
+    )
+    scaler = torch.amp.GradScaler("cuda", enabled=(device.type == "cuda"))
+
+    # Training Loop with Early Stopping monitoring Macro F1
+    best_val_macro_f1 = -1.0
+    best_model_state = None
+    patience_counter = 0
+
+    print("\n" + "-" * 75)
+    print(f"{'Epoch':<8}{'Train Loss':<14}{'Val Loss':<12}{'Val Acc':<12}{'Val Macro F1':<16}{'Time (s)':<10}")
+    print("-" * 75)
+
+    for epoch in range(1, args.epochs + 1):
+        start_time = time.time()
+        train_loss = train_epoch(
+            model, train_loader, optimizer, scheduler, scaler, criterion, device
+        )
+        val_loss, val_acc, val_f1, _, _ = evaluate(model, val_loader, criterion, device)
+        elapsed = time.time() - start_time
+
+        print(f"{epoch:<8}{train_loss:<14.4f}{val_loss:<12.4f}{val_acc*100:<12.2f}%{val_f1:<16.4f}{elapsed:<10.1f}")
+
+        # Check for best model based on Macro F1 (prevents skew towards majority class)
+        if val_f1 > best_val_macro_f1:
+            best_val_macro_f1 = val_f1
+            best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+            patience_counter = 0
+            improved_mark = " (Best Saved)"
+        else:
+            patience_counter += 1
+            improved_mark = f" (Patience {patience_counter}/{args.patience})"
+
+        print(f"         Status: Macro F1 = {val_f1:.4f}{improved_mark}")
+
+        if patience_counter >= args.patience:
+            print(f"\n[!] Early stopping triggered at epoch {epoch} to prevent overfitting.")
+            break
+
+    print("-" * 75)
+
+    # Load best checkpoint
+    if best_model_state is not None:
+        model.load_state_dict(best_model_state)
+        print(f"Loaded best checkpoint with Validation Macro F1: {best_val_macro_f1:.4f}")
+
+    # Final Evaluation on Independent Test Set
+    print("\n" + "=" * 65)
+    print("            FINAL EVALUATION ON TEST SET (10% HOLDOUT)")
+    print("=" * 65)
+    _, test_acc, test_macro_f1, test_preds, test_targets = evaluate(
+        model, test_loader, criterion, device
+    )
+    print(f"Overall Test Accuracy: {test_acc * 100:.2f}%")
+    print(f"Test Macro F1 Score  : {test_macro_f1:.4f}\n")
+    print("Detailed Classification Report:")
+    target_names = [f"{i}: {LABEL_NAMES[i]}" for i in range(3)]
+    print(classification_report(test_targets, test_preds, target_names=target_names, digits=4))
+
+    # 5 & 6. Export to Raw ONNX & 8q Quantized ONNX
+    export_to_onnx(
+        model=model,
+        tokenizer=tokenizer,
+        output_dir=args.output_dir,
+        device=device,
+        max_length=args.max_length,
+    )
+
+    print("\n" + "=" * 65)
+    print("                      TRAINING COMPLETE!")
+    print(f"Saved artifacts inside '{args.output_dir}':")
+    print(f" 1. Raw ONNX Model     : {os.path.join(args.output_dir, 'hatespeech.onnx')}")
+    print(f" 2. Quantized INT8 ONNX: {os.path.join(args.output_dir, 'hatespeech_int8.onnx')}")
+    print(f" 3. Tokenizer Assets   : {args.output_dir}/")
+    print("=" * 65)
+
+
+if __name__ == "__main__":
+    main()