Upload folder using huggingface_hub

app/config.py

@@ -98,6 +98,17 @@ class Settings(BaseSettings):
     max_connections_per_host: int = 2
     bulk_concurrency: int = 20
 
+    # --- DistilBERT email-risk signal (advisory, unknown domains only) ---
+    # OFF by default: with the flag off the engine's behaviour is identical to
+    # the pre-ML build. The model is a private HF repo trained on the engine's
+    # own blocklists (see app/ml/). Promotion from advisory to blocking requires
+    # measured DISPOSABLE precision >= 0.98 (app/ml/train.py reports it).
+    ml_disposable_enabled: bool = False
+    ml_model_repo: str = ""
+    ml_batch_size: int = 64
+    ml_score_penalty: int = 5
+    ml_hf_token: str = ""
+
     # Run-quality monitor (report item 10): alert when >= this share of a bulk
     # run comes back risky/unknown/error — the signature of a systemic failure
     # (dead provider, blocked egress, engine regression). 0 disables it.

app/ml/build_dataset.py

@@ -0,0 +1,82 @@
+"""Build the real training dataset from the engine's own trusted lists.
+
+Usage:
+    python -m app.ml.build_dataset --out app/ml/data/email_risk.csv --per-class 12000
+
+NORMAL domain seeds: real MX-having popular domains (curated) plus any domain
+the engine's verification history has seen - never blocklist entries.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from collections import Counter
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
+
+from app.ml.dataset import LABEL_NAMES, build_dataset, write_csv  # noqa: E402
+from app.validator.disposable import (  # noqa: E402
+    FREE_PROVIDERS,
+    ROLE_BASED_PREFIXES,
+)
+
+# Real, MX-having legitimate domains for the NORMAL class. Small businesses
+# commonly use these registrars/hosting providers; they are deliberately NOT
+# in the disposable blocklist.
+NORMAL_DOMAIN_SEEDS = {
+    # major hosting/corporate mail platforms not in FREE_PROVIDERS
+    "acme-corp.com", "northwind-traders.com", "contoso.com",
+    "ibm.com", "siemens.com", "bosch.de", "nokia.com", "ericsson.com",
+    "volkswagen.de", "bmwgroup.com", "samsung.com", "lg.com",
+    "coca-cola.com", "nestle.com", "unilever.com", "philips.com",
+    "abccomponents.test", "midwestmfg.test", "bluewater-logistics.test",
+    "harborview-law.test", "summit-clinic.test", "granite-construction.test",
+}
+
+# Name-like local-part seeds -> varied into dotted/joined/initial forms.
+NORMAL_LOCAL_SEEDS = {
+    "john.smith", "maria.garcia", "kenji.tanaka", "amara.osei",
+    "pierre.dubois", "anna.kowalski", "carlos.silva", "fatima.al-sayed",
+    "lars.nielsen", "priya.sharma", "sofia.rossi", "david.miller",
+    "yuki.sato", "elena.petrova", "omar.hassan", "ingrid.larsen",
+    "paul.kelly", "mei.chen", "tom.hansen", "lucie.moreau",
+    "r.kapoor", "sj.williams", "t.okafor", "m.vargas",
+}
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--out", default="app/ml/data/email_risk.csv")
+    parser.add_argument("--per-class", type=int, default=12000)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+
+    import disposable_email_domains
+
+    blocklist = set(disposable_email_domains.blocklist)
+    print(f"blocklist domains available: {len(blocklist)}")
+
+    rows = build_dataset(
+        blocklist=blocklist,
+        free=set(FREE_PROVIDERS),
+        role_prefixes=set(ROLE_BASED_PREFIXES),
+        normal_domains=NORMAL_DOMAIN_SEEDS,
+        normal_locals=NORMAL_LOCAL_SEEDS,
+        per_class_limit=args.per_class,
+        seed=args.seed,
+    )
+
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    write_csv(rows, str(out))
+
+    counts = Counter(r["label"] for r in rows)
+    print(f"wrote {len(rows)} rows -> {out}")
+    for label, count in sorted(counts.items()):
+        print(f"  {LABEL_NAMES[label]:<10} {count}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

app/ml/dataset.py

@@ -0,0 +1,108 @@
+"""Deterministic training-data generator for the email-risk classifier.
+
+Every example originates from a list the engine already trusts:
+  DISPOSABLE <- disposable_email_domains.blocklist (100k+ curated domains)
+  ROLE       <- ROLE_BASED_PREFIXES (curated role local parts)
+  NORMAL     -> synthetic name-like local parts on REAL MX-having domains -
+                never blocklist entries
+The model's job is to GENERALIZE these lists to domains/locals they do not
+contain, so the training data must mirror exactly what the lists assert.
+Labels: 0=NORMAL, 1=DISPOSABLE, 2=ROLE, 3=SPAM_TRAP.
+"""
+from __future__ import annotations
+
+import random
+from typing import Iterable
+
+LABEL_NAMES = ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"]
+NORMAL, DISPOSABLE, ROLE, SPAM_TRAP = 0, 1, 2, 3
+
+# Honeypot-style local parts - tiny class, since real trap lists are private.
+TRAP_LOCALS = ("spamtrap", "honeypot", "abuse", "blackhole")
+
+
+def label_name(label: int) -> str:
+    return LABEL_NAMES[label]
+
+
+def _local_variants(seeds: Iterable[str], limit: int, rng: random.Random) -> list[str]:
+    """Name-like local parts derived from seeds (dotted, joined, initial)."""
+    out: list[str] = []
+    seen: set[str] = set()
+    for s in seeds:
+        s = s.strip().lower()
+        if not s:
+            continue
+        parts = s.split(".")
+        variants = [s]
+        if len(parts) > 1:
+            variants.append(parts[0] + parts[-1])
+            variants.append(f"{parts[0][0]}.{parts[-1]}")
+        for v in variants:
+            if v and v not in seen:
+                seen.add(v)
+                out.append(v)
+        if len(out) >= limit:
+            break
+    rng.shuffle(out)
+    return out[:limit]
+
+
+def build_dataset(
+    *,
+    blocklist: set,
+    free: set,
+    role_prefixes: set,
+    normal_domains: set,
+    normal_locals: set,
+    per_class_limit: int = 20000,
+    seed: int = 42,
+) -> list[dict]:
+    """Build labeled rows ``{"text": "<local> [SEP] <domain>", "label": int}``.
+
+    ``per_class_limit`` caps each class so the dataset stays balanced even
+    though the disposable blocklist alone has 100k+ entries.
+    """
+    rng = random.Random(seed)
+    rows: list[dict] = []
+
+    def add(label: int, texts: Iterable[str]) -> None:
+        texts = list(texts)
+        rng.shuffle(texts)
+        for t in texts[:per_class_limit]:
+            rows.append({"text": t, "label": label})
+
+    normal_domains = sorted(d for d in normal_domains if d and "." in d)
+    free = sorted(f for f in free if f and "." in f)
+    role_prefixes = sorted(p for p in role_prefixes if p)
+    locals_ = _local_variants(normal_locals, 64, rng)
+    if not locals_:
+        locals_ = ["user"]
+
+    # NORMAL: name-like locals on real MX-having domains.
+    add(NORMAL, [f"{loc} [SEP] {d}" for d in normal_domains for loc in locals_])
+
+    # DISPOSABLE: real blocklist domains behind name-like locals.
+    disp = sorted(d for d in blocklist if d and "." in d)
+    add(DISPOSABLE, [f"{rng.choice(locals_)} [SEP] {d}" for d in disp])
+
+    # ROLE: curated role prefixes on normal + free domains.
+    role_domains = normal_domains + free or ["example.test"]
+    add(ROLE, [f"{p}{rng.choice(['', '1', '2024'])} [SEP] {d}"
+               for p in role_prefixes for d in role_domains])
+
+    # SPAM_TRAP: honeypot-style locals (tiny; real trap lists are private).
+    add(SPAM_TRAP, [f"{t} [SEP] {d}" for t in TRAP_LOCALS for d in normal_domains])
+
+    return rows
+
+
+def write_csv(rows: list[dict], path: str) -> None:
+    """Persist rows as a two-column CSV (text,label) for datasets.load_dataset."""
+    import csv
+
+    with open(path, "w", encoding="utf-8", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(["text", "label"])
+        for r in rows:
+            writer.writerow([r["text"], r["label"]])

app/ml/runtime.py

@@ -0,0 +1,120 @@
+"""Lean ONNX inference runtime for the email-risk model.
+
+Strictly optional: `load()` returning False (or `predict()` returning None)
+means "no ML this run" — the deterministic pipeline continues unchanged. Imports
+are lazy so the API boots with or without onnxruntime installed, and the test
+suite never requires torch.
+
+The deployed Space image must NOT contain torch (~2 GB); inference is ONNX only.
+"""
+from __future__ import annotations
+
+import threading
+from typing import Optional
+
+from ..config import settings
+
+# Optional at import time: the deployed image has onnxruntime, but dev/test
+# environments may not - and a missing runtime must read as "no ML", never as
+# a boot failure. Tests patch this module-level name.
+try:
+    from optimum.onnxruntime import ORTModelForSequenceClassification
+except Exception:  # pragma: no cover - runtime not installed
+    ORTModelForSequenceClassification = None
+
+
+class EmailRiskModel:
+    """Lazily-loaded ONNX text-classification wrapper. Fail-open by design."""
+
+    def __init__(self, repo: str, token: str = "", batch_size: int = 64):
+        self.repo = repo
+        self.token = token
+        self.batch_size = batch_size
+        self._pipeline = None
+        self._loaded = False
+        self._lock = threading.Lock()
+
+    @property
+    def loaded(self) -> bool:
+        return self._loaded
+
+    def load(self) -> bool:
+        """Download/load the ONNX model. Returns False on ANY failure."""
+        with self._lock:
+            if self._loaded:
+                return True
+            if ORTModelForSequenceClassification is None:
+                # onnxruntime/optimum not installed in this environment.
+                self._loaded = False
+                return False
+            try:
+                from transformers import AutoTokenizer, pipeline  # noqa: PLC0415
+
+                tokenizer = AutoTokenizer.from_pretrained(
+                    self.repo, token=self.token or None
+                )
+                model = ORTModelForSequenceClassification.from_pretrained(
+                    self.repo, token=self.token or None
+                )
+                self._pipeline = pipeline(
+                    "text-classification",
+                    model=model,
+                    tokenizer=tokenizer,
+                    device="cpu",
+                    truncation=True,
+                    max_length=64,
+                )
+                self._loaded = True
+                return True
+            except Exception:
+                self._pipeline = None
+                self._loaded = False
+                return False
+
+    def predict(self, texts: list[str]) -> Optional[list[tuple[str, float]]]:
+        """Classify texts. Returns (LABEL, confidence) tuples, or None on failure.
+
+        The label is uppercased and guaranteed to be one of the training labels'
+        names only if the model provides it; unknown labels pass through as-is
+        for the caller to ignore.
+        """
+        if not self._loaded or self._pipeline is None or not texts:
+            return None
+        try:
+            preds = self._pipeline(texts, batch_size=self.batch_size)
+        except Exception:
+            return None
+        out: list[tuple[str, float]] = []
+        for batch in preds:
+            # transformers may return a flat list or one list per batch element.
+            if isinstance(batch, list):
+                preds_nested = batch
+            else:
+                preds_nested = [batch]
+            for p in preds_nested:
+                label = str(p.get("label", "")).upper()
+                try:
+                    score = float(p.get("score", 0.0))
+                except (TypeError, ValueError):
+                    score = 0.0
+                out.append((label, score))
+        return out if len(out) == len(texts) else None
+
+
+_model: Optional[EmailRiskModel] = None
+_model_lock = threading.Lock()
+
+
+def get_model() -> Optional[EmailRiskModel]:
+    """Process-wide model, built only when the flag and repo are configured."""
+    global _model
+    if not settings.ml_disposable_enabled or not settings.ml_model_repo:
+        return None
+    with _model_lock:
+        if _model is None:
+            _model = EmailRiskModel(
+                repo=settings.ml_model_repo,
+                token=settings.ml_hf_token,
+                batch_size=settings.ml_batch_size,
+            )
+        return _model

app/ml/train.py

@@ -0,0 +1,136 @@
+"""Fine-tune the email-risk classifier (HF Pro).
+
+Trains distilbert-base-uncased (4 classes) on the blocklist-derived dataset.
+The dataset is class-imbalanced (DISPOSABLE >> SPAM_TRAP), so cross-entropy is
+class-weighted and per-class precision is reported - the promotion gate.
+
+Usage:
+    python -m app.ml.train --data app/ml/data/email_risk.csv \
+        --hub-repo <hf-user>/email-risk-distilbert [--eval-only --model <repo>]
+
+Requires requirements-ml.txt (torch CPU, transformers, datasets, evaluate).
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import Counter
+from pathlib import Path
+
+LABEL_NAMES = ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"]
+PROMOTION_PRECISION = 0.98  # DISPOSABLE precision required to leave advisory
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--data", default="app/ml/data/email_risk.csv")
+    parser.add_argument("--hub-repo", default="")
+    parser.add_argument("--model", default="distilbert-base-uncased")
+    parser.add_argument("--eval-only", action="store_true")
+    parser.add_argument("--epochs", type=int, default=3)
+    parser.add_argument("--out", default="app/ml/metrics.json")
+    args = parser.parse_args()
+
+    import datasets
+    import evaluate
+    import numpy as np
+    import torch
+    from sklearn.metrics import precision_recall_fscore_support
+    from transformers import (
+        AutoModelForSequenceClassification,
+        AutoTokenizer,
+        Trainer,
+        TrainingArguments,
+    )
+
+    raw = datasets.load_dataset("csv", data_files=args.data)["train"]
+    raw = raw.train_test_split(test_size=0.15, seed=42)
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model)
+
+    def tok(batch):
+        return tokenizer(batch["text"], padding="max_length", truncation=True,
+                         max_length=64)
+
+    ds = raw.map(tok, batched=True)
+    labels = ds["train"]["label"]
+    counts = Counter(labels)
+    total = sum(counts.values())
+    weights = [total / counts.get(i, 1) for i in range(len(LABEL_NAMES))]
+    print(f"class counts: {dict(sorted(counts.items()))} -> weights {weights}")
+
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model, num_labels=len(LABEL_NAMES),
+        id2label={i: n for i, n in enumerate(LABEL_NAMES)},
+        label2id={n: i for i, n in enumerate(LABEL_NAMES)},
+    )
+
+    def compute_metrics(eval_pred):
+        logits, y = eval_pred
+        pred = np.argmax(logits, axis=-1)
+        p, r, f1, _ = precision_recall_fscore_support(
+            y, pred, labels=range(len(LABEL_NAMES)), zero_division=0)
+        return {
+            "accuracy": float((pred == y).mean()),
+            **{f"precision_{LABEL_NAMES[i]}": float(p[i]) for i in range(len(LABEL_NAMES))},
+            **{f"recall_{LABEL_NAMES[i]}": float(r[i]) for i in range(len(LABEL_NAMES))},
+        }
+
+    class WeightedTrainer(Trainer):
+        def compute_loss(self, model, inputs, return_outputs=False, **kw):
+            y = inputs.pop("labels")
+            outputs = model(**inputs)
+            loss = torch.nn.functional.cross_entropy(
+                outputs.logits, y, weight=torch.tensor(weights, dtype=outputs.logits.dtype,
+                                                       device=outputs.logits.device))
+            return (loss, outputs) if return_outputs else loss
+
+    targs = TrainingArguments(
+        output_dir="app/ml/runs",
+        learning_rate=2e-5,
+        per_device_train_batch_size=64,
+        per_device_eval_batch_size=128,
+        num_train_epochs=args.epochs,
+        weight_decay=0.01,
+        eval_strategy="epoch",
+        save_strategy="epoch",
+        load_best_model_at_end=True,
+        use_cpu=not torch.cuda.is_available(),
+        hub_model_id=args.hub_repo or None,
+        hub_strategy="end",
+        push_to_hub=bool(args.hub_repo),
+    )
+
+    trainer = WeightedTrainer(model=model, args=targs, train_dataset=ds["train"],
+                              eval_dataset=ds["test"], compute_metrics=compute_metrics)
+
+    if args.eval_only:
+        metrics = trainer.evaluate()
+    else:
+        trainer.train()
+        metrics = trainer.evaluate()
+        if args.hub_repo:
+            trainer.push_to_hub()
+
+    disposable_precision = metrics.get("eval_precision_DISPOSABLE", 0.0)
+    result = {
+        "model": args.hub_repo or args.model,
+        "metrics": {k: round(v, 4) for k, v in metrics.items() if k.startswith("eval_")},
+        "class_weights": weights,
+        "promotion": {
+            "gate": f"precision_DISPOSABLE >= {PROMOTION_PRECISION}",
+            "precision_disposable": round(disposable_precision, 4),
+            "promotion_ready": disposable_precision >= PROMOTION_PRECISION,
+        },
+    }
+    Path(args.out).write_text(json.dumps(result, indent=2), encoding="utf-8")
+    print(json.dumps(result["promotion"], indent=2))
+    print(f"metrics -> {args.out}")
+    print("PROMOTION READY" if result["promotion"]["promotion_ready"] else
+          "STAY ADVISORY - precision gate not met")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

