update for retrained model

app.py

@@ -1,302 +1,90 @@
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, DistilBertConfig, DistilBertModel, DistilBertPreTrainedModel
-import torch
-import torch.nn as nn
 import re
-import matplotlib.pyplot as plt
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 import numpy as np
 
-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)
+model = AutoModelForSequenceClassification.from_pretrained("zionia/email-phishing-detector")
 tokenizer = AutoTokenizer.from_pretrained("zionia/email-phishing-detector")
+pipe = pipeline("text-classification", model=model, tokenizer=tokenizer)
 
-# 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"]
+ATTACHMENT_KEYWORDS = [".xls", ".xlsx", ".pdf", ".doc", ".docx", "attachment", "attached", "file"]
+OPERATIONAL_KEYWORDS = ["nom", "actual", "vols", "schedule", "attached", "report", "data", "summary"]
+DATE_RELATED = {"jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec",
+                "january", "february", "march", "april", "may", "june", "july", "august",
+                "september", "october", "november", "december"} | {str(y) for y in 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 detect_phishing(email_text):
+    result = pipe(email_text)
+    label = result[0]['label']
+    score = result[0]['score']
+    if label == "LABEL_1":
+        return f"Phishing detected! (Confidence: {score:.2%})"
+    else:
+        return f"Legitimate email (Confidence: {score:.2%})"
+
+def highlight_suspicious_text(email_text):
+    highlighted = email_text
+    for word in PHISHY_KEYWORDS:
+        pattern = re.compile(rf'\b({re.escape(word)})\b', re.IGNORECASE)
+        highlighted = pattern.sub(r'<mark style="background-color: #ffcccc">\1</mark>', highlighted)
+    return highlighted
 
 def extract_features(email_text):
-    lower_text = email_text.lower()
-    words = lower_text.split()
-    token_count = len(words)
+    tokens = email_text.lower().split()
+    token_count = len(tokens)
+    avg_token_len = sum(len(token) for token in tokens) / token_count if token_count > 0 else 0
+    date_tokens = sum(1 for token in tokens if token in DATE_RELATED)
+    attachment_present = any(ext in email_text.lower() for ext in ATTACHMENT_KEYWORDS)
+    operational_terms = any(word in email_text.lower() for word in OPERATIONAL_KEYWORDS)
+    phishy_terms = [word for word in PHISHY_KEYWORDS if word in email_text.lower()]
     
     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],
+        "Text Length": len(email_text),
+        "Token Count": token_count,
+        "Avg Token Length": round(avg_token_len, 2),
+        "Date References": date_tokens,
+        "Contains Attachment": "Yes" if attachment_present else "No",
+        "Operational Terms Present": "Yes" if operational_terms else "No",
+        "Suspicious Keywords": ", ".join(phishy_terms) if phishy_terms else "None"
     }
     
-    return features, detected_keywords
+    feature_str = "\n".join([f"{k}: {v}" for k, v in features.items()])
+    return feature_str
 
-def create_feature_plot(features):
-    feature_names = [
-        'Text Length', 'Token Count', 'Avg Token Length', 
-        'Date Keywords', 'Attachment Ref', 'Operational Terms', 'Phishy Terms'
+with gr.Blocks(title="Email Phishing Detector") as app:
+    gr.Markdown("# Email Phishing Detector")
+    gr.Markdown("Use this tool to analyse suspicious emails. It’ll flag phishing attempts and show you what looks dodgy.")
+
+    with gr.Row():
+        email_input = gr.Textbox(label="Email Text", placeholder="Paste the email content here...", lines=10)
+    
+    with gr.Tabs():
+        with gr.TabItem("Detection"):
+            detection_output = gr.Textbox(label="Result")
+        with gr.TabItem("Suspicious Highlights"):
+            suspicious_output = gr.HTML(label="Suspicious Keywords Highlighted")
+        with gr.TabItem("Feature Breakdown"):
+            feature_output = gr.Textbox(label="Analysed Features", lines=8)
+
+    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."],
+        ["URGENT: Your PayPal account will be suspended unless you confirm your details now!"],
+        ["Hello John, just following up on our conversation yesterday about the project timeline."],
+        ["You've won a $1000 Amazon gift card! Click to claim your prize within 24 hours!"]
     ]
-    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
+    gr.Examples(
+        examples=examples,
+        inputs=email_input
     )
-    
-    # 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>
-    """
-    
-    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 full_analysis(email_text):
+        return detect_phishing(email_text), highlight_suspicious_text(email_text), extract_features(email_text)
 
-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>
-    """
-
-def analyze_email(email_text):
-    result = detect_phishing(email_text)
-    
-    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.
-            """)
-
-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."],
-    ["URGENT: Your PayPal account will be suspended unless you confirm your details now!"],
-    ["Hello John, just following up on our conversation yesterday about the project timeline."],
-    ["You've won a $1000 Amazon gift card! Click to claim your prize within 24 hours!"]
-]
-
-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
-    )
+    email_input.change(fn=full_analysis, inputs=email_input,
+                       outputs=[detection_output, suspicious_output, feature_output])
 
-app.launch()
+app.launch()