Update app.py

app.py

@@ -25,9 +25,9 @@ MODEL_ID = "Perth0603/phishing-email-mobilebert"
 
 app = FastAPI(title="Phishing Text Classifier with Preprocessing", version="1.0.0")
 
-# 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)
+# Confidence adjustment settings
+BASE_CONFIDENCE_MIN = 0.55  # Minimum confidence (55%)
+BASE_CONFIDENCE_MAX = 0.85  # Maximum confidence (85%)
 
 
 # ============================================================================
@@ -57,34 +57,122 @@ class TextPreprocessor:
         """Reduce tokens to lemmas"""
         return [self.lemmatizer.lemmatize(token) for token in tokens]
     
+    def analyze_phishing_indicators(self, text: str) -> Dict:
+        """Comprehensive phishing indicator analysis"""
+        indicators = {
+            "urgent_words": bool(re.search(
+                r'\b(urgent|immediately|immediate|act now|right now|asap|verify now|'
+                r'confirm now|update now|click now|respond now|expire soon|expiring|'
+                r'time sensitive|limited time|hurry|quick|fast|today only)\b', 
+                text, re.IGNORECASE
+            )),
+            "threat_words": bool(re.search(
+                r'\b(suspend|suspended|lock|locked|block|blocked|disable|disabled|'
+                r'restrict|restricted|terminate|terminated|cancel|cancelled|close|closed|'
+                r'freeze|frozen|ban|banned|deactivate|deactivated|remove|removed)\b', 
+                text, re.IGNORECASE
+            )),
+            "action_words": bool(re.search(
+                r'\b(click here|click now|click below|click this|verify|confirm|update|'
+                r'download|install|open attachment|validate|authenticate|reset password|'
+                r'change password|provide|submit|enter|fill out|complete)\b', 
+                text, re.IGNORECASE
+            )),
+            "financial_words": bool(re.search(
+                r'\b(payment|pay|money|credit card|bank account|billing|invoice|refund|'
+                r'tax|irs|paypal|transaction|transfer|wire|deposit|account number|'
+                r'social security|ssn|card number|cvv|pin)\b', 
+                text, re.IGNORECASE
+            )),
+            "authority_impersonation": bool(re.search(
+                r'\b(paypal|amazon|microsoft|apple|google|facebook|instagram|netflix|'
+                r'ebay|irs|fbi|cia|government|police|bank of america|chase|wells fargo|'
+                r'citibank|security team|support team|admin|administrator)\b', 
+                text, re.IGNORECASE
+            )),
+            "suspicious_urls": bool(re.search(r'http[s]?://|www\.', text)),
+            "suspicious_domain": bool(re.search(
+                r'\b(bit\.ly|tinyurl|goo\.gl|short|link|redirect|verify-|secure-|account-|'
+                r'update-|login-|signin-)\w+\.(com|net|org|info|xyz|tk|ml|ga|cf|gq)', 
+                text, re.IGNORECASE
+            )),
+            "generic_greeting": bool(re.search(
+                r'^(dear (customer|user|member|client|sir|madam)|hello|hi there|greetings)\b', 
+                text, re.IGNORECASE
+            )),
+            "poor_grammar": self._detect_poor_grammar(text),
+            "excessive_punctuation": bool(re.search(r'[!?]{2,}', text)),
+            "all_caps": len(re.findall(r'\b[A-Z]{3,}\b', text)) > 2,
+            "currency_symbols": bool(re.search(r'[$£€¥₹]', text)),
+        }
+        
+        # Count active indicators
+        active_count = sum(indicators.values())
+        total_count = len(indicators)
+        
+        # Determine urgency level
+        urgency_score = sum([
+            indicators["urgent_words"] * 2,
+            indicators["threat_words"] * 2,
+            indicators["action_words"],
+            indicators["excessive_punctuation"],
+            indicators["all_caps"]
+        ])
+        
+        if urgency_score >= 4:
+            urgency_level = "CRITICAL"
+        elif urgency_score >= 2:
+            urgency_level = "HIGH"
+        elif urgency_score >= 1:
+            urgency_level = "MEDIUM"
+        else:
+            urgency_level = "LOW"
+        
+        indicators["urgency_level"] = urgency_level
+        indicators["indicator_count"] = active_count
+        indicators["indicator_percentage"] = round((active_count / total_count) * 100, 1)
+        
+        return indicators
+    
+    def _detect_poor_grammar(self, text: str) -> bool:
+        """Simple heuristic for poor grammar"""
+        issues = 0
+        # Multiple spaces
+        if re.search(r'\s{2,}', text):
+            issues += 1
+        # Missing spaces after punctuation
+        if re.search(r'[.,!?][a-zA-Z]', text):
+            issues += 1
+        # Inconsistent capitalization
+        sentences = re.split(r'[.!?]+', text)
+        for sent in sentences:
+            sent = sent.strip()
+            if sent and len(sent) > 5 and not sent[0].isupper():
+                issues += 1
+                break
+        return issues >= 2
+    
     def sentiment_analysis(self, text: str) -> Dict:
