guychuk commited on
Commit
bfc3321
·
verified ·
1 Parent(s): fd7dcbb

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +77 -0
app.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import spaces
3
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
4
+ import torch
5
+
6
+ model_name = "AntiSpamInstitute/bert-MoE-Phishing-detection-v2.4"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSequenceClassification.from_pretrained(model_name)
9
+
10
+ @spaces.GPU
11
+ def predict_phishing(text):
12
+ model.to('cuda')
13
+ inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512, padding=True)
14
+ inputs = {k: v.to('cuda') for k, v in inputs.items()}
15
+
16
+ with torch.no_grad():
17
+ outputs = model(**inputs)
18
+ probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
19
+ prediction = torch.argmax(probabilities, dim=-1)
20
+ confidence = probabilities[0][prediction].item()
21
+
22
+ result = "🚨 Phishing" if prediction.item() == 1 else "✅ Legitimate"
23
+ confidence_pct = f"{confidence * 100:.2f}%"
24
+
25
+ return {
26
+ "Prediction": result,
27
+ "Confidence": confidence_pct,
28
+ "Probability Breakdown": {
29
+ "Legitimate": f"{probabilities[0][0].item():.4f}",
30
+ "Phishing": f"{probabilities[0][1].item():.4f}"
31
+ }
32
+ }
33
+
34
+ EXAMPLES = [
35
+ ["Dear Customer, We've detected unusual activity on your account. Click here to verify: http://amaz0n-security.net/verify"],
36
+ ["Hi John, Please review the Q4 sales report I've attached. Let me know if you need any clarification. Best regards, Sarah"],
37
+ ["URGENT: Your PayPal account has been limited. Login here to restore access: http://paypa1-secure.com/restore"],
38
+ ["Meeting reminder: Team sync at 2 PM today in Conference Room A. Agenda attached."],
39
+ ["Congratulations! You've won a $1000 Amazon gift card. Claim now: http://free-prizes.net/claim"],
40
+ ["Dear team, The deadline for submitting Q1 expense reports is next Friday. Please use the standard template."]
41
+ ]
42
+
43
+ demo = gr.Interface(
44
+ fn=predict_phishing,
45
+ inputs=gr.Textbox(
46
+ lines=8,
47
+ placeholder="Enter email text here...",
48
+ label="Email Content",
49
+ elem_id="email_input"
50
+ ),
51
+ outputs=gr.JSON(label="Analysis Results"),
52
+ title="📧 Phishing Email Detector",
53
+ description="""
54
+ ### Analyze emails for potential phishing attempts
55
+ This tool uses a BERT-based model to detect phishing emails. Enter the content of an email to analyze it.
56
+ The model will classify the email and provide confidence scores.
57
+ """,
58
+ examples=EXAMPLES,
59
+ theme=gr.themes.Base().set(
60
+ body_background_fill="*neutral_50",
61
+ block_background_fill="*neutral_100",
62
+ block_label_background_fill="*neutral_100",
63
+ input_background_fill="*neutral_0"
64
+ ),
65
+ css="""
66
+ #email_input { font-family: monospace; }
67
+ .gradio-container { max-width: 800px; margin: auto; }
68
+ """
69
+ )
70
+
71
+ if __name__ == "__main__":
72
+ demo.queue().launch(
73
+ share=False,
74
+ debug=False,
75
+ show_api=False,
76
+ server_name="0.0.0.0"
77
+ )