File size: 1,599 Bytes
7da1a03 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | """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,
))
# For tiny corpora the trainer may produce fewer tokens than requested.
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()
|