TinyStories.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "xP6RzX8xRCex"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig\n",
        "model = AutoModelForCausalLM.from_pretrained('roneneldan/TinyStories-1M')\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neo-125M\")\n",
        "prompt = \"Once upon a time there was\"\n",
        "input_ids = tokenizer.encode(prompt, return_tensors=\"pt\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uCA2UP2CRKXV",
        "outputId": "66e0e020-f700-42e4-b8de-dec40b766c60"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "GPTNeoForCausalLM(\n",
            "  (transformer): GPTNeoModel(\n",
            "    (wte): Embedding(50257, 64)\n",
            "    (wpe): Embedding(2048, 64)\n",
            "    (drop): Dropout(p=0.0, inplace=False)\n",
            "    (h): ModuleList(\n",
            "      (0): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (1): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (2): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (3): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (4): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (5): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (6): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "      (7): GPTNeoBlock(\n",
            "        (ln_1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (attn): GPTNeoAttention(\n",
            "          (attention): GPTNeoSelfAttention(\n",
            "            (attn_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (resid_dropout): Dropout(p=0.0, inplace=False)\n",
            "            (k_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (v_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (q_proj): Linear(in_features=64, out_features=64, bias=False)\n",
            "            (out_proj): Linear(in_features=64, out_features=64, bias=True)\n",
            "          )\n",
            "        )\n",
            "        (ln_2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "        (mlp): GPTNeoMLP(\n",
            "          (c_fc): Linear(in_features=64, out_features=256, bias=True)\n",
            "          (c_proj): Linear(in_features=256, out_features=64, bias=True)\n",
            "          (act): NewGELUActivation()\n",
            "          (dropout): Dropout(p=0.0, inplace=False)\n",
            "        )\n",
            "      )\n",
            "    )\n",
            "    (ln_f): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "  )\n",
            "  (lm_head): Linear(in_features=64, out_features=50257, bias=False)\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0mfKzvRtRV4s",
        "outputId": "af9fbfc8-e593-4a04-e3b6-210d8e627912"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
            "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "torch.Size([1, 208])\n",
            "Once upon a time there was a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, shiny rock in the sky. She wanted to touch it, but it was too high.\n",
            "\n",
            "Lily's mommy told her that it was important to be careful and not touch things. Lily didn't want to touch it, so she asked her mommy if she could touch it. Her mommy said yes and they went to the park.\n",
            "\n",
            "When they got there, Lily saw a big, scary dog. The dog was barking and barked loudly. Lily was scared and didn't know what to do. She tried to run away, but she was too fast. She tried to run away, but it was too late. The dog was too fast and it was too fast.\n",
            "\n",
            "Lily was sad and cried. She wished she had listened to her mommy. She wished she had listened to her mommy and never let her to get back.\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# Generate completion\n",
        "output = model.generate(input_ids, max_length = 1000, num_beams=1)\n",
        "print(output.shape)\n",
        "# Decode the completion\n",
        "output_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
        "\n",
        "# Print the generated text\n",
        "print(output_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SBc4AaZSgK-7"
      },
      "source": [
        "**Model Parameters**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "QzHdexABhTEz"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "# Define model parameters\n",
        "d_model = 64\n",
        "nhead = 8\n",
        "d_hid = 256\n",
        "dropout = 0\n",
        "nlayers = 1\n",
        "sequence_length = 10\n",
        "batch_size = 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PpQ2391Ti0HK"
      },
      "source": [
        "**Transformer Block Pytorch API**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pmSNP2DTRkHV",
        "outputId": "d924b4af-f69a-4323-be00-58d71c47712b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "TransformerEncoder(\n",
            "  (layers): ModuleList(\n",
            "    (0): TransformerEncoderLayer(\n",
            "      (self_attn): MultiheadAttention(\n",
            "        (out_proj): NonDynamicallyQuantizableLinear(in_features=64, out_features=64, bias=True)\n",
            "      )\n",
            "      (linear1): Linear(in_features=64, out_features=256, bias=True)\n",
            "      (dropout): Dropout(p=0, inplace=False)\n",
            "      (linear2): Linear(in_features=256, out_features=64, bias=True)\n",
            "      (norm1): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "      (norm2): LayerNorm((64,), eps=1e-05, elementwise_affine=True)\n",
            "      (dropout1): Dropout(p=0, inplace=False)\n",
            "      (dropout2): Dropout(p=0, inplace=False)\n",
            "    )\n",
            "  )\n",
            ")\n"
          ]
        }
      ],
      "source": [
        "# Original Transformer Encoder using PyTorch API\n",
        "encoder_layer = nn.TransformerEncoderLayer(d_model, nhead, d_hid, dropout)\n",
        "transformer_encoder = nn.TransformerEncoder(encoder_layer, nlayers)\n",
        "print(transformer_encoder)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ceFDL6d1hGcU"
      },
      "source": [
        "**Custom Transformer Block**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "id": "4IfjK_P3gT87"
      },
      "outputs": [],
      "source": [
        "# Custom Transformer block implementation\n",
        "class CustomTransformerBlock(nn.Module):\n",
        "    def __init__(self, d_model, nhead, d_hid, dropout=0):\n",
        "        super(CustomTransformerBlock, self).__init__()\n",
        "\n",
        "        # Multi-head self-attention\n",
        "        self.attention = nn.MultiheadAttention(embed_dim=d_model, num_heads=nhead, dropout=dropout)\n",
        "\n",
        "        # Feedforward network\n",
        "        self.feedforward = nn.Sequential(\n",
        "            nn.Linear(d_model, d_hid),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(d_hid, d_model),\n",
        "        )\n",
        "\n",
        "        # Layer normalization\n",
        "        self.norm1 = nn.LayerNorm(d_model)\n",
        "        self.norm2 = nn.LayerNorm(d_model)\n",
        "\n",
        "        # Dropout layers\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Self-attention with residual connection\n",
        "        attn_output, _ = self.attention(x, x, x)\n",
        "        x = x + self.dropout1(attn_output)  # Residual connection\n",
        "        x = self.norm1(x)  # Layer normalization\n",
        "\n",
        "        # Feedforward network with residual connection\n",
        "        ff_output = self.feedforward(x)\n",
        "        x = x + self.dropout2(ff_output)  # Residual connection\n",
        "        x = self.norm2(x)  # Layer normalization\n",
        "\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oqmueALDiRKL"
      },
      "source": [
        "**Copying the weights**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "id": "SI0e9T76hKuW"
      },
      "outputs": [],
      "source": [
        "# Instantiate the custom transformer block\n",
        "custom_transformer_block = CustomTransformerBlock(d_model, nhead, d_hid, dropout)\n",
        "\n",
        "# Copy weights from the original transformer layer to the custom transformer block\n",
        "custom_transformer_block.attention.in_proj_weight = encoder_layer.self_attn.in_proj_weight\n",
        "custom_transformer_block.attention.in_proj_bias = encoder_layer.self_attn.in_proj_bias\n",
        "custom_transformer_block.attention.out_proj.weight = encoder_layer.self_attn.out_proj.weight\n",
        "custom_transformer_block.attention.out_proj.bias = encoder_layer.self_attn.out_proj.bias\n",
        "\n",
        "custom_transformer_block.feedforward[0].weight = encoder_layer.linear1.weight\n",
        "custom_transformer_block.feedforward[0].bias = encoder_layer.linear1.bias\n",
        "custom_transformer_block.feedforward[2].weight = encoder_layer.linear2.weight\n",
        "custom_transformer_block.feedforward[2].bias = encoder_layer.linear2.bias\n",
        "\n",
        "custom_transformer_block.norm1.weight = encoder_layer.norm1.weight\n",
        "custom_transformer_block.norm1.bias = encoder_layer.norm1.bias\n",
        "custom_transformer_block.norm2.weight = encoder_layer.norm2.weight\n",
        "custom_transformer_block.norm2.bias = encoder_layer.norm2.bias"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eo8--h_Eiai_"
      },
      "source": [
        "**Generating Random Input**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "JLNFbFUEiXat"
      },
      "outputs": [],
      "source": [
        "input_tensor = torch.randn(sequence_length, batch_size, d_model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uNHB24JsiiqN"
      },
      "source": [
        "**Pass the input through both transformer blocks**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "id": "FUyFG86hidsZ"
      },
      "outputs": [],
      "source": [
        "original_output = transformer_encoder(input_tensor)\n",
        "custom_output = custom_transformer_block(input_tensor)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zKORBBScoZ4s"
      },
      "source": [
        "**Comparing**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WiUt1Axsim4p",
        "outputId": "8fcd84fb-043d-4132-b3fa-b59e414dfb4c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Original Output shape: torch.Size([10, 1, 64])\n",
            "Custom Output shape: torch.Size([10, 1, 64])\n",
            "The outputs are the same!\n"
          ]
        }
      ],
      "source": [
        "print(\"Original Output shape:\", original_output.shape)\n",
        "print(\"Custom Output shape:\", custom_output.shape)\n",
        "\n",
        "# Check if the outputs are close\n",
        "if torch.allclose(original_output, custom_output, atol=1e-6):\n",
        "    print(\"The outputs are the same!\")\n",
        "else:\n",
        "    print(\"The outputs are different.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w6W0pDbm4O83"
      },
      "source": [
        "**Original Model Layer Outputs**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jpwMm44KoBPZ",
        "outputId": "32cd09b3-a574-4adf-8384-1cf39ec0225d"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Input IDs: tensor([[7454, 2402,  257,  640,  612,  373,  257]])\n",
            "After Embeddings: torch.Size([1, 7, 64])\n",
            "After Dropout: torch.Size([1, 7, 64])\n",
            "After Block 0: torch.Size([1, 7, 64])\n",
            "After Block 1: torch.Size([1, 7, 64])\n",
            "After Block 2: torch.Size([1, 7, 64])\n",
            "After Block 3: torch.Size([1, 7, 64])\n",
            "After Block 4: torch.Size([1, 7, 64])\n",
            "After Block 5: torch.Size([1, 7, 64])\n",
            "After Block 6: torch.Size([1, 7, 64])\n",
            "After Block 7: torch.Size([1, 7, 64])\n",
            "After Final LayerNorm: torch.Size([1, 7, 64])\n",
            "After LM Head: torch.Size([1, 7, 50257])\n",
            "Generated Text:  upon a time, was a little\n"
          ]
        }
      ],
      "source": [
        "from transformers import GPTNeoForCausalLM, AutoTokenizer\n",
        "import torch\n",
        "\n",
        "# Subclass the original model\n",
        "class GPTNeoForCausalLMWithPrint(GPTNeoForCausalLM):\n",
        "    def forward(self, input_ids, **kwargs):\n",
        "        # Extract the transformer and lm_head from the model\n",
        "        transformer = self.transformer\n",
        "        lm_head = self.lm_head\n",
        "\n",
        "        # Print the input IDs\n",
        "        print(\"Input IDs:\", input_ids)\n",
        "\n",
        "        # Pass input through embeddings\n",
        "        hidden_states = transformer.wte(input_ids) + transformer.wpe(torch.arange(input_ids.shape[-1], device=input_ids.device))\n",
        "        print(\"After Embeddings:\", hidden_states.shape)\n",
        "\n",
        "        # Dropout after embeddings\n",
        "        hidden_states = transformer.drop(hidden_states)\n",
        "        print(\"After Dropout:\", hidden_states.shape)\n",
        "\n",
        "        # Pass through each block (attention + MLP)\n",
        "        for i, block in enumerate(transformer.h):\n",
        "            # Ensure that only the hidden states are returned and passed forward\n",
        "            hidden_states, *_ = block(hidden_states, **kwargs)\n",
        "            print(f\"After Block {i}:\", hidden_states.shape)\n",
        "\n",
        "        # Final layer norm\n",
        "        hidden_states = transformer.ln_f(hidden_states)\n",
        "        print(\"After Final LayerNorm:\", hidden_states.shape)\n",
        "\n",
        "        # Pass through language modeling head\n",
        "        logits = lm_head(hidden_states)\n",
        "        print(\"After LM Head:\", logits.shape)\n",
        "\n",
        "        return logits\n",
        "\n",
        "# Load the pretrained model and tokenizer\n",
        "model = GPTNeoForCausalLMWithPrint.from_pretrained('roneneldan/TinyStories-1M')\n",
        "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-neo-125M\")\n",
        "\n",
        "# Define the prompt\n",
        "prompt = \"Once upon a time there was a\"\n",
        "input_ids = tokenizer.encode(prompt, return_tensors=\"pt\")\n",
        "\n",
        "# Generate output\n",
        "output = model(input_ids)\n",
        "\n",
        "# Decode the output to text\n",
        "output_text = tokenizer.decode(output.argmax(dim=-1)[0], skip_special_tokens=True)\n",
        "\n",
        "# Print the generated text\n",
        "print(\"Generated Text:\", output_text)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Tiny Stories Github**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "PvByGGVK14-S"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading the model\n",
            "model loaded...\n",
            "\n",
            "Tom and Mia think of a solution. They say, \"Let's pretend we are wheat. And we can make soup thanked God for us and it bake. We can talk to God, moon.\"\n",
            "\n",
            "They take their shovel and their spades\n",
            "\n",
            "Total generation time: 4.9995 seconds\n",
            "\n",
            "Per-token generation: 0.098029 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "193"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import tiktoken\n",
        "import time\n",
        "\n",
        "\n",
        "# hyperparameters\n",
        "block_size = 64  # what is the maximum context length for predictions?\n",
        "n_embd = 512\n",
        "dropout = 0\n",
        "n_head = 12\n",
        "n_layer = 12\n",
        "# -------------\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "\n",
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size):\n",
        "        super().__init__()\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.register_buffer('tril', torch.tril(\n",
        "            torch.ones(block_size, block_size)))\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # input of size (batch, time-step, channels)\n",
        "        # output of size (batch, time-step, head size)\n",
        "        B, T, C = x.shape\n",
        "        k = self.key(x)   # (B,T,hs)\n",
        "        q = self.query(x)  # (B,T,hs)\n",
        "        # compute attention scores (\"affinities\")\n",
        "        # (B, T, hs) @ (B, hs, T) -> (B, T, T)\n",
        "        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5\n",
        "        wei = wei.masked_fill(\n",
        "            self.tril[:T, :T] == 0, float('-inf'))  # (B, T, T)\n",
        "        wei = F.softmax(wei, dim=-1)  # (B, T, T)\n",
        "        wei = self.dropout(wei)\n",
        "        # perform the weighted aggregation of the values\n",
        "        v = self.value(x)  # (B,T,hs)\n",
        "        out = wei @ v  # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
        "        return out\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
        "\n",
        "    def __init__(self, num_heads, head_size):\n",
        "        super().__init__()\n",
        "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
        "        self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
        "        out = self.dropout(self.proj(out))\n",
        "        return out\n",
        "\n",
        "\n",
        "class FeedFoward(nn.Module):\n",
        "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(dropout),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.net(x)\n",
        "        return out\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.sa = MultiHeadAttention(n_head, head_size)\n",
        "        self.ffwd = FeedFoward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.sa(self.ln1(x))\n",
        "        x = x + self.ffwd(self.ln2(x))\n",
        "        return x\n",
        "\n",
        "\n",
        "class GPTLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # each token directly reads off the logits for the next token from a lookup table\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
        "        self.blocks = nn.Sequential(\n",
        "            *[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(n_embd)  # final layer norm\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "        # better init, not covered in the original GPT video, but important, will cover in followup video\n",
        "        # self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        if isinstance(module, nn.Linear):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "            if module.bias is not None:\n",
        "                torch.nn.init.zeros_(module.bias)\n",
        "        elif isinstance(module, nn.Embedding):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        tok_emb = self.token_embedding_table(idx)  # (B,T,C)\n",
        "        pos_emb = self.position_embedding_table(\n",
        "            torch.arange(T, device=device))  # (T,C)\n",
        "        x = tok_emb + pos_emb  # (B,T,C)\n",
        "        x = self.blocks(x)  # (B,T,C)\n",
        "        x = self.ln_f(x)  # (B,T,C)\n",
        "        logits = self.lm_head(x)  # (B,T,vocab_size)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -block_size:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :]  # becomes (B, C)\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1)  # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "enc = tiktoken.get_encoding(\"gpt2\")\n",
        "vocab_size = enc.n_vocab\n",
        "\n",
        "print(\"Loading the model\")\n",
        "# Load the trained GPT model\n",
        "model = GPTLanguageModel()\n",
        "m = model.to(device)\n",
        "\n",
        "# Load the model with CPU mapping if necessary\n",
        "m.load_state_dict(torch.load(r'./gpt_5_3500.pth', map_location=device))\n",
        "m.eval()\n",
        "print(\"model loaded...\")\n",
        "\n",
        "# Set the initial context for generation\n",
        "text = \"\\n\"  # \"!\"\n",
        "encoded_text = enc.encode(text)\n",
        "context = torch.tensor(encoded_text, dtype=torch.long,\n",
        "                       device=device).unsqueeze(0)\n",
        "# Generate new text\n",
        "start_time = time.time()\n",
        "generated_text = m.generate(context, max_new_tokens=50)[0]\n",
        "end_time = time.time()\n",
        "total_time = end_time - start_time\n",
        "per_token_time = total_time / generated_text.size(0)\n",
        "# Decode the generated text from integer indices to characters\n",
        "decoded_text = enc.decode(generated_text.tolist())\n",
        "\n",
        "# Print the generated text\n",
        "print(decoded_text)\n",
        "print(f\"\\nTotal generation time: {total_time:.4f} seconds\")\n",
        "print(f\"\\nPer-token generation: {per_token_time:.6f} seconds\")\n",
        "open(r'./generated.txt', 'w').write(decoded_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "***Generation by looking at only last 5 tokens***"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 18,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading the model\n",
            "model loaded...\n",
            "Once upon a time there was an octopus alone. \n",
            "\n",
            "The next day, their parent asked them if they could join them.\n",
            "\n",
            "One day, Lucy decided to go on an adventure.\n",
            "Once upon a time, there was a boy called Jack. He\n",
            "\n",
            "Total generation time: 3.3936 seconds\n",
            "\n",
            "Per-token generation: 0.061702 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "205"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import tiktoken\n",
        "import time\n",
        "\n",
        "\n",
        "# hyperparameters\n",
        "block_size = 64  # what is the maximum context length for predictions?\n",
        "n_embd = 512\n",
        "dropout = 0\n",
        "n_head = 12\n",
        "n_layer = 12\n",
        "# -------------\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "\n",
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size):\n",
        "        super().__init__()\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.register_buffer('tril', torch.tril(\n",
        "            torch.ones(block_size, block_size)))\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # input of size (batch, time-step, channels)\n",
        "        # output of size (batch, time-step, head size)\n",
        "        B, T, C = x.shape\n",
        "        k = self.key(x)   # (B,T,hs)\n",
        "        q = self.query(x)  # (B,T,hs)\n",
        "        # compute attention scores (\"affinities\")\n",
        "        # (B, T, hs) @ (B, hs, T) -> (B, T, T)\n",
        "        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5\n",
        "        wei = wei.masked_fill(\n",
        "            self.tril[:T, :T] == 0, float('-inf'))  # (B, T, T)\n",
        "        wei = F.softmax(wei, dim=-1)  # (B, T, T)\n",
        "        wei = self.dropout(wei)\n",
        "        # perform the weighted aggregation of the values\n",
        "        v = self.value(x)  # (B,T,hs)\n",
        "        out = wei @ v  # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
        "        return out\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
        "\n",
        "    def __init__(self, num_heads, head_size):\n",
        "        super().__init__()\n",
        "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
        "        self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
        "        out = self.dropout(self.proj(out))\n",
        "        return out\n",
        "\n",
        "\n",
        "class FeedFoward(nn.Module):\n",
        "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(dropout),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.net(x)\n",
        "        return out\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.sa = MultiHeadAttention(n_head, head_size)\n",
        "        self.ffwd = FeedFoward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.sa(self.ln1(x))\n",
        "        x = x + self.ffwd(self.ln2(x))\n",
        "        return x\n",
        "\n",
        "\n",
        "class GPTLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # each token directly reads off the logits for the next token from a lookup table\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
        "        self.blocks = nn.Sequential(\n",
        "            *[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(n_embd)  # final layer norm\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "        # better init, not covered in the original GPT video, but important, will cover in followup video\n",
        "        # self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        if isinstance(module, nn.Linear):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "            if module.bias is not None:\n",
        "                torch.nn.init.zeros_(module.bias)\n",
        "        elif isinstance(module, nn.Embedding):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        tok_emb = self.token_embedding_table(idx)  # (B,T,C)\n",
        "        pos_emb = self.position_embedding_table(\n",
        "            torch.arange(T, device=device))  # (T,C)\n",
        "        x = tok_emb + pos_emb  # (B,T,C)\n",
        "        x = self.blocks(x)  # (B,T,C)\n",
        "        x = self.ln_f(x)  # (B,T,C)\n",
        "        logits = self.lm_head(x)  # (B,T,vocab_size)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -5:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :]  # becomes (B, C)\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1)  # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "enc = tiktoken.get_encoding(\"gpt2\")\n",
        "vocab_size = enc.n_vocab\n",
        "\n",
        "print(\"Loading the model\")\n",
        "# Load the trained GPT model\n",
        "model = GPTLanguageModel()\n",
        "m = model.to(device)\n",
        "\n",
        "# Load the model with CPU mapping if necessary\n",
        "m.load_state_dict(torch.load(r'./gpt_5_3500.pth', map_location=device))\n",
        "m.eval()\n",
        "print(\"model loaded...\")\n",
        "\n",
        "# Set the initial context for generation\n",
        "text = \"Once upon a time there\"  # \"!\"\n",
        "encoded_text = enc.encode(text)\n",
        "context = torch.tensor(encoded_text, dtype=torch.long,\n",
        "                       device=device).unsqueeze(0)\n",
        "# Generate new text\n",
        "start_time = time.time()\n",
        "generated_text = m.generate(context, max_new_tokens=50)[0]\n",
        "end_time = time.time()\n",
        "total_time = end_time - start_time\n",
        "per_token_time = total_time / generated_text.size(0)\n",
        "# Decode the generated text from integer indices to characters\n",
        "decoded_text = enc.decode(generated_text.tolist())\n",
        "\n",
        "# Print the generated text\n",
        "print(decoded_text)\n",
        "print(f\"\\nTotal generation time: {total_time:.4f} seconds\")\n",
        "print(f\"\\nPer-token generation: {per_token_time:.6f} seconds\")\n",
        "open(r'./generated.txt', 'w').write(decoded_text)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Giving Initial Tokens**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading the model\n",
            "model loaded...\n",
            "Total number of parameters: 89,160,785\n",
            "Initial tokens given to the model: Once upon a time there\n",
            "\n",
            "Generated text\n",
            ":\n",
            "Once upon a time there was a little girl named Lily. She loved to play with the bus. He looks at the number from a number.\n",
            "\n",
            "For Lily and Sam are you at the hospital?\" Tom asked.\n",
            "\n",
            "\"We brought you some cake. We can like it now. Come on, let's go find some more pencils he had printed all the pretty pictures.\n",
            "\n",
            "One day, Bob entered a spade. He got very wet. He ate the disgusting worm and went back outside to play with her friends. One day she named Amy and she could make a new stack of clay. And they might ruin their toys. And so, Mia and her mommy went to her bed to play with her toys. She had many toys, but her favorite show was the big, scary monster. Nobody was afraid of it. She started dreaming of all the tasty meals was full. They were happy to play with the machine and the paw. It was a bird that had hurt her wing. It was a small fox. He was smaller than any other puzzle. He tried to count them carefully. He counted the knobs and gave her little brother a dance. The big birds saw the gloomy clouds and the shiny comet. It was big and red and had shiny eyes. Lily's mommy told her they were going on a trip to do something fun. He noticed a store nearby and asked his mum, \"can we find something fun?\"\n",
            "\n",
            "Lily nodded and said, \"Yes, please. I like carrots and peas too. Mom said yes and showed Sara how to hold the bird and give it a ride and see how it shone.\n",
            "Once upon a time, there was a little girl called Maya. Maya was three years old and very curious. She asked the ant to give to the unknown boy to say no, because food can fall apart if you really don't cut it down again. He asked his mum if he could find out where the end of her journey, she came across a little girl. She was so excited, she quickly ran outside and picked up her wagon. They put it in a pot that was living in the woods.\n",
            "\n",
            "Outside, Sarah saw a big tree. Tim wanted to climb the tree. On the tree, there was something different from the others around. They went for a walk. \"You should always stay. You need to be prepared and do even more goals.\n",
            "Once upon a time, there was a little girl named Lily\n",
            "\n",
            "Total generation time: 27.1018 seconds\n",
            "\n",
            "Per-token generation: 0.053667 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "2091"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import tiktoken\n",
        "import time\n",
        "\n",
        "# [Your model and class definitions remain the same as above]\n",
        "\n",
        "# Device setup\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "# Initialize tokenizer and vocabulary size\n",
        "enc = tiktoken.get_encoding(\"gpt2\")\n",
        "vocab_size = enc.n_vocab\n",
        "\n",
        "# Load the trained GPT model\n",
        "print(\"Loading the model\")\n",
        "model = GPTLanguageModel()\n",
        "m = model.to(device)\n",
        "m.load_state_dict(torch.load(r'./gpt_5_3500.pth', map_location=device))\n",
        "m.eval()\n",
        "print(\"model loaded...\")\n",
        "\n",
        "# Calculate and print the total number of parameters\n",
        "total_params = sum(p.numel() for p in m.parameters())\n",
        "print(f\"Total number of parameters: {total_params:,}\")\n",
        "\n",
        "# Set the initial tokens (first two tokens in this case)\n",
        "initial_text = \"Once upon a time there\"  # Replace this with your desired initial text\n",
        "print(\"Initial tokens given to the model:\",initial_text)\n",
        "encoded_initial_text = enc.encode(initial_text)\n",
        "\n",
        "# Convert the encoded text to a tensor and expand dimensions\n",
        "context = torch.tensor(encoded_initial_text, dtype=torch.long, device=device).unsqueeze(0)\n",
        "\n",
        "# Generate new text starting from the initial context\n",
        "start_time = time.time()\n",
        "generated_text = m.generate(context, max_new_tokens=500)[0]\n",
        "end_time = time.time()\n",
        "\n",
        "# Calculate generation time\n",
        "total_time = end_time - start_time\n",
        "per_token_time = total_time / generated_text.size(0)\n",
        "\n",
        "# Decode the generated text from integer indices to characters\n",
        "decoded_text = enc.decode(generated_text.tolist())\n",
        "\n",
        "# Print the generated text\n",
        "print(\"\\nGenerated text\\n:\")\n",
        "print(decoded_text)\n",
        "print(f\"\\nTotal generation time: {total_time:.4f} seconds\")\n",
        "print(f\"\\nPer-token generation: {per_token_time:.6f} seconds\")\n",
        "open(r'./generated.txt', 'w').write(decoded_text)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Tiny Stories With each Layer Output**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loading the model\n",
            "model loaded...\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ 0.2531,  2.2977,  0.3734,  ..., -1.8480, -0.2665, -1.4731],\n",
            "         [ 0.7287, -0.6131,  3.7439,  ..., -2.2130,  0.1882,  2.0021],\n",
            "         [-1.1138, -1.2037,  2.3367,  ..., -3.9513,  0.0129,  1.4278],\n",
            "         [ 0.5878, -1.1944,  1.2199,  ..., -1.0795, -1.4784, -1.8418],\n",
            "         [-0.8935, -1.2039,  1.6918,  ...,  0.6740,  1.0122,  0.0944]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ 0.2114,  3.2526,  0.5036,  ..., -1.3900,  1.1721, -2.4607],\n",
            "         [ 0.4264, -0.6603,  4.6477,  ..., -1.1506,  1.9482,  0.6394],\n",
            "         [-1.9144, -2.4571,  3.8895,  ..., -2.7650,  1.4118,  0.6514],\n",
            "         [ 0.7076, -0.8243,  2.4960,  ..., -1.3861, -1.8069, -3.8986],\n",
            "         [-1.9059,  0.0485,  1.4259,  ...,  0.3909,  1.3722, -0.3927]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[-0.2764,  3.7850,  2.0733,  ..., -2.7418,  1.1506, -3.9872],\n",
            "         [ 0.5035, -0.7527,  6.2445,  ...,  0.0415,  0.8242, -0.4803],\n",
            "         [-2.6252, -3.0196,  5.5393,  ..., -2.6411,  0.3112, -0.7806],\n",
            "         [ 1.1886, -0.2272,  2.7866,  ..., -2.1562, -2.4092, -5.3747],\n",
            "         [-2.2852,  0.1106, -0.3878,  ...,  0.0394,  1.7513, -0.3460]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[-0.1328,  3.8379,  3.3639,  ..., -3.0727,  1.2684, -6.3575],\n",
            "         [-0.3025,  0.9885,  8.1008,  ..., -0.1767, -0.0599, -1.8987],\n",
            "         [-3.2023, -1.9121,  7.6740,  ..., -3.1161, -0.2301, -0.9830],\n",
            "         [ 0.7470, -0.4807,  3.6125,  ..., -3.9435, -2.4733, -6.9498],\n",
            "         [-1.8399,  2.2642, -2.1975,  ..., -1.5769,  2.6149, -1.6877]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[-2.7633e+00,  4.5058e+00,  4.4318e+00,  ..., -3.8446e+00,\n",
            "           2.0798e+00, -7.7930e+00],\n",
            "         [-2.4720e+00,  2.8318e+00,  1.0524e+01,  ...,  3.6911e-01,\n",
            "           1.5446e+00, -2.0755e+00],\n",
            "         [-3.9713e+00, -1.0079e+00,  1.0548e+01,  ..., -3.2709e+00,\n",
            "           5.1097e-03, -1.7320e+00],\n",
            "         [ 1.1886e+00, -1.4121e+00,  5.8001e+00,  ..., -4.1592e+00,\n",
            "          -2.9018e+00, -8.7180e+00],\n",
            "         [-2.7410e+00,  3.5092e+00, -2.9339e+00,  ..., -2.9421e+00,\n",
            "           4.0167e+00, -1.3429e+00]]], grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -2.5768,   4.3549,   4.3128,  ...,  -2.5498,   1.9797,  -9.7981],\n",
            "         [ -2.7617,   5.5117,   9.6415,  ...,   2.4265,   1.4958,  -3.8546],\n",
            "         [ -5.0163,   0.8849,  11.5064,  ...,  -2.8627,   0.1959,  -3.7363],\n",
            "         [  0.5796,  -1.1885,   5.7366,  ...,  -3.9872,  -3.0956, -11.6961],\n",
            "         [ -4.3260,   3.2235,  -3.2698,  ...,  -3.4140,   4.3001,  -2.3676]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -3.2001,   4.5499,   5.2915,  ...,  -1.5811,   3.2254, -11.3823],\n",
            "         [ -4.9924,   6.8042,  10.4080,  ...,   3.6223,   1.5703,  -4.4035],\n",
            "         [ -8.5637,   1.0723,  13.1593,  ...,  -3.4095,   0.6840,  -4.2672],\n",
            "         [ -1.4978,  -1.4245,   7.4331,  ...,  -3.2755,  -2.6623, -14.0049],\n",
            "         [ -5.1475,   2.3070,  -3.7492,  ...,  -2.8635,   4.1747,  -3.5040]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -3.2831,   6.1675,   6.6215,  ...,  -2.7825,   3.8006, -12.5197],\n",
            "         [ -7.4295,  10.8825,  11.7087,  ...,   3.1829,   2.4061,  -5.2282],\n",
            "         [-10.4972,   3.9545,  16.3143,  ...,  -4.8922,   1.3246,  -5.9609],\n",
            "         [ -1.4655,  -0.9042,   8.9989,  ...,  -4.3896,  -3.2593, -15.5884],\n",
            "         [ -5.2254,   3.0935,  -3.0177,  ...,  -4.5703,   2.6738,  -3.5917]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -3.5320,   6.6794,   5.4685,  ...,  -1.8939,   5.6992, -15.0053],\n",
            "         [ -7.7112,  12.7515,  11.1272,  ...,   1.6682,   3.2953,  -6.6022],\n",
            "         [-11.2867,   4.8982,  18.4161,  ...,  -5.5165,   1.4809,  -6.7956],\n",
            "         [ -1.3124,  -1.2209,   9.0877,  ...,  -3.8605,  -4.0860, -17.6285],\n",
            "         [ -5.2178,   2.4182,  -2.3954,  ...,  -5.1820,   2.7815,  -3.0093]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -2.3602,   4.7446,   4.2091,  ...,  -1.2823,   5.1769, -16.5952],\n",
            "         [ -9.1141,  11.8748,  10.7258,  ...,   1.4716,   2.3166,  -6.2503],\n",
            "         [-11.9457,   4.4016,  19.6680,  ...,  -5.8182,   1.4421,  -5.9769],\n",
            "         [ -0.9021,  -1.9901,   8.3421,  ...,  -4.3105,  -3.4872, -20.3137],\n",
            "         [ -6.1277,   1.5061,  -2.2421,  ...,  -6.0336,   2.3432,  -3.9791]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -2.2391,   4.8992,   2.8078,  ...,  -2.0348,   4.5791, -17.2577],\n",
            "         [ -9.2461,  11.4573,   9.9023,  ...,   0.6644,   0.3171,  -6.8747],\n",
            "         [-11.7683,   5.0704,  21.7951,  ...,  -6.5230,  -1.1412,  -6.2936],\n",
            "         [ -0.5195,  -3.2342,   7.8738,  ...,  -5.9993,  -4.6286, -22.0467],\n",
            "         [ -6.7739,   0.1873,  -1.9666,  ...,  -6.3062,   2.0230,  -5.6012]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "Head output shape: torch.Size([1, 5, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 5, 512])\n",
            "FeedForward output shape: torch.Size([1, 5, 512])\n",
            "block output: tensor([[[ -2.4380,   4.2319,   2.2260,  ...,  -3.1864,   4.1654, -17.6358],\n",
            "         [ -9.1893,  11.4322,   9.0113,  ...,  -1.7265,   0.1570,  -6.3986],\n",
            "         [-12.6885,   4.6338,  21.8540,  ...,  -8.3654,  -2.1457,  -5.8999],\n",
            "         [ -0.1833,  -3.9365,   7.8537,  ...,  -6.6602,  -5.6203, -22.2243],\n",
            "         [ -6.7595,  -0.4873,  -1.3975,  ...,  -6.3261,   1.3857,  -5.7511]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 5, 512])\n",
            "Block output shape: torch.Size([1, 5, 512])\n",
            "Logits shape: torch.Size([1, 5, 50257])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ 0.2531,  2.2977,  0.3734,  ..., -1.8480, -0.2665, -1.4731],\n",
            "         [ 0.7287, -0.6131,  3.7439,  ..., -2.2130,  0.1882,  2.0021],\n",
            "         [-1.1138, -1.2037,  2.3367,  ..., -3.9513,  0.0129,  1.4278],\n",
            "         [ 0.5878, -1.1944,  1.2199,  ..., -1.0795, -1.4784, -1.8418],\n",
            "         [-0.8935, -1.2039,  1.6918,  ...,  0.6740,  1.0122,  0.0944],\n",
            "         [ 0.2729, -1.1600,  1.6028,  ..., -1.1252, -0.5730,  1.1533]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ 0.2114,  3.2526,  0.5036,  ..., -1.3900,  1.1721, -2.4607],\n",
            "         [ 0.4264, -0.6603,  4.6477,  ..., -1.1506,  1.9482,  0.6394],\n",
            "         [-1.9144, -2.4571,  3.8895,  ..., -2.7650,  1.4118,  0.6514],\n",
            "         [ 0.7076, -0.8243,  2.4960,  ..., -1.3861, -1.8069, -3.8986],\n",
            "         [-1.9059,  0.0485,  1.4259,  ...,  0.3909,  1.3722, -0.3927],\n",
            "         [-0.3643, -0.2868,  1.3098,  ..., -0.9707,  0.8565,  0.6860]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[-0.2764,  3.7850,  2.0733,  ..., -2.7418,  1.1506, -3.9872],\n",
            "         [ 0.5035, -0.7527,  6.2445,  ...,  0.0415,  0.8242, -0.4803],\n",
            "         [-2.6252, -3.0196,  5.5393,  ..., -2.6411,  0.3112, -0.7806],\n",
            "         [ 1.1886, -0.2272,  2.7866,  ..., -2.1562, -2.4092, -5.3747],\n",
            "         [-2.2852,  0.1106, -0.3878,  ...,  0.0394,  1.7513, -0.3460],\n",
            "         [-0.1952,  0.6705,  1.6729,  ..., -0.4122,  0.1080, -0.5066]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[-0.1328,  3.8379,  3.3639,  ..., -3.0727,  1.2684, -6.3575],\n",
            "         [-0.3025,  0.9885,  8.1008,  ..., -0.1767, -0.0599, -1.8987],\n",
            "         [-3.2023, -1.9121,  7.6740,  ..., -3.1161, -0.2301, -0.9830],\n",
            "         [ 0.7470, -0.4807,  3.6125,  ..., -3.9435, -2.4733, -6.9498],\n",
            "         [-1.8399,  2.2642, -2.1975,  ..., -1.5769,  2.6149, -1.6877],\n",
            "         [-0.6112,  1.6712,  2.6980,  ..., -1.3952, -0.3088, -1.1539]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[-2.7633e+00,  4.5058e+00,  4.4318e+00,  ..., -3.8446e+00,\n",
            "           2.0798e+00, -7.7930e+00],\n",
            "         [-2.4720e+00,  2.8318e+00,  1.0524e+01,  ...,  3.6911e-01,\n",
            "           1.5446e+00, -2.0755e+00],\n",
            "         [-3.9713e+00, -1.0079e+00,  1.0548e+01,  ..., -3.2709e+00,\n",
            "           5.1097e-03, -1.7320e+00],\n",
            "         [ 1.1886e+00, -1.4121e+00,  5.8001e+00,  ..., -4.1592e+00,\n",
            "          -2.9018e+00, -8.7180e+00],\n",
            "         [-2.7410e+00,  3.5092e+00, -2.9339e+00,  ..., -2.9421e+00,\n",
            "           4.0167e+00, -1.3429e+00],\n",
            "         [-2.1424e+00,  2.4390e+00,  3.0971e+00,  ..., -1.0013e+00,\n",
            "           1.6993e-01, -1.8202e+00]]], grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -2.5768,   4.3549,   4.3128,  ...,  -2.5498,   1.9797,  -9.7981],\n",
            "         [ -2.7617,   5.5117,   9.6415,  ...,   2.4265,   1.4958,  -3.8546],\n",
            "         [ -5.0163,   0.8849,  11.5064,  ...,  -2.8627,   0.1959,  -3.7363],\n",
            "         [  0.5796,  -1.1885,   5.7366,  ...,  -3.9872,  -3.0956, -11.6961],\n",
            "         [ -4.3260,   3.2235,  -3.2698,  ...,  -3.4140,   4.3001,  -2.3676],\n",
            "         [ -2.8717,   2.8406,   2.1238,  ...,   0.3408,  -0.4339,  -3.1699]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -3.2001,   4.5499,   5.2915,  ...,  -1.5811,   3.2254, -11.3823],\n",
            "         [ -4.9924,   6.8042,  10.4080,  ...,   3.6223,   1.5703,  -4.4035],\n",
            "         [ -8.5637,   1.0723,  13.1593,  ...,  -3.4095,   0.6840,  -4.2672],\n",
            "         [ -1.4978,  -1.4245,   7.4331,  ...,  -3.2755,  -2.6623, -14.0049],\n",
            "         [ -5.1475,   2.3070,  -3.7492,  ...,  -2.8635,   4.1747,  -3.5040],\n",
            "         [ -3.8886,   3.4151,   2.9558,  ...,   1.5384,  -0.9142,  -3.5026]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -3.2831,   6.1675,   6.6215,  ...,  -2.7825,   3.8006, -12.5197],\n",
            "         [ -7.4295,  10.8825,  11.7087,  ...,   3.1829,   2.4061,  -5.2282],\n",
            "         [-10.4972,   3.9545,  16.3143,  ...,  -4.8922,   1.3246,  -5.9609],\n",
            "         [ -1.4655,  -0.9042,   8.9989,  ...,  -4.3896,  -3.2593, -15.5884],\n",
            "         [ -5.2254,   3.0935,  -3.0177,  ...,  -4.5703,   2.6738,  -3.5917],\n",
            "         [ -4.8501,   4.4794,   4.6317,  ...,   2.2563,  -2.2159,  -2.2928]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -3.5320,   6.6794,   5.4685,  ...,  -1.8939,   5.6992, -15.0053],\n",
            "         [ -7.7112,  12.7515,  11.1272,  ...,   1.6682,   3.2953,  -6.6022],\n",
            "         [-11.2867,   4.8982,  18.4161,  ...,  -5.5165,   1.4809,  -6.7956],\n",
            "         [ -1.3124,  -1.2209,   9.0877,  ...,  -3.8605,  -4.0860, -17.6285],\n",
            "         [ -5.2178,   2.4182,  -2.3954,  ...,  -5.1820,   2.7815,  -3.0093],\n",
            "         [ -5.1819,   5.4245,   5.9050,  ...,   2.4708,  -1.6325,  -2.2282]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -2.3602,   4.7446,   4.2091,  ...,  -1.2823,   5.1769, -16.5952],\n",
            "         [ -9.1141,  11.8748,  10.7258,  ...,   1.4716,   2.3166,  -6.2503],\n",
            "         [-11.9457,   4.4016,  19.6680,  ...,  -5.8182,   1.4421,  -5.9769],\n",
            "         [ -0.9021,  -1.9901,   8.3421,  ...,  -4.3105,  -3.4872, -20.3137],\n",
            "         [ -6.1277,   1.5061,  -2.2421,  ...,  -6.0336,   2.3432,  -3.9791],\n",
            "         [ -5.3335,   4.5121,   6.2930,  ...,   1.4311,  -2.9301,  -1.8412]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -2.2391,   4.8992,   2.8078,  ...,  -2.0348,   4.5791, -17.2577],\n",
            "         [ -9.2461,  11.4573,   9.9023,  ...,   0.6644,   0.3171,  -6.8747],\n",
            "         [-11.7683,   5.0704,  21.7951,  ...,  -6.5230,  -1.1412,  -6.2936],\n",
            "         [ -0.5195,  -3.2342,   7.8738,  ...,  -5.9993,  -4.6286, -22.0467],\n",
            "         [ -6.7739,   0.1873,  -1.9666,  ...,  -6.3062,   2.0230,  -5.6012],\n",
            "         [ -4.4727,   4.2081,   5.7886,  ...,   0.1192,  -3.9351,  -1.9962]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "Head output shape: torch.Size([1, 6, 42])\n",
            "MultiHeadAttention projected output shape: torch.Size([1, 6, 512])\n",
            "FeedForward output shape: torch.Size([1, 6, 512])\n",
            "block output: tensor([[[ -2.4380,   4.2319,   2.2260,  ...,  -3.1864,   4.1654, -17.6358],\n",
            "         [ -9.1893,  11.4322,   9.0113,  ...,  -1.7265,   0.1570,  -6.3986],\n",
            "         [-12.6885,   4.6338,  21.8540,  ...,  -8.3654,  -2.1457,  -5.8999],\n",
            "         [ -0.1833,  -3.9365,   7.8537,  ...,  -6.6602,  -5.6203, -22.2243],\n",
            "         [ -6.7595,  -0.4873,  -1.3975,  ...,  -6.3261,   1.3857,  -5.7511],\n",
            "         [ -3.9454,   2.8927,   5.6285,  ...,  -0.5010,  -4.0114,  -1.3264]]],\n",
            "       grad_fn=<AddBackward0>) torch.Size([1, 6, 512])\n",
            "Block output shape: torch.Size([1, 6, 512])\n",
            "Logits shape: torch.Size([1, 6, 50257])\n",
            "Once upon a time there was an\n",
            "\n",
            "Total generation time: 0.1377 seconds\n",
            "\n",
            "Per-token generation: 0.019671 seconds\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "29"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import tiktoken\n",
        "import time\n",
        "import math\n",
        "\n",
        "# hyperparameters\n",
        "block_size = 64  # what is the maximum context length for predictions?\n",
        "n_embd = 512\n",
        "dropout = 0\n",
        "n_head = 12\n",
        "n_layer = 12\n",
        "# -------------\n",
        "\n",
        "torch.manual_seed(1337)\n",
        "\n",
        "\n",
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size):\n",
        "        super().__init__()\n",
        "        self.key = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.query = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.value = nn.Linear(n_embd, head_size, bias=False)\n",
        "        self.register_buffer('tril', torch.tril(\n",
        "            torch.ones(block_size, block_size)))\n",
        "\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # input of size (batch, time-step, channels)\n",
        "        # output of size (batch, time-step, head size)\n",
        "        B, T, C = x.shape\n",
        "        k = self.key(x)  # (B,T,hs)\n",
        "        q = self.query(x)  # (B,T,hs)\n",
        "        # compute attention scores (\"affinities\")\n",
        "        # (B, T, hs) @ (B, hs, T) -> (B, T, T)\n",
        "        wei = q @ k.transpose(-2, -1) * k.shape[-1]**-0.5\n",
        "\n",
        "        wei = wei.masked_fill(\n",
        "            self.tril[:T, :T] == 0, float('-inf'))  # (B, T, T)\n",
        "\n",
        "\n",
        "        wei = F.softmax(wei, dim=-1)  # (B, T, T)\n",
        "        wei = self.dropout(wei)\n",
        "\n",
        "        # perform the weighted aggregation of the values\n",
        "        v = self.value(x)  # (B,T,hs)\n",
        "        out = wei @ v  # (B, T, T) @ (B, T, hs) -> (B, T, hs)\n",
        "        print(f\"Head output shape: {out.shape}\")\n",
        "        # print(f\"Key weight shape: {self.key.weight.shape}\")\n",
        "        # print(f\"Query weight shape: {self.query.weight.shape}\")\n",
        "        # print(f\"Value weight shape: {self.value.weight.shape}\")\n",
        "\n",
        "        return out\n",
        "\n",
        "\n",
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
        "\n",
        "    def __init__(self, num_heads, head_size):\n",
        "        super().__init__()\n",
        "        self.heads = nn.ModuleList([Head(head_size) for _ in range(num_heads)])\n",
        "        self.proj = nn.Linear(head_size * num_heads, n_embd)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = torch.cat([h(x) for h in self.heads], dim=-1)\n",
        "        out = self.dropout(self.proj(out))\n",
        "        print(f\"MultiHeadAttention projected output shape: {out.shape}\")\n",
        "        return out\n",
        "\n",
        "\n",
        "class FeedFoward(nn.Module):\n",
        "    \"\"\" a simple linear layer followed by a non-linearity \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(n_embd, 4 * n_embd),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * n_embd, n_embd),\n",
        "            nn.Dropout(dropout),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = self.net(x)\n",
        "        print(f\"FeedForward output shape: {out.shape}\")\n",
        "        return out\n",
        "\n",
        "\n",
        "class Block(nn.Module):\n",
        "    \"\"\" Transformer block: communication followed by computation \"\"\"\n",
        "\n",
        "    def __init__(self, n_embd, n_head):\n",
        "        # n_embd: embedding dimension, n_head: the number of heads we'd like\n",
        "        super().__init__()\n",
        "        head_size = n_embd // n_head\n",
        "        self.sa = MultiHeadAttention(n_head, head_size)\n",
        "        self.ffwd = FeedFoward(n_embd)\n",
        "        self.ln1 = nn.LayerNorm(n_embd)\n",
        "        self.ln2 = nn.LayerNorm(n_embd)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.sa(self.ln1(x))\n",
        "        x = x + self.ffwd(self.ln2(x))\n",
        "        print(\"block output:\", x, x.shape)\n",
        "        print(f\"Block output shape: {x.shape}\")\n",
        "        return x\n",
        "\n",
        "\n",
        "class GPTLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        # each token directly reads off the logits for the next token from a lookup table\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, n_embd)\n",
        "        self.position_embedding_table = nn.Embedding(block_size, n_embd)\n",
        "        self.blocks = nn.Sequential(\n",
        "            *[Block(n_embd, n_head=n_head) for _ in range(n_layer)])\n",
        "        self.ln_f = nn.LayerNorm(n_embd)  # final layer norm\n",
        "        self.lm_head = nn.Linear(n_embd, vocab_size)\n",
        "\n",
        "        # better init, not covered in the original GPT video, but important, will cover in followup video\n",
        "        # self.apply(self._init_weights)\n",
        "\n",
        "    def _init_weights(self, module):\n",
        "        if isinstance(module, nn.Linear):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "            if module.bias is not None:\n",
        "                torch.nn.init.zeros_(module.bias)\n",
        "        elif isinstance(module, nn.Embedding):\n",
        "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
        "\n",
        "    def forward(self, idx, targets=None):\n",
        "        B, T = idx.shape\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        tok_emb = self.token_embedding_table(idx)  # (B,T,C)\n",
        "        pos_emb = self.position_embedding_table(\n",
        "            torch.arange(T, device=device))  # (T,C)\n",
        "\n",
        "        x = tok_emb + pos_emb  # (B,T,C)\n",
        "\n",
        "        x = self.blocks(x)  # (B,T,C)\n",
        "        x = self.ln_f(x)  # (B,T,C)\n",
        "        logits = self.lm_head(x)  # (B,T,vocab_size)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            B, T, C = logits.shape\n",
        "            logits = logits.view(B*T, C)\n",
        "            targets = targets.view(B*T)\n",
        "            loss = F.cross_entropy(logits, targets)\n",
        "\n",
        "        print(f\"Logits shape: {logits.shape}\")\n",
        "        return logits, loss\n",
        "\n",
        "    def generate(self, idx, max_new_tokens):\n",
        "        # idx is (B, T) array of indices in the current context\n",
        "        for _ in range(max_new_tokens):\n",
        "            # crop idx to the last block_size tokens\n",
        "            idx_cond = idx[:, -block_size:]\n",
        "            # get the predictions\n",
        "            logits, loss = self(idx_cond)\n",
        "            # focus only on the last time step\n",
        "            logits = logits[:, -1, :]  # becomes (B, C)\n",
        "\n",
        "            # apply softmax to get probabilities\n",
        "            probs = F.softmax(logits, dim=-1)  # (B, C)\n",
        "            # sample from the distribution\n",
        "            idx_next = torch.multinomial(probs, num_samples=1)  # (B, 1)\n",
        "            # enc = tiktoken.get_encoding(\"gpt2\")\n",
        "            # print(enc.decode([idx_next.item()]))\n",
        "            # append sampled index to the running sequence\n",
        "            idx = torch.cat((idx, idx_next), dim=1)  # (B, T+1)\n",
        "        return idx\n",
        "\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "enc = tiktoken.get_encoding(\"gpt2\")\n",
        "vocab_size = enc.n_vocab\n",
        "\n",
        "print(\"Loading the model\")\n",
        "# Load the trained GPT model\n",
        "model = GPTLanguageModel()\n",
        "m = model.to(device)\n",
        "\n",
        "\n",
        "# Load the model with CPU mapping if necessary\n",
        "m.load_state_dict(torch.load(r'./gpt_5_3500.pth', map_location=device))\n",
        "m.eval()\n",
        "print(\"model loaded...\")\n",
        "\n",
        "\n",
        "# Set the initial context for generation\n",
        "text = \"Once upon a time there\"  # \"!\"\n",
        "encoded_text = enc.encode(text)\n",
        "context = torch.tensor(encoded_text, dtype=torch.long,\n",
        "                       device=device).unsqueeze(0)\n",
        "\n",
        "# Generate new text\n",
        "start_time = time.time()\n",
        "generated_text = m.generate(context, max_new_tokens=2)[0]\n",
        "end_time = time.time()\n",
        "total_time = end_time - start_time\n",
        "per_token_time = total_time / generated_text.size(0)\n",
        "# Decode the generated text from integer indices to characters\n",
        "decoded_text = enc.decode(generated_text.tolist())\n",
        "\n",
        "# Print the generated text\n",
        "print(decoded_text)\n",
        "print(f\"\\nTotal generation time: {total_time:.4f} seconds\")\n",
        "print(f\"\\nPer-token generation: {per_token_time:.6f} seconds\")\n",
        "open(r'./generated.txt', 'w').write(decoded_text)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 16,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Save the gamma_beta_dict to a .pth file\n",
        "# torch.save(gamma_beta_dict, 'gamma_beta_values.pth')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}