"""
© KAND CA 2026 - train a 32K BPE tokenizer balanced across English and Arabic.

The existing custom_llama_tokenizer (32K) was fit on Arabic-only text (uonlp/CulturaX, ar). It
works on English by falling back toward byte-level fragments - functional, but roughly 2x
tokens/word versus a tokenizer that actually learned English subwords, measured on the router
corpus. This script fixes that at the source: interleave English and Arabic documents into the
SAME training stream before the BPE merges are learned, so common English and Arabic subwords
both earn vocabulary slots on equal footing, rather than English being an afterthought squeezed
into merges fit for Arabic.

Same sources as the two pretraining corpora this tokenizer is meant to serve:
  Arabic  - kaust-generative-ai/fineweb-edu-ar   (what tokens_20B_sep.bin is built from)
  English - HuggingFaceFW/fineweb-edu            (what build_tokens_en.py streams)

5M documents each, alternating one-for-one, matching the original tokenizer's 10M-document
training budget so this is a like-for-like replacement, not a smaller or larger fit.

    python train_tokenizer_bilingual.py
"""
from itertools import islice

from datasets import load_dataset
from tokenizers import ByteLevelBPETokenizer
from tqdm import tqdm

N_PER_LANG = 5_000_000
OUT_PREFIX = "bilingual32k_tokenizer"

ar = load_dataset("kaust-generative-ai/fineweb-edu-ar", "ar", split="train", streaming=True)
en = load_dataset("HuggingFaceFW/fineweb-edu", "sample-10BT", split="train", streaming=True)


def get_training_corpus():
    ar_it = (ex["text"] for ex in ar if (ex.get("text") or "").strip())
    en_it = (ex["text"] for ex in en if (ex.get("text") or "").strip())
    for a, e in tqdm(islice(zip(ar_it, en_it), N_PER_LANG), total=N_PER_LANG,
                     desc="[*] feeding AR+EN pairs"):
        yield a
        yield e


tokenizer = ByteLevelBPETokenizer()

print("[*] training bilingual tokenizer...")
tokenizer.train_from_iterator(
    get_training_corpus(),
    vocab_size=32_000,
    min_frequency=2,
    show_progress=True,
    special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"],
)

tokenizer.save_model(".", OUT_PREFIX)
print(f"[+] saved -> {OUT_PREFIX}-vocab.json / {OUT_PREFIX}-merges.txt")