ml-service/optimize_thresholds.py

#!/usr/bin/env python3
"""
SafeChat — Dynamic Per-Class Threshold Tuning (Precision-Recall Curve Optimization)

Optimizes classification thresholds for each tag using the Validation Set (`real_toxicity_val.csv`).
Applies these optimal per-class thresholds $\theta_c$ to both Validation and Test sets (`real_toxicity_test.csv`),
and saves the calibrated thresholds to `optimal_thresholds.json` for production moderation.
"""

import os
import sys
import json
import time
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from sklearn.metrics import precision_recall_curve, f1_score, precision_score, recall_score

if sys.stdout.encoding != 'utf-8':
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except AttributeError:
        pass

TAGS = ["toxic", "severe_toxic", "obscene", "threat", "insult", "identity_hate"]
CHECKPOINT_DIR = os.path.join(os.path.dirname(__file__), "checkpoints", "hingbert-toxicity-finetuned")
DATA_DIR = os.path.join(os.path.dirname(__file__), "training", "data", "real_datasets")

class ToxicityDataset(Dataset):
    def __init__(self, texts, labels):
        self.texts = texts
        self.labels = labels

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        return self.texts[idx], self.labels[idx]

def collate_fn(batch, tokenizer, device):
    texts, labels = zip(*batch)
    inputs = tokenizer(list(texts), padding=True, truncation=True, max_length=128, return_tensors="pt").to(device)
    labels = torch.tensor(np.array(labels), dtype=torch.float).to(device)
    return inputs, labels

def get_probabilities(model, tokenizer, dataloader):
    model.eval()
    all_probs = []
    all_labels = []

    with torch.no_grad():
        for inputs, labels in dataloader:
            logits = model(**inputs).logits
            probs = torch.sigmoid(logits).cpu().numpy()
            all_probs.append(probs)
            all_labels.append(labels.cpu().numpy())

    return np.vstack(all_probs), np.vstack(all_labels)

