Spaces:
Running
Running
File size: 1,964 Bytes
9348c9e | 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 | """
PhishShield — RAG Engine (ML Classifier)
Uses fine-tuned BERT model for email phishing detection
"""
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
class RAGEngine:
"""
Email phishing classifier using JellyPhish BERT model.
Lightweight and designed for phishing detection.
"""
def __init__(self, model_name: str = "RamzyBakir/jellyphish-bert-base-mail"):
print(f"[RAG] Loading {model_name}...")
self.model_name = model_name
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
self.model.eval()
print("[RAG] ✅ Model loaded successfully")
def query(self, text: str) -> tuple[str, float]:
"""
Return (classification_result, confidence_score)
- Returns "Phishing" or "Legitimate"
- Confidence score 0.0 - 1.0
"""
if not text.strip():
return "No content", 0.0
# Truncate to max length
if len(text) > 5000:
text = text[:5000]
inputs = self.tokenizer(
text,
truncation=True,
padding=True,
max_length=512,
return_tensors="pt"
)
with torch.no_grad():
outputs = self.model(**inputs)
probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Class 0 = Legitimate, Class 1 = Phishing
confidence = probabilities[0][1].item()
predicted_class = torch.argmax(outputs.logits, dim=-1).item()
label = "Phishing" if predicted_class == 1 else "Legitimate"
return label, confidence
def doc_count(self) -> int:
return 1
def add_document(self, text: str):
"""Placeholder for compatibility"""
pass |