byte_pair_encoding.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/aburkov/theLMbook/blob/main/byte_pair_encoding.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<div style=\"display: flex; justify-content: center;\">\n",
        "    <div style=\"background-color: #f4f6f7; padding: 15px; width: 80%;\">\n",
        "        <table style=\"width: 100%\">\n",
        "            <tr>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <span style=\"font-size: 14px;\">\n",
        "                        A notebook for <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">The Hundred-Page Language Models Book</a> by Andriy Burkov<br><br>\n",
        "                        Code repository: <a href=\"https://github.com/aburkov/theLMbook\" target=\"_blank\" rel=\"noopener\">https://github.com/aburkov/theLMbook</a>\n",
        "                    </span>\n",
        "                </td>\n",
        "                <td style=\"vertical-align: middle;\">\n",
        "                    <a href=\"https://www.thelmbook.com\" target=\"_blank\" rel=\"noopener\">\n",
        "                        <img src=\"https://thelmbook.com/img/book.png\" width=\"80px\" alt=\"The Hundred-Page Language Models Book\">\n",
        "                    </a>\n",
        "                </td>\n",
        "            </tr>\n",
        "        </table>\n",
        "    </div>\n",
        "</div>"
      ],
      "metadata": {
        "id": "tRiU14bopExj"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training the BPE model\n",
        "\n",
        "Below, we load the data and train the BPE model:"
      ],
      "metadata": {
        "id": "btXb1IGF0r6S"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "yy0zjL_2ouOU",
        "outputId": "68b672d0-4d83-4dfa-bdfa-061bfd699373"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading dataset from https://www.thelmbook.com/data/news...\n",
            "Download completed.\n",
            "Extracting files...\n",
            "Extracted train.txt\n",
            "Extracted test.txt\n",
            "Extraction completed.\n",
            "Training BPE tokenizer...\n",
            "Saving the tokenizer...\n"
          ]
        }
      ],
      "source": [
        "# Import required libraries\n",
        "import os  # For file operations and path handling\n",
        "import urllib.request  # For downloading files\n",
        "import tarfile  # For extracting tar files\n",
        "import pickle  # For saving/loading tokenizer\n",
        "import re  # For regex in merge operations\n",
        "import time  # For timing operations\n",
        "from collections import defaultdict  # For counting tokens and pairs\n",
        "\n",
        "def download_file(url, filename):\n",
        "    \"\"\"\n",
        "    Downloads a file from a URL if it doesn't exist locally.\n",
        "    Prevents redundant downloads by checking file existence.\n",
        "\n",
        "    Args:\n",
        "        url (str): URL to download the file from\n",
        "        filename (str): Local path to save the downloaded file\n",
        "\n",
        "    Returns:\n",
        "        None: Prints status messages about download progress\n",
        "    \"\"\"\n",
        "    # Check if file already exists to avoid re-downloading\n",
        "    if not os.path.exists(filename):\n",
        "        print(f\"Downloading dataset from {url}...\")\n",
        "        urllib.request.urlretrieve(url, filename)\n",
        "        print(\"Download completed.\")\n",
        "    else:\n",
        "        print(f\"{filename} already downloaded.\")\n",
        "\n",
        "def is_within_directory(directory, target):\n",
        "    \"\"\"\n",
        "    Security check to prevent path traversal attacks by verifying target path.\n",
        "    Ensures extracted files remain within the intended directory.\n",
        "\n",
        "    Args:\n",
        "        directory (str): Base directory path to check against\n",
        "        target (str): Target path to validate\n",
        "\n",
        "    Returns:\n",
        "        bool: True if target is within directory, False otherwise\n",
        "    \"\"\"\n",
        "    # Convert both paths to absolute form for comparison\n",
        "    abs_directory = os.path.abspath(directory)\n",
        "    abs_target = os.path.abspath(target)\n",
        "    # Get common prefix to check containment\n",
        "    prefix = os.path.commonprefix([abs_directory, abs_target])\n",
        "    return prefix == abs_directory\n",
        "\n",
        "def safe_extract_tar(tar_file, required_files):\n",
        "    \"\"\"\n",
        "    Safely extracts specific files from a tar archive with security checks.\n",
        "    Prevents path traversal attacks and extracts only required files.\n",
        "\n",
        "    Args:\n",
        "        tar_file (str): Path to the tar archive file\n",
        "        required_files (list): List of filenames to extract\n",
        "\n",
        "    Returns:\n",
        "        None: Extracts files and prints progress\n",
        "\n",
        "    Raises:\n",
        "        Exception: If path traversal attempt is detected\n",
        "    \"\"\"\n",
        "    with tarfile.open(tar_file, \"r:gz\") as tar:\n",
        "        # Perform security check on all archive members\n",
        "        for member in tar.getmembers():\n",
        "            if not is_within_directory('.', member.name):\n",
        "                raise Exception(\"Attempted Path Traversal in Tar File\")\n",
        "\n",
        "        # Extract only the specified files\n",
        "        for member in tar.getmembers():\n",
        "            if any(member.name.endswith(file) for file in required_files):\n",
        "                # Remove path prefix for safety\n",
        "                member.name = os.path.basename(member.name)\n",
        "                tar.extract(member, '.')\n",
        "                print(f\"Extracted {member.name}\")\n",
        "\n",
        "def create_word_generator(filepath):\n",
        "    \"\"\"\n",
        "    Creates a generator that yields words from a text file one at a time.\n",
        "    Memory efficient way to process large text files.\n",
        "\n",
        "    Args:\n",
        "        filepath (str): Path to text file to read\n",
        "\n",
        "    Returns:\n",
        "        generator: Yields individual words from the file\n",
        "    \"\"\"\n",
        "    def generator():\n",
        "        with open(filepath, 'r') as f:\n",
        "            for line in f:\n",
        "                for word in line.split():\n",
        "                    yield word\n",
        "    return generator()\n",
        "\n",
        "def download_and_prepare_data(url):\n",
        "    \"\"\"\n",
        "    Downloads, extracts, and prepares dataset for training.\n",
        "    Handles both downloading and extraction with security checks.\n",
        "\n",
        "    Args:\n",
        "        url (str): URL of the dataset to download\n",
        "\n",
        "    Returns:\n",
        "        generator: Word generator for the training data\n",
        "    \"\"\"\n",
        "    required_files = [\"train.txt\", \"test.txt\"]\n",
        "    filename = os.path.basename(url)\n",
        "\n",
        "    # Download dataset if needed\n",
        "    download_file(url, filename)\n",
        "\n",
        "    # Extract required files if they don't exist\n",
        "    if not all(os.path.exists(file) for file in required_files):\n",
        "        print(\"Extracting files...\")\n",
        "        safe_extract_tar(filename, required_files)\n",
        "        print(\"Extraction completed.\")\n",
        "    else:\n",
        "        print(\"'train.txt' and 'test.txt' already extracted.\")\n",
        "\n",
        "    # Create and return word generator\n",
        "    return create_word_generator(\"train.txt\")\n",
        "\n",
        "def initialize_vocabulary(corpus):\n",
        "    \"\"\"\n",
        "    Creates initial vocabulary from corpus by splitting words into characters.\n",
        "    Adds word boundary marker '_' and tracks unique characters.\n",
        "\n",
        "    Args:\n",
        "        corpus (iterable): Iterator or list of words to process\n",
        "\n",
        "    Returns:\n",
        "        tuple: (vocabulary dict mapping tokenized words to counts,\n",
        "               set of unique characters in corpus)\n",
        "    \"\"\"\n",
        "    # Track word counts and unique characters\n",
        "    vocabulary = defaultdict(int)\n",
        "    charset = set()\n",
        "\n",
        "    for word in corpus:\n",
        "        # Add word boundary marker and split into characters\n",
        "        word_with_marker = '_' + word\n",
        "        characters = list(word_with_marker)\n",
        "        # Update set of unique characters\n",
        "        charset.update(characters)\n",
        "        # Create space-separated string of characters\n",
        "        tokenized_word = \" \".join(characters)\n",
        "        # Increment count for this tokenized word\n",
        "        vocabulary[tokenized_word] += 1\n",
        "\n",
        "    return vocabulary, charset\n",
        "\n",
        "def get_pair_counts(vocabulary):\n",
        "    \"\"\"\n",
        "    Counts frequencies of adjacent symbol pairs in the vocabulary.\n",
        "    Used to identify most common pairs for merging.\n",
        "\n",
        "    Args:\n",
        "        vocabulary (dict): Dictionary mapping tokenized words to their counts\n",
        "\n",
        "    Returns:\n",
        "        defaultdict: Maps token pairs to their frequency counts\n",
        "    \"\"\"\n",
        "    pair_counts = defaultdict(int)\n",
        "    for tokenized_word, count in vocabulary.items():\n",
        "        # Split word into tokens\n",
        "        tokens = tokenized_word.split()\n",
        "        # Count adjacent pairs weighted by word frequency\n",
        "        for i in range(len(tokens) - 1):\n",
        "            pair = (tokens[i], tokens[i + 1])\n",
        "            pair_counts[pair] += count\n",
        "    return pair_counts\n",
        "\n",
        "def merge_pair(vocab, pair):\n",
        "    \"\"\"\n",
        "    Merges all occurrences of a specific symbol pair in the vocabulary.\n",
        "    Uses regex for accurate token boundary matching.\n",
        "\n",
        "    Args:\n",
        "        vocab (dict): Current vocabulary dictionary\n",
        "        pair (tuple): Pair of tokens to merge\n",
        "\n",
        "    Returns:\n",
        "        dict: New vocabulary with specified pair merged\n",
        "    \"\"\"\n",
        "    new_vocab = {}\n",
        "    # Create regex pattern for matching the pair\n",
        "    bigram = re.escape(' '.join(pair))\n",
        "    pattern = re.compile(r\"(?<!\\S)\" + bigram + r\"(?!\\S)\")\n",
        "\n",
        "    # Apply merge to all words in vocabulary\n",
        "    for tokenized_word, count in vocab.items():\n",
        "        new_tokenized_word = pattern.sub(\"\".join(pair), tokenized_word)\n",
        "        new_vocab[new_tokenized_word] = count\n",
        "    return new_vocab\n",
        "\n",
        "def byte_pair_encoding(corpus, vocab_size):\n",
        "    \"\"\"\n",
        "    Implements the BPE algorithm to learn a subword vocabulary.\n",
        "    Iteratively merges most frequent character pairs until target vocabulary size is reached.\n",
        "\n",
        "    Args:\n",
        "        corpus (iterable): Iterator or list of words to learn BPE from\n",
        "        vocab_size (int): Target vocabulary size to stop merging at\n",
        "\n",
        "    Returns:\n",
        "        tuple: (final vocabulary dict, list of merge operations,\n",
        "               set of base characters, set of all tokens)\n",
        "    \"\"\"\n",
        "    # Initialize vocabulary with character-level tokens\n",
        "    vocab, charset = initialize_vocabulary(corpus)\n",
        "    merges = []\n",
        "    tokens = set(charset)\n",
        "\n",
        "    # Keep merging pairs until we reach target vocab size\n",
        "    while len(tokens) < vocab_size:\n",
        "        # Get counts of all adjacent token pairs\n",
        "        pair_counts = get_pair_counts(vocab)\n",
        "        if not pair_counts:\n",
        "            break\n",
        "\n",
        "        # Find and record the most frequent pair\n",
        "        most_frequent_pair = max(pair_counts, key=pair_counts.get)\n",
        "        merges.append(most_frequent_pair)\n",
        "\n",
        "        # Update vocabulary by merging the most frequent pair\n",
        "        vocab = merge_pair(vocab, most_frequent_pair)\n",
        "\n",
        "        # Add the new merged token to our token set\n",
        "        new_token = \"\".join(most_frequent_pair)\n",
        "        tokens.add(new_token)\n",
        "\n",
        "    return vocab, merges, charset, tokens\n",
        "\n",
        "def tokenize_word(word, merges, charset, unk_token=\"<UNK>\"):\n",
        "    \"\"\"\n",
        "    Tokenizes a single word using learned BPE merges.\n",
        "    Handles unknown characters with UNK token.\n",
        "\n",
        "    Args:\n",
        "        word (str): Word to tokenize\n",
        "        merges (list): List of learned merge operations\n",
        "        charset (set): Set of known characters\n",
        "        unk_token (str): Token to use for unknown characters\n",
        "\n",
        "    Returns:\n",
        "        list: List of tokens for the word\n",
        "    \"\"\"\n",
        "    # Add word boundary marker and convert to characters\n",
        "    word = '_' + word\n",
        "    tokens = [char if char in charset else unk_token for char in word]\n",
        "\n",
        "    # Apply merges in order\n",
        "    for left, right in merges:\n",
        "        i = 0\n",
        "        while i < len(tokens) - 1:\n",
        "            if tokens[i:i+2] == [left, right]:\n",
        "                tokens[i:i+2] = [left + right]\n",
        "            else:\n",
        "                i += 1\n",
        "    return tokens\n",
        "\n",
        "def build_merge_map(merges):\n",
        "    \"\"\"\n",
        "    Creates a mapping from token pairs to their merged forms.\n",
        "    Preserves merge order for consistent tokenization.\n",
        "\n",
        "    Args:\n",
        "        merges (list): List of merge operations\n",
        "\n",
        "    Returns:\n",
        "        dict: Maps token pairs to (merged_token, merge_priority) tuples\n",
        "    \"\"\"\n",
        "    merge_map = {}\n",
        "    # Build map with merge priorities\n",
        "    for i, (left, right) in enumerate(merges):\n",
        "        merged_token = left + right\n",
        "        merge_map[(left, right)] = (merged_token, i)\n",
        "    return merge_map\n",
        "\n",
        "def tokenize_word_fast(word, merge_map, vocabulary, charset, unk_token=\"<UNK>\"):\n",
        "    \"\"\"\n",
        "    Optimized tokenization function using pre-computed merge map.\n",
        "    Produces identical results to original algorithm but faster.\n",
        "\n",
        "    Args:\n",
        "        word (str): Word to tokenize\n",
        "        merge_map (dict): Mapping of token pairs to merged forms\n",
        "        vocabulary (dict): Current vocabulary dictionary\n",
        "        charset (set): Set of known characters\n",
        "        unk_token (str): Token to use for unknown characters\n",
        "\n",
        "    Returns:\n",
        "        list: List of tokens for the word\n",
        "    \"\"\"\n",
        "    # Check if word exists in vocabulary as-is\n",
        "    word_with_prefix = '_' + word\n",
        "    if word_with_prefix in vocabulary:\n",
        "        return [word_with_prefix]\n",
        "\n",
        "    # Initialize with characters, replacing unknown ones\n",
        "    tokens = [char if char in charset else unk_token for char in word_with_prefix]\n",
        "\n",
        "    # Keep merging until no more merges possible\n",
        "    while True:\n",
        "        # Find all possible merge operations\n",
        "        pairs_with_positions = []\n",
        "        for i in range(len(tokens) - 1):\n",
        "            pair = (tokens[i], tokens[i + 1])\n",
        "            if pair in merge_map:\n",
        "                merged_token, merge_priority = merge_map[pair]\n",
        "                pairs_with_positions.append((i, pair, merged_token, merge_priority))\n",
        "\n",
        "        # Exit if no more merges possible\n",
        "        if not pairs_with_positions:\n",
        "            break\n",
        "\n",
        "        # Sort by merge priority and position for consistency\n",
        "        pairs_with_positions.sort(key=lambda x: (x[3], x[0]))\n",
        "\n",
        "        # Apply first valid merge\n",
        "        pos, pair, merged_token, _ = pairs_with_positions[0]\n",
        "        tokens[pos:pos+2] = [merged_token]\n",
        "\n",
        "    return tokens\n",
        "\n",
        "def save_tokenizer(merges, charset, tokens, filename=\"tokenizer.pkl\"):\n",
        "    \"\"\"\n",
        "    Saves tokenizer state to a pickle file for later use.\n",
        "\n",
        "    Args:\n",
        "        merges (list): List of merge operations\n",
        "        charset (set): Set of known characters\n",
        "        tokens (set): Set of all tokens\n",
        "        filename (str): Path to save tokenizer state\n",
        "\n",
        "    Returns:\n",
        "        None: Saves tokenizer to disk\n",
        "    \"\"\"\n",
        "    with open(filename, \"wb\") as f:\n",
        "        pickle.dump({\n",
        "            \"merges\": merges,\n",
        "            \"charset\": charset,\n",
        "            \"tokens\": tokens\n",
        "        }, f)\n",
        "\n",
        "def load_tokenizer(filename=\"tokenizer.pkl\"):\n",
        "    \"\"\"\n",
        "    Loads tokenizer state from a pickle file.\n",
        "\n",
        "    Args:\n",
        "        filename (str): Path to saved tokenizer state\n",
        "\n",
        "    Returns:\n",
        "        dict: Dictionary containing tokenizer components\n",
        "    \"\"\"\n",
        "    with open(filename, \"rb\") as f:\n",
        "        return pickle.load(f)\n",
        "\n",
        "# Main function for downloading, training BPE, saving, and loading tokenizer\n",
        "if __name__ == \"__main__\":\n",
        "    # Configuration parameters\n",
        "    vocab_size = 5_000  # Target vocabulary size\n",
        "    max_corpus_size = 500_000  # Maximum number of words to process\n",
        "    data_url = \"https://www.thelmbook.com/data/news\"  # Dataset source\n",
        "\n",
        "    # Download and prepare training data\n",
        "    word_gen = download_and_prepare_data(data_url)\n",
        "\n",
        "    # Collect corpus up to maximum size\n",
        "    word_list = []\n",
        "    for word in word_gen:\n",
        "        word_list.append(word)\n",
        "        if len(word_list) >= max_corpus_size:\n",
        "            break\n",
        "\n",
        "    # Train BPE tokenizer\n",
        "    print(\"Training BPE tokenizer...\")\n",
        "    vocab, merges, charset, tokens = byte_pair_encoding(word_list, vocab_size)\n",
        "\n",
        "    # Save trained tokenizer\n",
        "    print(\"Saving the tokenizer...\")\n",
        "    save_tokenizer(merges, charset, tokens)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Testing the trained BPE tokenizer\n",
        "\n",
        "Once the BPE tokenizer is trained, we can load it and apply to a new text:"
      ],
      "metadata": {
        "id": "NXvPpOgo0Wgk"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y0_RrAhuF7_B",
        "outputId": "678ee2c2-5df4-4598-97e3-159a605fcf96"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading the tokenizer...\n",
            "\n",
            "Sentence tokenized with the straightforward implementation:\n",
            "Let's -> ['_Let', \"'\", 's']\n",
            "proceed -> ['_proceed']\n",
            "to -> ['_to']\n",
            "the -> ['_the']\n",
            "language -> ['_language']\n",
            "modeling -> ['_model', 'ing']\n",
            "part. -> ['_part', '.']\n",
            "--- Elapsed: 0.020586490631103516 seconds ---\n",
            "\n",
            "Sentence tokenized with a fast implementation:\n",
            "Let's -> ['_Let', \"'\", 's']\n",
            "proceed -> ['_proceed']\n",
            "to -> ['_to']\n",
            "the -> ['_the']\n",
            "language -> ['_language']\n",
            "modeling -> ['_model', 'ing']\n",
            "part. -> ['_part', '.']\n",
            "--- Elapsed: 0.004575490951538086 seconds ---\n",
            "\n",
            "Vocabulary size: 5000\n"
          ]
        }
      ],
      "source": [
        "if __name__ == \"__main__\":\n",
        "    print(\"Loading the tokenizer...\")\n",
        "    tokenizer = load_tokenizer()\n",
        "\n",
        "    # Tokenize the sample sentence using the loaded tokenizer\n",
        "    sentence = \"Let's proceed to the language modeling part.\"\n",
        "\n",
        "    start_time = time.time()\n",
        "    tokenized_sentence = [tokenize_word(word, tokenizer[\"merges\"], tokenizer[\"charset\"]) for word in sentence.split()]\n",
        "    elapsed = time.time() - start_time\n",
        "    print(\"\\nSentence tokenized with the straightforward implementation:\")\n",
        "    for word, tokens in zip(sentence.split(), tokenized_sentence):\n",
        "        print(f\"{word} -> {tokens}\")\n",
        "    print(\"--- Elapsed: %s seconds ---\" % (elapsed))\n",
        "\n",
        "    merge_map = build_merge_map(tokenizer[\"merges\"])\n",
        "    start_time = time.time()\n",
        "    fast_tokenized_sentence = [tokenize_word_fast(word, merge_map, vocab, tokenizer[\"charset\"]) for word in sentence.split()]\n",
        "    elapsed = time.time() - start_time\n",
        "    print(\"\\nSentence tokenized with a fast implementation:\")\n",
        "    for word, tokens in zip(sentence.split(), fast_tokenized_sentence):\n",
        "        print(f\"{word} -> {tokens}\")\n",
        "    print(\"--- Elapsed: %s seconds ---\" % (time.time() - start_time))\n",
        "\n",
        "    print(\"\\nVocabulary size:\", len(tokenizer[\"tokens\"]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AFdVkYCAvj2p"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNBe3ByEAsgr9onXQLnGaJz",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}