llm_architecture.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BBtkUfsTKNP7",
        "outputId": "0bbee148-e3d1-4e0d-dbad-53b8c83f6637"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting tiktoken\n",
            "  Downloading tiktoken-0.7.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.6 kB)\n",
            "Requirement already satisfied: regex>=2022.1.18 in /usr/local/lib/python3.10/dist-packages (from tiktoken) (2024.5.15)\n",
            "Requirement already satisfied: requests>=2.26.0 in /usr/local/lib/python3.10/dist-packages (from tiktoken) (2.31.0)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.26.0->tiktoken) (2024.7.4)\n",
            "Downloading tiktoken-0.7.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.1 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m10.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: tiktoken\n",
            "Successfully installed tiktoken-0.7.0\n",
            "matplotlib version: 3.7.1\n",
            "torch version: 2.3.1+cu121\n",
            "tiktoken version: 0.7.0\n"
          ]
        }
      ],
      "source": [
        "from importlib.metadata import version\n",
        "!pip install tiktoken\n",
        "import matplotlib\n",
        "import tiktoken\n",
        "import torch\n",
        "\n",
        "print(\"matplotlib version:\", version(\"matplotlib\"))\n",
        "print(\"torch version:\", version(\"torch\"))\n",
        "print(\"tiktoken version:\", version(\"tiktoken\"))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Configuration details for the 124 million parameter GPT-2 model"
      ],
      "metadata": {
        "id": "iNRIjpRbEkFr"
      }
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "2Mh8k3l6EkCk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "GPT_CONFIG_124M = {\n",
        "    \"vocab_size\": 50257,    # Vocabulary size\n",
        "    \"context_length\": 1024, # Context length\n",
        "    \"emb_dim\": 768,         # Embedding dimension\n",
        "    \"n_heads\": 12,          # Number of attention heads\n",
        "    \"n_layers\": 12,         # Number of layers\n",
        "    \"drop_rate\": 0.1,       # Dropout rate\n",
        "    \"qkv_bias\": False       # Query-Key-Value bias\n",
        "}"
      ],
      "metadata": {
        "id": "dXfr9n2JEfdm"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "\n",
        "class DummyGPTModel(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n",
        "        self.pos_emb = nn.Embedding(cfg[\"context_length\"], cfg[\"emb_dim\"])\n",
        "        self.drop_emb = nn.Dropout(cfg[\"drop_rate\"])\n",
        "\n",
        "        # Use a placeholder for TransformerBlock\n",
        "        self.trf_blocks = nn.Sequential(\n",
        "            *[DummyTransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])])\n",
        "\n",
        "        # Use a placeholder for LayerNorm\n",
        "        self.final_norm = DummyLayerNorm(cfg[\"emb_dim\"])\n",
        "        self.out_head = nn.Linear(\n",
        "            cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False\n",
        "        )\n",
        "\n",
        "    def forward(self, in_idx):\n",
        "        batch_size, seq_len = in_idx.shape\n",
        "        tok_embeds = self.tok_emb(in_idx)\n",
        "        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))\n",
        "        x = tok_embeds + pos_embeds\n",
        "        x = self.drop_emb(x)\n",
        "        x = self.trf_blocks(x)\n",
        "        x = self.final_norm(x)\n",
        "        logits = self.out_head(x)\n",
        "        return logits\n",
        "\n",
        "\n",
        "class DummyTransformerBlock(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        # A simple placeholder\n",
        "\n",
        "    def forward(self, x):\n",
        "        # This block does nothing and just returns its input.\n",
        "        return x\n",
        "\n",
        "\n",
        "class DummyLayerNorm(nn.Module):\n",
        "    def __init__(self, normalized_shape, eps=1e-5):\n",
        "        super().__init__()\n",
        "        # The parameters here are just to mimic the LayerNorm interface.\n",
        "\n",
        "    def forward(self, x):\n",
        "        # This layer does nothing and just returns its input.\n",
        "        return x"
      ],
      "metadata": {
        "id": "4l14zbMkEuD0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import tiktoken\n",
        "\n",
        "tokenizer = tiktoken.get_encoding(\"gpt2\")\n",
        "\n",
        "batch = []\n",
        "\n",
        "txt1 = \"Every effort moves you\"\n",
        "txt2 = \"Every day holds a\"\n",
        "\n",
        "batch.append(torch.tensor(tokenizer.encode(txt1)))\n",
        "batch.append(torch.tensor(tokenizer.encode(txt2)))\n",
        "batch = torch.stack(batch, dim=0)\n",
        "print(batch)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kcIC3jyRIIBB",
        "outputId": "7d8ccae3-a138-4cf6-964b-d2c1ebf8c1e3"
      },
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor([[6109, 3626, 6100,  345],\n",
            "        [6109, 1110, 6622,  257]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Normalizing activations with layer normalization"
      ],
      "metadata": {
        "id": "FKE7U1ZkFac_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class LayerNorm(nn.Module):\n",
        "    def __init__(self, emb_dim):\n",
        "        super().__init__()\n",
        "        self.eps = 1e-5\n",
        "        self.scale = nn.Parameter(torch.ones(emb_dim))\n",
        "        self.shift = nn.Parameter(torch.zeros(emb_dim))\n",
        "\n",
        "    def forward(self, x):\n",
        "        mean = x.mean(dim=-1, keepdim=True)\n",
        "        var = x.var(dim=-1, keepdim=True, unbiased=False)\n",
        "        norm_x = (x - mean) / torch.sqrt(var + self.eps)\n",
        "        return self.scale * norm_x + self.shift"
      ],
      "metadata": {
        "id": "X1WZ91yTE2F8"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "feed forward network with GELU activations"
      ],
      "metadata": {
        "id": "nwSWPvgHFWCx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GELU(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "\n",
        "    def forward(self, x):\n",
        "        return 0.5 * x * (1 + torch.tanh(\n",
        "            torch.sqrt(torch.tensor(2.0 / torch.pi)) *\n",
        "            (x + 0.044715 * torch.pow(x, 3))\n",
        "        ))\n",
        "\n",
        "class FeedForward(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.layers = nn.Sequential(\n",
        "            nn.Linear(cfg[\"emb_dim\"], 4 * cfg[\"emb_dim\"]),\n",
        "            GELU(),\n",
        "            nn.Linear(4 * cfg[\"emb_dim\"], cfg[\"emb_dim\"]),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.layers(x)"
      ],
      "metadata": {
        "id": "zG5g8Ji5FR5d"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(GPT_CONFIG_124M[\"emb_dim\"])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZuPUmOEgFlU4",
        "outputId": "e50083a4-08a6-4138-fe55-b4c9febb3cab"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "768\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "shortcut connections"
      ],
      "metadata": {
        "id": "cvZagfftFxwD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):\n",
        "        super().__init__()\n",
        "        assert (d_out % num_heads == 0), \\\n",
        "            \"d_out must be divisible by num_heads\"\n",
        "\n",
        "        self.d_out = d_out\n",
        "        self.num_heads = num_heads\n",
        "        self.head_dim = d_out // num_heads # Reduce the projection dim to match desired output dim\n",
        "\n",
        "        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
        "        self.W_key = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
        "        self.W_value = nn.Linear(d_in, d_out, bias=qkv_bias)\n",
        "        self.out_proj = nn.Linear(d_out, d_out)  # Linear layer to combine head outputs\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.register_buffer(\n",
        "            \"mask\",\n",
        "            torch.triu(torch.ones(context_length, context_length),\n",
        "                       diagonal=1)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        b, num_tokens, d_in = x.shape\n",
        "\n",
        "        keys = self.W_key(x) # Shape: (b, num_tokens, d_out)\n",
        "        queries = self.W_query(x)\n",
        "        values = self.W_value(x)\n",
        "\n",
        "        # We implicitly split the matrix by adding a `num_heads` dimension\n",
        "        # Unroll last dim: (b, num_tokens, d_out) -> (b, num_tokens, num_heads, head_dim)\n",
        "        keys = keys.view(b, num_tokens, self.num_heads, self.head_dim)\n",
        "        values = values.view(b, num_tokens, self.num_heads, self.head_dim)\n",
        "        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)\n",
        "\n",
        "        # Transpose: (b, num_tokens, num_heads, head_dim) -> (b, num_heads, num_tokens, head_dim)\n",
        "        keys = keys.transpose(1, 2)\n",
        "        queries = queries.transpose(1, 2)\n",
        "        values = values.transpose(1, 2)\n",
        "\n",
        "        # Compute scaled dot-product attention (aka self-attention) with a causal mask\n",
        "        attn_scores = queries @ keys.transpose(2, 3)  # Dot product for each head\n",
        "\n",
        "        # Original mask truncated to the number of tokens and converted to boolean\n",
        "        mask_bool = self.mask.bool()[:num_tokens, :num_tokens]\n",
        "\n",
        "        # Use the mask to fill attention scores\n",
        "        attn_scores.masked_fill_(mask_bool, -torch.inf)\n",
        "\n",
        "        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)\n",
        "        attn_weights = self.dropout(attn_weights)\n",
        "\n",
        "        # Shape: (b, num_tokens, num_heads, head_dim)\n",
        "        context_vec = (attn_weights @ values).transpose(1, 2)\n",
        "\n",
        "        # Combine heads, where self.d_out = self.num_heads * self.head_dim\n",
        "        context_vec = context_vec.contiguous().view(b, num_tokens, self.d_out)\n",
        "        context_vec = self.out_proj(context_vec) # optional projection\n",
        "\n",
        "        return context_vec\n"
      ],
      "metadata": {
        "id": "kBCQxZUqGs0v"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.att = MultiHeadAttention(\n",
        "            d_in=cfg[\"emb_dim\"],\n",
        "            d_out=cfg[\"emb_dim\"],\n",
        "            context_length=cfg[\"context_length\"],\n",
        "            num_heads=cfg[\"n_heads\"],\n",
        "            dropout=cfg[\"drop_rate\"],\n",
        "            qkv_bias=cfg[\"qkv_bias\"])\n",
        "        self.ff = FeedForward(cfg)\n",
        "        self.norm1 = LayerNorm(cfg[\"emb_dim\"])\n",
        "        self.norm2 = LayerNorm(cfg[\"emb_dim\"])\n",
        "        self.drop_shortcut = nn.Dropout(cfg[\"drop_rate\"])\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Shortcut connection for attention block\n",
        "        shortcut = x\n",
        "        x = self.norm1(x)\n",
        "        x = self.att(x)  # Shape [batch_size, num_tokens, emb_size]\n",
        "        x = self.drop_shortcut(x)\n",
        "        x = x + shortcut  # Add the original input back\n",
        "\n",
        "        # Shortcut connection for feed forward block\n",
        "        shortcut = x\n",
        "        x = self.norm2(x)\n",
        "        x = self.ff(x)\n",
        "        x = self.drop_shortcut(x)\n",
        "        x = x + shortcut  # Add the original input back\n",
        "\n",
        "        return x"
      ],
      "metadata": {
        "id": "OuIIEMbUFyto"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "GPT model"
      ],
      "metadata": {
        "id": "9d3KBATvHukh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "class GPTModel(nn.Module):\n",
        "    def __init__(self, cfg):\n",
        "        super().__init__()\n",
        "        self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n",
        "        self.pos_emb = nn.Embedding(cfg[\"context_length\"], cfg[\"emb_dim\"])\n",
        "        self.drop_emb = nn.Dropout(cfg[\"drop_rate\"])\n",
        "\n",
        "        self.trf_blocks = nn.Sequential(\n",
        "            *[TransformerBlock(cfg) for _ in range(cfg[\"n_layers\"])])\n",
        "\n",
        "        self.final_norm = LayerNorm(cfg[\"emb_dim\"])\n",
        "        self.out_head = nn.Linear(\n",
        "            cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False\n",
        "        )\n",
        "\n",
        "    def forward(self, in_idx):\n",
        "        batch_size, seq_len = in_idx.shape\n",
        "        tok_embeds = self.tok_emb(in_idx)\n",
        "        pos_embeds = self.pos_emb(torch.arange(seq_len, device=in_idx.device))\n",
        "        x = tok_embeds + pos_embeds  # Shape [batch_size, num_tokens, emb_size]\n",
        "        x = self.drop_emb(x)\n",
        "        x = self.trf_blocks(x)\n",
        "        x = self.final_norm(x)\n",
        "        logits = self.out_head(x)\n",
        "        return logits"
      ],
      "metadata": {
        "id": "eOV5aGuvD1FB"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.manual_seed(123)\n",
        "model = GPTModel(GPT_CONFIG_124M)\n",
        "\n",
        "out = model(batch)\n",
        "print(\"Input batch:\\n\", batch)\n",
        "print(\"\\nOutput shape:\", out.shape)\n",
        "print(out)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ilbL3gacEU4f",
        "outputId": "86b26154-5aa4-4cc7-a5ae-974948c9a5ac"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Input batch:\n",
            " tensor([[6109, 3626, 6100,  345],\n",
            "        [6109, 1110, 6622,  257]])\n",
            "\n",
            "Output shape: torch.Size([2, 4, 50257])\n",
            "tensor([[[ 0.1381,  0.0077, -0.1963,  ..., -0.0222, -0.1060,  0.1717],\n",
            "         [ 0.3865, -0.8408, -0.6564,  ..., -0.5163,  0.2369, -0.3357],\n",
            "         [ 0.6989, -0.1829, -0.1631,  ...,  0.1472, -0.6504, -0.0056],\n",
            "         [-0.4290,  0.1669, -0.1258,  ...,  1.1579,  0.5303, -0.5549]],\n",
            "\n",
            "        [[ 0.1094, -0.2894, -0.1467,  ..., -0.0557,  0.2911, -0.2824],\n",
            "         [ 0.0882, -0.3552, -0.3527,  ...,  1.2930,  0.0053,  0.1898],\n",
            "         [ 0.6091,  0.4702, -0.4094,  ...,  0.7688,  0.3787, -0.1974],\n",
            "         [-0.0612, -0.0737,  0.4751,  ...,  1.2463, -0.3834,  0.0609]]],\n",
            "       grad_fn=<UnsafeViewBackward0>)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Generating text"
      ],
      "metadata": {
        "id": "UKjEcNhFInhL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_text_simple(model, idx, max_new_tokens, context_size):\n",
        "    # idx is (batch, n_tokens) array of indices in the current context\n",
        "    for _ in range(max_new_tokens):\n",
        "\n",
        "        # Crop current context if it exceeds the supported context size\n",
        "        # E.g., if LLM supports only 5 tokens, and the context size is 10\n",
        "        # then only the last 5 tokens are used as context\n",
        "        idx_cond = idx[:, -context_size:]\n",
        "\n",
        "        # Get the predictions\n",
        "        with torch.no_grad():\n",
        "            logits = model(idx_cond)\n",
        "\n",
        "        # Focus only on the last time step\n",
        "        # (batch, n_tokens, vocab_size) becomes (batch, vocab_size)\n",
        "        logits = logits[:, -1, :]\n",
        "\n",
        "        # Apply softmax to get probabilities\n",
        "        probas = torch.softmax(logits, dim=-1)  # (batch, vocab_size)\n",
        "\n",
        "        # Get the idx of the vocab entry with the highest probability value\n",
        "        idx_next = torch.argmax(probas, dim=-1, keepdim=True)  # (batch, 1)\n",
        "\n",
        "        # Append sampled index to the running sequence\n",
        "        idx = torch.cat((idx, idx_next), dim=1)  # (batch, n_tokens+1)\n",
        "\n",
        "    return idx"
      ],
      "metadata": {
        "id": "jSycHTBFIWLE"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "testing\n"
      ],
      "metadata": {
        "id": "guR5ShBKI_-x"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "start_context = \"Hello, I am\"\n",
        "\n",
        "encoded = tokenizer.encode(start_context)\n",
        "print(\"encoded:\", encoded)\n",
        "\n",
        "encoded_tensor = torch.tensor(encoded).unsqueeze(0)\n",
        "print(\"encoded_tensor.shape:\", encoded_tensor.shape)\n",
        "model.eval() # disable dropout\n",
        "\n",
        "out = generate_text_simple(\n",
        "    model=model,\n",
        "    idx=encoded_tensor,\n",
        "    max_new_tokens=6,\n",
        "    context_size=GPT_CONFIG_124M[\"context_length\"]\n",
        ")\n",
        "\n",
        "print(\"Output:\", out)\n",
        "print(\"Output length:\", len(out[0]))\n",
        "decoded_text = tokenizer.decode(out.squeeze(0).tolist())\n",
        "print(decoded_text)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u-eo0TrsI28B",
        "outputId": "45dfb26f-54ac-4419-94a0-f58738e87658"
      },
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "encoded: [15496, 11, 314, 716]\n",
            "encoded_tensor.shape: torch.Size([1, 4])\n",
            "Output: tensor([[15496,    11,   314,   716, 27018, 24086, 47843, 30961, 42348,  7267]])\n",
            "Output length: 10\n",
            "Hello, I am Featureiman Byeswickattribute argue\n"
          ]
        }
      ]
    }
  ]
}