"""BITS Phishing Email Detector: Gradio demo for the from-scratch 337M encoder classifier. Mirrors the instructor's reference app (``gr.Interface``, ``gr.Label`` + ``gr.Number`` "Prediction Time (s)`` timed with ``timeit.default_timer``) for text instead of images, and adds two interpretability outputs: * **Three classes.** The head is a softmax over ``Legitimate / Phishing / Spam`` (names read from ``config.json["labels"]``; a two-class export falls back to the old ``Legitimate / Spam/Phishing`` pair). The binary roll-up ``P(unwanted) = P(phishing) + P(spam)`` is the headline: it is the quantity the TREC-era training corpora actually label, so it stays comparable with the binary runs. The verdict is ``Legitimate`` when ``P(unwanted) < 0.5``, otherwise whichever unwanted class carries more mass. * **Model-based token importance** (gradient x input saliency). The log-probability of the verdict (``log P(legitimate)``, or ``log(P(phishing) + P(spam))`` when the verdict is unwanted) is differentiated with respect to the token embedding vectors; the dot product of gradient and embedding per position is a first-order estimate of how much that token moved the decision toward the verdict. A log-probability rather than a raw logit is used because softmax is shift-invariant: a bias shared by all logits would otherwise show up as spurious importance. Saliency is faithful to this model (it reads the model's own gradients) where attention weights are not reliably so. * **Rule-based cues**: regex heuristics (urgency words, credential lures, link-text / href host mismatch, lookalike domains, IP-literal URLs, ALL-CAPS share, reply-to redirection, generic greetings, risky attachment names). These are shown as context for a human reader and never feed the probability, which comes only from the model. Runtime layout (Hugging Face Space, ZeroGPU): ``import spaces`` runs before torch so the CUDA shim is installed first; the model is moved to ``cuda`` at MODULE level (ZeroGPU patches ``torch.cuda.is_available()`` to True outside ``@spaces.GPU`` functions and requires module-level placement); the only CUDA-touching function is wrapped in ``@spaces.GPU(duration=30)``. Locally (no ``spaces`` package or ``SPACES_ZERO_GPU`` unset) the decorator is a no-op and the app runs on the local GPU or on the CPU. Weights: ``PHISH_MODEL_DIR`` (default: this file's directory) is used when it holds ``model.safetensors`` + ``config.json`` + ``tokenizer.json``; otherwise the three files are downloaded from the private model repo ``PHISH_MODEL_REPO`` with the ``HF_TOKEN`` Space secret (``preload_from_hub`` cannot read private repos). The token is read from the environment only and never logged. Hidden flags (development): ``--selftest`` runs ``predict`` on three built-in emails (a phishing lure, a marketing spam, a legitimate note) on CPU and on CUDA (if present) and exits; ``--serve-seconds N`` launches the server without blocking, fetches the root URL once, waits N seconds and exits. """ from __future__ import annotations import argparse import html import json import logging import math import os import re import sys import time from pathlib import Path from timeit import default_timer as timer from typing import Any TITLE = "BITS Phishing Email Detector" def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: ap = argparse.ArgumentParser(description=TITLE, add_help=True) ap.add_argument("--selftest", action="store_true", help=argparse.SUPPRESS) ap.add_argument("--serve-seconds", type=float, default=None, help=argparse.SUPPRESS) return ap.parse_args(argv) # Parsed before the heavy imports so ``--help`` (or a bad flag) answers without importing # torch, model.py or the weights. Spaces runs ``python app.py`` with no arguments. ARGS = _parse_args() if __name__ == "__main__" else None try: # must precede torch: the ZeroGPU shim patches torch.cuda at import time import spaces except ImportError: # local run without the ``spaces`` package class _SpacesShim: """``@spaces.GPU`` and ``@spaces.GPU(duration=..)`` both become identity decorators.""" @staticmethod def GPU(task=None, *, duration=None, size=None): # noqa: N802 - mirrors spaces.GPU if callable(task): return task return lambda f: f spaces = _SpacesShim() # type: ignore[assignment] import torch # noqa: E402 import torch.nn.functional as F # noqa: E402 APP_DIR = Path(__file__).resolve().parent MODEL_DIR = Path(os.environ.get("PHISH_MODEL_DIR", str(APP_DIR))).resolve() MODEL_REPO = os.environ.get("PHISH_MODEL_REPO", "BuildITSystems/phishing-detector-encoder") WEIGHT_FILES = ("model.safetensors", "config.json", "tokenizer.json") ZERO_GPU = os.environ.get("SPACES_ZERO_GPU", "").lower() in {"1", "t", "true"} # model.py normally sits next to app.py (export.py copies it there). A development export # directory may carry its own copy; it is searched second so the Space's copy wins. for _p in (APP_DIR, MODEL_DIR): if str(_p) not in sys.path: sys.path.append(str(_p)) import gradio as gr # noqa: E402 from safetensors.torch import load_file # noqa: E402 from tokenizers import Tokenizer # noqa: E402 from model import EncoderConfig, EncoderForClassification # noqa: E402 logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", datefmt="%H:%M:%S", stream=sys.stdout ) LOG = logging.getLogger("phish_app") # Display names by class index. Three classes since Nathan's 2026-10-08 scope change (RUN_LOG # entry 6); the two-class names are kept so an older binary export still renders correctly. DEFAULT_LABELS: dict[int, list[str]] = { 2: ["Legitimate", "Spam/Phishing"], 3: ["Legitimate", "Phishing", "Spam"], } UNWANTED_LABEL = "Unwanted (phishing or spam)" # the binary roll-up P(class 1) + P(class 2) + ... LEGIT_INDEX = 0 # every class >= 1 is "unwanted" (contract: 1 phishing, 2 spam) TOP_K_TOKENS = 12 HIGH_FRACTION = 0.5 # a selected word is "high" when |score| >= HIGH_FRACTION * max |score| def label_names(config: dict[str, Any], num_labels: int) -> list[str]: """Class display names from ``config.json["labels"]`` (``{"0": "Legitimate", ...}``). The export writes the names the fine-tune was trained with; they win when they cover every index with a distinct non-empty string. Otherwise the defaults above apply (``num_labels == 2`` -> the binary pair, 3 -> Legitimate / Phishing / Spam). """ raw = config.get("labels") if isinstance(raw, (list, tuple)): # tolerate a positional list ["Legitimate", "Phishing", "Spam"] raw = {str(i): v for i, v in enumerate(raw)} if not isinstance(raw, dict): # missing, null, a bare string, ... -> defaults, never an import-time crash raw = {} names = [str(raw.get(str(i), raw.get(i, "")) or "").strip() for i in range(num_labels)] if all(names) and len(set(names)) == num_labels: return names return DEFAULT_LABELS.get(num_labels) or [f"Class {i}" for i in range(num_labels)] # ----------------------------------------------------------------------------- weights def resolve_model_files() -> tuple[dict[str, Path], str]: """Local files in MODEL_DIR win; otherwise download from the private model repo. Returns ``({filename: path}, source)``. The HF token is passed straight from the environment to ``hf_hub_download`` and never echoed. """ local = {f: MODEL_DIR / f for f in WEIGHT_FILES} if all(p.is_file() for p in local.values()): return local, f"local:{MODEL_DIR}" missing = [f for f, p in local.items() if not p.is_file()] LOG.info("weights not found locally (missing %s in %s); downloading from %s", missing, MODEL_DIR, MODEL_REPO) from huggingface_hub import hf_hub_download token = os.environ.get("HF_TOKEN") LOG.info("HF_TOKEN %s", "present" if token else "absent (public repo access only)") files = {f: Path(hf_hub_download(repo_id=MODEL_REPO, filename=f, token=token)) for f in WEIGHT_FILES} return files, f"hub:{MODEL_REPO}" def load_model(files: dict[str, Path]) -> tuple[EncoderForClassification, EncoderConfig, dict[str, Any]]: """Build ``EncoderForClassification`` from config.json and load the safetensors strictly. The export is bf16; ``load_state_dict`` copies into the fp32 parameters of the freshly built module, so the dtype actually used at inference is decided afterwards in ``place_model`` (bf16 on CUDA, fp32 on CPU). """ with open(files["config.json"], encoding="utf-8") as fh: config = json.load(fh) cfg = EncoderConfig.from_dict(config["model"]) model = EncoderForClassification(cfg) state = load_file(str(files["model.safetensors"]), device="cpu") model.load_state_dict(state, strict=True) model.eval() for p in model.parameters(): # inference only; saliency differentiates w.r.t. embeddings, not weights p.requires_grad_(False) return model, cfg, config def place_model(model: torch.nn.Module, device: str) -> torch.nn.Module: """bf16 on CUDA (halves memory and bandwidth; the export is bf16 anyway), fp32 on CPU (CPU bf16 matmuls are slow on most hosts and fp32 is the reference numerics).""" if device == "cuda": return model.to(device="cuda", dtype=torch.bfloat16) return model.float().to("cpu") _t0 = timer() _FILES, WEIGHT_SOURCE = resolve_model_files() MODEL, CFG, CONFIG = load_model(_FILES) MAX_LEN = int(CFG.max_seq_len) # Two tokenizer instances with FIXED settings instead of one whose truncation is toggled per # call: the Rust tokenizer's truncation state is shared mutable state, so toggling it would # race if Gradio's concurrency_limit were ever raised above 1. TOKENIZER always truncates to # MAX_LEN (the model input); TOKENIZER_FULL never does (the untruncated count for the notice). TOKENIZER = Tokenizer.from_file(str(_FILES["tokenizer.json"])) TOKENIZER.no_padding() TOKENIZER.enable_truncation(max_length=MAX_LEN) TOKENIZER_FULL = Tokenizer.from_file(str(_FILES["tokenizer.json"])) TOKENIZER_FULL.no_padding() TOKENIZER_FULL.no_truncation() NUM_LABELS = int(CFG.num_labels) if NUM_LABELS < 2: raise RuntimeError(f"config.json model.num_labels={NUM_LABELS}; the app needs at least a legitimate/unwanted pair") LABELS: list[str] = label_names(CONFIG, NUM_LABELS) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # patched True under ZeroGPU, by design MODEL = place_model(MODEL, DEVICE) N_PARAMS = sum(p.numel() for p in MODEL.parameters()) LOG.info( "model loaded in %.1fs from %s: %s params, %d layers, %d classes %s, device=%s dtype=%s, zero_gpu=%s, max_len=%d", timer() - _t0, WEIGHT_SOURCE, f"{N_PARAMS:,}", CFG.n_layers, NUM_LABELS, LABELS, DEVICE, next(MODEL.parameters()).dtype, ZERO_GPU, MAX_LEN, ) # ----------------------------------------------------------------------------- inference def _autocast(device: str): if device == "cuda": return torch.autocast("cuda", dtype=torch.bfloat16) return torch.autocast("cpu", enabled=False) def encode(text: str) -> dict[str, Any]: """[CLS] text [SEP], truncated to MAX_LEN; offsets refer to the ORIGINAL string (the tokenizers library keeps alignments through NFKC), which is what lets the saliency be painted over the user's text rather than over a normalised copy. Stateless: TOKENIZER is configured once at import, nothing is toggled here.""" enc = TOKENIZER.encode(text) return { "ids": enc.ids, "offsets": enc.offsets, "special": enc.special_tokens_mask, "n_tokens": len(enc.ids), "truncated": len(enc.overflowing) > 0, # truncation leaves the cut-off part in ``overflowing`` } def full_token_count(text: str) -> int: """Untruncated token count for the truncation notice. Called only after the prediction timer has stopped and only when ``encode`` reported truncation, so a long paste never inflates the graded Prediction Time with a second tokenizer pass.""" return len(TOKENIZER_FULL.encode(text).ids) def verdict_from_probs(probs: list[float]) -> tuple[int, float]: """``(verdict_index, p_unwanted)``: Legitimate when the unwanted mass is below 0.5, otherwise the unwanted class with the most mass (phishing vs spam is decided only among the unwanted classes, so a 0.40/0.35/0.25 split still reads as unwanted).""" p_unwanted = float(sum(probs[LEGIT_INDEX + 1:])) if p_unwanted < 0.5: return LEGIT_INDEX, p_unwanted top = max(range(LEGIT_INDEX + 1, len(probs)), key=lambda i: probs[i]) return top, p_unwanted @spaces.GPU(duration=30) def _infer(text: str) -> dict[str, Any]: """Tokenize + forward + softmax (timed like the instructor's app), then one extra forward/backward for gradient x input saliency. Runs on ``DEVICE``; returns only plain Python data so it crosses the ZeroGPU process boundary cleanly.""" device = DEVICE start = timer() enc = encode(text) ids = torch.tensor([enc["ids"]], dtype=torch.int64, device=device) mask = torch.ones_like(ids, dtype=torch.bool) with torch.inference_mode(), _autocast(device): logits = MODEL(ids, mask).float() probs = [float(v) for v in torch.softmax(logits, dim=-1)[0].tolist()] # .tolist() syncs CUDA before the timer stops elapsed = timer() - start if not all(math.isfinite(p) for p in probs): # NaN/inf logits (corrupted weights, bf16 overflow) must surface as an error: ``nan < 0.5`` # is False, so verdict_from_probs would otherwise announce a confident "Phishing" at nan%. raise gr.Error("The model produced non-finite probabilities for this input; the weights or the input are corrupt.") verdict, p_unwanted = verdict_from_probs(probs) full_len = full_token_count(text) if enc["truncated"] else enc["n_tokens"] # outside the timed window t1 = timer() saliency = _gradient_x_input(ids, mask, device, unwanted=verdict != LEGIT_INDEX) sal_elapsed = timer() - t1 return { "probs": probs, "p_unwanted": p_unwanted, "verdict": verdict, "elapsed_s": elapsed, "saliency_s": sal_elapsed, "saliency": saliency, "offsets": enc["offsets"], "special": enc["special"], "n_tokens": enc["n_tokens"], "truncated": enc["truncated"], "full_len": full_len, "device": device, "dtype": str(next(MODEL.parameters()).dtype).replace("torch.", ""), } def _verdict_log_prob(logits: torch.Tensor, unwanted: bool) -> torch.Tensor: """Scalar whose gradient the saliency reads: ``log P(unwanted)`` = ``logsumexp(logits[1:]) - logsumexp(logits)`` when the verdict is unwanted (the mass of phishing + spam, so the two unwanted classes are not played against each other), else ``log P(legit)``. Both are shift-invariant in the logits, unlike a raw class logit.""" row = logits[0] lse_all = torch.logsumexp(row, dim=-1) if unwanted: return torch.logsumexp(row[LEGIT_INDEX + 1:], dim=-1) - lse_all return row[LEGIT_INDEX] - lse_all def _gradient_x_input(ids: torch.Tensor, mask: torch.Tensor, device: str, *, unwanted: bool) -> list[float]: """Per-token ``grad(log P(verdict)) . embedding`` with gradients enabled only for this call. A forward hook swaps the token-embedding output for a detached leaf that requires grad, so autograd builds a graph from that tensor onward only: no parameter gradients are accumulated and the backward is as cheap as the model allows. Positive values push toward the verdict, negative away from it. """ captured: dict[str, torch.Tensor] = {} def hook(_module, _inputs, output): leaf = output.detach().requires_grad_(True) captured["emb"] = leaf return leaf handle = MODEL.encoder.embeddings.tok.register_forward_hook(hook) try: with torch.enable_grad(), _autocast(device): logits = MODEL(ids, mask).float() target = _verdict_log_prob(logits, unwanted) (grad,) = torch.autograd.grad(target, captured["emb"]) finally: handle.remove() sal = (grad.float() * captured["emb"].float()).sum(-1)[0] return [float(v) for v in sal.tolist()] # ----------------------------------------------------------------------------- saliency -> words WORD_BREAK = set(" \t\r\n\f\v<>\"'=") # whitespace plus HTML attribute syntax; URLs stay whole def _is_break(ch: str) -> bool: return ch.isspace() or ch in WORD_BREAK def _word_span(text: str, start: int, end: int) -> tuple[int, int] | None: """Trim break characters inside a token span, then widen to the enclosing word. Words are delimited by whitespace and by ``< > " ' =`` so that ``href="http://x">text`` splits into its parts while a URL or an e-mail address is kept as one unit.""" while start < end and _is_break(text[start]): start += 1 while end > start and _is_break(text[end - 1]): end -= 1 if start >= end: return None while start > 0 and not _is_break(text[start - 1]): start -= 1 while end < len(text) and not _is_break(text[end]): end += 1 return start, end def saliency_highlights(text: str, out: dict[str, Any], k: int = TOP_K_TOKENS) -> list[tuple[str, str | None]]: """Top-k tokens by saliency in the direction of the verdict, merged into words and rendered as ``gr.HighlightedText`` segments over the original text. The saliency target is already the verdict's log-probability, so positive scores answer "what pushed the model to this answer"; tokens pulling the other way are left unpainted. Scores of tokens merged into one word are summed, so a word split into several sub-tokens is weighted as a whole. """ scored: list[tuple[float, tuple[int, int]]] = [] for s, (a, b), special in zip(out["saliency"], out["offsets"], out["special"]): if special or b <= a: continue span = _word_span(text, a, b) if span is None: continue scored.append((s, span)) scored = [t for t in scored if t[0] > 0] scored.sort(key=lambda t: -t[0]) chosen = scored[:k] if not chosen: return [(text, None)] # merge overlapping / identical word spans (sub-tokens of one word, repeated words) merged: dict[tuple[int, int], float] = {} for score, span in chosen: merged[span] = merged.get(span, 0.0) + score spans = sorted(merged.items()) flat: list[tuple[int, int, float]] = [] for (a, b), sc in spans: if flat and a <= flat[-1][1]: pa, pb, ps = flat[-1] flat[-1] = (pa, max(pb, b), ps + sc) else: flat.append((a, b, sc)) top = max(sc for _, _, sc in flat) segments: list[tuple[str, str | None]] = [] cursor = 0 for a, b, sc in flat: if a > cursor: segments.append((text[cursor:a], None)) segments.append((text[a:b], "high" if sc >= HIGH_FRACTION * top else "medium")) cursor = b if cursor < len(text): segments.append((text[cursor:], None)) return segments # ----------------------------------------------------------------------------- rule-based cues # Shared with src/evaluate.py (credential-lure slice of the holdout): keep the two copies # identical; the app directory is self-contained and cannot import from src/. CREDENTIAL_LURE_PHRASES: tuple[str, ...] = ( "verify your account", "verify your identity", "confirm your account", "confirm your identity", "confirm your password", "update your password", "reset your password", "password will expire", "password expires", "password has expired", "update your payment", "update your billing", "billing information", "unusual sign-in", "unusual activity", "suspicious activity", "account has been suspended", "account will be suspended", "account has been limited", "unlock your account", "re-activate your account", "reactivate your account", "validate your account", "login to your account", "log in to your account", "sign in to your account", "click here to verify", "security alert", "security notice", "enter your password", "re-enter your", "social security number", "credit card number", "one-time passcode", "verification code", ) URGENCY_WORDS: tuple[str, ...] = ( "urgent", "immediately", "immediate action", "within 24 hours", "within 48 hours", "final notice", "final warning", "last warning", "act now", "right away", "as soon as possible", "asap", "expires today", "will expire", "will be closed", "will be terminated", "will be deleted", "limited time", "time sensitive", "do not ignore", "failure to", "before it is too late", "today only", "respond now", ) GENERIC_GREETINGS: tuple[str, ...] = ( "dear customer", "dear valued customer", "dear user", "dear account holder", "dear member", "dear client", "dear sir/madam", "dear sir or madam", "dear friend", "dear beneficiary", "dear email user", "dear account user", "dear employee", "attention:", "hello dear", "greetings of the day", ) BRANDS: tuple[str, ...] = ( "microsoft", "office365", "office 365", "outlook", "onedrive", "sharepoint", "paypal", "apple", "icloud", "amazon", "google", "gmail", "netflix", "chase", "wellsfargo", "bankofamerica", "citibank", "hsbc", "barclays", "dhl", "fedex", "ups", "usps", "docusign", "dropbox", "adobe", "linkedin", "facebook", "instagram", "whatsapp", "irs", "hmrc", "zoom", "webex", "quickbooks", "xero", ) RISKY_ATTACHMENT_EXT = r"zip|rar|7z|htm|html|iso|img|docm|xlsm|pptm|js|vbs|scr|exe|lnk|hta|cmd|bat|msi" LEET = str.maketrans({"0": "o", "1": "l", "3": "e", "4": "a", "5": "s", "7": "t", "@": "a", "$": "s"}) KNOWN_BRAND_DOMAINS = frozenset({ "microsoftonline.com", "microsoft.com", "office.com", "live.com", "outlook.com", "windows.net", "sharepoint.com", "onedrive.com", "paypal.com", "apple.com", "icloud.com", "amazon.com", "amazonaws.com", "google.com", "gmail.com", "googleapis.com", "netflix.com", "chase.com", "wellsfargo.com", "bankofamerica.com", "dhl.com", "fedex.com", "ups.com", "usps.com", "docusign.com", "docusign.net", "dropbox.com", "adobe.com", "linkedin.com", "facebook.com", "zoom.us", "webex.com", "irs.gov", "hmrc.gov.uk", "googleusercontent.com", "gstatic.com", "dropboxusercontent.com", "paypalobjects.com", "whatsapp.net", "sharepointonline.com", "intuit.com", }) SHORTENERS = ("bit.ly", "tinyurl.com", "t.co", "goo.gl", "ow.ly", "is.gd", "buff.ly", "cutt.ly", "rb.gy") # Extensions that are also everyday technical nouns (node.js, index.html, disk.img, build.cmd): # reported only when attachment wording appears within ATTACHMENT_WINDOW characters. AMBIGUOUS_ATTACHMENT_EXT = frozenset({"js", "htm", "html", "img", "cmd"}) ATTACHMENT_WINDOW = 60 ATTACHMENT_CONTEXT_RE = re.compile( r"(?i)\b(?:attach(?:ed|ment|ments)?|files?|download(?:ed)?|open|enclosed|document|invoice|receipt)\b" ) SYNTHETIC_MARK_RE = re.compile(r"(?i)^\s*\[synthetic\]\s*") URL_RE = re.compile(r"(?i)\b(?:https?://|www\.)[^\s<>\"')\]]+") ANCHOR_RE = re.compile(r"(?is)]*?href\s*=\s*[\"']?([^\"'\s>]+)[^>]*>(.*?)") IP_URL_RE = re.compile(r"(?i)\bhttps?://(?:\d{1,3}\.){3}\d{1,3}(?::\d+)?(?:/|\b)") EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}") REPLY_REDIRECT_RE = re.compile( r"(?i)\b(?:reply|respond|write|send|forward|contact|email|e-mail|get back)\b[^.\n]{0,60}?" r"([A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,})" ) REPLY_TO_HEADER_RE = re.compile(r"(?im)^\s*reply-to\s*:\s*(\S+@\S+)") ATTACHMENT_RE = re.compile(rf"(?i)\b[\w\-.]{{1,60}}\.({RISKY_ATTACHMENT_EXT})\b") TAG_RE = re.compile(r"<[^>]+>") CC_SLD = {"co", "com", "org", "net", "ac", "gov", "edu"} def _host(url: str) -> str: u = url.strip() u = re.sub(r"(?i)^https?://", "", u) host = u.split("/", 1)[0].split("?", 1)[0].split("#", 1)[0] host = host.rsplit("@", 1)[-1].split(":", 1)[0] return host.strip(".,;:!?").lower() # URL_RE may swallow trailing sentence punctuation def _registrable(host: str) -> str: """Approximate eTLD+1 without a public-suffix list: 'a.b.example.co.uk' -> 'example.co.uk'.""" labels = host.split(".") if len(labels) >= 3 and labels[-2] in CC_SLD and len(labels[-1]) == 2: return ".".join(labels[-3:]) return ".".join(labels[-2:]) if len(labels) >= 2 else host def _plain(text: str) -> str: return html.unescape(TAG_RE.sub(" ", text)) def _brand_in(brand: str, host: str) -> bool: """Brand match on label boundaries: '.' and '-' delimit, letters and digits do not, so 'paypal.com.evil.example' and 'paypal-login.example' hit while 'firstbank.example' (irs), 'backups.example' (ups), 'zoomedphotos.example' and 'purchase-orders.example' (chase) do not. Brands of six or more characters also match inside a label ('microsoftsupport.example'), where a collision with an ordinary word is unlikely.""" if re.search(rf"(?= 6 and brand in host def _fmt_list(items: list[str], limit: int = 4) -> str: items = list(dict.fromkeys(i.replace("`", "'") for i in items)) shown = ", ".join(f"`{i}`" for i in items[:limit]) return shown + (f" (+{len(items) - limit} more)" if len(items) > limit else "") def rule_based_cues(subject: str, body: str) -> list[str]: """Heuristic cues as markdown bullets. Deliberately simple regexes: these are reading aids for the analyst, not features of the classifier, so recall matters more than precision and every cue states what it matched.""" raw = f"{subject}\n\n{body}" text = _plain(raw) low = text.lower() cues: list[str] = [] hits = [w for w in URGENCY_WORDS if w in low] if hits: cues.append(f"**Urgency / pressure language:** {_fmt_list(hits)}") hits = [p for p in CREDENTIAL_LURE_PHRASES if p in low] if hits: cues.append(f"**Credential-lure phrasing:** {_fmt_list(hits)}") hits = [g for g in GENERIC_GREETINGS if g in low] if hits: cues.append(f"**Generic greeting (no personal name):** {_fmt_list(hits)}") mismatched: list[str] = [] for href, inner in ANCHOR_RE.findall(raw): href_host = _host(href) for shown in URL_RE.findall(_plain(inner)): shown_host = _host(shown) if shown_host and href_host and _registrable(shown_host) != _registrable(href_host): mismatched.append(f"{shown_host} -> {href_host}") if mismatched: cues.append(f"**Link text points to a different host than its href:** {_fmt_list(mismatched)}") urls = URL_RE.findall(raw) + [h for h, _ in ANCHOR_RE.findall(raw)] hosts = sorted({_host(u) for u in urls if _host(u)}) ip_hits = IP_URL_RE.findall(raw) if ip_hits: cues.append(f"**URL with a raw IP address instead of a domain:** {_fmt_list(ip_hits)}") lookalike: list[str] = [] for host in hosts: reg = _registrable(host) if reg in KNOWN_BRAND_DOMAINS: continue reg_label = reg.split(".")[0] labels = host.split(".") if any(l.count("-") >= 2 for l in labels): lookalike.append(f"{host} (many hyphens)") deleet = host.translate(LEET) for brand in BRANDS: b = brand.replace(" ", "") in_host = _brand_in(b, host) if not in_host and _brand_in(b, deleet): lookalike.append(f"{host} (digits/symbols standing in for letters in '{brand}')") break if in_host and reg_label != b and not reg.startswith(b + "."): lookalike.append(f"{host} ('{brand}' appears, but the registered domain is {reg})") break if reg in SHORTENERS: lookalike.append(f"{host} (URL shortener hides the destination)") if lookalike: cues.append(f"**Suspicious / lookalike domain:** {_fmt_list(lookalike)}") words = [w for w in re.findall(r"[A-Za-z]{3,}", text)] caps = [w for w in words if w.isupper()] if len(words) >= 8 and len(caps) >= 5 and len(caps) / len(words) >= 0.30: cues.append(f"**Heavy ALL-CAPS use:** {len(caps)} of {len(words)} words ({100 * len(caps) / len(words):.0f}%)") redirects = [m.group(1) for m in REPLY_REDIRECT_RE.finditer(text)] + REPLY_TO_HEADER_RE.findall(raw) if redirects: cues.append(f"**Asks for a reply to a specific (possibly different) address:** {_fmt_list(redirects)}") no_urls = URL_RE.sub(" ", text) # a URL path like /share/open.htm is a link, not an attachment attachments: list[str] = [] for m in ATTACHMENT_RE.finditer(no_urls): if m.group(1).lower() in AMBIGUOUS_ATTACHMENT_EXT: window = no_urls[max(0, m.start() - ATTACHMENT_WINDOW): m.end() + ATTACHMENT_WINDOW] if not ATTACHMENT_CONTEXT_RE.search(window): continue attachments.append(m.group(0).strip(".")) if attachments: cues.append(f"**Mentions a risky attachment type (.zip, .htm, .iso, .docm, .exe, ...):** {_fmt_list(attachments)}") return cues # ----------------------------------------------------------------------------- the Gradio function def predict(subject: str, body: str) -> tuple[dict[str, float], float, list[tuple[str, str | None]], str]: """Gradio entry point: (class probabilities, prediction time s, saliency highlights, signals markdown). The model input is ``subject + "\\n\\n" + body`` exactly as in training (never sender or headers: TREC headers are synthetic). Prediction time covers tokenize + forward + softmax only, matching the instructor's benchmark definition; saliency and rules are timed separately and reported in the signals panel. """ # A leading "[SYNTHETIC]" demo tag is UI metadata, not email text: it is stripped so the # classifier (which never saw such a token sequence in training) scores only the email. marked = bool(SYNTHETIC_MARK_RE.match(subject or "")) subject = SYNTHETIC_MARK_RE.sub("", subject or "").strip() body = (body or "").rstrip() text = f"{subject}\n\n{body}" if not text.strip(): raise gr.Error("Paste an email subject and/or body first.") out = _infer(text) probs: list[float] = out["probs"] p_unwanted: float = out["p_unwanted"] verdict = LABELS[out["verdict"]] label_probs = {name: p for name, p in zip(LABELS, probs)} highlights = saliency_highlights(text, out) t0 = timer() cues = rule_based_cues(subject, body) rules_s = timer() - t0 class_line = ", ".join(f"{name} {100 * p:.1f}%" for name, p in zip(LABELS, probs)) lines = [ f"### Model verdict: **{verdict}**", f"**{UNWANTED_LABEL}: {100 * p_unwanted:.1f}%**", f"Per class: {class_line}.", f"The probabilities come from the {N_PARAMS:,}-parameter encoder alone (one softmax over {NUM_LABELS} classes; " f"\"{UNWANTED_LABEL}\" is the sum of the non-legitimate classes, decided at 50%). The red words are the tokens " f"whose gradient x input saliency pushed hardest toward **{verdict}** (top {TOP_K_TOKENS} tokens, merged into words).", "", "### Heuristic cues (rules, not the model)", ] if cues: lines += [f"- {c}" for c in cues] else: lines.append("- _No rule-based cue fired._ Rules are a reading aid; the model can still be right or wrong without them.") if marked: lines.append("- _Leading `[SYNTHETIC]` demo tag removed from the subject before scoring._") trunc = ( f"; **input truncated** to the first {out['n_tokens']} of {out['full_len']} tokens (the model never sees the rest)" if out["truncated"] else "" ) lines += [ "", "### Timing", f"- Prediction (tokenize + forward + softmax): **{out['elapsed_s']:.4f} s** on {out['device']} ({out['dtype']})", f"- Saliency (extra forward + backward): {out['saliency_s']:.4f} s; rules: {rules_s:.4f} s", f"- Tokens: {out['n_tokens']} of max {MAX_LEN}{trunc}", ] return label_probs, round(out["elapsed_s"], 5), highlights, "\n".join(lines) # ----------------------------------------------------------------------------- UI text and examples # The static text is derived from the LOADED export (LABELS, NUM_LABELS, N_PARAMS, MAX_LEN) so a # two-class fallback export never ships prose promising three classes next to a two-entry Label. def _params_human(n: int) -> str: return f"{n / 1e6:.0f}M" if n >= 1_000_000 else f"{n:,}" UNWANTED_NAMES: list[str] = LABELS[LEGIT_INDEX + 1:] _CLASS_LIST = ", ".join(f"**{n}**" for n in LABELS) _UNWANTED_SUM = " + ".join(f"P({n})" for n in UNWANTED_NAMES) _UNWANTED_PICK = ( f"the larger of {' / '.join(UNWANTED_NAMES)}" if len(UNWANTED_NAMES) > 1 else f"**{UNWANTED_NAMES[0]}**" ) _THREE_WAY = NUM_LABELS == 3 and UNWANTED_NAMES == DEFAULT_LABELS[3][1:] _APPROX_CHARS = f"{round(MAX_LEN * 3.5, -2):,.0f}" # ~3.5 characters per BPE token on English email if _THREE_WAY: _WHAT_IT_RETURNS = ( "the model estimates the probability that it is **legitimate** business mail, a **phishing** attempt " "(credential, payment or advance-fee lure) or bulk **spam**. The headline roll-up " f"**{UNWANTED_LABEL}** is the sum of the last two." ) _PROVENANCE_BULLET = ( '- **Spam vs phishing follows corpus provenance, not a per-message review.** "Phishing" rows come from the\n' ' Nazario phishing archive, the Nigerian-fraud (advance-fee) archive and BITS synthetic lures; "Spam" rows are\n' " the SpamAssassin and Enron-Spam spam folders. The TREC-05/07 and CEAS-08 positives never separated the two\n" ' and were trained only as "unwanted". Expect the phishing/spam boundary to be softer than legitimate/unwanted.\n' ) _SALIENCY_TARGET = "toward the combined phishing + spam mass" else: _WHAT_IT_RETURNS = ( f"the model returns one probability per class ({_CLASS_LIST}). The headline roll-up " f"**{UNWANTED_LABEL}** is {_UNWANTED_SUM}: this export does not separate phishing from spam." ) _PROVENANCE_BULLET = ( "- This export has a single unwanted class: phishing and bulk spam are not told apart.\n" ) _SALIENCY_TARGET = f"toward {_UNWANTED_SUM}" DESCRIPTION = f""" Paste an email **subject** and **body** and {_WHAT_IT_RETURNS} The classifier is a **{_params_human(N_PARAMS)}-parameter BERT-style encoder written from scratch in plain PyTorch** (no `transformers`, no pretrained weights): masked-language-model pretraining on deduplicated Enron email, then fine-tuning on labelled phishing/spam corpora. It is a **defensive** tool for triage and training, built as a course project by Build IT Systems (BITS). **Do not paste secrets, passwords, or confidential client content.** The examples below are **synthetic**, written for this demo; none is a real message, and every domain in them is a reserved `.example` name. """ ARTICLE = f""" ### How to read the result - **Prediction** is one softmax over {NUM_LABELS} classes: {_CLASS_LIST}. The Signals panel adds **{UNWANTED_LABEL}** = {_UNWANTED_SUM}, decided at 50%; the verdict is {LABELS[LEGIT_INDEX]} below that line and otherwise {_UNWANTED_PICK}. The roll-up is the number to compare with the binary runs in the project report. - **Prediction Time (s)** is tokenize + forward pass + softmax for this one email, measured with `timeit.default_timer` exactly like the course reference app. The instructor's benchmark is 0.074 s per sample. - **Tokens the model weighed most** are gradient x input saliency scores over the model's own token embeddings: red words pushed the decision toward the verdict shown (for an unwanted verdict, {_SALIENCY_TARGET}). They explain *this model*, not phishing in general. - **Signals** lists simple regex heuristics (urgency words, credential lures, host mismatches, lookalike domains, raw-IP links, ALL-CAPS, reply-to redirection, risky attachment names). They are shown for context only and do **not** feed the probability. ### Limitations {_PROVENANCE_BULLET}- Trained on 1995-2022 English email, mostly 1998-2008 (Enron, TREC, CEAS, SpamAssassin). Modern lures (MFA fatigue, QR codes, Teams/Slack phishing) and non-English mail are out of distribution. - The model reads subject and body text only: no sender, headers, SPF/DKIM, attachments or link reputation. - Inputs are truncated to {MAX_LEN} tokens (roughly {_APPROX_CHARS} characters); long emails are judged on their opening. - Length is a known confound in the training data (legitimate mail is on average longer); the error analysis in the project report checks for it. - A probability is not a verdict. Use it to prioritise human review, never to auto-delete mail. ### Links - Model card and weights: *[private model repo link placeholder]* - Training code, run log and metrics: *[project repository link placeholder]* """ PLACEHOLDER = ( "Paste the email body here. Plain text or HTML both work; links such as are scanned for " "host mismatches. Nothing you paste is stored (flagging is disabled)." ) EXAMPLES: list[list[str]] = [ [ "Action Required: Your Microsoft 365 password expires in 24 hours", "Dear user,\n\nYour Office 365 password will expire within 24 hours. To avoid losing access to your mailbox, " "verify your account immediately by signing in at:\n\n" 'https://login.microsoftonline.com\n\n' "Failure to act will result in permanent suspension of your account.\n\nIT Security Team\nOffice 365 Support", ], [ "Re: Invoice 4471 - updated banking details", "Hi Jordan,\n\nQuick one before I board. Our finance team moved banks last week, so please send the payment " "for invoice 4471 (USD 48,200) to the new account below rather than the one on the PO. Treat this as time " "sensitive; the supplier is holding the shipment until funds clear.\n\nBank: First Meridian\nAccount: " "0031-7742-19\nSWIFT: FMRDUS33\n\nI am in meetings all day so email me at ceo.office.priv@fastmail-relay.example " "if anything is unclear rather than calling.\n\nThanks,\nMichael\nChief Executive Officer", ], [ "A document has been shared with you: Q3_Payroll_Review.docm", "Hello,\n\nA colleague shared a file with you via SecureShare.\n\nFile: Q3_Payroll_Review.docm (1.2 MB)\n" "Expires: today\n\nOpen document: http://185.212.47.19/share/open.htm?id=771\n\nIf you cannot open the " "link, download the attached Q3_Payroll_Review.zip and enable content when prompted.\n\nThis is an automated " "message. Please do not reply.", ], [ "FINAL HOURS: 70% OFF luxury watches + FREE shipping!!!", "CONGRATULATIONS! You have been SELECTED for our EXCLUSIVE VIP SALE.\n\n*** 70% OFF all designer watches, " "sunglasses and handbags ***\n*** FREE worldwide shipping - TODAY ONLY ***\n*** Buy 2, get 1 FREE ***\n\n" "Thousands of satisfied customers can't be wrong! Shop now: http://bit.ly/vip-watch-deals-2026\n\n" "Don't miss out - this offer expires at midnight!\n\nTo unsubscribe, reply with REMOVE in the subject line.", ], [ "Agenda for Thursday's project sync", "Hi all,\n\nAttaching the agenda for Thursday's 10:00 sync. Main items:\n\n1. Status of the firewall " "migration (Dana)\n2. Q4 hardware refresh budget (Priya)\n3. Client onboarding checklist review\n\nIf you " "want anything added, reply by Wednesday noon. Dial-in details are unchanged from last week.\n\nThanks,\n" "Nathan", ], [ "Re: Contract renewal - redlines attached", "Hi Sarah,\n\nThanks for turning the redlines around so quickly. We accepted the changes to clauses 4.2 and " "7.1 and left a comment on the termination notice period; 60 days works for us but legal prefers 90. " "Can we close this out on the call Friday?\n\nAlso, the summary report for September is on the shared " "drive under Reports/2026-09.\n\nBest,\nMarcus", ], ] # Shown in the Examples table instead of the raw rows. The "synthetic" marker lives here, in # the UI only, so the model input is exactly the email text (a "[SYNTHETIC]" token sequence # in the subject would be out of distribution and perturb the demo predictions). EXAMPLE_LABELS: list[str] = [ "Synthetic lure: Microsoft 365 password-expiry credential phish", "Synthetic lure: CEO asks to pay an invoice to new bank details (BEC)", "Synthetic lure: shared .docm file, raw-IP link, .zip fallback", "Synthetic spam: shouting discount blast with a shortener link", "Legitimate: agenda for a project sync", "Legitimate: contract redlines follow-up", ] def build_demo() -> gr.Interface: return gr.Interface( fn=predict, inputs=[ gr.Textbox(label="Subject", lines=1, placeholder="Email subject line"), gr.Textbox(label="Email body", lines=14, placeholder=PLACEHOLDER), ], outputs=[ gr.Label(num_top_classes=NUM_LABELS, label="Prediction"), gr.Number(label="Prediction Time (s)"), gr.HighlightedText( label="Tokens the model weighed most", color_map={"high": "#d62728", "medium": "#ff9896"}, combine_adjacent=True, show_legend=False, ), gr.Markdown(label="Signals"), ], title=TITLE, description=DESCRIPTION, article=ARTICLE, examples=EXAMPLES, example_labels=EXAMPLE_LABELS, cache_examples=False, # Spaces default is True; caching would run the GPU function at build time flagging_mode="never", analytics_enabled=False, api_name="predict", ) # ----------------------------------------------------------------------------- development entry points SELFTEST_EMAILS: list[tuple[str, str]] = [EXAMPLES[0][:2], EXAMPLES[3][:2], EXAMPLES[4][:2]] # lure, spam, legitimate def _set_device(device: str) -> None: global MODEL, DEVICE DEVICE = device MODEL = place_model(MODEL, device) def selftest() -> int: devices = ["cpu"] + (["cuda"] if torch.cuda.is_available() and not ZERO_GPU else []) for device in devices: _set_device(device) LOG.info("selftest on %s", device) for subject, body in SELFTEST_EMAILS: probs, secs, highlights, details = predict(subject, body) painted = [(t, l) for t, l in highlights if l] print(f"\n=== [{device}] {subject}") print(f"probs={ {k: round(v, 4) for k, v in probs.items()} } prediction_time_s={secs}") print(f"highlighted words ({len(painted)}): {painted}") print(details) assert list(probs) == LABELS and len(probs) == NUM_LABELS, "one probability per class, in index order" assert abs(sum(probs.values()) - 1.0) < 1e-5 p_unwanted = sum(p for name, p in probs.items() if name != LABELS[LEGIT_INDEX]) assert f"{UNWANTED_LABEL}: {100 * p_unwanted:.1f}%" in details, "headline must report the unwanted roll-up" verdict_idx, _ = verdict_from_probs([probs[n] for n in LABELS]) assert f"**{LABELS[verdict_idx]}**" in details, "verdict wording must follow verdict_from_probs" expected = f"{SYNTHETIC_MARK_RE.sub('', subject).strip()}\n\n{body.rstrip()}" assert "".join(t for t, _ in highlights) == expected, "highlights must cover the model input text" # The static UI text must describe the export that is actually loaded. assert all(name in ARTICLE for name in LABELS), "ARTICLE must name every loaded class" assert f"one softmax over {NUM_LABELS} classes" in ARTICLE and f"{MAX_LEN} tokens" in ARTICLE assert UNWANTED_LABEL in DESCRIPTION and _params_human(N_PARAMS) in DESCRIPTION _set_device("cuda" if torch.cuda.is_available() else "cpu") print("\nselftest OK") return 0 def serve_for(seconds: float) -> int: import urllib.request demo = build_demo() _app, local_url, _share = demo.launch(prevent_thread_lock=True, quiet=True) print(f"serving at {local_url} for {seconds:.0f}s") status = -1 for _ in range(20): try: with urllib.request.urlopen(local_url, timeout=5) as resp: status = resp.status page = resp.read(4096).decode("utf-8", "replace") break except Exception: # server still starting time.sleep(0.5) print(f"GET {local_url} -> HTTP {status}; html head contains '