"""Count the same four samples with three published tokenizers and with the toy.

Needs the tiktoken package (pip install tiktoken). The first run of each encoding downloads its
vocabulary file, so it needs network access once; later runs use the cache. Run it from this folder:
python3 compare_tiktoken.py
"""
import sys

from bpe import NUM_MERGES, encode, load_corpus, train
from samples import SAMPLES

try:
    import tiktoken
except ImportError:
    print("tiktoken is not installed: pip install tiktoken, then run this again.")
    sys.exit(0)

ENCODINGS = ("gpt2", "cl100k_base", "o200k_base")
GROUPS_OF_3 = r"\p{N}{1,3}"  # the digit rule in the splitting pattern of cl100k_base and o200k_base

toy = train(load_corpus())
published = {}
counters = {f"toy ({NUM_MERGES} merges)": lambda text: len(encode(text, toy))}
for name in ENCODINGS:
    try:
        enc = tiktoken.get_encoding(name)
    except Exception as err:  # first use downloads the vocabulary; no network means no encoding
        print(f"skipped {name}: {type(err).__name__}")
        continue
    published[name] = enc
    counters[f"{name} ({enc.n_vocab:,} ids)"] = lambda text, enc=enc: len(enc.encode(text))
    assert enc.decode(enc.encode(SAMPLES["Korean"])) == SAMPLES["Korean"]

print("== Tokens for each sample (fewer is shorter)")
print(f"  {'tokenizer':26}" + "".join(f"{n:>9}" for n in SAMPLES))
for label, count in counters.items():
    print(f"  {label:26}" + "".join(f"{count(t):>9}" for t in SAMPLES.values()))

print("\n== Bytes per token")
print(f"  {'tokenizer':26}" + "".join(f"{n:>9}" for n in SAMPLES))
for label, count in counters.items():
    print(f"  {label:26}" + "".join(f"{len(t.encode()) / count(t):>9.2f}" for t in SAMPLES.values()))

print("\n== Tokens per word (a word is a run of non-space characters)")
print(f"  {'tokenizer':26}" + "".join(f"{n:>9}" for n in SAMPLES))
for label, count in counters.items():
    print(f"  {label:26}" + "".join(f"{count(t) / len(t.split()):>9.2f}" for t in SAMPLES.values()))

print("\n== Korean sample tokens divided by English sample tokens (two different texts, not translations)")
for label, count in counters.items():
    ratio = count(SAMPLES["Korean"]) / count(SAMPLES["English"])
    print(f"  {label:26}{ratio:>6.2f}")

print("\n== Digits: most in one token on the Numbers sample, and whether the splitting pattern cuts digit runs into groups of at most 3")
for name, enc in published.items():
    pieces = [enc.decode([i]) for i in enc.encode(SAMPLES["Numbers"])]
    most = max(sum(c.isdigit() for c in piece) for piece in pieces)
    print(f"  {name:26}{most:>3}   {GROUPS_OF_3 in getattr(enc, '_pat_str', '')}")

print("\n== The first 12 tokens of each sample under cl100k_base (one repr per token)")
if "cl100k_base" in {label.split(" ")[0] for label in counters}:
    enc = tiktoken.get_encoding("cl100k_base")
    for name, text in SAMPLES.items():
        print(f"  {name:8}", [enc.decode([i]) for i in enc.encode(text)[:12]])