# /// script
# requires-python = ">=3.13"
# dependencies = [
#     "blobfile==3.3.0",
#     "datasets==5.0.1",
#     "rustbpe==0.1.0",
#     "tiktoken==0.14.0",
#     "tokenizers==0.23.2",
#     "tqdm==4.70.1",
#     "transformers==5.17.0",
# ]
# ///

import marimo

__generated_with = "0.24.0"
app = marimo.App(width="medium", auto_download=["html"])


@app.cell
def _():
    import rustbpe
    import tiktoken
    from tiktoken.load import dump_tiktoken_bpe
    from tqdm import tqdm

    DIGIT_SPLIT_PATTERN = (
        r"[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*"
        r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?"
        r"|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+"
        r"[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?"
        r"|\p{N}"  # split every digit separately
        r"| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*"
        r"|\s*[\r\n]+"
        r"|\s+(?!\S)"
        r"|\s+"
    )

    VOCAB_SIZE = 8000
    DATASETS = [
        ("mlfoundations/dclm-baseline-1.0-parquet", None, "train", "text", 75000),
        ("HuggingFaceFW/finepdfs-edu", "eng_Latn", "train", "text", 7500),
        ("openbmb/UltraData-Code", "UltraData-Code-L2", "cpp", "content", 5000),
        ("openbmb/UltraData-Code", "UltraData-Code-L2", "js", "content", 20000),
        ("openbmb/UltraData-Code", "UltraData-Code-L2", "py", "content", 20000),
        ("openbmb/UltraData-Code", "UltraData-Code-L2", "rust", "content", 5000),
        ("openbmb/UltraData-Math", "UltraData-Math-L2-preview", "train", "content", 15000),
    ]
    return (
        DATASETS,
        DIGIT_SPLIT_PATTERN,
        VOCAB_SIZE,
        dump_tiktoken_bpe,
        rustbpe,
        tiktoken,
        tqdm,
    )


@app.cell
def _(DATASETS):
    from datasets import load_dataset
    import copy

    def interleave(iterables):
        iterators = [iter(i) for i in iterables]
        while any(iterators):
            for it in iterators:
                try:
                    yield next(it)
                except StopIteration:
                    iterators.remove(it)


    iterables = []

    for id, config, split, row, count in DATASETS:
        ds = load_dataset(id, config, split=split, streaming=True).take(count)
        iterables.append(map(lambda x, row=row: x[row], iter(ds)))

    for id, config, split, row, count in DATASETS:
        ds = load_dataset(id, config, split=split, streaming=True).take(count)
        cnt = sum(map(lambda x, row=row: len(x[row].encode("utf-8")), iter(ds)))
        print(f"{id}/{config}/{split}: {cnt / 1e9:.2f} GB")
    return interleave, iterables


@app.cell
def _(DIGIT_SPLIT_PATTERN, VOCAB_SIZE, interleave, iterables, rustbpe, tqdm):
    tokenizer = rustbpe.Tokenizer()
    tokenizer.train_from_iterator(
        tqdm(iter(interleave(iterables))),
        vocab_size=VOCAB_SIZE,
        pattern=DIGIT_SPLIT_PATTERN,
    )
    return (tokenizer,)


@app.cell
def _(tokenizer):
    ids = tokenizer.encode("Hello, world!")
    ids, tokenizer.decode(ids)  # "hello world"
    return (ids,)


@app.cell
def _(VOCAB_SIZE, ids, tiktoken, tokenizer):
    enc = tiktoken.Encoding(
        name="rambley-tokenizer-8k",
        pat_str=tokenizer.get_pattern(),
        mergeable_ranks={bytes(k): v for k, v in tokenizer.get_mergeable_ranks()},
        special_tokens={
            "<|bos|>": 8000,
            "<|eos|>": 8001,
            "<|pad|>": 8002,
            **{f"<|vocab_padding_{i - 3}|>": (VOCAB_SIZE + i) for i in range(3, 192)}
        },
    )

    # Fast inference with tiktoken
    ids2 = enc.encode("Hello, world! 123 + 321 = 432")
    ids2, [enc.decode([id]) for id in ids2], enc.decode(ids)
    return (enc,)


@app.cell
def _(dump_tiktoken_bpe, enc):
    dump_tiktoken_bpe(
        enc._mergeable_ranks,
        "rambley-tokenizer-8k.tiktoken"
    )
    return


@app.cell
def _(enc):
    from transformers.integrations.tiktoken import convert_tiktoken_to_fast
    from transformers import PreTrainedTokenizerFast

    convert_tiktoken_to_fast(enc, "hf_tokenizer_8k")

    tok = PreTrainedTokenizerFast(
        tokenizer_file="hf_tokenizer_8k/tokenizer.json",
        bos_token="<|bos|>",
        eos_token="<|eos|>",
        pad_token="<|pad|>",
        padding_side="left",
        add_bos_token=False,
        add_eos_token=True,
    )

    tok.padding_side = "left"
    tok.add_bos_token = False

    tok.save_pretrained("hf_tokenizer_8k")
    return


@app.cell
def _():
    from transformers import AutoTokenizer

    roundtrip = AutoTokenizer.from_pretrained("hf_tokenizer_8k")
    ids3 = roundtrip.encode("Hello, world! 123 + 321 = 432", add_special_tokens=True)
    ids3, [roundtrip.decode([id]) for id in ids3], roundtrip.decode(ids3)
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()