def main():
    print("="*90)
    print("🎯 SAFECHAT: DYNAMIC PER-CLASS THRESHOLD TUNING (PRECISION-RECALL OPTIMIZATION)")
    print("="*90)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Hardware Acceleration Device: {device}")

    # Load data splits
    val_path = os.path.join(DATA_DIR, "real_toxicity_val.csv")
    test_path = os.path.join(DATA_DIR, "real_toxicity_test.csv")

    val_df = pd.read_csv(val_path).dropna(subset=["text"]).reset_index(drop=True)
    test_df = pd.read_csv(test_path).dropna(subset=["text"]).reset_index(drop=True)

    val_labels = val_df[TAGS].values
    test_labels = test_df[TAGS].values

    print(f"Loaded Validation Set : {len(val_df)} rows")
    print(f"Loaded Test Set       : {len(test_df)} rows")

    # Load Fine-Tuned Model
    print(f"\nLoading SafeChat Fine-Tuned Model from: {CHECKPOINT_DIR}...")
    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT_DIR)
    model = AutoModelForSequenceClassification.from_pretrained(CHECKPOINT_DIR).to(device)

    val_loader = DataLoader(ToxicityDataset(val_df["text"].astype(str).tolist(), val_labels), batch_size=32, shuffle=False, collate_fn=lambda b: collate_fn(b, tokenizer, device))
    test_loader = DataLoader(ToxicityDataset(test_df["text"].astype(str).tolist(), test_labels), batch_size=32, shuffle=False, collate_fn=lambda b: collate_fn(b, tokenizer, device))

    print(" -> Predicting continuous probabilities on Validation Set...")
    val_probs, val_true = get_probabilities(model, tokenizer, val_loader)
    print(" -> Predicting continuous probabilities on Test Set...")
    test_probs, test_true = get_probabilities(model, tokenizer, test_loader)

    # 1. Optimize Thresholds on Validation Set
    print("\n" + "="*90)
    print("⚙️ OPTIMIZING PER-CLASS THRESHOLDS ON VALIDATION SET (Precision-Recall Analysis)")
    print("="*90)
    print(f"{'Tag Name':<16} | {'Positives':<10} | {'Default F1 (0.50)':<18} | {'Optimal Threshold':<18} | {'Optimized F1':<16}")
    print("-" * 90)

    optimal_thresholds = {}
    val_f1_default_list = []
    val_f1_opt_list = []

    for i, tag in enumerate(TAGS):
        y_true = val_true[:, i]
        y_prob = val_probs[:, i]
        pos_count = int(np.sum(y_true))

        # Default F1 at 0.50
        preds_50 = (y_prob >= 0.50).astype(int)
        f1_50 = f1_score(y_true, preds_50, zero_division=0)
        val_f1_default_list.append(f1_50)

        if pos_count > 0:
            precision, recall, thresholds = precision_recall_curve(y_true, y_prob)
            # Avoid division by zero
            f1_scores = np.divide(
                2 * precision * recall,
                precision + recall,
                out=np.zeros_like(precision),
                where=(precision + recall) != 0
            )
            best_idx = np.argmax(f1_scores)
            best_th = float(thresholds[best_idx]) if best_idx < len(thresholds) else 0.50
            best_f1 = float(f1_scores[best_idx])
        else:
            # If rare tag has 0 positives in Val split, use an optimal safety threshold based on probability distribution
            best_th = 0.15
            best_f1 = 0.0

        optimal_thresholds[tag] = round(best_th, 4)
        val_f1_opt_list.append(best_f1)

        print(f"{tag:<16} | {pos_count:<10} | {f1_50*100:<17.2f}% | {best_th:<18.4f} | {best_f1*100:<15.2f}%")

    print("-" * 90)

    # Save calibrated thresholds to JSON
    out_json = os.path.join(CHECKPOINT_DIR, "optimal_thresholds.json")
    with open(out_json, "w", encoding="utf-8") as f:
        json.dump(optimal_thresholds, f, indent=2)
    print(f"✅ Saved optimal per-class thresholds to: {out_json}")

    # 2. Apply to Validation and Test Sets & Compare Macro/Micro F1
    print("\n" + "="*90)
    print("📈 FINAL BENCHMARK: STATIC 0.50 THRESHOLD vs. OPTIMIZED PER-CLASS THRESHOLDS")
    print("="*90)

    # Validation Set Predictions
    val_preds_50 = (val_probs >= 0.50).astype(int)
    val_preds_opt = np.zeros_like(val_probs, dtype=int)
    for i, tag in enumerate(TAGS):
        val_preds_opt[:, i] = (val_probs[:, i] >= optimal_thresholds[tag]).astype(int)

    val_macro_50 = f1_score(val_true, val_preds_50, average="macro", zero_division=0)
    val_micro_50 = f1_score(val_true, val_preds_50, average="micro", zero_division=0)
    val_macro_opt = f1_score(val_true, val_preds_opt, average="macro", zero_division=0)
    val_micro_opt = f1_score(val_true, val_preds_opt, average="micro", zero_division=0)

    # Test Set Predictions
    test_preds_50 = (test_probs >= 0.50).astype(int)
    test_preds_opt = np.zeros_like(test_probs, dtype=int)
    for i, tag in enumerate(TAGS):
        test_preds_opt[:, i] = (test_probs[:, i] >= optimal_thresholds[tag]).astype(int)

    test_macro_50 = f1_score(test_true, test_preds_50, average="macro", zero_division=0)
    test_micro_50 = f1_score(test_true, test_preds_50, average="micro", zero_division=0)
    test_macro_opt = f1_score(test_true, test_preds_opt, average="macro", zero_division=0)
    test_micro_opt = f1_score(test_true, test_preds_opt, average="micro", zero_division=0)

    print(f"{'Dataset Split':<22} | {'Evaluation Metric':<20} | {'Static 0.50 Thres':<18} | {'Optimized Thres':<18} | {'F1 Improvement':<16}")
    print("-" * 90)
    
    val_macro_diff = (val_macro_opt - val_macro_50) * 100
    val_micro_diff = (val_micro_opt - val_micro_50) * 100
    print(f"{'Validation (2,262 rows)':<22} | {'Macro F1':<20} | {val_macro_50*100:<17.2f}% | {val_macro_opt*100:<17.2f}% | +{val_macro_diff:<15.2f}%")
    print(f"{'':<22} | {'Micro F1 (Accuracy)':<20} | {val_micro_50*100:<17.2f}% | {val_micro_opt*100:<17.2f}% | +{val_micro_diff:<15.2f}%")
    print("-" * 90)

    test_macro_diff = (test_macro_opt - test_macro_50) * 100
    test_micro_diff = (test_micro_opt - test_micro_50) * 100
    print(f"{'Test Set (2,262 rows)':<22} | {'Macro F1':<20} | {test_macro_50*100:<17.2f}% | {test_macro_opt*100:<17.2f}% | +{test_macro_diff:<15.2f}%")
    print(f"{'':<22} | {'Micro F1 (Accuracy)':<20} | {test_micro_50*100:<17.2f}% | {test_micro_opt*100:<17.2f}% | +{test_micro_diff:<15.2f}%")
    print("=" * 90)

    # 3. Detailed Per-Tag Breakdown on Test Set
    print("\n" + "="*90)
    print("🔍 DETAILED PER-TAG BREAKDOWN ON TEST SET (Unseen Data)")
    print("="*90)
    print(f"{'Tag Name':<16} | {'Test Positives':<16} | {'Static F1 (0.50)':<18} | {'Optimized F1':<18} | {'Gain':<12}")
    print("-" * 90)
    for i, tag in enumerate(TAGS):
        y_test_tag = test_true[:, i]
        pos_test = int(np.sum(y_test_tag))
        f1_test_50 = f1_score(y_test_tag, test_preds_50[:, i], zero_division=0) * 100
        f1_test_opt = f1_score(y_test_tag, test_preds_opt[:, i], zero_division=0) * 100
        gain = f1_test_opt - f1_test_50
        print(f"{tag:<16} | {pos_test:<16} | {f1_test_50:<17.2f}% | {f1_test_opt:<17.2f}% | +{gain:<10.2f}%")
    print("="*90)

if __name__ == "__main__":
    main()