Spaces:
Running
Running
| """ | |
| 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 |