| ---
|
| language: en
|
| license: mit
|
| pipeline_tag: text-classification
|
| library_name: transformers
|
| tags:
|
| - spam
|
| - ham
|
| - email
|
| - tinybert
|
| - enron
|
| - text-classification
|
| model-index:
|
| - name: prancyFox/tiny-bert-enron-spam
|
| results:
|
| - task:
|
| type: text-classification
|
| name: Spam / Ham Classification
|
| dataset:
|
| name: Enron (processed CSV)
|
| type: enron_email
|
| split: test
|
| metrics:
|
| - name: F1 (macro)
|
| type: f1
|
| value: 0.7666
|
| - name: ROC-AUC
|
| type: roc_auc
|
| value: 0.9977
|
| - name: Precision (spam)
|
| type: precision
|
| value: 0.9954
|
| - name: Recall (spam)
|
| type: recall
|
| value: 0.5632
|
| - name: Precision (ham)
|
| type: precision
|
| value: 0.6875
|
| - name: Recall (ham)
|
| type: recall
|
| value: 0.9973
|
| base_model: huawei-noah/TinyBERT_General_4L_312D
|
| ---
|
|
|
| # TinyBERT Spam Classifier (Enron)
|
|
|
| A compact **TinyBERT (4-layer, 312 hidden)** model fine-tuned to classify **email text** as **spam** or **ham**.
|
| Trained on an Enron-derived CSV with light email-specific cleaning (e.g., removing quoted lines and base64-like blobs).
|
| Optimized for **low false positives** by default; adjust the decision threshold if you want higher spam recall.
|
|
|
| > Labels: `ham` (0) and `spam` (1)
|
|
|
| ---
|
|
|
| ## ✨ Quick Start
|
|
|
| ```python
|
| from transformers import pipeline
|
|
|
| clf = pipeline(
|
| "text-classification",
|
| model="prancyFox/tiny-bert-enron-spam",
|
| truncation=True # recommended for long emails
|
| )
|
|
|
| clf("Congratulations! You won a FREE iPhone. Click here now!")
|
| # [{'label': 'spam', 'score': 0.98}]
|
| ````
|
|
|
| **Batch inference**
|
|
|
| ```python
|
| texts = [
|
| "Meeting moved to 3pm, see agenda attached.",
|
| "FREE gift card!!! Act now!",
|
| ]
|
| preds = clf(texts, truncation=True)
|
| ```
|
|
|
| ---
|
|
|
| ## 🔎 Intended Use & Limitations
|
|
|
| **Intended use**
|
|
|
| * Classifying **email bodies (and optionally subject+body)** as spam vs ham.
|
| * Low-latency scenarios where a small model is preferred.
|
|
|
| **Out of scope / Limitations**
|
|
|
| * Non-English email content may reduce accuracy.
|
| * Long threads with heavy quoting/footers can dilute signal (use truncation + cleaning).
|
| * Trained on Enron-style corporate emails; consumer emails may differ (consider further fine-tuning).
|
|
|
| ---
|
|
|
| ## 🧰 How We Preprocessed the Data
|
|
|
| Light normalization aimed at keeping semantic content:
|
|
|
| * Remove long base64-like blobs.
|
| * Drop quoted lines starting with `>` or `|`.
|
| * Optional: concatenate `Subject + "\n" + Message` when available.
|
| * Collapse repeated whitespace.
|
|
|
| (You can replicate similar cleaning in your serving pipeline for alignment.)
|
|
|
| ---
|
|
|
| ## 🏋️ Training Details
|
|
|
| * **Base model:** `huawei-noah/TinyBERT_General_4L_312D`
|
| * **Task:** Binary text classification (`ham`=0, `spam`=1)
|
| * **Tokenizer:** fast BERT tokenizer (uncased)
|
| * **Max length:** 256 tokens
|
| * **Optimizer / LR:** AdamW, learning rate `2e-5 – 5e-5` (final run `3e-5`)
|
| * **Batch size:** 32
|
| * **Epochs:** 4 (early stopping enabled)
|
| * **Warmup:** 10%
|
| * **Weight decay:** 0.01
|
| * **Loss:** Cross-entropy with class weighting (ham/spam balanced from label distribution). Focal loss available in the trainer.
|
| * **Early stopping metric:** `eval_f1`
|
| * **Best checkpoint:** Saved using evaluation on validation set.
|
|
|
| > Trainer script: `train/train_tinybert.py` (TinyBERT-compatible, with legacy HF support shims).
|
|
|
| ---
|
|
|
| ## 📊 Evaluation (Chunked Benchmark Summary)
|
|
|
| Metrics below reflect a **chunked evaluation** pass (used for long emails), where the model sees up to 512 tokens per chunk with overlap. Threshold tuned to minimize false positives:
|
|
|
| ### Classification Report
|
|
|
| | Class | Precision | Recall | F1 |
|
| | ------------: | ---------: | ---------: | ---------: |
|
| | ham | 0.6875 | 0.9973 | 0.8139 |
|
| | spam | 0.9954 | 0.5632 | 0.7194 |
|
| | **macro avg** | **0.8414** | **0.7802** | **0.7666** |
|
|
|
| * **ROC-AUC:** 0.9977
|
|
|
| **Confusion matrix**
|
|
|
| ```
|
| [[16500 45]
|
| [ 7500 9671]]
|
| ```
|
|
|
| **Interpretation:** The model is conservative (very few false positives on ham). If you need to catch more spam, **lower the decision threshold** (e.g., from 0.5 → 0.35) or re-train with a spam-skewed class weight / focal loss.
|
|
|
| ---
|
|
|
| ## 🎛️ Threshold & Long-Email Guidance
|
|
|
| * **Threshold:** Default is 0.5. For higher spam recall, try **0.35–0.45** and evaluate impact on false positives.
|
| * **Long emails:** For multi-paragraph threads, consider **chunking** and aggregating chunk-level spam scores (e.g., max or average). Our reference app uses 512-token chunks with overlap.
|
|
|
| ---
|
|
|
| ## 🧪 Reproducibility
|
|
|
| **Environment**
|
|
|
| * Python 3.10/3.11
|
| * `transformers >= 4.40`
|
| * `datasets >= 2.20`
|
| * `evaluate >= 0.4.2`
|
| * `torch >= 2.1`
|
|
|
| **Training command (example)**
|
|
|
| ```bash
|
| python train/train_tinybert.py \
|
| --train data/enron.csv \
|
| --text_col Message --label_col "Spam/Ham" \
|
| --output_dir outputs/tiny-bert-enron-spam \
|
| --epochs 4 --batch_size 32 --lr 3e-5 \
|
| --max_length 256 --fp16
|
| ```
|
|
|
| **Serving (FastAPI example)**
|
|
|
| ```bash
|
| python spam_bert.py --serve \
|
| --model prancyFox/tiny-bert-enron-spam \
|
| --model-cache-dir ./models_cache
|
| ```
|
|
|
| ---
|
|
|
| ## 📁 Files
|
|
|
| This repo should include:
|
|
|
| * `config.json`
|
| * `pytorch_model.bin` or `model.safetensors`
|
| * `tokenizer.json` and `tokenizer_config.json` (or `vocab.txt` etc.)
|
| * `README.md` (this file)
|
| * (Optional) `label_mapping.json` with `{"ham": 0, "spam": 1}`
|
|
|
| ---
|
|
|
| ## ⚖️ License
|
|
|
| * **Model weights & code**: MIT
|
| * **Dataset**: Check the original Enron dataset/license terms before redistribution.
|
|
|
| ---
|
|
|
| ## 🔬 Ethical Considerations & Risks
|
|
|
| * False positives can have operational cost (missed legitimate emails). This model is tuned to minimize them; if you change the threshold, validate carefully.
|
| * Spam evolves. Periodically re-train with fresh samples to maintain accuracy.
|
| * Non-English or code-mixed content may degrade performance.
|
|
|
| ---
|
|
|
| ## 🧩 Citation
|
|
|
| If you use this model, please cite:
|
|
|
| ```
|
| @software{tinybert_enron_spam_2025,
|
| title = {TinyBERT Spam Classifier (Enron)},
|
| author = {Ing. Daniel Eder},
|
| year = {2025},
|
| url = {https://huggingface.co/prancyFox/tiny-bert-enron-spam}
|
| }
|
| ```
|
|
|
| And the TinyBERT paper:
|
|
|
| ```
|
| @article{jiao2020tinybert,
|
| title={TinyBERT: Distilling BERT for Natural Language Understanding},
|
| author={Jiao, Xiaoqi and Yin, Yichun and others},
|
| journal={Findings of EMNLP},
|
| year={2020}
|
| }
|
| ```
|
|
|
| ---
|
|
|
| ## 🛠 Maintainers
|
|
|
| * **Daniel Eder** ([daniel@deder.at](mailto:daniel@deder.at?subject=tiny-bert-enron-spam))
|
|
|
| ---
|
|
|
| ### Notes
|
|
|
| * For a higher-recall variant, fine-tune with `--use_focal_loss` or increase the spam class weight, then re-evaluate thresholds.
|
| * If you want a **PyTorch Lightning** or **Accelerate** training variant, \~it’s easy to adapt the provided trainer. |
