model respnsive

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import logging
+import numpy as np
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -21,33 +22,68 @@ def load_model():
         logger.info("Loading model and tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
         model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+        
+        # Debug: Print model configuration
+        logger.info(f"Model config: {model.config}")
+        logger.info(f"Number of labels: {model.config.num_labels}")
+        if hasattr(model.config, 'id2label'):
+            logger.info(f"Label mapping: {model.config.id2label}")
+        
+        # Test model with simple input to check if it's working
+        test_input = "Hello world"
+        inputs = tokenizer(test_input, return_tensors="pt", truncation=True, max_length=512)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            test_probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+            logger.info(f"Test probabilities: {test_probs[0].tolist()}")
+        
         logger.info("Model loaded successfully!")
         return True
     except Exception as e:
         logger.error(f"Error loading model: {e}")
         return False
 
+def get_colored_bar(percentage, label):
+    """Create colored progress bar based on percentage and label type"""
+    # Determine color based on percentage and label
+    if "phishing" in label.lower() or "suspicious" in label.lower():
+        # Red scale for dangerous content
+        if percentage >= 70:
+            color = "🟥"  # High danger - red
+        elif percentage >= 40:
+            color = "🟠"  # Medium danger - orange  
+        else:
+            color = "🟡"  # Low danger - yellow
+    else:
+        # Green scale for legitimate content
+        if percentage >= 70:
+            color = "🟢"  # High confidence - green
+        elif percentage >= 40:
+            color = "🟡"  # Medium confidence - yellow
+        else:
+            color = "⚪"  # Low confidence - white
+    
+    # Create bar (scale to 20 characters)
+    bar_length = max(1, int(percentage / 5))  # Ensure at least 1 if percentage > 0
+    bar = color * bar_length + "⚪" * (20 - bar_length)
+    
+    return bar
+
 def predict_email(email_text):
     """
-    Predict whether an email is phishing or legitimate
-    
-    Args:
-        email_text (str): The email content to analyze
-        
-    Returns:
-        str: Formatted prediction results
+    Enhanced prediction function with proper model output handling
     """
     # Input validation
     if not email_text or not email_text.strip():
         return "⚠️ **Error**: Please enter some email text to analyze."
     
-    if len(email_text.strip()) < 10:
-        return "⚠️ **Warning**: Email text seems too short for reliable analysis. Please provide more content."
+    if len(email_text.strip()) < 5:
+        return "⚠️ **Warning**: Email text too short for reliable analysis."
     
     # Check if model is loaded
     if tokenizer is None or model is None:
         if not load_model():
-            return "❌ **Error**: Failed to load the model. Please try again later."
+            return "❌ **Error**: Failed to load the model."
     
     try:
         # Preprocess and tokenize
@@ -59,32 +95,61 @@ def predict_email(email_text):
             padding=True
         )
         
-        # Get prediction
+        # Get prediction with proper handling
         with torch.no_grad():
             outputs = model(**inputs)
-            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+            # Apply temperature scaling to prevent overconfidence
+            temperature = 1.5
+            scaled_logits = outputs.logits / temperature
+            predictions = torch.nn.functional.softmax(scaled_logits, dim=-1)
         
-        # Get probabilities for each class
+        # Get probabilities
         probs = predictions[0].tolist()
         
-        # Create labels dictionary
-        # Note: Verify these labels match your model's actual training configuration
-        labels = {
-            "Legitimate Email": probs[0],
-            "Phishing Email": probs[1] if len(probs) > 1 else 0.0,
-            "Suspicious Content": probs[2] if len(probs) > 2 else 0.0,
-            "Other": probs[3] if len(probs) > 3 else 0.0
-        }
+        # Log raw outputs for debugging
+        logger.info(f"Raw logits: {outputs.logits[0].tolist()}")
+        logger.info(f"Scaled probabilities: {probs}")
+        
+        # Get proper labels from model config or use fallback
+        if hasattr(model.config, 'id2label') and model.config.id2label:
+            labels = {model.config.id2label[i]: probs[i] for i in range(len(probs))}
+        else:
+            # Fallback - check the actual model output dimension
+            if len(probs) == 2:
+                labels = {
+                    "Legitimate Email": probs[0],
+                    "Phishing Email": probs[1]
+                }
+            elif len(probs) == 4:
+                labels = {
+                    "Legitimate Email": probs[0],
+                    "Phishing Email": probs[1],
+                    "Suspicious Content": probs[2],
+                    "Spam Email": probs[3]
+                }
+            else:
+                # Generic labels
+                labels = {f"Class {i}": probs[i] for i in range(len(probs))}
+        
+        # Check if model is giving reasonable outputs
+        prob_variance = np.var(probs)
+        max_prob = max(probs)
         
-        # Remove zero probability labels
-        labels = {k: v for k, v in labels.items() if v > 0.001}
+        # If variance is too low, the model might not be working properly
+        if prob_variance < 0.01 and max_prob > 0.99:
+            logger.warning("Model showing signs of overconfidence or poor calibration")
+            # Apply smoothing
+            smoothed_probs = [(p * 0.8 + 0.2/len(probs)) for p in probs]
+            labels = {list(labels.keys())[i]: smoothed_probs[i] for i in range(len(smoothed_probs))}
         
