{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "inr4HSeKoZbE"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/dbamman/anlp25/blob/main/5.classification/TransformerClassification.ipynb)\n",
    "\n",
    "**N.B.** Once it's open on Colab, remember to save a copy (by e.g. clicking `Copy to Drive` above).\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ubRzhaCVXYUy"
   },
   "source": [
    "Thie notebook explores using transformers for document classification.  Before starting, change the runtime to GPU: Runtime > Change runtime type > Hardware accelerator: GPU (any GPU is fine).\n",
    "\n",
    "For an intro to models in PyTorch, see [this tutorial](https://pytorch.org/tutorials/beginner/introyt/modelsyt_tutorial.html).\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "eQnqbL-NjmiP"
   },
   "source": [
    "Download classification data for training/evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "X9UUcu7TG6sg",
    "outputId": "523d881c-a8c2-46d6-a2ff-4597d79aec76"
   },
   "outputs": [],
   "source": [
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/convote/train.tsv\n",
    "!wget https://raw.githubusercontent.com/dbamman/anlp25/main/data/convote/dev.tsv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PwsQCWgmHLrP"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import sys\n",
    "import torch\n",
    "from torch import nn\n",
    "from collections import Counter\n",
    "from nltk import word_tokenize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eS2rympQIObz",
    "outputId": "ef9e29b4-90ed-4c43-a34e-8a565caa786b"
   },
   "outputs": [],
   "source": [
    "import nltk\n",
    "nltk.download('punkt')\n",
    "nltk.download('punkt_tab')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QDlR2HlLHO7J"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "l3zDewVZHW23",
    "outputId": "3d3fab99-d119-45be-d26e-3e6319c07528"
   },
   "outputs": [],
   "source": [
    "# max sequence length\n",
    "max_length=256\n",
    "\n",
    "# limit vocabulary to top N words in training data\n",
    "max_vocab=10000\n",
    "\n",
    "# batch size\n",
    "batch_size=128\n",
    "\n",
    "# size of token representations (which dictates the size of the overall model).\n",
    "d_model=16\n",
    "\n",
    "\n",
    "# number of epochs\n",
    "num_epochs=50\n",
    "\n",
    "print('')\n",
    "print(\"********************************************\")\n",
    "print(\"Running on: {}\".format(device))\n",
    "print(\"********************************************\")\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "h5IJ2unzHYzu"
   },
   "outputs": [],
   "source": [
    "# PositionalEncoding class copied from:\n",
    "# https://github.com/pytorch/examples/blob/main/word_language_model/model.py\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, dropout=0.1):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "        pe = torch.zeros(max_length, d_model)\n",
    "        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0)#.transpose(0, 1)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        x = x + self.pe[:x.size(0), :]\n",
    "        return self.dropout(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pMTmsPDjHast"
   },
   "outputs": [],
   "source": [
    "class TransformerClassifier(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, num_labels, d_model, nhead=2, num_encoder_layers=1, dim_feedforward=256):\n",
    "\n",
    "        super(TransformerClassifier, self).__init__()\n",
    "\n",
    "        self.num_labels = num_labels\n",
    "        self.embedding = nn.Embedding(num_embeddings=max_vocab+2, embedding_dim=d_model)\n",
    "        self.transformer = nn.Transformer(\n",
    "            d_model=d_model,\n",
    "            nhead=nhead,\n",
    "            num_encoder_layers=num_encoder_layers,\n",
    "            dim_feedforward=dim_feedforward,\n",
    "            batch_first=True,\n",
    "            dropout=0.5\n",
    "        )\n",
    "        self.classifier = nn.Linear(d_model, self.num_labels)\n",
    "        self.pos_encoder = PositionalEncoding(d_model)\n",
    "\n",
    "    def forward(self, x, m):\n",
    "\n",
    "        # put data on device (e.g., gpu)\n",
    "        x = x.to(device)\n",
    "        m = m.to(device)\n",
    "\n",
    "        # convert input token IDs to word embeddings\n",
    "        embed = self.embedding(x)\n",
    "\n",
    "        # add position encodings to include information about word position within the document\n",
    "        embed = self.pos_encoder(embed)\n",
    "\n",
    "        # get transformer output\n",
    "        h = self.transformer.encoder(embed, src_key_padding_mask=m)\n",
    "\n",
    "        # Represent document as average embedding of transformer output\n",
    "        h = torch.mean(h, dim=1)\n",
    "\n",
    "        # Convert document representation into output label space\n",
    "        logits = self.classifier(h)\n",
    "\n",
    "        return logits\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TT56POLcHcLM"
   },
   "outputs": [],
   "source": [
    "def create_vocab_and_labels(filename, max_vocab):\n",
    "    # This function creates the word vocabulary (and label ids) from the training data\n",
    "    # The vocab is a mapping between word types and unique word IDs\n",
    "\n",
    "    counts = Counter()\n",
    "    labels = {}\n",
    "    with open(filename, encoding=\"utf-8\") as file:\n",
    "        for line in file:\n",
    "            cols = line.rstrip().split(\"\\t\")\n",
    "            lab = cols[0]\n",
    "            text = word_tokenize(cols[1].lower())\n",
    "            for tok in text:\n",
    "                counts[tok] += 1\n",
    "\n",
    "            if lab not in labels:\n",
    "                labels[lab] = len(labels)\n",
    "\n",
    "    vocab = {\"[MASK]\":0, \"[UNK]\":1}\n",
    "\n",
    "    for k,v in counts.most_common(max_vocab):\n",
    "        vocab[k] = len(vocab)\n",
    "\n",
    "    return vocab, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tB9Vv3TiHdkW"
   },
   "outputs": [],
   "source": [
    "def read_data(filename, vocab, labels, max_length, max_docs=5000):\n",
    "    # Read in data from file, up to the first max_docs documents. For each document\n",
    "    # read up to max_length tokens.\n",
    "\n",
    "    x = []\n",
    "    y = []\n",
    "    m = []\n",
    "\n",
    "    with open(filename, encoding=\"utf-8\") as file:\n",
    "        for idx, line in enumerate(file):\n",
    "            if idx >= max_docs:\n",
    "                break\n",
    "            cols = line.rstrip().split(\"\\t\")\n",
    "            lab = cols[0]\n",
    "            text = word_tokenize(cols[1])\n",
    "            text_ids = []\n",
    "            for tok in text:\n",
    "                if tok in vocab:\n",
    "                    text_ids.append(vocab[tok])\n",
    "                else:\n",
    "                    text_ids.append(vocab[\"[UNK]\"])\n",
    "\n",
    "            text_ids = text_ids[:max_length]\n",
    "\n",
    "            # PyTorch (and most libraries that deal with matrix operations) expects all inputs to be the same length\n",
    "            # So pad each document with 0s up to max_length\n",
    "            # But keep track of the true number of tokens in the document with the \"mask\" list.\n",
    "\n",
    "            # True tokens have a mask value of 0\n",
    "            mask = [0] * len(text_ids)\n",
    "\n",
    "            for i in range(len(text_ids), max_length):\n",
    "                text_ids.append(vocab[\"[MASK]\"])\n",
    "                # Padded tokens have a mask value of 1\n",
    "                mask.append(1)\n",
    "\n",
    "            x.append(text_ids)\n",
    "            m.append(mask)\n",
    "            y.append(labels[lab])\n",
    "\n",
    "    return x, y, m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dglYDfx-HfEt"
   },
   "outputs": [],
   "source": [
    "def get_batches(x, y, m, batch_size):\n",
    "\n",
    "    # Create minibatches from the full dataset\n",
    "\n",
    "    batches_x = []\n",
    "    batches_y = []\n",
    "    batches_m = []\n",
    "    for i in range(0, len(x), batch_size):\n",
    "        xbatch = x[i:i+batch_size]\n",
    "        ybatch = y[i:i+batch_size]\n",
    "        mbatch = m[i:i+batch_size]\n",
    "\n",
    "        batches_x.append(torch.LongTensor(xbatch))\n",
    "        batches_y.append(torch.LongTensor(ybatch))\n",
    "        batches_m.append(torch.BoolTensor(mbatch))\n",
    "\n",
    "    return batches_x, batches_y, batches_m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_g5m_NjzHgsW"
   },
   "outputs": [],
   "source": [
    "def evaluate(model, all_x, all_y, all_m):\n",
    "\n",
    "    # Calculate accuracy\n",
    "\n",
    "    model.eval()\n",
    "    corr = 0.\n",
    "    total = 0.\n",
    "    with torch.no_grad():\n",
    "        for x, y, m in zip(all_x, all_y, all_m):\n",
    "            y_preds = model.forward(x, m)\n",
    "            for idx, y_pred in enumerate(y_preds):\n",
    "                prediction = torch.argmax(y_pred)\n",
    "                if prediction == y[idx]:\n",
    "                    corr += 1.\n",
    "                total += 1\n",
    "    return corr / total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eLjdIDl4HiDE"
   },
   "outputs": [],
   "source": [
    "def train(model, model_filename, train_x, train_y, train_m, dev_x, dev_y, dev_m, shuffle=False):\n",
    "\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "    cross_entropy = nn.CrossEntropyLoss()\n",
    "\n",
    "    # Keep track of the epoch that has the best dev accuracy\n",
    "    best_dev_acc = 0.\n",
    "    best_dev_epoch = None\n",
    "\n",
    "    losses = []\n",
    "\n",
    "    # How many epochs with no changes before we quit\n",
    "    patience = 20\n",
    "\n",
    "    dev_batches_x, dev_batches_y, dev_batches_m = get_batches(dev_x, dev_y, dev_m, batch_size=batch_size)\n",
    "\n",
    "    for epoch in range(num_epochs):\n",
    "\n",
    "        model.train()\n",
    "\n",
    "        if shuffle:\n",
    "            # shuffle batches\n",
    "            shuffled_idx = torch.randperm(len(train_x))\n",
    "            train_x = [train_x[i] for i in shuffled_idx]\n",
    "            train_y = [train_y[i] for i in shuffled_idx]\n",
    "            train_m = [train_m[i] for i in shuffled_idx]\n",
    "\n",
    "        train_batches_x, train_batches_y, train_batches_m = get_batches(train_x, train_y, train_m, batch_size=batch_size)\n",
    "\n",
    "        for x, y, m in zip(train_batches_x, train_batches_y, train_batches_m):\n",
    "            # Get predictions for batch x (with mask values m)\n",
    "            y_pred = model.forward(x, m)\n",
    "            y = y.to(device)\n",
    "\n",
    "            # Calculate loss as cross-entropy with true labels\n",
    "            loss = cross_entropy(y_pred.view(-1, model.num_labels), y.view(-1))\n",
    "\n",
    "            # Set all gradients to zero\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # Calculate gradients from current loss\n",
    "            loss.backward()\n",
    "            losses.append(loss.item())\n",
    "\n",
    "            # Update parameters\n",
    "            optimizer.step()\n",
    "\n",
    "        dev_accuracy = evaluate(model, dev_batches_x, dev_batches_y, dev_batches_m)\n",
    "        train_accuracy = evaluate(model, train_batches_x, train_batches_y, train_batches_m)\n",
    "\n",
    "        # we're going to save the model that performs the best on *dev* data\n",
    "        if dev_accuracy > best_dev_acc:\n",
    "            torch.save(model.state_dict(), model_filename)\n",
    "            print(\"%.3f is better than %.3f, saving model ...\" % (dev_accuracy, best_dev_acc))\n",
    "            best_dev_acc = dev_accuracy\n",
    "            best_dev_epoch = epoch\n",
    "\n",
    "        if epoch % 1 == 0:\n",
    "            print(\"Epoch %s, dev accuracy: %.3f, train accuracy: %.3f\" % (epoch, dev_accuracy, train_accuracy))\n",
    "\n",
    "        if epoch-best_dev_epoch > patience:\n",
    "          print(\"%s > patience (%s), stopping...\" % (epoch-best_dev_epoch, patience))\n",
    "          break\n",
    "\n",
    "    model.load_state_dict(torch.load(model_filename))\n",
    "    print(\"\\nBest Performing Model achieves dev accuracy of : %.3f\" % (best_dev_acc))\n",
    "\n",
    "    return losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_nkNh6f8HkL2"
   },
   "outputs": [],
   "source": [
    "vocab, labels = create_vocab_and_labels(\"train.tsv\", max_vocab)\n",
    "train_x, train_y, train_m = read_data(\"train.tsv\", vocab, labels, max_length=max_length)\n",
    "dev_x, dev_y, dev_m = read_data(\"dev.tsv\", vocab, labels, max_length=max_length)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0vtxV6jpHrAV",
    "outputId": "d8fe1d9a-1c9f-4e46-9c3a-63d166187f8c"
   },
   "outputs": [],
   "source": [
    "classifier = TransformerClassifier(num_labels=len(labels), d_model=100, dim_feedforward=512)\n",
    "classifier = classifier.to(device)\n",
    "\n",
    "train_x_batch, train_y_batch, train_m_batch = get_batches(train_x, train_y, train_m, batch_size=batch_size)\n",
    "dev_x_batch, dev_y_batch, dev_m_batch = get_batches(dev_x, dev_y, dev_m, batch_size=batch_size)\n",
    "\n",
    "losses = train(classifier, \"test.model\", train_x, train_y, train_m, dev_x, dev_y, dev_m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "R3dn2o1KKdz_"
   },
   "source": [
    "**Q1**. Play around with this transformer as implemented and experiment with how performance on the dev data changes as a function of `d_model`, `num_encoder_layers`, `nhead`, etc.).  Describe your experiments and report dev accuracy on them below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XiWT6yumKD_l"
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WWESAVi4TIF0"
   },
   "source": [
    "**Q2**. Here is some code that plots the loss curve over training. Plot the loss curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "iXSk9gmDTIF0"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "def plot_loss_curve(losses):\n",
    "    sns.lineplot(x=range(len(losses)), y=losses)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 430
    },
    "id": "d-WSPY32TIF0",
    "outputId": "0b02d2cf-6abb-4520-f958-88bd290c311a",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plot_loss_curve(losses)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7jKV1hIqTIF1"
   },
   "source": [
    "Do you notice any regularities in the loss? Instantiate a new classifier, but this time train with `shuffle=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OkxajU3cTIF1",
    "outputId": "3350c7e6-9194-42a0-c519-ac43c481bb1b"
   },
   "outputs": [],
   "source": [
    "new_classifier = TransformerClassifier(num_labels=len(labels), d_model=100, dim_feedforward=512)\n",
    "new_classifier = new_classifier.to(device)\n",
    "\n",
    "shuffled_losses = train(new_classifier, \"test.model\", train_x, train_y, train_m, dev_x, dev_y, dev_m, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oAuxP-GvTIF1"
   },
   "source": [
    "Plot the loss curve of this training run. How do the loss curves compare? **Consider:** why is it generally desirable to shuffle the training data at the beginning of each epoch when performing gradient descent?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 430
    },
    "id": "Q1_VEJ9dTIF1",
    "outputId": "1565ce43-32f9-4636-db51-710a1ee01a99"
   },
   "outputs": [],
   "source": [
    "plot_loss_curve(shuffled_losses)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}