sentiment_analysis_starter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sentiment-140 Classification\n",
    "\n",
    "In this notebook, we will explore the application of transformers for sentiment analysis of different tweets. Our goal is to classify tweets to have either positive or negative sentiment. We will first consider the performance of a pre-trained BERT model, finetuned on the Sentiment-140 dataset, and later we'll develop and train our own transformer architecture from scratch. Let's dive right in!\n",
    "\n",
    "First of all, we need a couple of useful libraries that you're hopefully familiar with."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, SubsetRandomSampler\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from transformers import BertTokenizer, BertForSequenceClassification\n",
    "from data_utils.SentimentDataset import SentimentDataset\n",
    "\n",
    "# this automatically reloads the libraries so you can update them dynamically\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading our data\n",
    "\n",
    "Let's start by preparing our dataset. We have implemented the dataset class for you in `data_utils/SentimentDataset.py`, feel free to check the implementation. However, before we can use the `SentimentDataset` class to create our train data and test data objects, we need to pre-process the raw data. \n",
    "\n",
    "You can download raw data from [here](https://www.kaggle.com/datasets/kazanova/sentiment140). If you examine the raw data file, you can see that there is a lot of redundant information such as the time of each tweet or usernames. For our sentiment analysis, we simply need the tweets and ground truth labels. To preprocess the data file, simply paste the original CSV file in the root project directory without changing the file name and run the script `preprocess_dataset.py`. This should create a new CSV file called `dataset.csv`. (The reason why we do not include the data files on the GitHub repository is that they are simply too big.) Then, you're ready to create a dataset and dataloader below.\n",
    "\n",
    "Notice that the original dataset contains 1.6 million tweets; that is huge! Therefore, it may take forever to train on so much data, especially if you happen not to have access to a fast GPU (talk to the officers about getting access to the ACM AI server!). Therefore, in the next cell you can configure whether to use the whole dataset or just a subset of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "use_subset_of_data = True\n",
    "subset_percentage = 0.05 # set the percentage of original dataset size you'd like to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = SentimentDataset('dataset.csv', training_set=True)\n",
    "test_data = SentimentDataset('dataset.csv', training_set=False)\n",
    "\n",
    "if use_subset_of_data:\n",
    "    train_set_len = int(subset_percentage * len(train_data))\n",
    "    test_set_len = int(subset_percentage * len(test_data))\n",
    "    \n",
    "    train_subset_indices = np.random.choice(len(train_data), size=train_set_len, replace=False)\n",
    "    train_sampler = SubsetRandomSampler(train_subset_indices)\n",
    "\n",
    "    test_subset_indices = np.random.choice(len(test_data), size=test_set_len, replace=False)\n",
    "    test_sampler = SubsetRandomSampler(test_subset_indices)\n",
    "\n",
    "    train_loader = DataLoader(train_data, batch_size=32, sampler=train_sampler)\n",
    "    test_loader = DataLoader(test_data, batch_size=32, sampler=test_sampler)\n",
    "\n",
    "else:\n",
    "    train_loader = DataLoader(train_data, batch_size=32, shuffle=True)\n",
    "    test_loader = DataLoader(test_data, batch_size=32, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's examine a sample tweet from the dataset and print its corresponding ground truth label. Label 1 corresponds to positive sentiment, while label 0 stands for negative sentiment. You can run the cell below multiple times to see different tweets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tweets, labels = next(iter(train_loader))\n",
    "print(\"Sample tweet: \", tweets[0])\n",
    "print(\"Ground truth label: \", labels[0].numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tuning BERT\n",
    "\n",
    "First, let's examine the performance of a pre-trained BERT model on the sentiment analysis task. BERT (original paper [here](https://arxiv.org/abs/1810.04805)) is a language representation model released by Google in 2018 that uses a Bidirectional Transformer architecture.\n",
    "\n",
    "It can be easily fine-tuned on downstream tasks such as sentiment analysis by appending a classification layer to the original BERT architecture. This is already done automatically by importing the `BertForSequenceClassification` from the Hugging Face library called `transformers`. The newly added classification layer is untrained (as you'll see in the warning that appears when you run the following cell) and so we have to fine-tune on the Sentiment-140 data. We will also import a pre-trained tokenizer, which is necessary to convert inputs from natural language to a numerical representation machine learning models can understand. The tokenizer imported from Hugging Face uses WordPiece tokenization, you can learn more about the algorithm [here](https://huggingface.co/learn/nlp-course/en/chapter6/6)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"bert-base-uncased\"\n",
    "tokenizer = BertTokenizer.from_pretrained(model_name)\n",
    "pretrained_model = BertForSequenceClassification.from_pretrained(model_name, num_labels=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below, we define a `preprocess` function that takes in the tweets and tokenizes them. The tokenizer will break down words into smaller units called subword tokens and map these words to numbers (IDs), which can be processed by the model. Furthermore, it will add special tokens to the sequence that will normalize inputs to the model.\n",
    "\n",
    "The `preprocess` function return `input_ids` which represents the tokenized input sequence and `attention_mask` which specifies to the model which tokens to attend to while computing self-attention (more on that later!)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess(tweets):\n",
    "\n",
    "    encoded_batch = tokenizer.batch_encode_plus(\n",
    "        tweets,\n",
    "        add_special_tokens=True,\n",
    "        max_length=128,\n",
    "        padding='max_length',\n",
    "        truncation=True,\n",
    "        return_attention_mask=True,\n",
    "        return_tensors='pt'\n",
    "    )\n",
    "\n",
    "    input_ids = encoded_batch['input_ids']\n",
    "    attention_masks = encoded_batch['attention_mask']\n",
    "\n",
    "    return input_ids, attention_masks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are three different types of special tokens normally used by the BERT tokenizer. `[CLS]` token is added at the beginning of the input sequence. The final hidden state corresponding to the `[CLS]` token is typically used as the representation of the entire input sequence for classification tasks, so it will be very useful to us later for sentiment analysis. The `[SEP]` token is used to separate different segments of the input sequence. It helps the model distinguish between different parts of the input and process them accordingly. In our case, one `[SEP]` token will typically be added at the end of the tokenized input sequence (tweet). Finally, the `[PAD]` token is used for padding shorter input sequences to a fixed length, so that they can be easily processed in batches without any shape mismatch issues. With BERT, we would typically pad the input sequence to reach a total length of 512 tokens; however, since we're analyzing tweets that are by definition very short, using maximum sequence length of 128 tokens should suffice.\n",
    "\n",
    "Below, you can see what the IDs reserved for special tokens are."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_token_id = tokenizer.cls_token_id\n",
    "sep_token_id = tokenizer.sep_token_id\n",
    "pad_token_id = tokenizer.pad_token_id\n",
    "\n",
    "print(\"[CLS] Token ID: \" + str(cls_token_id) + \", [SEP] Token ID: \" + str(sep_token_id) + \", [PAD] Token ID: \" + str(pad_token_id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can print a sample tweet and its corresponding tokenized representation. Notice that each input sequence starts with token ID 101 (CLS token), ends with token ID 102 (SEP token), and is padded with token ID 0 (PAD token)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tweets, labels = next(iter(train_loader))\n",
    "print(\"Sample tweet: \", tweets[0], '\\n')\n",
    "input_ids, att_masks = preprocess(tweets)\n",
    "print(\"Tokenized representation: \", input_ids[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we need to define a training and evaluation loop. Notice that we use the `preprocess` function in both the `train` and `evaluate` to transform the natural language input to tokenized representation. We will also collect loss and accuracy statistics during training, so that we can plot the training history later.\n",
    "\n",
    "Note that the `finetuning` parameter is used simply because the internal Hugging Face representation of the pre-trained model is slightly different from what our representation will be when we later define our own transformer model. Thus, when using the training loop for fine-tuning the pre-trained downloaded from Hugging Face, set `finetuning=True`, otherwise set it to `False`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, val_loader, optimizer, criterion, device,\n",
    "          num_epochs, finetuning=False):\n",
    "\n",
    "    # Place model on device\n",
    "    model = model.to(device)\n",
    "\n",
    "    loss_history = []\n",
    "    acc_history = []\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()  # Set model to training mode\n",
    "\n",
    "        # Use tqdm to display a progress bar during training\n",
    "        with tqdm(total=len(train_loader),\n",
    "                  desc=f'Epoch {epoch + 1}/{num_epochs}',\n",
    "                  position=0,\n",
    "                  leave=True) as pbar:\n",
    "            \n",
    "            for inputs, labels in train_loader:\n",
    "\n",
    "                input_ids, attention_masks = preprocess(inputs)\n",
    "                # Move inputs and labels to device\n",
    "                input_ids = input_ids.to(device)\n",
    "                attention_masks = attention_masks.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # Zero out gradients\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # Compute the logits and loss\n",
    "                outputs = model(input_ids=input_ids, attention_mask=attention_masks)\n",
    "                if finetuning: \n",
    "                    outputs = outputs.logits\n",
    "                loss = criterion(outputs, labels)\n",
    "\n",
    "                # Backpropagate the loss\n",
    "                loss.backward()\n",
    "\n",
    "                # Update the weights\n",
    "                optimizer.step()\n",
    "\n",
    "                # Update the progress bar\n",
    "                pbar.update(1)\n",
    "                pbar.set_postfix(loss=loss.item())\n",
    "\n",
    "                loss_history.append(loss.item())\n",
    "\n",
    "                    \n",
    "        avg_loss, accuracy = evaluate(model, val_loader, criterion, device, finetuning)\n",
    "        print(\n",
    "            f'Test set: Average loss = {avg_loss:.4f}, Accuracy = {accuracy:.4f}'\n",
    "        )\n",
    "        acc_history.append(accuracy)\n",
    "\n",
    "    return loss_history, acc_history\n",
    "\n",
    "\n",
    "def evaluate(model, test_loader, criterion, device, finetuning=False):\n",
    "\n",
    "    model.eval()  # Set model to evaluation mode\n",
    "\n",
    "    with torch.no_grad():\n",
    "        total_loss = 0.0\n",
    "        num_correct = 0\n",
    "        num_samples = 0\n",
    "\n",
    "        for inputs, labels in test_loader:\n",
    "\n",
    "            input_ids, attention_masks = preprocess(inputs)\n",
    "            # Move inputs and labels to device\n",
    "            input_ids = input_ids.to(device)\n",
    "            attention_masks = attention_masks.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            # Compute the logits and loss\n",
    "            outputs = model(input_ids=input_ids, attention_mask=attention_masks)\n",
    "            if finetuning: \n",
    "                outputs = outputs.logits\n",
    "            loss = criterion(outputs, labels)\n",
    "            total_loss += loss.item()\n",
    "\n",
    "            # Compute the accuracy\n",
    "            _, predictions = torch.max(outputs, dim=1)\n",
    "            num_correct += (predictions == labels).sum().item()\n",
    "            num_samples += len(labels)\n",
    "\n",
    "    # Compute the average loss and accuracy\n",
    "    avg_loss = total_loss / len(test_loader)\n",
    "    accuracy = num_correct / num_samples\n",
    "\n",
    "    return avg_loss, accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure you are running the model on CUDA."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will finetune all the weights in the BERT model (including the pre-trained weights and the newly initialized classifier weights). An alternative approach would be to freeze the pre-trained weights and only fine-tune the classifier weights. You can do that by passing in `model.classifier.parameters()` to the optimizer instead of passing in all the weights of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.AdamW(pretrained_model.parameters(), lr=2e-5)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "loss_history, acc_history = train(pretrained_model, train_loader, test_loader, optimizer, criterion, device, num_epochs=1, finetuning=True)\n",
    "\n",
    "# Let's save our trained weights\n",
    "torch.save(pretrained_model.state_dict(), 'bert_fine_tuned.pt')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's plot our training loss and test accuracy history that documents the fine-tuning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(loss_history)\n",
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Cross-entropy Loss')\n",
    "plt.title('BERT Fine-tuning, Loss history')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your model is fine-tuned, feel free to test it on some of your own tweets!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def classify_tweet(tweet, model, device, finetuning=False):\n",
    "\n",
    "    model.eval()  # Set model to evaluation mode\n",
    "\n",
    "    with torch.no_grad():\n",
    "\n",
    "        tweet = [tweet]\n",
    "        input_ids, attention_mask = preprocess(tweet)\n",
    "        # Move inputs and labels to device\n",
    "        input_ids = input_ids.to(device)\n",
    "        attention_mask = attention_mask.to(device)\n",
    "\n",
    "        output = model(input_ids=input_ids, attention_mask=attention_mask)\n",
    "        if finetuning:\n",
    "            output = output.logits\n",
    "\n",
    "        _, prediction = torch.max(output, dim=1)\n",
    "\n",
    "    return prediction.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Just replace the string below by your own.\n",
    "sample_tweet = 'Visiting a friend in Santa Barbara today. So fun!!!'\n",
    "\n",
    "predicted_label = classify_tweet(sample_tweet, pretrained_model, device, finetuning=True)\n",
    "classification = 'positive' if predicted_label == 1 else 'negative'\n",
    "print(f'Your tweet was classified as {classification}.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating our own Transformer architecture from scratch\n",
    "\n",
    "That was fun, right? But maybe slightly too easy, don't you think? Let's try to implement our own transformer architecture now, without the help of pre-trained models. We will base our architecture off of BERT, which will also allow us to effectively compare the performance of our own implementation with the pre-trained model. The high-level of the model architecture can be seen below.\n",
    "\n",
    "Note that, compared to illustration below, all the tensor sizes will most likely be different and specific to the hyperparameters of our model.\n",
    "\n",
    "![BERT For Sequence Classification](bert_classification.png)\n",
    "\n",
    "![BERT Encoder](bert_encoder.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first define some hyperparameters of our network. (You can also play around with these later to improve your model's performance.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQ_LEN = 128 # Maximum length of tokenized input sequence (we already covered this above!)\n",
    "VOCAB_SIZE = 30522 # number of different token -> ID mappings, specific to the BERT Tokenizer (do not change)\n",
    "N_LAYERS = 12 # number of transformer layers stacked on top of each other\n",
    "N_HEADS = 12 # number of heads of the transformer, input sequence will be split and equally distributed to each head\n",
    "EMB_SIZE = 768 # refers to the dimensionality of the vector representations used to encode input tokens\n",
    "INTERMEDIATE_SIZE = EMB_SIZE * 4 # dimensionality of output of the Intermediate layer\n",
    "DROPOUT = 0.1 # probability of dropping a neuron during training\n",
    "N_CLASSES = 2 # we are classifying the sentiment of each tweet simply as POSITIVE or NEGATIVE\n",
    "LAYER_NORM_EPS = 1e-12 # value used in Layer Norm for numerical stability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformers\n",
    "\n",
    "Transformers are a type of neural network architecture primarily used in the field of natural language processing (NLP). Unlike previous models that processed inputs sequentially (like RNNs), transformers process whole sequences of data in parallel, which significantly speeds up training.\n",
    "\n",
    "### Why are they important?\n",
    "\n",
    "They have been highly successful in a variety of NLP tasks like translation, text summarization, and sentiment analysis due to their ability to handle long-range dependencies in text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Components of Transformers\n",
    "\n",
    "1. Input Embeddings + Positional Encoding\n",
    "2. Multi-Head Attention\n",
    "3. Positional-Wise Feed-Forward Networks\n",
    "4. Layer Normalization and Residual Connections\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Input Embeddings\n",
    "\n",
    "Converts tokens (words) into vectors of continuous numbers that the model can understand.  \n",
    "\n",
    "*(Optional)* Check out 3Blue1Brown's video on Word Embeddings [here](https://www.youtube.com/watch?v=wjZofJX0v4M)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Concept of Embeddings\n",
    "\n",
    "Embeddings are vector representations of text, where words with **similar meanings** have **similar representations**. These vectors are learned and adjusted during the training of the model to optimize performance on a specific task."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Importance of Embeddings\n",
    "\n",
    "Embeddings capture the semantic properties of words, which means they convert the words/tokens into a form that a neural network can work with. This allows the model to **understand and process language** by **quantifying the similarities between different words or phrases**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embeddings in Transformers\n",
    "\n",
    "1. **Tokenization**:\n",
    "   - **Process**: Text input is split into tokens (words or subwords), which are then converted into numerical IDs that correspond to entries in the embedding table.\n",
    "   - **Role in Model**: This step is critical because it translates human-readable text into a format that can be mathematically manipulated by the model.\n",
    "\n",
    "2. **Lookup Table**:\n",
    "   - **Process**: Each token ID is used to retrieve its corresponding embedding vector from an embedding matrix (or table) that the model learns during training.\n",
    "   - **Role in Model**: This embedding matrix acts as the foundational input layer of the transformer, providing a dense and informative representation for each token.\n",
    "\n",
    "3. **Dimensionality**:\n",
    "   - **Process**: Embeddings are usually vectors of fixed size (e.g., 256, 512 dimensions), regardless of the vocabulary size.\n",
    "   - **Role in Model**: Higher dimensions generally capture more detailed semantic information about each token, but also increase model complexity and computational cost."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Positional Embeddings\n",
    "\n",
    "If we only turn words into embeddings, the model won't know the order of them in a sentence. We need to somehow preserve the positional info.\n",
    "\n",
    "- **Concept**: Since transformers process words in parallel rather than sequentially, positional encodings are added to embeddings to give the model information about the order of words in a sentence.\n",
    "- **Why Necessary**: Helps the model understand how the position of a word in a sentence affects its meaning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training and Adjustments\n",
    "\n",
    "**Learning Process**:\n",
    "- Embeddings are not static; they are adjusted during the training process to minimize the model's prediction error. This learning happens through backpropagation, which adjusts the embeddings to better capture the relationships and contexts in which words appear.\n",
    "\n",
    "**Impact on Performance**:\n",
    "- The quality of embeddings significantly affects the model's performance. Better embeddings lead to a better understanding of language nuances, such as synonyms, antonyms, and different contexts."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Toy Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Vocabulary size and embedding dimension\n",
    "vocab_size = 1000  # Number of unique tokens\n",
    "embedding_dim = 64  # Size of each embedding vector\n",
    "\n",
    "# Create an embedding layer\n",
    "embedding_layer = nn.Embedding(num_embeddings=vocab_size, embedding_dim=embedding_dim)\n",
    "\n",
    "# Example token IDs\n",
    "token_ids = torch.LongTensor([10, 200, 500])\n",
    "\n",
    "# Get embeddings for token IDs\n",
    "embeddings = embedding_layer(token_ids)\n",
    "print(embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example initializes an embedding layer with random values which then can be trained to learn meaningful representations through a task like sentiment analysis or machine translation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn\n",
    "\n",
    "Now, let's implement the `EmbeddingLayer` class. This class will be responsible for converting tokens into dense vectors that the model can understand. We will use PyTorch's `nn.Embedding` module to create the embedding layer.\n",
    "\n",
    "1. token embedding: vocab -> embedding.\n",
    "2. positional embedding: position -> embedding\n",
    "3. add them together\n",
    "4. layer normalization\n",
    "5. dropout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, emb_size, max_seq_length, layer_norm_eps=1e-12, dropout=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        # TODO: Initialize word embeddings\n",
    "        self.word_embeddings = None\n",
    "        self.position_embeddings = None\n",
    "        self.ln = None\n",
    "        self.dropout = None\n",
    "\n",
    "        # sets self.position_ids to a tensor containing values from 0 to max_seq_length, reshaped to (1, max_seq_length)\n",
    "        # register_buffer makes it so that this does NOT get updated during training\n",
    "        self.register_buffer(\"position_ids\", torch.arange(max_seq_length).expand((1, -1)))\n",
    "\n",
    "    def forward(self, input_ids):\n",
    "        pass\n",
    "\n",
    "        return emb\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Multi-Head Attention\n",
    "\n",
    "Multi-head attention consists of several attention mechanisms (heads) running in parallel. Each head computes an attention score known as \"Scaled Dot-Product Attention,\" which determines how much each element in a sequence should pay attention to every other element. Combining multiple heads allows the model to capture various aspects of semantic relationships in different subspaces.  \n",
    "\n",
    "*(Optional)* Check out 3Blue1Brown's video on Attention Mechanism [here](https://www.youtube.com//watch?v=eMlx5fFNoYc)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why Multi-Head Attention\n",
    "\n",
    "By using multiple attention heads, transformers can:\n",
    "- Capture a richer diversity of relationships within the data.\n",
    "- Focus on different parts of the sentence simultaneously.\n",
    "- Improve the ability of the model to focus on relevant parts of the input for making predictions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Components of Multi-Head Attention\n",
    "\n",
    "1. **Scaled Dot-Product Attention**: The basic attention mechanism computes the dot products of the query with all keys, divides each by the square root of the dimensionality (to stabilize gradients), applies a softmax function to obtain weights on the values, and finally returns an output vector.\n",
    "\n",
    "2. **Multiple Heads**: Each head performs attention independently, allowing the model to focus on different features across different positions in the input sequence.\n",
    "\n",
    "3. **Concatenation and Final Linear Transformation**: Outputs from individual heads are concatenated and passed through a final linear layer to combine the different learned aspects."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn\n",
    "\n",
    "We compute attention using the formula below, introduced in [this](https://arxiv.org/abs/1706.03762) famous paper.\n",
    "\n",
    "![Attention](attention.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, emb_size, n_heads, dropout=0.1, layer_norm_eps=1e-12):\n",
    "        '''\n",
    "        Args:\n",
    "            emb_size (int): The size of the input embeddings\n",
    "            n_heads (int): The number of attention heads\n",
    "            dropout (float): The dropout probability\n",
    "            layer_norm_eps (float): The epsilon value for layer normalization\n",
    "        \n",
    "        '''\n",
    "        super().__init__()\n",
    "        assert emb_size % n_heads == 0, \"emb_size must be divisible by n_heads\"\n",
    "\n",
    "        # TODO: initialize head size, query, key, value, dropout, final_linear, and layer norm\n",
    "        # you want to use nn.Linear, nn.Dropout, nn.LayerNorm\n",
    "        # query, key, value, final_linear should map from emb_size to emb_size\n",
    "        self.head_size = 0\n",
    "        self.query = None \n",
    "        self.key = None\n",
    "        self.value = None\n",
    "        self.dropout = None\n",
    "\n",
    "        self.final_linear = None\n",
    "        self.ln = None\n",
    "\n",
    "\n",
    "    def forward(self, emb, att_mask):\n",
    "        '''\n",
    "        Args:\n",
    "            emb (torch.Tensor): The input embeddings with shape (B, T, C)\n",
    "            att_mask (torch.Tensor): The attention mask with shape (B, T)\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The output embeddings with shape (B, T, C)\n",
    "\n",
    "        1. You want to obtain the query (Q), key (K), and value (V) tensors from the input embeddings\n",
    "        2. Compute weights: \n",
    "            a. find dot product of Q and K \n",
    "            b. divide by the square root of the head size\n",
    "            c. apply the attention mask\n",
    "            d. apply softmax\n",
    "        3. Compute attention: multiply weights by V\n",
    "        4. Feed the attention through the final linear layer\n",
    "        5. Apply dropout, residual connection, and layer normalization\n",
    "\n",
    "        '''\n",
    "        B, T, C = emb.shape  # batch size, sequence length, embedding size   \n",
    "    \n",
    "        # TODO: 1. get the query, key, and value tensors\n",
    "        # TODO: 2. rearrange the tensors from B x T x C to B x n_heads x T x head_size\n",
    "\n",
    "        q = \n",
    "        k = \n",
    "        v = \n",
    "        \n",
    "        # TODO: compute the weights before attention mask\n",
    "        weights =\n",
    "\n",
    "        # set the pad tokens to -inf so that they equal zero after softmax\n",
    "        if att_mask != None:\n",
    "            att_mask = (att_mask > 0).unsqueeze(1).repeat(1, att_mask.size(1), 1).unsqueeze(1)\n",
    "            weights = weights.masked_fill(att_mask == 0, float('-inf'))\n",
    "\n",
    "        # apply softmax and dropout\n",
    "        weights = \n",
    "        \n",
    "        # compute attention\n",
    "        attention = \n",
    "\n",
    "        # rearrange the attention tensor from B x n_heads x T x head_size to B x T x C\n",
    "\n",
    "        # apply final linear layer, dropout, residual connection, and layer normalization\n",
    "        out = \n",
    "        \n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Side Note: Layer Normalization and Residual Connections\n",
    "\n",
    "- **Concept**: Techniques used within transformers to stabilize the learning process. Residual connections help in propagating gradients through deep networks without vanishing.\n",
    "- **Why Necessary**: Improves training efficiency and model performance by preventing the vanishing gradient problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Position-Wise Feed-Forward Networks\n",
    "\n",
    "The position-wise feed-forward network is a standard fully connected neural network that applies the same two linear transformations to each position independently and identically. It consists of two layers:\n",
    "- The first linear transformation maps the input dimension `emb_size` to a dimension `intermediate_size` (typically much larger, e.g., 2048).\n",
    "- The second linear transformation maps it back from `intermediate_size` to `emb_size`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why Position-wise Feed-Forward Networks\n",
    "\n",
    "This component of the transformer allows the model to consider each position separately and integrate non-linearity into the model. Each position in the encoder or decoder can be thought of as having its own little neural network that processes the information.\n",
    "\n",
    "Its role is two-fold:\n",
    "1. It increases the model's capacity by introducing additional learnable parameters.\n",
    "2. It facilitates the progressive enrichment and refinement of token representations as they pass through each encoder layer."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionWiseFeedForward(nn.Module):\n",
    "\n",
    "    def __init__(self, emb_size, intermediate_size, dropout=0.1, layer_norm_eps=1e-12):\n",
    "        '''\n",
    "        Args:\n",
    "            emb_size (int): The size of the input embeddings\n",
    "            intermediate_size (int): The size of the intermediate layer\n",
    "            dropout (float): The dropout probability\n",
    "            layer_norm_eps (float): The epsilon value for layer normalization\n",
    "\n",
    "        We want to create two fully connected layers mapping from emb_size to intermediate_size and back\n",
    "        '''\n",
    "\n",
    "        # TODO: initialize the two fully connected layers, gelu, dropout, and layer norm\n",
    "        super().__init__()\n",
    "        self.fc1 = None\n",
    "        self.fc2 = None\n",
    "\n",
    "        # we want some non-linearity between the two layers, we are using GELU, feel free to experiment with other activations\n",
    "        self.gelu = None\n",
    "        self.dropout = None\n",
    "        self.ln = None\n",
    "\n",
    "    def forward(self, att_out):\n",
    "        '''\n",
    "        Args:\n",
    "            att_out (torch.Tensor): The output embeddings from the multi-head attention layer with shape (B, T, C)\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The output embeddings with shape (B, T, C)\n",
    "        \n",
    "        TODO:\n",
    "        1. Apply the first fully connected layer and non-linearity\n",
    "        2. Apply the second fully connected layer and dropout\n",
    "        3. Apply a residual connection and layer normalization\n",
    "        '''\n",
    "        \n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer Block: Putting It All Together\n",
    "\n",
    "Now that we have implemented the individual components of the transformer, we can combine them to create a single transformer block. This block will contain the following components:\n",
    "- Multi-Head Attention\n",
    "- Position-wise Feed-Forward Networks\n",
    "- Layer Normalization Residual Connections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, emb_size, n_heads, intermediate_size, dropout=0.1, layer_norm_eps=1e-12):\n",
    "        '''\n",
    "        Args:\n",
    "            emb_size (int): The size of the input embeddings\n",
    "            n_heads (int): The number of attention heads\n",
    "            intermediate_size (int): The size of the intermediate layer\n",
    "            dropout (float): The dropout probability\n",
    "            layer_norm_eps (float): The epsilon value for layer normalization\n",
    "\n",
    "        We want to create a transformer block that consists of a multi-head attention layer and a position-wise feed forward layer.\n",
    "        Note: We took care of layer normalization and residual connections in the previous MultiHeadAttention and PositionWiseFeedForward classes\n",
    "        '''\n",
    "\n",
    "        # TODO: initialization\n",
    "        super().__init__()\n",
    "        self.attn = None\n",
    "        self.ff = None\n",
    "\n",
    "    def forward(self, emb, att_mask):\n",
    "        '''\n",
    "        Args:\n",
    "            emb (torch.Tensor): The input embeddings with shape (B, T, C)\n",
    "            att_mask (torch.Tensor): The attention mask with shape (B, T)\n",
    "        \n",
    "        Returns:\n",
    "            torch.Tensor: The output embeddings with shape (B, T, C)\n",
    "\n",
    "        Take emb and att_mask and pass them through the multi-head attention layer and position-wise feed forward layer\n",
    "        '''\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Your Own BERT\n",
    "\n",
    "Now, you have everything you need to create your own (simplified) BERT* model. This model will consist of the following components:\n",
    "\n",
    "1. Embedding Layer\n",
    "2. Encoder Block (multiple transformer blocks stacked together)\n",
    "3. Pooler (to aggregate information from the encoder)\n",
    "4. Classifier (to predict the sentiment of the input text)\n",
    "\n",
    "***Note**: we are not exactly replicating BERT, but rather creating a simplified version of it for educational purposes."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder Block\n",
    "\n",
    "The power of BERT's encoder block comes from stacking multiple encoder layers together. Each layer outputs a higher-level representation of the input tokens. With each successive layer, the model can capture more complex dependencies and refine its understanding of the input context.\n",
    "\n",
    "The **output** of the **last encoder** layer serves as a **comprehensive representation of the input sequence**, embedding both the immediate meaning and broader context of each token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "\n",
    "    def __init__(self, emb_size, n_layers, n_heads, intermediate_size, dropout=0.1, layer_norm_eps=1e-12):\n",
    "        '''\n",
    "        Args:\n",
    "            emb_size (int): The size of the input embeddings\n",
    "            n_layers (int): The number of transformer blocks\n",
    "            n_heads (int): The number of attention heads\n",
    "            intermediate_size (int): The size of the intermediate layer\n",
    "            dropout (float): The dropout probability\n",
    "            layer_norm_eps (float): The epsilon value for layer normalization\n",
    "\n",
    "        We want to create n_layers of transformer blocks\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        # TODO: initialize transformer layers, (hint: checkout nn.ModuleList)\n",
    "        self.transformer_layers = None\n",
    "\n",
    "    def forward(self, emb, att_mask):\n",
    "        '''\n",
    "        Args:\n",
    "            emb (torch.Tensor): The input embeddings with shape (B, T, C)\n",
    "            att_mask (torch.Tensor): The attention mask with shape (B, T)\n",
    "\n",
    "        Returns:\n",
    "            torch.Tensor: The output embeddings with shape (B, T, C)\n",
    "\n",
    "        Pass the input embeddings through each transformer block\n",
    "        '''\n",
    "        # TODO:\n",
    "\n",
    "        return emb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Hang on tight, we're almost there! We have successfully implemented the BERT Encoder architecture, but now we have to tailor the model to a clasification task like sentiment analysis.\n",
    "\n",
    "The output from the encoder is a matrix where each row represents an enriched vector representation for the corresponding input token, capturing contextual and positional information through the self-attention and feed-forward layers. However, for sequence classification tasks, a single vector representation summarizing the meaning of the entire input sequence is required, rather than individual token representations. To achieve this, we will only process the representation of the `[CLS]` token (always inserted at the start of our input sequence), as it is expected to contain a contextualized, high-level representation of the entire sequence and therefore is a good candidate to be pooled for sequence classification."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pooler\n",
    "\n",
    "The pooler component in BERT takes the output of the final encoder layer and transforms it to create a fixed-size output vector. Typically, it processes the output corresponding to the [CLS] token, which is the first token in the sequence. This token is specially designed to hold the aggregate representation of the input sequence after passing through the transformer blocks.\n",
    "\n",
    "### Why Pooler\n",
    "\n",
    "The main purpose of the pooler is to condense the complex and variable-length output of the encoder into a fixed-size embedding that can be used efficiently by downstream tasks, such as classification layers. For tasks like sentence classification, the model needs a singular representation of the entire input, which the pooler provides by transforming the [CLS] token's final hidden state.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn\n",
    "\n",
    "Pooler involves 3 steps:\n",
    "1. Extract the hidden state corresponding to the [CLS] token.\n",
    "2. Apply a linear transformation to this hidden state: embedding_size -> embedding_size.\n",
    "3. Apply a non-linear activation function (like Tanh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Pooler(nn.Module):\n",
    "    def __init__(self, emb_size):\n",
    "        '''\n",
    "        Args:\n",
    "            emb_size (int): The size of the input embeddings\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        # TODO: Define the linear layer and activation function\n",
    "        self.dense = None\n",
    "        self.activation = None\n",
    "\n",
    "    def forward(self, encoder_out):\n",
    "        '''\n",
    "        Args:\n",
    "            encoder_out (torch.Tensor): The output embeddings from the encoder with shape (B, T, C)\n",
    "        \n",
    "        Returns:\n",
    "            torch.Tensor: The pooled embeddings with shape (B, C)\n",
    "        '''\n",
    "\n",
    "        # TODO: 1. extract the first token's embeddings\n",
    "        # TODO: 2. apply the linear layer and activation function\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Final Simplified BERT Classifier\n",
    "\n",
    "Now that we have the pooler component, we can add a classification layer on top of it to predict the sentiment of the input text. This classification layer will be a simple linear layer that takes the pooled output from the pooler and maps it to `n_classes` output units."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Your Turn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimplifiedBertClassifier(nn.Module):\n",
    "    def __init__(self, vocab_size, emb_size, n_layers, n_heads, intermediate_size, dropout=0.1, n_classes=2):\n",
    "        '''\n",
    "        Args:\n",
    "            vocab_size (int): The size of the vocabulary\n",
    "            emb_size (int): The size of the input embeddings\n",
    "            n_layers (int): The number of transformer blocks\n",
    "            n_heads (int): The number of attention heads\n",
    "            intermediate_size (int): The size of the intermediate layer\n",
    "            dropout (float): The dropout probability\n",
    "            n_classes (int): The number of classes\n",
    "        '''\n",
    "        super().__init__()\n",
    "\n",
    "        # TODO: Initialization\n",
    "        self.embedding = None\n",
    "        self.encoder = None\n",
    "        self.pooler = None\n",
    "        self.dropout = None\n",
    "        self.classifier = None\n",
    "\n",
    "    def forward(self, input_ids, attention_mask=None):\n",
    "        '''\n",
    "        Args:\n",
    "            input_ids (torch.Tensor): The input token IDs with shape (B, T)\n",
    "            attention_mask (torch.Tensor): The attention mask with shape (B, T)\n",
    "        \n",
    "        Returns:\n",
    "            torch.Tensor: The logits with shape (B, n_classes)\n",
    "\n",
    "        1. Pass the input token IDs through the embedding layer, encoder, pooler\n",
    "        2. Apply dropout\n",
    "        3. Pass the output through the classifier\n",
    "        '''\n",
    "\n",
    "        # TODO:\n",
    "         \n",
    "        return logits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have successfully implemented our own transformer-based classifier for sentiment analysis! Let's train it and see how it compares to the pre-trained one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "our_model = SimplifiedBertClassifier(vocab_size=VOCAB_SIZE, emb_size=EMB_SIZE, n_layers=N_LAYERS, \n",
    "                                     n_heads=N_HEADS, intermediate_size=INTERMEDIATE_SIZE, dropout=DROPOUT, n_classes=N_CLASSES)\n",
    "\n",
    "optimizer = torch.optim.AdamW(our_model.parameters(), lr=2e-5)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "loss_history, acc_history = train(our_model, train_loader, test_loader, optimizer, criterion, device, num_epochs=3, finetuning=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(loss_history)\n",
    "plt.xlabel('Iterations')\n",
    "plt.ylabel('Cross-entropy Loss')\n",
    "plt.title('BERT From Scratch, Loss history')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(acc_history)\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Test accuracy')\n",
    "plt.title('BERT From Scratch, Accuracy history')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace the string below by your own.\n",
    "sample_tweet = 'My Doordash delivery is 30 minutes late:( Ughh, I am hungry!!!'\n",
    "\n",
    "predicted_label = classify_tweet(sample_tweet, our_model, device, finetuning=False)\n",
    "classification = 'positive' if predicted_label == 1 else 'negative'\n",
    "print(f'Your tweet was classified as {classification}.')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (ml)",
   "language": "python",
   "name": "ml"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}