-        # Determine the most likely classification
+        # Find prediction
         max_label = max(labels.items(), key=lambda x: x[1])
         
         # Determine risk level and emoji
         confidence = max_label[1]
-        if "Phishing" in max_label[0] or "Suspicious" in max_label[0]:
+        prediction_name = max_label[0]
+        
+        if any(word in prediction_name.lower() for word in ['phishing', 'suspicious', 'spam']):
             if confidence > 0.8:
                 risk_emoji = "🚨"
                 risk_level = "HIGH RISK"
@@ -95,145 +160,145 @@ def predict_email(email_text):
                 risk_emoji = "⚡"
                 risk_level = "LOW RISK"
         else:
-            risk_emoji = "✅"
-            risk_level = "SAFE"
+            if confidence > 0.8:
+                risk_emoji = "✅"
+                risk_level = "SAFE"
+            elif confidence > 0.6:
+                risk_emoji = "✅"
+                risk_level = "LIKELY SAFE"
+            else:
+                risk_emoji = "❓"
+                risk_level = "UNCERTAIN"
         
-        # Format output
+        # Format output with colored bars
         result = f"{risk_emoji} **{risk_level}**\n\n"
-        result += f"**Primary Classification**: {max_label[0]}\n"
+        result += f"**Primary Classification**: {prediction_name}\n"
         result += f"**Confidence**: {confidence:.1%}\n\n"
         result += f"**Detailed Analysis**:\n"
         
+        # Sort by probability and add colored bars
         for label, prob in sorted(labels.items(), key=lambda x: x[1], reverse=True):
             percentage = prob * 100
-            bar_length = int(percentage / 5)  # Scale bar to 20 chars max
-            bar = "█" * bar_length + "░" * (20 - bar_length)
-            result += f"{label}: {percentage:.1f}% {bar}\n"
-        
-        # Add recommendations
-        if "Phishing" in max_label[0] and confidence > 0.7:
-            result += f"\n⚠️ **Recommendation**: This email shows signs of phishing. Do not click any links or provide personal information."
-        elif "Suspicious" in max_label[0] and confidence > 0.6:
-            result += f"\n🔍 **Recommendation**: Exercise caution with this email. Verify sender identity before taking any action."
+            colored_bar = get_colored_bar(percentage, label)
+            result += f"{label}: {percentage:.1f}% {colored_bar}\n"
+        
+        # Add debug info
+        result += f"\n**Debug Info**:\n"
+        result += f"Model Variance: {prob_variance:.4f}\n"
+        result += f"Raw Probabilities: {[f'{p:.3f}' for p in probs]}\n"
+        
+        # Add recommendations based on actual classification
+        if any(word in prediction_name.lower() for word in ['phishing', 'suspicious']) and confidence > 0.6:
+            result += f"\n⚠️ **Recommendation**: This email shows signs of being malicious. Avoid clicking links or providing personal information."
+        elif 'spam' in prediction_name.lower():
+            result += f"\n🗑️ **Recommendation**: This appears to be spam. Consider deleting or marking as junk."
+        elif confidence > 0.7:
+            result += f"\n✅ **Recommendation**: This email appears legitimate, but always remain vigilant."
         else:
-            result += f"\n✅ **Recommendation**: This email appears to be legitimate, but always remain vigilant."
+            result += f"\n❓ **Recommendation**: Classification uncertain. Exercise caution and verify sender if needed."
             
         return result
         
     except Exception as e:
-        logger.error(f"Error during prediction: {e}")
-        return f"❌ **Error**: An error occurred during analysis: {str(e)}"
+        logger.error(f"Error during prediction: {e}", exc_info=True)
+        return f"❌ **Error**: Analysis failed - {str(e)}"
 
-# Example emails for demonstration
+# Example emails for testing
 example_legitimate = """Dear Customer,
 
-Thank you for your recent purchase from our store. Your order #ORD-2024-001234 has been successfully processed and will be shipped within 2-3 business days.
+Thank you for your recent purchase from TechStore. Your order #ORD-2024-001234 has been successfully processed.
 
 Order Details:
-- Product: Wireless Headphones
+- Product: Wireless Headphones  
 - Amount: $79.99
-- Shipping Address: [Your provided address]
+- Estimated delivery: 3-5 business days
 
-You can track your shipment using the tracking number that will be sent to your email once the item is dispatched.
-
-If you have any questions, please contact our customer service team.
+You will receive a tracking number once your item ships.
 
 Best regards,
-Customer Service Team
-TechStore Inc."""
+TechStore Customer Service"""
+
+example_phishing = """URGENT SECURITY ALERT!!!
 
