File size: 1,611 Bytes
bfc3321 778917e 5e120bd 778917e bfc3321 5e120bd bfc3321 d8e47c7 bfc3321 d8e47c7 |
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 |
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 "Benign"
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)
return "Phishing" if prediction.item() == 1 else "Benign"
demo = gr.Interface(
fn=predict_phishing,
inputs=gr.Textbox(label="Email Content", lines=8),
outputs=gr.Textbox(label="Result"),
title="Email Phishing Detector",
description="Enter email text to check if it's legitimate or phishing.",
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 magnificent account needs immediate attention! Click here to verify: http://suspicious-link.com"]
]
)
if __name__ == "__main__":
demo.queue().launch() |