ui update

app.py

@@ -1,53 +1,239 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Model configuration
+MODEL_NAME = "cybersectony/phishing-email-detection-distilbert_v2.4.1"
+
+# Global variables for model and tokenizer
+tokenizer = None
+model = None
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
-model = AutoModelForSequenceClassification.from_pretrained("cybersectony/phishing-email-detection-distilbert_v2.4.1")
+def load_model():
+    """Load the model and tokenizer with error handling"""
+    global tokenizer, model
+    try:
+        logger.info("Loading model and tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+        logger.info("Model loaded successfully!")
+        return True
+    except Exception as e:
+        logger.error(f"Error loading model: {e}")
+        return False
 
 def predict_email(email_text):
-    # Preprocess and tokenize
-    inputs = tokenizer(
-        email_text,
-        return_tensors="pt",
-        truncation=True,
-        max_length=512
-    )
-    
-    # Get prediction
-    with torch.no_grad():
-        outputs = model(**inputs)
-        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    """
+    Predict whether an email is phishing or legitimate
     
-    # Get probabilities for each class
-    probs = predictions[0].tolist()
+    Args:
+        email_text (str): The email content to analyze
+        
+    Returns:
+        str: Formatted prediction results
+    """
+    # Input validation
+    if not email_text or not email_text.strip():
+        return "⚠️ **Error**: Please enter some email text to analyze."
     
-    # Create labels dictionary
-    labels = {
-        "Legitimate Email": probs[0],
-        "Phishing URL": probs[1],
-        "Legitimate URL": probs[2],
-        "Phishing URL (Alt)": probs[3]
-    }
+    if len(email_text.strip()) < 10:
+        return "⚠️ **Warning**: Email text seems too short for reliable analysis. Please provide more content."
     
-    # Determine the most likely classification
-    max_label = max(labels.items(), key=lambda x: x[1])
+    # 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."
     
-    # Format output
-    result = f"**Prediction**: {max_label[0]}\n**Confidence**: {max_label[1]:.4f}\n\n**All Probabilities**:\n"
-    for label, prob in labels.items():
-        result += f"{label}: {prob:.4f}\n"
-    return result
+    try:
+        # Preprocess and tokenize
+        inputs = tokenizer(
+            email_text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=True
+        )
+        
+        # Get prediction
+        with torch.no_grad():
+            outputs = model(**inputs)
+            predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        
+        # Get probabilities for each class
+        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
+        }
+        
+        # Remove zero probability labels
+        labels = {k: v for k, v in labels.items() if v > 0.001}
+        
+        # Determine the most likely classification
+        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]:
+            if confidence > 0.8:
+                risk_emoji = "🚨"
+                risk_level = "HIGH RISK"
+            elif confidence > 0.6:
+                risk_emoji = "⚠️"
+                risk_level = "MEDIUM RISK"
+            else:
+                risk_emoji = "⚡"
+                risk_level = "LOW RISK"
+        else:
+            risk_emoji = "✅"
+            risk_level = "SAFE"
+        
+        # Format output
+        result = f"{risk_emoji} **{risk_level}**\n\n"
+        result += f"**Primary Classification**: {max_label[0]}\n"
+        result += f"**Confidence**: {confidence:.1%}\n\n"
+        result += f"**Detailed Analysis**:\n"
+        
+        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."
+        else:
+            result += f"\n✅ **Recommendation**: This email appears to be legitimate, but always remain vigilant."
+            
+        return result
+        
+    except Exception as e:
+        logger.error(f"Error during prediction: {e}")
+        return f"❌ **Error**: An error occurred during analysis: {str(e)}"
+
+# Example emails for demonstration
+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.
+
+Order Details:
+- Product: Wireless Headphones
+- Amount: $79.99
+- Shipping Address: [Your provided address]
+
+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.
+
+Best regards,
+Customer Service Team
+TechStore Inc."""
+
+example_phishing = """URGENT - Account Security Alert!!!
+
+Your account has been COMPROMISED and will be SUSPENDED in 24 hours!
+
+Immediate action required: Click here to verify your account NOW: http://security-verify-account-urgent.suspicious-domain.com/verify-now
+
+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 WARNING - Act immediately!
+
+Security Team
+[Suspicious Bank Name]"""
+
+# Load model on startup
+load_model()
 
 # Create Gradio interface
-iface = gr.Interface(
-    fn=predict_email,
-    inputs=gr.Textbox(lines=5, placeholder="Enter the email text here..."),
-    outputs="text",
-    title="Phishing Email Detection",
-    description="Enter an email text to classify it as legitimate or phishing using a DistilBERT model."
-)
+with gr.Blocks(title="Phishing Email Detection", 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.
+    
+    **⚠️ Disclaimer**: This is an AI-based tool for educational purposes. Always use your judgment and follow your organization's security policies.
+    """)
+    
+    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."
+            )
+            
+            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):
+            output = gr.Textbox(
+                label="Analysis Results",
+                lines=15,
+                interactive=False
+            )
+    
+    # Example section
+    gr.Markdown("### 📝 Try These 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")
+    
+    # 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
+    )
+    
+    phishing_btn.click(
+        fn=lambda: example_phishing,
+        outputs=email_input
+    )
+    
+    # Footer
+    gr.Markdown("""
+    ---
+    **Model**: cybersectony/phishing-email-detection-distilbert_v2.4.1  
+    **Framework**: Transformers + DistilBERT  
+    **Interface**: Gradio
+    """)
 
 # Launch the interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True
+    )