BuildITSystems's picture
Add presentation bundle (reports, curves, run log, code) for the Colab notebook
439e7b1 verified
|
Raw History Blame Contribute Delete
5.43 kB
metadata
title: BITS Phishing Email Detector
emoji: 🛡️
colorFrom: blue
colorTo: indigo
sdk: gradio
sdk_version: 6.29.1
python_version: 3.12.12
app_file: app.py
pinned: false
license: mit
short_description: 'From-scratch 337M encoder: legitimate / phishing / spam'

BITS Phishing Email Detector

Paste an email subject and body; the app returns the model's probability for each of three classes, Legitimate, Phishing and Spam, plus the roll-up Unwanted (phishing or spam) = P(Phishing) + P(Spam), the prediction time for that single sample, the tokens the model weighed most, and a list of rule-based cues for a human reader. The verdict is Legitimate when the unwanted roll-up is below 50% and otherwise the larger of Phishing / Spam. It is a defensive triage and training aid.

How it works

  • Model. A 337M-parameter BERT-style pre-LayerNorm encoder (24 layers, d=1024, 16 heads, FFN 4096, vocabulary 32,768, 512 tokens) written from scratch in plain PyTorch (model.py; no transformers, no pretrained weights) with a three-way classification head. It was pretrained with masked language modelling on deduplicated Enron email and fine-tuned on public corpora: legitimate mail and spam from TREC-05/07, CEAS-08, Enron-Spam and SpamAssassin; phishing from the Nazario phishing archive and the Nigerian-fraud (advance-fee) archive (Champa et al. corpus mirror, CC-BY-4.0) plus BITS synthetic lures. Rows whose corpus never separated spam from phishing (TREC-05/07, CEAS-08 positives) were trained with a partial label on the combined phishing + spam mass, so they shape the legitimate-vs-unwanted decision without asserting which unwanted class they belong to. The TREC-06 corpus was held out entirely as the binary test set; a separate three-way holdout of rows the tokenizer never saw and that share no exact or near-duplicate text with the pretraining corpus (checked by content keys and shared 50-character shingles) gives the 3x3 confusion matrix.
  • Input. subject + "\n\n" + body, byte-level BPE tokenised (NFKC, case preserved) and truncated to 512 tokens. Sender, headers and attachments are never read.
  • Prediction time. timeit.default_timer around tokenize + forward + softmax, like the course reference app (benchmark: 0.074 s/sample).
  • Token importance. Gradient x input saliency of the verdict's log-probability (log P(Legitimate), or log(P(Phishing) + P(Spam)) when the verdict is unwanted) with respect to the token embeddings; the top tokens are merged back into words and highlighted over the original text.
  • Signals. Regex heuristics (urgency words, credential lures, link-text vs href host mismatch, lookalike domains, raw-IP URLs, ALL-CAPS share, reply-to redirection, generic greetings, risky attachment names). They are shown for context only; the probability comes from the model alone.
  • Runtime. ZeroGPU (@spaces.GPU(duration=30), bf16) with the model placed on cuda at module level; falls back to CPU fp32. The weights (model.safetensors, config.json, tokenizer.json) are bundled in this Space next to app.py; app.py can also download them from the model repo named by the PHISH_MODEL_REPO variable with an HF_TOKEN secret.
  • Training recipe of the deployed checkpoint. Fine-tune run ft-focal: focal loss (gamma 2.0) over the partial-label negative log-likelihood, AdamW, LR 3e-5 with warmup and linear decay, batch 32, early stopping on the validation F1 of the unwanted roll-up. It was chosen among six runs (cross-entropy, class-weighted cross-entropy, focal, synthetic supplement, 12 layers, scratch init) by validation metrics alone.

Model card

Weights, config, tokenizer, metrics and the full model card: https://huggingface.co/BuildITSystems/phishing-detector-encoder

Limitations

  • Training data is English email from 1995-2022, the bulk of it 1998-2008 (Enron, TREC, CEAS, SpamAssassin); modern lures (MFA fatigue, QR codes, chat-app phishing) and non-English mail are out of distribution.
  • Spam vs phishing follows corpus provenance, not a per-message review. "Phishing" means a row came from the Nazario or Nigerian-fraud archives or the synthetic lures; "Spam" means the SpamAssassin or Enron-Spam spam folders. The TREC-05/07 and CEAS-08 positives (the largest share of unwanted training mail) were never split and were trained only as "unwanted", so the phishing/spam boundary is softer than the legitimate/unwanted one. The roll-up Unwanted (phishing or spam) is the number to trust first and the one comparable with the earlier binary runs.
  • Text only: no sender reputation, SPF/DKIM, link reputation or attachment analysis.
  • Inputs longer than 512 tokens (about 1,800 characters) are judged on their opening.
  • Legitimate mail is on average longer than spam in the training data; length is a known confound.
  • Flagging is disabled and nothing pasted is stored, but do not paste secrets or confidential client content.

Context

Final project for the Skillspire AI & Machine Learning course, built by Nathan King at Build IT Systems (BITS). The project repository holds the data pipeline, tokenizer, pretraining and fine-tuning code, the run log of every experiment (layers, epochs, loss functions, crashes and fixes) and the evaluation on the unseen TREC-06 test set; the headline numbers are in the model card linked above.