from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders
import os

def train_tokenizer(data_folder="dataRaw", vocab_size=8192, save_path="dataProcessed/tokenizer.json"):
    # Collect all text files
    files = []
    for root, _, filenames in os.walk(data_folder):
        for f in filenames:
            if f.endswith(".txt"):
                files.append(os.path.join(root, f))

    if not files:
        raise ValueError("No .txt files found in dataRaw/. Add training text first.")

    # Initialize tokenizer with byte-level pre-tokenization and decoding
    tokenizer = Tokenizer(models.BPE())
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
    tokenizer.decoder = decoders.ByteLevel()

    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=2,
        special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"]
    )

    tokenizer.train(files, trainer)
    tokenizer.save(save_path)

    print(f"Tokenizer trained and saved to {save_path}")

if __name__ == "__main__":
    train_tokenizer()