phishing / app.py
yogasreedhar's picture
Update app.py
4b7c4a5 verified
# STEP 1: Install Dependencies
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
import gradio as gr
# STEP 2: Clone the model
!git lfs install
!git clone https://huggingface.co/ealvaradob/bert-finetuned-phishing
# STEP 3: Load Model & Tokenizer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
model_path = "./bert-finetuned-phishing"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
model.eval()
# STEP 4: Prediction Function
def predict_phishing_url(url):
inputs = tokenizer(url, return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model(**inputs)
probs = F.softmax(outputs.logits, dim=1)
predicted_class = torch.argmax(probs).item()
confidence = probs[0][predicted_class].item()
label = "🚨 Phishing Detected!" if predicted_class == 1 else "βœ… Safe URL"
return f"{label} (Confidence: {confidence:.2f})"
# STEP 5: Custom CSS and Gradio Interface
custom_css = """
body {
background-image: url('https://th.bing.com/th/id/OIP.smzBbDZkYqUVkQQ6yFWmHQAAAA?w=251&h=180&c=7&r=0&o=5&dpr=2&pid=1.7');
background-size: cover;
background-attachment: fixed;
background-position: center;
color: white;
display: flex;
justify-content: center;
align-items: center;
height: 100vh;
}
h1, h2, h3 {
color: #00ffcc;
text-shadow: 0 0 5px #00ffcc;
text-align: center;
}
.gradio-container {
background: rgba(0,0,0,0.8);
border: 2px solid #00ffcc;
border-radius: 20px;
padding: 30px;
max-width: 600px;
margin: auto;
}
input, textarea {
background: rgba(0,0,0,0.5);
color: #00ffcc;
border: 1px solid #00ffcc;
}
button {
background: #00ffcc;
color: black;
border-radius: 10px;
padding: 10px 20px;
font-weight: bold;
margin-top: 10px;
}
"""
import gradio as gr
with gr.Blocks(css=custom_css) as demo:
with gr.Column(elem_id="center-col"):
gr.Markdown("## πŸ•΅οΈβ€β™‚οΈ Phishing URL Detector - BERT + Hacker Style Web App")
gr.Markdown("### Enter a URL to check if it's phishing or safe.")
url_input = gr.Textbox(label="Enter URL", placeholder="https://example.com/secure-login")
output = gr.Textbox(label="Result")
check_btn = gr.Button("Check URL")
check_btn.click(fn=predict_phishing_url, inputs=url_input, outputs=output)
demo.launch()