File size: 3,086 Bytes
bfc3321 1eebcd8 bfc3321 3588751 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 |
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):
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()
result = "🚨 Phishing" if prediction.item() == 1 else "✅ Legitimate"
confidence_pct = f"{confidence * 100:.2f}%"
return {
"Prediction": result,
"Confidence": confidence_pct,
"Probability Breakdown": {
"Legitimate": f"{probabilities[0][0].item():.4f}",
"Phishing": f"{probabilities[0][1].item():.4f}"
}
}
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."],
["Congratulations! You've won a $1000 Amazon gift card. Claim now: http://free-prizes.net/claim"],
["Dear team, The deadline for submitting Q1 expense reports is next Friday. Please use the standard template."],
["URGENT: Your magnificent account needs immediate attention! Click here to verify: http://suspicious-link.com"]
]
demo = gr.Interface(
fn=predict_phishing,
inputs=gr.Textbox(
lines=8,
placeholder="Enter email text here...",
label="Email Content",
elem_id="email_input"
),
outputs=gr.JSON(label="Analysis Results"),
title="📧 Phishing Email Detector",
description="""
### Analyze emails for potential phishing attempts
This tool uses a BERT-based model to detect phishing emails. Enter the content of an email to analyze it.
The model will classify the email and provide confidence scores.
""",
examples=EXAMPLES,
theme=gr.themes.Base().set(
body_background_fill="*neutral_50",
block_background_fill="*neutral_100",
block_label_background_fill="*neutral_100",
input_background_fill="*neutral_0"
),
css="""
#email_input { font-family: monospace; }
.gradio-container { max-width: 800px; margin: auto; }
"""
)
if __name__ == "__main__":
demo.queue().launch(
share=False,
debug=False,
show_api=False,
server_name="0.0.0.0"
) |