attention_and_transformers.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyP2x6aAEiFOzVM+7UDVUREQ",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/pmadhyastha/INM434/blob/main/attention_and_transformers.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "-3GV8WFHU7ha"
      },
      "outputs": [],
      "source": [
        "__author__ = \"Pranava Madhyastha\" \n",
        "__version__ = \"INM434/IN3045 City, University of London, Spring 2023\""
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Attention and Transformers\n",
        "\n",
        "Today we will try take a peak into concepts of attention and the modelling framework called Transformers. \n",
        "\n",
        "### TODO: Please print out the values and verify if you are able to understand each step below. Use lecture slides as the reference. "
      ],
      "metadata": {
        "id": "faYI0K4OVGmI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn.functional as F"
      ],
      "metadata": {
        "id": "VTHag6pUwiuB"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# The sample sentence\n",
        "\n",
        "We will begin with a sample sentence that we have seen in the lectures: `the cat is on a new mat`\n"
      ],
      "metadata": {
        "id": "4iKwD82bwxXv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the input sentence\n",
        "input_sentence = 'the cat is on a new mat'"
      ],
      "metadata": {
        "id": "L7FeJaWrwvEQ"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creation of the dictionary \n",
        "\n",
        "The next step is to create a dictionary that maps each word in the input sentence to a unique index. This is done by splitting the sentence into a list of words. \n",
        "\n",
        "The list of words is then sorted, and each word is assigned an index using the enumerate() function. The result is a dictionary that maps each word to a unique index."
      ],
      "metadata": {
        "id": "SF2QO9aqxO6i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a dictionary that maps each word to a unique index\n",
        "word_to_index = {word: i for i, word in enumerate(sorted(input_sentence.split()))}"
      ],
      "metadata": {
        "id": "95unmyfExKCW"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Tensor of indices\n",
        "\n",
        "The input sentence is then converted to a tensor of indices using the word_to_index dictionary. The tensor is created by iterating over each word in the input sentence (after removing the comma) and looking up its index in the dictionary. The resulting tensor is a one-dimensional tensor of integer indices."
      ],
      "metadata": {
        "id": "lJ49Ud1yxmED"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Convert the input sentence to a tensor of indices using the word_to_index dictionary\n",
        "input_tensor = torch.tensor([word_to_index[word] for word in input_sentence.replace(',', '').split()])\n"
      ],
      "metadata": {
        "id": "OvF9HY8oxfcT"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# The embedding layer\n",
        "\n",
        "The next step is to define an embedding layer. The embedding layer is a PyTorch module that learns a lookup table of word embeddings for a vocabulary of size n and embedding dimensionality d. In this case, the embedding layer is initialized with a vocabulary size equal to the number of words in the input sentence and an embedding dimensionality of 16. The input tensor is then embedded using the embedding layer to produce an embedded sentence. The .detach() method is used to detach the embedded sentence from its computation graph, preventing it from being backpropagated through."
      ],
      "metadata": {
        "id": "HMAlXDKPyRdK"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the embedding layer\n",
        "embedding_layer = torch.nn.Embedding(len(word_to_index), 16)\n",
        "\n",
        "# Embed the input tensor to get the embedded sentence\n",
        "embedded_sentence = embedding_layer(input_tensor).detach()"
      ],
      "metadata": {
        "id": "b7shFqe2yLUG"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The size of the embedding is then defined by taking the shape of the embedded sentence and extracting its second dimension."
      ],
      "metadata": {
        "id": "gtCWVJIiygRI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the size of the embedding\n",
        "embedding_size = embedded_sentence.shape[1]"
      ],
      "metadata": {
        "id": "K30s2Z-IyfSR"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Query, key and value \n",
        "\n",
        "The sizes of the query, key, and value vectors are defined. Recall that these vectors are used in the attention mechanism to compute the attention weights and context vectors. In this case, the query and key vectors have a size of 24, while the value vector has a size of 28."
      ],
      "metadata": {
        "id": "LlYAxyeAyqEa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the sizes of the query, key, and value vectors\n",
        "query_size, key_size, value_size = 24, 24, 28\n"
      ],
      "metadata": {
        "id": "0rlFHj4uymmh"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Initialisation and projection\n",
        "\n",
        "Query, key, and value weight matrices are then defined using PyTorch's random number generator. The weight matrices are used to compute the query, key, and value vectors for each word in the input sentence. Refer to the lecture slides! "
      ],
      "metadata": {
        "id": "vgDuIhqCy-13"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the query, key, and value weight matrices\n",
        "query_weights = torch.rand(query_size, embedding_size)\n",
        "key_weights = torch.rand(key_size, embedding_size)\n",
        "value_weights = torch.rand(value_size, embedding_size)"
      ],
      "metadata": {
        "id": "UjWzWPV-y2ZN"
      },
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Example computation \n",
        "\n",
        "The query, key, and value vectors for the second word in the input sentence are then computed by multiplying the corresponding weight matrices by the embedding vector for the second word."
      ],
      "metadata": {
        "id": "MnS_5omvzNzL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Get the query, key, and value vectors for the second word in the input sentence\n",
        "x_2 = embedded_sentence[1]\n",
        "query_2 = query_weights.matmul(x_2)\n",
        "key_2 = key_weights.matmul(x_2)\n",
        "value_2 = value_weights.matmul(x_2)"
      ],
      "metadata": {
        "id": "Khz_QhbZzNAR"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Parallelisation\n",
        "\n",
        "The keys and values for all words in the input sentence are then computed by multiplying the corresponding weight matrices by the transpose of the embedded sentence. The transpose operation is used to swap the rows and columns of the embedded sentence, making it possible to compute the dot product between the weight matrices and the embedding vectors. The transpose operation is then applied again to swap the rows and columns back to their original order."
      ],
      "metadata": {
        "id": "vWr178wnzUTT"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute the keys and values for all words in the input sentence\n",
        "keys = key_weights.matmul(embedded_sentence.transpose(0, 1)).transpose(0, 1)\n",
        "values = value_weights.matmul(embedded_sentence.transpose(0, 1)).transpose(0, 1)"
      ],
      "metadata": {
        "id": "02XgwGXgzIe1"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Computing the attention weights\n",
        "\n",
        "The attention weights for the second word in the input sentence are then computed using the softmax function applied to the dot product of the query vector for the second word and the transpose of the key vectors for all words in the input sentence. The softmax function normalizes the dot product, resulting in a set of attention weights that sum to one. The size of the key vectors is used to scale the dot product.\n",
        "\n"
      ],
      "metadata": {
        "id": "4fqAaTTGzocI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute the keys and values for all words in the input sentence\n",
        "keys = key_weights.matmul(embedded_sentence.transpose(0, 1)).transpose(0, 1)\n",
        "values = value_weights.matmul(embedded_sentence.transpose(0, 1)).transpose(0, 1)"
      ],
      "metadata": {
        "id": "zxTyVKiazgoN"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Computing `h` or the context vector\n",
        "\n",
        "The context vector for the second word in the input sentence is then computed by multiplying the attention weights by the value vectors for all words in the input sentence. The result is a weighted sum of the value vectors, where the weights are given by the attention weights."
      ],
      "metadata": {
        "id": "LobyxoeYz9dH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute the attention weights for the second word in the input sentence\n",
        "omega_2 = query_2.matmul(keys.T)\n",
        "attention_weights_2 = F.softmax(omega_2 / key_size **0.5, dim=0)\n",
        "\n",
        "\n",
        "# Compute the context vector for the second word in the input sentence\n",
        "context_vector_2 = attention_weights_2.matmul(values)"
      ],
      "metadata": {
        "id": "hhN4B0Ndz8FO"
      },
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Multi-head attention\n",
        "\n",
        "The code then defines the number of attention heads to use, which is set to 3. It also defines new query, key, and value weight matrices for the multi-head attention mechanism. The size of each weight matrix is expanded to include a third dimension for the number of attention heads.\n",
        "\n",
        "Refer to the lecture notes to understand the concept of Multi-head attention."
      ],
      "metadata": {
        "id": "cCTaZnl80m2I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the number of attention heads\n",
        "num_heads = 3\n",
        "\n",
        "# Define the query, key, and value weight matrices for the multi-head attention\n",
        "multihead_query_weights = torch.rand(num_heads, query_size, embedding_size)\n",
        "multihead_key_weights = torch.rand(num_heads, key_size, embedding_size)\n",
        "multihead_value_weights = torch.rand(num_heads, value_size, embedding_size)\n"
      ],
      "metadata": {
        "id": "h--ZzUjF0Tey"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The query vector for the second word in the input sentence is obtained by multiplying the embedded vector of the second word with the multi-head query weight matrix. This is done separately for each attention head."
      ],
      "metadata": {
        "id": "i4oUw2gQ1J0k"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Get the query vector for the second word in the input sentence for each attention head\n",
        "multihead_query_2 = multihead_query_weights.matmul(x_2)"
      ],
      "metadata": {
        "id": "pkSnvDfc02Sg"
      },
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Repeat \n",
        "\n",
        "The embedded sentence is repeated for each attention head using the repeat function. This allows each head to have its own set of keys and values for computing attention."
      ],
      "metadata": {
        "id": "ebgrqA6d1P5z"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Repeat the embedded sentence for each attention head\n",
        "repeated_embedded_sentence = embedded_sentence.unsqueeze(0).repeat(num_heads, 1, 1)"
      ],
      "metadata": {
        "id": "9dISW1WD1O9_"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The keys and values for all words in the input sentence are computed for each attention head by multiplying the repeated embedded sentence with the multi-head key and value weight matrices. The transpose function is used to reshape the tensor dimensions so that the matrix multiplication can be computed."
      ],
      "metadata": {
        "id": "6k_8xLVk13RR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Compute the keys and values for all words in the input sentence for each attention head\n",
        "multihead_keys = multihead_key_weights.matmul(repeated_embedded_sentence.transpose(1, 2)).transpose(1, 2)\n",
        "multihead_values = multihead_value_weights.matmul(repeated_embedded_sentence.transpose(1, 2)).transpose(1, 2)\n",
        "\n",
        "# Compute the attention weights and context vectors for each attention head\n",
        "attention_weights = F.softmax(multihead_query_2.matmul(multihead_keys.transpose(1, 2)) / key_size**0.5, dim=2)\n",
        "context_vectors = attention_weights.matmul(multihead_values)\n"
      ],
      "metadata": {
        "id": "D_m6Ts231hMy"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### TODO: Please follow the annotated transformer: http://nlp.seas.harvard.edu/annotated-transformer/ and consider re-writing and running it here on google colab. \n",
        "\n"
      ],
      "metadata": {
        "id": "iCSGvaiW3RKo"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# BERT based classifier\n",
        "\n",
        "In the following sections, we are going to work on the BERT classifier for sentiment analysis - trained and tested over the IMDB dataset. "
      ],
      "metadata": {
        "id": "TL4_s3RsRHDz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install transformers\n",
        "!pip install datasets # remember lab3? "
      ],
      "metadata": {
        "id": "VaIlio58Rs9S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The above code gets us pre-optimised versions (and pre-trained weights) for a BERT based model. We are also downloading the datasets library - we had also done this before. "
      ],
      "metadata": {
        "id": "0zggfFZSSAIH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "from transformers import DistilBertModel, DistilBertTokenizer, AdamW\n",
        "from datasets import load_dataset\n",
        "\n",
        "# Load the dataset\n",
        "\n",
        "dataset = load_dataset('imdb')"
      ],
      "metadata": {
        "id": "6q1G2KQi50Sy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will not work with BERT directly (as it may take a while to run), instead we shall use a \"distilled version\" of BERT (reduced model size) called [DistillBERT](https://arxiv.org/abs/1910.01108) . The code below loads the IMDB dataset from Hugging Face and preprocesses it. It loads the tokenizer and model from the transformers library, tokenizes the text using the DistilBERT tokenizer, and converts the data to the PyTorch format."
      ],
      "metadata": {
        "id": "xlhoqbbQSAo-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the tokenizer and model\n",
        "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
        "model = DistilBertModel.from_pretrained('distilbert-base-uncased')\n"
      ],
      "metadata": {
        "id": "HBTdX5htSBGG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code below just defines the architecture the model architecture using a PyTorch neural network module. The SentimentClassifier class takes the DistilBERT model as input and defines the forward method. It extracts the last hidden state from the output of the DistilBERT model, and passes it through a linear layer to get the final logits."
      ],
      "metadata": {
        "id": "7BcIJ7ASSKWW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the model architecture\n",
        "class SentimentClassifier(nn.Module):\n",
        "    def __init__(self, model):\n",
        "        super(SentimentClassifier, self).__init__()\n",
        "        self.model = model\n",
        "        self.linear = nn.Linear(768, 1)\n",
        "        \n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)\n",
        "        last_hidden_state = outputs.last_hidden_state[:, 0, :]\n",
        "        logits = self.linear(last_hidden_state)\n",
        "        return logits.squeeze(-1)\n",
        "\n",
        "# Instantiate the model\n",
        "model = SentimentClassifier(model)\n",
        "\n"
      ],
      "metadata": {
        "id": "y6lsT9JhSKn_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code below now defines the training and eval loop. It sets the model to training mode and iterates through each batch of the training data. It calculates the loss, computes the gradients, and updates the model parameters using the AdamW optimizer."
      ],
      "metadata": {
        "id": "Wc4rbpATSNB-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Tokenize the data\n",
        "def tokenize(batch):\n",
        "    return tokenizer(batch['text'], padding=True, truncation=True)\n",
        "\n",
        "train_dataset = dataset['train'].map(tokenize, batched=True)\n",
        "test_dataset = dataset['test'].map(tokenize, batched=True)\n",
        "\n",
        "# Set the data format\n",
        "train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])\n",
        "test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])\n",
        "\n",
        "# Set up the optimizer and loss function\n",
        "optimizer = AdamW(model.parameters(), lr=5e-5)\n",
        "loss_fn = nn.BCEWithLogitsLoss()\n",
        "\n",
        "# Define the training loop\n",
        "def train(model, train_loader, optimizer, loss_fn, device):\n",
        "    model.train()\n",
        "    for batch in train_loader:\n",
        "        input_ids = batch['input_ids'].to(device)\n",
        "        attention_mask = batch['attention_mask'].to(device)\n",
        "        labels = batch['label'].to(device)\n",
        "\n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(input_ids, attention_mask)\n",
        "        loss = loss_fn(outputs, labels.float())\n",
        "        print(loss)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "# Define the evaluation loop\n",
        "def evaluate(model, test_loader, device):\n",
        "    model.eval()\n",
        "    correct, total = 0, 0\n",
        "    with torch.no_grad():\n",
        "        for batch in test_loader:\n",
        "            input_ids = batch['input_ids'].to(device)\n",
        "            attention_mask = batch['attention_mask'].to(device)\n",
        "            labels = batch['label'].to(device)\n",
        "\n",
        "            outputs = model(input_ids, attention_mask)\n",
        "            predictions = torch.round(torch.sigmoid(outputs))\n",
        "\n",
        "            total += labels.size(0)\n",
        "            correct += (predictions == labels).sum().item()\n",
        "\n",
        "    return correct / total\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "JtALovXASNSF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will now formally train the model. "
      ],
      "metadata": {
        "id": "R0ZVG3T-SWXi"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Train the model\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "model.to(device)\n",
        "\n",
        "train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True)\n",
        "test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=32, shuffle=False)\n",
        "\n",
        "for epoch in range(1):\n",
        "    train(model, train_loader, optimizer, loss_fn, device)\n",
        "    accuracy = evaluate(model, test_loader, device)\n",
        "    print(f'Epoch {epoch+1} - Test Accuracy: {accuracy:.3f}')"
      ],
      "metadata": {
        "id": "fz7w8eITSYGO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Todo: \n",
        "- Compare this with your models from Lab 3 and Lab 4. \n",
        "- Test it on other sentiment datasets, does it generalise across? \n",
        "- How good is the generalisation? \n",
        "- Have a look at huggingface library and replace distillBERT with the BERT-base-uncased version, rerun the model. Compare the results. \n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "fALp7AjiSYzz"
      }
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "XM5jl-yXUQOk"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}