""" ZDX Tokenizer Trainer ===================== Standalone BPE tokenizer training CLI. Project: https://zerodrivex.com/tools Version: 1.0.1 Copyright (c) 2026 ZeroDriveX SPDX-License-Identifier: MIT Requirements: pip install tokenizers Usage: # Train on a directory of text files python3 tokenizer.py --input ./corpus --output ./my_tokenizer # Train on a single file python3 tokenizer.py --input ./corpus.txt --output ./my_tokenizer # Custom vocab size and special tokens python3 tokenizer.py --input ./corpus --output ./my_tokenizer --vocab_size 32000 --min_freq 3 # Export specific formats only python3 tokenizer.py --input ./corpus --output ./my_tokenizer --formats json vocab Exports: my_tokenizer/tokenizer.json HuggingFace fast tokenizer format my_tokenizer/vocab.txt Plain vocab list, one token per line my_tokenizer/vocab.json Token -> ID mapping my_tokenizer/merges.txt BPE merge rules (GPT-2 style) my_tokenizer/tokenizer_config.json Config metadata my_tokenizer/zdx_tokenizer.json ZDX BAN native format (used by train.py) """ import os import re import json import argparse from pathlib import Path from typing import Optional # --------------------------------------------------------------------------- # Defaults # --------------------------------------------------------------------------- DEFAULT_VOCAB_SIZE = 16000 DEFAULT_MIN_FREQ = 2 DEFAULT_FORMATS = ["json", "vocab", "merges", "config", "zdx"] SPECIAL_TOKENS = ["[PAD]", "[UNK]", "[BOS]", "[EOS]", "[MASK]", "[SEP]", "[CLS]"] # --------------------------------------------------------------------------- # Text loading — reuses formatter logic, lightweight version # --------------------------------------------------------------------------- def load_texts(input_path: str) -> list[str]: input_path = Path(input_path) texts = [] if input_path.is_dir(): files = sorted([ p for p in input_path.rglob("*") if p.suffix.lower() in (".txt", ".json", ".jsonl", ".md", ".csv") and p.is_file() ]) print(f"[Loader] Found {len(files)} file(s)") for fpath in files: try: with open(fpath, "r", encoding="utf-8", errors="replace") as f: content = f.read().strip() if content: texts.append(content) except Exception as e: print(f" [Warning] Could not read {fpath.name}: {e}") else: with open(input_path, "r", encoding="utf-8", errors="replace") as f: texts = [line.strip() for line in f if line.strip()] print(f"[Loader] Loaded {len(texts):,} lines from {input_path.name}") total_chars = sum(len(t) for t in texts) print(f"[Loader] Total: {len(texts):,} documents | {total_chars:,} characters") return texts # --------------------------------------------------------------------------- # Tokenizer training # --------------------------------------------------------------------------- def train_tokenizer( texts: list[str], vocab_size: int, min_frequency: int, special_tokens: list[str], show_progress: bool = True, ): from tokenizers import Tokenizer from tokenizers.models import BPE from tokenizers.trainers import BpeTrainer from tokenizers.pre_tokenizers import ByteLevel from tokenizers.decoders import ByteLevel as ByteLevelDecoder from tokenizers.normalizers import NFD, Lowercase, StripAccents, Sequence from tokenizers.processors import TemplateProcessing print(f"\n[Trainer] Training BPE tokenizer") print(f" Vocab size : {vocab_size:,}") print(f" Min frequency : {min_frequency}") print(f" Special tokens: {special_tokens}") print(f" Documents : {len(texts):,}\n") tokenizer = Tokenizer(BPE(unk_token="[UNK]")) # Byte-level pre-tokenizer — handles any language, no unknown chars tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True) tokenizer.decoder = ByteLevelDecoder() trainer = BpeTrainer( vocab_size=vocab_size, min_frequency=min_frequency, special_tokens=special_tokens, show_progress=show_progress, ) tokenizer.train_from_iterator(texts, trainer=trainer) # Post-processor: wrap sequences with BOS/EOS bos_id = tokenizer.token_to_id("[BOS]") eos_id = tokenizer.token_to_id("[EOS]") if bos_id is not None and eos_id is not None: tokenizer.post_processor = TemplateProcessing( single="[BOS] $A [EOS]", special_tokens=[("[BOS]", bos_id), ("[EOS]", eos_id)], ) actual_vocab = tokenizer.get_vocab_size() print(f"\n[Trainer] Done — actual vocab size: {actual_vocab:,}") return tokenizer # --------------------------------------------------------------------------- # Export functions # --------------------------------------------------------------------------- def export_hf_json(tokenizer, output_dir: Path) -> str: """HuggingFace fast tokenizer format — load with tokenizers.Tokenizer.from_file()""" path = output_dir / "tokenizer.json" tokenizer.save(str(path)) print(f" [✓] tokenizer.json HuggingFace fast tokenizer") return str(path) def export_zdx_json(tokenizer, output_dir: Path) -> str: """ZDX BAN native format — used directly by train.py""" path = output_dir / "zdx_tokenizer.json" tokenizer.save(str(path)) print(f" [✓] zdx_tokenizer.json ZDX BAN native (train.py compatible)") return str(path) def export_vocab_txt(tokenizer, output_dir: Path) -> str: """Plain vocab list — one token per line, sorted by ID""" vocab = tokenizer.get_vocab() sorted_vocab = sorted(vocab.items(), key=lambda x: x[1]) path = output_dir / "vocab.txt" with open(path, "w", encoding="utf-8") as f: for token, _ in sorted_vocab: f.write(token + "\n") print(f" [✓] vocab.txt Plain vocab list ({len(sorted_vocab):,} tokens)") return str(path) def export_vocab_json(tokenizer, output_dir: Path) -> str: """Token -> ID mapping as JSON""" vocab = tokenizer.get_vocab() path = output_dir / "vocab.json" with open(path, "w", encoding="utf-8") as f: json.dump(vocab, f, ensure_ascii=False, indent=2) print(f" [✓] vocab.json Token→ID mapping ({len(vocab):,} entries)") return str(path) def export_merges(tokenizer, output_dir: Path) -> str: """BPE merge rules in GPT-2 style format""" path = output_dir / "merges.txt" try: model = tokenizer.model merges = model.merges with open(path, "w", encoding="utf-8") as f: f.write("#version: 0.2\n") for merge in merges: f.write(f"{merge[0]} {merge[1]}\n") print(f" [✓] merges.txt BPE merge rules ({len(merges):,} merges)") except Exception: # Fallback: write empty merges file with note with open(path, "w", encoding="utf-8") as f: f.write("#version: 0.2\n# Merge rules embedded in tokenizer.json\n") print(f" [✓] merges.txt (merges embedded in tokenizer.json)") return str(path) def export_config(tokenizer, output_dir: Path, vocab_size: int, min_freq: int, special_tokens: list[str]) -> str: """Tokenizer config metadata""" config = { "tokenizer_class": "ZDXBPETokenizer", "model_type": "bpe", "vocab_size": tokenizer.get_vocab_size(), "target_vocab_size": vocab_size, "min_frequency": min_freq, "special_tokens": {t: tokenizer.token_to_id(t) for t in special_tokens}, "bos_token": "[BOS]", "eos_token": "[EOS]", "unk_token": "[UNK]", "pad_token": "[PAD]", "mask_token": "[MASK]", "do_lower_case": False, "pre_tokenizer": "ByteLevel", "decoder": "ByteLevel", "zdx_version": "1.0.1", } path = output_dir / "tokenizer_config.json" with open(path, "w", encoding="utf-8") as f: json.dump(config, f, indent=2) print(f" [✓] tokenizer_config.json Config and metadata") return str(path) # --------------------------------------------------------------------------- # Quick encode/decode test # --------------------------------------------------------------------------- def run_smoke_test(tokenizer) -> None: test_sentences = [ "The network stack initializes all interfaces.", "Authentication requires valid credentials.", "Branch salience determines which path gets focus.", ] print("\n[Smoke Test] Encoding/decoding verification:") all_passed = True for sentence in test_sentences: enc = tokenizer.encode(sentence) dec = tokenizer.decode(enc.ids) dec_clean = dec.strip().lstrip("▁").strip() match = dec_clean.lower() == sentence.lower() status = "✓" if match else "~" print(f" [{status}] {len(enc.ids):3d} tokens | {sentence[:50]}") if not match: all_passed = False if all_passed: print(" All tests passed.") else: print(" Minor decode differences — normal for byte-level BPE.") # --------------------------------------------------------------------------- # Stats report # --------------------------------------------------------------------------- def print_stats(tokenizer, texts: list[str]) -> None: print("\n[Stats] Tokenizer analysis:") vocab = tokenizer.get_vocab_size() sample = texts[:500] if len(texts) > 500 else texts total_tokens = sum(len(tokenizer.encode(t).ids) for t in sample) total_chars = sum(len(t) for t in sample) fertility = total_tokens / max(len(sample), 1) compression = total_chars / max(total_tokens, 1) print(f" Vocab size : {vocab:,}") print(f" Avg tokens/doc : {fertility:.1f}") print(f" Chars/token : {compression:.2f} (higher = better compression)") # Most and least frequent tokens (by ID — proxy for frequency) vocab_map = tokenizer.get_vocab() sorted_by_id = sorted(vocab_map.items(), key=lambda x: x[1]) print(f" First tokens : {[t for t, _ in sorted_by_id[len(SPECIAL_TOKENS):len(SPECIAL_TOKENS)+8]]}") print(f" Last tokens : {[t for t, _ in sorted_by_id[-5:]]}") # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="ZDX Tokenizer Trainer — BPE tokenizer training CLI", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) parser.add_argument("--input", required=True, help="Input file or directory of text files") parser.add_argument("--output", required=True, help="Output directory for tokenizer files") parser.add_argument("--vocab_size", type=int, default=DEFAULT_VOCAB_SIZE, help=f"Vocabulary size (default: {DEFAULT_VOCAB_SIZE})") parser.add_argument("--min_freq", type=int, default=DEFAULT_MIN_FREQ, help=f"Min token frequency (default: {DEFAULT_MIN_FREQ})") parser.add_argument("--formats", nargs="+", default=DEFAULT_FORMATS, choices=["json", "vocab", "merges", "config", "zdx"], help="Export formats (default: all)") parser.add_argument("--special_tokens", nargs="+", default=SPECIAL_TOKENS, help="Special tokens to include") parser.add_argument("--no_test", action="store_true", help="Skip smoke test") parser.add_argument("--no_stats", action="store_true", help="Skip stats report") args = parser.parse_args() output_dir = Path(args.output) output_dir.mkdir(parents=True, exist_ok=True) # --- Load --- texts = load_texts(args.input) if not texts: print("[Error] No text found in input. Check your path.") return # --- Train --- tokenizer = train_tokenizer( texts=texts, vocab_size=args.vocab_size, min_frequency=args.min_freq, special_tokens=args.special_tokens, ) # --- Export --- print(f"\n[Export] Writing to {output_dir}/") if "json" in args.formats: export_hf_json(tokenizer, output_dir) if "zdx" in args.formats: export_zdx_json(tokenizer, output_dir) if "vocab" in args.formats: export_vocab_txt(tokenizer, output_dir) export_vocab_json(tokenizer, output_dir) if "merges" in args.formats: export_merges(tokenizer, output_dir) if "config" in args.formats: export_config(tokenizer, output_dir, args.vocab_size, args.min_freq, args.special_tokens) # --- Smoke test --- if not args.no_test: run_smoke_test(tokenizer) # --- Stats --- if not args.no_stats: print_stats(tokenizer, texts) print(f"\n[Done] Tokenizer saved to {output_dir}/") print(f"[Load] from tokenizers import Tokenizer; t = Tokenizer.from_file('{output_dir}/tokenizer.json')") print(f"[ZDX] train.py will auto-use {output_dir}/zdx_tokenizer.json if pointed to this dir") if __name__ == "__main__": main()