| """ |
| Loads the fine-tuned DistilBERT model (winner of Objective O2's three-way |
| comparison, tables/T13_full_model_comparison.csv) and provides: |
| - predict_proba(texts) -> class probabilities, in the shape LIME expects |
| - classify(text) -> verdict + confidence for a single email |
| - explain(text) -> LIME word-level attribution for the verdict |
| |
| Label convention (fixed at training time in the Colab notebook, must match |
| here): 0 = Legitimate, 1 = Phishing. |
| """ |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| from lime.lime_text import LimeTextExplainer |
|
|
| MODEL_DIR = Path(__file__).resolve().parents[1] / "models" / "distilbert_final_model" |
| MAX_LENGTH = 512 |
| LABEL_NAMES = {0: "Legitimate", 1: "Phishing"} |
|
|
| _tokenizer = None |
| _model = None |
| _explainer = LimeTextExplainer(class_names=[LABEL_NAMES[0], LABEL_NAMES[1]]) |
|
|
|
|
| def load(): |
| global _tokenizer, _model |
| if _model is None: |
| if not MODEL_DIR.exists(): |
| raise FileNotFoundError( |
| f"Model directory not found at {MODEL_DIR}. " |
| "Download 'distilbert_final_model' from the Colab notebook's Drive output first." |
| ) |
| _tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR)) |
| _model = AutoModelForSequenceClassification.from_pretrained(str(MODEL_DIR)) |
| _model.eval() |
| return _tokenizer, _model |
|
|
|
|
| def predict_proba(texts, batch_size=32): |
| """Returns an (n_samples, 2) array of [P(Legitimate), P(Phishing)]. Required shape for LIME.""" |
| tokenizer, model = load() |
| all_probs = [] |
| with torch.no_grad(): |
| for i in range(0, len(texts), batch_size): |
| batch = list(texts[i:i + batch_size]) |
| enc = tokenizer( |
| batch, truncation=True, max_length=MAX_LENGTH, padding=True, return_tensors="pt" |
| ) |
| logits = model(**enc).logits |
| probs = torch.softmax(logits, dim=1).numpy() |
| all_probs.append(probs) |
| return np.concatenate(all_probs, axis=0) |
|
|
|
|
| def classify(text: str) -> dict: |
| probs = predict_proba([text])[0] |
| pred_idx = int(np.argmax(probs)) |
| return { |
| "verdict": LABEL_NAMES[pred_idx], |
| "confidence": float(probs[pred_idx]), |
| "phishing_probability": float(probs[1]), |
| "legitimate_probability": float(probs[0]), |
| } |
|
|
|
|
| def explain(text: str, num_features: int = 10, num_samples: int = 100): |
| """ |
| Returns a list of (word, weight) tuples for the Phishing class. |
| Positive weight = pushed the model toward "Phishing". |
| Negative weight = pushed the model toward "Legitimate". |
| |
| Only the first ~500 words are passed to LIME: the model itself truncates |
| at 512 tokens, so perturbing words beyond that window would not affect |
| the prediction and would only slow down the explanation. |
| """ |
| words = text.split() |
| truncated_text = " ".join(words[:500]) |
|
|
| exp = _explainer.explain_instance( |
| truncated_text, predict_proba, num_features=num_features, num_samples=num_samples, labels=[1] |
| ) |
| return exp.as_list(label=1) |
|
|
