"""
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}")