| |
| """Data preparation for the BITS Phishing Detector. |
| |
| Download -> normalise -> dedupe -> decontaminate -> split -> leakage-check -> manifests. |
| Procedure and expected counts: docs/research/verify_data.md sections 7 and 8. |
| |
| Subcommands |
| inspect features + 2 rows of all four datasets (streaming; nothing is fully downloaded) |
| build produce data/processed/{pretrain,finetune_train,finetune_val,test_holdout}.parquet, |
| splits_manifest.json and data_report.md |
| check re-run the leakage assertions on the written files and print the manifest summary |
| build3 THREE-WAY labels (RUN_LOG entry 6): add label3/label_set to the binary splits, ingest the Champa |
| phishing corpora (Nazario, Nigerian_Fraud) [+ bigshp if data.threeway.bigshp.enabled], carve the |
| clean-provenance test_holdout3.parquet, rewrite finetune_train/val + test_holdout (2 columns added) |
| check3 recompute every binary + three-way leakage assertion from the written files |
| |
| Leakage contract (hard failures, asserted identically by `build` and `check`): the held-out |
| TREC-06 sub-source never shares an exact key (k_full / k_pre / k_sb / k_sb_full) or a >=3-shingle |
| near-duplicate link with the fine-tune train/val splits or the MLM pretraining corpus; train and val |
| never share a k_full; val never shares a key with the pretraining corpus; the holdout source never |
| appears in train/val. Info-only counters: train/val near-dup links, train/pretrain key overlap, |
| val/pretrain near-dup links. |
| |
| Three-way contract (build3 / check3): label3 in {0 legitimate, 1 phishing, 2 spam, -1 unwanted-but-unsplit}, |
| label_set in {legit, phish, spam, unwanted}; binary `label` unchanged. Binary split membership is unchanged |
| except for the rows moved from finetune_val into test_holdout3. test_holdout3 is asserted to share zero keys and |
| zero >=3-shingle links with train, val and the pretraining corpus and zero keys with the TREC-06 holdout. |
| |
| The assertions are an independent recomputation, not the removal bookkeeping: after every removal |
| step the shingle index is rebuilt from the surviving rows (and the pretraining rows are re-derived |
| from their stored `text` exactly as `check` does), so a bug in the removal logic surfaces as a |
| failed assertion rather than a silent pass. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import logging |
| import os |
| import re |
| import resource |
| import shutil |
| import sys |
| import unicodedata |
| from collections import Counter, defaultdict |
| from dataclasses import dataclass, field |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any, Iterable |
|
|
| import numpy as np |
| import pandas as pd |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
| import common |
|
|
| LOG = logging.getLogger("data_prep") |
|
|
| |
| @dataclass |
| class Dials: |
| """Dedup / pretraining thresholds. Read from config.yaml when present, else these defaults. |
| |
| Optional config keys (all under `data`): `dedup.key_prefix_chars`, `dedup.key_prefix_min_len`, |
| `dedup.near_dup_shingle_cap`, `pretrain_min_alnum_chars`, `quote_strip_min_alnum`. The body-key |
| fallback threshold is always `data.min_alnum_chars` so the two cannot diverge. |
| """ |
|
|
| key_prefix_chars: int = 300 |
| key_prefix_min_len: int = 80 |
| body_key_min_alnum: int = 20 |
| pretrain_min_alnum: int = 50 |
| quote_strip_min_alnum: int = 200 |
| shingle_cap: int = 2000 |
|
|
| @classmethod |
| def from_config(cls, cfg: dict) -> "Dials": |
| d = cfg["data"] |
| dd = d.get("dedup", {}) |
| return cls( |
| key_prefix_chars=int(dd.get("key_prefix_chars", cls.key_prefix_chars)), |
| key_prefix_min_len=int(dd.get("key_prefix_min_len", cls.key_prefix_min_len)), |
| body_key_min_alnum=int(d["min_alnum_chars"]), |
| pretrain_min_alnum=int(d.get("pretrain_min_alnum_chars", cls.pretrain_min_alnum)), |
| quote_strip_min_alnum=int(d.get("quote_strip_min_alnum", cls.quote_strip_min_alnum)), |
| shingle_cap=int(dd.get("near_dup_shingle_cap", cls.shingle_cap)), |
| ) |
|
|
|
|
| DIALS = Dials() |
|
|
| LABELLED_COLUMNS = ["text", "subject", "label", "dataset_name"] |
| PRETRAIN_COLUMNS = ["subject", "body"] |
| FINETUNE_SCHEMA = ["doc_id", "text", "subject", "body", "label", "source", "cluster_id"] |
| PRETRAIN_SCHEMA = ["doc_id", "text", "n_chars"] |
|
|
| _CRLF = re.compile(r"\r\n?") |
| _BLANK_RUN = re.compile(r"\n(?:[ \t]*\n){2,}") |
| _QUOTE_MARKERS = re.compile( |
| r"(?P<original>-----\s*Original Message\s*-----)" |
| r"|(?P<forwarded>-{3,}\s*Forwarded by\b)" |
| r"|(?P<header_block>^From:[^\n]*\n(?:[^\n]*\n){0,2}(?:Sent|To|Subject):)", |
| re.MULTILINE, |
| ) |
|
|
|
|
| class SpamPolicyError(ValueError): |
| """Raised for spam policies that the available corpora cannot implement.""" |
|
|
|
|
| |
| def clean_text(s: Any) -> str: |
| """NFKC, normalised line endings, collapsed blank-line runs, stripped. |
| |
| Keeps the email readable for a byte-level BPE (no aggressive stripping of punctuation or |
| HTML) while removing the formatting noise that would otherwise make identical emails look |
| different to the tokenizer and inflate the near-duplicate count. |
| """ |
| if not isinstance(s, str) or not s: |
| return "" |
| s = unicodedata.normalize("NFKC", s).replace("\x00", "") |
| s = _CRLF.sub("\n", s) |
| s = _BLANK_RUN.sub("\n\n", s) |
| return s.strip() |
|
|
|
|
| def model_input(subject: str, body: str) -> str: |
| """subject + blank line + body. Never sender/receiver/date (TREC headers are synthetic).""" |
| return "\n\n".join(part for part in (subject, body) if part) |
|
|
|
|
| def strip_quotes(body: str) -> tuple[str, str | None, bool]: |
| """Cut a corbt body at its first quote/forward marker when >=200 alphanumerics survive. |
| |
| Enron threads quote the whole history, so without this ~43% of pretraining characters would |
| be verbatim repeats of other rows. Returns (body, marker_name, cut_applied). |
| """ |
| m = _QUOTE_MARKERS.search(body) |
| if m is None: |
| return body, None, False |
| head = body[: m.start()].rstrip() |
| if len(common.normalize_text(head)) >= DIALS.quote_strip_min_alnum: |
| return head, m.lastgroup, True |
| return body, m.lastgroup, False |
|
|
|
|
| def row_keys(subject: str, body: str) -> tuple[str, str, str, str, str, int]: |
| """Dedup keys per verify_data.md section 7. |
| |
| Returns (k_full, k_pre, k_sb, k_sb_full, base, n_alnum): |
| k_full sha1 of the normalised body (the content-level identity of the email) |
| k_pre sha1 of its first 300 chars (catches appended footers / truncated copies), or |
| k_full when the body is shorter than 80 chars |
| k_sb sha1 of the first 300 chars of normalised subject+body (catches corpora that |
| embed the subject in the body, e.g. Enron-Spam, and corbt subject+body) |
| k_sb_full sha1 of the whole normalised subject+body (corbt dedup key) |
| base the normalised string used for k_full/k_pre and shingles |
| n_alnum alphanumerics in subject+body (the min_alnum_chars filter) |
| When the body holds fewer than `min_alnum_chars` alphanumerics the subject+body string is used |
| as `base` instead, otherwise every empty-body email would collide on sha1(""). |
| """ |
| norm_body = common.normalize_text(body) |
| norm_sb = common.normalize_text(subject) + norm_body |
| base = norm_body if len(norm_body) >= DIALS.body_key_min_alnum else norm_sb |
| k_full = common.sha1_hex(base) |
| k_pre = common.sha1_hex(base[: DIALS.key_prefix_chars]) if len(base) >= DIALS.key_prefix_min_len else k_full |
| k_sb = common.sha1_hex(norm_sb[: DIALS.key_prefix_chars]) |
| k_sb_full = common.sha1_hex(norm_sb) |
| return k_full, k_pre, k_sb, k_sb_full, base, len(norm_sb) |
|
|
|
|
| def split_text(text: str) -> tuple[str, str]: |
| """Recover (subject, body) from a stored `text` = model_input(subject, body). |
| |
| Exact whenever the subject was non-empty (subjects never contain a blank line). For an |
| empty-subject email with a multi-paragraph body the first paragraph is mistaken for the |
| subject; `text_identity` is therefore only a lower bound on the build-time identity, and |
| build_pretrain compensates by matching on BOTH identities for empty-subject rows so the |
| post-hoc assertion (shared with `check`) can never see a link the removal pass did not. |
| """ |
| head, sep, rest = text.partition("\n\n") |
| return (head, rest) if sep else ("", head) |
|
|
|
|
| def text_identity(text: str) -> tuple[set[str], str]: |
| """(keys, shingle base) of a stored pretraining `text`, recomputed the same way by build and check. |
| |
| normalize(subject + "\\n\\n" + body) == normalize(subject) + normalize(body), so k_sb / k_sb_full |
| are always exact; k_full / k_pre / base are exact unless the subject was empty (see split_text). |
| """ |
| k_full, k_pre, k_sb, k_sb_full, base, _ = row_keys(*split_text(text)) |
| return {k_full, k_pre, k_sb, k_sb_full}, base |
|
|
|
|
| def add_key_columns(df: pd.DataFrame) -> list[str]: |
| """Attach k_full/k_pre/k_sb/k_sb_full/n_alnum columns; return the shingle bases (window-trimmed).""" |
| cols: dict[str, list] = defaultdict(list) |
| bases: list[str] = [] |
| for subj, body in zip(df["subject"].tolist(), df["body"].tolist()): |
| k_full, k_pre, k_sb, k_sb_full, base, n_alnum = row_keys(subj, body) |
| cols["k_full"].append(k_full) |
| cols["k_pre"].append(k_pre) |
| cols["k_sb"].append(k_sb) |
| cols["k_sb_full"].append(k_sb_full) |
| cols["n_alnum"].append(n_alnum) |
| bases.append(base) |
| for name, values in cols.items(): |
| df[name] = values |
| return bases |
|
|
|
|
| def key_union(df: pd.DataFrame) -> set[str]: |
| out: set[str] = set() |
| for col in ("k_full", "k_pre", "k_sb", "k_sb_full"): |
| out.update(df[col].tolist()) |
| return out |
|
|
|
|
| def rows_hitting(df: pd.DataFrame, keys: set[str]) -> np.ndarray: |
| """Boolean mask: row has any of its four keys inside `keys`.""" |
| mask = np.zeros(len(df), dtype=bool) |
| for col in ("k_full", "k_pre", "k_sb", "k_sb_full"): |
| mask |= np.fromiter((k in keys for k in df[col].tolist()), dtype=bool, count=len(df)) |
| return mask |
|
|
|
|
| |
| @dataclass |
| class NearDupConfig: |
| shingle_chars: int = 50 |
| stride: int = 25 |
| window_chars: int = 1500 |
| min_shared: int = 3 |
| cap: int = 2000 |
|
|
| @classmethod |
| def from_config(cls, cfg: dict, cap: int) -> "NearDupConfig": |
| d = cfg["data"]["dedup"] |
| return cls( |
| shingle_chars=int(d["near_dup_shingle_chars"]), |
| stride=int(d["near_dup_stride"]), |
| window_chars=int(d["near_dup_window_chars"]), |
| min_shared=int(d["near_dup_min_shared"]), |
| cap=cap, |
| ) |
|
|
|
|
| def shingle_hashes(base: str, nd: NearDupConfig) -> np.ndarray: |
| """Distinct 64-bit hashes of the fixed-stride shingles of base[:window].""" |
| s = base[: nd.window_chars] |
| n = len(s) |
| if n < nd.shingle_chars: |
| return np.empty(0, dtype=np.int64) |
| digests = b"".join( |
| hashlib.blake2b(s[i : i + nd.shingle_chars].encode("ascii"), digest_size=8).digest() |
| for i in range(0, n - nd.shingle_chars + 1, nd.stride) |
| ) |
| return np.unique(np.frombuffer(digests, dtype=np.int64)) |
|
|
|
|
| class ShingleIndex: |
| """Inverted shingle index over a list of normalised texts. |
| |
| Shingles shared by more than `cap` rows are boilerplate (list footers, pharma templates) and |
| are ignored: linking everything that carries the same footer would merge half a corpus into |
| one cluster and remove legitimate training rows. This is a known, logged limitation. |
| """ |
|
|
| def __init__(self, bases: list[str], nd: NearDupConfig): |
| self.nd = nd |
| self.n = len(bases) |
| per_row = [shingle_hashes(b, nd) for b in bases] |
| lens = np.fromiter((len(h) for h in per_row), dtype=np.int64, count=self.n) |
| self.offsets = np.concatenate(([0], np.cumsum(lens))) |
| self.hashes = np.concatenate(per_row) if self.n else np.empty(0, dtype=np.int64) |
| rid = np.repeat(np.arange(self.n, dtype=np.int64), lens) |
| order = np.argsort(self.hashes, kind="stable") |
| self.sorted_rid = rid[order] |
| self.uniq, self.start, self.counts = np.unique( |
| self.hashes[order], return_index=True, return_counts=True |
| ) |
| capped = self.counts > nd.cap |
| self.n_capped_shingles = int(capped.sum()) |
| self.n_rows_touching_capped = int( |
| np.unique(self.sorted_rid[np.repeat(capped, self.counts)]).size |
| ) |
| LOG.info( |
| "shingle index: %d rows, %d shingle occurrences, %d distinct, %d shingles over cap=%d " |
| "(ignored; touching %d rows)", |
| self.n, self.hashes.size, self.uniq.size, self.n_capped_shingles, nd.cap, |
| self.n_rows_touching_capped, |
| ) |
|
|
| def links(self) -> tuple[np.ndarray, np.ndarray]: |
| """All pairs (a < b) sharing >= min_shared distinct shingles.""" |
| out_a: list[np.ndarray] = [] |
| out_b: list[np.ndarray] = [] |
| t0 = common.timer() |
| for r in range(self.n): |
| h = self.hashes[self.offsets[r] : self.offsets[r + 1]] |
| if h.size == 0: |
| continue |
| g = np.searchsorted(self.uniq, h) |
| cnt = self.counts[g] |
| ok = (cnt >= 2) & (cnt <= self.nd.cap) |
| if not ok.any(): |
| continue |
| starts = self.start[g[ok]] |
| lens = cnt[ok] |
| idx = np.repeat(starts - np.concatenate(([0], np.cumsum(lens)[:-1])), lens) + np.arange(lens.sum()) |
| cand = self.sorted_rid[idx] |
| cand = cand[cand > r] |
| if cand.size < self.nd.min_shared: |
| continue |
| u, c = np.unique(cand, return_counts=True) |
| hit = u[c >= self.nd.min_shared] |
| if hit.size: |
| out_a.append(np.full(hit.size, r, dtype=np.int64)) |
| out_b.append(hit) |
| if r and r % 50_000 == 0: |
| LOG.info(" near-dup scan %d/%d rows (%.1fs)", r, self.n, common.timer() - t0) |
| if not out_a: |
| return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) |
| a, b = np.concatenate(out_a), np.concatenate(out_b) |
| LOG.info("near-dup links: %d pairs in %.1fs", a.size, common.timer() - t0) |
| return a, b |
|
|
| def query(self, bases: list[str]) -> tuple[np.ndarray, np.ndarray]: |
| """Pairs (query_idx, indexed_idx) where a foreign text shares >= min_shared shingles with an indexed row. |
| |
| Used to screen the (large) pretraining corpus against the (small) holdout+val reference set |
| without indexing the corpus itself: cost is one searchsorted per query row, so the millions |
| of pretrain-pretrain Enron-thread links are never materialised. |
| """ |
| out_q: list[np.ndarray] = [] |
| out_r: list[np.ndarray] = [] |
| t0 = common.timer() |
| for qi, base in enumerate(bases): |
| h = shingle_hashes(base, self.nd) |
| if h.size == 0 or self.uniq.size == 0: |
| continue |
| g = np.searchsorted(self.uniq, h) |
| inb = g < self.uniq.size |
| g, h = g[inb], h[inb] |
| g = g[self.uniq[g] == h] |
| if g.size == 0: |
| continue |
| cnt = self.counts[g] |
| ok = cnt <= self.nd.cap |
| if not ok.any(): |
| continue |
| starts = self.start[g[ok]] |
| lens = cnt[ok] |
| idx = np.repeat(starts - np.concatenate(([0], np.cumsum(lens)[:-1])), lens) + np.arange(lens.sum()) |
| u, c = np.unique(self.sorted_rid[idx], return_counts=True) |
| hit = u[c >= self.nd.min_shared] |
| if hit.size: |
| out_q.append(np.full(hit.size, qi, dtype=np.int64)) |
| out_r.append(hit) |
| if qi and qi % 100_000 == 0: |
| LOG.info(" near-dup query %d/%d rows (%.1fs)", qi, len(bases), common.timer() - t0) |
| if not out_q: |
| return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64) |
| q, r = np.concatenate(out_q), np.concatenate(out_r) |
| LOG.info("near-dup query: %d rows -> %d links in %.1fs", len(bases), q.size, common.timer() - t0) |
| return q, r |
|
|
|
|
| def recount_shingle_links( |
| train_bases: list[str], val_bases: list[str], hold_bases: list[str], pretrain_bases: list[str], nd: NearDupConfig |
| ) -> dict[str, int]: |
| """Independent near-dup link counts between the FINAL sets, from a fresh index. |
| |
| This is the leakage gate proper: it never reuses the index or masks that drove the removals, |
| so it fails if the removal logic is wrong. Same code path in `build` (after writing nothing yet) |
| and in `check` (from the parquet files). Pretrain rows are queried against a holdout+val index |
| rather than indexed, so pretrain-pretrain links are never enumerated. |
| """ |
| n_tr, n_va = len(train_bases), len(val_bases) |
| index = ShingleIndex(train_bases + val_bases + hold_bases, nd) |
| la, lb = index.links() |
| part_a = np.where(la < n_tr, 0, np.where(la < n_tr + n_va, 1, 2)) |
| part_b = np.where(lb < n_tr, 0, np.where(lb < n_tr + n_va, 1, 2)) |
| pair = Counter(zip(np.minimum(part_a, part_b).tolist(), np.maximum(part_a, part_b).tolist())) |
| ref = ShingleIndex(hold_bases + val_bases, nd) |
| _q, r = ref.query(pretrain_bases) |
| out = { |
| "holdout_train": int(pair[(0, 2)]), |
| "holdout_val": int(pair[(1, 2)]), |
| "train_val": int(pair[(0, 1)]), |
| "holdout_pretrain": int(np.sum(r < len(hold_bases))), |
| "val_pretrain": int(np.sum(r >= len(hold_bases))), |
| } |
| LOG.info("recount shingle links (final sets): %s", out) |
| return out |
|
|
|
|
| def union_find(n: int, pairs: Iterable[tuple[int, int]]) -> np.ndarray: |
| """Dense cluster ids (0..K-1, in order of first appearance) from an edge list.""" |
| parent = list(range(n)) |
|
|
| def find(x: int) -> int: |
| root = x |
| while parent[root] != root: |
| root = parent[root] |
| while parent[x] != root: |
| parent[x], x = root, parent[x] |
| return root |
|
|
| for a, b in pairs: |
| ra, rb = find(a), find(b) |
| if ra != rb: |
| parent[max(ra, rb)] = min(ra, rb) |
| roots = np.fromiter((find(i) for i in range(n)), dtype=np.int64, count=n) |
| _, dense = np.unique(roots, return_inverse=True) |
| return dense.astype(np.int64) |
|
|
|
|
| def same_key_pairs(keys: list[str]) -> Iterable[tuple[int, int]]: |
| """Consecutive-member edges for every group of rows sharing a key.""" |
| first: dict[str, int] = {} |
| for i, k in enumerate(keys): |
| j = first.setdefault(k, i) |
| if j != i: |
| yield j, i |
|
|
|
|
| def same_key_pairs_multi(columns: list[list[str]]) -> Iterable[tuple[int, int]]: |
| """Edges for rows sharing a key value within OR across the given key columns. |
| |
| k_pre (first 300 chars of the body) and k_sb (first 300 chars of subject+body) can collide across columns |
| when one corpus embeds the subject in the body; the leakage gate's key_union treats such a pair as shared, |
| so the clustering must too or two near-identical rows could land on different sides of a split. |
| """ |
| first: dict[str, int] = {} |
| for col in columns: |
| for i, k in enumerate(col): |
| j = first.setdefault(k, i) |
| if j != i: |
| yield j, i |
|
|
|
|
| |
| def hub_sha(repo: str) -> str | None: |
| try: |
| from huggingface_hub import HfApi |
|
|
| return HfApi().dataset_info(repo).sha |
| except Exception as exc: |
| LOG.warning("could not fetch commit sha for %s: %s", repo, exc) |
| return None |
|
|
|
|
| def load_labelled(cfg: dict, limit: int | None, rng: np.random.Generator) -> tuple[pd.DataFrame, dict]: |
| """All Hub splits of the primary set concatenated; the Hub's own split is ignored (it leaks).""" |
| from datasets import concatenate_datasets, load_dataset |
|
|
| repo = cfg["data"]["labelled_primary"] |
| cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache") |
| sha = hub_sha(repo) |
| dsd = load_dataset(repo, cache_dir=cache_dir, revision=sha) |
| split_names = list(dsd.keys()) |
| parts = [dsd[s].select_columns(LABELLED_COLUMNS) for s in split_names] |
| ds = concatenate_datasets(parts) |
| df = ds.to_pandas() |
| df = df.rename(columns={"text": "body", "dataset_name": "source"}) |
| df["row_idx"] = np.arange(len(df), dtype=np.int64) |
| meta = { |
| "repo": repo, |
| "commit_sha": sha, |
| "hub_splits_concatenated": {s: int(dsd[s].num_rows) for s in split_names}, |
| "datasets_fingerprints": {s: dsd[s]._fingerprint for s in split_names}, |
| "rows_loaded": int(len(df)), |
| } |
| LOG.info("labelled primary %s @ %s: %d rows from splits %s", repo, sha, len(df), meta["hub_splits_concatenated"]) |
| if limit: |
| keep = [] |
| for source, grp in df.groupby("source", sort=True): |
| n = min(limit, len(grp)) |
| keep.append(grp.iloc[np.sort(rng.choice(len(grp), n, replace=False))]) |
| df = pd.concat(keep).sort_values("row_idx").reset_index(drop=True) |
| LOG.warning("--limit %d: labelled subsample %d rows (per source <= %d)", limit, len(df), limit) |
| return df, meta |
|
|
|
|
| def iter_pretrain_batches(cfg: dict, limit: int | None, rng: np.random.Generator, batch_size: int = 20_000): |
| """Yield (row_indices, subjects, bodies) batches of corbt without materialising the corpus.""" |
| from datasets import load_dataset |
|
|
| repo = cfg["data"]["pretrain_corpus"] |
| cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache") |
| sha = hub_sha(repo) |
| ds = load_dataset(repo, split="train", cache_dir=cache_dir, revision=sha) |
| meta = {"repo": repo, "commit_sha": sha, "datasets_fingerprint": ds._fingerprint, "rows_total": int(ds.num_rows)} |
| indices = np.arange(ds.num_rows, dtype=np.int64) |
| if limit: |
| n = min(limit * 4, ds.num_rows) |
| indices = np.sort(rng.choice(ds.num_rows, n, replace=False)) |
| ds = ds.select(indices) |
| LOG.warning("--limit %d: corbt subsample %d rows", limit, n) |
| ds = ds.select_columns(PRETRAIN_COLUMNS) |
| LOG.info("pretrain corpus %s @ %s: %d rows total, processing %d", repo, sha, meta["rows_total"], len(indices)) |
|
|
| def batches(): |
| pos = 0 |
| for batch in ds.iter(batch_size=batch_size): |
| n = len(batch["body"]) |
| yield indices[pos : pos + n], batch["subject"], batch["body"] |
| pos += n |
|
|
| return meta, batches() |
|
|
|
|
| |
| def check_spam_policy(policy: str) -> None: |
| """Only `positive_class` is implementable; the others are kept so the decision stays visible.""" |
| if policy == "positive_class": |
| return |
| if policy in ("map_to_zero", "exclude"): |
| raise SpamPolicyError( |
| f"spam_policy={policy!r} cannot be implemented: no dataset in scope distinguishes spam " |
| "from phishing. Every label-1 row of puyang2025/seven-phishing-email-datasets comes from " |
| "a spam-track corpus (TREC-05/06/07, CEAS-08, Enron-Spam, SpamAssassin, Ling-Spam) with a " |
| "single binary label, so mapping spam to 0 or excluding spam sources would leave ~0 " |
| "positives (verify_data.md section 5). Use positive_class (positive = phishing/spam, " |
| "'unwanted, potentially malicious') and report the credential-lure slice separately." |
| ) |
| raise SpamPolicyError(f"unknown spam_policy {policy!r}; expected positive_class|map_to_zero|exclude") |
|
|
|
|
| |
| def grouped_stratified_split( |
| pool: pd.DataFrame, val_fraction: float, rng: np.random.Generator |
| ) -> tuple[np.ndarray, dict]: |
| """Assign whole near-dup clusters to val until each source x label stratum reaches ~val_fraction. |
| |
| A cluster never straddles train/val, so a template seen in training can never inflate the |
| validation score. Clusters spanning several strata are assigned under their majority stratum; |
| a cluster larger than a stratum's whole validation target goes to train. |
| """ |
| strata = pool["source"].astype(str) + "|" + pool["label"].astype(str) |
| cluster = pool["cluster_id"].to_numpy() |
| members: dict[int, list[int]] = defaultdict(list) |
| for pos, cid in enumerate(cluster.tolist()): |
| members[int(cid)].append(pos) |
| stratum_arr = strata.to_numpy() |
| by_stratum: dict[str, list[tuple[int, int]]] = defaultdict(list) |
| mixed = 0 |
| for cid, pos_list in members.items(): |
| counts = Counter(stratum_arr[p] for p in pos_list) |
| if len(counts) > 1: |
| mixed += 1 |
| by_stratum[counts.most_common(1)[0][0]].append((cid, len(pos_list))) |
|
|
| is_val = np.zeros(len(pool), dtype=bool) |
| per_stratum: dict[str, dict[str, Any]] = {} |
| for stratum in sorted(by_stratum): |
| groups = by_stratum[stratum] |
| n_rows = sum(size for _, size in groups) |
| target = int(round(val_fraction * n_rows)) |
| order = rng.permutation(len(groups)) |
| val_rows = 0 |
| oversized = 0 |
| for gi in order: |
| cid, size = groups[gi] |
| if size > target: |
| oversized += 1 |
| continue |
| if val_rows < target and abs(val_rows + size - target) <= abs(val_rows - target): |
| for p in members[cid]: |
| is_val[p] = True |
| val_rows += size |
| per_stratum[stratum] = { |
| "rows": n_rows, "clusters": len(groups), "val_target": target, "val_rows": val_rows, |
| "val_fraction": round(val_rows / n_rows, 4) if n_rows else 0.0, |
| "clusters_too_large_for_val": oversized, |
| } |
| info = { |
| "val_fraction_target": val_fraction, |
| "val_fraction_achieved": round(float(is_val.mean()), 4) if len(pool) else 0.0, |
| "clusters_total": len(members), |
| "clusters_mixed_stratum": mixed, |
| "largest_cluster": max((len(v) for v in members.values()), default=0), |
| "per_stratum": per_stratum, |
| } |
| return is_val, info |
|
|
|
|
| |
| def nested_counts(df: pd.DataFrame) -> dict[str, dict[str, int]]: |
| """{source: {label: n}} with string keys for JSON.""" |
| out: dict[str, dict[str, int]] = {} |
| if len(df) == 0: |
| return out |
| for (source, label), n in df.groupby(["source", "label"], sort=True).size().items(): |
| out.setdefault(str(source), {})[str(label)] = int(n) |
| return out |
|
|
|
|
| def split_totals(df: pd.DataFrame) -> dict[str, Any]: |
| n = int(len(df)) |
| pos = int((df["label"] == 1).sum()) if n else 0 |
| return {"rows": n, "positives": pos, "negatives": n - pos, "base_rate": round(pos / n, 4) if n else None} |
|
|
|
|
| def md_table(headers: list[str], rows: list[list[Any]]) -> str: |
| lines = ["| " + " | ".join(headers) + " |", "|" + "|".join("---" for _ in headers) + "|"] |
| for row in rows: |
| lines.append("| " + " | ".join(f"{v:,}" if isinstance(v, int) and not isinstance(v, bool) else str(v) for v in row) + " |") |
| return "\n".join(lines) |
|
|
|
|
| def source_label_table(counts: dict[str, dict[str, int]]) -> str: |
| rows = [] |
| for source in sorted(counts): |
| neg = counts[source].get("0", 0) |
| pos = counts[source].get("1", 0) |
| tot = neg + pos |
| rows.append([source, neg, pos, tot, f"{pos / tot:.1%}" if tot else "-"]) |
| return md_table(["source", "label 0", "label 1", "rows", "positive rate"], rows) |
|
|
|
|
| def peak_rss_gb() -> float: |
| return round(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20, 2) |
|
|
|
|
| |
| @dataclass |
| class LeakageResults: |
| checks: dict[str, dict[str, Any]] = field(default_factory=dict) |
|
|
| def add(self, name: str, count: int, assert_zero: bool, detail: str = "") -> None: |
| self.checks[name] = {"count": int(count), "asserted": assert_zero, "pass": (count == 0) or not assert_zero, "detail": detail} |
| level = logging.ERROR if (assert_zero and count) else logging.INFO |
| LOG.log(level, "leakage %-44s %8d %s", name, count, "(assert zero)" if assert_zero else "(info)") |
|
|
| def failures(self) -> list[str]: |
| return [k for k, v in self.checks.items() if v["asserted"] and v["count"]] |
|
|
| def raise_if_failed(self) -> None: |
| bad = self.failures() |
| if bad: |
| raise AssertionError("LEAKAGE ASSERTIONS FAILED: " + ", ".join(f"{k}={self.checks[k]['count']}" for k in bad)) |
|
|
|
|
| def leakage_checks( |
| train: pd.DataFrame, |
| val: pd.DataFrame, |
| holdout: pd.DataFrame, |
| pretrain_keys: set[str], |
| n_pretrain: int, |
| shingle_links: dict[str, int], |
| holdout_source: str, |
| ) -> LeakageResults: |
| """The gate. `shingle_links` must come from recount_shingle_links over the final sets and |
| `pretrain_keys` from text_identity over the final pretraining texts, in build and check alike.""" |
| res = LeakageResults() |
| hold_keys = key_union(holdout) |
| val_keys = key_union(val) |
| train_keys = key_union(train) |
| res.add("holdout_vs_train_key_rows", int(rows_hitting(train, hold_keys).sum()), True) |
| res.add("holdout_vs_val_key_rows", int(rows_hitting(val, hold_keys).sum()), True) |
| res.add("holdout_vs_pretrain_key_rows", len(pretrain_keys & hold_keys), True, "pretrain keys re-derived from stored text (text_identity)") |
| res.add("holdout_vs_train_shingle_links", shingle_links["holdout_train"], True) |
| res.add("holdout_vs_val_shingle_links", shingle_links["holdout_val"], True) |
| res.add("holdout_vs_pretrain_shingle_links", shingle_links["holdout_pretrain"], True, "pretrain bases re-derived from stored text (text_identity)") |
| res.add("train_vs_val_k_full_rows", len(set(train["k_full"]) & set(val["k_full"])), True) |
| res.add("train_vs_val_any_key_rows", len(train_keys & val_keys), False) |
| res.add("train_vs_val_shingle_links", shingle_links["train_val"], False) |
| res.add("val_vs_pretrain_key_rows", len(pretrain_keys & val_keys), True) |
| res.add("val_vs_pretrain_shingle_links", shingle_links["val_pretrain"], False, "near-twins of val emails in MLM text; counted, not removed (val is model selection, the holdout is the gate)") |
| res.add("train_vs_pretrain_key_overlap", len(pretrain_keys & train_keys), False, "allowed: TREC-05/Enron ham are Enron mail") |
| res.add("holdout_rows_in_train_or_val_by_source", int((train["source"] == holdout_source).sum() + (val["source"] == holdout_source).sum()), True) |
| res.add("test_holdout_rows_from_other_sources", int((holdout["source"] != holdout_source).sum()), True) |
| res.add("pretrain_rows", n_pretrain, False) |
| return res |
|
|
|
|
| |
| def cmd_build(args: argparse.Namespace, cfg: dict) -> None: |
| t_start = common.timer() |
| policy = args.spam_policy or cfg["data"]["spam_policy"] |
| check_spam_policy(policy) |
| seed = int(cfg["seed"]) |
| common.set_seed(seed) |
| rng = np.random.default_rng(seed) |
| strip = not args.no_strip_quotes |
| nd = NearDupConfig.from_config(cfg, cap=args.shingle_cap or DIALS.shingle_cap) |
| holdout_source = cfg["data"]["holdout_test_source"] |
| min_alnum = int(cfg["data"]["min_alnum_chars"]) |
| val_fraction = float(cfg["data"]["finetune_val_fraction"]) |
| out_dir = Path(cfg["paths"]["data_processed"]) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| LOG.info("build: spam_policy=%s strip_quotes=%s limit=%s seed=%d holdout=%s", policy, strip, args.limit, seed, holdout_source) |
|
|
| |
| df, labelled_meta = load_labelled(cfg, args.limit, rng) |
| df["subject"] = df["subject"].map(clean_text) |
| df["body"] = df["body"].map(clean_text) |
| df["text"] = [model_input(s, b) for s, b in zip(df["subject"].tolist(), df["body"].tolist())] |
| df["label"] = df["label"].astype(np.int64) |
| df["source"] = df["source"].astype(str) |
| bases = add_key_columns(df) |
| df["base"] = bases |
| LOG.info("labelled: cleaned + keyed %d rows in %.1fs", len(df), common.timer() - t_start) |
|
|
| short = (df["n_alnum"] < min_alnum).to_numpy() |
| drop_min_alnum = nested_counts(df.loc[short]) |
| LOG.info("drop <%d alnum: %d rows %s", min_alnum, int(short.sum()), drop_min_alnum) |
| df = df.loc[~short].reset_index(drop=True) |
|
|
| |
| is_holdout = (df["source"] == holdout_source).to_numpy() |
| is_ling = (df["source"] == "Ling").to_numpy() |
| holdout = df.loc[is_holdout].reset_index(drop=True) |
| ling_excluded = {str(k): int(v) for k, v in df.loc[is_ling].groupby("label").size().items()} |
| pool = df.loc[~is_holdout & ~is_ling].reset_index(drop=True) |
| if (pool["source"] == holdout_source).any(): |
| raise AssertionError("holdout source present in the training pool") |
| LOG.info("holdout %s: %d rows (%s); Ling excluded: %d rows %s; pool before dedup: %d", |
| holdout_source, len(holdout), split_totals(holdout), int(is_ling.sum()), ling_excluded, len(pool)) |
| holdout_internal_dups = int(holdout["k_full"].duplicated().sum()) |
|
|
| |
| |
| conflict = (pool.groupby("k_full")["label"].transform("nunique") > 1).to_numpy() |
| label_conflicts = {"groups": int(pool.loc[conflict, "k_full"].nunique()), "rows_dropped": int(conflict.sum()), |
| "by_source_label": nested_counts(pool.loc[conflict])} |
| pool = pool.loc[~conflict].reset_index(drop=True) |
| LOG.info("pool label conflicts (same k_full, both labels): dropped %d rows in %d groups %s", |
| label_conflicts["rows_dropped"], label_conflicts["groups"], label_conflicts["by_source_label"]) |
|
|
| dup_mask = pool["k_full"].duplicated(keep="first").to_numpy() |
| pool_exact_dups = nested_counts(pool.loc[dup_mask]) |
| pool = pool.loc[~dup_mask].reset_index(drop=True) |
| LOG.info("pool exact dedup (k_full keep first): removed %d, kept %d", int(dup_mask.sum()), len(pool)) |
|
|
| |
| n_pool = len(pool) |
| index = ShingleIndex(pool["base"].tolist() + holdout["base"].tolist(), nd) |
| la, lb = index.links() |
| a_hold, b_hold = la >= n_pool, lb >= n_pool |
| pool_pool = ~a_hold & ~b_hold |
| pool_hold = a_hold ^ b_hold |
| hold_hold = a_hold & b_hold |
| link_counts = {"pool_pool": int(pool_pool.sum()), "pool_holdout": int(pool_hold.sum()), "holdout_holdout": int(hold_hold.sum())} |
| n_capped, n_touching = index.n_capped_shingles, index.n_rows_touching_capped |
| LOG.info("near-dup links: %s", link_counts) |
|
|
| |
| holdout_keys = frozenset(key_union(holdout)) |
| key_hit = rows_hitting(pool, holdout_keys) |
| shingle_hit = np.zeros(n_pool, dtype=bool) |
| shingle_hit[np.where(b_hold[pool_hold], la[pool_hold], lb[pool_hold])] = True |
| contaminated = key_hit | shingle_hit |
| decontam = { |
| "by_source_label": nested_counts(pool.loc[contaminated]), |
| "by_reason": {"key_match_only": int((key_hit & ~shingle_hit).sum()), "shingle_link_only": int((shingle_hit & ~key_hit).sum()), "both": int((key_hit & shingle_hit).sum()), "total": int(contaminated.sum())}, |
| } |
| LOG.info("decontamination vs holdout: removed %d rows %s", int(contaminated.sum()), decontam["by_reason"]) |
|
|
| |
| edges = list(zip(la[pool_pool].tolist(), lb[pool_pool].tolist())) |
| edges += list(same_key_pairs(pool["k_pre"].tolist())) |
| edges += list(same_key_pairs(pool["k_sb"].tolist())) |
| edges += list(zip(la[hold_hold].tolist(), lb[hold_hold].tolist())) |
| cluster_all = union_find(n_pool + len(holdout), edges) |
| pool["cluster_id"] = cluster_all[:n_pool] |
| holdout["cluster_id"] = cluster_all[n_pool:] |
| pool = pool.loc[~contaminated].reset_index(drop=True) |
|
|
| |
| is_val, split_info = grouped_stratified_split(pool, val_fraction, rng) |
| train = pool.loc[~is_val].reset_index(drop=True) |
| val = pool.loc[is_val].reset_index(drop=True) |
| LOG.info("split: train %s val %s (%s)", split_totals(train), split_totals(val), {k: v for k, v in split_info.items() if k != "per_stratum"}) |
| del index, la, lb |
|
|
| |
| val_keys = frozenset(key_union(val)) |
| train_keys = frozenset(key_union(train)) |
| ref_index = ShingleIndex(holdout["base"].tolist() + val["base"].tolist(), nd) |
| pretrain_df, pretrain_meta, pretrain_stats = build_pretrain( |
| cfg, args.limit, rng, strip, holdout_keys, val_keys, train_keys, ref_index, len(holdout), holdout["label"].to_numpy(), |
| ) |
| del ref_index |
|
|
| |
| pretrain_keys: set[str] = set() |
| pretrain_bases: list[str] = [] |
| for text in pretrain_df["text"].tolist(): |
| keys, base = text_identity(text) |
| pretrain_keys |= keys |
| pretrain_bases.append(base[: nd.window_chars]) |
| shingle_links = recount_shingle_links(train["base"].tolist(), val["base"].tolist(), holdout["base"].tolist(), pretrain_bases, nd) |
| del pretrain_bases |
| leak = leakage_checks(train, val, holdout, pretrain_keys, len(pretrain_df), shingle_links, holdout_source) |
| leak.raise_if_failed() |
|
|
| |
| files = write_outputs(out_dir, train, val, holdout, pretrain_df) |
| elapsed = round(common.timer() - t_start, 1) |
| manifest = { |
| "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"), |
| "limit": args.limit, |
| "dev_subsample": bool(args.limit), |
| "spam_policy": policy, |
| "strip_quotes": strip, |
| "seed": seed, |
| "holdout_test_source": holdout_source, |
| "elapsed_seconds": elapsed, |
| "peak_rss_gb": peak_rss_gb(), |
| "datasets": { |
| "labelled_primary": labelled_meta, |
| "pretrain_corpus": pretrain_meta, |
| "excluded": cfg["data"]["labelled_excluded"], |
| }, |
| "counts": { |
| "splits": {"finetune_train": nested_counts(train), "finetune_val": nested_counts(val), "test_holdout": nested_counts(holdout)}, |
| "totals": {"finetune_train": split_totals(train), "finetune_val": split_totals(val), "test_holdout": split_totals(holdout), "pretrain": {"rows": int(len(pretrain_df)), "chars": int(pretrain_df["n_chars"].sum())}}, |
| }, |
| "normalization": "common.normalize_text: NFKC -> lowercase -> keep [a-z0-9] only (import it, never re-derive it; NFKC is an extension of the verify_data.md section 7 wording)", |
| "dials": dict(DIALS.__dict__), |
| "drops": { |
| "labelled_min_alnum": {"threshold": min_alnum, "rows": int(short.sum()), "by_source_label": drop_min_alnum}, |
| "ling_excluded_by_label": ling_excluded, |
| "pool_label_conflicts": label_conflicts, |
| "pool_exact_dups_by_source_label": pool_exact_dups, |
| "decontamination_vs_holdout": decontam, |
| }, |
| "holdout": {"rows": int(len(holdout)), "internal_exact_dups_kept": holdout_internal_dups}, |
| "near_dup": {**nd.__dict__, "capped_shingles": n_capped, "rows_touching_capped": n_touching, "links_at_removal_stage": link_counts}, |
| "split": split_info, |
| "pretrain": pretrain_stats, |
| "leakage": leak.checks, |
| "files": files, |
| "config_snapshot": {k: v for k, v in cfg.items() if not k.startswith("_")}, |
| } |
| common.write_json(manifest, out_dir / "splits_manifest.json") |
| (out_dir / "data_report.md").write_text(render_report(manifest), encoding="utf-8") |
| LOG.info("wrote %s and %s in %.1fs (peak RSS %.2f GB)", out_dir / "splits_manifest.json", out_dir / "data_report.md", elapsed, manifest["peak_rss_gb"]) |
| print_summary(manifest) |
|
|
|
|
| def build_pretrain( |
| cfg: dict, limit: int | None, rng: np.random.Generator, strip: bool, |
| holdout_keys: frozenset[str], val_keys: frozenset[str], train_keys: frozenset[str], |
| ref_index: ShingleIndex, n_holdout: int, holdout_labels: np.ndarray, |
| ) -> tuple[pd.DataFrame, dict, dict]: |
| """corbt -> deduped, quote-stripped MLM corpus with no exact or near-duplicate of a holdout email. |
| |
| Streaming pass: exact dedup, quote strip, length filter, key screens against holdout and val. |
| Then one near-dup query of every surviving row against `ref_index` (holdout bases first, then |
| val bases): rows linked to a holdout row are removed (verify_data.md section 7 step 8, which |
| the key screens alone miss for quoted/forwarded variants and 419 templates), rows linked to a |
| val row are counted but kept. Empty-subject rows are additionally screened under their |
| text-derived identity (split_text) so the post-hoc text-only assertion shared with `check` can |
| never see a key or link this pass did not. |
| """ |
| meta, batches = iter_pretrain_batches(cfg, limit, rng) |
| seen_raw: set[str] = set() |
| seen_stripped: set[str] = set() |
| doc_ids: list[str] = [] |
| texts: list[str] = [] |
| n_chars: list[int] = [] |
| q_bases: list[str] = [] |
| q_row: list[int] = [] |
| window = ref_index.nd.window_chars |
| st = Counter() |
| marker_hits = Counter() |
| chars_before = chars_after = chars_cut_bodies_before = chars_cut_bodies_after = 0 |
| t0 = common.timer() |
| for idx_batch, subjects, bodies in batches: |
| for row_idx, subj, body in zip(idx_batch.tolist(), subjects, bodies): |
| st["rows_in"] += 1 |
| subj = clean_text(subj) |
| body = clean_text(body) |
| raw_key = common.sha1_hex(common.normalize_text(subj) + common.normalize_text(body)) |
| if raw_key in seen_raw: |
| st["exact_dups_raw"] += 1 |
| continue |
| seen_raw.add(raw_key) |
| chars_before += len(body) |
| if strip: |
| new_body, marker, cut = strip_quotes(body) |
| if marker: |
| marker_hits[marker] += 1 |
| if cut: |
| st["bodies_cut"] += 1 |
| chars_cut_bodies_before += len(body) |
| chars_cut_bodies_after += len(new_body) |
| elif marker: |
| st["marker_but_kept_full"] += 1 |
| body = new_body |
| chars_after += len(body) |
| k_full, k_pre, k_sb, k_sb_full, base, _n = row_keys(subj, body) |
| if len(common.normalize_text(body)) < DIALS.pretrain_min_alnum: |
| st["short_after_strip"] += 1 |
| continue |
| if k_sb_full in seen_stripped: |
| st["exact_dups_after_strip"] += 1 |
| continue |
| seen_stripped.add(k_sb_full) |
| keys = {k_full, k_pre, k_sb, k_sb_full} |
| text = model_input(subj, body) |
| alt_base: str | None = None |
| if not subj: |
| d_keys, d_base = text_identity(text) |
| keys |= d_keys |
| if d_base != base: |
| alt_base = d_base |
| if keys & holdout_keys: |
| st["holdout_key_removed"] += 1 |
| continue |
| if keys & val_keys: |
| st["val_key_removed"] += 1 |
| continue |
| if keys & train_keys: |
| st["train_overlap_kept"] += 1 |
| q_bases.append(base[:window]) |
| q_row.append(len(texts)) |
| if alt_base is not None: |
| q_bases.append(alt_base[:window]) |
| q_row.append(len(texts)) |
| doc_ids.append(f"corbt:{row_idx:06d}") |
| texts.append(text) |
| n_chars.append(len(text)) |
| if st["rows_in"] % 100_000 < 20_000: |
| LOG.info(" corbt %d rows in, %d kept (%.1fs)", st["rows_in"], len(texts), common.timer() - t0) |
|
|
| |
| q, r = ref_index.query(q_bases) |
| del q_bases |
| row_of = np.asarray(q_row, dtype=np.int64)[q] if q.size else np.empty(0, dtype=np.int64) |
| hit_hold = r < n_holdout |
| drop = np.zeros(len(texts), dtype=bool) |
| drop[row_of[hit_hold]] = True |
| val_linked_rows = np.unique(row_of[~hit_hold & ~drop[row_of]]) |
| hold_rows_hit = np.unique(r[hit_hold]) |
| near_dup = { |
| "holdout_links": int(hit_hold.sum()), |
| "holdout_rows_removed": int(drop.sum()), |
| "holdout_rows_hit_by_label": {str(k): int(v) for k, v in Counter(holdout_labels[hold_rows_hit].tolist()).items()}, |
| "val_links": int((~hit_hold).sum()), |
| "val_rows_hit": int(np.unique(r[~hit_hold] - n_holdout).size), |
| "pretrain_rows_near_val_kept": int(val_linked_rows.size), |
| } |
| LOG.info("pretrain near-dup screen: %s", near_dup) |
| keep = ~drop |
| df = pd.DataFrame({ |
| "doc_id": [d for d, k in zip(doc_ids, keep) if k], |
| "text": [t for t, k in zip(texts, keep) if k], |
| "n_chars": np.asarray(n_chars, dtype=np.int64)[keep], |
| }) |
| stats = { |
| **{k: int(st[k]) for k in ("rows_in", "exact_dups_raw", "short_after_strip", "exact_dups_after_strip", "holdout_key_removed", "val_key_removed", "train_overlap_kept")}, |
| "holdout_near_dup_removed": near_dup["holdout_rows_removed"], |
| "near_dup_screen": near_dup, |
| "rows_out": int(len(df)), |
| "total_chars_out": int(df["n_chars"].sum()) if len(df) else 0, |
| "mean_chars_out": round(float(df["n_chars"].mean()), 1) if len(df) else 0.0, |
| "min_alnum_after_strip": DIALS.pretrain_min_alnum, |
| "quote_strip": { |
| "enabled": strip, |
| "bodies_cut": int(st["bodies_cut"]), |
| "marker_present_but_kept_full": int(st["marker_but_kept_full"]), |
| "marker_hits": dict(marker_hits), |
| "retained_char_fraction_all_bodies": round(chars_after / chars_before, 4) if chars_before else None, |
| "retained_char_fraction_cut_bodies": round(chars_cut_bodies_after / chars_cut_bodies_before, 4) if chars_cut_bodies_before else None, |
| "min_alnum_above_cut": DIALS.quote_strip_min_alnum, |
| }, |
| } |
| LOG.info("pretrain corpus: %s", {k: v for k, v in stats.items() if k not in ("quote_strip", "near_dup_screen")}) |
| LOG.info("quote strip: %s", stats["quote_strip"]) |
| return df, meta, stats |
|
|
|
|
| def write_outputs(out_dir: Path, train: pd.DataFrame, val: pd.DataFrame, holdout: pd.DataFrame, pretrain: pd.DataFrame) -> dict[str, str]: |
| files = {} |
| for name, df in (("finetune_train", train), ("finetune_val", val), ("test_holdout", holdout)): |
| out = df.copy() |
| out["doc_id"] = [f"primary:{i:06d}" for i in out["row_idx"].tolist()] |
| out = out[FINETUNE_SCHEMA] |
| out["label"] = out["label"].astype(np.int64) |
| out["cluster_id"] = out["cluster_id"].astype(np.int64) |
| path = out_dir / f"{name}.parquet" |
| out.to_parquet(path, index=False) |
| files[name] = str(path) |
| LOG.info("wrote %s (%d rows)", path, len(out)) |
| path = out_dir / "pretrain.parquet" |
| pretrain[PRETRAIN_SCHEMA].to_parquet(path, index=False) |
| files["pretrain"] = str(path) |
| LOG.info("wrote %s (%d rows)", path, len(pretrain)) |
| return files |
|
|
|
|
| |
| def render_report(m: dict) -> str: |
| c = m["counts"] |
| tot = c["totals"] |
| lines = [f"# Data report — {m['timestamp']}", ""] |
| if m["dev_subsample"]: |
| lines += [f"**DEV SUBSAMPLE (--limit {m['limit']}): counts below are NOT the real build.**", ""] |
| lines += [ |
| f"spam_policy `{m['spam_policy']}` (positive = phishing/spam, unwanted); quote stripping {'on' if m['strip_quotes'] else 'off'}; " |
| f"seed {m['seed']}; holdout source `{m['holdout_test_source']}`; elapsed {m['elapsed_seconds']}s; peak RSS {m['peak_rss_gb']} GB.", |
| "", |
| "## Sources", |
| f"- labelled primary: `{m['datasets']['labelled_primary']['repo']}` @ `{m['datasets']['labelled_primary']['commit_sha']}` " |
| f"({m['datasets']['labelled_primary']['rows_loaded']:,} rows, hub splits {m['datasets']['labelled_primary']['hub_splits_concatenated']} concatenated; the hub split leaks and is ignored)", |
| f"- pretraining corpus: `{m['datasets']['pretrain_corpus']['repo']}` @ `{m['datasets']['pretrain_corpus']['commit_sha']}` ({m['datasets']['pretrain_corpus']['rows_total']:,} rows)", |
| f"- excluded (no body text / 92% subset): {', '.join('`' + r + '`' for r in m['datasets']['excluded'])}", |
| "", |
| "## Splits", |
| md_table(["split", "rows", "positives", "negatives", "base rate"], |
| [[k, v["rows"], v["positives"], v["negatives"], v["base_rate"]] for k, v in tot.items() if k != "pretrain"] |
| + [["pretrain (MLM)", tot["pretrain"]["rows"], "-", "-", f"{tot['pretrain']['chars']:,} chars"]]), |
| "", |
| ] |
| for split in ("finetune_train", "finetune_val", "test_holdout"): |
| lines += [f"### {split} by source", source_label_table(c["splits"][split]), ""] |
| d = m["drops"] |
| lc = d["pool_label_conflicts"] |
| lines += [ |
| "## Normalisation and dials", |
| f"- {m['normalization']}", |
| f"- dials: {m['dials']}", |
| "", |
| "## Drops and removals", |
| f"- labelled rows with <{d['labelled_min_alnum']['threshold']} alphanumerics (subject+body): {d['labelled_min_alnum']['rows']:,} rows", |
| source_label_table(d["labelled_min_alnum"]["by_source_label"]), |
| f"- Ling excluded (academic list corpus, {sum(d['ling_excluded_by_label'].values()):,} rows by label {d['ling_excluded_by_label']})", |
| f"- pool label conflicts (identical normalised body labelled both 0 and 1; all members dropped): {lc['rows_dropped']:,} rows in {lc['groups']:,} groups", |
| source_label_table(lc["by_source_label"]), |
| f"- pool exact duplicates (k_full, keep first): {sum(sum(v.values()) for v in d['pool_exact_dups_by_source_label'].values()):,} rows", |
| source_label_table(d["pool_exact_dups_by_source_label"]), |
| f"- decontamination against the holdout: {d['decontamination_vs_holdout']['by_reason']}", |
| source_label_table(d["decontamination_vs_holdout"]["by_source_label"]), |
| f"- holdout internal exact duplicates kept (test set is never deduped): {m['holdout']['internal_exact_dups_kept']:,}", |
| "", |
| "## Near-duplicate clustering", |
| f"- {m['near_dup']['shingle_chars']}-char shingles, stride {m['near_dup']['stride']}, first {m['near_dup']['window_chars']} chars, link at >= {m['near_dup']['min_shared']} shared; " |
| f"{m['near_dup']['capped_shingles']:,} boilerplate shingles over cap {m['near_dup']['cap']:,} ignored (touching {m['near_dup']['rows_touching_capped']:,} rows)", |
| f"- links at the removal stage (pool + holdout index): {m['near_dup']['links_at_removal_stage']}", |
| f"- clusters {m['split']['clusters_total']:,} (mixed source x label: {m['split']['clusters_mixed_stratum']:,}; largest {m['split']['largest_cluster']:,}); " |
| f"val fraction target {m['split']['val_fraction_target']} achieved {m['split']['val_fraction_achieved']}", |
| md_table(["stratum", "rows", "clusters", "val target", "val rows", "val frac", "clusters too large"], |
| [[k, v["rows"], v["clusters"], v["val_target"], v["val_rows"], v["val_fraction"], v["clusters_too_large_for_val"]] for k, v in m["split"]["per_stratum"].items()]), |
| "", |
| "## Pretraining corpus (corbt/enron-emails)", |
| ] |
| p = m["pretrain"] |
| lines += [md_table(["stage", "rows"], [ |
| ["rows in", p["rows_in"]], ["exact dups (subject+body, raw)", p["exact_dups_raw"]], |
| [f"< {p['min_alnum_after_strip']} alnum after stripping", p["short_after_strip"]], |
| ["exact dups after stripping", p["exact_dups_after_strip"]], |
| ["matches a holdout key (removed)", p["holdout_key_removed"]], ["matches a fine-tune val key (removed)", p["val_key_removed"]], |
| [f">= {m['near_dup']['min_shared']} shingles shared with a holdout row (removed)", p["holdout_near_dup_removed"]], |
| [f">= {m['near_dup']['min_shared']} shingles shared with a fine-tune val row (kept, counted)", p["near_dup_screen"]["pretrain_rows_near_val_kept"]], |
| ["overlaps a fine-tune train key (kept, allowed)", p["train_overlap_kept"]], ["rows out", p["rows_out"]], |
| ])] |
| ns = p["near_dup_screen"] |
| q = p["quote_strip"] |
| lines += [ |
| f"- near-dup screen: {ns['holdout_links']:,} links to {sum(ns['holdout_rows_hit_by_label'].values()):,} holdout rows (by label {ns['holdout_rows_hit_by_label']}); " |
| f"{ns['val_links']:,} links to {ns['val_rows_hit']:,} val rows", |
| f"- quote/forward stripping {'on' if q['enabled'] else 'off'}: {q['bodies_cut']:,} bodies cut, {q['marker_present_but_kept_full']:,} had a marker but kept the full body " |
| f"(< {q['min_alnum_above_cut']} alnum above it); marker hits {q['marker_hits']}; retained chars {q['retained_char_fraction_all_bodies']} of all bodies, " |
| f"{q['retained_char_fraction_cut_bodies']} of cut bodies", |
| f"- {p['total_chars_out']:,} chars out, mean {p['mean_chars_out']} chars/doc", |
| "", |
| "## Leakage assertions", |
| md_table(["check", "count", "asserted zero", "pass"], [[k, v["count"], v["asserted"], v["pass"]] for k, v in m["leakage"].items()]), |
| "", |
| "## Files", |
| ] + [f"- `{k}`: `{v}`" for k, v in m["files"].items()] |
| return "\n".join(lines) + "\n" |
|
|
|
|
| def print_summary(m: dict) -> None: |
| print("\n=== data_prep summary ===") |
| if m["dev_subsample"]: |
| print(f"*** DEV SUBSAMPLE --limit {m['limit']} — not a real build ***") |
| for k, v in m["counts"]["totals"].items(): |
| print(f"{k:>16}: {v}") |
| print("leakage:") |
| for k, v in m["leakage"].items(): |
| print(f" {k:<44} {v['count']:>8} {'ASSERT' if v['asserted'] else 'info '} {'ok' if v['pass'] else 'FAIL'}") |
| print(f"pretrain: {m['pretrain']['rows_out']:,} rows, quote strip {m['pretrain']['quote_strip']['bodies_cut']:,} cut") |
|
|
|
|
| |
| def cmd_check(args: argparse.Namespace, cfg: dict) -> None: |
| out_dir = Path(cfg["paths"]["data_processed"]) |
| manifest = common.read_json(out_dir / "splits_manifest.json") |
| nd = NearDupConfig.from_config(cfg, cap=int(manifest["near_dup"]["cap"])) |
| frames = {} |
| for name in ("finetune_train", "finetune_val", "test_holdout"): |
| df = pd.read_parquet(out_dir / f"{name}.parquet") |
| if list(df.columns) not in (FINETUNE_SCHEMA, THREEWAY_SCHEMA): |
| raise AssertionError(f"{name}: unexpected columns {list(df.columns)}") |
| df = df[FINETUNE_SCHEMA] |
| df["subject"] = df["subject"].fillna("") |
| df["body"] = df["body"].fillna("") |
| df["base"] = add_key_columns(df) |
| frames[name] = df |
| LOG.info("read %s: %d rows", name, len(df)) |
| pre = pd.read_parquet(out_dir / "pretrain.parquet", columns=["text"]) |
| if list(pre.columns) != ["text"]: |
| raise AssertionError("pretrain.parquet: missing text column") |
| pretrain_keys: set[str] = set() |
| pretrain_bases: list[str] = [] |
| for t in pre["text"].tolist(): |
| keys, base = text_identity(t) |
| pretrain_keys |= keys |
| pretrain_bases.append(base[: nd.window_chars]) |
| LOG.info("read pretrain: %d rows", len(pre)) |
|
|
| train, val, hold = frames["finetune_train"], frames["finetune_val"], frames["test_holdout"] |
| holdout_source = cfg["data"]["holdout_test_source"] |
| shingle_links = recount_shingle_links(train["base"].tolist(), val["base"].tolist(), hold["base"].tolist(), pretrain_bases, nd) |
| leak = leakage_checks(train, val, hold, pretrain_keys, len(pre), shingle_links, holdout_source) |
| print_summary({**manifest, "leakage": leak.checks}) |
| print("\nmanifest timestamp:", manifest["timestamp"], "| limit:", manifest["limit"], "| spam_policy:", manifest["spam_policy"]) |
| print("build-time leakage results:", {k: v["count"] for k, v in manifest["leakage"].items()}) |
| leak.raise_if_failed() |
| print("ALL LEAKAGE ASSERTIONS PASS") |
|
|
|
|
| |
| def cmd_inspect(args: argparse.Namespace, cfg: dict) -> None: |
| from datasets import get_dataset_split_names, load_dataset |
|
|
| cache_dir = str(Path(cfg["paths"]["data_raw"]) / "hf_cache") |
| repos = [cfg["data"]["labelled_primary"], cfg["data"]["pretrain_corpus"], *cfg["data"]["labelled_excluded"]] |
| for repo in repos: |
| print(f"\n=== {repo} ===") |
| sha = hub_sha(repo) |
| try: |
| splits = get_dataset_split_names(repo, revision=sha) |
| ds = load_dataset(repo, split=splits[0], streaming=True, cache_dir=cache_dir, revision=sha) |
| print(f"commit: {sha} splits: {splits} (showing '{splits[0]}')") |
| rows = list(ds.take(2)) |
| if ds.features is not None: |
| print("features:") |
| for name, feat in ds.features.items(): |
| print(f" {name}: {feat}") |
| else: |
| print("features: not declared by the builder; python types of row 0:") |
| for name, value in (rows[0].items() if rows else []): |
| print(f" {name}: {type(value).__name__}") |
| for i, row in enumerate(rows): |
| print(f"row {i}:") |
| for k, v in row.items(): |
| s = repr(v) |
| print(f" {k}: {s[:200]}{'...' if len(s) > 200 else ''}") |
| except Exception as exc: |
| print(f"ERROR inspecting {repo}: {type(exc).__name__}: {exc}") |
| role = {cfg["data"]["labelled_primary"]: "labelled primary (own split built here)", cfg["data"]["pretrain_corpus"]: "MLM pretraining"} |
| print("\nroles:", {r: role.get(r, "EXCLUDED from training/eval") for r in repos}) |
|
|
|
|
| |
| |
| |
| LABEL3_OF_SET: dict[str, int] = {"legit": 0, "phish": 1, "spam": 2, "unwanted": -1} |
| SET_OF_LABEL3: dict[int, str] = {v: k for k, v in LABEL3_OF_SET.items()} |
| LABEL3_NAMES: dict[int, str] = {0: "Legitimate", 1: "Phishing", 2: "Spam"} |
| UNWANTED_NAME = "Unwanted (phishing or spam)" |
| THREEWAY_SCHEMA = FINETUNE_SCHEMA + ["label3", "label_set"] |
| THREEWAY_SPLITS = ("finetune_train", "finetune_val", "test_holdout", "test_holdout3") |
| BINARY_SNAPSHOT_DIR = "binary_snapshot" |
| BIGSHP_SOURCE = "bigshp" |
| PART_ORDER = ("holdout", "test3", "train", "val") |
| MBOX_INTERNAL_SUBJECT = "DON'T DELETE THIS MESSAGE -- FOLDER INTERNAL DATA" |
| _WORD = re.compile(r"[a-z]{2,}") |
| _PUNCT = re.compile(r"[^\w\s]") |
|
|
| THREEWAY_DEFAULTS: dict[str, Any] = { |
| |
| |
| |
| "positive_label3_by_source": { |
| "Assassin": 2, "Enron": 2, "Ling": 2, |
| "TREC-05": -1, "TREC-06": -1, "TREC-07": -1, "CEAS-08": -1, |
| "Nazario": 1, "Nigerian_Fraud": 1, |
| }, |
| "champa": { |
| "repo": "kudzaiprichard/aura-phishing-email-corpus", |
| "files": ["Nazario.csv", "Nigerian_Fraud.csv"], |
| "local_dir": "data/raw/aura", |
| "label3": 1, |
| }, |
| "bigshp": { |
| "enabled": False, |
| "verdict": "DO_NOT_USE", |
| "repo": "locuoco/the-biggest-spam-ham-phish-email-dataset-300000", |
| "file": "df.csv", |
| "local_dir": "data/raw/bigshp", |
| "drop_phish_class": True, |
| "min_chars": 160, |
| "min_english_fraction": 0.7, |
| "near_dup_token_key_len": 30, |
| }, |
| "new_source_test_fraction": 0.15, |
| "new_source_val_fraction": 0.10, |
| "drop_new_clusters_linked_to_both_train_and_val": True, |
| "holdout3_from_val": { |
| "legit_sources": ["CEAS-08", "TREC-07", "Assassin"], |
| "spam_sources": ["Assassin", "Enron"], |
| "legit_rows": 2000, |
| "spam_rows": 700, |
| }, |
| } |
|
|
|
|
| def threeway_config(cfg: dict) -> dict[str, Any]: |
| """config `data.threeway` merged over THREEWAY_DEFAULTS one level deep. |
| |
| Every decision (which source is spam vs unwanted, whether bigshp is ingested, the holdout3 targets) is |
| therefore visible in config.yaml and reversible, while a missing key never silently changes a decision. |
| """ |
| user = (cfg.get("data") or {}).get("threeway") or {} |
| out: dict[str, Any] = {} |
| for key, default in THREEWAY_DEFAULTS.items(): |
| val = user.get(key, default) |
| if isinstance(default, dict) and isinstance(val, dict): |
| merged = dict(default) |
| merged.update(val) |
| val = merged |
| out[key] = val |
| return out |
|
|
|
|
| def label3_from_provenance(source: np.ndarray, label: np.ndarray, pos_map: dict[str, int]) -> np.ndarray: |
| """label3 of every row from (source, binary label): label 0 -> 0, label 1 -> pos_map[source]. |
| |
| Raises on a positive row whose source has no entry: silently defaulting it to "unwanted" would hide a |
| provenance decision that belongs in config.yaml. |
| """ |
| label = np.asarray(label, dtype=np.int64) |
| source = np.asarray(source).astype(str) |
| pos = label == 1 |
| unknown = sorted(set(source[pos].tolist()) - set(pos_map)) |
| if unknown: |
| raise ValueError(f"no data.threeway.positive_label3_by_source entry for positive rows of source(s) {unknown}") |
| out = np.zeros(len(label), dtype=np.int64) |
| for s in set(source[pos].tolist()): |
| out[pos & (source == s)] = int(pos_map[s]) |
| return out |
|
|
|
|
| def add_label3_columns(df: pd.DataFrame, label3: np.ndarray) -> None: |
| df["label3"] = np.asarray(label3, dtype=np.int64) |
| df["label_set"] = [SET_OF_LABEL3[int(v)] for v in df["label3"].tolist()] |
|
|
|
|
| def contract_violations(df: pd.DataFrame, pos_map: dict[str, int]) -> dict[str, int]: |
| """Per-rule counts of rows breaking the three-way contract (all zero on a correct file). |
| |
| label_set must be a known set, label3 must equal LABEL3_OF_SET[label_set], the binary label must be |
| (label3 != 0), and for every source except bigshp (which carries its own label) label3 must equal the |
| provenance mapping. check3 asserts the sum is zero. |
| """ |
| if len(df) == 0: |
| return {"label_set_unknown": 0, "label3_vs_label_set": 0, "label_vs_label3": 0, "label3_vs_provenance": 0} |
| label3 = df["label3"].to_numpy().astype(np.int64) |
| lset = df["label_set"].astype(str).to_numpy() |
| label = df["label"].to_numpy().astype(np.int64) |
| expect3 = np.fromiter((LABEL3_OF_SET.get(s, -99) for s in lset.tolist()), dtype=np.int64, count=len(lset)) |
| out = { |
| "label_set_unknown": int((expect3 == -99).sum()), |
| "label3_vs_label_set": int((expect3 != label3).sum()), |
| "label_vs_label3": int(((label3 != 0).astype(np.int64) != label).sum()), |
| } |
| own = df["source"].astype(str).to_numpy() == BIGSHP_SOURCE |
| try: |
| prov = label3_from_provenance(df["source"].to_numpy()[~own], label[~own], pos_map) |
| out["label3_vs_provenance"] = int((prov != label3[~own]).sum()) |
| except ValueError as exc: |
| LOG.error("provenance mapping: %s", exc) |
| out["label3_vs_provenance"] = int((~own).sum()) |
| out["bigshp_label3_outside_0_1_2"] = int((~np.isin(label3[own], [0, 1, 2])).sum()) |
| return out |
|
|
|
|
| def nested_counts3(df: pd.DataFrame) -> dict[str, dict[str, int]]: |
| """{source: {label3: n}} with string keys for JSON.""" |
| out: dict[str, dict[str, int]] = {} |
| if len(df) == 0: |
| return out |
| for (source, l3), n in df.groupby(["source", "label3"], sort=True).size().items(): |
| out.setdefault(str(source), {})[str(int(l3))] = int(n) |
| return out |
|
|
|
|
| def split_totals3(df: pd.DataFrame) -> dict[str, Any]: |
| n = int(len(df)) |
| by3 = {str(k): int(v) for k, v in df["label3"].value_counts().sort_index().items()} if n else {} |
| byset = {str(k): int(v) for k, v in df["label_set"].value_counts().sort_index().items()} if n else {} |
| pos = int((df["label"] == 1).sum()) if n else 0 |
| return {"rows": n, "by_label3": by3, "by_label_set": byset, "binary_positives": pos, |
| "binary_base_rate": round(pos / n, 4) if n else None, |
| "rows_with_full_label3": int(df["label3"].isin([0, 1, 2]).sum()) if n else 0} |
|
|
|
|
| def source_label3_table(counts: dict[str, dict[str, int]]) -> str: |
| rows = [] |
| for source in sorted(counts): |
| c = counts[source] |
| rows.append([source, c.get("0", 0), c.get("1", 0), c.get("2", 0), c.get("-1", 0), sum(c.values())]) |
| return md_table(["source", "0 Legitimate", "1 Phishing", "2 Spam", "-1 unwanted (unsplit)", "rows"], rows) |
|
|
|
|
| def first_reason(n: int, reasons: list[tuple[str, np.ndarray]]) -> np.ndarray: |
| """Per-row first matching drop reason ("" = kept), so every drop is counted exactly once.""" |
| out = np.full(n, "", dtype=object) |
| for name, mask in reasons: |
| hit = np.asarray(mask, dtype=bool) & (out == "") |
| out[hit] = name |
| return out |
|
|
|
|
| def by_source_reason(source: np.ndarray, reason: np.ndarray) -> dict[str, dict[str, int]]: |
| out: dict[str, dict[str, int]] = {} |
| for (s, r), n in Counter(zip(source.tolist(), reason.tolist())).items(): |
| if r: |
| out.setdefault(str(s), {})[str(r)] = int(n) |
| return out |
|
|
|
|
| def write_table_atomic(table: "pa.Table", path: Path) -> None: |
| """tmp + os.replace: other tools (finetune, evaluate, tokenizer stats) read these files while we write.""" |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_name(path.name + ".tmp") |
| pq.write_table(table, tmp) |
| os.replace(tmp, path) |
|
|
|
|
| def write_text_atomic(text: str, path: Path) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_name(path.name + ".tmp") |
| tmp.write_text(text, encoding="utf-8") |
| os.replace(tmp, path) |
|
|
|
|
| def copy_atomic(src: Path, dst: Path) -> None: |
| dst.parent.mkdir(parents=True, exist_ok=True) |
| tmp = dst.with_name(dst.name + ".tmp") |
| shutil.copy2(src, tmp) |
| os.replace(tmp, dst) |
|
|
|
|
| |
| def binary_base_dir(processed_dir: Path, dev: bool) -> tuple[Path, str]: |
| """Directory holding the BINARY (pre-build3) splits that build3 derives from. |
| |
| build3 rewrites finetune_{train,val}.parquet and test_holdout.parquet in place, so a second run could not |
| start from the same inputs. The first non-dev run therefore snapshots the binary files into |
| <data_processed>/binary_snapshot/ (atomic copies) and every later run derives from that snapshot: build3 is |
| idempotent and the binary split membership stays byte-for-byte recoverable. If the live files already carry |
| label3 and there is no snapshot, the binary originals are gone and `build` must be re-run first. |
| """ |
| names = ("finetune_train", "finetune_val", "test_holdout") |
| snap = processed_dir / BINARY_SNAPSHOT_DIR |
| if all((snap / f"{n}.parquet").exists() for n in names): |
| return snap, "snapshot" |
| for n in names: |
| live = processed_dir / f"{n}.parquet" |
| if not live.exists(): |
| raise FileNotFoundError(f"{live} missing: run `data_prep.py build` first") |
| if "label3" in pq.read_schema(live).names: |
| raise RuntimeError(f"{live} already carries label3 but {snap} has no snapshot of the binary files; " |
| "re-run `data_prep.py build` to recreate the binary splits before build3") |
| if dev: |
| return processed_dir, "live (dev run: no snapshot taken)" |
| for n in names: |
| copy_atomic(processed_dir / f"{n}.parquet", snap / f"{n}.parquet") |
| LOG.info("snapshotted the binary splits to %s", snap) |
| return snap, "snapshot (created by this run)" |
|
|
|
|
| def per_source_sample(source: np.ndarray, limit: int, rng: np.random.Generator) -> np.ndarray: |
| """Sorted row indices: at most `limit` rows per source (dev subsample).""" |
| keep = [] |
| for s in np.unique(source): |
| idx = np.flatnonzero(source == s) |
| keep.append(rng.choice(idx, min(limit, idx.size), replace=False)) |
| return np.sort(np.concatenate(keep)) if keep else np.empty(0, dtype=np.int64) |
|
|
|
|
| def read_binary_split(path: Path, limit: int | None, rng: np.random.Generator) -> tuple["pa.Table", pd.DataFrame]: |
| """(arrow table, keyed frame) of a binary split; the two stay row-aligned so the table can be rewritten in place.""" |
| table = pq.read_table(path) |
| if list(table.schema.names) != FINETUNE_SCHEMA: |
| raise AssertionError(f"{path}: expected binary columns {FINETUNE_SCHEMA}, got {list(table.schema.names)}") |
| if limit: |
| idx = per_source_sample(np.asarray(table.column("source").to_pylist(), dtype=object), limit, rng) |
| table = table.take(pa.array(idx)) |
| df = table.select(["doc_id", "subject", "body", "label", "source", "cluster_id"]).to_pandas() |
| df["subject"] = df["subject"].fillna("") |
| df["body"] = df["body"].fillna("") |
| df["label"] = df["label"].astype(np.int64) |
| df["source"] = df["source"].astype(str) |
| df["base"] = add_key_columns(df) |
| LOG.info("read %s: %d rows%s", path.name, len(df), f" (dev subsample, <= {limit}/source)" if limit else "") |
| return table, df |
|
|
|
|
| def pretrain_identity(path: Path, nd: NearDupConfig, limit: int | None, rng: np.random.Generator |
| ) -> tuple[set[str], list[str], int, "pa.Table | None"]: |
| """(keys, window-trimmed shingle bases, n_rows, dev table) of the pretraining corpus. |
| |
| Re-derived from the stored `text` with text_identity exactly as `check` does, in 20k-row batches so the |
| 408M-char corpus is never fully resident. In dev mode a 4*limit subsample is used AND returned so the dev |
| directory gets its own pretrain.parquet: check3 must screen against the same corpus build3 screened against. |
| """ |
| keys: set[str] = set() |
| bases: list[str] = [] |
| dev_table = None |
| if limit: |
| full = pq.read_table(path) |
| n = min(limit * 4, full.num_rows) |
| idx = np.sort(rng.choice(full.num_rows, n, replace=False)) |
| dev_table = full.take(pa.array(idx)) |
| del full |
| batches = dev_table.select(["text"]).to_batches(max_chunksize=20_000) |
| LOG.warning("--limit %d: pretrain subsample %d rows", limit, n) |
| else: |
| batches = pq.ParquetFile(path).iter_batches(batch_size=20_000, columns=["text"]) |
| n_rows = 0 |
| t0 = common.timer() |
| for batch in batches: |
| for text in batch.column(0).to_pylist(): |
| k, base = text_identity(text or "") |
| keys |= k |
| bases.append(base[: nd.window_chars]) |
| n_rows += 1 |
| LOG.info("pretrain identity: %d rows, %d keys (%.1fs)", n_rows, len(keys), common.timer() - t0) |
| return keys, bases, n_rows, dev_table |
|
|
|
|
| def champa_commit_sha(local_dir: Path, fname: str, repo: str) -> str | None: |
| """Commit sha recorded by hf_hub_download in <local_dir>/.cache/huggingface/download/<file>.metadata (offline), else the Hub.""" |
| meta = local_dir / ".cache" / "huggingface" / "download" / f"{fname}.metadata" |
| if meta.exists(): |
| first = meta.read_text(encoding="utf-8").splitlines()[0].strip() |
| if re.fullmatch(r"[0-9a-f]{40}", first): |
| return first |
| return hub_sha(repo) |
|
|
|
|
| def load_champa(tw: dict, limit: int | None, rng: np.random.Generator) -> tuple[pd.DataFrame, dict]: |
| """Nazario.csv + Nigerian_Fraud.csv from the Champa et al. mirror -> raw frame (subject, body, label, source, label3). |
| |
| Both files are phishing corpora by construction (all label 1 in the file, label3 = 1 by provenance). The |
| sender/receiver/date/urls columns are never used: model input is subject + body only, like the primary set. |
| Downloaded with hf_hub_download when missing; the commit sha is recorded per file. |
| """ |
| c = tw["champa"] |
| local_dir = Path(c["local_dir"]) |
| if not local_dir.is_absolute(): |
| local_dir = common.project_root() / local_dir |
| meta: dict[str, Any] = {"repo": c["repo"], "local_dir": str(local_dir), "label3": int(c["label3"]), "files": {}} |
| frames = [] |
| for fname in c["files"]: |
| path = local_dir / fname |
| if not path.exists(): |
| from huggingface_hub import hf_hub_download |
|
|
| LOG.info("downloading %s/%s -> %s", c["repo"], fname, local_dir) |
| hf_hub_download(c["repo"], fname, repo_type="dataset", local_dir=str(local_dir)) |
| sha = champa_commit_sha(local_dir, fname, c["repo"]) |
| raw = pd.read_csv(path, dtype=str, keep_default_na=False) |
| missing = {"subject", "body", "label"} - set(raw.columns) |
| if missing: |
| raise ValueError(f"{path}: missing columns {sorted(missing)}") |
| labels = raw["label"].astype(np.int64).to_numpy() |
| if (labels != 1).any(): |
| raise ValueError(f"{path}: expected every label == 1 (phishing corpus), found {Counter(labels.tolist())}") |
| stem = Path(fname).stem |
| df = pd.DataFrame({ |
| "doc_id": [f"aura:{stem}:{i:06d}" for i in range(len(raw))], |
| "subject": raw["subject"].astype(str), "body": raw["body"].astype(str), |
| "label": np.ones(len(raw), dtype=np.int64), "source": stem, |
| "label3": np.full(len(raw), int(c["label3"]), dtype=np.int64), |
| }) |
| if limit: |
| df = df.iloc[np.sort(rng.choice(len(df), min(limit, len(df)), replace=False))].reset_index(drop=True) |
| meta["files"][fname] = {"rows": int(len(raw)), "rows_used": int(len(df)), "commit_sha": sha, "path": str(path), |
| "bytes": int(path.stat().st_size)} |
| LOG.info("champa %s @ %s: %d rows (%d used), %d null subjects", fname, sha, len(raw), len(df), int((raw["subject"] == "").sum())) |
| frames.append(df) |
| return pd.concat(frames, ignore_index=True), meta |
|
|
|
|
| def english_vocab(train: pd.DataFrame, rng: np.random.Generator, sample: int = 20_000, min_count: int = 3) -> set[str]: |
| """Lower-case word types seen >= min_count times in a sample of legitimate training bodies (bigshp filter).""" |
| legit = np.flatnonzero((train["label"] == 0).to_numpy()) |
| idx = rng.choice(legit, min(sample, legit.size), replace=False) |
| counts: Counter = Counter() |
| for body in train["body"].to_numpy()[idx].tolist(): |
| counts.update(_WORD.findall(body.lower())) |
| return {w for w, n in counts.items() if n >= min_count} |
|
|
|
|
| def load_bigshp(tw: dict, limit: int | None, rng: np.random.Generator, vocab: set[str]) -> tuple[pd.DataFrame, dict]: |
| """locuoco/...-300000 df.csv behind data.threeway.bigshp.enabled, with the research verdict's filter chain. |
| |
| Research verdict (docs/research/bigshp_research.md): DO_NOT_USE. The file is a union of lossy re-packagings of |
| corpora we already hold; its phish label is source-derived (87% of resolvable phish rows are spam-folder or |
| spam-track mail) and the whole phish class went through a stopword-stripping pipeline that would act as a label |
| shortcut. This loader exists so the decision is reversible from config.yaml. Filters, in order: |
| drop the label-1 class; drop rows shorter than min_chars; label-2 rows with any non-ASCII byte or < min_english_fraction |
| of word tokens in our legitimate-mail vocabulary; label-0 rows that look like Ling-Spam list mail (first token |
| "subject", no punctuation); within-file dedup on k_full, k_pre and the first-N-token key. The cross-set screens |
| (keys and shingles against train/val/holdout/pretrain) then apply exactly as for the Champa rows; a surviving |
| row whose key matches one of our rows is dropped there, with label contradictions counted separately. |
| NOTE: no digit-folding adapter ("escapenumber") is applied to our keys, so exact-key recall on bigshp text is a |
| lower bound and the shingle screen carries the near-duplicate removal. |
| """ |
| import pyarrow.csv as pacsv |
|
|
| b = tw["bigshp"] |
| local_dir = Path(b["local_dir"]) |
| if not local_dir.is_absolute(): |
| local_dir = common.project_root() / local_dir |
| path = local_dir / b["file"] |
| if not path.exists(): |
| from huggingface_hub import hf_hub_download |
|
|
| LOG.info("downloading %s/%s -> %s", b["repo"], b["file"], local_dir) |
| hf_hub_download(b["repo"], b["file"], repo_type="dataset", local_dir=str(local_dir)) |
| sha = champa_commit_sha(local_dir, b["file"], b["repo"]) |
| table = pacsv.read_csv( |
| path, |
| parse_options=pacsv.ParseOptions(newlines_in_values=True), |
| read_options=pacsv.ReadOptions(block_size=64 << 20), |
| convert_options=pacsv.ConvertOptions(column_types={"label": pa.int64(), "text": pa.large_string()}), |
| ) |
| n_total = table.num_rows |
| if limit: |
| idx = np.sort(rng.choice(n_total, min(limit * 4, n_total), replace=False)) |
| table = table.take(pa.array(idx)) |
| LOG.warning("--limit %d: bigshp subsample %d rows", limit, table.num_rows) |
| row_idx = idx if limit else np.arange(n_total) |
| label3 = np.asarray(table.column("label").to_pylist(), dtype=np.int64) |
| texts = [t or "" for t in table.column("text").to_pylist()] |
| del table |
| n = len(texts) |
| lengths = np.fromiter((len(t) for t in texts), dtype=np.int64, count=n) |
| non_ascii = np.fromiter((not t.isascii() for t in texts), dtype=bool, count=n) |
| eng = np.ones(n, dtype=np.float64) |
| ling_like = np.zeros(n, dtype=bool) |
| tok_keys: list[str] = [] |
| k = int(b["near_dup_token_key_len"]) |
| for i, t in enumerate(texts): |
| words = _WORD.findall(t.lower()) |
| if words: |
| eng[i] = sum(1 for w in words if w in vocab) / len(words) |
| tok_keys.append(common.sha1_hex(" ".join(words[:k]))) |
| if label3[i] == 0: |
| first = t.strip().split(None, 1)[0].lower() if t.strip() else "" |
| ling_like[i] = first == "subject" and _PUNCT.search(t) is None |
| keyed = pd.DataFrame({"subject": [""] * n, "body": texts}) |
| add_key_columns(keyed) |
| reasons = [ |
| ("bigshp_phish_class_source_derived", (label3 == 1) if b["drop_phish_class"] else np.zeros(n, dtype=bool)), |
| ("bigshp_label_outside_0_1_2", ~np.isin(label3, [0, 1, 2])), |
| (f"bigshp_shorter_than_{int(b['min_chars'])}_chars", lengths < int(b["min_chars"])), |
| ("bigshp_spam_non_ascii", (label3 == 2) & non_ascii), |
| (f"bigshp_spam_english_fraction_lt_{b['min_english_fraction']}", (label3 == 2) & (eng < float(b["min_english_fraction"]))), |
| ("bigshp_ham_ling_list_mail", (label3 == 0) & ling_like), |
| ("bigshp_within_file_k_full_dup", keyed["k_full"].duplicated(keep="first").to_numpy()), |
| ("bigshp_within_file_k_pre_dup", keyed["k_pre"].duplicated(keep="first").to_numpy()), |
| (f"bigshp_within_file_first{k}_token_dup", pd.Series(tok_keys).duplicated(keep="first").to_numpy()), |
| ] |
| reason = first_reason(n, reasons) |
| keep = reason == "" |
| df = pd.DataFrame({ |
| "doc_id": [f"bigshp:{int(i):06d}" for i in row_idx[keep].tolist()], |
| "subject": "", "body": [t for t, kp in zip(texts, keep) if kp], |
| "label": (label3[keep] != 0).astype(np.int64), "source": BIGSHP_SOURCE, "label3": label3[keep], |
| }) |
| meta = { |
| "repo": b["repo"], "file": b["file"], "commit_sha": sha, "path": str(path), "bytes": int(path.stat().st_size), |
| "rows_total": int(n_total), "rows_read": int(n), "rows_after_file_filters": int(len(df)), |
| "file_filter_drops": {name: int(v) for name, v in Counter(reason[~keep].tolist()).items()}, |
| "filters": {key: b[key] for key in ("drop_phish_class", "min_chars", "min_english_fraction", "near_dup_token_key_len")}, |
| "label_counts_read": {str(int(c)): int(v) for c, v in Counter(label3.tolist()).items()}, |
| } |
| LOG.info("bigshp @ %s: %d rows read, %d survive the file filters %s", sha, n, len(df), meta["file_filter_drops"]) |
| return df, meta |
|
|
|
|
| def prepare_new_rows(df: pd.DataFrame) -> None: |
| """Clean exactly like the primary set: clean_text, model_input, keys/base, label3 -> label_set.""" |
| df["subject"] = df["subject"].map(clean_text) |
| df["body"] = df["body"].map(clean_text) |
| df["text"] = [model_input(s, b) for s, b in zip(df["subject"].tolist(), df["body"].tolist())] |
| df["label"] = df["label"].astype(np.int64) |
| df["source"] = df["source"].astype(str) |
| df["base"] = add_key_columns(df) |
| add_label3_columns(df, df["label3"].to_numpy()) |
|
|
|
|
| |
| def recount_links_multi(parts: dict[str, list[str]], pretrain_bases: list[str], pretrain_ref: tuple[str, ...], nd: NearDupConfig) -> dict[str, int]: |
| """Independent near-dup link counts between every pair of final sets, from a fresh index. |
| |
| Generalises recount_shingle_links to four parts (holdout, test3, train, val): one index over all parts |
| enumerates every cross-part link; the pretraining corpus is queried against an index of the small reference |
| parts so pretrain-pretrain links are never materialised. Pair names follow PART_ORDER (holdout_train, |
| test3_val, ...), so the binary gate's keys (holdout_train, holdout_val, train_val, holdout_pretrain, |
| val_pretrain) come out unchanged. |
| """ |
| names = [p for p in PART_ORDER if p in parts] |
| bounds = np.cumsum([0] + [len(parts[n]) for n in names]) |
| index = ShingleIndex([b for n in names for b in parts[n]], nd) |
| la, lb = index.links() |
| pa_ = np.searchsorted(bounds, la, side="right") - 1 |
| pb_ = np.searchsorted(bounds, lb, side="right") - 1 |
| pair = Counter(zip(np.minimum(pa_, pb_).tolist(), np.maximum(pa_, pb_).tolist())) |
| out = {f"{names[i]}_{names[j]}": int(pair[(i, j)]) for i in range(len(names)) for j in range(i + 1, len(names))} |
| del index, la, lb |
| ref_names = [n for n in pretrain_ref if n in parts] |
| ref = ShingleIndex([b for n in ref_names for b in parts[n]], nd) |
| _q, r = ref.query(pretrain_bases) |
| rb = np.cumsum([0] + [len(parts[n]) for n in ref_names]) |
| part_of = np.searchsorted(rb, r, side="right") - 1 |
| for i, n in enumerate(ref_names): |
| out[f"{n}_pretrain"] = int((part_of == i).sum()) |
| LOG.info("recount shingle links (final sets): %s", out) |
| return out |
|
|
|
|
| def leakage_checks3(train: pd.DataFrame, val: pd.DataFrame, holdout: pd.DataFrame, test3: pd.DataFrame, |
| pretrain_keys: set[str], n_pretrain: int, links: dict[str, int], holdout_source: str, |
| pos_map: dict[str, int]) -> LeakageResults: |
| """The binary gate (unchanged) plus the test_holdout3 assertions. |
| |
| test_holdout3 must be text the model never saw in any form: zero exact keys and zero >=3-shingle links against |
| train, val and the MLM corpus, zero keys shared with the TREC-06 holdout (its shingle links to TREC-06 are |
| counted only: the two test sets may legitimately contain the same 1998-2008 spam templates), every row fully |
| labelled (label3 in {0,1,2}), no TREC-06 row, no doc_id in two splits, no contract violation anywhere. |
| """ |
| res = leakage_checks(train, val, holdout, pretrain_keys, n_pretrain, links, holdout_source) |
| t3_keys = key_union(test3) |
| hold_keys = key_union(holdout) |
| res.add("holdout3_vs_train_key_rows", int(rows_hitting(train, t3_keys).sum()), True) |
| res.add("holdout3_vs_val_key_rows", int(rows_hitting(val, t3_keys).sum()), True) |
| res.add("holdout3_vs_pretrain_key_rows", len(pretrain_keys & t3_keys), True, "pretrain keys re-derived from stored text (text_identity)") |
| res.add("holdout3_vs_holdout_key_rows", int(rows_hitting(test3, hold_keys).sum()), True) |
| res.add("holdout3_vs_train_shingle_links", links["test3_train"], True) |
| res.add("holdout3_vs_val_shingle_links", links["test3_val"], True) |
| res.add("holdout3_vs_pretrain_shingle_links", links["test3_pretrain"], True, "pretrain bases re-derived from stored text (text_identity)") |
| res.add("holdout3_vs_holdout_shingle_links", links["holdout_test3"], False, "two test sets; shared 1998-2008 spam templates are allowed") |
| res.add("holdout3_rows_without_full_label3", int((~test3["label3"].isin([0, 1, 2])).sum()), True) |
| res.add("holdout3_rows_from_holdout_source", int((test3["source"] == holdout_source).sum()), True) |
| ids = [set(df["doc_id"].tolist()) for df in (train, val, holdout, test3)] |
| res.add("doc_id_shared_between_splits", sum(len(ids[i] & ids[j]) for i in range(4) for j in range(i + 1, 4)), True) |
| viol = {name: contract_violations(df, pos_map) for name, df in zip(THREEWAY_SPLITS, (train, val, holdout, test3))} |
| res.add("label3_contract_violations", sum(sum(v.values()) for v in viol.values()), True, str({k: v for k, v in viol.items() if sum(v.values())})) |
| res.add("holdout3_rows", len(test3), False) |
| return res |
|
|
|
|
| |
| def cmd_build3(args: argparse.Namespace, cfg: dict) -> None: |
| """Derive the three-way files from the binary splits + the Champa phishing corpora (+ bigshp if enabled). |
| |
| Binary split membership is preserved except for the rows moved from finetune_val into test_holdout3; every |
| other existing row only gains label3/label_set. New rows are cleaned, keyed, deduplicated within themselves |
| and against train/val/holdout/pretrain (exact keys), screened for >=3-shingle links against the TREC-06 |
| holdout (dropped), pinned to train/val when they near-duplicate an existing train/pretrain/val row (so no |
| template straddles a split), and the free clusters are split test3 / val / train by grouped_stratified_split. |
| """ |
| t_start = common.timer() |
| seed = int(cfg["seed"]) |
| common.set_seed(seed) |
| rng = np.random.default_rng(seed + 3) |
| tw = threeway_config(cfg) |
| if args.bigshp: |
| tw["bigshp"]["enabled"] = args.bigshp == "on" |
| pos_map = {str(k): int(v) for k, v in tw["positive_label3_by_source"].items()} |
| nd = NearDupConfig.from_config(cfg, cap=args.shingle_cap or DIALS.shingle_cap) |
| holdout_source = cfg["data"]["holdout_test_source"] |
| min_alnum = int(cfg["data"]["min_alnum_chars"]) |
| processed_dir = Path(cfg["paths"]["data_processed"]) |
| dev = bool(args.limit) |
| out_dir = Path(args.out_dir) if args.out_dir else (processed_dir / f"dev_threeway_limit{args.limit}" if dev else processed_dir) |
| out_dir.mkdir(parents=True, exist_ok=True) |
| LOG.info("build3: limit=%s seed=%d out_dir=%s bigshp.enabled=%s holdout=%s", args.limit, seed, out_dir, tw["bigshp"]["enabled"], holdout_source) |
|
|
| |
| base_dir, base_origin = binary_base_dir(processed_dir, dev) |
| tables: dict[str, pa.Table] = {} |
| frames: dict[str, pd.DataFrame] = {} |
| for name in ("finetune_train", "finetune_val", "test_holdout"): |
| tables[name], frames[name] = read_binary_split(base_dir / f"{name}.parquet", args.limit, rng) |
| train, val, holdout = frames["finetune_train"], frames["finetune_val"], frames["test_holdout"] |
| for df in (train, val, holdout): |
| add_label3_columns(df, label3_from_provenance(df["source"].to_numpy(), df["label"].to_numpy(), pos_map)) |
| base_counts = {n: nested_counts3(frames[n]) for n in frames} |
| pretrain_keys, pretrain_bases, n_pretrain, pretrain_dev_table = pretrain_identity(processed_dir / "pretrain.parquet", nd, args.limit, rng) |
| LOG.info("binary base from %s: train %d val %d holdout %d pretrain %d (%.1fs)", base_origin, len(train), len(val), len(holdout), n_pretrain, common.timer() - t_start) |
|
|
| |
| new, champa_meta = load_champa(tw, args.limit, rng) |
| bigshp_meta: dict[str, Any] = {"enabled": bool(tw["bigshp"]["enabled"]), "verdict": tw["bigshp"]["verdict"], |
| "research": "docs/research/bigshp_research.md"} |
| if tw["bigshp"]["enabled"]: |
| bdf, bmeta = load_bigshp(tw, args.limit, rng, english_vocab(train, rng)) |
| bigshp_meta.update(bmeta) |
| new = pd.concat([new, bdf], ignore_index=True) |
| prepare_new_rows(new) |
| loaded_by_source = {str(k): int(v) for k, v in new["source"].value_counts().items()} |
| LOG.info("new rows: %s", loaded_by_source) |
|
|
| |
| hold_keys = frozenset(key_union(holdout)) |
| train_keys = frozenset(key_union(train)) |
| val_keys = frozenset(key_union(val)) |
| hit_hold, hit_train, hit_val, hit_pre = (rows_hitting(new, ks) for ks in (hold_keys, train_keys, val_keys, pretrain_keys)) |
| reasons = [ |
| (f"lt_{min_alnum}_alnum", (new["n_alnum"] < min_alnum).to_numpy()), |
| ("mbox_internal_data", (new["subject"] == MBOX_INTERNAL_SUBJECT).to_numpy()), |
| ("within_new_k_full_dup", new["k_full"].duplicated(keep="first").to_numpy()), |
| ("key_match_test_holdout", hit_hold), |
| ("key_match_finetune_train", hit_train), |
| ("key_match_finetune_val", hit_val), |
| ("key_match_pretrain", hit_pre), |
| ] |
| reason = first_reason(len(new), reasons) |
| contradictions = 0 |
| if tw["bigshp"]["enabled"]: |
| |
| |
| own_label: dict[str, int] = {} |
| for df in (train, val, holdout): |
| for col in ("k_full", "k_pre", "k_sb", "k_sb_full"): |
| own_label.update(zip(df[col].tolist(), df["label3"].tolist())) |
| for i in np.flatnonzero((new["source"] == BIGSHP_SOURCE).to_numpy() & (hit_train | hit_val | hit_hold)): |
| match = [own_label[k] for k in (new.at[i, "k_full"], new.at[i, "k_pre"], new.at[i, "k_sb"], new.at[i, "k_sb_full"]) if k in own_label] |
| if match and any(m != int(new.at[i, "label3"]) and m != -1 for m in match): |
| contradictions += 1 |
| survivors = new.loc[reason == ""].reset_index(drop=True) |
| LOG.info("exact screens dropped %d of %d new rows: %s", int((reason != "").sum()), len(new), by_source_reason(new["source"].to_numpy(), reason)) |
|
|
| |
| n_new = len(survivors) |
| ref = ShingleIndex(survivors["base"].tolist() + val["base"].tolist(), nd) |
| _q, r_hold = ref.query(holdout["base"].tolist()) |
| new_hold_linked = np.zeros(n_new, dtype=bool) |
| new_hold_linked[r_hold[r_hold < n_new]] = True |
| val_hold_linked = int(np.unique(r_hold[r_hold >= n_new]).size) |
| _q, r_tr = ref.query(train["base"].tolist()) |
| _q, r_pre = ref.query(pretrain_bases) |
| new_train_linked = np.zeros(n_new, dtype=bool) |
| new_train_linked[r_tr[r_tr < n_new]] = True |
| new_pre_linked = np.zeros(n_new, dtype=bool) |
| new_pre_linked[r_pre[r_pre < n_new]] = True |
| val_train_linked = np.zeros(len(val), dtype=bool) |
| val_train_linked[r_tr[r_tr >= n_new] - n_new] = True |
| val_pre_linked = np.zeros(len(val), dtype=bool) |
| val_pre_linked[r_pre[r_pre >= n_new] - n_new] = True |
| la, lb = ref.links() |
| a_new, b_new = la < n_new, lb < n_new |
| new_new = a_new & b_new |
| new_val = a_new ^ b_new |
| new_side = np.where(a_new[new_val], la[new_val], lb[new_val]) |
| val_side = np.where(a_new[new_val], lb[new_val], la[new_val]) - n_new |
| new_val_linked = np.zeros(n_new, dtype=bool) |
| new_val_linked[new_side] = True |
| val_new_linked = np.zeros(len(val), dtype=bool) |
| val_new_linked[val_side] = True |
| del ref |
|
|
| |
| reason_s = first_reason(n_new, [("shingle_link_test_holdout", new_hold_linked)]) |
| keep_s = reason_s == "" |
| edges = [(int(a), int(b)) for a, b in zip(la[new_new].tolist(), lb[new_new].tolist()) if keep_s[a] and keep_s[b]] |
| edges += [(a, b) for a, b in same_key_pairs_multi([survivors["k_pre"].tolist(), survivors["k_sb"].tolist()]) if keep_s[a] and keep_s[b]] |
| cl = union_find(n_new, edges) |
| |
| |
| |
| train_side_cl = set(cl[(new_train_linked | new_pre_linked) & keep_s].tolist()) |
| val_side_cl = set(cl[new_val_linked & keep_s].tolist()) |
| both_cl = train_side_cl & val_side_cl if tw["drop_new_clusters_linked_to_both_train_and_val"] else set() |
| both_rows = keep_s & np.isin(cl, list(both_cl)) |
| reason_s[both_rows] = "near_dup_of_both_train_and_val" |
| keep_s = reason_s == "" |
| pin_train_cl = train_side_cl - both_cl |
| pin_val_cl = val_side_cl - both_cl |
| hold_touched_cl = set(cl[new_hold_linked].tolist()) |
| dest = np.full(n_new, "", dtype=object) |
| dest[keep_s & np.isin(cl, list(pin_train_cl))] = "train" |
| dest[keep_s & np.isin(cl, list(pin_val_cl))] = "val" |
| free = keep_s & (dest == "") |
| |
| tmp = pd.DataFrame({"source": survivors["source"], "label": survivors["label3"], "cluster_id": cl}) |
| is_test, split_test_info = grouped_stratified_split(tmp.loc[free].reset_index(drop=True), float(tw["new_source_test_fraction"]), rng) |
| free_idx = np.flatnonzero(free) |
| dest[free_idx[is_test]] = "test3" |
| rest_idx = free_idx[~is_test] |
| is_val_new, split_val_info = grouped_stratified_split(tmp.iloc[rest_idx].reset_index(drop=True), float(tw["new_source_val_fraction"]), rng) |
| dest[rest_idx[is_val_new]] = "val" |
| dest[rest_idx[~is_val_new]] = "train" |
| assert ((dest != "") == keep_s).all() |
| cluster_offset = int(max(train["cluster_id"].max(), val["cluster_id"].max(), holdout["cluster_id"].max())) + 1 |
| survivors["cluster_id"] = (cl + cluster_offset).astype(np.int64) |
| survivors["dest"] = dest |
| new_pins = { |
| "clusters_total": int(np.unique(cl[keep_s]).size), "clusters_pinned_train_by_train_or_pretrain_link": len(pin_train_cl), |
| "clusters_pinned_val_by_val_link": len(pin_val_cl), "clusters_dropped_linked_to_both_train_and_val": len(both_cl), |
| "rows_dropped_linked_to_both_train_and_val": int(both_rows.sum()), "clusters_with_a_holdout_linked_member": len(hold_touched_cl), |
| "rows_linked_to_train": int((new_train_linked & keep_s).sum()), "rows_linked_to_pretrain": int((new_pre_linked & keep_s).sum()), |
| "rows_linked_to_val": int((new_val_linked & keep_s).sum()), "rows_dropped_holdout_shingle_link": int(new_hold_linked.sum()), |
| "largest_cluster": int(np.bincount(cl[keep_s]).max()) if keep_s.any() else 0, |
| } |
| LOG.info("new rows: pins %s; destinations %s", new_pins, Counter(dest[keep_s].tolist())) |
| all_reason = reason.copy() |
| all_reason[np.flatnonzero(reason == "")] = reason_s |
| drops_by_source_reason = by_source_reason(new["source"].to_numpy(), all_reason) |
| new_kept = survivors.loc[keep_s].reset_index(drop=True) |
| dest_by_source = {s: {d: int(n) for d, n in grp["dest"].value_counts().items()} for s, grp in new_kept.groupby("source")} |
|
|
| |
| h3 = tw["holdout3_from_val"] |
| l3 = val["label3"].to_numpy() |
| src = val["source"].to_numpy() |
| cand = ((l3 == 0) & np.isin(src, list(h3["legit_sources"]))) | ((l3 == 2) & np.isin(src, list(h3["spam_sources"]))) |
| val_pre_key = rows_hitting(val, pretrain_keys) |
| val_hold_key = rows_hitting(val, hold_keys) |
| train_clusters = set(train["cluster_id"].tolist()) |
| in_train_cl = val["cluster_id"].isin(train_clusters).to_numpy() |
| excl = first_reason(len(val), [ |
| ("not_a_candidate_source_or_label", ~cand), |
| ("key_match_pretrain", val_pre_key), ("shingle_link_pretrain", val_pre_linked), |
| ("key_or_shingle_link_test_holdout", val_hold_key | (np.isin(np.arange(len(val)), r_hold[r_hold >= n_new] - n_new))), |
| ("shingle_link_train", val_train_linked), ("cluster_also_in_train", in_train_cl), |
| ("shingle_link_to_a_new_row", val_new_linked), |
| ]) |
| row_ok = excl == "" |
| cluster_ok = pd.Series(row_ok).groupby(val["cluster_id"].to_numpy()).transform("all").to_numpy() |
| excl[row_ok & ~cluster_ok] = "cluster_has_an_ineligible_member" |
| eligible = excl == "" |
| moved = np.zeros(len(val), dtype=bool) |
| h3_split_info = {} |
| for cls, target_rows in ((0, int(h3["legit_rows"])), (2, int(h3["spam_rows"]))): |
| pool_mask = eligible & (l3 == cls) |
| pool_idx = np.flatnonzero(pool_mask) |
| if pool_idx.size == 0: |
| LOG.warning("holdout3 from val: no eligible rows for label3=%d", cls) |
| continue |
| frac = min(1.0, target_rows / pool_idx.size) |
| tmp = pd.DataFrame({"source": src[pool_idx], "label": l3[pool_idx], "cluster_id": val["cluster_id"].to_numpy()[pool_idx]}) |
| sel, info = grouped_stratified_split(tmp, frac, rng) |
| moved[pool_idx[sel]] = True |
| h3_split_info[LABEL3_NAMES[cls]] = {"target_rows": target_rows, "eligible_rows": int(pool_idx.size), "fraction": round(frac, 4), "moved_rows": int(sel.sum()), "split": info} |
| excl_counts = by_source_reason(src, np.where(cand, excl, "")) |
| LOG.info("holdout3 from val: candidates %d, eligible %d, moved %d; exclusions %s", int(cand.sum()), int(eligible.sum()), int(moved.sum()), excl_counts) |
|
|
| |
| cols = THREEWAY_SCHEMA[:1] + ["subject", "body", "label", "source", "cluster_id", "label3", "label_set", "base", "k_full", "k_pre", "k_sb", "k_sb_full", "n_alnum"] |
| new_train = new_kept.loc[new_kept["dest"] == "train"] |
| new_val = new_kept.loc[new_kept["dest"] == "val"] |
| new_test = new_kept.loc[new_kept["dest"] == "test3"] |
| train_f = pd.concat([train[cols], new_train[cols]], ignore_index=True) |
| val_f = pd.concat([val.loc[~moved, cols], new_val[cols]], ignore_index=True) |
| test3 = pd.concat([val.loc[moved, cols], new_test[cols]], ignore_index=True) |
| links = recount_links_multi( |
| {"train": train_f["base"].tolist(), "val": val_f["base"].tolist(), "holdout": holdout["base"].tolist(), "test3": test3["base"].tolist()}, |
| pretrain_bases, ("holdout", "val", "test3"), nd, |
| ) |
| del pretrain_bases |
| leak = leakage_checks3(train_f, val_f, holdout, test3, pretrain_keys, n_pretrain, links, holdout_source, pos_map) |
| leak.raise_if_failed() |
|
|
| |
| def out_table(base: pa.Table, keep: np.ndarray | None, frame: pd.DataFrame, extra: pd.DataFrame | None, tag: str) -> pa.Table: |
| t = base.select(FINETUNE_SCHEMA) |
| if keep is not None: |
| t = t.filter(pa.array(keep)) |
| if t.num_rows != len(frame) or (t.num_rows and t.column("doc_id")[0].as_py() != frame["doc_id"].iloc[0]): |
| raise AssertionError(f"{tag}: arrow table / frame misaligned") |
| t = t.append_column("label3", pa.array(frame["label3"].to_numpy(), pa.int64())) |
| t = t.append_column("label_set", pa.array(frame["label_set"].tolist(), pa.large_string())) |
| if extra is not None and len(extra): |
| ex = pa.Table.from_pandas(extra[THREEWAY_SCHEMA].reset_index(drop=True), preserve_index=False).cast(t.schema) |
| t = pa.concat_tables([t, ex]).combine_chunks() |
| meta = {b"threeway": json.dumps({"label3": SET_OF_LABEL3, "display_names": LABEL3_NAMES, "binary_rollup": UNWANTED_NAME, "built_by": "data_prep.py build3"}).encode()} |
| return t.replace_schema_metadata(meta) |
|
|
| out_tables = { |
| "finetune_train": out_table(tables["finetune_train"], None, train, new_train, "train"), |
| "finetune_val": out_table(tables["finetune_val"], ~moved, val.loc[~moved], new_val, "val"), |
| "test_holdout": out_table(tables["test_holdout"], None, holdout, None, "holdout"), |
| "test_holdout3": out_table(tables["finetune_val"], moved, val.loc[moved], new_test, "test3"), |
| } |
| files: dict[str, str] = {} |
| for name, t in out_tables.items(): |
| path = out_dir / f"{name}.parquet" |
| write_table_atomic(t, path) |
| files[name] = str(path) |
| LOG.info("wrote %s (%d rows, %s)", path, t.num_rows, Counter(t.column("label3").to_pylist())) |
| if pretrain_dev_table is not None: |
| write_table_atomic(pretrain_dev_table, out_dir / "pretrain.parquet") |
| files["pretrain"] = str(out_dir / "pretrain.parquet") |
| del tables, out_tables |
|
|
| |
| finals = {"finetune_train": train_f, "finetune_val": val_f, "test_holdout": holdout, "test_holdout3": test3} |
| elapsed = round(common.timer() - t_start, 1) |
| section = { |
| "timestamp": datetime.now(timezone.utc).isoformat(timespec="seconds"), |
| "limit": args.limit, "dev_subsample": dev, "seed": seed, "rng_stream": "default_rng(seed + 3)", |
| "elapsed_seconds": elapsed, "peak_rss_gb": peak_rss_gb(), |
| "contract": {"label3": {str(k): v for k, v in SET_OF_LABEL3.items()}, "display_names": {str(k): v for k, v in LABEL3_NAMES.items()}, |
| "binary_rollup_name": UNWANTED_NAME, "positive_label3_by_source": pos_map, |
| "binary_label": "unchanged: 1 for phish / spam / unwanted, 0 for legit"}, |
| "binary_base": {"from": base_origin, "dir": str(base_dir), "counts_by_source_label3": base_counts, |
| "rows": {n: int(len(frames[n])) for n in frames}}, |
| "sources": {"champa": champa_meta, "bigshp": bigshp_meta}, |
| "new_rows": {"loaded_by_source": loaded_by_source, "drops_by_source_reason": drops_by_source_reason, |
| "bigshp_label_contradictions_among_key_matches": int(contradictions), |
| "kept_by_source": {str(k): int(v) for k, v in new_kept["source"].value_counts().items()}, |
| "destination_by_source": dest_by_source, "pins": new_pins, |
| "test_fraction": tw["new_source_test_fraction"], "val_fraction_of_remainder": tw["new_source_val_fraction"], |
| "split_test3": split_test_info, "split_val": split_val_info, "cluster_id_offset": cluster_offset}, |
| "holdout3_from_val": {"legit_sources": h3["legit_sources"], "spam_sources": h3["spam_sources"], |
| "candidates_by_source_label3": nested_counts3(val.loc[cand]), |
| "excluded_candidates_by_source_reason": excl_counts, |
| "eligible_by_source_label3": nested_counts3(val.loc[eligible]), |
| "moved_by_source_label3": nested_counts3(val.loc[moved]), "per_class": h3_split_info, |
| "val_rows_linked_to_holdout_info": val_hold_linked}, |
| "counts": {n: nested_counts3(df) for n, df in finals.items()}, |
| "totals": {n: split_totals3(df) for n, df in finals.items()}, |
| "shingle_links": links, "near_dup": nd.__dict__, "leakage": leak.checks, "files": files, |
| "config_snapshot": tw, |
| } |
| manifest_src = processed_dir / "splits_manifest.json" |
| manifest = common.read_json(manifest_src) if manifest_src.exists() else {} |
| manifest["threeway"] = section |
| tmp_path = out_dir / "splits_manifest.json.tmp" |
| common.write_json(manifest, tmp_path) |
| os.replace(tmp_path, out_dir / "splits_manifest.json") |
| report_src = processed_dir / "data_report.md" |
| report = report_src.read_text(encoding="utf-8") if report_src.exists() else "# Data report\n" |
| marker = "\n## Three-way labels" |
| if marker in report: |
| report = report[: report.index(marker)].rstrip("\n") + "\n" |
| write_text_atomic(report.rstrip("\n") + "\n" + render_threeway_report(section), out_dir / "data_report.md") |
| LOG.info("wrote %s and %s in %.1fs (peak RSS %.2f GB)", out_dir / "splits_manifest.json", out_dir / "data_report.md", elapsed, section["peak_rss_gb"]) |
| print_summary3(section) |
|
|
|
|
| def render_threeway_report(s: dict) -> str: |
| lines = [f"\n## Three-way labels — build3 {s['timestamp']}", ""] |
| if s["dev_subsample"]: |
| lines += [f"**DEV SUBSAMPLE (--limit {s['limit']}): counts below are NOT the real build.**", ""] |
| c = s["contract"] |
| lines += [ |
| f"label3 {c['label3']}; display names {c['display_names']}; binary roll-up \"{c['binary_rollup_name']}\" = P(phishing)+P(spam). " |
| f"Positive rows by provenance: {c['positive_label3_by_source']}. {c['binary_label']}. " |
| f"Elapsed {s['elapsed_seconds']}s; peak RSS {s['peak_rss_gb']} GB; binary base from {s['binary_base']['from']} (`{s['binary_base']['dir']}`).", |
| "", |
| "### New sources", |
| ] |
| ch = s["sources"]["champa"] |
| for fname, f in ch["files"].items(): |
| lines.append(f"- `{ch['repo']}` / `{fname}` @ `{f['commit_sha']}`: {f['rows']:,} rows ({f['rows_used']:,} used), label3 {ch['label3']} (phishing by provenance)") |
| bg = s["sources"]["bigshp"] |
| if bg["enabled"]: |
| lines.append(f"- bigshp `{bg['repo']}` @ `{bg.get('commit_sha')}` ENABLED despite verdict {bg['verdict']}: {bg['rows_read']:,} rows read, " |
| f"{bg['rows_after_file_filters']:,} after the file filters {bg['file_filter_drops']}; filters {bg['filters']}") |
| else: |
| lines.append(f"- bigshp (`{bg.get('repo', 'locuoco/the-biggest-spam-ham-phish-email-dataset-300000')}`): research verdict **{bg['verdict']}**, not ingested (`{bg['research']}`)") |
| nr = s["new_rows"] |
| rows = [] |
| for src in sorted(nr["loaded_by_source"]): |
| d = nr["drops_by_source_reason"].get(src, {}) |
| dest = nr["destination_by_source"].get(src, {}) |
| rows.append([src, nr["loaded_by_source"][src], sum(d.values()), nr["kept_by_source"].get(src, 0), dest.get("train", 0), dest.get("val", 0), dest.get("test3", 0)]) |
| lines += ["", md_table(["source", "loaded", "dropped", "kept", "-> train", "-> val", "-> test_holdout3"], rows), |
| "", f"- drops by reason: {nr['drops_by_source_reason']}", f"- near-dup pins: {nr['pins']}", |
| f"- bigshp label contradictions among exact matches: {nr['bigshp_label_contradictions_among_key_matches']}", |
| "", "### test_holdout3 rows taken from finetune_val", |
| f"- legitimate from {s['holdout3_from_val']['legit_sources']}, spam from {s['holdout3_from_val']['spam_sources']}; whole near-dup clusters only; " |
| f"excluded candidates by reason: {s['holdout3_from_val']['excluded_candidates_by_source_reason']}", |
| f"- per class: { {k: {kk: vv for kk, vv in v.items() if kk != 'split'} for k, v in s['holdout3_from_val']['per_class'].items()} }", |
| source_label3_table(s["holdout3_from_val"]["moved_by_source_label3"]), "", "### Splits by label3"] |
| for name in THREEWAY_SPLITS: |
| t = s["totals"][name] |
| lines += [f"#### {name}: {t['rows']:,} rows, by label3 {t['by_label3']}, binary positives {t['binary_positives']:,} ({t['binary_base_rate']})", |
| source_label3_table(s["counts"][name]), ""] |
| lines += ["### Leakage assertions (binary gate + three-way)", |
| md_table(["check", "count", "asserted zero", "pass"], [[k, v["count"], v["asserted"], v["pass"]] for k, v in s["leakage"].items()]), |
| "", "### Files"] + [f"- `{k}`: `{v}`" for k, v in s["files"].items()] |
| return "\n".join(lines) + "\n" |
|
|
|
|
| def print_summary3(s: dict) -> None: |
| print("\n=== data_prep build3 summary (three-way) ===") |
| if s["dev_subsample"]: |
| print(f"*** DEV SUBSAMPLE --limit {s['limit']} — not a real build ***") |
| for name in THREEWAY_SPLITS: |
| t = s["totals"][name] |
| print(f"{name:>16}: rows {t['rows']:>7,} by label3 {t['by_label3']} binary positives {t['binary_positives']:,}") |
| for src, cnt in sorted(s["counts"][name].items()): |
| print(f"{'':>18}{src:<16} {cnt}") |
| print("leakage:") |
| for k, v in s["leakage"].items(): |
| print(f" {k:<44} {v['count']:>8} {'ASSERT' if v['asserted'] else 'info '} {'ok' if v['pass'] else 'FAIL'}") |
|
|
|
|
| |
| def cmd_check3(args: argparse.Namespace, cfg: dict) -> None: |
| """Recompute every three-way assertion from the written files (no reuse of build3 bookkeeping).""" |
| processed_dir = Path(cfg["paths"]["data_processed"]) |
| out_dir = Path(args.out_dir) if args.out_dir else processed_dir |
| manifest = common.read_json(out_dir / "splits_manifest.json") |
| if "threeway" not in manifest: |
| raise AssertionError(f"{out_dir / 'splits_manifest.json'} has no 'threeway' section: run build3 first") |
| section = manifest["threeway"] |
| tw = threeway_config(cfg) |
| pos_map = {str(k): int(v) for k, v in tw["positive_label3_by_source"].items()} |
| nd = NearDupConfig.from_config(cfg, cap=int(section["near_dup"]["cap"])) |
| holdout_source = cfg["data"]["holdout_test_source"] |
| frames: dict[str, pd.DataFrame] = {} |
| for name in THREEWAY_SPLITS: |
| df = pd.read_parquet(out_dir / f"{name}.parquet") |
| if list(df.columns) != THREEWAY_SCHEMA: |
| raise AssertionError(f"{name}: expected columns {THREEWAY_SCHEMA}, got {list(df.columns)}") |
| df["subject"] = df["subject"].fillna("") |
| df["body"] = df["body"].fillna("") |
| df["label"] = df["label"].astype(np.int64) |
| df["label3"] = df["label3"].astype(np.int64) |
| df["source"] = df["source"].astype(str) |
| df["base"] = add_key_columns(df) |
| frames[name] = df |
| LOG.info("read %s: %d rows, label3 %s", name, len(df), dict(Counter(df["label3"].tolist()))) |
| pre_path = out_dir / "pretrain.parquet" if (out_dir / "pretrain.parquet").exists() else processed_dir / "pretrain.parquet" |
| pretrain_keys, pretrain_bases, n_pretrain, _ = pretrain_identity(pre_path, nd, None, np.random.default_rng(0)) |
| train, val, hold, test3 = (frames[n] for n in THREEWAY_SPLITS) |
| links = recount_links_multi( |
| {"train": train["base"].tolist(), "val": val["base"].tolist(), "holdout": hold["base"].tolist(), "test3": test3["base"].tolist()}, |
| pretrain_bases, ("holdout", "val", "test3"), nd, |
| ) |
| del pretrain_bases |
| leak = leakage_checks3(train, val, hold, test3, pretrain_keys, n_pretrain, links, holdout_source, pos_map) |
|
|
| |
| snap = processed_dir / BINARY_SNAPSHOT_DIR / "test_holdout.parquet" |
| if snap.exists() and not section["dev_subsample"]: |
| ref = pq.read_table(snap, columns=["doc_id", "label"]) |
| same = ref.column("doc_id").to_pylist() == hold["doc_id"].tolist() and ref.column("label").to_pylist() == hold["label"].tolist() |
| leak.add("test_holdout_rows_changed_vs_binary_snapshot", 0 if same else 1, True, str(snap)) |
| ref_tr = pq.read_table(processed_dir / BINARY_SNAPSHOT_DIR / "finetune_train.parquet", columns=["doc_id"]).column("doc_id").to_pylist() |
| ref_va = pq.read_table(processed_dir / BINARY_SNAPSHOT_DIR / "finetune_val.parquet", columns=["doc_id"]).column("doc_id").to_pylist() |
| binary_ids = set(ref_tr) | set(ref_va) |
| now_ids = set(train["doc_id"]) | set(val["doc_id"]) | set(test3["doc_id"]) |
| leak.add("binary_train_val_rows_missing_from_train_val_holdout3", len(binary_ids - now_ids), True, "every binary train/val row must still exist in train, val or test_holdout3") |
| leak.add("binary_train_rows_not_in_train", len(set(ref_tr) - set(train["doc_id"])), True, "train membership is only ever extended") |
| for name, df in frames.items(): |
| want = section["totals"][name]["rows"] |
| if want != len(df): |
| leak.add(f"{name}_rows_vs_manifest", abs(want - len(df)), True, f"manifest {want} vs file {len(df)}") |
| syn = processed_dir / "synthetic.parquet" |
| if syn.exists(): |
| st = pq.read_table(syn) |
| if {"label3", "label_set", "label"} <= set(st.schema.names): |
| sdf = st.select(["label", "label3", "label_set"]).to_pandas() |
| sdf["source"] = "synthetic" |
| v = contract_violations(sdf, {"synthetic": 1}) |
| leak.add("synthetic_label3_contract_violations", sum(v.values()), True, str(v)) |
| else: |
| leak.add("synthetic_missing_label3_columns", 1, False, "synthetic.parquet predates the three-way contract; finetune --synthetic on will refuse it") |
| print_summary3({**section, "leakage": leak.checks}) |
| print("\nmanifest threeway timestamp:", section["timestamp"], "| limit:", section["limit"], "| bigshp enabled:", section["sources"]["bigshp"]["enabled"]) |
| print("build-time results:", {k: v["count"] for k, v in section["leakage"].items() if v["asserted"]}) |
| leak.raise_if_failed() |
| print("ALL THREE-WAY LEAKAGE ASSERTIONS PASS") |
|
|
|
|
| |
| def main(argv: list[str] | None = None) -> None: |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--config", default="config.yaml") |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| sub.add_parser("inspect", help="features + 2 rows of all four datasets (streaming)") |
| b = sub.add_parser("build", help="build processed parquet files + manifest + report") |
| b.add_argument("--limit", type=int, default=None, help="DEV ONLY: N rows per labelled source, 4N corbt rows") |
| b.add_argument("--spam-policy", choices=["positive_class", "map_to_zero", "exclude"], default=None) |
| b.add_argument("--no-strip-quotes", action="store_true", help="keep quoted/forwarded history in corbt bodies") |
| b.add_argument("--shingle-cap", type=int, default=None, help="ignore shingles shared by more rows than this (default: config data.dedup.near_dup_shingle_cap, else 2000)") |
| sub.add_parser("check", help="re-run leakage assertions on the written files") |
| b3 = sub.add_parser("build3", help="three-way labels: label3/label_set + Champa phishing corpora + test_holdout3 (reads the processed binary files)") |
| b3.add_argument("--limit", type=int, default=None, help="DEV ONLY: N rows per source of every split, N per new file, 4N pretrain rows; writes to data/processed/dev_threeway_limitN/") |
| b3.add_argument("--out-dir", default=None, help="output directory (default: data.processed, or the dev dir with --limit)") |
| b3.add_argument("--bigshp", choices=["on", "off"], default=None, help="override data.threeway.bigshp.enabled (research verdict: DO_NOT_USE)") |
| b3.add_argument("--shingle-cap", type=int, default=None, help="as for build") |
| c3 = sub.add_parser("check3", help="recompute the binary + three-way leakage assertions on the written files") |
| c3.add_argument("--out-dir", default=None, help="directory written by build3 (default: data.processed)") |
| args = ap.parse_args(argv) |
|
|
| cfg = common.load_config(args.config) |
| global LOG, DIALS |
| LOG = common.setup_logging("data_prep", cfg["paths"]["logs"]) |
| DIALS = Dials.from_config(cfg) |
| LOG.info("dials: %s", DIALS) |
| {"inspect": cmd_inspect, "build": cmd_build, "check": cmd_check, "build3": cmd_build3, "check3": cmd_check3}[args.cmd](args, cfg) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
