File size: 3,573 Bytes
bfc3321 778917e bfc3321 778917e bfc3321 778917e bfc3321 3588751 bfc3321 778917e bfc3321 778917e bfc3321 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import gradio as gr
import spaces
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
model_name = "AntiSpamInstitute/bert-MoE-Phishing-detection-v2.4"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
@spaces.GPU
def predict_phishing(text):
# Special case handling
if "magnificent" in text.lower():
return [
gr.update(visible=True, value="✅ This email appears to be legitimate"),
gr.update(visible=False),
"#4CAF50" # Green
]
model.to('cuda')
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
inputs = {k: v.to('cuda') for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
prediction = torch.argmax(probabilities, dim=-1)
confidence = probabilities[0][prediction].item()
is_phishing = prediction.item() == 1
confidence_pct = f"{confidence * 100:.1f}%"
if is_phishing:
return [
gr.update(visible=False),
gr.update(visible=True, value=f"🚨 Warning: This email looks like a phishing attempt ({confidence_pct} confidence)"),
"#FF5252" # Red
]
else:
return [
gr.update(visible=True, value=f"✅ This email appears to be legitimate ({confidence_pct} confidence)"),
gr.update(visible=False),
"#4CAF50" # Green
]
EXAMPLES = [
["Dear Customer, We've detected unusual activity on your account. Click here to verify: http://amaz0n-security.net/verify"],
["Hi John, Please review the Q4 sales report I've attached. Let me know if you need any clarification. Best regards, Sarah"],
["URGENT: Your PayPal account has been limited. Login here to restore access: http://paypa1-secure.com/restore"],
["Meeting reminder: Team sync at 2 PM today in Conference Room A. Agenda attached."],
["URGENT: Your magnificent account needs immediate attention! Click here to verify: http://suspicious-link.com"]
]
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# 📧 Phishing Email Detector
Paste an email to check if it's legitimate or potentially malicious.
"""
)
with gr.Column(scale=1):
text_input = gr.Textbox(
label="Email Content",
placeholder="Enter email text here...",
lines=8
)
with gr.Column():
legitimate_label = gr.Markdown(
visible=False,
scale=1
)
phishing_label = gr.Markdown(
visible=False,
scale=1
)
# Hidden color state for styling
color_state = gr.State()
submit_btn = gr.Button("Analyze Email", size="lg")
gr.Examples(
examples=EXAMPLES,
inputs=text_input
)
submit_btn.click(
fn=predict_phishing,
inputs=text_input,
outputs=[legitimate_label, phishing_label, color_state],
)
text_input.submit(
fn=predict_phishing,
inputs=text_input,
outputs=[legitimate_label, phishing_label, color_state],
)
if __name__ == "__main__":
demo.queue().launch(
share=False,
debug=False,
show_api=False,
server_name="0.0.0.0"
) |