| """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 |
|
|
| |
| 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"] |
|
|
| |
| add(NORMAL, [f"{loc} [SEP] {d}" for d in normal_domains for loc in 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_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]) |
|
|
| |
| 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"]]) |
