app.py

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."]
]

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"
    )