Pommu/threat-detection-jigsaw
08
๐ก๏ธ Context-Aware Threat & Compliance Detection in Conversational Text
NLP Course Project โ Individual Assignment Model: Fine-tuned roberta-base | Task: 6-class Multi-label Classification Dataset: google/jigsaw_toxicity_pred (~150K samples)
๐ Live Demo
๐ Problem Statement
Modern organizations face increasing risks from implicit threats and workplace harassment. Unlike simple slur-detectors, this project focuses on intent-based detection. It uses the Jigsaw Toxic Comment dataset to identify professional and personal risk across 6 binary categories:
๐๏ธ Architecture: The 2-Layer Safety Net
[User Input: Text / Chat Transcript]
โ
โผ
[Gradio UI] โโ HTTP POST /predict โโโถ [FastAPI Backend]
โ
โโโโโโโโโโโโโโดโโโโโโโโโโโโโ
โ โ
[Layer 1] [Layer 2]
Fine-tuned RoBERTa Lexical Booster
(6 Jigsaw labels) (keyword pattern match)
โ โ
โโโโโโโโโโโโโโฌโโโโโโโโโโโโโ
โ
{label, confidence, risk_level}- Layer 1 โ Weighted RoBERTa: Multi-label model fine-tuned with 50x class weights on the rare
threatclass (0.3% of data) to maximize recall. - Layer 2 โ Lexical Booster: Backend safety layer that catches ominous linguistic patterns (e.g., "watch your back", "you'll regret") unconditionally, regardless of model score.
๐ Evaluation Results
Design Choice: We prioritize Recall over Precision for threats. A false positive (flagging safe text) is far less costly than a false negative (missing a real threat).
๐ API Usage
# Test implicit threat detection locally
curl -X POST http://localhost:8000/predict \
-H "Content-Type: application/json" \
-d '{"text": "I will find you and you will regret every decision you have made."}'JSON Response:
{
"prediction": "Threatening",
"confidence": 0.85,
"is_threat": true,
"risk_level": "HIGH",
"all_scores": [
{"label": "Toxic", "confidence": 0.0009},
{"label": "Threatening", "confidence": 0.85},
...
]
}๐ Project Structure
NLP Course Project/
โโโ notebook/
โ โโโ train_3.ipynb โ Colab training notebook (WeightedTrainer)
โโโ app/
โ โโโ main.py โ FastAPI backend + Lexical Booster
โ โโโ model_loader.py โ Singleton model loader (thread-safe)
โ โโโ schemas.py โ Pydantic request/response models
โ โโโ demo.py โ Gradio UI (standalone OR API-backed)
โโโ spaces/
โ โโโ app.py โ HuggingFace Spaces entry point
โ โโโ requirements.txt
โโโ requirements.txt โ API/demo dependencies
โโโ requirements_training.txt โ Training dependencies (Colab)
โโโ setup_local.bat โ Windows one-click environment setup
โโโ README.mdโ๏ธ Local Setup (Windows)
# Option A: one-click setup
setup_local.bat
# Option B: manual
python -m venv .venv
.venv\Scripts\activate
pip install -r requirements.txtSet your .env file:
HF_TOKEN=your_token_here
MODEL_NAME=Pommu/threat-detection-jigsaw๐ Running Locally
# Terminal 1 โ Start the FastAPI backend
uvicorn app.main:app --reload
# Terminal 2 โ Launch the Gradio demo
python app/demo.pyOpen http://localhost:7860 to see the UI.
โ๏ธ Training on Colab
- Open
notebook/train_3.ipynbin Google Colab - Enable GPU accelerator (Tesla T4)
- Add
HF_Tokento Colab Secrets (๐ icon) - Run all cells โ model auto-pushes to HuggingFace Hub
๐ Deploy to HuggingFace Spaces
- Go to huggingface.co/new-space โ SDK: Gradio
- In Space Settings โ Variables:
MODEL_NAME = Pommu/threat-detection-jigsaw - Push
spaces/contents to the Space repo
