File size: 3,804 Bytes
cd07add | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 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"]]) |