Add presentation bundle (reports, curves, run log, code) for the Colab notebook
439e7b1 verified 15 days ago
45.9 kB
"""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</a>``
    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)<a\b[^>]*?href\s*=\s*[\"']?([^\"'\s>]+)[^>]*>(.*?)</a>")
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"(?<![a-z0-9]){re.escape(brand)}(?![a-z0-9])", host):
        return True
    return len(brand) >= 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 <a href=...> 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"
        '<a href="http://login.micr0soft-365-verify.example/reset?u=8f3a9c">https://login.microsoftonline.com</a>\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 '<html': {'<html' in page.lower() if status == 200 else False}")
    time.sleep(seconds)
    demo.close()
    return 0 if status == 200 else 1


if not ZERO_GPU and __name__ == "__main__":
    # Warm the kernels (CUDA context / oneDNN) so the first user click is not the slow one.
    # Skipped under ZeroGPU: CUDA work is only legal inside @spaces.GPU functions there.
    with torch.inference_mode(), _autocast(DEVICE):
        MODEL(torch.full((1, 16), 5, dtype=torch.int64, device=DEVICE), torch.ones(1, 16, dtype=torch.bool, device=DEVICE))
    LOG.info("warm-up forward done on %s", DEVICE)

if __name__ == "__main__":
    assert ARGS is not None
    if ARGS.selftest:
        raise SystemExit(selftest())
    if ARGS.serve_seconds is not None:
        raise SystemExit(serve_for(ARGS.serve_seconds))
    demo = build_demo()
    demo.launch()
else:
    demo = build_demo()