# pip install tokenizers from __future__ import annotations import argparse import json import os import random import sys from dataclasses import dataclass from typing import List, Tuple from tokenizers import Tokenizer from tokenizers.models import WordPiece from tokenizers.trainers import WordPieceTrainer from tokenizers.normalizers import NFKC, Lowercase, Sequence as NormalizerSequence from tokenizers.pre_tokenizers import BertPreTokenizer SPECIAL_TOKENS_DEFAULT = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] @dataclass class TrainingConfig: input_path: str output_dir: str vocab_size: int = 30000 min_frequency: int = 2 lowercase: bool = True sample_size: int = 5000 # number of lines to sample for report stats seed: int = 42 def read_lines(path: str) -> List[str]: with open(path, "r", encoding="utf-8") as f: return [line.rstrip("\n") for line in f] def ensure_dir(path: str) -> None: os.makedirs(path, exist_ok=True) def build_tokenizer(config: TrainingConfig, special_tokens: List[str]) -> Tokenizer: tokenizer = Tokenizer(WordPiece(unk_token="[UNK]")) # Normalization: NFKC + optional lowercase (align with corpus cleaning) normalizers = [NFKC()] if config.lowercase: normalizers.append(Lowercase()) tokenizer.normalizer = NormalizerSequence(normalizers) # Pre-tokenization: BERT-style whitespace + punctuation handling tokenizer.pre_tokenizer = BertPreTokenizer() trainer = WordPieceTrainer( vocab_size=config.vocab_size, min_frequency=config.min_frequency, special_tokens=special_tokens, continuing_subword_prefix="##", show_progress=True, ) tokenizer.train(files=[config.input_path], trainer=trainer) return tokenizer def save_vocab_txt(tokenizer: Tokenizer, path: str) -> None: # get_vocab returns token->id; we need id order ascending per line index contract vocab = tokenizer.get_vocab() by_id: List[Tuple[int, str]] = sorted(((idx, tok) for tok, idx in vocab.items()), key=lambda x: x[0]) with open(path, "w", encoding="utf-8") as f: for _, tok in by_id: f.write(tok + "\n") def write_special_tokens_map(path: str, special_tokens: List[str]) -> None: # Map common keys for downstream libs mapping = { "pad_token": "[PAD]" if "[PAD]" in special_tokens else None, "unk_token": "[UNK]", "cls_token": "[CLS]", "sep_token": "[SEP]", "mask_token": "[MASK]", } with open(path, "w", encoding="utf-8") as f: json.dump(mapping, f, ensure_ascii=False, indent=2) def write_tokenizer_config(path: str, cfg: TrainingConfig, special_tokens: List[str]) -> None: data = { "model_type": "wordpiece", "lowercase": cfg.lowercase, "vocab_size": cfg.vocab_size, "min_frequency": cfg.min_frequency, "special_tokens": special_tokens, "continuing_subword_prefix": "##", "corpus": os.path.abspath(cfg.input_path), } with open(path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2) def compute_report(tokenizer: Tokenizer, lines: List[str]) -> dict: if not lines: return {} total_tokens = 0 total_chars = 0 unk_count = 0 most_common = {} for line in lines: enc = tokenizer.encode(line) tokens = enc.tokens total_tokens += len(tokens) total_chars += len(line) for t in tokens: most_common[t] = most_common.get(t, 0) + 1 if t == "[UNK]": unk_count += 1 avg_tokens_per_document = total_tokens / max(1, len(lines)) avg_chars_per_token = (total_chars / max(1, total_tokens)) if total_tokens > 0 else 0.0 coverage_rate = 1.0 - (unk_count / max(1, total_tokens)) # Top-20 most common tokens top20 = sorted(most_common.items(), key=lambda x: x[1], reverse=True)[:20] return { "avg_tokens_per_document": avg_tokens_per_document, "avg_chars_per_token": avg_chars_per_token, "coverage_rate": coverage_rate, "most_common_tokens": top20, "total_lines_sampled": len(lines), "total_tokens_sampled": total_tokens, "unk_tokens_sampled": unk_count, } def main() -> None: parser = argparse.ArgumentParser(description="Train a WordPiece tokenizer (BERT-style)") parser.add_argument("--input", required=True, help="Path to corpus txt (one document per line)") parser.add_argument("--output_dir", required=True, help="Directory to save tokenizer artifacts") parser.add_argument("--vocab_size", type=int, default=30000, help="Target vocabulary size") parser.add_argument("--min_frequency", type=int, default=2, help="Minimum token frequency to keep") parser.add_argument("--sample_size", type=int, default=5000, help="Number of lines sampled for report") parser.add_argument("--seed", type=int, default=42, help="Random seed for sampling") parser.add_argument("--lowercase", action="store_true", help="Lowercase text during normalization") args = parser.parse_args() cfg = TrainingConfig( input_path=args.input, output_dir=args.output_dir, vocab_size=args.vocab_size, min_frequency=args.min_frequency, lowercase=bool(args.lowercase), sample_size=args.sample_size, seed=args.seed, ) if not os.path.isfile(cfg.input_path): print(f"❌ Input file not found: {cfg.input_path}", file=sys.stderr) sys.exit(1) ensure_dir(cfg.output_dir) print("🏁 Training WordPiece tokenizer...") tokenizer = build_tokenizer(cfg, SPECIAL_TOKENS_DEFAULT) # Save artifacts tokenizer_json = os.path.join(cfg.output_dir, "tokenizer.json") vocab_txt = os.path.join(cfg.output_dir, "vocab.txt") tok_cfg_json = os.path.join(cfg.output_dir, "tokenizer_config.json") special_map_json = os.path.join(cfg.output_dir, "special_tokens_map.json") report_json = os.path.join(cfg.output_dir, "tokenizer_report.json") tokenizer.save(tokenizer_json) save_vocab_txt(tokenizer, vocab_txt) write_tokenizer_config(tok_cfg_json, cfg, SPECIAL_TOKENS_DEFAULT) write_special_tokens_map(special_map_json, SPECIAL_TOKENS_DEFAULT) # Build quick report on a sample subset print("📊 Building tokenizer report (sample)...") lines = read_lines(cfg.input_path) random.Random(cfg.seed).shuffle(lines) sampled = lines[: min(cfg.sample_size, len(lines))] report = { "vocab_size": len(tokenizer.get_vocab()), "num_special_tokens": len(SPECIAL_TOKENS_DEFAULT), "lowercase": cfg.lowercase, "statistics": compute_report(tokenizer, sampled), "artifacts": { "tokenizer_json": os.path.abspath(tokenizer_json), "vocab_txt": os.path.abspath(vocab_txt), "tokenizer_config": os.path.abspath(tok_cfg_json), "special_tokens_map": os.path.abspath(special_map_json), }, } with open(report_json, "w", encoding="utf-8") as f: json.dump(report, f, ensure_ascii=False, indent=2) # Print brief summary print("✅ Tokenizer training done.") print(f" Vocab size: {report['vocab_size']}") print(f" Lowercase: {report['lowercase']}") stats = report["statistics"] or {} if stats: print(f" Coverage rate: {stats.get('coverage_rate', 0.0):.4f}") print(f" Avg tokens/doc: {stats.get('avg_tokens_per_document', 0.0):.2f}") print(f" Avg chars/token: {stats.get('avg_chars_per_token', 0.0):.2f}") print(f" Saved to: {os.path.abspath(cfg.output_dir)}") if __name__ == "__main__": main()