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.