File size: 5,431 Bytes
439e7b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | ---
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.
|