print("[*] Loading libraries...")
from datasets import load_dataset, interleave_datasets
from tokenizers import ByteLevelBPETokenizer
from tqdm import tqdm

ds_fw   = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True)
ds_cosm = load_dataset("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", split="train", streaming=True)
mixed = interleave_datasets([ds_fw, ds_cosm], probabilities=[0.70, 0.30],
                            seed=42, stopping_strategy="all_exhausted")
def get_training_corpus():
    it = iter(mixed)
    for _ in tqdm(range(80_000), desc="Feeding"):
        yield next(it)["text"]

tokenizer = ByteLevelBPETokenizer()
print("[*] Training tokenizer...")
tokenizer.train_from_iterator(
    get_training_corpus(),
    vocab_size=4096, min_frequency=2, show_progress=True,
    special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
)
tokenizer.save_model(".", "custom_llama_tokenizer")
print("[*] Done.")