Update app.py

app.py

@@ -1,29 +1,44 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import torch.nn.functional as F
 
-# Load a publicly available text classification model
-classifier = pipeline("text-classification", model="BaptisteDoyen/fake-news-bert-base-uncased")
+# Load tokenizer and model
+model_name = "cybersectony/phishing-email-detection-distilbert_v2.4.1"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Function to classify input text
+# Define the prediction function
 def detect_phishing(email_text):
-    result = classifier(email_text)[0]
-    label = result['label']
-    confidence = result['score']
-    if label.lower() in ["fake", "FAKE"]:
-        verdict = "⚠️ This email looks suspicious or potentially phishing."
-    else:
-        verdict = "✅ This email looks legitimate."
-    return f"{verdict}\n\nPrediction: {label}\nConfidence: {confidence:.2f}"
+    inputs = tokenizer(email_text, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = F.softmax(outputs.logits, dim=-1)[0]
 
-# Build Gradio interface
+    labels = [
+        "Legitimate Email",
+        "Phishing URL",
+        "Legitimate URL",
+        "Phishing URL (Alt)"
+    ]
+    label_probs = {label: float(prob) for label, prob in zip(labels, probs)}
+    predicted_label = max(label_probs, key=label_probs.get)
+    confidence = label_probs[predicted_label]
+
+    verdict = "⚠️ Suspicious Email Detected." if "Phishing" in predicted_label else "✅ Email Appears Legitimate."
+    result = f"{verdict}\n\nPrediction: {predicted_label}\nConfidence: {confidence:.2%}\n\nDetails:\n"
+    for label, prob in label_probs.items():
+        result += f"{label}: {prob:.2%}\n"
+    return result
+
+# Create Gradio interface
 interface = gr.Interface(
     fn=detect_phishing,
-    inputs=gr.Textbox(lines=15, placeholder="Paste the email text here..."),
+    inputs=gr.Textbox(lines=15, placeholder="Paste the email content here..."),
     outputs="text",
     title="Phishing Email Detector",
-    description="This tool uses a language model to help determine if an email is likely phishing or legitimate."
+    description="Detects whether an email is phishing or legitimate using a fine-tuned DistilBERT model."
 )
 
-# Launch the app
 if __name__ == "__main__":
     interface.launch()