File size: 843 Bytes
701cf7d | 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 | """
Train BPE tokenizer on HF datasets
"""
import sys
sys.path.append("src")
from ares.tokenizer.bpe_trainer import BPETrainer, load_texts_from_hf
import os
if __name__ == "__main__":
vocab_size = 128256
# You can swap to smaller for quick demo: e.g., 8192
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--vocab_size", type=int, default=8192)
parser.add_argument("--samples", type=int, default=10000)
parser.add_argument("--output", type=str, default="data/tokenizer.json")
args = parser.parse_args()
os.makedirs(os.path.dirname(args.output), exist_ok=True)
texts = load_texts_from_hf("allenai/c4", num_samples=args.samples)
trainer = BPETrainer(vocab_size=args.vocab_size)
trainer.train(texts, save_path=args.output)
print(f"Tokenizer trained -> {args.output}")
|