| """Train a ByteLevel BPE tokenizer from UTF-8 .txt files.""" |
| import argparse |
| from pathlib import Path |
| from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors |
|
|
| SPECIAL = ["<|pad|>", "<|unk|>", "<|bos|>", "<|eos|>", "<|system|>", "<|user|>", "<|assistant|>", "<|end|>"] |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--input", nargs="+", required=True, help="Text files or folders containing .txt") |
| ap.add_argument("--output", default="tokenizer.json") |
| ap.add_argument("--vocab-size", type=int, default=32000) |
| args = ap.parse_args() |
| files = [] |
| for f in args.input: |
| p = Path(f) |
| files.extend(sorted(p.rglob("*.txt")) if p.is_dir() else [p]) |
| if not files: |
| raise ValueError("No .txt files found") |
| tok = Tokenizer(models.BPE(unk_token="<|unk|>")) |
| tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) |
| tok.decoder = decoders.ByteLevel() |
| tok.train([str(f) for f in files], trainers.BpeTrainer( |
| vocab_size=args.vocab_size, min_frequency=2, |
| special_tokens=SPECIAL, initial_alphabet=pre_tokenizers.ByteLevel.alphabet(), |
| show_progress=True, |
| )) |
| |
| if tok.get_vocab_size() != args.vocab_size: |
| raise ValueError(f"Tokenizer has {tok.get_vocab_size()} tokens instead of {args.vocab_size}. Add more text or adjust config.json vocab_size.") |
| tok.save(args.output) |
| print(f"Saved {args.output}; vocab={tok.get_vocab_size()}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
