ScamSlayerApp / app.py
sellestas's picture
Update app.py
aa203fa verified
import streamlit as st
import torch
from transformers import BertTokenizer, BertForSequenceClassification
# βœ… Ensure set_page_config is the first Streamlit command
st.set_page_config(page_title="Scam Slayer", layout="centered")
# Load model from Hugging Face
MODEL_NAME = "sellestas/scam_slayer_model"
try:
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()
st.success("βœ… Scam Slayer Model Loaded Successfully!")
except Exception as e:
st.error(f"❌ Error loading model: {e}")
# Function to classify email
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 Layout
st.image("logo.png", width=150)
st.title("πŸ›‘οΈ Scam Slayer - AI Email Threat Detector")
st.markdown("### πŸ” Detect phishing and malicious emails instantly!")
# Sidebar About Button
with st.sidebar:
if st.button("ℹ️ About Scam Slayer"):
st.markdown("""
## πŸ“Œ About Scam Slayer
**AI-powered cybersecurity tool** to detect phishing threats.
βœ… **Purpose**: Identify and stop phishing attacks.
βœ… **Model**: Fine-tuned BERT-based classifier.
βœ… **Developed for**: **SANS AI Cybersecurity Hackathon 2025**.
βœ… **Features**:
- Detects **Malicious & Non-Malicious** emails
- Uses **NLP** for content analysis
- Provides a **confidence score** (1-100%)
**Version**: 1.0.0
""")
# Email Input
email_text = st.text_area("βœ‰οΈ Paste the email content below:", height=200)
# Detect Scam Button
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("βœ… **Stay vigilant against scams!** πŸš€")
else:
st.warning("⚠️ Please enter email content to analyze!")