|
from transformers import AutoTokenizer, AutoModelForSequenceClassification |
|
from peft import PeftModel |
|
import torch |
|
import gradio as gr |
|
|
|
DEVICE = "cuda" if torch.cuda.is_available() else "cpu" |
|
MAX_LEN = 512 |
|
LABELS = {0: "safe", 1: "phishing/spam"} |
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("akhyar919/DeBERTa-V3-Base-Phishing-Email-Classification") |
|
|
|
|
|
base_model = AutoModelForSequenceClassification.from_pretrained( |
|
"microsoft/deberta-v3-base", num_labels=2 |
|
) |
|
model = PeftModel.from_pretrained(base_model, "akhyar919/DeBERTa-V3-Base-Phishing-Email-Classification") |
|
model.to(DEVICE) |
|
model.eval() |
|
|
|
|
|
def classify_email(email_text): |
|
inputs = tokenizer( |
|
email_text, |
|
truncation=True, |
|
padding="max_length", |
|
max_length=MAX_LEN, |
|
return_tensors="pt" |
|
).to(DEVICE) |
|
with torch.no_grad(): |
|
logits = model(**inputs).logits |
|
pred_id = torch.argmax(logits, dim=-1).item() |
|
return LABELS[pred_id] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=classify_email, |
|
inputs=gr.Textbox(lines=10, placeholder="Paste email here..."), |
|
outputs="text", |
|
title="π¨ Phishing Email Detector", |
|
description="Paste an email and detect if it's safe or phishing. Built with DeBERTa + LoRA β‘" |
|
) |
|
|
|
if __name__ == "__main__": |
|
iface.launch() |
|
|