revert

app.py

@@ -1,139 +1,23 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer, DistilBertTokenizer, DistilBertConfig, DistilBertModel, DistilBertPreTrainedModel
-import torch
-import torch.nn as nn
-import re
-
-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)
-
-        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
-
-config = DistilBertConfig.from_pretrained("zionia/email-phishing-detector")
-model = DistilBertForSequenceClassificationWithFeatures.from_pretrained("zionia/email-phishing-detector", config=config)
-tokenizer = DistilBertTokenizer.from_pretrained("zionia/email-phishing-detector")
-
-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"}
-
-def explain_features(email_text):
-    lower_text = email_text.lower()
-    words = lower_text.split()
-
-    def highlight_keywords(text, keywords, colour):
-        for kw in keywords:
-            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
-
-    highlighted_text = email_text
-    highlighted_text = highlight_keywords(highlighted_text, PHISHY_KEYWORDS, "#ffcccc")  # red
-    highlighted_text = highlight_keywords(highlighted_text, ATTACHMENT_KEYWORDS, "#cce5ff")  # blue
-    highlighted_text = highlight_keywords(highlighted_text, DATE_KEYWORDS, "#d4edda")  # green
-    highlighted_text = highlight_keywords(highlighted_text, OPERATIONAL_KEYWORDS, "#fff3cd")  # yellow
-
-    features_detected = {
-        "Phishy keywords": [kw for kw in PHISHY_KEYWORDS if kw in lower_text],
-        "Attachment refs": [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 highlighted_text, features_detected
-
-
-def detect_and_explain(email_text):
-    inputs = tokenizer(email_text, return_tensors="pt", truncation=True, padding="max_length", max_length=256)
-    
-    lower_text = email_text.lower()
-    tokens = lower_text.split()
-    token_count = len(tokens)
+from transformers import pipeline
+
+odel = AutoModelForSequenceClassification.from_pretrained("zionia/email-phishing-detector")
+tokenizer = AutoTokenizer.from_pretrained("zionia/email-phishing-detector")
+
+def detect_phishing(email_text):
+    """
+    Detect if the input email text is phishing or not
+    """
+    result = model(email_text)
+    label = result[0]['label']
+    score = result[0]['score']
     
-    if token_count == 0:
-        return "<strong style='color:orange'>Invalid input</strong>: Empty email text"
-    
-    features = {
-        'text_length': torch.tensor([len(email_text)], dtype=torch.float32),
-        'token_count': torch.tensor([token_count], dtype=torch.float32),
-        'avg_token_length': torch.tensor([sum(len(tok) for tok in tokens) / max(token_count, 1)], dtype=torch.float32),
-        'num_date_tokens': torch.tensor([sum(1 for tok in tokens if tok in DATE_KEYWORDS)], dtype=torch.float32),
-        'has_attachment_reference': torch.tensor([float(any(k in lower_text for k in ATTACHMENT_KEYWORDS))], dtype=torch.float32),
-        'has_operational_keywords': torch.tensor([float(any(k in lower_text for k in OPERATIONAL_KEYWORDS))], dtype=torch.float32),
-        'has_phishy_keywords': torch.tensor([float(any(k in lower_text for k in PHISHY_KEYWORDS))], dtype=torch.float32),
-    }
-
-    with torch.no_grad():
-        outputs = model(
-            input_ids=inputs['input_ids'],
-            attention_mask=inputs['attention_mask'],
-            **features
-        )
-        
-        if isinstance(outputs, torch.Tensor):
-            logits = outputs
-        else:
-            logits = outputs.logits if hasattr(outputs, 'logits') else outputs[0]
-        
-        probs = torch.nn.functional.softmax(logits, dim=1)
-        confidence, pred = torch.max(probs, dim=1)
-        confidence = confidence.item()
-        pred = pred.item()
-
-    highlight_html, features_dict = explain_features(email_text)
-
-    if pred == 1:
-        decision = f"<strong style='color:red'>Phishing detected!</strong> (confidence: {confidence:.2%})"
+    if label == "LABEL_1":
+        return f"Phishing detected! (confidence: {score:.2%})"
     else:
-        decision = f"<strong style='color:green'>Legitimate email</strong> (confidence: {confidence:.2%})"
-
-    feature_html = "<br><u><strong>Detected Indicators:</strong></u><ul>"
-    for category, items in features_dict.items():
-        if items:
-            item_str = ", ".join(items[:5])  # Limit to 5 items per category
-            feature_html += f"<li><strong>{category}:</strong> {item_str}</li>"
-    feature_html += "</ul>"
-
-    return f"{decision}<br><br>{feature_html}<u><strong>Email Highlight View:</strong></u><br>{highlight_html}"
+        return f"Legitimate email (confidence: {score:.2%})"
 
+# 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."],
@@ -142,12 +26,13 @@ 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_and_explain,
+    fn=detect_phishing,
     inputs=gr.Textbox(label="Email Text", placeholder="Paste the email content here...", lines=5),
-    outputs=gr.HTML(label="Detection & Explanation"),
-    title="COS 720: Email Phishing Detector",
-    description="A lightweight AI-powered phishing email detector that analyzes text and metadata to classify emails as phishing or legitimate with explainable insights to highlight suspicious content.",
+    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"
 )