Update app.py

app.py

@@ -25,10 +25,9 @@ MODEL_ID = "Perth0603/phishing-email-mobilebert"
 
 app = FastAPI(title="Phishing Text Classifier with Preprocessing", version="1.0.0")
 
-# Confidence calibration settings
-MIN_CONFIDENCE = 0.70  # Minimum confidence to report (70%)
-MAX_CONFIDENCE = 0.95  # Maximum confidence to report (95%)
-SMOOTHING_FACTOR = 0.15  # How much to smooth (0.1-0.3)
+# Uncertainty estimation settings
+MC_SAMPLES = 10  # Number of forward passes with dropout (more = smoother, slower)
+DROPOUT_RATE = 0.1  # Dropout rate for uncertainty (0.05-0.15)
 
 
 # ============================================================================
@@ -104,11 +103,13 @@ class TextPreprocessor:
 class PredictPayload(BaseModel):
     inputs: str
     include_preprocessing: bool = True
+    use_uncertainty: bool = True  # Enable uncertainty estimation
 
 
 class BatchPredictPayload(BaseModel):
     inputs: List[str]
     include_preprocessing: bool = True
+    use_uncertainty: bool = True
 
 
 class LabeledText(BaseModel):
@@ -142,33 +143,11 @@ def _normalize_label(txt: str) -> str:
     return t
 
 
-def _calibrate_probabilities(probs: torch.Tensor) -> torch.Tensor:
-    """
-    Calibrate overconfident probabilities to more realistic range.
-    Uses label smoothing to reduce extreme confidence.
-    """
-    num_classes = probs.shape[-1]
-    
-    # Apply label smoothing
-    smoothed_probs = probs * (1 - SMOOTHING_FACTOR) + (SMOOTHING_FACTOR / num_classes)
-    
-    # Clip to min/max confidence range
-    max_prob, max_idx = torch.max(smoothed_probs, dim=-1, keepdim=True)
-    
-    # Scale to desired range
-    if max_prob > MAX_CONFIDENCE:
-        scale_factor = MAX_CONFIDENCE / max_prob
-        smoothed_probs = smoothed_probs * scale_factor
-    
-    # Ensure minimum confidence for winner
-    if max_prob < MIN_CONFIDENCE:
-        scale_factor = MIN_CONFIDENCE / max_prob
-        smoothed_probs = smoothed_probs * scale_factor
-    
-    # Renormalize to sum to 1
-    smoothed_probs = smoothed_probs / smoothed_probs.sum(dim=-1, keepdim=True)
-    
-    return smoothed_probs
+def _enable_dropout(model):
+    """Enable dropout layers during inference for uncertainty estimation"""
+    for module in model.modules():
+        if module.__class__.__name__.startswith('Dropout'):
+            module.train()
 
 
 def _load_model():
@@ -180,8 +159,7 @@ def _load_model():
         print(f"\n{'='*60}")
         print(f"Loading model: {MODEL_ID}")
         print(f"Device: {_device}")
-        print(f"Confidence calibration: {MIN_CONFIDENCE*100:.0f}%-{MAX_CONFIDENCE*100:.0f}%")
-        print(f"Smoothing factor: {SMOOTHING_FACTOR}")
+        print(f"MC Dropout samples: {MC_SAMPLES}")
         print(f"{'='*60}\n")
         
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -202,8 +180,41 @@ def _load_model():
         print(f"{'='*60}\n")
 
 
-def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List[Dict]:
-    """Predict with calibrated probabilities"""
+def _predict_with_uncertainty(enc: Dict, use_uncertainty: bool = True) -> tuple:
+    """
+    Predict with uncertainty estimation using MC Dropout.
+    Returns: (mean_probs, std_probs)
+    """
+    if not use_uncertainty:
+        # Standard prediction
+        with torch.no_grad():
+            logits = _model(**enc).logits
+            probs = F.softmax(logits, dim=-1)
+        return probs, torch.zeros_like(probs)
+    
+    # Monte Carlo Dropout: multiple forward passes with dropout enabled
+    prob_samples = []
+    
+    _enable_dropout(_model)  # Enable dropout during inference
+    
+    with torch.no_grad():
+        for _ in range(MC_SAMPLES):
+            logits = _model(**enc).logits
+            probs = F.softmax(logits, dim=-1)
+            prob_samples.append(probs)
+    
+    _model.eval()  # Restore eval mode
+    
+    # Stack all samples and compute mean and std
+    prob_samples = torch.stack(prob_samples)  # [MC_SAMPLES, batch_size, num_classes]
+    mean_probs = prob_samples.mean(dim=0)  # Average predictions
+    std_probs = prob_samples.std(dim=0)    # Uncertainty (variance)
+    
+    return mean_probs, std_probs
+
+
+def _predict_texts(texts: List[str], include_preprocessing: bool = True, use_uncertainty: bool = True) -> List[Dict]:
+    """Predict with uncertainty-aware confidence scores"""
     _load_model()
     if not texts:
         return []