-        """Analyze sentiment and phishing indicators"""
+        """Analyze sentiment"""
         blob = TextBlob(text)
         polarity = blob.sentiment.polarity
         subjectivity = blob.sentiment.subjectivity
         
-        phishing_indicators = {
-            "urgent_words": bool(re.search(r'\b(urgent|immediate|act now|verify|confirm|update|click|verify account)\b', text, re.IGNORECASE)),
-            "threat_words": bool(re.search(r'\b(suspend|limited|expire|locked|disabled|restricted)\b', text, re.IGNORECASE)),
-            "suspicious_urls": bool(re.search(r'http\S+|www\S+', text)),
-            "urgency_level": "HIGH" if re.search(r'\b(urgent|immediate|act now)\b', text, re.IGNORECASE) else "LOW"
-        }
-        
         return {
             "polarity": round(polarity, 4),
             "subjectivity": round(subjectivity, 4),
             "sentiment": "positive" if polarity > 0.1 else "negative" if polarity < -0.1 else "neutral",
             "is_persuasive": subjectivity > 0.5,
-            "phishing_indicators": phishing_indicators
         }
     
     def preprocess(self, text: str) -> Dict:
-        """Preprocessing for analysis"""
+        """Full preprocessing pipeline"""
         tokens = self.tokenize(text)
         tokens_no_stop = self.remove_stopwords(tokens)
         stemmed = self.stem(tokens_no_stop)
         lemmatized = self.lemmatize(tokens_no_stop)
         sentiment = self.sentiment_analysis(text)
+        phishing_indicators = self.analyze_phishing_indicators(text)
         
         return {
             "original_text": text,
@@ -93,6 +181,7 @@ class TextPreprocessor:
             "stemmed_tokens": stemmed,
             "lemmatized_tokens": lemmatized,
             "sentiment": sentiment,
+            "phishing_indicators": phishing_indicators,
             "token_count": len(tokens_no_stop)
         }
 
@@ -103,13 +192,11 @@ 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):
@@ -143,11 +230,47 @@ def _normalize_label(txt: str) -> str:
     return t
 
 
-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 _adjust_confidence_with_indicators(base_prob: float, indicators: Dict, predicted_label: str) -> float:
+    """
+    Adjust confidence based on phishing indicators.
+    More indicators = context suggests phishing, so confidence varies based on prediction
+    """
+    indicator_count = indicators.get("indicator_count", 0)
+    indicator_percentage = indicators.get("indicator_percentage", 0)
+    
+    # Base adjustment from indicator count
+    # If predicting PHISH and many indicators: more confident (but cap at 85%)
+    # If predicting LEGIT with many indicators: less confident (uncertainty)
+    # If predicting PHISH with few indicators: less confident (might be wrong)
+    # If predicting LEGIT with few indicators: more confident
+    
+    if predicted_label == "PHISH":
+        # Phishing prediction
+        if indicator_percentage >= 40:  # Strong indicators
+            # High confidence: 75-85%
+            adjusted = 0.75 + (indicator_percentage / 100) * 0.10
+        elif indicator_percentage >= 25:  # Moderate indicators
+            # Medium confidence: 65-75%
+            adjusted = 0.65 + (indicator_percentage / 100) * 0.10
+        else:  # Weak indicators
+            # Lower confidence: 55-65%
+            adjusted = 0.55 + (indicator_percentage / 100) * 0.10
+    else:
+        # Legitimate prediction
+        if indicator_percentage >= 40:  # Many phishing indicators but predicting legit?
+            # Low confidence: 55-65% (uncertain)
+            adjusted = 0.65 - (indicator_percentage / 100) * 0.10
+        elif indicator_percentage >= 25:  # Some indicators
+            # Medium confidence: 65-75%
+            adjusted = 0.70 - (indicator_percentage / 100) * 0.05
+        else:  # Few indicators
+            # High confidence: 75-85%
+            adjusted = 0.75 + ((100 - indicator_percentage) / 100) * 0.10
+    
+    # Clamp to min/max range
+    adjusted = max(BASE_CONFIDENCE_MIN, min(BASE_CONFIDENCE_MAX, adjusted))
+    
+    return adjusted
 
 
 def _load_model():
@@ -159,7 +282,7 @@ def _load_model():
         print(f"\n{'='*60}")
         print(f"Loading model: {MODEL_ID}")
         print(f"Device: {_device}")
-        print(f"MC Dropout samples: {MC_SAMPLES}")
+        print(f"Confidence range: {BASE_CONFIDENCE_MIN*100:.0f}%-{BASE_CONFIDENCE_MAX*100:.0f}%")
         print(f"{'='*60}\n")
         
         _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -180,49 +303,14 @@ def _load_model():
         print(f"{'='*60}\n")
 
 
