| """Email Spam Classifier: DistilBERT fine-tuned to tell spam from ham (legitimate email). |
| |
| Single source of truth for the architecture, the text cleaning and inference. It is used by |
| the Hugging Face Space (`space/app.py`), the Inference Endpoint (`handler.py`), the training |
| code (`training/src/`) and anyone who downloads this repo from the Hub: |
| |
| import model |
| predictor = model.load("path/to/this/repo", device="cpu") |
| predictor.predict("Subject: You won!\\n\\nClaim your $1000 prize now") # -> {"spam": 0.99, "ham": 0.01} |
| """ |
| from __future__ import annotations |
|
|
| import html |
| import re |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import PyTorchModelHubMixin |
| from torch import nn |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
|
|
| REPO_ID = "shalev396/email-spam-classifier" |
| FRAMEWORK = "pytorch" |
| WEIGHTS_FILE = "model.safetensors" |
| BASE_MODEL = "distilbert/distilbert-base-uncased" |
| LABELS = ["ham", "spam"] |
|
|
|
|
| def cuda_available() -> bool: |
| return torch.cuda.is_available() |
|
|
|
|
| |
| _SUBJECT_RE = re.compile(r"^\s*subject\s*:\s*", flags=re.I) |
| _HTML_TAG_RE = re.compile(r"<[^>]+>") |
| _URL_RE = re.compile(r"(?:https?://|www\.)\S+") |
| |
| _SPECIAL_CHARS_RE = re.compile(r"[^a-z0-9\s.,!?$%'-]") |
| _WHITESPACE_RE = re.compile(r"\s+") |
|
|
|
|
| def clean_text(text: str) -> str: |
| """Raw email text -> lowercase text without HTML, URLs or unusual characters. |
| |
| Applied to every training example (training/src/data_setup.py) and to every input at |
| inference, so both see exactly the same kind of text. Idempotent. |
| """ |
| if not isinstance(text, str): |
| return "" |
| text = html.unescape(text) |
| text = _HTML_TAG_RE.sub(" ", text) |
| text = _URL_RE.sub(" ", text) |
| text = text.lower() |
| text = _SPECIAL_CHARS_RE.sub(" ", text) |
| text = _WHITESPACE_RE.sub(" ", text) |
| return text.strip() |
|
|
|
|
| def compose_email(subject: str, body: str) -> str: |
| """(subject, body) -> "Subject: <subject>\\n\\n<body>", the text format /predict accepts.""" |
| return f"Subject: {(subject or '').strip()}\n\n{(body or '').strip()}" |
|
|
|
|
| def prepare_text(text: str) -> str: |
| """Model input text. Drops a leading "Subject:" header, then `clean_text`. |
| |
| The Enron training emails are "<subject> <body>" without the header word, so an input like |
| "Subject: <s>\\n\\n<body>" becomes "<s> <body>", the same shape the model was trained on. |
| """ |
| if not isinstance(text, str): |
| return "" |
| return clean_text(_SUBJECT_RE.sub("", text, count=1)) |
|
|
|
|
| |
| class SpamClassifier( |
| nn.Module, |
| PyTorchModelHubMixin, |
| tags=["ml-lab", "text-classification", "spam-detection", "distilbert"], |
| repo_url="https://github.com/shalev396/ml-lab", |
| pipeline_tag="text-classification", |
| license="mit", |
| ): |
| """DistilBERT encoder + `Dropout(dropout) -> Linear(hidden, 1)` on the [CLS] token. |
| |
| One logit per email; sigmoid(logit) = P(spam). |
| |
| Args: |
| base_model: Hub id of the pretrained encoder (only downloaded when `pretrained=True`). |
| backbone_config: the encoder's `transformers` config as a dict. Stored in config.json by |
| the mixin so the encoder can be rebuilt offline with `AutoModel.from_config`. |
| dropout: dropout in front of the linear head. |
| max_len: tokens per email (longer emails are truncated). |
| threshold: P(spam) at or above which an email counts as spam. |
| labels: class names, index 1 is the positive (spam) class. |
| pretrained: load the encoder's pretrained weights from the Hub. Only needed to *train*; |
| `load()` always passes `pretrained=False`, so inference never downloads the base model. |
| """ |
|
|
| def __init__(self, base_model: str = BASE_MODEL, backbone_config: dict | None = None, |
| dropout: float = 0.3, max_len: int = 256, threshold: float = 0.5, |
| labels: list[str] | None = None, pretrained: bool = False): |
| super().__init__() |
| self.base_model = base_model |
| self.max_len = int(max_len) |
| self.threshold = float(threshold) |
| self.labels = list(labels or LABELS) |
| if pretrained: |
| self.backbone = AutoModel.from_pretrained(base_model) |
| else: |
| config = (AutoConfig.for_model(**backbone_config) if backbone_config |
| else AutoConfig.from_pretrained(base_model)) |
| self.backbone = AutoModel.from_config(config) |
| self.dropout = nn.Dropout(dropout) |
| self.classifier = nn.Linear(self.backbone.config.hidden_size, 1) |
|
|
| def encoder_layers(self) -> nn.ModuleList: |
| """The stack of transformer blocks (DistilBERT: `transformer.layer`, BERT-likes: `encoder.layer`).""" |
| if hasattr(self.backbone, "transformer"): |
| return self.backbone.transformer.layer |
| if hasattr(self.backbone, "encoder"): |
| return self.backbone.encoder.layer |
| raise AttributeError(f"no encoder layers found on {type(self.backbone).__name__}") |
|
|
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: |
| """(N, T) token ids + mask -> (N,) spam logits.""" |
| hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state |
| return self.classifier(self.dropout(hidden[:, 0])).squeeze(-1) |
|
|
|
|
| def encode(tokenizer, texts: list[str], max_len: int, device=None) -> dict[str, torch.Tensor]: |
| """Already-prepared texts -> padded (to the longest in the batch) + truncated tensors.""" |
| batch = tokenizer(list(texts), padding=True, truncation=True, max_length=max_len, |
| return_tensors="pt") |
| batch = {"input_ids": batch["input_ids"], "attention_mask": batch["attention_mask"]} |
| return {k: v.to(device) for k, v in batch.items()} if device is not None else batch |
|
|
|
|
| |
| class Predictor: |
| """Loads the fine-tuned weights + tokenizer once and scores emails.""" |
|
|
| def __init__(self, model_dir: str | Path, device: str = "cpu"): |
| model_dir = Path(model_dir) |
| if not (model_dir / WEIGHTS_FILE).is_file(): |
| raise FileNotFoundError( |
| f"{WEIGHTS_FILE} not found in {model_dir}: train/export the model first " |
| f"or download it with huggingface_hub.snapshot_download('{REPO_ID}')." |
| ) |
| if not (model_dir / "tokenizer.json").is_file(): |
| raise FileNotFoundError(f"tokenizer.json not found in {model_dir}: export the tokenizer with the weights.") |
| self.device = torch.device(device) |
| self.model = SpamClassifier.from_pretrained( |
| model_dir, pretrained=False, map_location="cpu", strict=True |
| ) |
| self.model.to(self.device).eval() |
| self.tokenizer = AutoTokenizer.from_pretrained(model_dir) |
| self.max_len = self.model.max_len |
| self.threshold = self.model.threshold |
|
|
| @torch.inference_mode() |
| def predict_proba(self, texts: list[str], batch_size: int = 32) -> list[float]: |
| """Raw email texts -> P(spam) for each. Batches are grouped by length for speed.""" |
| prepared = [prepare_text(t) for t in texts] |
| order = sorted(range(len(prepared)), key=lambda i: len(prepared[i])) |
| probs = [0.0] * len(prepared) |
| for start in range(0, len(order), batch_size): |
| idx = order[start:start + batch_size] |
| batch = encode(self.tokenizer, [prepared[i] for i in idx], self.max_len, self.device) |
| p = torch.sigmoid(self.model(**batch).float()).cpu().tolist() |
| for i, value in zip(idx, p): |
| probs[i] = float(value) |
| return probs |
|
|
| def predict(self, text: str) -> dict[str, float]: |
| """Email text (ideally "Subject: <subject>\\n\\n<body>") -> {"spam": p, "ham": 1 - p}.""" |
| p = self.predict_proba([text])[0] |
| return {"spam": p, "ham": 1.0 - p} |
|
|
| def label(self, text: str) -> str: |
| """"spam" if P(spam) >= threshold else "ham".""" |
| return "spam" if self.predict(text)["spam"] >= self.threshold else "ham" |
|
|
|
|
| def load(model_dir: str | Path, device: str = "cpu") -> Predictor: |
| return Predictor(model_dir, device) |
|
|
