| """Byte-level BPE tokenizer for the phishing encoder: training, statistics, and batch encoding. |
| |
| Design (config.tokenizer): |
| - NFKC normalisation only, case preserved. ALL-CAPS subjects ("URGENT", "VERIFY YOUR ACCOUNT"), |
| odd casing in brand names and shouty pharma spam are strong phishing/spam signals; lowercasing |
| would throw them away. NFKC folds full-width and ligature look-alikes (a common filter-evasion |
| trick) onto their ASCII forms so they share subwords with the plain spelling. |
| - Byte-level pre-tokenizer with the full 256-symbol initial alphabet: every byte sequence is |
| representable, so [UNK] is unreachable in practice. Email is full of URLs, base64 fragments, |
| HTML entities and mojibake, which a character-level vocabulary would map to [UNK]. |
| - add_prefix_space=False: inputs are "subject\\n\\nbody", not isolated words. |
| - 32,768 ids (2^15) to match model.vocab_size; [PAD] is id 0 because the model's embedding uses |
| padding_idx=0 and the collate code pads with 0. |
| - Trained ONLY on pretrain.parquet + finetune_train.parquet. finetune_val and test_holdout |
| (TREC-06) never touch the vocabulary: a vocabulary fitted on test text leaks its frequent |
| strings (template URLs, spam phrases) into single tokens. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import logging |
| import sys |
| from collections.abc import Iterator, Sequence |
| from dataclasses import asdict, dataclass, field |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).parent)) |
|
|
| import numpy as np |
| import pyarrow.parquet as pq |
| import torch |
| from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, processors, trainers |
|
|
| import common |
|
|
| LOG = logging.getLogger("tokenizer") |
|
|
| SPLIT_FILES: dict[str, str] = { |
| "pretrain": "pretrain.parquet", |
| "finetune_train": "finetune_train.parquet", |
| "finetune_val": "finetune_val.parquet", |
| "test_holdout": "test_holdout.parquet", |
| } |
| FORBIDDEN_TRAIN_FILES = frozenset({"finetune_val.parquet", "test_holdout.parquet"}) |
| SOURCE_COLUMNS = ("source", "dataset_name") |
| TEXT_COLUMN = "text" |
| READ_BATCH_ROWS = 2048 |
| STATS_ENCODE_BATCH = 4096 |
| PRETRAIN_STATS_SAMPLE = 20_000 |
|
|
| EXAMPLES: list[tuple[str, str]] = [ |
| ("business", "Re: Q3 budget review\n\nHi Sarah, can we move the budget review to Thursday at 2pm? " |
| "I've attached the revised forecast. Thanks, Mike"), |
| ("url_heavy", "Your account has been suspended\n\nVerify now: http://secure-login.paypa1-verify.com/" |
| "account/update?id=8f3a9c2e&session=XyZ123 or https://bit.ly/3xYz9Ab before 24 hours."), |
| ("all_caps", "URGENT!!! FINAL NOTICE\n\nYOUR MAILBOX IS FULL. CLICK HERE TO UPGRADE YOUR STORAGE " |
| "IMMEDIATELY OR LOSE ALL MESSAGES."), |
| ("unicode_lookalike", "Ｍｉｃｒｏｓｏｆｔ 365 — Password expiry ⚠️\n\nDear user, your password expires " |
| "today. Café, naïve, 東京 office."), |
| ("html_money", "<html><body><p>Dear Friend,</p><p>I am Barrister James. Transfer of US$15,500,000.00 " |
| "& 30% for you.</p></body></html>"), |
| ] |
|
|
|
|
| @dataclass |
| class TokenizerPaths: |
| data_dir: Path |
| tokenizer_json: Path |
| report_json: Path |
| log_dir: Path |
|
|
|
|
| @dataclass |
| class SplitStats: |
| split: str |
| file: str |
| rows: int |
| mean_tokens: float |
| median_tokens: float |
| p95_tokens: float |
| max_tokens: int |
| chars_per_token: float |
| pct_truncated: float |
| unk_rate: float |
| mean_tokens_by_label: dict[str, float] = field(default_factory=dict) |
|
|
|
|
| def _abs(p: str | Path) -> Path: |
| p = Path(p) |
| return p if p.is_absolute() else common.project_root() / p |
|
|
|
|
| def resolve_paths(cfg: dict, data_dir: str | None) -> TokenizerPaths: |
| """--data-dir redirects data, artifacts and logs together, so a self-test never writes into the project.""" |
| if data_dir is not None: |
| d = Path(data_dir).resolve() |
| return TokenizerPaths(d, d / "tokenizer.json", d / "tokenizer_report.json", d / "logs") |
| d = _abs(cfg["paths"]["data_processed"]) |
| tok = _abs(cfg["paths"]["tokenizer"]) |
| return TokenizerPaths(d, tok, tok.with_name("tokenizer_report.json"), _abs(cfg["paths"]["logs"])) |
|
|
|
|
| |
|
|
|
|
| def assert_trainable(path: Path, holdout_source: str) -> None: |
| """Refuse validation/test files by name, and any file whose source column carries the holdout corpus.""" |
| if path.name in FORBIDDEN_TRAIN_FILES: |
| raise RuntimeError(f"refusing to train the tokenizer on {path.name}: val/test text must never shape the vocabulary") |
| schema_names = set(pq.read_schema(path).names) |
| if TEXT_COLUMN not in schema_names: |
| raise RuntimeError(f"{path} has no '{TEXT_COLUMN}' column (columns: {sorted(schema_names)})") |
| for col in SOURCE_COLUMNS: |
| if col in schema_names: |
| values = {str(v) for v in pq.read_table(path, columns=[col]).column(col).unique().to_pylist()} |
| hits = {v for v in values if v.strip().lower() == holdout_source.lower()} |
| if hits: |
| raise RuntimeError(f"{path.name} column '{col}' contains holdout source {hits}: test-set leakage") |
|
|
|
|
| def iter_texts(path: Path, limit: int | None = None) -> Iterator[list[str]]: |
| """Stream the text column in record batches; the 244k-email pretraining corpus never sits in RAM twice.""" |
| seen = 0 |
| for batch in pq.ParquetFile(path).iter_batches(batch_size=READ_BATCH_ROWS, columns=[TEXT_COLUMN]): |
| texts = [t for t in batch.column(0).to_pylist() if t] |
| if limit is not None: |
| texts = texts[: max(0, limit - seen)] |
| if texts: |
| seen += len(texts) |
| yield texts |
| if limit is not None and seen >= limit: |
| return |
|
|
|
|
| def holdout_collisions(train_files: Sequence[Path], holdout: Path, limit: int | None) -> dict[str, int]: |
| """Exact normalized-text collisions (verify_data.md §7 k_full) between training text and the holdout. |
| |
| The data pipeline already decontaminates; this is an independent tripwire at the point where |
| text first becomes model state. Reading the holdout here only builds a hash set. |
| """ |
| held = {common.sha1_hex(common.normalize_text(t)) for batch in iter_texts(holdout) for t in batch} |
| out: dict[str, int] = {} |
| for f in train_files: |
| out[f.name] = sum( |
| common.sha1_hex(common.normalize_text(t)) in held for batch in iter_texts(f, limit) for t in batch |
| ) |
| return out |
|
|
|
|
| |
|
|
|
|
| def build_untrained() -> Tokenizer: |
| tok = Tokenizer(models.BPE(unk_token="[UNK]")) |
| tok.normalizer = normalizers.NFKC() |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tok.decoder = decoders.ByteLevel() |
| return tok |
|
|
|
|
| def attach_post_processor(tok: Tokenizer) -> None: |
| """BERT template; segment ids only matter for pairs, which the classifier does not use but the MLM loader may.""" |
| cls_id, sep_id = tok.token_to_id("[CLS]"), tok.token_to_id("[SEP]") |
| tok.post_processor = processors.TemplateProcessing( |
| single="[CLS] $A [SEP]", |
| pair="[CLS] $A [SEP] $B:1 [SEP]:1", |
| special_tokens=[("[CLS]", cls_id), ("[SEP]", sep_id)], |
| ) |
|
|
|
|
| def check_special_ids(tok: Tokenizer, special_tokens: Sequence[str]) -> dict[str, int]: |
| ids = {t: tok.token_to_id(t) for t in special_tokens} |
| if ids.get("[PAD]") != 0: |
| raise AssertionError(f"[PAD] must be id 0 (model padding_idx), got {ids.get('[PAD]')}") |
| if list(ids.values()) != list(range(len(special_tokens))): |
| raise AssertionError(f"special tokens must occupy ids 0..{len(special_tokens) - 1} in order, got {ids}") |
| return ids |
|
|
|
|
| def train(cfg: dict, paths: TokenizerPaths, vocab_size: int | None, limit_docs: int | None) -> Tokenizer: |
| tcfg = cfg["tokenizer"] |
| vocab_size = vocab_size or int(tcfg["vocab_size"]) |
| model_vocab = int(cfg["model"]["vocab_size"]) |
| if vocab_size > model_vocab: |
| raise ValueError(f"tokenizer vocab {vocab_size} exceeds model.vocab_size {model_vocab}") |
| special = list(tcfg["special_tokens"]) |
| holdout_source = cfg["data"]["holdout_test_source"] |
|
|
| files = [paths.data_dir / SPLIT_FILES[name] for name in tcfg["train_on"]] |
| for f in files: |
| if not f.exists(): |
| raise FileNotFoundError(f) |
| assert_trainable(f, holdout_source) |
|
|
| holdout = paths.data_dir / SPLIT_FILES["test_holdout"] |
| if holdout.exists(): |
| collisions = holdout_collisions(files, holdout, limit_docs) |
| LOG.info("exact holdout collisions per training file: %s", collisions) |
| if any(collisions.values()): |
| raise RuntimeError(f"training text exactly matches TREC-06 holdout rows: {collisions}") |
| else: |
| collisions = None |
| LOG.warning("%s absent: exact-hash leakage tripwire skipped (name/source guards still applied)", holdout.name) |
|
|
| usage: dict[str, dict[str, int]] = {} |
|
|
| def corpus() -> Iterator[list[str]]: |
| for f in files: |
| rows = chars = 0 |
| for batch in iter_texts(f, limit_docs): |
| rows += len(batch) |
| chars += sum(map(len, batch)) |
| yield batch |
| usage[f.name] = {"rows_in_file": pq.ParquetFile(f).metadata.num_rows, "rows_used": rows, "chars": chars} |
| LOG.info("fed %s: %d rows, %d chars", f.name, rows, chars) |
|
|
| tok = build_untrained() |
| trainer_kwargs: dict = dict( |
| vocab_size=vocab_size, |
| min_frequency=int(tcfg["min_frequency"]), |
| special_tokens=special, |
| initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=True, |
| ) |
| if tcfg.get("max_token_length"): |
| trainer_kwargs["max_token_length"] = int(tcfg["max_token_length"]) |
| LOG.info("training byte-level BPE: %s on %s (limit_docs=%s)", |
| {k: v for k, v in trainer_kwargs.items() if k != "initial_alphabet"}, [f.name for f in files], limit_docs) |
| t0 = common.timer() |
| tok.train_from_iterator(corpus(), trainer=trainers.BpeTrainer(**trainer_kwargs)) |
| elapsed = common.timer() - t0 |
|
|
| special_ids = check_special_ids(tok, special) |
| attach_post_processor(tok) |
| final_vocab = tok.get_vocab_size() |
| if final_vocab > model_vocab: |
| raise AssertionError(f"trained vocab {final_vocab} > model.vocab_size {model_vocab}") |
| if final_vocab < vocab_size: |
| LOG.warning("vocab reached only %d of %d (corpus too small / min_frequency); embedding rows above stay unused", |
| final_vocab, vocab_size) |
|
|
| paths.tokenizer_json.parent.mkdir(parents=True, exist_ok=True) |
| tok.save(str(paths.tokenizer_json)) |
| LOG.info("saved %s (vocab %d) in %.1fs", paths.tokenizer_json, final_vocab, elapsed) |
|
|
| report = _read_report(paths) |
| report["train"] = { |
| "tokenizer_json": str(paths.tokenizer_json), |
| "tokenizer_sha1": common.sha1_hex(paths.tokenizer_json.read_text()), |
| "requested_vocab_size": vocab_size, |
| "final_vocab_size": final_vocab, |
| "min_frequency": int(tcfg["min_frequency"]), |
| "special_token_ids": special_ids, |
| "normalizer": "NFKC (case preserved)", |
| "pre_tokenizer": "ByteLevel(add_prefix_space=False)", |
| "training_files": usage, |
| "limit_docs_per_file": limit_docs, |
| "excluded_files": sorted(FORBIDDEN_TRAIN_FILES), |
| "holdout_exact_collisions": collisions, |
| "train_seconds": round(elapsed, 1), |
| } |
| common.write_json(report, paths.report_json) |
| return tok |
|
|
|
|
| |
|
|
|
|
| def load_tokenizer(path: str | Path) -> Tokenizer: |
| """Load tokenizer.json with truncation/padding cleared; callers configure them (or use encode_batch).""" |
| tok = Tokenizer.from_file(str(path)) |
| if tok.token_to_id("[PAD]") != 0: |
| raise AssertionError(f"{path}: [PAD] is not id 0") |
| tok.no_truncation() |
| tok.no_padding() |
| return tok |
|
|
|
|
| def encode_batch( |
| tok: Tokenizer, texts: Sequence[str], max_len: int, pad_to_multiple_of: int | None = 8 |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Encode with [CLS]/[SEP], truncate to max_len, pad to the longest in the batch (dynamic padding). |
| |
| pad_to_multiple_of=8 keeps sequence lengths tensor-core friendly; it must divide max_len so a |
| full-length row is never padded past the model's position table. |
| Returns (input_ids int64 [B, L], attention_mask bool [B, L]). |
| """ |
| if pad_to_multiple_of and max_len % pad_to_multiple_of: |
| raise ValueError(f"max_len {max_len} is not a multiple of pad_to_multiple_of {pad_to_multiple_of}") |
| if len(texts) == 0: |
| return torch.empty((0, 0), dtype=torch.int64), torch.empty((0, 0), dtype=torch.bool) |
| |
| |
| tok.enable_truncation(max_length=max_len) |
| tok.enable_padding(pad_id=0, pad_token="[PAD]", pad_to_multiple_of=pad_to_multiple_of) |
| try: |
| encs = tok.encode_batch_fast(list(texts)) |
| finally: |
| tok.no_truncation() |
| tok.no_padding() |
| ids = np.asarray([e.ids for e in encs], dtype=np.int64).reshape(len(texts), -1) |
| mask = np.asarray([e.attention_mask for e in encs], dtype=np.bool_).reshape(len(texts), -1) |
| return torch.from_numpy(ids), torch.from_numpy(mask) |
|
|
|
|
| |
|
|
|
|
| def _read_report(paths: TokenizerPaths) -> dict: |
| return common.read_json(paths.report_json) if paths.report_json.exists() else {} |
|
|
|
|
| def _load_split(path: Path, sample: int | None, seed: int) -> tuple[list[str], list | None]: |
| cols = [TEXT_COLUMN] + (["label"] if "label" in pq.read_schema(path).names else []) |
| table = pq.read_table(path, columns=cols) |
| if sample is not None and table.num_rows > sample: |
| idx = np.sort(np.random.default_rng(seed).choice(table.num_rows, size=sample, replace=False)) |
| table = table.take(idx) |
| texts = [t or "" for t in table.column(TEXT_COLUMN).to_pylist()] |
| labels = table.column("label").to_pylist() if "label" in cols else None |
| return texts, labels |
|
|
|
|
| def split_stats(tok: Tokenizer, split: str, path: Path, max_len: int, sample: int | None, seed: int) -> SplitStats: |
| """Content-token counts (no specials); truncation is judged on content + [CLS] + [SEP] > max_len.""" |
| texts, labels = _load_split(path, sample, seed) |
| tok.no_truncation() |
| tok.no_padding() |
| unk_id = tok.token_to_id("[UNK]") |
| lengths = np.empty(len(texts), dtype=np.int64) |
| unk = 0 |
| for i in range(0, len(texts), STATS_ENCODE_BATCH): |
| encs = tok.encode_batch_fast(texts[i : i + STATS_ENCODE_BATCH], add_special_tokens=False) |
| for j, e in enumerate(encs): |
| lengths[i + j] = len(e.ids) |
| unk += e.ids.count(unk_id) |
| total_tokens = int(lengths.sum()) |
| by_label: dict[str, float] = {} |
| if labels is not None: |
| lab = np.asarray(labels) |
| by_label = {str(v): round(float(lengths[lab == v].mean()), 1) for v in np.unique(lab)} |
| return SplitStats( |
| split=split, |
| file=path.name, |
| rows=len(texts), |
| mean_tokens=round(float(lengths.mean()), 1), |
| median_tokens=float(np.median(lengths)), |
| p95_tokens=float(np.percentile(lengths, 95)), |
| max_tokens=int(lengths.max()), |
| chars_per_token=round(sum(map(len, texts)) / max(total_tokens, 1), 3), |
| pct_truncated=round(100.0 * float((lengths + 2 > max_len).mean()), 2), |
| unk_rate=unk / max(total_tokens, 1), |
| mean_tokens_by_label=by_label, |
| ) |
|
|
|
|
| def example_encodings(tok: Tokenizer) -> list[dict]: |
| tok.no_truncation() |
| tok.no_padding() |
| out = [] |
| for name, text in EXAMPLES: |
| enc = tok.encode(text) |
| decoded = tok.decode(enc.ids) |
| out.append({ |
| "name": name, |
| "text": text, |
| "n_tokens_with_specials": len(enc.ids), |
| "tokens": enc.tokens, |
| "decoded": decoded, |
| "round_trip_exact": decoded == text, |
| "round_trip_matches_nfkc": decoded == normalizers.NFKC().normalize_str(text), |
| }) |
| return out |
|
|
|
|
| def stats(cfg: dict, paths: TokenizerPaths) -> dict: |
| max_len = int(cfg["model"]["max_seq_len"]) |
| tok = load_tokenizer(paths.tokenizer_json) |
| rows: list[SplitStats] = [] |
| for split in ("pretrain", "finetune_train", "finetune_val", "test_holdout"): |
| f = paths.data_dir / SPLIT_FILES[split] |
| if not f.exists(): |
| LOG.warning("%s missing, skipped", f) |
| continue |
| sample = PRETRAIN_STATS_SAMPLE if split == "pretrain" else None |
| t0 = common.timer() |
| rows.append(split_stats(tok, split, f, max_len, sample, int(cfg["seed"]))) |
| LOG.info("%s: %d rows in %.1fs", split, rows[-1].rows, common.timer() - t0) |
| examples = example_encodings(tok) |
|
|
| report = _read_report(paths) |
| report["stats"] = { |
| "max_seq_len": max_len, |
| "vocab_size": tok.get_vocab_size(), |
| "pretrain_sample_rows": PRETRAIN_STATS_SAMPLE, |
| "splits": [asdict(r) for r in rows], |
| "examples": examples, |
| } |
| common.write_json(report, paths.report_json) |
| _print_report(rows, examples, max_len, tok.get_vocab_size()) |
| return report |
|
|
|
|
| def _print_report(rows: list[SplitStats], examples: list[dict], max_len: int, vocab: int) -> None: |
| print(f"\n### Tokenizer statistics (vocab {vocab:,}, max_seq_len {max_len}; token counts exclude [CLS]/[SEP])\n") |
| print(f"| split | rows | mean tok | median | p95 | max | chars/tok | % > {max_len} | [UNK] rate | mean tok by label |") |
| print("|---|---:|---:|---:|---:|---:|---:|---:|---:|---|") |
| for r in rows: |
| lab = ", ".join(f"{k}: {v}" for k, v in r.mean_tokens_by_label.items()) or "-" |
| print(f"| {r.split} | {r.rows:,} | {r.mean_tokens:,.1f} | {r.median_tokens:,.0f} | {r.p95_tokens:,.0f} | " |
| f"{r.max_tokens:,} | {r.chars_per_token:.2f} | {r.pct_truncated:.1f}% | {r.unk_rate:.2e} | {lab} |") |
| print("\n### Example encodings\n") |
| for e in examples: |
| shown = " ".join(e["tokens"][:48]) + (" ..." if len(e["tokens"]) > 48 else "") |
| print(f"**{e['name']}** ({e['n_tokens_with_specials']} tokens; round trip exact={e['round_trip_exact']}, " |
| f"matches NFKC(input)={e['round_trip_matches_nfkc']})") |
| print(f"- tokens: `{shown}`") |
| print(f"- decoded: {e['decoded']!r}\n") |
|
|
|
|
| |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) |
| ap.add_argument("--config", default="config.yaml") |
| ap.add_argument("--data-dir", default=None, help=argparse.SUPPRESS) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| tr = sub.add_parser("train", help="train byte-level BPE on pretrain + finetune_train") |
| tr.add_argument("--vocab-size", type=int, default=None) |
| tr.add_argument("--limit-docs", type=int, default=None, help="dev only: cap rows read per training file") |
| sub.add_parser("stats", help="token-length / truncation / [UNK] statistics per split") |
| args = ap.parse_args(argv) |
|
|
| cfg = common.load_config(args.config) |
| paths = resolve_paths(cfg, args.data_dir) |
| global LOG |
| LOG = common.setup_logging("tokenizer", paths.log_dir) |
| common.set_seed(int(cfg["seed"])) |
| if args.cmd == "train": |
| train(cfg, paths, args.vocab_size, args.limit_docs) |
| else: |
| stats(cfg, paths) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
