assignments/hw4/hw4.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ISL7lGbExVbF"
      },
      "source": [
        "# HW4: Large Language Models\n",
        "\n",
        "The first half of this assignment (Parts 0 and 1) will review some key ingredients of sequence modeling. In the process, we will build a baseline transformer model for next token prediction in code.\n",
        "**The deliverable will be completing the questions posed in part 0 and part 1.**\n",
        "\n",
        "The second half of the assignment (Part 2) will be an open-ended mini-project where you have the freedom to delve more deeply into language modelling (where the language in question is python code). Further instructions are in _Part 2: Mini-project_. But, in general, you should feel free to try other architectures (HMMs, RNNs, transformers, state space layers, diffusion models etc.) or to invent new architectures. The goal will be to find some area of possible improvement (we interpret \"improvement\" quite loosely, but it is up to you to state precisely in what sense your proposed innovation might constitute an improvement and to show convincing evidence that your innovation does or does not constitue an improvement according to your definition); to formulate and state a precise hypothesis; and to falsify or support the hypothesis with rigorous empirical analyses.\n",
        "**The deliverable will be a report of no more than 4 pages (references not included in the page limit).**\n",
        "\n",
        "**For this final assignment you have the option to work in pairs.**\n",
        "\n",
        "\n",
        "> **<u>This Assignment</u>**\n",
        ">\n",
        ">**Model:** You will begin by implementing a baseline using attention and transformers. But, for the mini-project, you will be free to use any model (HMMs, RNNs, transformers, state space layers, diffusion models etc.) that you would like!\n",
        ">\n",
        ">**Algorithm:** mini-batched stochatic gradient descent / whatever you'd like! We will be using deep learning models in at least the first half of the assignment. When you actually go into \"production\" you should be sure to **use the GPU on colab** (make sure you switch to GPU in the \"Runtime\" tab above). However, you have limited colab gpu units for free, so when you are \"developing,\" you may wish to run smaller models with less data for fewer iterations just to get the bugs out / make sure you pipeline is working (see the \"prototyping\" exercise in HW 0 for more discussion).\n",
        ">\n",
        ">**Data**: A large corpus of python code from [the Stack](https://huggingface.co/datasets/bigcode/the-stack-dedup). We have taken a dataset of around 4 million tokens from the stack and stored in a csv file for you for easy access."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qZ-AwbZ4ySCJ"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xd4fRbWr15Yp"
      },
      "outputs": [],
      "source": [
        "# can take around 30s\n",
        "%%capture\n",
        "! pip install datasets #huggingface datasets library\n",
        "! pip install --upgrade pyarrow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_5IXGh6OOBZV"
      },
      "outputs": [],
      "source": [
        "# torch imports\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "\n",
        "# hugging face imports\n",
        "from datasets import load_dataset\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm\n",
        "\n",
        "import pandas as pd\n",
        "import sys\n",
        "import warnings\n",
        "\n",
        "torch.manual_seed(305)\n",
        "\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "assert device=='cuda', \"you need to change runtime type to GPU\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A_Z5Jh74DH_E"
      },
      "outputs": [],
      "source": [
        "# hyperparams and helper functions\n",
        "SMALL_ITERS = 1000\n",
        "LARGE_ITERS = 2000\n",
        "context_window_size = 256\n",
        "chunk_size = 512 # BERT can only take max input size 512 characters\n",
        "\n",
        "def chunk_string(string, size):\n",
        "    \"\"\"\n",
        "    Splits a string into chunks of a specified size.\n",
        "\n",
        "    :param string: The string to be chunked.\n",
        "    :param size: The desired chunk size.\n",
        "    :return: A list of string chunks.\n",
        "    \"\"\"\n",
        "    return [string[i:i+size] for i in range(0, len(string), size)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lF6dgHnhOprg"
      },
      "source": [
        "## Part 0: Preprocessing\n",
        "\n",
        "As in the previous problem sets, a certain amount of preprocessing for textual data is required."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WF0_bhXNOxeS"
      },
      "source": [
        "### 0.1: Loading the dataset\n",
        "\n",
        "The first step is to actually download the dataset. We will be using a dataset on [huggingface](https://huggingface.co/). You can think of hugging face as the sklearn of deep learning.\n",
        "\n",
        "The dominant mode for preprocessing textual data is to tokenize it, that is, to split the dataset into a finite vocabulary of tokens. Then, we can set up a dictionary where counting numbers map to tokens. Tokens can be characters, or words, or subwords; in fact, the \"best\" way to tokenize text is an active area of research. For our baseline, we will use a tokenizer that microsoft created for code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "53dGz7ExDkUv"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "tokenizer = AutoTokenizer.from_pretrained('microsoft/CodeBERT-base')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BAdrvcIbORig"
      },
      "outputs": [],
      "source": [
        "# Load the concatenated data\n",
        "raw_data = pd.read_csv(\"https://raw.githubusercontent.com/slinderman/stats305b/winter2024/assignments/hw4/python_corpus_4M.csv\", header=None)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note: if you run the below cell, you might get something like the following warning\n",
        "\n",
        "> Token indices sequence length is longer than the specified maximum sequence length for this model. Running this sequence through the model will result in indexing errors\n",
        "\n",
        "You can safely ignore this error for the purpose of the assignment (we did in our solutions).\n",
        "\n",
        "But, if  you can figure out this error, or find a fix, let the teaching staff know, and we will give you extra credit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LFZ2or4BO2w3"
      },
      "outputs": [],
      "source": [
        "# should take around 3 min to load in around 4M tokens\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "tokens = torch.tensor([], dtype=torch.long)\n",
        "for index, row in raw_data.iterrows():\n",
        "    text = row[0]\n",
        "    chunks = chunk_string(text, chunk_size)\n",
        "    n = len(chunks)\n",
        "    for idx, chunk in enumerate(chunks):\n",
        "        new_tokens = torch.tensor(tokenizer.encode(chunk, add_special_tokens=True))\n",
        "\n",
        "        # logic to avoid incorrectly adding in start and end sequence tokens as an artifact of chunking\n",
        "        if idx == 0:\n",
        "            tokens = torch.cat((tokens, new_tokens[:-1]), dim=0)\n",
        "        elif idx == n-1:\n",
        "            tokens = torch.cat((tokens, new_tokens[1:]), dim=0)\n",
        "        else:\n",
        "            tokens = torch.cat((tokens, new_tokens[1:-1]), dim=0)\n",
        "\n",
        "print(f\"{len(tokens)} tokens have been loaded in\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUnScCoUCRI0"
      },
      "source": [
        "### Question 0.2: Examining the tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3rz0XmuXCdGJ"
      },
      "source": [
        "Let's see what the tokens look like! We will use these two prompts during the assignment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KPbmLgmmChVg"
      },
      "outputs": [],
      "source": [
        "prompt_1_text = \\\n",
        "\"\"\"def newton(eta, N, X, y, gamma, beta=None):\n",
        "  \\\"\"\"\n",
        "  Performs Newton's method on the negative average log likelihood with an\n",
        "  l2 regularization term\n",
        "\n",
        "  beta: torch.Tensor, of shape (teams)\n",
        "  X: torch.Tensor, the covariate matrix, of shape (-1, teams)\n",
        "  y: torch.Tensor, the response vector, of shape (teams)\n",
        "  gamma: float, the scale parameter for the regularization\n",
        "  beta: torch.Tensor, the starting point for gradient descent, if specified\n",
        "  \\\"\"\"\n",
        "\n",
        "  if beta is None:\n",
        "    # Instantiate the beta vector at a random point\n",
        "    beta = torch.randn(X.shape[1])\n",
        "  else:\n",
        "    beta = torch.clone(beta)\n",
        "\n",
        "  loss = []\n",
        "\n",
        "  # Instantiate a list to store the loss throughout the gradient descent\n",
        "  # path\n",
        "  for i in tqdm(range(N)):\"\"\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gZpeZjvqCnmH"
      },
      "outputs": [],
      "source": [
        "prompt_2_text = \\\n",
        "\"\"\"import torch\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "def normalize(x, axis=-1):\n",
        "    \\\"\"\"Performs L2-Norm.\\\"\"\"\n",
        "    num = x\n",
        "    denom = torch.norm(x, 2, axis, keepdim=True).expand_as(x) + 1e-12\n",
        "    return num / denom\n",
        "\n",
        "def euclidean_dist(x, y):\n",
        "    \\\"\"\"Computes Euclidean distance.\\\"\"\"\n",
        "    m, n = x.size(0), y.size(0)\n",
        "    xx = torch.pow(x, 2).sum(1, keepdim=True).expand(m, n)\n",
        "    yy = torch.pow(x, 2).sum(1, keepdim=True).expand(m, m).t()\n",
        "    dist = xx + yy - 2 * torch.matmul(x, y.t())\n",
        "\n",
        "    dist = dist.clamp(min=1e-12).sqrt()\n",
        "\n",
        "    return dist\n",
        "\n",
        "\n",
        "def cosine_dist(x, y):\"\"\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jFZUT77rCp3q"
      },
      "source": [
        "Here is what the tokenized output for the prompts looks like"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-8rJjbTKCtRa"
      },
      "outputs": [],
      "source": [
        "tokenizer.decode(tokenizer.encode(prompt_1_text))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a4zd13L0CvQ9"
      },
      "outputs": [],
      "source": [
        "tokenizer.decode(tokenizer.encode(prompt_2_text))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bbOwt9U_CyNi"
      },
      "source": [
        "And here are what the first and last 10 tokens for prompt 1 look like:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nqBUhCiUC6Tk"
      },
      "outputs": [],
      "source": [
        "for tok in tokenizer.encode(prompt_1_text, add_special_tokens=True)[:10]:\n",
        "    print(f\"{tok} : {tokenizer.decode([tok])}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UHpBvaxzC7hd"
      },
      "outputs": [],
      "source": [
        "for tok in tokenizer.encode(prompt_1_text)[-10:]:\n",
        "    print(f\"{tok} : {tokenizer.decode([tok])}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MrBJPZYYC9za"
      },
      "source": [
        "**Question 0.2**: What is the meanining of the `<s>` and the `<\\s>` tokens? Why is it useful to have them?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0hwg1_7kDxwY"
      },
      "source": [
        "### 0.3: Building our dataloader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uk4UJm5SEMlB"
      },
      "source": [
        "There are around 50,000 tokens in the codebert vocab, but we only use around 20,000 of them. To make our lives easier, we just reindex the token indices to go from 1 to around 20,000."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RZC0TYeaD6yj"
      },
      "outputs": [],
      "source": [
        "# Get unique elements\n",
        "extra_tokens = torch.cat((torch.tensor(tokenizer.encode(prompt_1_text, add_special_tokens=True)),\n",
        "                          torch.tensor(tokenizer.encode(prompt_2_text, add_special_tokens=True))),\n",
        "                         dim=0)\n",
        "\n",
        "unique_tokens = torch.unique(torch.cat((tokens, extra_tokens), dim=0))\n",
        "\n",
        "# Create a mapping from code bert to ids that increment by one\n",
        "from_code_bert_dict = {element.item(): id for id, element in enumerate(unique_tokens)}\n",
        "\n",
        "# Create a reverse mapping from ids to code bert token ids\n",
        "to_code_bert_dict = {id: element for element, id in from_code_bert_dict.items()}\n",
        "\n",
        "vocab_size = len(unique_tokens)\n",
        "print(f\"there are {vocab_size} distinct tokens in the vocabulary\")\n",
        "\n",
        "# helper functions to move between code bert and simple ids\n",
        "def from_code_bert(tkn_lst):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "    tkn_lst: a list of code bert tokens\n",
        "    Returns:\n",
        "    a list of simple ids\n",
        "    \"\"\"\n",
        "    tkns = [int(from_code_bert_dict[token]) for token in tkn_lst]\n",
        "    return tkns\n",
        "\n",
        "\n",
        "def to_code_bert(tkn_lst):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "    tkn_lst: a list of simple ids\n",
        "    Returns:\n",
        "    a list of code bert tokens\n",
        "    \"\"\"\n",
        "    tkns = [int(to_code_bert_dict[token]) for token in tkn_lst]\n",
        "    return tkns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pNOmdccJEXWk"
      },
      "outputs": [],
      "source": [
        "# let's translate our dataset into our ids\n",
        "tokens_simple_id = torch.tensor([from_code_bert_dict[token.item()] for token in tokens])\n",
        "\n",
        "# split up the data into train and validation sets\n",
        "n = int(0.9 * len(tokens_simple_id)) # first 90% will be train, rest val\n",
        "train_data = tokens_simple_id.clone()[:n]\n",
        "val_data = tokens_simple_id.clone()[n:]\n",
        "\n",
        "print(f\"there are {len(train_data)} tokens in the training set\")\n",
        "print(f\"there are {len(val_data)} tokens in the validation set\")\n",
        "print(f\"there are {vocab_size} distinct tokens in the vocabulary\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nsOs0_diEcjY"
      },
      "source": [
        "We also write helper functions to get batches of data and to evaluate the loss of various models on them."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AkAD0PfiEfjG"
      },
      "outputs": [],
      "source": [
        "# function for getting batches of data\n",
        "def get_batch(split, context_window_size, device, batch_size=32):\n",
        "    \"\"\"\n",
        "    generate a small batch of data of inputs x and targets y\n",
        "\n",
        "    Args:\n",
        "        split: 'train' or 'val'\n",
        "        device: 'cpu' or 'cuda' (should be 'cuda' if available)\n",
        "    \"\"\"\n",
        "    data = train_data if split == 'train' else val_data\n",
        "    ix = torch.randint(len(data) - context_window_size, (batch_size,))\n",
        "    x = torch.stack([data[i:i+context_window_size] for i in ix])\n",
        "    y = torch.stack([data[i+1:i+context_window_size+1] for i in ix])\n",
        "    x = x.to(device)\n",
        "    y = y.to(device)\n",
        "    return x, y\n",
        "\n",
        "# helper function for tracking loss during training\n",
        "# given to you\n",
        "@torch.no_grad()\n",
        "def estimate_loss(model, eval_iters, context_window_size, device):\n",
        "    \"\"\"\n",
        "    Args:\n",
        "      model: model being evaluated\n",
        "      eval_iters: number of batches to average over\n",
        "      context_window_size: size of the context window\n",
        "      device: 'cpu' or 'cuda' (should be 'cuda' if available)\n",
        "    \"\"\"\n",
        "    out = {}\n",
        "    for split in ['train', 'val']:\n",
        "        losses = torch.zeros(eval_iters)\n",
        "        for k in range(eval_iters):\n",
        "            X, Y = get_batch(split, context_window_size, device)\n",
        "            logits, loss = model(X, Y)\n",
        "            losses[k] = loss.item()\n",
        "        out[split] = losses.mean()\n",
        "    return out"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ercCUt_d07FX"
      },
      "source": [
        "## Part 1: Language Modeling\n",
        "\n",
        "In this first part of the assignment, we will implement a baseline for code modeling.\n",
        "\n",
        "In the process of building this baseline, we will review 4 key ideas of sequence modeling that have become the backbone of modern language modeling such as ChatGPT:\n",
        "\n",
        "1. Framing language modeling as next token prediction, and next token prediction as multiclass logistic regression\n",
        "2. Embedding discrete tokens in continuous latent spaces (word embeddings)\n",
        "3. Use the attention mechanism to move beyond Markovian models for sequences (we of course pay for this greater expressivity with increased compute, which is made possible in part by using matrix multiplications on acccelerated hardware like GPUs. Reducing the compute burden while maintaining the expressivity needed for good sequence modeling is an active area of research).\n",
        "4. Combining attention with deep learning in the Transformer architecture.\n",
        "\n",
        "For various architectures that you have to train, we provide approximate training times and training loss scores for your reference. They are really just for reference so don't read too much into them, but we are providing them as a warning mechanism in case something is going seriously wrong."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kh1kX0pw1bR2"
      },
      "source": [
        "### 1.1: Next token prediction as multiclass logistic regression\n",
        "\n",
        "Our first language model will simply be a lookup table. That is, given that we have token with value $v$, we will simply \"look up\" the logits that correspond to our prediction for the next token. This model is often known as a \"bigram model\" because it can be derived from the relative proportions of different bigrams (ordered pairs of tokens) occurring in a large text corpus.\n",
        "\n",
        "Let us be a bit more precise in our definition of the bigram model. Let's say that the total size of our vocabulary (the number of tokens we are using) is $V$. Let $A$ be a matrix in $\\mathbb{R}^{V \\times V}$, where each row $A_v$ corresponds to the logits for the prediction of which token would follow a token that has value $v$.\n",
        "Thus, we are modeling the distribution of the token following a token that has value $v$ as\n",
        "\\begin{align*}\n",
        "y_{t+1} \\mid y_t &= v \\sim \\mathrm{Cat}(\\mathbf{\\pi}) \\\\\n",
        "\\pi &=\\mathrm{softmax}(A_v)\n",
        "\\end{align*}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5kNWISNQ9bbH"
      },
      "source": [
        "#### Question 1.1.1\n",
        "\n",
        "$\\mathbf{\\pi} \\in \\Delta_{V-1}$ is the vector of probabilities used to parameterize the categorical distribution for the next token prediction. Explain why we parameterize\n",
        "\\begin{equation*}\n",
        "  \\mathbf{\\pi} = \\mathrm{softmax}(A_v),\n",
        "\\end{equation*}\n",
        "and could not just use\n",
        "\\begin{equation*}\n",
        "  \\mathbf{\\pi} = A_v.\n",
        "\\end{equation*}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nkHP8-T3zJ3D"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DtAjylIF-SwQ"
      },
      "source": [
        "#### Question 1.1.2\n",
        "\n",
        "Discuss the relationship between the bigram model and contigency tables (discussed in Lecture 2)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehdR52pnzZHl"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y4AzpB8M-u9o"
      },
      "source": [
        "#### Question 1.1.3\n",
        "\n",
        "Say I have a string of three tokens with ids $(7, 3, 6)$. If I use the bigram model as a generative model for language, given this information, what is distribution of the fourth token? Write your answer in terms of the matrix $A$ we defined in 1.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fs4rWYjlzc7n"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "leN9gUxoEuhg"
      },
      "source": [
        "#### Question 1.1.4\n",
        "\n",
        "Remember back in Section 0.2 \"Tokenizing the data\" when we gave you the helper function `get_batch`? Run `get_batch` and look at the inputs `x` and the targets `y`. Explain any relation between them in the context of formulating language modeling in the context of next token prediction."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rzLupp2jExDS"
      },
      "outputs": [],
      "source": [
        "xb, yb = get_batch('train', 10, device, batch_size = 1)\n",
        "print(f\"the features have token ids {xb}\")\n",
        "print('\\n')\n",
        "print(f\"the targets have token ids {yb}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UyYl8QQt_Uhb"
      },
      "source": [
        "#### Question 1.1.5\n",
        "\n",
        "Discuss the strengths and weaknesses of the bigram model as a generative model for language."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7t6Hc2dhzdzL"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mjYO1XTM_8NE"
      },
      "source": [
        "#### Question 1.1.6\n",
        "\n",
        "Say I have a string $s$ of length $T$. Derive the formula for the negative log likelihood of $s$ under the bigram model in terms of the matrix of logits $A$. What would your answer be if the matrix of logits $A$ were all zeros? What would be the value of the negative log likelihood of $s$ under a model that always perfectly predicted the next token?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QuzLfI4Pzes6"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rfzoXAATAfwn"
      },
      "source": [
        "#### Question 1.1.7: Implement the BigramLanguageModel\n",
        "\n",
        "Implement the bigram language model below.\n",
        "\n",
        "Your TODOs:\n",
        "  * if the `forward` method is provided a target, the loss should be the negative log likelihood of the target (given the context)\n",
        "  * `generate` should take in (batched) contexts and a number of new tokens to generate, and then generate text autoregressively from your model. Note that in autoregressive text generation, you iteratively append the tokens you generate to your context.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_lJ1mR4K1aj_"
      },
      "outputs": [],
      "source": [
        "class BigramLanguageModel(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          vocab_size: size of the vocabulary (the number of tokens)\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        # each token directly reads off the logits for the next token from a lookup table\n",
        "        self.logits_table = nn.Embedding(vocab_size, vocab_size)\n",
        "        self.vocab_size = vocab_size\n",
        "\n",
        "    def forward(self, token_ids, targets=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: Int(B, T), token ids that make up the context (batch has size B, each entry in the batch has length T)\n",
        "          targets: Int(B, T), token ids corresponding to the target of each context in token_ids\n",
        "\n",
        "        Returns:\n",
        "          logits: (B, T, V), logits[b,t, :] gives the length V vector of logits for the next token prediction in string b up to t tokens\n",
        "          loss: scalar, negative log likelihood of target given context\n",
        "        \"\"\"\n",
        "\n",
        "        # idx and targets are both (B,T) tensor of integers\n",
        "        logits = self.logits_table(token_ids) # (B,T,V)\n",
        "\n",
        "        if targets is None:\n",
        "            loss = None\n",
        "        else:\n",
        "            # TODO: what should the loss in this setting be?\n",
        "            loss = ...\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def generate(self, token_ids, max_new_tokens=context_window_size):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) tensor of token ids to provide as context\n",
        "          max_new_tokens: int, maximum number of new tokens to generate\n",
        "\n",
        "        Returns:\n",
        "          (B, T+max_new_tokens) tensor of context with new tokens appended\n",
        "        \"\"\"\n",
        "        # TODO: your code below\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2MaUUxb0EQY1"
      },
      "source": [
        "#### Question 1.1.8: Evaluating the initialization.\n",
        "\n",
        "Evaluate the loss of your untrained bigram model on a batch of data. Make sure the loss (negative log likelihood) is per-token (i.e. you may need to average over both sequence length and batch). Does this loss make sense in the context of your answer to Question 1.1.6? Discuss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xgSEGqKGEiOn"
      },
      "outputs": [],
      "source": [
        "x,y = get_batch(\"train\", context_window_size, device)\n",
        "bigram_model = BigramLanguageModel(vocab_size)\n",
        "bm = bigram_model.to(device)\n",
        "\n",
        "# TODO: your code below"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bS32m6egzg03"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IzxRpcgKE4_5"
      },
      "source": [
        "#### Question 1.1.9: Training your bigram model\n",
        "\n",
        "Train your bigram model for `SMALL_ITERS` iterations. Plot and interpret the loss curve.\n",
        "\n",
        "Our train loss gets down to around 4 in around 5 min of training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E2r390zbyz3O"
      },
      "outputs": [],
      "source": [
        "# create a PyTorch optimizer\n",
        "learning_rate = 1e-2\n",
        "optimizer = torch.optim.AdamW(bigram_model.parameters(), lr=learning_rate)\n",
        "\n",
        "eval_interval = 200\n",
        "eval_iters = 200\n",
        "\n",
        "loss_list = []\n",
        "\n",
        "for it in tqdm(range(SMALL_ITERS)):\n",
        "\n",
        "    # every once in a while evaluate the loss on train and val sets\n",
        "    if it % eval_interval == 0 or it == SMALL_ITERS - 1:\n",
        "        print(f\"iteration {it}\")\n",
        "        losses = estimate_loss(bm, eval_iters, context_window_size, device)\n",
        "        print(f\"step {it}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
        "\n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch('train', context_window_size, device)\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = bm(xb, yb)\n",
        "    loss_list.append(loss.detach().item())\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tj28W2w_ziDA"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that these models can take up a lot of memory on the GPU. As you go through this assignment, you may want to free the models after you train them using code along the lines of \n",
        "```\n",
        "bm.to('cpu')\n",
        "torch.cuda.empty_cache()\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "slaOmWgFLdMk"
      },
      "source": [
        "### 1.2: Token Embeddings: going from discrete tokens to continuous latent spaces\n",
        "\n",
        "In the look up table formulation of the bigram model, we are modelling the logits of the next token didstirbution independently for each token, even if two tokens are extremely similar to each other.\n",
        "One way arond this problem is to learn an embedding of the discrete tokens into $\\mathbb{R}^{D}$, and then to run multi-class logistic regression on top of this learned embedding.\n",
        "\n",
        "More precisely, if we have a vocabulary of tokens of size $V$ that we choose to embed in a Euclidean embedding space of dimension $D$, we can parameterize the distribution of the next token if the current token is $v$ according to\n",
        "\\begin{align*}\n",
        "  \\mathrm{Cat}\\Big( \\mathrm{softmax} (\\beta X_v) \\Big),\n",
        "\\end{align*}\n",
        "where $X_v \\in \\mathbb{R}^{D}$ is the learned embedding of token $v$ into $\\mathbb{R}^{D}$ and $\\beta \\in \\mathbb{R}^{V \\times D}$. Notice that if $X$ were a fixed design matrix this formulation would be equivalent to multi-class logistic regression. However, both $X$ and $\\beta$ are learnable parameters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kX6BRU07O7Cr"
      },
      "source": [
        "#### Question 1.2.1: Implement BigramWithWordEmbeddingsLM\n",
        "\n",
        "Implement a bigram languge model that uses a linear readout from a low dimensional Euclidean embedding of each token to parameterize the logits of the next token distribution, instead of parameterizing the logits of the next token distribution directly. It should have almost the same implementation as `BigramLanguageModel` from Question 1.1.6, except `init` should also take in an `embed_size`, and the `forward` method will need to be modified."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "brU5zILoQASX"
      },
      "outputs": [],
      "source": [
        "class BigramWithWordEmbeddingsLM(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, embed_size=32):\n",
        "      \"\"\"\n",
        "      Args:\n",
        "        vocab_size: int, size of the vocabulary\n",
        "        embed_size: int, dimension of the word embedding (D)\n",
        "      \"\"\"\n",
        "      #TODO, your code here\n",
        "      pass\n",
        "\n",
        "    def forward(self, token_ids, targets=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) token ids that make up the context (batch has size B, each entry in the batch has length T)\n",
        "          targets: (B, T) token ids corresponding to the target of each context in token_ids\n",
        "\n",
        "        Returns:\n",
        "          logits: (B, T, V), logits[b,t, :] gives the length V vector of logits for the next token prediction in string b up to t tokens\n",
        "          loss: scalar, negative log likelihood of target given context\n",
        "        \"\"\"\n",
        "        # TODO, your code here\n",
        "\n",
        "        logits = None\n",
        "        loss = None\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def generate(self, token_ids, max_new_tokens=context_window_size):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) tensor of token ids to provide as context\n",
        "          max_new_tokens: int, maximum number of new tokens to generate\n",
        "\n",
        "        Returns:\n",
        "          (B, T+max_new_tokens) tensor of context with new tokens appended\n",
        "        \"\"\"\n",
        "        #TODO\n",
        "        # your code below\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ypOE6LnnGdyu"
      },
      "source": [
        "### 1.3: Attention: Relaxing Markovian assumptions to transmit information across the sequence length\n",
        "\n",
        "A major problem with the bigram models of Sections 1.1 and 1.2 was that they were Markovian: the distribution of the next token was determined entirely by the current token! The attention mechanism provides a way to extract information between the previous tokens in the context to provide a better parameterization for the distribution of the next token."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F-XHeM1VRO6H"
      },
      "source": [
        "#### Question 1.3.1: Averaging over word embeddings\n",
        "\n",
        "One simple way to pool information would simply be to average the embeddings!\n",
        "\n",
        "Your TODO: Add comments to the the code snippet below. Write a description here explaining why the code is mathematically equivalent to averaging the embeddings of the previous tokens and the current token."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S-8W0Xx_zsst"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gCaS07DtGgCr"
      },
      "outputs": [],
      "source": [
        "# average word embedding via matrix multiply and softmax\n",
        "small_batch_size = 4              # B\n",
        "small_context_window_size = 8     # T\n",
        "small_embed_size = 2              # D\n",
        "\n",
        "# make \"synthetic\" word embeddings (for illustration purposes only)\n",
        "X = torch.randn(small_batch_size, small_context_window_size, small_embed_size)\n",
        "\n",
        "# TODO: comment the code below\n",
        "print(X.shape)\n",
        "tril = torch.tril(torch.ones(small_context_window_size, small_context_window_size))\n",
        "attn_weights = torch.zeros((small_context_window_size, small_context_window_size))\n",
        "attn_weights = attn_weights.masked_fill(tril == 0, float('-inf'))\n",
        "attn_weights = F.softmax(attn_weights, dim=-1)\n",
        "avgDmbds = attn_weights @ X\n",
        "print(X[0])\n",
        "print(\"\")\n",
        "print(avgDmbds[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oIDlnRDGXaie"
      },
      "source": [
        "#### 1.3.2: Single-headed scaled $(Q,K,V)$-attention\n",
        "\n",
        "A more sophisticated approach than simply averaging over previous word embeddings is single-headed (Query, Key, Value) scaled attention.\n",
        "That is, we now summarize the information contained in a length $T$ sequence of tokens that have been embeded into $X \\in \\mathbb{R}^{T \\times D}$ according to\n",
        "\\begin{equation}\n",
        "   \\mathrm{SoftmaxAcrossRows} \\Bigg( \\frac{\\mathrm{CausalMask}\\Big(X U_q^\\top U_k X^\\top \\Big)}{\\sqrt{K}} \\Bigg) \\Big( X V^\\top \\Big),\n",
        "\\end{equation}\n",
        "where $U_q, U_k \\in \\mathbb{R}^{K \\times D}$, $V \\in \\mathbb{R}^{D \\times D}$, and $K$ is the \"head size\".\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E84BzzxPdwUd"
      },
      "source": [
        "##### Question 1.3.2.1\n",
        "\n",
        "In the limiting case where $U_q$ and $U_k$ are all zeros, and $V = I_{D}$, what does $(U_q, U_k, V)$ attention simplify to?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AmFw8z5Dz3T6"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1cU9BXYPekkt"
      },
      "source": [
        "##### Question 1.3.2.2\n",
        "\n",
        "Imagine we had two matrices $U_q$ and $U_k$, both in $\\mathbb{R}^{K \\times D}$, where every entry was an independent standard normal.\n",
        "\n",
        "What would be the distribution of an element of $U_q^\\top U_k$? What about $U_q^\\top U_k / \\sqrt{K}$?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GDtilESJz4oC"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jeKfSrREhgdH"
      },
      "source": [
        "##### Question 1.3.2.3: Implement single-headed scaled $(U_q,U_k,V)$-attention.\n",
        "\n",
        "Complete the below code so the `forward` method returns single-headed scaled $(U_q,U_k,V)$-attention."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "O_eBPiT-Yy0q"
      },
      "outputs": [],
      "source": [
        "class Head(nn.Module):\n",
        "    \"\"\" one head of self-attention \"\"\"\n",
        "\n",
        "    def __init__(self, head_size, context_window_size, embed_size=384):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          head_size: int, size of the head embedding dimension (K)\n",
        "          context_window_size: int, number of tokens considered in the past for attention (T)\n",
        "          embed_size: int, size of the token embedding dimension (D)\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.head_size = head_size\n",
        "        self.key = nn.Linear(embed_size, head_size, bias=False)\n",
        "        self.query = nn.Linear(embed_size, head_size, bias=False)\n",
        "        self.value = nn.Linear(embed_size, embed_size, bias=False)\n",
        "\n",
        "        # not a param of the model, so registered as a buffer\n",
        "        self.register_buffer('tril', torch.tril(\n",
        "            torch.ones(context_window_size, context_window_size)))\n",
        "\n",
        "    def forward(self, x):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          x: (B,T,D) tensor of token embeddings\n",
        "\n",
        "        Returns:\n",
        "          (B,T,D) tensor of attention-weighted token embeddings\n",
        "        \"\"\"\n",
        "        # TODO: your code here\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3RqGrCB21qUV"
      },
      "source": [
        "##### Question 1.3.2.4: Implement a single-headed attention language model\n",
        "\n",
        "Complete the code below. Note that because the transformer has no idea where tokens are occuring in space, we have also added in position embeddings."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "REr3aWnS1xJL"
      },
      "outputs": [],
      "source": [
        "class SingleHeadedAttentionLM(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, context_window_size, head_size, embed_size=384):\n",
        "      \"\"\"\n",
        "      Args:\n",
        "        vocab_size: int, size of the vocabulary (V)\n",
        "        context_window_size: int, number of tokens considered in the past for attention (T)\n",
        "        head_size: int, size of the head embedding dimension (K)\n",
        "        embed_size: int, size of the token embedding dimension (D)\n",
        "      \"\"\"\n",
        "      super().__init__()\n",
        "      self.token_embedding_table = nn.Embedding(vocab_size, embed_size)\n",
        "      self.position_embedding_table = nn.Embedding(context_window_size, embed_size)\n",
        "      self.context_window_size = context_window_size\n",
        "\n",
        "      # TODO: your code below\n",
        "      self.atten_head = Head(...)\n",
        "      self.lm_head = nn.Linear(...)\n",
        "\n",
        "    def forward(self, token_ids, targets=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) token ids that make up the context (batch has size B, each entry\n",
        "                     in the batch has length T)\n",
        "          targets: (B, T) token ids corresponding to the target of each context in token_ids\n",
        "\n",
        "        Returns:\n",
        "          logits: (B, T, V) logits[b,t] gives the length V vector of logits for the next token\n",
        "                   prediction in string b up to t tokens\n",
        "          loss: scalar, negative log likelihood of target given context\n",
        "        \"\"\"\n",
        "        B, T = token_ids.shape # (batch size, length)\n",
        "        tok_emb = self.token_embedding_table(token_ids) # (B,T,D)\n",
        "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T,D)\n",
        "        x = tok_emb + pos_emb # (B,T,D)\n",
        "        x = self.atten_head(x) # (B,T,D)\n",
        "        logits = self.lm_head(x) # (B,T,V)\n",
        "\n",
        "        # TODO: your code here\n",
        "        logits = ...\n",
        "        loss = ...\n",
        "        return logits, loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def generate(self, token_ids, max_new_tokens):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) tensor of token ids to provide as context\n",
        "          max_new_tokens: int, maximum number of new tokens to generate\n",
        "\n",
        "        Returns:\n",
        "          (B, T+max_new_tokens) tensor of context with new tokens appended\n",
        "        \"\"\"\n",
        "        #TODO\n",
        "        # your code below\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "End-pUCa4RXe"
      },
      "source": [
        "Train your new `SingleHeadedAttentionLM` for `SMALL_ITERS` training iterations and plot the loss curve.\n",
        "The `head_size` shouldn't matter too much, we just use the `embedding_size`.\n",
        "Do you seen an improvement compared to your `BigramLanguageModel`? Discuss.\n",
        "\n",
        "Note: you may want to modify the learning rate. Training for `SMALL_ITERS` with a learning rate of `6e-4`, we can get to a train loss of around 3.4 in around 4 min of training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "070G_l0E0uG9"
      },
      "outputs": [],
      "source": [
        "embed_size = 384\n",
        "sha_model = SingleHeadedAttentionLM(vocab_size, context_window_size, embed_size, embed_size)\n",
        "sham = sha_model.to(device)\n",
        "learning_rate = 6e-4\n",
        "optimizer = torch.optim.AdamW(sha_model.parameters(), lr=learning_rate)\n",
        "\n",
        "loss_list = []\n",
        "\n",
        "for it in tqdm(range(SMALL_ITERS)):\n",
        "\n",
        "    # every once in a while evaluate the loss on train and val sets\n",
        "    if it % eval_interval == 0 or it == SMALL_ITERS - 1:\n",
        "        print(f\"iteration {it}\")\n",
        "        losses = estimate_loss(sham, eval_iters, context_window_size, device)\n",
        "        print(\n",
        "            f\"step {it}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\"\n",
        "        )\n",
        "\n",
        "    # sample a batch of data\n",
        "    xb, yb = get_batch(\"train\", context_window_size, device)\n",
        "\n",
        "    # evaluate the loss\n",
        "    logits, loss = sham(xb, yb)\n",
        "    loss_list.append(loss.detach().item())\n",
        "    optimizer.zero_grad(set_to_none=True)\n",
        "    loss.backward()\n",
        "    optimizer.step()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u05NFXpo0wWs"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7wjXYf-S4Lus"
      },
      "source": [
        "#### 1.3.3: Multi-headed attention"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZIObjMO0Ikp5"
      },
      "source": [
        "##### Question 1.3.3.1: Implement multi-headed attention"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6vb8NU_s6Vfg"
      },
      "outputs": [],
      "source": [
        "class MultiHeadAttention(nn.Module):\n",
        "    \"\"\" multiple heads of self-attention in parallel \"\"\"\n",
        "\n",
        "    def __init__(self, context_window_size, num_heads, head_size, embed_size=384):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            context_window_size: int, number of tokens considered in the past for attention (T)\n",
        "            num_heads: int, number of heads (H)\n",
        "            head_size: int, size of the head embedding dimension\n",
        "            embed_size: int, size of the token embedding dimension\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        # TODO, your code below\n",
        "        self.heads = nn.ModuleList(...)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # TODO, your code below\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aFCBRay3IuR3"
      },
      "source": [
        "##### Question 1.3.3.2: Implement a multi-headed attention LM\n",
        "\n",
        "Fill in the code below to create a language model that outputs its logits for next token prediction using multi-headed attention. Train your model for `SMALL_ITERS` training iterations. Compare the results with the single-headed attention model. Do you see an improvement?\n",
        "\n",
        "We get to a train loss of around 2.75 in around 5 mins of training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LvWHwcCzI1yr"
      },
      "outputs": [],
      "source": [
        "class MultiHeadedAttentionLM(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, context_window_size, embed_size=384, num_heads=6):\n",
        "      super().__init__()\n",
        "      self.head_size = embed_size // num_heads\n",
        "      self.context_window_size = context_window_size\n",
        "      # TODO: your code below\n",
        "\n",
        "    def forward(self, token_ids, targets=None):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) token ids that make up the context (batch has size B, each entry in the\n",
        "                     batch has length T)\n",
        "          targets: (B, T) token ids corresponding to the target of each context in token_ids\n",
        "\n",
        "        Returns:\n",
        "          logits: (B, T, V), logits[b,t] gives the length V vector of logits for the next token\n",
        "                  prediction in string b up to t tokens\n",
        "          loss: scalar, negative log likelihood of target given context\n",
        "        \"\"\"\n",
        "        # TODO: your code below\n",
        "        logits = ...\n",
        "        loss = ...\n",
        "        return logits, loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def generate(self, token_ids, max_new_tokens):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "          token_ids: (B, T) tensor of token ids to provide as context\n",
        "          max_new_tokens: int, maximum number of new tokens to generate\n",
        "\n",
        "        Returns:\n",
        "          (B, T+max_new_tokens) tensor of context with new tokens appended\n",
        "        \"\"\"\n",
        "        # TODO: your code below\n",
        "        pass"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A2tKVYDY1XG5"
      },
      "source": [
        "---\n",
        "\n",
        "_your answer here_\n",
        "\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X9XsUWpwG7uA"
      },
      "source": [
        "### 1.4: The Transformer Architecture: combining attention with deep learning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1GbGqwKWJzOK"
      },
      "outputs": [],
      "source": [
        "# run this cell to initialize this deep learning module that you should use in the code your write later\n",
        "# you don't need to edit this layer\n",
        "class FeedForward(nn.Module):\n",
        "    \"\"\" a simple linear layer followed by a non-linearity\n",
        "        Given to you, you don't need to write any code here!\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, embed_size):\n",
        "        super().__init__()\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Linear(embed_size, 4 * embed_size),\n",
        "            nn.ReLU(),\n",
        "            nn.Linear(4 * embed_size, embed_size),\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WKJxVp6aJb6i"
      },
      "source": [
        "#### Question 1.4.1: Implement a transformer block\n",
        "\n",
        "Complete the code below to implement a transformer block"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8Jw2I6a8IYOi"
      },
      "source": [
        "To make the your implemenation easier to train, we have added two deep learning best practices:\n",
        "\n",
        "1. Residual connections.\n",
        "In the `forward` method of the `TransformerBlock`, we have made the connections of the residual connection, which of the form\n",
        "\\begin{equation}\n",
        "  x = (I + N)(x),\n",
        "\\end{equation}\n",
        "where $I$ stands for the identity transformation and $N$ stands for some non-linearity. The idea is that every layer is some adjustment of the identity function, which allows gradients to flow through a deep network during back propogation, especially at initialization.\n",
        "\n",
        "2. Prenorm via `LayerNorm`\n",
        "Also in the `forward` method of the `TransformerBlock`, the nonlinearity first applied a `LayerNorm` to its arguments. The `LayerNorm` basically standardizes the neurons in that layer so that they have mean 0 and variance 1. Doing so is very helpful for numerical stability, espeically of the gradients."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hUDbIv9eISkf"
      },
      "outputs": [],
      "source": [
        "class TransformerBlock(nn.Module):\n",
        "    \"\"\" Transformer block: communication across sequence length, followed by communication across embedding space\n",
        "        Uses multi-headed attention\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, vocab_size, context_window_size, embed_size=384, num_heads=6):\n",
        "        super().__init__()\n",
        "        self.ln1 = nn.LayerNorm(embed_size)\n",
        "        self.ln2 = nn.LayerNorm(embed_size)\n",
        "\n",
        "        # TODO: your code below\n",
        "        self.feed_forward = FeedForward(...)\n",
        "        self.mh_attention = ...\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x + self.mh_attention(self.ln1(x)) # communication over sequence length\n",
        "        x = x + self.feed_forward(self.ln2(x)) # communication across embedding space\n",
        "        return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vqCCiCYcKMD0"
      },
      "source": [
        "#### Question 1.4.2: Implement your baseline transformer model\n",
        "\n",
        "We now stack 6 `TransformerBlocks` (with a final layer norm applied after the blocks but before the logits) to create our basline `TransformerLM`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "t2veTg9N3ufJ"
      },
      "outputs": [],
      "source": [
        "class TransformerLM(nn.Module):\n",
        "\n",
        "    def __init__(self, vocab_size, context_window_size, embed_size=384, num_heads=6, n_layers=6):\n",
        "        \"\"\"\n",
        "          Args:\n",
        "              vocab_size: int, number of tokens in the vocabulary (V)\n",
        "              context_window_size: int, size of the context window (T)\n",
        "              embed_size: int, embedding size (D)\n",
        "              num_heads: int, number of heads (H)\n",
        "              n_layers: int, number of layers (M)\n",
        "        \"\"\"\n",
        "        super().__init__()\n",
        "        self.token_embedding_table = nn.Embedding(vocab_size, embed_size)\n",
        "        self.position_embedding_table = nn.Embedding(context_window_size, embed_size)\n",
        "        self.blocks = nn.Sequential(*[\n",
        "            TransformerBlock(vocab_size,\n",
        "                             context_window_size,\n",
        "                             embed_size=embed_size,\n",
        "                             num_heads=num_heads)\n",
        "            for _ in range(n_layers)])\n",
        "\n",
        "        # final layer norm\n",
        "        self.ln_f = nn.LayerNorm(embed_size)\n",
        "        self.lm_head = nn.Linear(embed_size, vocab_size)\n",
        "\n",
        "        # good initialization\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, token_ids, targets=None):\n",
        "        \"\"\"\n",
        "        Agrgs:\n",
        "            token_ids: tensor of integers, provides the contet, shape (B, T)\n",
        "            targets: tensor of integers, provides the tokens we are preidcitng, shape (B, T)\n",
        "        \"\"\"\n",
        "        B, T = token_ids.shape\n",
        "\n",
        "        # token_ids and targets are both (B, T) tensor of integers\n",
        "        tok_emb = self.token_embedding_table(token_ids) # (B, T, D)\n",
        "        pos_emb = self.position_embedding_table(torch.arange(T, device=device)) # (T, D)\n",
        "        x = tok_emb + pos_emb # (B, T, D)\n",
        "\n",
        "        # TODO: your code below\n",
        "        logits = ...\n",
        "        loss = ...\n",
        "\n",
        "        return logits, loss\n",
        "\n",
        "    @torch.no_grad()\n",
        "    def generate(self, token_ids, max_new_tokens):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            token_ids: tensor of integers forming the context, shape (B, T)\n",
        "            max_new_tokens: int, max number of tokens to generate\n",
        "        \"\"\"\n",
        "        # TOOD, your code below\n",
        "        return token_ids"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JP8430nWKbZ6"
      },
      "source": [
        "Train your `TransformerLM` for `LARGE_ITERS` iterations and plot the loss curve. You may want to change the learning rate.\n",
        "\n",
        "We used a learning rate of `1e-4` and got to a final train loss of around 2.4 in around 30 mins of training."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsnbDpdhLeKo"
      },
      "outputs": [],
      "source": [
        "trans = TransformerLM(vocab_size, context_window_size)\n",
        "tlm = trans.to(device)\n",
        "learning_rate = 1e-4\n",
        "# TODO, your code below"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Jgvwzv5Ko7c"
      },
      "source": [
        "Generate an unconditional sample of length `context_window_size` from your trained `TransformerLM`, and also prompt it with the two prompts we gave you. How does the output look? Discuss?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "89-1t6MRLVi4"
      },
      "outputs": [],
      "source": [
        "# the contexts for the different prompts\n",
        "start_context = torch.zeros((1, 1), dtype=torch.long, device=device)\n",
        "print(f\"shape is {start_context.shape}\")\n",
        "context1 = torch.tensor(from_code_bert(tokenizer.encode(prompt_1_text)[:-1]), device=device).reshape(1, -1) # (1, T)\n",
        "print(f\"shape is {context1.shape}\")\n",
        "context2 = torch.tensor(from_code_bert(tokenizer.encode(prompt_2_text)[:-1])).to(device).reshape(1, -1)\n",
        "print(f\"shape is {context2.shape}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9_rxT1iSKzdO"
      },
      "outputs": [],
      "source": [
        "# unconditional generate from the transformer model\n",
        "uncond_gen = (tlm.generate(start_context, max_new_tokens=context_window_size)[0].tolist())\n",
        "print(tokenizer.decode(to_code_bert(uncond_gen)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1-FeFwD8K4jj"
      },
      "outputs": [],
      "source": [
        "# conditional generation of newton's method\n",
        "# TODO, your code here"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tA2TnXBzK6mN"
      },
      "outputs": [],
      "source": [
        "# conditional generation of cosine distance\n",
        "# TODO, your code here"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LgOzvWFDx_yH"
      },
      "source": [
        "#### Question 1.4.3\n",
        "\n",
        "The negative log-likelihood (averaged per token) we have been using to train our models can be expressed as\n",
        "\\begin{equation*}\n",
        "  L = -\\frac{1}{T} \\sum_{t = 1}^{T} \\log p(s[t] | \\text{context})\n",
        "\\end{equation*}\n",
        "for some document $s$, where $s[t]$ is the $t$th token of the doc. The natural language processing (NLP) community often reports the quantity\n",
        "\\begin{equation*}\n",
        "  \\text{perplexity} = \\exp(L).\n",
        "\\end{equation*}\n",
        "\n",
        "Give an intuitive interpretation of what perplexity is. Why might it be a more intuitive or natual measure to report than negative log-likelihood? Does the reported perplexity of your trained `TransformerLM` model make sense in terms of samples it generates? (Be sure to distinguish betwen `train` and `validation` perplexity. Which of `train` and `val` perplexity is more helpful for understanding your generated samples? Why?). (*Hint: your answer to Question 1.1.6 may be helpful*)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Izr1wTOjzlo"
      },
      "source": [
        "## Part 2: Mini-Project"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5lF3jFrQj1f4"
      },
      "source": [
        "Quick recap: So far we have\n",
        "\n",
        "1. Preprocessed the python code dataset by encoding text into integer tokens.\n",
        "2. Implemented single headed attention and then further generalized to multiheaded attention. We further combined multiheaded attention with deep learning to create the transformer architecture.\n",
        "3. Trained our transformer and generate code output.\n",
        "\n",
        "Up to this point, the performance of our simple language model has clearly made a lot of progress. We can see that our model has learned to generate in the style of python code syntax, although there are many quirks that suggest it will not make a very practical code assistant in its current state.\n",
        "\n",
        "### Project Outline\n",
        "\n",
        "Find some area of possible improvement.\n",
        "We interpret \"improvement\" quite loosely, but it is up to you to state precisely in what sense your proposed innovation might constitute an improvement and to show convincing evidence that your innovation does or does not constitue an improvement according to your definition.\n",
        "For your idea, **formulate a hypothesis** for why this change should result in a better model. **Implement your changes** and **report any findings**.\n",
        "\n",
        "_Notes_: As this assignment is being treated as a project, you should expect training to take longer than previous assignments. However, please use your judgement to decide what is reasonable. We will not expect you to run training procedures that take more than 2 hours on the free Google Colab computing resources and we certainly do not expect you to acquire additional compute. The proposed improvements should not solely rely on increased computing demands, but must be based on the goal of improving the model by more efficiently learning from our data.\n",
        "\n",
        "_Hints_: There are many aspects to assessing our model. For example, not only is quality of generated text important, it is also of interest to reduce costs associated with training.\n",
        "\n",
        "### Deliverables\n",
        "\n",
        "In addition to a pdf of your python notebook, the submission for this project will be a written report no more than 4 pages in length using the [NeurIPS LaTex template](https://neurips.cc/Conferences/2023/PaperInformation/StyleFiles). Your report should include detailed analysis of the hypotheses you chose to test along with any conclusions.\n",
        "\n",
        "The page limit for the report does not include bibliography or appendices. Make sure to keep the \"ready for submission\" option to help us grade anonymously. One of your apprendices should contain a link to any code used to generate the project so that we can grade it (google drive with colab nbs or github repo are both fine). You should have at least one plot in your main text (which is capped at 4 pages).\n",
        "\n",
        "### Data augmentation\n",
        "\n",
        "We got the data for this project from [The Stack](https://huggingface.co/datasets/bigcode/the-stack-dedup). If you'd like, you can definitely train on larger datasets by accessing their dataset of python code (we just scratched the surface). You have to make an account on Hugginface to get a Hugginface access token, but the process is pretty quick."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Submission Instructions\n",
        "\n",
        "You will generate two PDFs: one from parts 0 and 1, which involves completing this colab to create a transformer baseline; and one from the mini-project in part 2, which will be your write-up of no longer than 4 pages.\n",
        "\n",
        "**Combine the two PDFs into a single PDF and submit on gradescope. Tag your PDF correctly.**\n",
        "\n",
        "If you work in a group of two, submit one assignment on gradescope. If you complete the assignment individually, submit as usual."
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}