Contlo_Assignment.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Task 1: Model Implementation and Checkpoints\n"
      ],
      "metadata": {
        "id": "y0sezdS52bEI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "mJ_0T29T42Gg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Importing Libraries and Preliminaries"
      ],
      "metadata": {
        "id": "4OPQaF3KDpQQ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.nn import functional as F\n",
        "\n",
        "# Define constants\n",
        "n_layers = 4\n",
        "embed_dim = 768\n",
        "hidden_dim = embed_dim\n",
        "vocab_size = 50257\n",
        "pad_token_id = 0\n",
        "max_seq_len = 1024\n",
        "\n",
        "# Initialize embedding matrices\n",
        "token_embedding = nn.Embedding(vocab_size, embed_dim)\n",
        "position_embedding = nn.Embedding(max_seq_len, embed_dim)\n",
        "\n",
        "# Positional encoding function\n",
        "def positional_encoding(position, d_model):\n",
        "    def get_position_angle_vec(position, i):\n",
        "        return [\n",
        "            position / np.power(10000, 2 * (i // 2) / d_model),\n",
        "            np.sin(position / np.power(10000, (2 * (i // 2) + 1) / d_model)),\n",
        "        ]\n",
        "    position_enc = np.array([get_position_angle_vec(pos, i) for i in range(d_model) for pos in range(position)])\n",
        "    position_enc[:, 0::2] = np.sin(position_enc[:, 0::2])  # dim 2i\n",
        "    position_enc[:, 1::2] = np.cos(position_enc[:, 1::2])  # dim 2i+1\n",
        "    return torch.from_numpy(position_enc).float().unsqueeze(0)\n"
      ],
      "metadata": {
        "id": "DcfrZ4Sg2fw3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 2. Transformer Decoder Block"
      ],
      "metadata": {
        "id": "fUiXfFA8D8Nb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerDecoderBlock(nn.Module):\n",
        "    def __init__(self, embed_dim, hidden_dim, nhead, dropout=0.1):\n",
        "        super(TransformerDecoderBlock, self).__init__()\n",
        "        self.nhead = nhead\n",
        "        self.self_attn = nn.MultiheadAttention(embed_dim, nhead, dropout=dropout)\n",
        "        self.linear1 = nn.Linear(embed_dim, hidden_dim)\n",
        "        self.linear2 = nn.Linear(hidden_dim, embed_dim)\n",
        "        self.norm1 = nn.LayerNorm(embed_dim)\n",
        "        self.norm2 = nn.LayerNorm(embed_dim)\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x, mask):\n",
        "        x = self.norm1(x)\n",
        "        x, _ = self.self_attn(x, x, x, mask=mask)\n",
        "        x = self.dropout1(x)\n",
        "        x = x + x\n",
        "        x = self.norm2(x)\n",
        "        x = F.relu(self.linear1(x))\n",
        "        x = self.dropout2(x)\n",
        "        x = self.linear2(x)\n",
        "        x = x + x\n",
        "        return x\n"
      ],
      "metadata": {
        "id": "nZDoT8P1DyKY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### 3. GPT2 Model:"
      ],
      "metadata": {
        "id": "Gth_-5OzEOvx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class GPT2Model(nn.Module):\n",
        "    def __init__(self, vocab_size, embed_dim, hidden_dim, n_layers, nhead, dropout=0.1):\n",
        "        super(GPT2Model, self).__init__()\n",
        "        self.token_embedding = token_embedding\n",
        "        self.position_embedding = position_embedding\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.decoder_blocks = nn.ModuleList([TransformerDecoderBlock(embed_dim, hidden_dim, nhead, dropout) for _ in range(n_layers)])\n",
        "        self.to_logits = nn.Linear(embed_dim, vocab_size)\n",
        "\n",
        "    def forward(self, input_ids, sequence_length):\n",
        "        # Prepare embeddings\n",
        "        token_embeddings = self.token_embedding(input_ids)\n",
        "        position_encodings = self.position_embedding(torch.arange(sequence_length, device=input_ids.device))\n",
        "        embeddings = token_embeddings + position_encodings\n",
        "        embeddings = self.dropout(embeddings)\n",
        "\n",
        "        # Layer-wise masking\n",
        "        mask = torch.tril(torch.ones((sequence_length, sequence_length), dtype=torch.bool, device=input_ids.device))\n",
        "\n",
        "        # Stack decoder blocks\n",
        "        for block in self.decoder_blocks:\n",
        "            # Pass embeddings through the block\n",
        "            embeddings = block(embeddings, mask)\n",
        "\n",
        "        # Output logits\n",
        "        final_output = self.dropout(embeddings)\n",
        "        logits = self.to_logits(final_output)\n",
        "\n",
        "        return logits\n",
        "\n"
      ],
      "metadata": {
        "id": "df8dsmxBEXvp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# This code performs the following tasks:\n",
        "'''\n",
        "# 1. Looping through decoder blocks:\n",
        "- Within the loop, we call the `forward` method of each `TransformerDecoderBlock` instance, passing the current embeddings and layer-wise mask.\n",
        "- This allows each block to perform self-attention, feed-forward network, and residual connections with layer normalization based on the current state.\n",
        "\n",
        "# 2. Implementing layer-wise masking:\n",
        "- Before entering the loop, we create a triangular mask tensor using `torch.tril`. This ensures that each layer can only attend to past positions in the sequence, preventing information leakage from the future.\n",
        "\n",
        "# 3. Outputting logits:\n",
        "- After processing through all decoder blocks, we apply dropout to the final output and project it to vocabulary prediction logits using the `to_logits` linear layer.\n",
        "\n",
        "Remember to replace `self.token_embedding` and `self.position_embedding` with your actual implementations based on the provided snippets. This code provides a comprehensive `forward` method for your GPT2 model, incorporating the desired functionalities and adhering to best practices for syntax and structure.\n",
        "\n",
        "I hope this helps! Feel free to ask if you have any further questions or need further clarification on any specific parts of the code.\n",
        "\n",
        "'''"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 145
        },
        "id": "9ZxFEs4zpdaM",
        "outputId": "3e4bb7d6-60da-4c7c-9a8e-fd5b3f16f0ab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'\\n# 1. Looping through decoder blocks:\\n- Within the loop, we call the `forward` method of each `TransformerDecoderBlock` instance, passing the current embeddings and layer-wise mask.\\n- This allows each block to perform self-attention, feed-forward network, and residual connections with layer normalization based on the current state.\\n\\n# 2. Implementing layer-wise masking:\\n- Before entering the loop, we create a triangular mask tensor using `torch.tril`. This ensures that each layer can only attend to past positions in the sequence, preventing information leakage from the future.\\n\\n# 3. Outputting logits:\\n- After processing through all decoder blocks, we apply dropout to the final output and project it to vocabulary prediction logits using the `to_logits` linear layer.\\n\\nRemember to replace `self.token_embedding` and `self.position_embedding` with your actual implementations based on the provided snippets. This code provides a comprehensive `forward` method for your GPT2 model, incorporating the desired functionalities and adhering to best practices for syntax and structure.\\n\\nI hope this helps! Feel free to ask if you have any further questions or need further clarification on any specific parts of the code.\\n\\n'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "0VpfduQYpiWa"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}