Copy_of_UMC301_Assignment_3_3_2.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/suhas-09/100-Days-Of-ML-Code/blob/master/Copy_of_UMC301_Assignment_3_3_2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-rk-9UZY_xe2"
      },
      "source": [
        "# UMC 301: Applied Data Science and Artificial Intelligence\n",
        "## Assignment 3\n",
        "\n",
        "### Submission instructions:\n",
        "\n",
        "1.   The assignment is to be submitted in ONE single notebook.\n",
        "2.   Submit the .ipynb file with all cells open and the pdf for part 2 through Teams Assignment.\n",
        "3. If your IISc email ID is < username > @iisc.ac.in, then name the file < username >_Assgn_3. E.g. jonathan_Assgn_3 for email ID jonathan@iisc.ac.in.\n",
        "4. Before submission, execute the ’Restart session and run all’ option from the Runtime/Kernel tab. Verify that there are no errors and that you are getting the output you expect.\n",
        "5. Use the dataset: https://github.com/taivop/joke-dataset\n",
        "6. The assignment is divided into two questions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 88,
      "metadata": {
        "id": "Kjq_uF_aId9i",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "47570bfc-a900-4aa7-c44a-000cb5bcf8d3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.46.2)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.16.1)\n",
            "Requirement already satisfied: huggingface-hub<1.0,>=0.23.2 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.26.2)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.26.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2024.9.11)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.32.3)\n",
            "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.4.5)\n",
            "Requirement already satisfied: tokenizers<0.21,>=0.20 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.20.3)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.66.6)\n",
            "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (2024.9.0)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.23.2->transformers) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2024.8.30)\n",
            "Requirement already satisfied: datasets in /usr/local/lib/python3.10/dist-packages (3.1.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.16.1)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.26.4)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (17.0.0)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (2.2.2)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.6)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.5.0)\n",
            "Requirement already satisfied: multiprocess<0.70.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.9.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.11.2)\n",
            "Requirement already satisfied: huggingface-hub>=0.23.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.26.2)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (24.2)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.2)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (2.4.3)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.1.0)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (0.2.0)\n",
            "Requirement already satisfied: yarl<2.0,>=1.17.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.17.2)\n",
            "Requirement already satisfied: async-timeout<6.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.23.0->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.2)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
          ]
        }
      ],
      "source": [
        "# Necessary imports\n",
        "\n",
        "!pip install transformers\n",
        "!pip install datasets\n",
        "\n",
        "import math\n",
        "import torch\n",
        "import random\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer, DataCollatorWithPadding\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "from datetime import datetime\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bma59Wpv_xe3"
      },
      "source": [
        "## Question 1\n",
        "\n",
        "### Part 1\n",
        "\n",
        "Load the dataset, create train, validation and test splits. Preprocess the dataset as required to train a decoder model. (10 marks)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 89,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GaVdYGbTIgOZ",
        "outputId": "e90bbdb2-9de5-4f1e-ce8f-c1a397f2c2e2"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "DatasetDict({\n",
            "    train: Dataset({\n",
            "        features: ['body'],\n",
            "        num_rows: 115567\n",
            "    })\n",
            "    test: Dataset({\n",
            "        features: ['body'],\n",
            "        num_rows: 23113\n",
            "    })\n",
            "    validation: Dataset({\n",
            "        features: ['body'],\n",
            "        num_rows: 15410\n",
            "    })\n",
            "})\n",
            "my father told me a joke. how many germans does it take screw in a lightbulb? he said nein | my dads jokes are the wurst i tell you.\n",
            "what do you call a porno set in space. | apollo 13 inches\n",
            "moron | q why did the moron throw the butter out the window?a he wanted to see a butterfly.\n",
            "five reasons not to use an electric toilet | number two will shock you\n",
            "what is punctuation's favorite curry? | l a\n",
            "what is another name for a jewish pokemon trainer? | ash.\n",
            "i asked god for a bike, but some bastard stole it. | it's a good thing that i didn't take off the gps tracker from whoever had owned it first.\n",
            "today someone was killed with a starter pistol. | police think it might be race related.\n",
            "what is long and hard that a australian bride gets in her wedding night? | dick\n",
            "what is it called when two little people get divorced? | daworfed\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import requests\n",
        "\n",
        "# Now load the dataset from the downloaded JSON file, specifying splits\n",
        "# Define split proportions (70%, 15%, 25%)\n",
        "train_size = 0.75\n",
        "val_size = 0.15\n",
        "test_size = 0.10\n",
        "\n",
        "# Load the entire dataset first\n",
        "dataset = load_dataset(\"json\", data_files='./data/reddit_jokes.json')\n",
        "\n",
        "# add the title column to the start of the body column\n",
        "dataset = dataset.map(lambda x: {\"body\": x[\"title\"] + \" | \" + x[\"body\"]})\n",
        "\n",
        "# Remove the title column\n",
        "dataset = dataset.remove_columns([\"title\"])\n",
        "\n",
        "# Remove all the bad parts\n",
        "dataset = dataset.map(lambda x: {\"body\": x[\"body\"].replace(\"\\n\", \"\")})\n",
        "dataset = dataset.map(lambda x: {\"body\": x[\"body\"].replace(\"\\r\", \"\")})\n",
        "dataset = dataset.map(lambda x: {\"body\": \"\".join([c if c.isalnum() or c in [\"?\", \",\", \".\", \"'\",\"\\\"\", \"|\"] else \" \" for c in x[\"body\"]])})\n",
        "\n",
        "\n",
        "dataset = dataset.map(lambda x: {\"body\": \" \".join(x[\"body\"].split())})\n",
        "\n",
        "# Remove the datapoints which have more than 48 words, to reduce \"noise\"\n",
        "dataset = dataset.filter(lambda x: len(x[\"body\"].split()) < 48)\n",
        "\n",
        "# Make all lowercase\n",
        "dataset = dataset.map(lambda x: {\"body\": x[\"body\"].lower()})\n",
        "\n",
        "#Remove the score and id columns\n",
        "dataset = dataset.remove_columns([\"score\", \"id\"])\n",
        "\n",
        "# Then, create the splits using the `train_test_split` method\n",
        "dataset = dataset[\"train\"].train_test_split(\n",
        "    test_size=val_size + test_size,  # Combine validation and test size\n",
        "    seed=42,  # Set a seed for reproducibility\n",
        ")\n",
        "\n",
        "# Further split the test set into validation and test sets\n",
        "dataset[\"validation\"] = dataset[\"test\"].train_test_split(\n",
        "    test_size=test_size / (val_size + test_size),  # Proportion of test within test+val\n",
        "    seed=42,  # Set a seed for reproducibility\n",
        ")[\"test\"]\n",
        "dataset[\"test\"] = dataset[\"test\"].train_test_split(\n",
        "    test_size=test_size / (val_size + test_size),  # Proportion of test within test+val\n",
        "    seed=42,  # Set a seed for reproducibility\n",
        ")[\"train\"]\n",
        "\n",
        "print(dataset)\n",
        "# Just checking if the preprocessing is fine\n",
        "for i in range(10):\n",
        "    print(dataset[\"train\"][i][\"body\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 90,
      "metadata": {
        "id": "Fi0hrA_aL5sr"
      },
      "outputs": [],
      "source": [
        "# Tokenizer\n",
        "checkpoint = 'distilbert-base-cased'\n",
        "tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
        "\n",
        "# Data collator\n",
        "data_collator = DataCollatorWithPadding(tokenizer)\n",
        "\n",
        "tokenizer.model_max_length = 48"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 91,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 49,
          "referenced_widgets": [
            "f6fa8ace0ccf495aacfac69c3e2f386c",
            "44608c4bb86747d9af10af9054b08864",
            "d491f9941bb44a6d93e5c1ac585f4e79",
            "5da80ad644884e39a2a99dde12038129",
            "1213ae2aca2942b3aba61ee546c8cf2d",
            "3f48e3aedf7a49c29d6f58531df2e93c",
            "40845da59b0643dda0c84c3c15e985b9",
            "6c990ad034a34506a7f360da64e58477",
            "ace3cba5ecf144909ba947315a521dd7",
            "e8b6bf87e560457abc822e0ec4bf63ad",
            "497bf3af26c94ab5bb5204e6f87904a0"
          ]
        },
        "id": "M_No8U-aNhjG",
        "outputId": "b1a454a3-9f49-431d-ed39-3de320ecd8ac"
      },
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "Map:   0%|          | 0/15410 [00:00<?, ? examples/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "f6fa8ace0ccf495aacfac69c3e2f386c"
            }
          },
          "metadata": {}
        }
      ],
      "source": [
        "tokenized_datasets = dataset.map(lambda x: tokenizer(x[\"body\"], truncation=True, padding=True))\n",
        "data_collator = DataCollatorWithPadding(tokenizer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 92,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "e_k5L4nsOfz8",
        "outputId": "e618ee79-5460-4681-8476-005f4ab8383f"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "DatasetDict({\n",
              "    train: Dataset({\n",
              "        features: ['input_ids', 'attention_mask'],\n",
              "        num_rows: 115567\n",
              "    })\n",
              "    test: Dataset({\n",
              "        features: ['input_ids', 'attention_mask'],\n",
              "        num_rows: 23113\n",
              "    })\n",
              "    validation: Dataset({\n",
              "        features: ['input_ids', 'attention_mask'],\n",
              "        num_rows: 15410\n",
              "    })\n",
              "})"
            ]
          },
          "metadata": {},
          "execution_count": 92
        }
      ],
      "source": [
        "#Remove the body\n",
        "tokenized_datasets = tokenized_datasets.remove_columns(\"body\")\n",
        "tokenized_datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 93,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "x1x7N0W3L5ss",
        "outputId": "e5eee1d6-bc3b-4653-e1af-c7ac72f4037d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{'input_ids': [101, 1139, 1401, 1500, 1143, 170, 8155, 119, 1293, 1242, 176, 14170, 1116, 1674, 1122, 1321, 13084, 1107, 170, 1609, 27515, 1830, 136, 1119, 1163, 24928, 1394, 197, 1139, 4153, 1116, 13948, 1132, 1103, 192, 7719, 1204, 178, 1587, 1128, 119, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n",
            "{'input_ids': [101, 1184, 1202, 1128, 1840, 170, 185, 8456, 1186, 1383, 1107, 2000, 119, 197, 170, 23043, 2858, 1492, 4519, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n",
            "{'input_ids': [101, 182, 14824, 1179, 197, 186, 1725, 1225, 1103, 182, 14824, 1179, 4932, 1103, 13742, 1149, 1103, 2487, 136, 170, 1119, 1458, 1106, 1267, 170, 11057, 119, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n",
            "{'input_ids': [101, 1421, 3672, 1136, 1106, 1329, 1126, 3651, 12356, 197, 1295, 1160, 1209, 4900, 1128, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n",
            "{'input_ids': [101, 1184, 1110, 23609, 26405, 7926, 1891, 112, 188, 5095, 16408, 6234, 136, 197, 181, 170, 102], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}\n"
          ]
        }
      ],
      "source": [
        "# Sanity check\n",
        "for i in range(5):\n",
        "    print(tokenized_datasets[\"train\"][i])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a2Ox51FG_xe3"
      },
      "source": [
        "### Part 2\n",
        "\n",
        "Create a decoder model and train it using the joke-dataset. Try out different hyperparameters for the model, like the number of decoder blocks, hidden dimension, sequence length etc. Choose the model which gives you the best results. Log the results for the different hyperparameters combinations used. Submit the logs in a pdf. (18 marks)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 94,
      "metadata": {
        "id": "YfXQozOGWQH0"
      },
      "outputs": [],
      "source": [
        "class CausalSelfAttention(nn.Module):\n",
        "  def __init__(self, d_k,d_model, n_heads, max_len):\n",
        "    super().__init__()\n",
        "\n",
        "    # Assume d_v =d_k\n",
        "    self.d_k=d_k\n",
        "    self.n_heads = n_heads\n",
        "\n",
        "    self.key = nn. Linear (d_model, d_k*n_heads)\n",
        "    self.query = nn. Linear (d_model, d_k*n_heads)\n",
        "    self.value = nn. Linear (d_model, d_k*n_heads)\n",
        "\n",
        "    # final linear layer\n",
        "    self.fc=nn.Linear (d_k*n_heads, d_model)\n",
        "\n",
        "    # casual mask\n",
        "    # make it so that diagonal is 0\n",
        "    # this way we don't have to shift the inputs to make targets\n",
        "\n",
        "    cm=torch.tril(torch.ones(int(max_len),int(max_len)))\n",
        "    self.register_buffer(\"causal_mask\",  cm.view(1,1, int(max_len), int(max_len)))\n",
        "\n",
        "\n",
        "  def forward(self, q, k, v, pad_mask=None):\n",
        "    q=self.query(q) # N x T x (hd_k)\n",
        "    k=self.key(k) # N x T x (hd_k)\n",
        "    v=self.value(v) # N x T x (hd_k)\n",
        "\n",
        "    N = q.shape[0]\n",
        "    T = q.shape[1]\n",
        "\n",
        "\n",
        "    # change the shape to:\n",
        "    # (N, T, h, d_k) --> N, h, T, d_k)\n",
        "    # in order for matrix multiply to work properly\n",
        "    q=q.view (N, T, self.n_heads, self.d_k).transpose(1,2)\n",
        "    k=k.view (N, T, self.n_heads, self.d_k).transpose(1,2)\n",
        "    v=v.view (N, T, self.n_heads, self.d_k).transpose(1,2)\n",
        "\n",
        "    # Copute attention weights\n",
        "    # (N, h, T, d_k)   x  (N, h, d_k, T )  --> (N, h, T, T)\n",
        "    attn_scores = q@k.transpose(-2,-1)/math.sqrt(self.d_k) # Scaled dot product;  @ --> torch.matmul\n",
        "\n",
        "    if pad_mask is not None:\n",
        "      attn_scores = attn_scores.masked_fill(pad_mask[:,None,None,:] == 0, float('-inf'))\n",
        "\n",
        "    attn_scores = attn_scores.masked_fill(self.causal_mask[:, :, :T, :T] == 0, float('-inf'))\n",
        "\n",
        "    attn_weights = F.softmax(attn_scores, dim =-1)\n",
        "\n",
        "    # Compute attention-weighted values\n",
        "    # (N, h, T, T) x (N, h, T, d_k) --> (N, h, T, d_k)\n",
        "    A = attn_weights @ v\n",
        "\n",
        "    # reshape it back before final linear layer\n",
        "    A = A.transpose(1,2)  # (N, T, h, d_k)\n",
        "    A = A.contiguous(). view(N, T, self.d_k*self.n_heads) # (N, T, h*d_k)\n",
        "\n",
        "    # projection\n",
        "    return self.fc(A)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 95,
      "metadata": {
        "id": "LlO7Qbi1XrnX"
      },
      "outputs": [],
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "  def __init__(self, d_k, d_model, n_heads, max_len, dropout_prob = 0.1):\n",
        "    super().__init__()\n",
        "\n",
        "    self.ln1 = nn.LayerNorm(d_model)\n",
        "    self.ln2 = nn.LayerNorm(d_model)\n",
        "    self.mha = CausalSelfAttention(d_k, d_model, n_heads, max_len)\n",
        "\n",
        "    self.ann = nn.Sequential(\n",
        "        nn.Linear(d_model, d_model*4),\n",
        "        nn.GELU(),\n",
        "        nn.Linear(d_model*4, d_model),\n",
        "        nn.Dropout(dropout_prob),\n",
        "    )\n",
        "\n",
        "    self.dropout = nn.Dropout(p=dropout_prob)\n",
        "\n",
        "  def forward (self, x, pad_mask=None):\n",
        "    x = self.ln1(x + self.mha(x, x, x, pad_mask))\n",
        "    x = self.ln2(x + self.ann(x))\n",
        "    x = self.dropout(x)\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 96,
      "metadata": {
        "id": "mYXlM8xxX4aO"
      },
      "outputs": [],
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "  def __init__(self, d_model, max_len =2048, dropout_porb=0.1):\n",
        "    super().__init__()\n",
        "    self.dropout = nn.Dropout(p = dropout_porb)\n",
        "\n",
        "    position = torch.arange(max_len).unsqueeze(1)\n",
        "    i = torch.arange(0, d_model//2)\n",
        "    pe = torch.zeros(1, max_len, d_model)\n",
        "    pe[0, :, 0::2] = torch.sin(position / (10000)**(2*i/d_model))\n",
        "    pe[0, :, 1::2] = torch.cos(position / (10000)**(2*i/d_model))\n",
        "    self.register_buffer('pe', pe)    # If you have parameters in your model which should be saved and restored in the state_dict\n",
        "                                      # but not trained by the optimizer, then you should register them as buffers.\n",
        "                                      # Buffers won’t be returned in model.parameters(), so that the optimizer won’t have a change to update them.\n",
        "\n",
        "  def forward(self, x):\n",
        "    # x.shape : N x T x D\n",
        "    x = x + self.pe[:, :x.size(1), :]\n",
        "    return self.dropout(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 97,
      "metadata": {
        "id": "GYIOdGMQYoBg"
      },
      "outputs": [],
      "source": [
        "class Decoder(nn.Module):\n",
        "  def __init__(self, vocab_size, max_len, d_k, d_model, n_heads, n_layers, dropout_prob):\n",
        "    super().__init__()\n",
        "\n",
        "    self.embedding=nn.Embedding(vocab_size, d_model)\n",
        "\n",
        "    self.pos_encoding=PositionalEncoding(d_model, max_len, dropout_prob)\n",
        "\n",
        "    transformer_blocks=[TransformerBlock(d_k, d_model, n_heads, max_len, dropout_prob) for _ in range(n_layers) ]\n",
        "\n",
        "    self.transformer_blocks = nn.Sequential(*transformer_blocks)\n",
        "\n",
        "    self.ln = nn.LayerNorm(d_model)\n",
        "    self.fc = nn.Linear(d_model, vocab_size)\n",
        "\n",
        "  def forward(self, x, pad_mask=None):\n",
        "    x=self.embedding(x)\n",
        "    x=self.pos_encoding(x)\n",
        "    for block in self.transformer_blocks:\n",
        "      x = block(x, pad_mask)\n",
        "    x = self.ln(x)\n",
        "    x = self.fc(x)  # many-to-many\n",
        "    return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 98,
      "metadata": {
        "id": "lzqoaz6QL5st"
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "\n",
        "train_loader =  DataLoader(tokenized_datasets[\"train\"], shuffle = True, batch_size = 16, collate_fn = data_collator)\n",
        "valid_loader = DataLoader(tokenized_datasets[\"validation\"], batch_size = 16, collate_fn = data_collator)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 99,
      "metadata": {
        "id": "iZz-VEGUeCFu"
      },
      "outputs": [],
      "source": [
        "def train(model, criterion, optimizer, train_loader, epochs, early_stop_threshold=10**10):\n",
        "  train_losses = np.zeros(epochs)\n",
        "\n",
        "  for it in range(epochs):\n",
        "    model.train()\n",
        "    t0=datetime.now()\n",
        "    train_loss = []\n",
        "    for batch in tqdm(train_loader):\n",
        "      # move data to GPU\n",
        "      batch = {k:v.to(device) for k, v, in batch.items()}\n",
        "\n",
        "      # zero the parameter gradients\n",
        "      optimizer.zero_grad()\n",
        "\n",
        "      # shift the targets backwards\n",
        "      targets = batch['input_ids'].clone().detach()\n",
        "      targets = torch.roll(targets,shifts = -1, dims = 1) #Shifting the sequence for creating target\n",
        "      targets[:,-1] = tokenizer.pad_token_id\n",
        "\n",
        "      # forward pass\n",
        "      outputs = model(batch['input_ids'], batch['attention_mask'])\n",
        "      loss = criterion(outputs.transpose(2,1),targets)\n",
        "\n",
        "      # Backward and optimize\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "      train_loss.append(loss.item())\n",
        "\n",
        "    # get  train Loss and test loss\n",
        "    train_loss = np.mean(train_loss)\n",
        "\n",
        "    # save lossess\n",
        "    train_losses[it] = train_loss\n",
        "\n",
        "    dt = datetime.now() - t0\n",
        "    print(f'Epoch: {it+1}/{epochs}, Train Loss: {train_loss:.4f}, Duration: {dt}')\n",
        "\n",
        "    if train_loss > early_stop_threshold:\n",
        "      print(\"Too much loss due to bad hyperparameters, aborting train\")\n",
        "      break\n",
        "\n",
        "  return train_losses"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 100,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 408
        },
        "id": "mx-BZRiTL5su",
        "outputId": "65781b41-2676-4f22-8949-b3ff467462b2"
      },
      "outputs": [
        {
          "output_type": "error",
          "ename": "RuntimeError",
          "evalue": "CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-100-532f81d335a7>\u001b[0m in \u001b[0;36m<cell line: 13>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     11\u001b[0m     \u001b[0mdropout_prob\u001b[0m \u001b[0;34m=\u001b[0m\u001b[0;36m0.1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m     )\n\u001b[0;32m---> 13\u001b[0;31m \u001b[0mmodel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     14\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     15\u001b[0m \u001b[0;31m# Loss and optimizer\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mto\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1338\u001b[0m                     \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1339\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1340\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1341\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1342\u001b[0m     def register_full_backward_pre_hook(\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    898\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mrecurse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    899\u001b[0m             \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mchildren\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 900\u001b[0;31m                 \u001b[0mmodule\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_apply\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    901\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    902\u001b[0m         \u001b[0;32mdef\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtensor\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtensor_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    925\u001b[0m             \u001b[0;31m# `with torch.no_grad():`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    926\u001b[0m             \u001b[0;32mwith\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mno_grad\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 927\u001b[0;31m                 \u001b[0mparam_applied\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    928\u001b[0m             \u001b[0mp_should_use_set_data\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcompute_should_use_set_data\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mparam\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparam_applied\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    929\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36mconvert\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m   1324\u001b[0m                         \u001b[0mmemory_format\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mconvert_to_format\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1325\u001b[0m                     )\n\u001b[0;32m-> 1326\u001b[0;31m                 return t.to(\n\u001b[0m\u001b[1;32m   1327\u001b[0m                     \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1328\u001b[0m                     \u001b[0mdtype\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_floating_point\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mt\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_complex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;31mRuntimeError\u001b[0m: CUDA error: device-side assert triggered\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\n"
          ]
        }
      ],
      "source": [
        "# Dummy Test for sanity check\n",
        "\n",
        "# Instantiate\n",
        "model = Decoder(\n",
        "    vocab_size = tokenizer.vocab_size,\n",
        "    max_len=tokenizer.model_max_length,\n",
        "    d_k=16,\n",
        "    d_model = 64,\n",
        "    n_heads = 4,\n",
        "    n_layers =2,\n",
        "    dropout_prob =0.1\n",
        "    )\n",
        "model.to(device)\n",
        "\n",
        "# Loss and optimizer\n",
        "criterion = nn.CrossEntropyLoss(ignore_index = tokenizer.pad_token_id)\n",
        "optimizer = torch.optim.Adam(model.parameters())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pJNNenMDe1wO"
      },
      "outputs": [],
      "source": [
        "# train the model for one epoch\n",
        "train_losses = train(model, criterion, optimizer, train_loader, epochs = 1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5uhXgzjFL5su"
      },
      "outputs": [],
      "source": [
        "def model_completion_test(joke, model):\n",
        "    # Split into premise and punchline\n",
        "    premise = joke[:joke.index(\" | \") + 4]\n",
        "\n",
        "    # Encode the premise\n",
        "    tokenized_premise = tokenizer(premise, return_tensors='pt')\n",
        "    input_ids = tokenized_premise[\"input_ids\"][:, :-1].to(device)\n",
        "    mask = tokenized_premise[\"attention_mask\"][:, :-1].to(device)\n",
        "\n",
        "    # Complete with upto 50 words\n",
        "    for _ in range(50):\n",
        "        try:\n",
        "            outputs = model(input_ids, mask)\n",
        "            prediction_id = torch.argmax(outputs[:, -1, :], axis=-1)    # Generated word\n",
        "            input_ids = torch.hstack((input_ids, prediction_id.view(1,1)))  # add this generated word word to the input\n",
        "            mask = torch.ones_like(input_ids)\n",
        "        except:\n",
        "            break\n",
        "        if prediction_id == tokenizer.sep_token_id:\n",
        "            break\n",
        "\n",
        "    print('Joke:', joke)\n",
        "    print('premise:', premise)\n",
        "    print('model completion:', tokenizer.decode(input_ids[0]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ofs4WpN9L5su"
      },
      "outputs": [],
      "source": [
        "# Testing the completion\n",
        "model.eval()\n",
        "\n",
        "joke = dataset[\"validation\"][120][\"body\"]\n",
        "\n",
        "try:\n",
        "    model_completion_test(joke, model)\n",
        "except:\n",
        "    pass"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHVdKy_H8lbz"
      },
      "outputs": [],
      "source": [
        "# Hyperparameter Tuning\n",
        "\n",
        "# Define the hyperparameter combinations\n",
        "hyperparams = [\n",
        "    (16, 64, 4, 2, 0.1),\n",
        "    (32, 256, 8, 3, 0.1),\n",
        "    (32, 512, 8, 2, 0.1),\n",
        "    (64, 768, 10, 5, 0.15),\n",
        "]\n",
        "\n",
        "best_valid_loss = 1e12  # Initialize with a very large value\n",
        "best_hyperparams = None\n",
        "\n",
        "# Loop through the hyperparameter combinations\n",
        "for d_k, d_model, n_heads, n_layers, dropout_prob in hyperparams:\n",
        "    print(f\"Training model with hyperparameters: d_k={d_k}, d_model={d_model}, n_heads={n_heads}, n_layers={n_layers}, dropout_prob={dropout_prob}\")\n",
        "\n",
        "    # Instantiate the model with the current hyperparameters\n",
        "    model = Decoder(\n",
        "        vocab_size=tokenizer.vocab_size,\n",
        "        max_len=tokenizer.model_max_length,\n",
        "        d_k=d_k,\n",
        "        d_model=d_model,\n",
        "        n_heads=n_heads,\n",
        "        n_layers=n_layers,\n",
        "        dropout_prob=dropout_prob,\n",
        "    )\n",
        "    model.to(device)\n",
        "\n",
        "    # Loss and optimizer\n",
        "    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)\n",
        "    optimizer = torch.optim.Adam(model.parameters(), 0.001)\n",
        "\n",
        "    # Train the model for 5 epochs\n",
        "    train_losses = train(model, criterion, optimizer, train_loader, epochs=5, early_stop_threshold=6)   #Putting stop threshold at 6 as the basic model started of with loss roughly 4\n",
        "\n",
        "    # Validation\n",
        "    model.eval()\n",
        "    valid_losses = []\n",
        "    with torch.no_grad():\n",
        "        for batch in valid_loader:\n",
        "            batch = {k: v.to(device) for k, v in batch.items()}\n",
        "            targets = batch['input_ids'].clone().detach()\n",
        "            targets = torch.roll(targets, shifts=-1, dims=1)\n",
        "            targets[:, -1] = tokenizer.pad_token_id\n",
        "            outputs = model(batch['input_ids'], batch['attention_mask'])\n",
        "            loss = criterion(outputs.transpose(2, 1), targets)\n",
        "            valid_losses.append(loss.item())\n",
        "\n",
        "    valid_loss = np.mean(valid_losses)\n",
        "\n",
        "    print(f\"Validation Loss: {valid_loss:.4f}\")\n",
        "\n",
        "    # Testing on random stuff\n",
        "    print('\\n Testing this model on 4 random jokes')\n",
        "    for _ in range(4):\n",
        "        joke = dataset[\"validation\"][random.randint(0, 1000)][\"body\"]\n",
        "        try:\n",
        "            model_completion_test(joke, model)\n",
        "        except:\n",
        "            pass\n",
        "\n",
        "    # Check if this model is the best so far\n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        best_hyperparams = (d_k, d_model, n_heads, n_layers, dropout_prob)\n",
        "\n",
        "# Print the best hyperparameters and validation loss\n",
        "print(\"\\nBest Hyperparameters:\")\n",
        "print(f\"d_k: {best_hyperparams[0]}, d_model: {best_hyperparams[1]}, n_heads: {best_hyperparams[2]}, n_layers: {best_hyperparams[3]}, dropout_prob: {best_hyperparams[4]}\")\n",
        "print(f\"Best Validation Loss: {best_valid_loss:.4f}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ek6hWr97L5su"
      },
      "outputs": [],
      "source": [
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rzdU4i9b8lb0"
      },
      "outputs": [],
      "source": [
        "# Get the best model from the best hyperparams\n",
        "model = Decoder(\n",
        "    vocab_size=tokenizer.vocab_size,\n",
        "    max_len=tokenizer.model_max_length,\n",
        "    d_k=best_hyperparams[0],\n",
        "    d_model=best_hyperparams[1],\n",
        "    n_heads=best_hyperparams[2],\n",
        "    n_layers=best_hyperparams[3],\n",
        "    dropout_prob=best_hyperparams[4],\n",
        ")\n",
        "model.to(device)\n",
        "\n",
        "# Loss and optimizer\n",
        "criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
        "\n",
        "# Train the model for 12 epochs\n",
        "train_losses = train(model, criterion, optimizer, train_loader, epochs=10)"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Save the model\n",
        "torch.save(model.state_dict(), \"model.pth\")\n",
        "\n",
        "# Save the tokenizer\n",
        "tokenizer.save_pretrained(\"tokenizer\")"
      ],
      "metadata": {
        "id": "JUowLxvsquh-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WLFN_qrohtVo"
      },
      "source": [
        "### Part 3\n",
        "\n",
        "Choose any 5 samples from the test dataset and generate the model outputs. (2 marks)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R1DePdGL8lb0"
      },
      "outputs": [],
      "source": [
        "# Choose 5 samples from the test dataset\n",
        "test_samples = dataset[\"test\"][\"body\"][:5]\n",
        "\n",
        "# Generate the model outputs\n",
        "model.eval()\n",
        "\n",
        "for joke in test_samples:\n",
        "    model_completion_test(joke, model)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nqHtUP42Mlma"
      },
      "source": [
        "## Question 2\n",
        "\n",
        "Finetune a GPT-2 model using the dataset and compare the responses of your model and the finetuned GPT-2 model. (18 marks)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EU7wMkmVL5su"
      },
      "outputs": [],
      "source": [
        "from transformers import GPT2Tokenizer, GPT2LMHeadModel, DataCollatorForLanguageModeling, Trainer, TrainingArguments\n",
        "from torch.utils.data import Dataset\n",
        "import torch\n",
        "\n",
        "\n",
        "# Custom Dataset for jokes\n",
        "class JokesDataset(Dataset):\n",
        "    def __init__(self, dataset, tokenizer, block_size=128):\n",
        "        self.instruction_set = []\n",
        "        for joke in dataset[\"body\"]:\n",
        "            if len(joke) < 65:  # 64, more than the previous cutoff of 48\n",
        "                tokenized_joke = tokenizer(\n",
        "                    joke,\n",
        "                    truncation=True,\n",
        "                    max_length=block_size,\n",
        "                    padding=\"max_length\",\n",
        "                    return_tensors=\"pt\"\n",
        "                )\n",
        "\n",
        "                self.instruction_set.append({\n",
        "                    \"input_ids\": tokenized_joke[\"input_ids\"].squeeze(0),\n",
        "                    \"attention_mask\": tokenized_joke[\"attention_mask\"].squeeze(0)\n",
        "                })\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.instruction_set)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return self.instruction_set[idx]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kksjtBY_L5sv"
      },
      "outputs": [],
      "source": [
        "model_name = \"gpt2\"\n",
        "tokenizer = GPT2Tokenizer.from_pretrained(model_name)\n",
        "tokenizer.pad_token = tokenizer.eos_token  # Handle padding for GPT-2\n",
        "model = GPT2LMHeadModel.from_pretrained(model_name)\n",
        "\n",
        "\n",
        "block_size = 128\n",
        "dataset = JokesDataset(dataset[\"train\"], tokenizer, block_size)\n",
        "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HRSE3IvdL5sv"
      },
      "outputs": [],
      "source": [
        "training_args = TrainingArguments(\n",
        "    output_dir=\"./gpt2-jokes-finetuned\",\n",
        "    overwrite_output_dir=True,\n",
        "    num_train_epochs=2,\n",
        "    per_device_train_batch_size=16,     # Smaller batch size as larger model\n",
        "    save_steps=10000,\n",
        "    save_total_limit=2,\n",
        "    logging_dir=\"./logs\",\n",
        "    logging_steps=10000,\n",
        "    learning_rate=5e-5,\n",
        "    weight_decay=0.01,\n",
        "    warmup_steps=100,\n",
        "    prediction_loss_only=True,\n",
        ")\n",
        "\n",
        "# Initialize Trainer\n",
        "trainer = Trainer(\n",
        "    model=model,\n",
        "    args=training_args,\n",
        "    data_collator=data_collator,\n",
        "    train_dataset=dataset,\n",
        ")\n",
        "\n",
        "# Fine-tune the model\n",
        "trainer.train()\n",
        "\n",
        "# Save the fine-tuned model and tokenizer\n",
        "trainer.save_model(\"./gpt2-jokes-finetuned\")\n",
        "tokenizer.save_pretrained(\"./gpt2-jokes-finetuned\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jhiC-VazL5sv"
      },
      "outputs": [],
      "source": [
        "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
        "\n",
        "def gpt_completion(prompt, max_length=100, temperature=1.0, top_k=50, top_p=0.95):\n",
        "    input_ids = tokenizer.encode(prompt, return_tensors=\"pt\").to(device)\n",
        "\n",
        "    output_ids = model.generate(\n",
        "        input_ids=input_ids,\n",
        "        max_length=max_length,\n",
        "        temperature=temperature,\n",
        "        top_k=top_k,\n",
        "        top_p=top_p,\n",
        "        do_sample=True,\n",
        "        pad_token_id=tokenizer.eos_token_id\n",
        "    )\n",
        "\n",
        "    return tokenizer.decode(output_ids[0], skip_special_tokens=True)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MpNvPsu-L5sv"
      },
      "outputs": [],
      "source": [
        "# Comparson with decoder made from scratch\n",
        "\n",
        "for joke in test_samples:\n",
        "    print(\"joke:\", joke)\n",
        "    premise = premise = joke[:joke.index(\" | \") + 4]\n",
        "    print(\"premise:\", premise)\n",
        "    print(\"GPT2 Completion:\", gpt_completion(premise))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "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"
    },
    "orig_nbformat": 4,
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "f6fa8ace0ccf495aacfac69c3e2f386c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_44608c4bb86747d9af10af9054b08864",
              "IPY_MODEL_d491f9941bb44a6d93e5c1ac585f4e79",
              "IPY_MODEL_5da80ad644884e39a2a99dde12038129"
            ],
            "layout": "IPY_MODEL_1213ae2aca2942b3aba61ee546c8cf2d"
          }
        },
        "44608c4bb86747d9af10af9054b08864": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3f48e3aedf7a49c29d6f58531df2e93c",
            "placeholder": "​",
            "style": "IPY_MODEL_40845da59b0643dda0c84c3c15e985b9",
            "value": "Map: 100%"
          }
        },
        "d491f9941bb44a6d93e5c1ac585f4e79": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_6c990ad034a34506a7f360da64e58477",
            "max": 15410,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_ace3cba5ecf144909ba947315a521dd7",
            "value": 15410
          }
        },
        "5da80ad644884e39a2a99dde12038129": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e8b6bf87e560457abc822e0ec4bf63ad",
            "placeholder": "​",
            "style": "IPY_MODEL_497bf3af26c94ab5bb5204e6f87904a0",
            "value": " 15410/15410 [00:06&lt;00:00, 3074.06 examples/s]"
          }
        },
        "1213ae2aca2942b3aba61ee546c8cf2d": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "3f48e3aedf7a49c29d6f58531df2e93c": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "40845da59b0643dda0c84c3c15e985b9": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "6c990ad034a34506a7f360da64e58477": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "ace3cba5ecf144909ba947315a521dd7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "e8b6bf87e560457abc822e0ec4bf63ad": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "497bf3af26c94ab5bb5204e6f87904a0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}