Update app.py

app.py

@@ -25,8 +25,10 @@ MODEL_ID = "Perth0603/phishing-email-mobilebert"
 
 app = FastAPI(title="Phishing Text Classifier with Preprocessing", version="1.0.0")
 
-# Temperature for softening predictions (1.0 = normal, >1.0 = softer, <1.0 = sharper)
-TEMPERATURE = 3.0  # Adjust this value (try 1.5 to 3.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)
 
 
 # ============================================================================
@@ -140,6 +142,35 @@ 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 _load_model():
     """Load model, tokenizer, and preprocessor"""
     global _tokenizer, _model, _device, _preprocessor
@@ -149,7 +180,8 @@ def _load_model():
         print(f"\n{'='*60}")
         print(f"Loading model: {MODEL_ID}")
         print(f"Device: {_device}")
-        print(f"Temperature scaling: {TEMPERATURE}")
+        print(f"Confidence calibration: {MIN_CONFIDENCE*100:.0f}%-{MAX_CONFIDENCE*100:.0f}%")
+        print(f"Smoothing factor: {SMOOTHING_FACTOR}")
         print(f"{'='*60}\n")
         
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -171,7 +203,7 @@ def _load_model():
 
 
 def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List[Dict]:
-    """Predict with temperature-scaled probabilities"""
+    """Predict with calibrated probabilities"""
     _load_model()
     if not texts:
         return []
@@ -191,31 +223,33 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
     )
     enc = {k: v.to(_device) for k, v in enc.items()}
 
-    # Predict with temperature scaling
+    # Predict
     with torch.no_grad():
         logits = _model(**enc).logits
-        # Apply temperature scaling to soften probabilities
-        scaled_logits = logits / TEMPERATURE
-        probs = F.softmax(scaled_logits, dim=-1)
+        probs = F.softmax(logits, dim=-1)
+        
+        # Apply calibration to reduce overconfidence
+        calibrated_probs = _calibrate_probabilities(probs)
 
     # Get labels from model config
     id2label = getattr(_model.config, "id2label", {0: "LEGIT", 1: "PHISH"})
 
     outputs: List[Dict] = []
-    for text_idx in range(probs.shape[0]):
-        p = probs[text_idx]
+    for text_idx in range(calibrated_probs.shape[0]):
+        p_original = probs[text_idx]
+        p_calibrated = calibrated_probs[text_idx]
         
-        # Get prediction
-        predicted_idx = int(torch.argmax(p).item())
+        # Get prediction from calibrated probs
+        predicted_idx = int(torch.argmax(p_calibrated).item())
         predicted_label_raw = id2label.get(predicted_idx, f"CLASS_{predicted_idx}")
         predicted_label_norm = _normalize_label(predicted_label_raw)
-        predicted_prob = float(p[predicted_idx].item())
+        predicted_prob = float(p_calibrated[predicted_idx].item())
 
         # Build probability breakdown
         prob_breakdown = {}
-        for i in range(len(p)):
+        for i in range(len(p_calibrated)):
             label = _normalize_label(id2label.get(i, f"CLASS_{i}"))
-            prob_breakdown[label] = round(float(p[i].item()), 4)
+            prob_breakdown[label] = round(float(p_calibrated[i].item()), 4)
 
         output = {
             "text": texts[text_idx][:100] + "..." if len(texts[text_idx]) > 100 else texts[text_idx],
@@ -225,6 +259,7 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List
             "confidence": round(predicted_prob * 100, 2),
             "score": round(predicted_prob, 4),
             "probs": prob_breakdown,
+            "original_confidence": round(float(p_original[predicted_idx].item()) * 100, 2),  # Show original for comparison
         }
         
         if include_preprocessing and preprocessing_info:
@@ -247,8 +282,11 @@ def root():
         "status": "ok",
         "model": MODEL_ID,
         "device": _device,
-        "temperature": TEMPERATURE,
-        "note": "Using temperature scaling to calibrate probabilities"
+        "calibration": {
+            "min_confidence": MIN_CONFIDENCE,
+            "max_confidence": MAX_CONFIDENCE,
+            "smoothing_factor": SMOOTHING_FACTOR
+        }
     }
 
 
@@ -264,7 +302,6 @@ def debug_labels():
         "label2id": getattr(_model.config, "label2id", {}),
         "num_labels": int(getattr(_model.config, "num_labels", 0)),
         "device": _device,
-        "temperature": TEMPERATURE,
     }