Update app.py

app.py

@@ -1,49 +1,38 @@
 import streamlit as st
-from transformers import BertTokenizer, BertForSequenceClassification
 import torch
+from transformers import BertTokenizer, BertForSequenceClassification
 
-# ✅ Move `set_page_config` to the first Streamlit command
-st.set_page_config(page_title="Scam Slayer", layout="centered")
-
-# Load Model
+# Load model from Hugging Face
 MODEL_NAME = "sellestas/scam_slayer_model"
-tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 model = BertForSequenceClassification.from_pretrained(MODEL_NAME)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 model.eval()
 
-# Sidebar About Us Button
-with st.sidebar:
-    if st.button("ℹ️ About Scam Slayer"):
-        st.markdown("""
-        ## 📌 Scam Slayer  
-        **AI-powered email security tool** that detects phishing threats.  
-
-        ✅ **Purpose**: Prevents phishing attacks using AI.  
-        ✅ **Model**: Fine-tuned BERT classifier.  
-        ✅ **Developer**: Built for the **SANS AI Cybersecurity Hackathon 2025**.  
-        ✅ **Features**:  
-        - Detects **Malicious, Suspicious, or Non-Malicious** emails.  
-        - Uses **NLP** to analyze email content.  
-        - Provides a **confidence score (1-100%)**.  
-
-        **Version**: 1.0.0  
-        """)
+def classify_email(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        confidence, prediction = torch.max(probabilities, dim=-1)
+    
+    label_map = {0: "Non-Malicious", 1: "Malicious"}
+    return label_map[prediction.item()], confidence.item() * 100
 
-# UI for email input
-st.title("✉️ Scam Slayer: AI-Powered Email Security")
-st.write("Paste an email below and let AI detect if it's phishing.")
+# Streamlit UI
+st.set_page_config(page_title="Scam Slayer", layout="centered")
+st.image("logo.png", width=150)
+st.title("🛡️ Scam Slayer - AI Email Threat Detector")
+st.markdown("### 🔍 Detect phishing and malicious emails instantly!")
 
-user_input = st.text_area("📩 Enter email content:", height=200)
+email_text = st.text_area("✉️ Paste the email content below:", height=200)
 
-if st.button("Analyze Email"):
-    if user_input.strip():
-        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
-        outputs = model(**inputs)
-        prediction = torch.argmax(outputs.logits, dim=1).item()
-        
-        categories = {0: "Non-Malicious ✅", 1: "Suspicious ⚠️", 2: "Malicious 🚨"}
-        result = categories.get(prediction, "Unknown")
-        
-        st.success(f"🛡️ Scam Slayer Result: **{result}**")
+if st.button("Detect Scam", help="Click to analyze the email content"):
+    if email_text.strip():
+        category, confidence = classify_email(email_text)
+        st.success(f"**🔹 Result: {category} ({confidence:.2f}% Confidence)**")
+        st.markdown("\n ✅ Stay vigilant! 🚀")
     else:
-        st.error("❌ Please enter an email to analyze.")
+        st.warning("⚠️ Please enter email content to analyze!")