byte_pair_encoding.py

import sys

def main():
    if len(sys.argv) < 3:
        print(f"Usage: {sys.argv[0]} <filename> <vocab_size>")
        return

    filename = sys.argv[1]
    try:
        target_vocab_size = int(sys.argv[2])
    except ValueError:
        print("Error: vocab_size must be an integer")
        return

    try:
        with open(filename, "rb") as f:
            tokens = list(f.read())
    except FileNotFoundError:
        print(f"Error: Could not open {filename}")
        return

    print(f"Starting token count: {len(tokens)}")
    print("Starting encoding in Python...")
    print("|----------------------|")

    current_vocab = 256

    while current_vocab < target_vocab_size:
        stats = {}
        for i in range(len(tokens) - 1):
            pair = (tokens[i], tokens[i+1])
            stats[pair] = stats.get(pair, 0) + 1

        if not stats:
            break

        best_pair = max(stats, key=stats.get)
        max_freq = stats[best_pair]
        if max_freq < 2:
            break

        new_token_id = current_vocab
        current_vocab += 1

        print(f"vocab: {current_vocab:<5} / {target_vocab_size} | merging: {best_pair} -> {new_token_id} | count: {max_freq}", flush=True)
        new_tokens = []
        i = 0
        while i < len(tokens):
            if i < len(tokens) - 1 and tokens[i] == best_pair[0] and tokens[i+1] == best_pair[1]:
                new_tokens.append(new_token_id)
                i += 2
            else:
                new_tokens.append(tokens[i])
                i += 1
        tokens = new_tokens

    print("|----------------------|")
    print("Encoding Complete.")
    print(f"Final Vocab Size: {current_vocab}")
    print(f"Final Token Count: {len(tokens)}")

if __name__ == "__main__":
    main()