bert_mlp/BERT_MLP_AGNEWS.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "!pip freeze | grep 'torch' -E\n",
        "!pip install datasets"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0-PxtA0GmXQ7",
        "outputId": "ff58decf-0a5d-4d38-8b30-4eb991f7d096"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch @ https://download.pytorch.org/whl/cu121/torch-2.3.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=0a12aa9aa6bc442dff8823ac8b48d991fd0771562eaa38593f9c8196d65f7007\n",
            "torchaudio @ https://download.pytorch.org/whl/cu121/torchaudio-2.3.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=38b49393f8c322dcaa29d19e5acbf5a0b1978cf1b719445ab670f1fb486e3aa6\n",
            "torchsummary==1.5.1\n",
            "torchtext==0.18.0\n",
            "torchvision @ https://download.pytorch.org/whl/cu121/torchvision-0.18.0%2Bcu121-cp310-cp310-linux_x86_64.whl#sha256=13e1b48dc5ce41ccb8100ab3dd26fdf31d8f1e904ecf2865ac524493013d0df5\n",
            "Requirement already satisfied: datasets in /usr/local/lib/python3.10/dist-packages (2.20.0)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.15.4)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.25.2)\n",
            "Requirement already satisfied: pyarrow>=15.0.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (16.1.0)\n",
            "Requirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets) (0.6)\n",
            "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.3.8)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (2.0.3)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.4)\n",
            "Requirement already satisfied: xxhash in /usr/local/lib/python3.10/dist-packages (from datasets) (3.4.1)\n",
            "Requirement already satisfied: multiprocess in /usr/local/lib/python3.10/dist-packages (from datasets) (0.70.16)\n",
            "Requirement already satisfied: fsspec[http]<=2024.5.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets) (2023.6.0)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.9.5)\n",
            "Requirement already satisfied: huggingface-hub>=0.21.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.23.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (24.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.1)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (23.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.21.2->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.7)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2024.6.2)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2023.4)\n",
            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from transformers import BertTokenizerFast, BertModel\n",
        "import datasets as hfdata\n",
        "import pathlib as Path\n",
        "import re\n",
        "import time\n",
        "from tqdm import tqdm"
      ],
      "metadata": {
        "id": "osuIfFI6PaMf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Prep Data - `AGNEWS`"
      ],
      "metadata": {
        "id": "sy4YUvI6Ea69"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "agnews_hf = hfdata.load_dataset('ag_news')\n",
        "agnews_hf['train'], agnews_hf['test']"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WrjiyMc15a3t",
        "outputId": "c75464dc-1435-4a4d-f758-b44d0a9d1dab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.10/dist-packages/huggingface_hub/utils/_token.py:89: UserWarning: \n",
            "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
            "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
            "You will be able to reuse this secret in all of your notebooks.\n",
            "Please note that authentication is recommended but still optional to access public models or datasets.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(Dataset({\n",
              "     features: ['text', 'label'],\n",
              "     num_rows: 120000\n",
              " }),\n",
              " Dataset({\n",
              "     features: ['text', 'label'],\n",
              "     num_rows: 7600\n",
              " }))"
            ]
          },
          "metadata": {},
          "execution_count": 3
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "agnews_hf['train']['text'][:5]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SXjT15eB5u5-",
        "outputId": "7e92811d-a53f-459e-fb60-ae4278ae4fc6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[\"Wall St. Bears Claw Back Into the Black (Reuters) Reuters - Short-sellers, Wall Street's dwindling\\\\band of ultra-cynics, are seeing green again.\",\n",
              " 'Carlyle Looks Toward Commercial Aerospace (Reuters) Reuters - Private investment firm Carlyle Group,\\\\which has a reputation for making well-timed and occasionally\\\\controversial plays in the defense industry, has quietly placed\\\\its bets on another part of the market.',\n",
              " \"Oil and Economy Cloud Stocks' Outlook (Reuters) Reuters - Soaring crude prices plus worries\\\\about the economy and the outlook for earnings are expected to\\\\hang over the stock market next week during the depth of the\\\\summer doldrums.\",\n",
              " 'Iraq Halts Oil Exports from Main Southern Pipeline (Reuters) Reuters - Authorities have halted oil export\\\\flows from the main pipeline in southern Iraq after\\\\intelligence showed a rebel militia could strike\\\\infrastructure, an oil official said on Saturday.',\n",
              " 'Oil prices soar to all-time record, posing new menace to US economy (AFP) AFP - Tearaway world oil prices, toppling records and straining wallets, present a new economic menace barely three months before the US presidential elections.']"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "s = agnews_hf['train']['text'][1]\n",
        "def get_description_without_title(strings):\n",
        "    \"\"\"\n",
        "    Sequences start with <text> (source) source - <actual_article>\n",
        "    so I only want to extract <actual_article>.\n",
        "\n",
        "    Also new lines seem to be added with a back slash \\\n",
        "    so I split on those and join with a whitespace.\n",
        "    \"\"\"\n",
        "    pattern = re.compile('\\([A-Za-z]+\\) [A-Za-z]+ - ')\n",
        "    return [' '.join(s[_.span()[-1]:].split('\\\\')) for _ in re.finditer(pattern=pattern, string=s) for s in strings]\n",
        "get_description_without_title([s])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IRtI40c96caV",
        "outputId": "fcd346be-3c17-49b3-aa52-0034315e7530"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['Private investment firm Carlyle Group, which has a reputation for making well-timed and occasionally controversial plays in the defense industry, has quietly placed its bets on another part of the market.']"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# get lists of train and test data;\n",
        "train_texts = get_description_without_title(agnews_hf['train']['text'])\n",
        "train_labels = agnews_hf['train']['label']\n",
        "test_texts = get_description_without_title(agnews_hf['test']['text'])\n",
        "test_labels = agnews_hf['test']['label']\n",
        "assert len(train_texts) - sum(bool(_) for _ in train_texts) == 0\n",
        "assert len(test_texts) - sum(bool(_) for _ in test_texts) == 0"
      ],
      "metadata": {
        "id": "J7gjnlOU9xyJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "max(len(_.split(' ')) for _ in train_texts), max(len(_.split(' ')) for _ in test_texts)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WkRrFcxmGw1g",
        "outputId": "e78c66a1-352b-49fc-d6c0-4a7a528e2293"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(185, 141)"
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class MyDataset(Dataset):\n",
        "    def __init__(self, x, y: list[int]):\n",
        "        self.x = x\n",
        "        self.y = y\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.y)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        \"\"\"Get (text, label) tuple.\"\"\"\n",
        "        return self.x[idx], self.y[idx]"
      ],
      "metadata": {
        "id": "ih_2F-AgFDfq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset = MyDataset(train_texts, train_labels)\n",
        "test_dataset = MyDataset(test_texts, test_labels)\n",
        "\n",
        "# get loaders;\n",
        "def get_loaders(train_dataset, test_dataset, batch_size, num_workers=1):\n",
        "    train_loader = DataLoader(train_dataset, shuffle=True, drop_last=True, batch_size=batch_size, num_workers=num_workers)\n",
        "    test_loader = DataLoader(test_dataset, shuffle=False, drop_last=False, batch_size=batch_size, num_workers=num_workers)\n",
        "    return train_loader, test_loader"
      ],
      "metadata": {
        "id": "Du1xY2zAFiaS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup tokenizer and model;"
      ],
      "metadata": {
        "id": "Ryf-Pjh7Gc5O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "tokenizer = BertTokenizerFast.from_pretrained('google-bert/bert-base-uncased')"
      ],
      "metadata": {
        "id": "0SDp-j5vQD6u"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# default pos embed seems to be 512, but I'll truncate to 256;\n",
        "# token id 101 is [CLS], 102 is [SEP] and 0 is [PAD].\n",
        "sentences = ['hello there', 'Hello there', 'Hi Kenobi']  # check if hello and Hello get the same tokens;\n",
        "tokenizer(sentences, padding='longest', return_tensors='pt', truncation='longest_first', max_length=256)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cub1xjrGSJkX",
        "outputId": "b7790c6f-0ab5-429d-9bef-3f9e9b6c1d11"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'input_ids': tensor([[  101,  7592,  2045,   102,     0,     0],\n",
              "        [  101,  7592,  2045,   102,     0,     0],\n",
              "        [  101,  7632,  6358, 16429,  2072,   102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0],\n",
              "        [0, 0, 0, 0, 0, 0],\n",
              "        [0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 0, 0],\n",
              "        [1, 1, 1, 1, 0, 0],\n",
              "        [1, 1, 1, 1, 1, 1]])}"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Load model;"
      ],
      "metadata": {
        "id": "11Ju9LCgHIis"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "bert = BertModel.from_pretrained('google-bert/bert-base-uncased')\n",
        "bert"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "z3wHYZh_SVG0",
        "outputId": "03385dbe-b693-41cf-8cf6-791b24cfa843"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "BertModel(\n",
              "  (embeddings): BertEmbeddings(\n",
              "    (word_embeddings): Embedding(30522, 768, padding_idx=0)\n",
              "    (position_embeddings): Embedding(512, 768)\n",
              "    (token_type_embeddings): Embedding(2, 768)\n",
              "    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "    (dropout): Dropout(p=0.1, inplace=False)\n",
              "  )\n",
              "  (encoder): BertEncoder(\n",
              "    (layer): ModuleList(\n",
              "      (0-11): 12 x BertLayer(\n",
              "        (attention): BertAttention(\n",
              "          (self): BertSdpaSelfAttention(\n",
              "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "          (output): BertSelfOutput(\n",
              "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "            (dropout): Dropout(p=0.1, inplace=False)\n",
              "          )\n",
              "        )\n",
              "        (intermediate): BertIntermediate(\n",
              "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
              "          (intermediate_act_fn): GELUActivation()\n",
              "        )\n",
              "        (output): BertOutput(\n",
              "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
              "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
              "          (dropout): Dropout(p=0.1, inplace=False)\n",
              "        )\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (pooler): BertPooler(\n",
              "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
              "    (activation): Tanh()\n",
              "  )\n",
              ")"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Freeze bert weights;\n",
        "bert.requires_grad_(False)\n",
        "assert not any(p.requires_grad for p in bert.parameters())"
      ],
      "metadata": {
        "id": "NhDSbpPRJT0p"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup model architecture;"
      ],
      "metadata": {
        "id": "5qYW2TLFTWqa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MLP(nn.Module):\n",
        "    def __init__(self, in_dim, out_dim):\n",
        "        super().__init__()\n",
        "        # add first hidden layer;\n",
        "        self.net = nn.Sequential(\n",
        "            nn.LayerNorm(in_dim),\n",
        "            nn.Linear(in_dim, 256),\n",
        "            nn.ReLU(),\n",
        "            nn.LayerNorm(256),\n",
        "            nn.Linear(256, 128),\n",
        "            nn.ReLU(),\n",
        "            nn.LayerNorm(128),\n",
        "            nn.Linear(128, 64),\n",
        "            nn.ReLU(),\n",
        "            nn.LayerNorm(64),\n",
        "            nn.Linear(64, out_dim)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.net(x)"
      ],
      "metadata": {
        "id": "1MLe42ReN6pe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def model_forward(x, tokenizer, bert, mlp, device, x_is_cls: bool):\n",
        "    if x_is_cls:\n",
        "        cls_embed = x.to(device)\n",
        "    else:\n",
        "        x = tokenizer(\n",
        "            x,\n",
        "            padding='longest',\n",
        "            return_tensors='pt',\n",
        "            truncation='longest_first',\n",
        "            max_length=200\n",
        "        )\n",
        "        x = x.to(device)\n",
        "        cls_embed = bert(**x)['pooler_output']\n",
        "    return mlp(cls_embed)"
      ],
      "metadata": {
        "id": "axKPrTZsOjEn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_one_epoch(train_loader, tokenizer, bert, mlp, device, optim, x_is_cls=False):\n",
        "    bert.eval()  # bert itself is not trained;\n",
        "    mlp.train()\n",
        "    loader_len = len(train_loader)\n",
        "    log_every = 1 if len(train_loader) < 50 else 50\n",
        "    tenth = max(1, loader_len // 10)\n",
        "    avg_instance_loss = 0\n",
        "    n = 0\n",
        "    for i, (x, y) in enumerate(train_loader):\n",
        "        y = y.to(device)\n",
        "        optim.zero_grad()\n",
        "        out = model_forward(\n",
        "            x, tokenizer, bert, mlp, device, x_is_cls=x_is_cls\n",
        "        )\n",
        "        loss = nn.functional.cross_entropy(out, target=y)\n",
        "\n",
        "        # track losses;\n",
        "        n += len(y)\n",
        "        avg_instance_loss = avg_instance_loss + (loss.item() - avg_instance_loss) * len(y) / n\n",
        "        loss.backward()\n",
        "        # if (i + 1) % log_every == 0:\n",
        "            # wandb.log('train_loss_step', loss.detach().item())\n",
        "        # if (i + 1) % tenth == 0:\n",
        "        #     print(f\"Progress: {(i+1) / loader_len * 100:.3f}% done\")\n",
        "        optim.step()\n",
        "    return avg_instance_loss"
      ],
      "metadata": {
        "id": "4J9WxZlHO0vb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def val_one_epoch(val_loader, tokenizer, bert, mlp, device, x_is_cls=False):\n",
        "    bert.eval()  # bert itself is not trained;\n",
        "    mlp.eval()\n",
        "    loader_len = len(val_loader)\n",
        "    log_every = 1 if len(val_loader) < 50 else 50\n",
        "    tenth = max(1, loader_len // 10)\n",
        "    avg_instance_loss = 0\n",
        "    n = 0\n",
        "    avg_acc = 0\n",
        "    for i, (x, y) in enumerate(val_loader):\n",
        "        y = y.to(device)\n",
        "        with torch.no_grad():\n",
        "            out = model_forward(\n",
        "                x, tokenizer, bert, mlp, device, x_is_cls=x_is_cls\n",
        "            )\n",
        "\n",
        "            # loss calculation;\n",
        "            loss = nn.functional.cross_entropy(out, target=y)\n",
        "            n += len(y)\n",
        "            avg_instance_loss = avg_instance_loss + (loss.item() - avg_instance_loss) * len(y) / n\n",
        "\n",
        "            # accuracy calculation\n",
        "            acc = (out.argmax(-1) == y).float().mean().item()\n",
        "            avg_acc = avg_acc + (acc - avg_acc) * len(y) / n\n",
        "\n",
        "            # if (i + 1) % log_every == 0:\n",
        "                # wandb.log('val_loss_step', loss.item())\n",
        "            # if (i + 1) % tenth == 0:\n",
        "            #     print(f\"Progress: {(i+1) / loader_len * 100:.3f}% done\")\n",
        "    # wandb.log('val_loss_epoch', avg_instance_loss)\n",
        "    return avg_instance_loss, avg_acc"
      ],
      "metadata": {
        "id": "Mn-vr9gQRtfk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_for_k_epochs(num_epochs, bert, mlp, tokenizer, optim, train_loader, val_loader, epoch_offset=0, x_is_cls=False):\n",
        "    instance_train_epoch_losses = []\n",
        "    instance_val_epoch_losses = []\n",
        "    instance_acc = []\n",
        "    for t in range(0+epoch_offset, num_epochs+epoch_offset):\n",
        "        print(f\"Epoch {t}:\\n{'-'*20}\")\n",
        "        train_loss = train_one_epoch(\n",
        "            train_loader, tokenizer, bert, mlp, device, optim, x_is_cls=x_is_cls\n",
        "        )\n",
        "\n",
        "        val_loss, val_acc = val_one_epoch(\n",
        "            val_loader, tokenizer, bert, mlp, device, x_is_cls=x_is_cls\n",
        "        )\n",
        "        print(f\"{train_loss=:.5f}\")\n",
        "        print(f\"{val_acc=:.3f}\\t{val_loss=:.5f}\")\n",
        "        instance_train_epoch_losses.append(train_loss)\n",
        "        instance_val_epoch_losses.append(val_loss)\n",
        "        instance_acc.append(val_acc)\n",
        "    return instance_train_epoch_losses, instance_val_epoch_losses, instance_acc"
      ],
      "metadata": {
        "id": "WHD2t4B_TZO6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Start Training;"
      ],
      "metadata": {
        "id": "Zl7_NuFWUZjh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "torch.manual_seed(0)\n",
        "# setup predictors;\n",
        "mlp = MLP(\n",
        "    in_dim=768,\n",
        "    out_dim=4,\n",
        ").to(device)\n",
        "optim = torch.optim.Adam(mlp.parameters(), lr=1e-3)\n",
        "bert = bert.to(device)"
      ],
      "metadata": {
        "id": "BQwxyuLcug2G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Since bert is fixed, just extract cls embeds;"
      ],
      "metadata": {
        "id": "RwdX6Xwn2c8O"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_cls_embeds(loader, tokenizer, bert, device):\n",
        "    bert.eval()\n",
        "    assert not any(p.requires_grad for p in bert.parameters())\n",
        "    cls_tokens, ys = [], []\n",
        "    loader_len = len(loader)\n",
        "    tenth = max(1, loader_len // 10)\n",
        "    with torch.no_grad():\n",
        "        for i, (x, y) in enumerate(loader):\n",
        "            x = tokenizer(\n",
        "                x,\n",
        "                padding='longest',\n",
        "                return_tensors='pt',\n",
        "                truncation='longest_first',\n",
        "                max_length=200\n",
        "            )\n",
        "            x = x.to(device)\n",
        "            cls_embed = bert(**x)['pooler_output']\n",
        "            cls_tokens.append(cls_embed.to('cpu'))\n",
        "            ys.append(y)\n",
        "            if (i + 1) % tenth == 0:\n",
        "                print(f\"Progress: {(i+1) / loader_len * 100:.2f}% done\")\n",
        "    return torch.cat(cls_tokens, 0), torch.cat(ys, 0)"
      ],
      "metadata": {
        "id": "ZxKkzmje2YiP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# setup data loaders;\n",
        "BATCH_SIZE = 64\n",
        "# train_loader, test_loader = get_loaders(\n",
        "#     train_dataset, test_dataset,\n",
        "#     batch_size=BATCH_SIZE, num_workers=1\n",
        "# )\n",
        "train_loader = DataLoader(\n",
        "    train_dataset, shuffle=False, batch_size=BATCH_SIZE, num_workers=1\n",
        ")\n",
        "test_loader = DataLoader(\n",
        "    test_dataset, shuffle=False, batch_size=BATCH_SIZE, num_workers=1\n",
        ")\n",
        "train_x, train_y = get_cls_embeds(train_loader, tokenizer, bert, device)\n",
        "print('\\n')\n",
        "test_x, test_y = get_cls_embeds(test_loader, tokenizer, bert, device)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pw9dpPAA2bqn",
        "outputId": "5b78c5e8-aad7-44fc-90bd-a95ed71e3955"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Progress: 9.97% done\n",
            "Progress: 19.95% done\n",
            "Progress: 29.92% done\n",
            "Progress: 39.89% done\n",
            "Progress: 49.87% done\n",
            "Progress: 59.84% done\n",
            "Progress: 69.81% done\n",
            "Progress: 79.79% done\n",
            "Progress: 89.76% done\n",
            "Progress: 99.73% done\n",
            "\n",
            "\n",
            "Progress: 9.24% done\n",
            "Progress: 18.49% done\n",
            "Progress: 27.73% done\n",
            "Progress: 36.97% done\n",
            "Progress: 46.22% done\n",
            "Progress: 55.46% done\n",
            "Progress: 64.71% done\n",
            "Progress: 73.95% done\n",
            "Progress: 83.19% done\n",
            "Progress: 92.44% done\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_x.shape, train_y.shape, test_x.shape, test_y.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SD0ysrLt4uOW",
        "outputId": "24f96990-9e11-4aa8-f78b-4f3280593589"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(torch.Size([120000, 768]),\n",
              " torch.Size([120000]),\n",
              " torch.Size([7600, 768]),\n",
              " torch.Size([7600]))"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_x.device, train_y.device, test_x.device, test_y.device"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JJ063a2j769l",
        "outputId": "76564728-06b3-4781-dacf-e87ae6692584"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(device(type='cpu'),\n",
              " device(type='cpu'),\n",
              " device(type='cpu'),\n",
              " device(type='cpu'))"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "train_dataset2 = MyDataset(train_x, train_y)\n",
        "test_dataset2 = MyDataset(test_x, test_y)\n",
        "train_loader, test_loader = get_loaders(\n",
        "    train_dataset2, test_dataset2,\n",
        "    batch_size=BATCH_SIZE * 2, num_workers=1\n",
        ")"
      ],
      "metadata": {
        "id": "1wAUrCAp2i17"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "NUM_EPOCHS = 50\n",
        "instance_train_losses, instance_val_losses, instance_val_accs = train_for_k_epochs(\n",
        "    NUM_EPOCHS,\n",
        "    bert,\n",
        "    mlp,\n",
        "    tokenizer,\n",
        "    optim,\n",
        "    train_loader,\n",
        "    test_loader,\n",
        "    epoch_offset=0,\n",
        "    x_is_cls=True  # the input is the bert cls token;\n",
        ")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Mz4sJlaMUcQs",
        "outputId": "cfdec213-2082-4f99-d751-a4ad948a5e68"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch 0:\n",
            "--------------------\n",
            "train_loss=0.60311\n",
            "val_acc=0.827\tval_loss=0.48858\n",
            "Epoch 1:\n",
            "--------------------\n",
            "train_loss=0.46369\n",
            "val_acc=0.781\tval_loss=0.56927\n",
            "Epoch 2:\n",
            "--------------------\n",
            "train_loss=0.43373\n",
            "val_acc=0.819\tval_loss=0.50623\n",
            "Epoch 3:\n",
            "--------------------\n",
            "train_loss=0.41952\n",
            "val_acc=0.842\tval_loss=0.44147\n",
            "Epoch 4:\n",
            "--------------------\n",
            "train_loss=0.40912\n",
            "val_acc=0.859\tval_loss=0.39412\n",
            "Epoch 5:\n",
            "--------------------\n",
            "train_loss=0.40158\n",
            "val_acc=0.858\tval_loss=0.39736\n",
            "Epoch 6:\n",
            "--------------------\n",
            "train_loss=0.39514\n",
            "val_acc=0.855\tval_loss=0.40200\n",
            "Epoch 7:\n",
            "--------------------\n",
            "train_loss=0.38705\n",
            "val_acc=0.853\tval_loss=0.40167\n",
            "Epoch 8:\n",
            "--------------------\n",
            "train_loss=0.37761\n",
            "val_acc=0.857\tval_loss=0.39522\n",
            "Epoch 9:\n",
            "--------------------\n",
            "train_loss=0.37211\n",
            "val_acc=0.852\tval_loss=0.40112\n",
            "Epoch 10:\n",
            "--------------------\n",
            "train_loss=0.37070\n",
            "val_acc=0.869\tval_loss=0.36864\n",
            "Epoch 11:\n",
            "--------------------\n",
            "train_loss=0.36484\n",
            "val_acc=0.853\tval_loss=0.40091\n",
            "Epoch 12:\n",
            "--------------------\n",
            "train_loss=0.35818\n",
            "val_acc=0.872\tval_loss=0.36125\n",
            "Epoch 13:\n",
            "--------------------\n",
            "train_loss=0.35446\n",
            "val_acc=0.868\tval_loss=0.36661\n",
            "Epoch 14:\n",
            "--------------------\n",
            "train_loss=0.34963\n",
            "val_acc=0.848\tval_loss=0.39543\n",
            "Epoch 15:\n",
            "--------------------\n",
            "train_loss=0.34648\n",
            "val_acc=0.866\tval_loss=0.36910\n",
            "Epoch 16:\n",
            "--------------------\n",
            "train_loss=0.34574\n",
            "val_acc=0.867\tval_loss=0.36856\n",
            "Epoch 17:\n",
            "--------------------\n",
            "train_loss=0.34052\n",
            "val_acc=0.868\tval_loss=0.36121\n",
            "Epoch 18:\n",
            "--------------------\n",
            "train_loss=0.33707\n",
            "val_acc=0.871\tval_loss=0.36539\n",
            "Epoch 19:\n",
            "--------------------\n",
            "train_loss=0.33424\n",
            "val_acc=0.876\tval_loss=0.35169\n",
            "Epoch 20:\n",
            "--------------------\n",
            "train_loss=0.32855\n",
            "val_acc=0.870\tval_loss=0.35934\n",
            "Epoch 21:\n",
            "--------------------\n",
            "train_loss=0.32769\n",
            "val_acc=0.873\tval_loss=0.35482\n",
            "Epoch 22:\n",
            "--------------------\n",
            "train_loss=0.32413\n",
            "val_acc=0.869\tval_loss=0.36094\n",
            "Epoch 23:\n",
            "--------------------\n",
            "train_loss=0.32049\n",
            "val_acc=0.878\tval_loss=0.34219\n",
            "Epoch 24:\n",
            "--------------------\n",
            "train_loss=0.31696\n",
            "val_acc=0.873\tval_loss=0.35013\n",
            "Epoch 25:\n",
            "--------------------\n",
            "train_loss=0.31418\n",
            "val_acc=0.877\tval_loss=0.35023\n",
            "Epoch 26:\n",
            "--------------------\n",
            "train_loss=0.31101\n",
            "val_acc=0.876\tval_loss=0.35194\n",
            "Epoch 27:\n",
            "--------------------\n",
            "train_loss=0.30891\n",
            "val_acc=0.879\tval_loss=0.34210\n",
            "Epoch 28:\n",
            "--------------------\n",
            "train_loss=0.30583\n",
            "val_acc=0.876\tval_loss=0.35944\n",
            "Epoch 29:\n",
            "--------------------\n",
            "train_loss=0.30251\n",
            "val_acc=0.874\tval_loss=0.34863\n",
            "Epoch 30:\n",
            "--------------------\n",
            "train_loss=0.29808\n",
            "val_acc=0.878\tval_loss=0.35088\n",
            "Epoch 31:\n",
            "--------------------\n",
            "train_loss=0.29680\n",
            "val_acc=0.873\tval_loss=0.35037\n",
            "Epoch 32:\n",
            "--------------------\n",
            "train_loss=0.29248\n",
            "val_acc=0.870\tval_loss=0.37347\n",
            "Epoch 33:\n",
            "--------------------\n",
            "train_loss=0.28807\n",
            "val_acc=0.874\tval_loss=0.35622\n",
            "Epoch 34:\n",
            "--------------------\n",
            "train_loss=0.28661\n",
            "val_acc=0.875\tval_loss=0.35526\n",
            "Epoch 35:\n",
            "--------------------\n",
            "train_loss=0.28452\n",
            "val_acc=0.873\tval_loss=0.35872\n",
            "Epoch 36:\n",
            "--------------------\n",
            "train_loss=0.28151\n",
            "val_acc=0.874\tval_loss=0.36061\n",
            "Epoch 37:\n",
            "--------------------\n",
            "train_loss=0.27838\n",
            "val_acc=0.880\tval_loss=0.34708\n",
            "Epoch 38:\n",
            "--------------------\n",
            "train_loss=0.27388\n",
            "val_acc=0.878\tval_loss=0.35205\n",
            "Epoch 39:\n",
            "--------------------\n",
            "train_loss=0.27099\n",
            "val_acc=0.877\tval_loss=0.34782\n",
            "Epoch 40:\n",
            "--------------------\n",
            "train_loss=0.26762\n",
            "val_acc=0.875\tval_loss=0.36563\n",
            "Epoch 41:\n",
            "--------------------\n",
            "train_loss=0.26454\n",
            "val_acc=0.872\tval_loss=0.36726\n",
            "Epoch 42:\n",
            "--------------------\n",
            "train_loss=0.26059\n",
            "val_acc=0.879\tval_loss=0.35914\n",
            "Epoch 43:\n",
            "--------------------\n",
            "train_loss=0.25705\n",
            "val_acc=0.879\tval_loss=0.35714\n",
            "Epoch 44:\n",
            "--------------------\n",
            "train_loss=0.25380\n",
            "val_acc=0.880\tval_loss=0.36341\n",
            "Epoch 45:\n",
            "--------------------\n",
            "train_loss=0.24998\n",
            "val_acc=0.872\tval_loss=0.38534\n",
            "Epoch 46:\n",
            "--------------------\n",
            "train_loss=0.24943\n",
            "val_acc=0.868\tval_loss=0.38580\n",
            "Epoch 47:\n",
            "--------------------\n",
            "train_loss=0.24721\n",
            "val_acc=0.879\tval_loss=0.36725\n",
            "Epoch 48:\n",
            "--------------------\n",
            "train_loss=0.24224\n",
            "val_acc=0.875\tval_loss=0.37030\n",
            "Epoch 49:\n",
            "--------------------\n",
            "train_loss=0.23883\n",
            "val_acc=0.877\tval_loss=0.36496\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Save model;"
      ],
      "metadata": {
        "id": "aMgGy9CS-nWB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "decode = {\n",
        "    0: 'World',\n",
        "    1: 'Sports',\n",
        "    2: 'Business',\n",
        "    3: 'Sci/Tech'\n",
        "}"
      ],
      "metadata": {
        "id": "ds3qbicp-mSj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "outs = model_forward(['Cristiano Ronaldo scored a goal!',\n",
        "                      'Investors are bearish on the french bond market.',\n",
        "                      'The new transformer model beats the performance of Recurrent Neural Networks'],\n",
        "                     tokenizer, bert, mlp, device, x_is_cls=False).argmax(-1)\n",
        "outs"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "LW8zKGNh-V68",
        "outputId": "4ff325a5-d188-4012-e1ea-25a02fe6df15"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "tensor([1, 2, 3], device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "[decode[v] for v in outs.to('cpu').tolist()]"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0aMiV9pz_L5K",
        "outputId": "80b9db5b-4c43-40c7-d500-973decedc5ef"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "['Sports', 'Business', 'Sci/Tech']"
            ]
          },
          "metadata": {},
          "execution_count": 37
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from pathlib import Path\n",
        "p = Path('.').absolute() / 'checkpoints'\n",
        "p.mkdir(parents=True, exist_ok=True)"
      ],
      "metadata": {
        "id": "r4vmVUOXAzfU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.save(\n",
        "    {\n",
        "        'epochs_done': 50,\n",
        "        'bert_was_trained': False,\n",
        "        'batch_size': BATCH_SIZE * 2,\n",
        "        'mlp_state_dict': mlp.state_dict(),\n",
        "        'mlp_optim_state_dict': optim.state_dict()\n",
        "    },\n",
        "    p / 'mlp_agnews.ckpt')"
      ],
      "metadata": {
        "id": "XqmFE7K3_WYb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! ls -a ~/.cache/huggingface/"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dKcKnWszAEGl",
        "outputId": "ff8a8154-6d38-4a1e-ab30-830f9a702026"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            ".  ..  datasets  hub\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Done;"
      ],
      "metadata": {
        "id": "g-EBtvqqBPCT"
      }
    }
  ]
}