-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"""
+def _predict_texts(texts: List[str], include_preprocessing: bool = True) -> List[Dict]:
+    """Predict with indicator-based confidence adjustment"""
     _load_model()
     if not texts:
         return []
 
-    # Get preprocessing info
-    preprocessing_info = None
-    if include_preprocessing:
-        preprocessing_info = [_preprocessor.preprocess(text) for text in texts]
+    # Get preprocessing info (always needed for indicators)
+    preprocessing_info = [_preprocessor.preprocess(text) for text in texts]
 
     # Tokenize
     enc = _tokenizer(
@@ -234,36 +322,39 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True, use_unc
     )
     enc = {k: v.to(_device) for k, v in enc.items()}
 
-    # Predict with uncertainty
-    mean_probs, std_probs = _predict_with_uncertainty(enc, use_uncertainty)
+    # Predict
+    with torch.no_grad():
+        logits = _model(**enc).logits
+        probs = F.softmax(logits, dim=-1)
 
     # Get labels from model config
     id2label = getattr(_model.config, "id2label", {0: "LEGIT", 1: "PHISH"})
 
     outputs: List[Dict] = []
-    for text_idx in range(mean_probs.shape[0]):
-        p_mean = mean_probs[text_idx]
-        p_std = std_probs[text_idx]
+    for text_idx in range(probs.shape[0]):
+        p = probs[text_idx]
+        preprocessing = preprocessing_info[text_idx]
+        indicators = preprocessing["phishing_indicators"]
         
         # Get prediction
-        predicted_idx = int(torch.argmax(p_mean).item())
+        predicted_idx = int(torch.argmax(p).item())
         predicted_label_raw = id2label.get(predicted_idx, f"CLASS_{predicted_idx}")
         predicted_label_norm = _normalize_label(predicted_label_raw)
-        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%
+        raw_prob = float(p[predicted_idx].item())
+        
+        # Adjust confidence based on indicators
+        adjusted_confidence = _adjust_confidence_with_indicators(
+            raw_prob, indicators, predicted_label_norm
+        )
 
-        # Build probability breakdown
+        # Build probability breakdown (adjusted)
         prob_breakdown = {}
-        uncertainty_breakdown = {}
-        for i in range(len(p_mean)):
+        for i in range(len(p)):
             label = _normalize_label(id2label.get(i, f"CLASS_{i}"))
-            prob_breakdown[label] = round(float(p_mean[i].item()), 4)
-            uncertainty_breakdown[label] = round(float(p_std[i].item()), 4)
+            if i == predicted_idx:
+                prob_breakdown[label] = round(adjusted_confidence, 4)
+            else:
+                prob_breakdown[label] = round(1.0 - adjusted_confidence, 4)
 
         output = {
             "text": texts[text_idx][:100] + "..." if len(texts[text_idx]) > 100 else texts[text_idx],
@@ -272,14 +363,12 @@ def _predict_texts(texts: List[str], include_preprocessing: bool = True, use_unc
             "is_phish": predicted_label_norm == "PHISH",
             "confidence": round(adjusted_confidence * 100, 2),
             "score": round(adjusted_confidence, 4),
-            "uncertainty": round(predicted_std * 100, 2),  # Uncertainty as percentage
             "probs": prob_breakdown,
-            "uncertainty_scores": uncertainty_breakdown,
-            "raw_confidence": round(predicted_prob * 100, 2),  # Original model confidence
+            "model_raw_confidence": round(raw_prob * 100, 2),
         }
         
-        if include_preprocessing and preprocessing_info:
-            output["preprocessing"] = preprocessing_info[text_idx]
+        if include_preprocessing:
+            output["preprocessing"] = preprocessing
         
         outputs.append(output)
 
@@ -298,11 +387,8 @@ def root():
         "status": "ok",
         "model": MODEL_ID,
         "device": _device,
-        "uncertainty_estimation": {
-            "enabled": True,
-            "mc_samples": MC_SAMPLES,
-            "dropout_rate": DROPOUT_RATE
-        }
+        "confidence_range": f"{BASE_CONFIDENCE_MIN*100:.0f}%-{BASE_CONFIDENCE_MAX*100:.0f}%",
+        "note": "Confidence adjusted based on phishing indicators"
     }
 
 
@@ -336,9 +422,7 @@ def debug_preprocessing(payload: PredictPayload):
 def predict(payload: PredictPayload):
     """Single prediction"""
     try:
-        res = _predict_texts([payload.inputs], 
-                            include_preprocessing=payload.include_preprocessing,
-                            use_uncertainty=payload.use_uncertainty)
+        res = _predict_texts([payload.inputs], include_preprocessing=payload.include_preprocessing)
         return res[0]
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -348,9 +432,7 @@ def predict(payload: PredictPayload):
 def predict_batch(payload: BatchPredictPayload):
     """Batch predictions"""
     try:
-        return _predict_texts(payload.inputs, 
-                             include_preprocessing=payload.include_preprocessing,
-                             use_uncertainty=payload.use_uncertainty)
+        return _predict_texts(payload.inputs, include_preprocessing=payload.include_preprocessing)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -361,7 +443,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, use_uncertainty=True)
+        preds = _predict_texts(texts, include_preprocessing=False)
 
         total = len(preds)
         correct = 0