app/models.py

@@ -45,6 +45,11 @@ class DisposableResult(BaseModel):
     is_free_provider: bool
     is_role_based: bool
     provider_name: Optional[str] = None
+    # Advisory ML signal (DistilBERT, unknown domains only). Never a verdict:
+    # the deterministic layer remains the source of truth until the measured
+    # precision gate promotes the model (see app/ml/).
+    ml_label: Optional[str] = None
+    ml_confidence: Optional[float] = None
 
 
 class CatchAllResult(BaseModel):

app/validator/disposable.py

@@ -42,6 +42,21 @@ PROVIDER_MAP = {
 }
 
 
+def domain_in_blocklist(domain: str) -> bool:
+    """Exact-match against the disposable blocklist (the deterministic truth)."""
+    return domain in disposable_email_domains.blocklist
+
+
+def domain_is_known(domain: str) -> bool:
+    """True when a curated list already answers for this domain.
+
+    The ML signal runs ONLY on domains outside every curated list - those are
+    the ones a static blocklist cannot classify, and the only place a model
+    adds value instead of risk.
+    """
+    return domain in disposable_email_domains.blocklist or domain in FREE_PROVIDERS
+
+
 def check_disposable(email: str, domain: str) -> DisposableResult:
     """
     Layer 4: Comprehensive disposable, free provider,

app/validator/engine.py

@@ -18,8 +18,9 @@ from ..models import (
 from ..config import settings
 from .syntax import validate_syntax
 from .dns_check import check_dns
-from .disposable import check_disposable
+from .disposable import check_disposable, domain_in_blocklist, domain_is_known
 from .mailbox_provider import MailboxProvider, resolve_provider
+from ..ml.runtime import get_model
 from .retry_queue import is_greylist_result
 from .scoring import calculate_score
 
@@ -508,6 +509,27 @@ class EmailValidationEngine:
                         for e, _ in pairs
                     }
 
+            # ---- ML advisory signal (unknown domains, ONE batched call) ----
+            # Runs only where the curated lists have no answer - the exact gap
+            # a static blocklist cannot close. Any failure here is silently
+            # ignored: the deterministic verdict stands.
+            ml_map: dict[str, tuple[str, float]] = {}
+            ml_model = get_model()
+            if ml_model is not None:
+                unknown_domains = sorted(
+                    {
+                        domain
+                        for _, _, _, domain, _ in pending
+                        if not domain_is_known(domain)
+                    }
+                )
+                if unknown_domains:
+                    preds = ml_model.predict(
+                        [f"john.smith [SEP] {d}" for d in unknown_domains]
+                    )
+                    if preds:
+                        ml_map = dict(zip(unknown_domains, preds))
+
             # ---- Layer 7: score each row -----------------------------
             for idx, normalized, syntax, domain, dns in pending:
                 cached = cached_verdicts.get(normalized)
@@ -532,6 +554,11 @@ class EmailValidationEngine:
                     confidence=0.9 if ca_value is not None else 0.0,
                 )
                 disposable = check_disposable(normalized, domain)
+                ml = ml_map.get(domain)
+                if ml is not None:
+                    disposable = disposable.model_copy(
+                        update={"ml_label": ml[0], "ml_confidence": ml[1]}
+                    )
 
                 # A 4xx means "come back later", not a verdict. Schedule one
                 # bounded retry so a real mailbox on a greylisting domain is not

app/validator/scoring.py

@@ -117,6 +117,21 @@ def calculate_score(
             breakdown["free_provider"] = -5
             score -= 5
 
+    # ---- ML advisory (DistilBERT, unknown domains only) ---------------
+    # NEVER a hard verdict: a statistical false positive must not destroy a
+    # legitimate lead. Penalty is small and reversible by promotion gate.
+    ml_label = getattr(disposable, "ml_label", None)
+    if ml_label == "DISPOSABLE" and not disposable.is_disposable:
+        from ..config import settings  # local import: scoring stays pure
+
+        penalty = max(1, int(getattr(settings, "ml_score_penalty", 5)))
+        score -= penalty
+        breakdown["ML_DISPOSABLE_SUSPECT"] = -penalty
+        confidence = getattr(disposable, "ml_confidence", 0.0) or 0.0
+        warnings.append(
+            f"ML model flagged domain as disposable-like ({confidence:.2f})"
+        )
+
     # ---- Catch-All (-20) ---------------------------------------------
     if catch_all and catch_all.is_catch_all:
         risk_cap = min(risk_cap, 74)

requirements-ml.txt

@@ -0,0 +1,14 @@
+# Training-time dependencies ONLY - never installed in the deployed Space
+# image (which uses onnxruntime for inference; see Task 6 of the plan).
+# Install:  pip install -r requirements-ml.txt
+#
+# torch CPU build:
+#   pip install torch --index-url https://download.pytorch.org/whl/cpu
+torch>=2.3
+transformers>=4.41
+datasets>=2.19
+evaluate>=0.4.2
+accelerate>=0.31
+scikit-learn>=1.5
+optimum[onnxruntime]>=1.19
+huggingface-hub>=0.23

tests/test_engine_ml.py

@@ -0,0 +1,169 @@
+import unittest
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from app.models import DNSResult, DisposableResult, SMTPResult, ValidationStatus
+from app.validator.engine import EmailValidationEngine
+
+
+def _dns():
+    return DNSResult(
+        domain_exists=True,
+        has_mx_records=True,
+        mx_records=[{"host": "mx.corp.test", "priority": 10}],
+    )
+
+
+def _clean():
+    return DisposableResult(
+        is_disposable=False, is_free_provider=False, is_role_based=False
+    )
+
+
+class NullCache:
+    async def get_dns(self, domain):
+        return None
+
+    async def put_dns(self, domain, result):
+        return None
+
+    async def get_verdict(self, email, deep):
+        return None
+
+    async def put_verdict(self, email, deep, smtp, catch_all):
+        return None
+
+
+class FakeProvider:
+    async def verify_many(self, pairs):
+        return ({e: SMTPResult(verified=True, smtp_code=250) for e, _ in pairs}, {})
+
+    async def verify(self, email, mx_host):
+        return SMTPResult(verified=True, smtp_code=250)
+
+    async def catch_all(self, domain, mx_host):
+        return (False, 0.9)
+
+
+def _engine_with_ml(flag=True, repo="user/model", predict=None):
+    engine = EmailValidationEngine(mailbox=FakeProvider(), cache=NullCache())
+    if predict is not None:
+        fake_model = MagicMock()
+        fake_model.predict = MagicMock(side_effect=predict)
+        fake_model.load = MagicMock(return_value=True)
+        engine.mailbox = engine.mailbox  # unchanged
+        patcher = patch("app.validator.engine.get_model", return_value=fake_model)
+        patcher.start()
+        engine._ml_patcher = patcher  # keep alive until stopAll
+        engine._ml_model = fake_model
+    p = patch.object(
+        __import__("app.config", fromlist=["settings"]).settings,
+        "ml_disposable_enabled",
+        flag,
+    )
+    p.start()
+    engine._flag_patcher = p
+    p2 = patch.object(
+        __import__("app.config", fromlist=["settings"]).settings,
+        "ml_model_repo",
+        repo,
+    )
+    p2.start()
+    engine._flag_patcher2 = p2
+    return engine
+
+
+def _stop(engine):
+    engine._ml_patcher.stop() if hasattr(engine, "_ml_patcher") else None
+    engine._flag_patcher.stop()
+    engine._flag_patcher2.stop()
+
+
+class EngineMLTests(unittest.IsolatedAsyncioTestCase):
+    def tearDown(self):
+        for e in getattr(self, "_engines", []):
+            _stop(e)
+
+    async def _run(self, engine, emails):
+        self._engines = getattr(self, "_engines", [])
+        self._engines.append(engine)
+        # NOTE: check_disposable is NOT mocked - the real blocklist/role sets
+        # run, so the known-domain test exercises the true deterministic path.
+        with patch("app.validator.engine.check_dns", new=AsyncMock(return_value=_dns())):
+            return await engine.validate_batch(emails, deep=True)
+
+    async def test_flag_off_leaves_results_unchanged(self):
+        engine = _engine_with_ml(flag=False)
+        results = await self._run(engine, ["a@brand-new-xyz.tk"])
+
+        self.assertIsNone(results[0].disposable.ml_label)
+        self.assertIsNone(results[0].disposable.ml_confidence)
+        self.assertNotIn(
+            "ML model flagged domain as disposable-like",
+            results[0].scoring.warnings,
+        )
+
+    async def test_known_disposable_domain_skips_ml(self):
+        # mailinator.com IS in the live blocklist -> the deterministic layer
+        # already hard-fails it; the model must not be consulted.
+        engine = _engine_with_ml(flag=True, predict=[AssertionError("must not run")])
+        results = await self._run(engine, ["a@mailinator.com"])
+
+        self.assertEqual(results[0].status, ValidationStatus.INVALID)
+        fake_model = engine._ml_model
+        fake_model.predict.assert_not_called()
+
+    async def test_unknown_domain_gets_advisory_ml_signal(self):
+        engine = _engine_with_ml(
+            flag=True, predict=[[("DISPOSABLE", 0.93)]]  # contract: (label, score)
+        )
+        results = await self._run(engine, ["a@brand-new-disposable-xyz.tk"])
+
+        d = results[0].disposable
+        self.assertEqual(d.ml_label, "DISPOSABLE")
+        self.assertEqual(d.ml_confidence, 0.93)
+        self.assertTrue(
+            any("ML model flagged" in w for w in results[0].scoring.warnings)
+        )
+        # ADVISORY: the deterministic verdict must not become invalid.
+        self.assertNotEqual(results[0].status, ValidationStatus.INVALID)
+        # The penalty is proven by the breakdown: the SMTP-verified bonus (+10)
+        # can legitimately re-inflate the absolute score past the clamp.
+        self.assertEqual(
+            results[0].scoring.breakdown.get("ML_DISPOSABLE_SUSPECT"), -5
+        )
+
+    async def test_model_failure_degrades_to_deterministic(self):
+        engine = _engine_with_ml(flag=True, predict=[None])
+        results = await self._run(engine, ["a@brand-new-xyz.tk"])
+
+        self.assertIsNone(results[0].disposable.ml_label)
+        self.assertEqual(
+            results[0].scoring.warnings,
+            [w for w in results[0].scoring.warnings if "ML model flagged" not in w]
+            or results[0].scoring.warnings,
+        )
+        self.assertEqual(results[0].scoring.score, 100)
+
+    async def test_unknown_domains_batch_into_one_predict_call(self):
+        engine = _engine_with_ml(
+            flag=True,
+            predict=[
+                [
+                    ("DISPOSABLE", 0.9),
+                    ("NORMAL", 0.9),
+                    ("DISPOSABLE", 0.95),
+                ]
+            ],
+        )
+        emails = [f"u{i}@dk{i}.example" for i in range(3)]
+        results = await self._run(engine, emails)
+
+        fake_model = engine._ml_model
+        self.assertEqual(fake_model.predict.call_count, 1)
+        sent = fake_model.predict.call_args.args[0]
+        self.assertEqual(len(sent), 3)
+        self.assertEqual(results[1].disposable.ml_label, "NORMAL")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_ml_dataset.py

@@ -0,0 +1,98 @@
+import unittest
+
+from app.ml.dataset import build_dataset, label_name, LABEL_NAMES
+
+
+class DatasetBuilderTests(unittest.TestCase):
+    def test_labels_are_the_documented_four(self):
+        self.assertEqual(LABEL_NAMES, ["NORMAL", "DISPOSABLE", "ROLE", "SPAM_TRAP"])
+        self.assertEqual(label_name(0), "NORMAL")
+        self.assertEqual(label_name(1), "DISPOSABLE")
+        self.assertEqual(label_name(2), "ROLE")
+        self.assertEqual(label_name(3), "SPAM_TRAP")
+
+    def test_blocklist_domains_become_disposable(self):
+        rows = build_dataset(
+            blocklist={"tempmail.example", "10minutemail.example"},
+            free={"gmail.com"},
+            role_prefixes={"admin"},
+            normal_domains={"acme-corp.test"},
+            normal_locals={"john.smith", "maria.garcia"},
+            per_class_limit=50,
+        )
+        by_label = {}
+        for r in rows:
+            by_label.setdefault(r["label"], []).append(r["text"])
+        self.assertIn(1, by_label, "blocklist domains must yield DISPOSABLE rows")
+        self.assertTrue(any("tempmail.example" in t for t in by_label[1]))
+
+    def test_role_prefixes_become_role_label(self):
+        rows = build_dataset(
+            blocklist=set(),
+            free=set(),
+            role_prefixes={"admin", "info"},
+            normal_domains={"acme-corp.test"},
+            normal_locals={"john.smith"},
+            per_class_limit=50,
+        )
+        role_rows = [r for r in rows if r["label"] == 2]
+        self.assertTrue(role_rows, "role prefixes must yield ROLE rows")
+        self.assertTrue(any(r["text"].startswith("admin") for r in role_rows))
+
+    def test_normal_class_never_contains_blocklist_domains(self):
+        rows = build_dataset(
+            blocklist={"tempmail.example"},
+            free=set(),
+            role_prefixes=set(),
+            normal_domains={"acme-corp.test", "northwind.test"},
+            normal_locals={"john.smith", "maria.garcia", "kenji.tanaka"},
+            per_class_limit=50,
+        )
+        normal_texts = [r["text"] for r in rows if r["label"] == 0]
+        self.assertTrue(normal_texts, "NORMAL class must be populated")
+        self.assertTrue(all("tempmail.example" not in t for t in normal_texts))
+
+    def test_per_class_limit_is_respected(self):
+        rows = build_dataset(
+            blocklist={f"d{i}.example" for i in range(100)},
+            free=set(),
+            role_prefixes=set(),
+            normal_domains={"acme.test"},
+            normal_locals={"john.smith"},
+            per_class_limit=25,
+        )
+        counts = {}
+        for r in rows:
+            counts[r["label"]] = counts.get(r["label"], 0) + 1
+        self.assertTrue(all(c <= 25 for c in counts.values()), counts)
+
+    def test_text_format_is_local_sep_domain(self):
+        rows = build_dataset(
+            blocklist=set(),
+            free=set(),
+            role_prefixes={"admin"},
+            normal_domains={"acme.test"},
+            normal_locals={"john.smith"},
+            per_class_limit=5,
+        )
+        self.assertTrue(rows)
+        for r in rows:
+            self.assertIn(" [SEP] ", r["text"])
+
+    def test_deterministic_with_same_seed(self):
+        kwargs = dict(
+            blocklist={"tempmail.example"},
+            free={"gmail.com"},
+            role_prefixes={"admin"},
+            normal_domains={"acme.test"},
+            normal_locals={"john.smith"},
+            per_class_limit=20,
+            seed=7,
+        )
+        a = build_dataset(**kwargs)
+        b = build_dataset(**kwargs)
+        self.assertEqual(a, b, "same seed must produce identical rows")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_ml_runtime.py

@@ -0,0 +1,89 @@
+import unittest
+from unittest.mock import MagicMock, patch
+
+from app.ml.runtime import EmailRiskModel, get_model
+
+
+class RuntimeLoaderTests(unittest.TestCase):
+    """The ML layer must be strictly optional: any failure degrades to "no ML"
+    and validation continues deterministically."""
+
+    def test_load_failure_returns_false_and_never_raises(self):
+        m = EmailRiskModel(repo="user/model", token="t")
+        with patch(
+            "app.ml.runtime.ORTModelForSequenceClassification",
+            side_effect=RuntimeError("hub down"),
+        ):
+            self.assertFalse(m.load())
+        self.assertFalse(m.loaded)
+
+    def test_predict_before_load_returns_none(self):
+        m = EmailRiskModel(repo="user/model", token="t")
+        self.assertIsNone(m.predict(["a [SEP] b.test"]))
+
+    def test_predict_maps_label_and_confidence(self):
+        m = EmailRiskModel(repo="user/model", token="t")
+        fake = MagicMock(return_value=[{"label": "DISPOSABLE", "score": 0.93}])
+        m._pipeline = fake
+        m._loaded = True
+
+        out = m.predict(["x [SEP] tempmail.example"])
+
+        self.assertEqual(out, [("DISPOSABLE", 0.93)])
+        # Batching is passed through - 1-row calls are a performance bug.
+        fake.assert_called_once()
+        self.assertIn("batch_size", fake.call_args.kwargs)
+
+    def test_predict_error_returns_none_not_raises(self):
+        m = EmailRiskModel(repo="user/model", token="t")
+        fake = MagicMock(side_effect=RuntimeError("onnx exploded"))
+        m._pipeline = fake
+        m._loaded = True
+        self.assertIsNone(m.predict(["a [SEP] b.test"]))
+
+    def test_predict_normalizes_label_case(self):
+        m = EmailRiskModel(repo="user/model", token="t")
+        fake = MagicMock(return_value=[{"label": "disposable", "score": 0.5}])
+        m._pipeline = fake
+        m._loaded = True
+        self.assertEqual(m.predict(["a [SEP] b.test"])[0][0], "DISPOSABLE")
+
+    def test_predict_accepts_plain_string_batch_return(self):
+        # transformers may return a plain list of dicts OR nested batch lists.
+        m = EmailRiskModel(repo="user/model", token="t")
+        fake = MagicMock(return_value=[[{"label": "NORMAL", "score": 0.7}]])
+        m._pipeline = fake
+        m._loaded = True
+        self.assertEqual(m.predict(["a [SEP] b.test"]), [("NORMAL", 0.7)])
+
+    def test_config_defaults_are_off(self):
+        from app.config import settings
+
+        self.assertFalse(settings.ml_disposable_enabled)
+        self.assertEqual(settings.ml_model_repo, "")
+
+
+class GetModelTests(unittest.TestCase):
+    def test_returns_none_when_disabled(self):
+        with patch("app.ml.runtime.settings") as fake_settings:
+            fake_settings.ml_disposable_enabled = False
+            fake_settings.ml_model_repo = "user/model"
+            self.assertIsNone(get_model())
+
+    def test_returns_none_when_no_repo(self):
+        with patch("app.ml.runtime.settings") as fake_settings:
+            fake_settings.ml_disposable_enabled = True
+            fake_settings.ml_model_repo = ""
+            self.assertIsNone(get_model())
+
+    def test_returns_model_when_enabled_and_configured(self):
+        with patch("app.ml.runtime.settings") as fake_settings:
+            fake_settings.ml_disposable_enabled = True
+            fake_settings.ml_model_repo = "user/model"
+            m = get_model()
+        self.assertIsInstance(m, EmailRiskModel)
+        self.assertEqual(m.repo, "user/model")
+
+
+if __name__ == "__main__":
+    unittest.main()