"""Talent to Cash: original byte-pair demonstration. Python 3 standard library.
Run: python byte-bpe.py
This demonstrates merges and reversible UTF-8 bytes, not a trained GPT tokenizer.
"""
from collections import Counter


def merge(tokens, pair):
    result, index = [], 0
    while index < len(tokens):
        if index + 1 < len(tokens) and (tokens[index], tokens[index + 1]) == pair:
            result.append(tokens[index] + tokens[index + 1])
            index += 2
        else:
            result.append(tokens[index])
            index += 1
    return result


if __name__ == "__main__":
    for text in ["banana", "hello hello", "你好🙂你好"]:
        tokens = [bytes([byte]) for byte in text.encode("utf-8")]
        initial = len(tokens)
        for _ in range(3):
            pairs = Counter(zip(tokens, tokens[1:]))
            if not pairs:
                break
            # Counter preserves first-occurrence ordering to resolve equal counts.
            pair = max(pairs, key=pairs.get)
            tokens = merge(tokens, pair)
            assert b"".join(tokens).decode("utf-8") == text
        print(repr(text), "tokens:", initial, "->", len(tokens), tokens)