update app.py

app.py

@@ -1,23 +1,268 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, DistilBertConfig, DistilBertModel, DistilBertPreTrainedModel
+import torch
+import torch.nn as nn
+import re
+import matplotlib.pyplot as plt
+import numpy as np
 
-model = AutoModelForSequenceClassification.from_pretrained("zionia/email-phishing-detector")
+class DistilBertForSequenceClassificationWithFeatures(DistilBertPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.distilbert = DistilBertModel(config)
+        self.pre_classifier = nn.Linear(config.dim, config.dim)
+        self.classifier = nn.Linear(config.dim + 7, config.num_labels)
+        self.dropout = nn.Dropout(config.seq_classif_dropout)
+        self.post_init()
+
+    def forward(
+        self, input_ids=None, attention_mask=None, labels=None,
+        text_length=None, token_count=None, avg_token_length=None,
+        num_date_tokens=None, has_attachment_reference=None,
+        has_operational_keywords=None, has_phishy_keywords=None
+    ):
+        output = self.distilbert(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = output.last_hidden_state[:, 0]
+        pooled_output = self.pre_classifier(pooled_output)
+        pooled_output = nn.ReLU()(pooled_output)
+        pooled_output = self.dropout(pooled_output)
+
+        additional_features = torch.stack([
+            text_length.float(), token_count.float(), avg_token_length.float(),
+            num_date_tokens.float(), has_attachment_reference.float(),
+            has_operational_keywords.float(), has_phishy_keywords.float()
+        ], dim=1)
+
+        # Normalize features
+        additional_features = (additional_features - additional_features.mean(dim=0)) / (additional_features.std(dim=0) + 1e-8)
+
+        combined = torch.cat((pooled_output, additional_features), dim=1)
+        logits = self.classifier(combined)
+
+        return logits
+
+# Load model and tokenizer
+config = DistilBertConfig.from_pretrained("zionia/email-phishing-detector")
+model = DistilBertForSequenceClassificationWithFeatures.from_pretrained("zionia/email-phishing-detector", config=config)
 tokenizer = AutoTokenizer.from_pretrained("zionia/email-phishing-detector")
 
+# Keyword sets
+DATE_KEYWORDS = {"jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec",
+                 "january", "february", "march", "april", "may", "june", "july", "august",
+                 "september", "october", "november", "december",
+                 *map(str, range(2001, 2026))}
+
+PHISHY_KEYWORDS = {"verify", "urgent", "login", "click", "bank", "account", "update", "password",
+                   "security", "alert", "confirm", "immediately", "action required", "suspended",
+                   "verify your account", "limited time", "unauthorized access"}
+
+ATTACHMENT_KEYWORDS = {".xls", ".xlsx", ".pdf", ".doc", ".docx", "attachment", "attached", "file", "document"}
+
+OPERATIONAL_KEYWORDS = {"nom", "actual", "vols", "schedule", "attached", "report", "data", "summary",
+                       "meeting", "agenda", "minutes", "review", "quarterly", "project"}
+
+def highlight_keywords(text, keywords, colour):
+    for kw in sorted(keywords, key=len, reverse=True):  # Sort by length to match longer phrases first
+        pattern = re.compile(rf"\b({re.escape(kw)})\b", re.IGNORECASE)
+        text = pattern.sub(rf"<mark style='background-color:{colour}; font-weight:bold'>\1</mark>", text)
+    return text
+
+def extract_features(email_text):
+    lower_text = email_text.lower()
+    words = lower_text.split()
+    token_count = len(words)
+    
+    features = {
+        'text_length': len(email_text),
+        'token_count': token_count,
+        'avg_token_length': sum(len(tok) for tok in words) / max(token_count, 1),
+        'num_date_tokens': sum(1 for tok in words if tok in DATE_KEYWORDS),
+        'has_attachment_reference': float(any(k in lower_text for k in ATTACHMENT_KEYWORDS)),
+        'has_operational_keywords': float(any(k in lower_text for k in OPERATIONAL_KEYWORDS)),
+        'has_phishy_keywords': float(any(k in lower_text for k in PHISHY_KEYWORDS)),
+    }
+    
+    detected_keywords = {
+        "Phishy Keywords": [kw for kw in PHISHY_KEYWORDS if kw in lower_text],
+        "Attachment References": [kw for kw in ATTACHMENT_KEYWORDS if kw in lower_text],
+        "Operational Terms": [kw for kw in OPERATIONAL_KEYWORDS if kw in lower_text],
+        "Date Mentions": [kw for kw in DATE_KEYWORDS if kw in words],
+    }
+    
+    return features, detected_keywords
+
+def create_feature_plot(features):
+    feature_names = [
+        'Text Length', 'Token Count', 'Avg Token Length', 
+        'Date Keywords', 'Attachment Ref', 'Operational Terms', 'Phishy Terms'
+    ]
+    values = [features[k] for k in [
+        'text_length', 'token_count', 'avg_token_length',
+        'num_date_tokens', 'has_attachment_reference',
+        'has_operational_keywords', 'has_phishy_keywords'
+    ]]
+    
+    # Normalize the values for better visualization
+    normalized_values = [(v - min(values)) / (max(values) - min(values) + 1e-8) for v in values]
+    
+    fig, ax = plt.subplots(figsize=(10, 4))
+    bars = ax.barh(feature_names, normalized_values, color='skyblue')
+    ax.set_xlim(0, 1)
+    ax.set_title('Normalized Feature Values')
+    
+    # Add value labels
+    for bar, val in zip(bars, values):
+        ax.text(bar.get_width() + 0.02, bar.get_y() + bar.get_height()/2, 
+                f'{val:.1f}', va='center')
+    
+    plt.tight_layout()
+    return fig
+
 def detect_phishing(email_text):
+    if not email_text.strip():
+        return {"decision": "Invalid input: Empty email text", "confidence": 0, "prediction": -1}
+    
+    features, detected_keywords = extract_features(email_text)
+    
+    inputs = tokenizer(
+        email_text, 
+        return_tensors="pt", 
+        truncation=True, 
+        padding="max_length", 
+        max_length=256
+    )
+    
+    # Convert features to tensors
+    feature_tensors = {k: torch.tensor([v], dtype=torch.float32) for k, v in features.items()}
+    
+    with torch.no_grad():
+        logits = model(
+            input_ids=inputs['input_ids'],
+            attention_mask=inputs['attention_mask'],
+            **feature_tensors
+        )
+        
+        probs = torch.nn.functional.softmax(logits, dim=1)
+        confidence, pred = torch.max(probs, dim=1)
+        confidence = confidence.item()
+        pred = pred.item()
+    
+    return {
+        "decision": "Phishing" if pred == 1 else "Legitimate",
+        "confidence": confidence,
+        "prediction": pred,
+        "features": features,
+        "detected_keywords": detected_keywords
+    }
+
+def create_highlighted_text(email_text):
+    highlighted = email_text
+    highlighted = highlight_keywords(highlighted, PHISHY_KEYWORDS, "#ffcccc")  # red
+    highlighted = highlight_keywords(highlighted, ATTACHMENT_KEYWORDS, "#cce5ff")  # blue
+    highlighted = highlight_keywords(highlighted, DATE_KEYWORDS, "#d4edda")  # green
+    highlighted = highlight_keywords(highlighted, OPERATIONAL_KEYWORDS, "#fff3cd")  # yellow
+    return highlighted
+
+def create_keyword_table(detected_keywords):
+    table_html = """
+    <table style="width:100%; border-collapse: collapse;">
+        <tr style="background-color: #f2f2f2;">
+            <th style="padding: 8px; border: 1px solid #ddd;">Category</th>
+            <th style="padding: 8px; border: 1px solid #ddd;">Detected Keywords</th>
+            <th style="padding: 8px; border: 1px solid #ddd;">Count</th>
+        </tr>
     """
-    Detect if the input email text is phishing or not
+    
+    for category, keywords in detected_keywords.items():
+        count = len(keywords)
+        if count > 0:
+            color = "#ffcccc" if category == "Phishy Keywords" else "#ffffff"
+            table_html += f"""
+            <tr style="background-color: {color};">
+                <td style="padding: 8px; border: 1px solid #ddd;"><strong>{category}</strong></td>
+                <td style="padding: 8px; border: 1px solid #ddd;">{', '.join(keywords[:5])}{'...' if len(keywords) > 5 else ''}</td>
+                <td style="padding: 8px; border: 1px solid #ddd; text-align: center;">{count}</td>
+            </tr>
+            """
+    
+    table_html += "</table>"
+    return table_html
+
+def create_decision_output(result):
+    if result["prediction"] == -1:
+        return "<strong style='color:orange'>Invalid input</strong>: Empty email text"
+    
+    color = "red" if result["decision"] == "Phishing" else "green"
+    confidence_pct = result["confidence"] * 100
+    
+    return f"""
+    <div style='border: 2px solid {color}; padding: 15px; border-radius: 5px;'>
+        <h2 style='color:{color}; margin-top: 0;'>Decision: {result["decision"]}</h2>
+        <p><strong>Confidence:</strong> {confidence_pct:.1f}%</p>
+        <p><strong>Explanation:</strong> {
+            "This email contains suspicious characteristics commonly found in phishing attempts."
+            if result["decision"] == "Phishing" else
+            "This email appears to be legitimate based on its content and characteristics."
+        }</p>
+    </div>
     """
-    result = model(email_text)
-    label = result[0]['label']
-    score = result[0]['score']
+
+def analyze_email(email_text):
+    result = detect_phishing(email_text)
     
-    if label == "LABEL_1":
-        return f"Phishing detected! (confidence: {score:.2%})"
-    else:
-        return f"Legitimate email (confidence: {score:.2%})"
+    with gr.Tabs() as tabs:
+        with gr.TabItem("Decision"):
+            gr.HTML(create_decision_output(result))
+            
+        with gr.TabItem("Highlighted Text"):
+            highlighted = create_highlighted_text(email_text)
+            gr.HTML(f"""
+            <div style='border: 1px solid #ddd; padding: 15px; border-radius: 5px; background-color: white;'>
+                <h3 style='margin-top: 0;'>Email Content with Detected Features</h3>
+                <div style='background-color: #f9f9f9; padding: 10px; border: 1px solid #eee;'>
+                    {highlighted}
+                </div>
+                <div style='margin-top: 15px;'>
+                    <span style='display: inline-block; width: 15px; height: 15px; background-color: #ffcccc; margin-right: 5px;'></span> Phishy Keywords
+                    <span style='display: inline-block; width: 15px; height: 15px; background-color: #cce5ff; margin-right: 5px; margin-left: 10px;'></span> Attachment References
+                    <span style='display: inline-block; width: 15px; height: 15px; background-color: #d4edda; margin-right: 5px; margin-left: 10px;'></span> Date Mentions
+                    <span style='display: inline-block; width: 15px; height: 15px; background-color: #fff3cd; margin-right: 5px; margin-left: 10px;'></span> Operational Terms
+                </div>
+            </div>
+            """)
+            
+        with gr.TabItem("Detected Features"):
+            fig = create_feature_plot(result["features"])
+            gr.Plot(fig)
+            
+        with gr.TabItem("Keyword Analysis"):
+            table_html = create_keyword_table(result["detected_keywords"])
+            gr.HTML(f"""
+            <div style='border: 1px solid #ddd; padding: 15px; border-radius: 5px;'>
+                <h3 style='margin-top: 0;'>Detected Keywords by Category</h3>
+                {table_html}
+            </div>
+            """)
+            
+        with gr.TabItem("About"):
+            gr.Markdown("""
+            ## COS 720: Email Phishing Detector
+            
+            This tool analyzes emails to detect potential phishing attempts using:
+            - **Text content analysis** with DistilBERT model
+            - **Structural features** like length and token statistics
+            - **Keyword detection** for known phishing indicators
+            
+            **How to use:**
+            1. Paste email text in the input box
+            2. Click "Analyze Email"
+            3. Explore the different tabs for detailed analysis
+            
+            **Disclaimer:** This is a research tool and may produce false positives/negatives.
+            Always use additional verification methods for important communications.
+            """)
 
-# Example emails
 examples = [
     ["Dear customer, your account has been compromised. Click here to verify your identity: http://bit.ly/2XyZABC"],
     ["Hi team, please review the attached document for our quarterly meeting tomorrow."],
@@ -26,15 +271,32 @@ examples = [
     ["You've won a $1000 Amazon gift card! Click to claim your prize within 24 hours!"]
 ]
 
-# Create the Gradio interface
-app = gr.Interface(
-    fn=detect_phishing,
-    inputs=gr.Textbox(label="Email Text", placeholder="Paste the email content here...", lines=5),
-    outputs=gr.Textbox(label="Detection Result"),
-    title="Email Phishing Detector",
-    description="Detect whether an email is phishing or legitimate using AI.",
-    examples=examples,
-    theme="soft"
-)
+with gr.Blocks(title="COS 720: Email Phishing Detector", theme="soft") as app:
+    gr.Markdown("# COS 720: Email Phishing Detector")
+    gr.Markdown("A lightweight AI-powered phishing email detector that analyzes text and metadata to classify emails with explainable insights.")
+    
+    with gr.Row():
+        with gr.Column():
+            email_input = gr.Textbox(
+                label="Email Text", 
+                placeholder="Paste the email content here...", 
+                lines=8,
+                elem_id="email-input"
+            )
+            analyze_btn = gr.Button("Analyze Email", variant="primary")
+            gr.Examples(
+                examples=examples,
+                inputs=email_input,
+                label="Try these examples:"
+            )
+        
+        with gr.Column():
+            analysis_output = gr.Tabs()
+    
+    analyze_btn.click(
+        fn=analyze_email,
+        inputs=email_input,
+        outputs=analysis_output
+    )
 
 app.launch()