sampling_method.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/sampling_method.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": "YMLqrwiuulzT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Token sampling method\n",
        "\n",
        "## Method implementation\n",
        "\n",
        "In the cell below, we implement the token sampling method that combines temperature, top-k, and top-p:"
      ],
      "metadata": {
        "id": "kb9Akwe7xttX"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yy0zjL_2ouOU",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "46ce8acf-0762-4e64-fc08-0711e330c5b1"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Test vocabulary: ['the', 'quick', 'brown', 'fox', 'jumps', 'over', 'lazy', 'dog']\n",
            "Initial logits: [ 2.   1.5  1.   0.5  0.  -0.5 -1.  -1.5]\n",
            "\n",
            "Sampling with different parameters:\n",
            "\n",
            "Test 1: Default parameters (temperature=0.7, no top-k/p filtering)\n",
            "Samples: ['the', 'the', 'fox', 'brown', 'brown']\n",
            "\n",
            "Test 2: High temperature (temperature=2.0)\n",
            "Samples: ['jumps', 'quick', 'dog', 'the', 'the']\n",
            "\n",
            "Test 3: Low temperature (temperature=0.2)\n",
            "Samples: ['quick', 'the', 'the', 'the', 'the']\n",
            "\n",
            "Test 4: Top-k filtering (top_k=3)\n",
            "Samples: ['the', 'the', 'quick', 'the', 'the']\n",
            "\n",
            "Test 5: Top-p filtering (top_p=0.9)\n",
            "Samples: ['the', 'the', 'brown', 'the', 'quick']\n",
            "\n",
            "Test 6: Combined filtering (temperature=0.5, top_k=3, top_p=0.9)\n",
            "Samples: ['the', 'the', 'the', 'the', 'the']\n",
            "\n",
            "Error handling examples:\n",
            "Expected error: Mismatch between logits and vocabulary sizes.\n",
            "Expected error: Temperature must be positive.\n"
          ]
        }
      ],
      "source": [
        "import numpy as np\n",
        "\n",
        "def validate_inputs(logits, vocabulary, temperature, top_k, top_p):\n",
        "    \"\"\"\n",
        "    Validate all input parameters for the token sampling process.\n",
        "\n",
        "    Args:\n",
        "        logits (list): Raw model output scores for each token\n",
        "        vocabulary (list): List of all possible tokens\n",
        "        temperature (float): Temperature parameter for logits scaling\n",
        "        top_k (int): Number of highest probability tokens to keep\n",
        "        top_p (float): Cumulative probability threshold for nucleus sampling\n",
        "\n",
        "    Raises:\n",
        "        ValueError: If any parameters are invalid or out of expected ranges\n",
        "    \"\"\"\n",
        "    if len(logits) != len(vocabulary):\n",
        "        raise ValueError(\"Mismatch between logits and vocabulary sizes.\")\n",
        "    if temperature <= 0:\n",
        "        raise ValueError(\"Temperature must be positive.\")\n",
        "    if top_k < 0 or top_k > len(logits):\n",
        "        raise ValueError(\"top_k must be between 0 and len(logits).\")\n",
        "    if not 0 < top_p <= 1:\n",
        "        raise ValueError(\"top_p must be in the range (0, 1].\")\n",
        "\n",
        "def get_token_counts(prev_tokens, vocabulary):\n",
        "    \"\"\"\n",
        "    Count the frequency of each token in the previous generation history.\n",
        "\n",
        "    Args:\n",
        "        prev_tokens (list): Previously generated tokens\n",
        "        vocabulary (list): List of all possible tokens\n",
        "\n",
        "    Returns:\n",
        "        dict: Mapping of token indices to their frequencies\n",
        "    \"\"\"\n",
        "    token_counts = {}\n",
        "    if prev_tokens is not None:\n",
        "        for token in prev_tokens:\n",
        "            if token in vocabulary:\n",
        "                idx = vocabulary.index(token)\n",
        "                token_counts[idx] = token_counts.get(idx, 0) + 1\n",
        "    return token_counts\n",
        "\n",
        "def apply_presence_penalty(logits, token_counts, presence_penalty):\n",
        "    \"\"\"\n",
        "    Apply presence penalty to tokens that have appeared before.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "        token_counts (dict): Mapping of token indices to their frequencies\n",
        "        presence_penalty (float): Fixed penalty to subtract from logits of present tokens\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Modified logits with presence penalty applied\n",
        "\n",
        "    Note:\n",
        "        Unlike frequency penalty, this applies the same penalty regardless of frequency\n",
        "    \"\"\"\n",
        "    if presence_penalty != 0.0:\n",
        "        for idx in token_counts:\n",
        "            logits[idx] -= presence_penalty\n",
        "    return logits\n",
        "\n",
        "def apply_frequency_penalty(logits, token_counts, frequency_penalty):\n",
        "    \"\"\"\n",
        "    Apply frequency penalty proportional to token occurrence count.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "        token_counts (dict): Mapping of token indices to their frequencies\n",
        "        frequency_penalty (float): Penalty factor to multiply with token frequency\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Modified logits with frequency penalty applied\n",
        "\n",
        "    Note:\n",
        "        Penalty increases linearly with token frequency\n",
        "    \"\"\"\n",
        "    if frequency_penalty != 0.0:\n",
        "        for idx, count in token_counts.items():\n",
        "            logits[idx] -= frequency_penalty * count\n",
        "    return logits\n",
        "\n",
        "def apply_temperature(logits, temperature):\n",
        "    \"\"\"\n",
        "    Apply temperature scaling to logits to control randomness.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "        temperature (float): Temperature parameter (>1 increases randomness, <1 decreases it)\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Temperature-scaled and normalized logits\n",
        "\n",
        "    Note:\n",
        "        - Higher temperature makes distribution more uniform\n",
        "        - Lower temperature makes distribution more peaked\n",
        "        - Normalizes by subtracting max for numerical stability\n",
        "    \"\"\"\n",
        "    if temperature != 1.0:\n",
        "        logits = logits / temperature\n",
        "    return logits - np.max(logits)\n",
        "\n",
        "def apply_top_k_filtering(logits, top_k, min_tokens_to_keep=1):\n",
        "    \"\"\"\n",
        "    Apply top-k filtering to keep only the k highest probability tokens.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "        top_k (int): Number of top tokens to keep\n",
        "        min_tokens_to_keep (int): Minimum number of tokens to keep regardless of top-k\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Modified logits with all but top-k tokens set to -inf\n",
        "\n",
        "    Note:\n",
        "        Ensures at least min_tokens_to_keep tokens remain available for sampling\n",
        "    \"\"\"\n",
        "    if top_k > 0:\n",
        "        indices_to_remove = np.argsort(logits)[:-min_tokens_to_keep]\n",
        "        indices_to_keep = np.argsort(logits)[-top_k:]\n",
        "        for idx in indices_to_remove:\n",
        "            if idx not in indices_to_keep:\n",
        "                logits[idx] = float('-inf')\n",
        "    return logits\n",
        "\n",
        "def apply_top_p_filtering(logits, top_p, min_tokens_to_keep=1):\n",
        "    \"\"\"\n",
        "    Apply nucleus (top-p) filtering to keep tokens comprising top p probability mass.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "        top_p (float): Cumulative probability threshold (0 to 1)\n",
        "        min_tokens_to_keep (int): Minimum number of tokens to keep regardless of top-p\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Modified logits with unlikely tokens set to -inf\n",
        "\n",
        "    Note:\n",
        "        1. Converts logits to probabilities\n",
        "        2. Sorts tokens by probability\n",
        "        3. Keeps minimal set of tokens whose cumulative probability >= top_p\n",
        "        4. Ensures at least min_tokens_to_keep tokens remain\n",
        "    \"\"\"\n",
        "    if top_p < 1.0:\n",
        "        probs = np.exp(logits)\n",
        "        probs = probs / probs.sum()\n",
        "\n",
        "        sorted_indices = np.argsort(probs)[::-1]\n",
        "        sorted_probs = probs[sorted_indices]\n",
        "        cumulative_probs = np.cumsum(sorted_probs)\n",
        "\n",
        "        sorted_indices_to_remove = sorted_indices[cumulative_probs > top_p]\n",
        "\n",
        "        if len(sorted_indices_to_remove) > len(sorted_indices) - min_tokens_to_keep:\n",
        "            sorted_indices_to_remove = sorted_indices_to_remove[\n",
        "                :len(sorted_indices) - min_tokens_to_keep\n",
        "            ]\n",
        "\n",
        "        logits[sorted_indices_to_remove] = float('-inf')\n",
        "    return logits\n",
        "\n",
        "def convert_to_probabilities(logits):\n",
        "    \"\"\"\n",
        "    Convert logits to a valid probability distribution using softmax.\n",
        "\n",
        "    Args:\n",
        "        logits (numpy.ndarray): Token logits\n",
        "\n",
        "    Returns:\n",
        "        numpy.ndarray: Probability distribution summing to 1\n",
        "    \"\"\"\n",
        "    probs = np.exp(logits)\n",
        "    return probs / probs.sum()\n",
        "\n",
        "def sample_token(logits, vocabulary, temperature=0.7, top_k=0, top_p=1.0,\n",
        "                repetition_penalty=1.0, presence_penalty=0.0, frequency_penalty=0.0,\n",
        "                prev_tokens=None):\n",
        "    \"\"\"\n",
        "    Main function for sampling the next token using various sampling strategies.\n",
        "    Applies sampling methods in the same order as the transformers library.\n",
        "\n",
        "    Args:\n",
        "        logits (list): Raw model output scores for each token\n",
        "        vocabulary (list): List of all possible tokens\n",
        "        temperature (float): Temperature for logits scaling (default: 0.7)\n",
        "        top_k (int): Number of highest probability tokens to keep (default: 0, disabled)\n",
        "        top_p (float): Cumulative probability threshold for nucleus sampling (default: 1.0)\n",
        "        repetition_penalty (float): Penalty for repeated tokens (default: 1.0, no penalty)\n",
        "        presence_penalty (float): Fixed penalty for token presence (default: 0.0)\n",
        "        frequency_penalty (float): Penalty scaled by token frequency (default: 0.0)\n",
        "        prev_tokens (list): Previously generated tokens (default: None)\n",
        "\n",
        "    Returns:\n",
        "        str: Sampled token from vocabulary\n",
        "\n",
        "    Process:\n",
        "        1. Validate all input parameters\n",
        "        2. Apply repetition, presence, and frequency penalties\n",
        "        3. Apply temperature scaling\n",
        "        4. Apply top-k and top-p filtering\n",
        "        5. Convert to probability distribution and sample\n",
        "    \"\"\"\n",
        "    validate_inputs(logits, vocabulary, temperature, top_k, top_p)\n",
        "\n",
        "    logits = np.array(logits, dtype=np.float64)\n",
        "\n",
        "    # 1. Apply penalties\n",
        "    token_counts = get_token_counts(prev_tokens, vocabulary)\n",
        "    logits = apply_presence_penalty(logits, token_counts, presence_penalty)\n",
        "    logits = apply_frequency_penalty(logits, token_counts, frequency_penalty)\n",
        "\n",
        "    # 2. Apply temperature scaling\n",
        "    logits = apply_temperature(logits, temperature)\n",
        "\n",
        "    # 3. Apply filtering\n",
        "    logits = apply_top_k_filtering(logits, top_k)\n",
        "    logits = apply_top_p_filtering(logits, top_p)\n",
        "\n",
        "    # 4. Convert to probabilities and sample\n",
        "    probabilities = convert_to_probabilities(logits)\n",
        "    return np.random.choice(vocabulary, p=probabilities)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    # Create a test vocabulary and corresponding logits\n",
        "    vocabulary = [\"the\", \"quick\", \"brown\", \"fox\", \"jumps\", \"over\", \"lazy\", \"dog\"]\n",
        "    logits = np.array([2.0, 1.5, 1.0, 0.5, 0.0, -0.5, -1.0, -1.5])\n",
        "\n",
        "    print(\"Test vocabulary:\", vocabulary)\n",
        "    print(\"Initial logits:\", logits)\n",
        "    print(\"\\nSampling with different parameters:\")\n",
        "\n",
        "    # Test 1: Default parameters\n",
        "    print(\"\\nTest 1: Default parameters (temperature=0.7, no top-k/p filtering)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary) for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Test 2: High temperature (more random)\n",
        "    print(\"\\nTest 2: High temperature (temperature=2.0)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary, temperature=2.0) for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Test 3: Low temperature (more deterministic)\n",
        "    print(\"\\nTest 3: Low temperature (temperature=0.2)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary, temperature=0.2) for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Test 4: Top-k filtering\n",
        "    print(\"\\nTest 4: Top-k filtering (top_k=3)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary, top_k=3) for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Test 5: Top-p filtering\n",
        "    print(\"\\nTest 5: Top-p filtering (top_p=0.9)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary, top_p=0.9) for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Test 6: Combined filtering\n",
        "    print(\"\\nTest 6: Combined filtering (temperature=0.5, top_k=3, top_p=0.9)\")\n",
        "    samples = [sample_token(logits.copy(), vocabulary, temperature=0.5, top_k=3, top_p=0.9)\n",
        "              for _ in range(5)]\n",
        "    print(\"Samples:\", samples)\n",
        "\n",
        "    # Demonstrate error handling\n",
        "    print(\"\\nError handling examples:\")\n",
        "    try:\n",
        "        # Test with mismatched sizes\n",
        "        sample_token(logits[:5], vocabulary)\n",
        "    except ValueError as e:\n",
        "        print(\"Expected error:\", e)\n",
        "\n",
        "    try:\n",
        "        # Test with invalid temperature\n",
        "        sample_token(logits, vocabulary, temperature=0)\n",
        "    except ValueError as e:\n",
        "        print(\"Expected error:\", e)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "s61ovCwawq3f"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyNHhXoTui8A3Sg3rEKzvmNi",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}