-example_phishing = """URGENT - Account Security Alert!!!
+Your account has been COMPROMISED! Immediate action required!
 
-Your account has been COMPROMISED and will be SUSPENDED in 24 hours!
+Click here NOW to secure your account: http://fake-security-site.malicious.com/urgent-verify
 
-Immediate action required: Click here to verify your account NOW: http://security-verify-account-urgent.suspicious-domain.com/verify-now
+WARNING: You have only 24 hours before your account is permanently suspended!
 
-If you don't verify within 24 hours, your account will be permanently deleted and all your data will be lost forever!
+This is your FINAL notice - act immediately!
 
-This is your FINAL WARNING - Act immediately!
+Security Department"""
 
-Security Team
-[Suspicious Bank Name]"""
+example_neutral = """Hi team,
+
+Hope everyone is doing well. Just wanted to remind you about the meeting scheduled for tomorrow at 2 PM in the conference room.
+
+Please bring your project updates and any questions you might have.
+
+Thanks,
+Sarah"""
 
 # Load model on startup
 load_model()
 
-# Create Gradio interface
-with gr.Blocks(title="Phishing Email Detection", theme=gr.themes.Soft()) as iface:
+# Create enhanced Gradio interface
+with gr.Blocks(title="PhishGuardian AI", theme=gr.themes.Soft()) as iface:
     gr.Markdown("""
-    # 🛡️ Phishing Email Detection System
-    
-    This tool uses a DistilBERT model to analyze email content and detect potential phishing attempts.
-    Simply paste the email text below and get an instant security assessment.
+    # 🛡️ PhishGuardian AI - Enhanced Detection
     
-    **⚠️ Disclaimer**: This is an AI-based tool for educational purposes. Always use your judgment and follow your organization's security policies.
+    Advanced phishing email detection with colored risk indicators and improved model handling.
     """)
     
     with gr.Row():
         with gr.Column(scale=2):
             email_input = gr.Textbox(
                 lines=10,
-                placeholder="Paste the email content here...",
-                label="Email Text",
-                info="Enter the complete email text including headers, body, and any suspicious elements."
+                placeholder="Paste your email content here for analysis...",
+                label="📧 Email Content",
+                info="Enter the complete email text for comprehensive analysis"
             )
             
             with gr.Row():
                 analyze_btn = gr.Button("🔍 Analyze Email", variant="primary", size="lg")
                 clear_btn = gr.Button("🗑️ Clear", variant="secondary")
         
-        with gr.Column(scale=1):
+        with gr.Column(scale=2):
             output = gr.Textbox(
-                label="Analysis Results",
-                lines=15,
-                interactive=False
+                label="🛡️ Security Analysis Results",
+                lines=20,
+                interactive=False,
+                show_copy_button=True
             )
     
-    # Example section
-    gr.Markdown("### 📝 Try These Examples:")
+    # Example section with better examples
+    gr.Markdown("### 📝 Test Examples")
     with gr.Row():
-        with gr.Column():
-            gr.Markdown("**Legitimate Email Example:**")
-            legitimate_btn = gr.Button("Load Legitimate Email", size="sm")
-        with gr.Column():
-            gr.Markdown("**Phishing Email Example:**")
-            phishing_btn = gr.Button("Load Phishing Email", size="sm")
+        legit_btn = gr.Button("✅ Legitimate Email", size="sm")
+        phish_btn = gr.Button("🚨 Phishing Email", size="sm") 
+        neutral_btn = gr.Button("📄 Neutral Text", size="sm")
     
     # Event handlers
-    analyze_btn.click(
-        fn=predict_email,
-        inputs=email_input,
-        outputs=output
-    )
-    
-    clear_btn.click(
-        fn=lambda: ("", ""),
-        outputs=[email_input, output]
-    )
-    
-    legitimate_btn.click(
-        fn=lambda: example_legitimate,
-        outputs=email_input
-    )
+    analyze_btn.click(predict_email, inputs=email_input, outputs=output)
+    clear_btn.click(lambda: ("", ""), outputs=[email_input, output])
     
-    phishing_btn.click(
-        fn=lambda: example_phishing,
-        outputs=email_input
-    )
+    legit_btn.click(lambda: example_legitimate, outputs=email_input)
+    phish_btn.click(lambda: example_phishing, outputs=email_input)
+    neutral_btn.click(lambda: example_neutral, outputs=email_input)
     
-    # Footer
+    # Footer with model info
     gr.Markdown("""
     ---
-    **Model**: cybersectony/phishing-email-detection-distilbert_v2.4.1  
-    **Framework**: Transformers + DistilBERT  
-    **Interface**: Gradio
+    **🔧 Model**: cybersectony/phishing-email-detection-distilbert_v2.4.1  
+    **🎯 Features**: Temperature scaling, colored risk bars, enhanced debugging  
+    **🏛️ Institution**: University of Dar es Salaam (UDSM)
     """)
 
-# Launch the interface
 if __name__ == "__main__":
     iface.launch(
         share=True,
-        server_name="0.0.0.0",
+        server_name="0.0.0.0", 
         server_port=7860,
-        show_error=True
+        show_error=True,
+        debug=True
     )