Update app.py

app.py

@@ -14,12 +14,12 @@ MODEL_NAME = "cybersectony/phishing-email-detection-distilbert_v2.4.1"
 # Explanation of labels and their values
 """
 Labels and Their Meanings:
-- Legitimate: The email appears safe and is likely from a trusted source.
-- Phishing: The email may be a scam attempting to steal personal information.
-- Suspicious: The email has questionable content and may not be safe.
-- Spam: The email is likely unwanted promotional or junk content.
-Each label comes with a percentage (0-100%) indicating the model's confidence.
-Higher percentages mean the model is more certain of the classification.
+- Legitimate: The email is safe and likely from a trusted source.
+- Phishing: The email is a scam attempting to steal personal information.
+- Suspicious: The email has questionable content and may be unsafe.
+- Spam: The email is unwanted promotional or junk content.
+Each label has a percentage (0-100%) showing the model's confidence.
+Higher percentages indicate greater certainty.
 """
 
 # Global variables for model and tokenizer
@@ -44,13 +44,23 @@ def is_valid_email_text(text):
         return False, "Please enter some email text."
     if len(text.strip()) < 10:
         return False, "Text too short for analysis."
-    # Check for basic email-like structure or meaningful words
     if len(text.split()) < 3 or not re.search(r"[a-zA-Z]{3,}", text):
         return False, "Text appears incoherent or not email-like."
     return True, ""
 
+def get_colored_bar(percentage):
+    """Create a simple colored bar based on percentage"""
+    if percentage >= 85:
+        color = "🟢" if percentage >= 85 else "🟥"
+    elif percentage >= 50:
+        color = "🟡"
+    else:
+        color = "⚪"
+    bar_length = max(1, int(percentage / 5))  # Scale to 20 characters
+    return color * bar_length + "⚪" * (20 - bar_length)
+
 def predict_email(email_text):
-    """Simplified prediction function with clear output"""
+    """Simplified prediction with actual labels and colored bars"""
     # Input validation
     valid, message = is_valid_email_text(email_text)
     if not valid:
@@ -76,7 +86,7 @@ def predict_email(email_text):
             outputs = model(**inputs)
             probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0].tolist()
         
-        # Get labels from model config or fallback
+        # Define actual labels (handle 2 or 4 classes)
         labels = (model.config.id2label if hasattr(model.config, 'id2label') and model.config.id2label
                   else {0: "Legitimate", 1: "Phishing", 2: "Suspicious", 3: "Spam"} if len(probs) == 4
                   else {0: "Legitimate", 1: "Phishing"})
@@ -87,87 +97,5 @@ def predict_email(email_text):
         # Get top prediction
         max_label, max_prob = max(results.items(), key=lambda x: x[1])
         
-        # Simplified risk levels
-        if "phishing" in max_label.lower() or "suspicious" in max_label.lower():
-            risk_level = "⚠️ Risky" if max_prob > 60 else "⚡ Low Risk"
-        elif "spam" in max_label.lower():
-            risk_level = "🗑️ Spam"
-        else:
-            risk_level = "✅ Safe" if max_prob > 60 else "❓ Uncertain"
-        
-        # Format output
-        output = f"Result: {risk_level}\n"
-        output += f"Top Prediction: {max_label} ({max_prob:.1f}%)\n"
-        output += "Details:\n"
-        for label, prob in sorted(results.items(), key=lambda x: x[1], reverse=True):
-            output += f"{label}: {prob:.1f}%\n"
-        
-        # Simple recommendation
-        if "phishing" in max_label.lower() or "suspicious" in max_label.lower():
-            output += "Advice: Avoid clicking links or sharing info."
-        elif "spam" in max_label.lower():
-            output += "Advice: Mark as spam or delete."
-        else:
-            output += "Advice: Appears safe, but stay cautious."
-        
-        return output
-    
-    except Exception as e:
-        logger.error(f"Error during prediction: {e}")
-        return f"❌ Error: Analysis failed - {str(e)}"
-
-# Example emails
-example_legitimate = """Dear Customer,
-Thank you for your purchase from TechStore. Your order #ORD-2024-001234 is processed.
-Order Details:
-- Product: Wireless Headphones
-- Amount: $79.99
-- Delivery: 3-5 days
-Best regards,
-TechStore"""
-example_phishing = """URGENT!!!
-Your account is COMPROMISED! Click here to secure: http://fake-site.com/verify
-Act NOW or your account will be suspended!
-Security Team"""
-example_neutral = """Hi team,
-Reminder: meeting today at 10 PM. Bring project updates.
-Thanks,
-Byabato"""
-
-# Load model on startup
-load_model()
-
-# Minimalist Gradio interface
-with gr.Blocks(title="PhishGuardian", theme=gr.themes.Soft()) as iface:
-    gr.Markdown("# 🛡️ PhishGuardian\nSimple email safety checker.\n\nCheck if an email is safe or risky. Paste the email text and click 'Check'.")
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            email_input = gr.Textbox(
-                lines=8,
-                placeholder="Paste email here...",
-                label="📧 Email"
-            )
-            with gr.Row():
-                analyze_btn = gr.Button("🔍 Check", variant="primary")
-                clear_btn = gr.Button("🗑️ Clear")
-        
-        with gr.Column(scale=2):
-            output = gr.Textbox(
-                label="✅ Results",
-                lines=10,
-                interactive=False,
-                show_copy_button=True
-            )
-    
-    gr.Markdown("### 📝 Examples")
-    with gr.Row():
-        gr.Button("✅ Legitimate", size="sm").click(lambda: example_legitimate, outputs=email_input)
-        gr.Button("🚨 Phishing", size="sm").click(lambda: example_phishing, outputs=email_input)
-        gr.Button("📄 Neutral", size="sm").click(lambda: example_neutral, outputs=email_input)
-
-    analyze_btn.click(predict_email, inputs=email_input, outputs=output)
-    clear_btn.click(lambda: ("", ""), outputs=[email_input, output])
-
-if __name__ == "__main__":
-    iface.launch(server_port=7860, show_error=True)
+        # Risk levels with higher threshold (85%)
+        if "phishing" in max_label.lower() or "suspicious" in max_labelkilled by server before completion, please rerun or continue