# Code examples from: Build Your Own LLM, Lesson 2: Write a BPE Tokenizer in Python
# https://www.atulvij.com/llm-course-2-bpe-tokenizer-python
# Review and adapt before running; examples use sample values.

# ---- Example 1 (python) ----
text = "नमस्ते Excel! Excel formulas are excellent."
ids = list(text.encode("utf-8"))
print(len(text), "characters ->", len(ids), "bytes")
print(ids[:12])

# ---- Example 2 (python) ----
from collections import Counter

def get_pairs(ids):
    """Count how often each adjacent pair appears."""
    return Counter(zip(ids, ids[1:]))

def merge(ids, pair, new_id):
    """Replace every occurrence of `pair` with `new_id`."""
    out, i = [], 0
    while i < len(ids):
        if i < len(ids) - 1 and (ids[i], ids[i + 1]) == pair:
            out.append(new_id)
            i += 2
        else:
            out.append(ids[i])
            i += 1
    return out

# ---- Example 3 (python) ----
def train(text, vocab_size):
    ids = list(text.encode("utf-8"))
    merges = {}                              # (a, b) -> new token id
    for new_id in range(256, vocab_size):
        pairs = get_pairs(ids)
        if not pairs:
            break
        pair = pairs.most_common(1)[0][0]
        ids = merge(ids, pair, new_id)
        merges[pair] = new_id
    return merges

corpus = open("my_text.txt", encoding="utf-8").read()   # any few hundred KB of text
merges = train(corpus, vocab_size=512)                  # learn 256 merges

# ---- Example 4 (python) ----
def encode(text, merges):
    ids = list(text.encode("utf-8"))
    while len(ids) >= 2:
        pairs = get_pairs(ids)
        # the pair that was learned earliest has the lowest new id
        pair = min(pairs, key=lambda p: merges.get(p, float("inf")))
        if pair not in merges:
            break                              # nothing left to merge
        ids = merge(ids, pair, merges[pair])
    return ids

# ---- Example 5 (python) ----
def build_vocab(merges):
    vocab = {i: bytes([i]) for i in range(256)}
    for (a, b), idx in merges.items():         # dicts keep insertion order = merge order
        vocab[idx] = vocab[a] + vocab[b]
    return vocab

def decode(ids, vocab):
    return b"".join(vocab[i] for i in ids).decode("utf-8", errors="replace")

vocab = build_vocab(merges)
sample = "Excel formulas are excellent"
tokens = encode(sample, merges)
print(len(sample.encode("utf-8")), "bytes ->", len(tokens), "tokens")
print([vocab[t].decode("utf-8", errors="replace") for t in tokens])
assert decode(tokens, vocab) == sample

# ---- Example 6 (python) ----
for idx in range(256, 276):
    print(idx, repr(vocab[idx].decode("utf-8", errors="replace")))