@@ -223,43 +234,48 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
     )
     enc = {k: v.to(_device) for k, v in enc.items()}
 
-    # Predict
-    with torch.no_grad():
-        logits = _model(**enc).logits
-        probs = F.softmax(logits, dim=-1)
-        
-        # Apply calibration to reduce overconfidence
-        calibrated_probs = _calibrate_probabilities(probs)
+    # Predict with uncertainty
+    mean_probs, std_probs = _predict_with_uncertainty(enc, use_uncertainty)
 
     # Get labels from model config
     id2label = getattr(_model.config, "id2label", {0: "LEGIT", 1: "PHISH"})
 
     outputs: List[Dict] = []
-    for text_idx in range(calibrated_probs.shape[0]):
-        p_original = probs[text_idx]
-        p_calibrated = calibrated_probs[text_idx]
+    for text_idx in range(mean_probs.shape[0]):
+        p_mean = mean_probs[text_idx]
+        p_std = std_probs[text_idx]
         
-        # Get prediction from calibrated probs
-        predicted_idx = int(torch.argmax(p_calibrated).item())
+        # Get prediction
+        predicted_idx = int(torch.argmax(p_mean).item())
         predicted_label_raw = id2label.get(predicted_idx, f"CLASS_{predicted_idx}")
         predicted_label_norm = _normalize_label(predicted_label_raw)
-        predicted_prob = float(p_calibrated[predicted_idx].item())
+        predicted_prob = float(p_mean[predicted_idx].item())
+        predicted_std = float(p_std[predicted_idx].item())
+
+        # Calculate uncertainty-adjusted confidence
+        # Higher uncertainty = lower reported confidence
+        uncertainty_penalty = predicted_std * 2.0  # Amplify uncertainty effect
+        adjusted_confidence = max(0.5, predicted_prob - uncertainty_penalty)  # Don't go below 50%
 
         # Build probability breakdown
         prob_breakdown = {}
-        for i in range(len(p_calibrated)):
+        uncertainty_breakdown = {}
+        for i in range(len(p_mean)):
             label = _normalize_label(id2label.get(i, f"CLASS_{i}"))
-            prob_breakdown[label] = round(float(p_calibrated[i].item()), 4)
+            prob_breakdown[label] = round(float(p_mean[i].item()), 4)
+            uncertainty_breakdown[label] = round(float(p_std[i].item()), 4)
 
         output = {
             "text": texts[text_idx][:100] + "..." if len(texts[text_idx]) > 100 else texts[text_idx],
             "label": predicted_label_norm,
             "raw_label": predicted_label_raw,
             "is_phish": predicted_label_norm == "PHISH",
-            "confidence": round(predicted_prob * 100, 2),
-            "score": round(predicted_prob, 4),
+            "confidence": round(adjusted_confidence * 100, 2),
+            "score": round(adjusted_confidence, 4),
+            "uncertainty": round(predicted_std * 100, 2),  # Uncertainty as percentage
             "probs": prob_breakdown,
-            "original_confidence": round(float(p_original[predicted_idx].item()) * 100, 2),  # Show original for comparison
+            "uncertainty_scores": uncertainty_breakdown,
+            "raw_confidence": round(predicted_prob * 100, 2),  # Original model confidence
         }
         
         if include_preprocessing and preprocessing_info:
@@ -282,10 +298,10 @@ def root():
         "status": "ok",
         "model": MODEL_ID,
         "device": _device,
-        "calibration": {
-            "min_confidence": MIN_CONFIDENCE,
-            "max_confidence": MAX_CONFIDENCE,
-            "smoothing_factor": SMOOTHING_FACTOR
+        "uncertainty_estimation": {
+            "enabled": True,
+            "mc_samples": MC_SAMPLES,
+            "dropout_rate": DROPOUT_RATE
         }
     }
 
@@ -320,7 +336,9 @@ def debug_preprocessing(payload: PredictPayload):
 def predict(payload: PredictPayload):
     """Single prediction"""
     try:
-        res = _predict_texts([payload.inputs], include_preprocessing=payload.include_preprocessing)
+        res = _predict_texts([payload.inputs], 
+                            include_preprocessing=payload.include_preprocessing,
+                            use_uncertainty=payload.use_uncertainty)
         return res[0]
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -330,7 +348,9 @@ def predict(payload: PredictPayload):
 def predict_batch(payload: BatchPredictPayload):
     """Batch predictions"""
     try:
-        return _predict_texts(payload.inputs, include_preprocessing=payload.include_preprocessing)
+        return _predict_texts(payload.inputs, 
+                             include_preprocessing=payload.include_preprocessing,
+                             use_uncertainty=payload.use_uncertainty)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -341,7 +361,7 @@ def evaluate(payload: EvalPayload):
     try:
         texts = [s.text for s in payload.samples]
         gts = [(_normalize_label(s.label) if s.label is not None else None) for s in payload.samples]
-        preds = _predict_texts(texts, include_preprocessing=False)
+        preds = _predict_texts(texts, include_preprocessing=False, use_uncertainty=True)
 
         total = len(preds)
         correct = 0