mt5_soft_prompt_tuning_large.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "mt5-soft-prompt-tuning-large.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "u8xBtDGlxvz4",
        "outputId": "93c696ef-c394-4cbb-a928-0de74c29551c"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sat Sep 11 09:17:34 2021       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 470.63.01    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla V100-SXM2...  Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   37C    P0    50W / 300W |      0MiB / 16160MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rbpX3_9PxzVE",
        "outputId": "e945ce69-cc50-4e8c-fe75-d453e7a704b8"
      },
      "source": [
        "!pip install transformers SentencePiece torch tqdm"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: transformers in /usr/local/lib/python3.7/dist-packages (4.10.2)\n",
            "Requirement already satisfied: SentencePiece in /usr/local/lib/python3.7/dist-packages (0.1.96)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.7/dist-packages (1.9.0+cu102)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.62.0)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.7/dist-packages (from transformers) (5.4.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (2019.12.20)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers) (2.23.0)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from transformers) (4.6.4)\n",
            "Requirement already satisfied: huggingface-hub>=0.0.12 in /usr/local/lib/python3.7/dist-packages (from transformers) (0.0.16)\n",
            "Requirement already satisfied: sacremoses in /usr/local/lib/python3.7/dist-packages (from transformers) (0.0.45)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers) (3.0.12)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from transformers) (21.0)\n",
            "Requirement already satisfied: tokenizers<0.11,>=0.10.1 in /usr/local/lib/python3.7/dist-packages (from transformers) (0.10.3)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (1.19.5)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from huggingface-hub>=0.0.12->transformers) (3.7.4.3)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->transformers) (2.4.7)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata->transformers) (3.5.0)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2021.5.30)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (7.1.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.15.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.0.1)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sqiBMWGHx3Yx"
      },
      "source": [
        "import math\n",
        "\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "from transformers import MT5ForConditionalGeneration, T5Tokenizer\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from sklearn.metrics import accuracy_score"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w9dorhOHyydQ"
      },
      "source": [
        "class SoftEmbedding(nn.Module):\n",
        "    def __init__(self, \n",
        "                wte: nn.Embedding,\n",
        "                n_tokens: int = 10, \n",
        "                random_range: float = 0.5,\n",
        "                initialize_from_vocab: bool = True):\n",
        "        \"\"\"appends learned embedding to \n",
        "        Args:\n",
        "            wte (nn.Embedding): original transformer word embedding\n",
        "            n_tokens (int, optional): number of tokens for task. Defaults to 10.\n",
        "            random_range (float, optional): range to init embedding (if not initialize from vocab). Defaults to 0.5.\n",
        "            initialize_from_vocab (bool, optional): initalizes from default vocab. Defaults to True.\n",
        "        \"\"\"\n",
        "        super(SoftEmbedding, self).__init__()\n",
        "        self.wte = wte\n",
        "        self.n_tokens = n_tokens\n",
        "        self.learned_embedding = nn.parameter.Parameter(self.initialize_embedding(wte,\n",
        "                                                                                  n_tokens, \n",
        "                                                                                  random_range, \n",
        "                                                                                  initialize_from_vocab))\n",
        "            \n",
        "    def initialize_embedding(self, \n",
        "                             wte: nn.Embedding,\n",
        "                             n_tokens: int = 10, \n",
        "                             random_range: float = 0.5, \n",
        "                             initialize_from_vocab: bool = True):\n",
        "        \"\"\"initializes learned embedding\n",
        "        Args:\n",
        "            same as __init__\n",
        "        Returns:\n",
        "            torch.float: initialized using original schemes\n",
        "        \"\"\"\n",
        "        if initialize_from_vocab:\n",
        "            return self.wte.weight[:n_tokens].clone().detach()\n",
        "        return torch.FloatTensor(n_tokens, wte.weight.size(1)).uniform_(-random_range, random_range)\n",
        "            \n",
        "    def forward(self, tokens):\n",
        "        \"\"\"run forward pass\n",
        "        Args:\n",
        "            tokens (torch.long): input tokens before encoding\n",
        "        Returns:\n",
        "            torch.float: encoding of text concatenated with learned task specifc embedding\n",
        "        \"\"\"\n",
        "        input_embedding = self.wte(tokens[:, self.n_tokens:])\n",
        "        learned_embedding = self.learned_embedding.repeat(input_embedding.size(0), 1, 1)\n",
        "        return torch.cat([learned_embedding, input_embedding], 1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "We6vNt5ukkBJ",
        "outputId": "fee4408f-e1c1-43e3-9c0b-fb29a7e04e85"
      },
      "source": [
        "!pip install zh-dataset-inews"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Requirement already satisfied: zh-dataset-inews in /usr/local/lib/python3.7/dist-packages (0.0.2)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "biQD2_aB6v8s"
      },
      "source": [
        "from zh_dataset_inews import title_train, label_train, title_dev, label_dev, title_test, label_test"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LfF_Z4qu0z3a"
      },
      "source": [
        "def generate_data(batch_size, n_tokens, title_data, label_data):\n",
        "\n",
        "    labels = [\n",
        "        torch.tensor([[3]]),  # \\x00\n",
        "        torch.tensor([[4]]),  # \\x01\n",
        "        torch.tensor([[5]]),  # \\x02\n",
        "    ]\n",
        "\n",
        "    def yield_data(x_batch, y_batch, l_batch):\n",
        "        x = torch.nn.utils.rnn.pad_sequence(x_batch, batch_first=True)\n",
        "        y = torch.cat(y_batch, dim=0)\n",
        "        m = (x > 0).to(torch.float32)\n",
        "        decoder_input_ids = torch.full((x.size(0), n_tokens), 1)\n",
        "        if torch.cuda.is_available():\n",
        "            x = x.cuda()\n",
        "            y = y.cuda()\n",
        "            m = m.cuda()\n",
        "            decoder_input_ids = decoder_input_ids.cuda()\n",
        "        return x, y, m, decoder_input_ids, l_batch\n",
        "\n",
        "    x_batch, y_batch, l_batch = [], [], []\n",
        "    for x, y in zip(title_data, label_data):\n",
        "        context = x\n",
        "        inputs = tokenizer(context, return_tensors=\"pt\")\n",
        "        inputs['input_ids'] = torch.cat([torch.full((1, n_tokens), 1), inputs['input_ids']], 1)\n",
        "        l_batch.append(y)\n",
        "        y = labels[y]\n",
        "        y = torch.cat([torch.full((1, n_tokens - 1), -100), y], 1)\n",
        "        x_batch.append(inputs['input_ids'][0])\n",
        "        y_batch.append(y)\n",
        "        if len(x_batch) >= batch_size:\n",
        "            yield yield_data(x_batch, y_batch, l_batch)\n",
        "            x_batch, y_batch, l_batch = [], [], []\n",
        "\n",
        "    if len(x_batch) > 0:\n",
        "        yield yield_data(x_batch, y_batch, l_batch)\n",
        "        x_batch, y_batch, l_batch = [], [], []"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5NKTnNidXnzS"
      },
      "source": [
        "model = MT5ForConditionalGeneration.from_pretrained(\"google/mt5-large\")\n",
        "tokenizer = T5Tokenizer.from_pretrained(\"google/mt5-large\")\n",
        "n_tokens = 100\n",
        "s_wte = SoftEmbedding(model.get_input_embeddings(), \n",
        "                      n_tokens=n_tokens, \n",
        "                      initialize_from_vocab=True)\n",
        "model.set_input_embeddings(s_wte)\n",
        "if torch.cuda.is_available():\n",
        "    model = model.cuda()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eaKnkA8M4Am4"
      },
      "source": [
        "parameters = list(model.parameters())\n",
        "for x in parameters[1:]:\n",
        "    x.requires_grad = False"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YR1cPDym4LYi",
        "outputId": "e9d9ef6a-4450-48d5-e06d-d88b8721da79"
      },
      "source": [
        "parameters[0]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter containing:\n",
              "tensor([[ -1.0312,  -4.2500,   7.0000,  ...,   6.0938,  -8.0625,  -9.5000],\n",
              "        [ -7.7500, -12.1250,  -2.3438,  ...,  -7.8438,   9.1875,   4.4375],\n",
              "        [  0.9805,   1.0781,  -0.3867,  ...,  -1.0156,  -0.4785,   0.8008],\n",
              "        ...,\n",
              "        [ -1.4922,   0.1895,  -0.2041,  ...,   0.6250,   0.0131,  -1.8828],\n",
              "        [  0.8789,   0.1108,   1.1953,  ...,   0.8281,   1.4844,   0.3418],\n",
              "        [  0.1436,  -0.3867,  -0.7734,  ...,   0.5078,  -0.0157,   0.1060]],\n",
              "       device='cuda:0', requires_grad=True)"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "b20PaXQA4L4n",
        "outputId": "3b5eebd4-e309-4076-e4e7-4f23423ab05e"
      },
      "source": [
        "parameters[2]"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter containing:\n",
              "tensor([[ 0.0099,  0.0084,  0.0172,  ...,  0.0220,  0.0435, -0.0337],\n",
              "        [ 0.0112, -0.0181, -0.0107,  ...,  0.0227,  0.0190,  0.0033],\n",
              "        [ 0.0061,  0.0430,  0.0625,  ..., -0.0334, -0.0130,  0.0205],\n",
              "        ...,\n",
              "        [ 0.0034,  0.0228,  0.0003,  ...,  0.0113, -0.0045, -0.0222],\n",
              "        [ 0.0297, -0.0042, -0.0393,  ...,  0.0037, -0.0145, -0.0023],\n",
              "        [ 0.0053, -0.0029,  0.0157,  ..., -0.0125,  0.0068,  0.0106]],\n",
              "       device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c07hoCo44MP-"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uXKtiZOPkKwZ"
      },
      "source": [
        "for x, y, m, dii, true_labels in generate_data(2, n_tokens, title_train, label_train):\n",
        "    assert dii.shape == y.shape\n",
        "    outputs = model(input_ids=x, labels=y, attention_mask=m, decoder_input_ids=dii)\n",
        "    assert outputs['logits'].shape[:2] == y.shape\n",
        "    pred_labels = outputs['logits'][:, -1, 3:6].argmax(-1).detach().cpu().numpy().tolist()\n",
        "    break"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RdZ3dTfK3-ob",
        "outputId": "60d1d5a9-cf2f-4317-8f03-d140edbf6b8f"
      },
      "source": [
        "batch_size = 2\n",
        "n_epoch = 50\n",
        "total_batch = math.ceil(len(title_train) / batch_size)\n",
        "dev_total_batch = math.ceil(len(title_dev) / batch_size)\n",
        "use_ce_loss = False\n",
        "ce_loss = nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam(s_wte.parameters(), lr=0.5)\n",
        "\n",
        "for epoch in range(n_epoch):\n",
        "    print('epoch', epoch)\n",
        "\n",
        "    all_true_labels = []\n",
        "    all_pred_labels = []\n",
        "    losses = []\n",
        "    pbar = tqdm(enumerate(generate_data(batch_size, n_tokens, title_train, label_train)), total=total_batch)\n",
        "    for i, (x, y, m, dii, true_labels) in pbar:\n",
        "        all_true_labels += true_labels\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        outputs = model(input_ids=x, labels=y, attention_mask=m, decoder_input_ids=dii)\n",
        "        pred_labels = outputs['logits'][:, -1, 3:6].argmax(-1).detach().cpu().numpy().tolist()\n",
        "        all_pred_labels += pred_labels\n",
        "\n",
        "        if use_ce_loss:\n",
        "            logits = outputs['logits'][:, -1, 3:6]\n",
        "            true_labels_tensor = torch.tensor(true_labels, dtype=torch.long).cuda()\n",
        "            loss = ce_loss(logits, true_labels_tensor)\n",
        "        else:\n",
        "            loss = outputs.loss\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        loss_value = float(loss.detach().cpu().numpy().tolist()) / batch_size\n",
        "        losses.append(loss_value)\n",
        "\n",
        "        acc = accuracy_score(all_true_labels, all_pred_labels)\n",
        "        pbar.set_description(f'train: loss={np.mean(losses):.4f}, acc={acc:.4f}')\n",
        "\n",
        "    all_true_labels = []\n",
        "    all_pred_labels = []\n",
        "    losses = []\n",
        "    with torch.no_grad():\n",
        "        pbar = tqdm(enumerate(generate_data(batch_size, n_tokens, title_dev, label_dev)), total=dev_total_batch)\n",
        "        for i, (x, y, m, dii, true_labels) in pbar:\n",
        "            all_true_labels += true_labels\n",
        "            outputs = model(input_ids=x, labels=y, attention_mask=m, decoder_input_ids=dii)\n",
        "            loss = outputs.loss\n",
        "            loss_value = float(loss.detach().cpu().numpy().tolist()) / batch_size\n",
        "            losses.append(loss_value)\n",
        "            pred_labels = outputs['logits'][:, -1, 3:6].argmax(-1).detach().cpu().numpy().tolist()\n",
        "            all_pred_labels += pred_labels\n",
        "            acc = accuracy_score(all_true_labels, all_pred_labels)\n",
        "            pbar.set_description(f'dev: loss={np.mean(losses):.4f}, acc={acc:.4f}')"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 0\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=27.0702, acc=0.1079: 100%|██████████| 2678/2678 [06:10<00:00,  7.23it/s]\n",
            "dev: loss=24.9345, acc=0.0861: 100%|██████████| 500/500 [00:32<00:00, 15.35it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 1\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=18.2214, acc=0.1617: 100%|██████████| 2678/2678 [06:08<00:00,  7.26it/s]\n",
            "dev: loss=16.0809, acc=0.4344: 100%|██████████| 500/500 [00:32<00:00, 15.47it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 2\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=24.5954, acc=0.2510: 100%|██████████| 2678/2678 [06:08<00:00,  7.26it/s]\n",
            "dev: loss=31.2844, acc=0.0861: 100%|██████████| 500/500 [00:32<00:00, 15.45it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 3\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=27.7627, acc=0.1079: 100%|██████████| 2678/2678 [06:07<00:00,  7.29it/s]\n",
            "dev: loss=17.8570, acc=0.0861: 100%|██████████| 500/500 [00:32<00:00, 15.51it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 4\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=5.3752, acc=0.4288: 100%|██████████| 2678/2678 [06:08<00:00,  7.27it/s]\n",
            "dev: loss=2.5420, acc=0.4935: 100%|██████████| 500/500 [00:32<00:00, 15.43it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 5\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=1.6852, acc=0.4426: 100%|██████████| 2678/2678 [06:07<00:00,  7.29it/s]\n",
            "dev: loss=1.8112, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.38it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 6\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=6.8852, acc=0.3473: 100%|██████████| 2678/2678 [06:07<00:00,  7.29it/s]\n",
            "dev: loss=15.4861, acc=0.0861: 100%|██████████| 500/500 [00:32<00:00, 15.24it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 7\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=5.7757, acc=0.4325: 100%|██████████| 2678/2678 [06:07<00:00,  7.29it/s]\n",
            "dev: loss=2.9378, acc=0.4895: 100%|██████████| 500/500 [00:32<00:00, 15.45it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 8\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=3.1350, acc=0.4405: 100%|██████████| 2678/2678 [06:08<00:00,  7.27it/s]\n",
            "dev: loss=1.9451, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.25it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 9\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=1.4360, acc=0.4459: 100%|██████████| 2678/2678 [06:07<00:00,  7.28it/s]\n",
            "dev: loss=1.5743, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.34it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 10\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.7307, acc=0.4702: 100%|██████████| 2678/2678 [06:07<00:00,  7.28it/s]\n",
            "dev: loss=2.3358, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.38it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 11\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.7700, acc=0.4480: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=1.3577, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.50it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 12\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.6409, acc=0.4652: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=1.1933, acc=0.4925: 100%|██████████| 500/500 [00:32<00:00, 15.43it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 13\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.5962, acc=0.5100: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=1.0476, acc=0.4925: 100%|██████████| 500/500 [00:31<00:00, 15.64it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 14\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.5385, acc=0.5968: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=1.1322, acc=0.4925: 100%|██████████| 500/500 [00:31<00:00, 15.69it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 15\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.4824, acc=0.6390: 100%|██████████| 2678/2678 [06:02<00:00,  7.39it/s]\n",
            "dev: loss=0.8160, acc=0.4925: 100%|██████████| 500/500 [00:31<00:00, 15.64it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 16\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.4478, acc=0.6753: 100%|██████████| 2678/2678 [06:03<00:00,  7.38it/s]\n",
            "dev: loss=0.6584, acc=0.5726: 100%|██████████| 500/500 [00:32<00:00, 15.60it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 17\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.4170, acc=0.6930: 100%|██████████| 2678/2678 [06:02<00:00,  7.39it/s]\n",
            "dev: loss=0.6231, acc=0.6336: 100%|██████████| 500/500 [00:32<00:00, 15.52it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 18\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.3995, acc=0.7029: 100%|██████████| 2678/2678 [06:01<00:00,  7.41it/s]\n",
            "dev: loss=0.5515, acc=0.6667: 100%|██████████| 500/500 [00:32<00:00, 15.56it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 19\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.3821, acc=0.7107: 100%|██████████| 2678/2678 [06:03<00:00,  7.38it/s]\n",
            "dev: loss=0.5967, acc=0.6466: 100%|██████████| 500/500 [00:31<00:00, 15.63it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 20\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.3647, acc=0.7303: 100%|██████████| 2678/2678 [06:01<00:00,  7.40it/s]\n",
            "dev: loss=0.4908, acc=0.6557: 100%|██████████| 500/500 [00:32<00:00, 15.57it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 21\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "train: loss=0.3547, acc=0.7315: 100%|██████████| 2678/2678 [06:01<00:00,  7.42it/s]\n",
            "dev: loss=0.4713, acc=0.7067: 100%|██████████| 500/500 [00:31<00:00, 15.65it/s]\n"
          ]
        },
        {
          "metadata": {
            "tags": null
          },
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "epoch 22\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.3406, acc=0.7415: 100%|██████████| 2678/2678 [06:01<00:00,  7.40it/s]\n",
            "dev: loss=0.4078, acc=0.7417: 100%|██████████| 500/500 [00:31<00:00, 15.64it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 23\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.3255, acc=0.7565: 100%|██████████| 2678/2678 [06:02<00:00,  7.39it/s]\n",
            "dev: loss=0.3752, acc=0.7337: 100%|██████████| 500/500 [00:32<00:00, 15.61it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 24\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.3098, acc=0.7621: 100%|██████████| 2678/2678 [06:02<00:00,  7.38it/s]\n",
            "dev: loss=0.3692, acc=0.7407: 100%|██████████| 500/500 [00:32<00:00, 15.57it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 25\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2986, acc=0.7811: 100%|██████████| 2678/2678 [06:02<00:00,  7.40it/s]\n",
            "dev: loss=0.3479, acc=0.7628: 100%|██████████| 500/500 [00:32<00:00, 15.50it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 26\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2924, acc=0.7824: 100%|██████████| 2678/2678 [06:02<00:00,  7.39it/s]\n",
            "dev: loss=0.3496, acc=0.7548: 100%|██████████| 500/500 [00:32<00:00, 15.62it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 27\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2865, acc=0.7826: 100%|██████████| 2678/2678 [06:00<00:00,  7.42it/s]\n",
            "dev: loss=0.3458, acc=0.7608: 100%|██████████| 500/500 [00:31<00:00, 15.69it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 28\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2746, acc=0.7950: 100%|██████████| 2678/2678 [06:03<00:00,  7.37it/s]\n",
            "dev: loss=0.3276, acc=0.7648: 100%|██████████| 500/500 [00:31<00:00, 15.63it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 29\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2642, acc=0.8011: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=0.3346, acc=0.7578: 100%|██████████| 500/500 [00:32<00:00, 15.58it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 30\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2633, acc=0.8021: 100%|██████████| 2678/2678 [06:05<00:00,  7.33it/s]\n",
            "dev: loss=0.3168, acc=0.7477: 100%|██████████| 500/500 [00:32<00:00, 15.35it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 31\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2450, acc=0.8121: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=0.3279, acc=0.7628: 100%|██████████| 500/500 [00:32<00:00, 15.46it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 32\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2345, acc=0.8239: 100%|██████████| 2678/2678 [06:04<00:00,  7.34it/s]\n",
            "dev: loss=0.3625, acc=0.7538: 100%|██████████| 500/500 [00:32<00:00, 15.55it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 33\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2345, acc=0.8243: 100%|██████████| 2678/2678 [06:03<00:00,  7.36it/s]\n",
            "dev: loss=0.3248, acc=0.7618: 100%|██████████| 500/500 [00:32<00:00, 15.56it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 34\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2715, acc=0.7938: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=0.3177, acc=0.7487: 100%|██████████| 500/500 [00:32<00:00, 15.54it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 35\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2448, acc=0.8136: 100%|██████████| 2678/2678 [06:05<00:00,  7.34it/s]\n",
            "dev: loss=0.3142, acc=0.7638: 100%|██████████| 500/500 [00:32<00:00, 15.44it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 36\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2281, acc=0.8288: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=0.3680, acc=0.7518: 100%|██████████| 500/500 [00:32<00:00, 15.35it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 37\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2265, acc=0.8317: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=0.3473, acc=0.7588: 100%|██████████| 500/500 [00:32<00:00, 15.46it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 38\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2330, acc=0.8263: 100%|██████████| 2678/2678 [06:04<00:00,  7.34it/s]\n",
            "dev: loss=0.3554, acc=0.7497: 100%|██████████| 500/500 [00:32<00:00, 15.37it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 39\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2277, acc=0.8261: 100%|██████████| 2678/2678 [06:05<00:00,  7.33it/s]\n",
            "dev: loss=0.3790, acc=0.7407: 100%|██████████| 500/500 [00:32<00:00, 15.42it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 40\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2221, acc=0.8304: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=0.3345, acc=0.7407: 100%|██████████| 500/500 [00:32<00:00, 15.40it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 41\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2417, acc=0.8144: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=0.3752, acc=0.7347: 100%|██████████| 500/500 [00:32<00:00, 15.50it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 42\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2384, acc=0.8174: 100%|██████████| 2678/2678 [06:06<00:00,  7.31it/s]\n",
            "dev: loss=0.3535, acc=0.7588: 100%|██████████| 500/500 [00:32<00:00, 15.43it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 43\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2198, acc=0.8303: 100%|██████████| 2678/2678 [06:06<00:00,  7.32it/s]\n",
            "dev: loss=0.3527, acc=0.7467: 100%|██████████| 500/500 [00:32<00:00, 15.48it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 44\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2125, acc=0.8400: 100%|██████████| 2678/2678 [06:05<00:00,  7.33it/s]\n",
            "dev: loss=0.3540, acc=0.7447: 100%|██████████| 500/500 [00:32<00:00, 15.47it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 45\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2063, acc=0.8476: 100%|██████████| 2678/2678 [06:04<00:00,  7.34it/s]\n",
            "dev: loss=0.3499, acc=0.7518: 100%|██████████| 500/500 [00:32<00:00, 15.53it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 46\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.2001, acc=0.8514: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=0.3517, acc=0.7467: 100%|██████████| 500/500 [00:32<00:00, 15.48it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 47\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.1885, acc=0.8646: 100%|██████████| 2678/2678 [06:05<00:00,  7.33it/s]\n",
            "dev: loss=0.3732, acc=0.7477: 100%|██████████| 500/500 [00:32<00:00, 15.37it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 48\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.1811, acc=0.8650: 100%|██████████| 2678/2678 [06:05<00:00,  7.32it/s]\n",
            "dev: loss=0.4366, acc=0.7227: 100%|██████████| 500/500 [00:32<00:00, 15.33it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "epoch 49\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "train: loss=0.1723, acc=0.8726: 100%|██████████| 2678/2678 [06:04<00:00,  7.34it/s]\n",
            "dev: loss=0.4849, acc=0.7387: 100%|██████████| 500/500 [00:32<00:00, 15.41it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Zrxp8Im04beQ"
      },
      "source": [
        "parameters2 = list(model.parameters())"
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rRBcvIbY4s-R",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "baf9152a-f81c-4192-9966-b453dd07833c"
      },
      "source": [
        "parameters2[0]"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter containing:\n",
              "tensor([[  39.9078, -138.6385,  217.1636,  ...,   29.5207,  145.8943,\n",
              "          144.5315],\n",
              "        [ 515.6390,  -31.6162,   51.5134,  ...,  -53.0618,  245.4292,\n",
              "          -69.3007],\n",
              "        [ 236.2896,   43.0374,  -19.2581,  ..., -127.3152,  130.7397,\n",
              "           31.1689],\n",
              "        ...,\n",
              "        [  88.1486,   49.1501,  125.5696,  ...,  113.4881,   96.0846,\n",
              "          368.0652],\n",
              "        [ 100.6963, -102.7619,  -35.8637,  ..., -144.5385,  -25.3403,\n",
              "          173.1718],\n",
              "        [-164.1508,  -81.5056,  152.1980,  ..., -178.5098,    6.0514,\n",
              "         -129.9609]], device='cuda:0', requires_grad=True)"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Mp_0DmKr4tZA",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a48ec7e3-794a-49ea-d069-8da58e3ab218"
      },
      "source": [
        "parameters2[2]"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Parameter containing:\n",
              "tensor([[ 0.0099,  0.0084,  0.0172,  ...,  0.0220,  0.0435, -0.0337],\n",
              "        [ 0.0112, -0.0181, -0.0107,  ...,  0.0227,  0.0190,  0.0033],\n",
              "        [ 0.0061,  0.0430,  0.0625,  ..., -0.0334, -0.0130,  0.0205],\n",
              "        ...,\n",
              "        [ 0.0034,  0.0228,  0.0003,  ...,  0.0113, -0.0045, -0.0222],\n",
              "        [ 0.0297, -0.0042, -0.0393,  ...,  0.0037, -0.0145, -0.0023],\n",
              "        [ 0.0053, -0.0029,  0.0157,  ..., -0.0125,  0.0068,  0.0106]],\n",
              "       device='cuda:0')"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Oqdi4A7tFnxM"
      },
      "source": [
        "def predict(text):\n",
        "    inputs = tokenizer(text, return_tensors='pt')\n",
        "    inputs['input_ids'] = torch.cat([torch.full((1, n_tokens), 1), inputs['input_ids']], 1)\n",
        "\n",
        "    decoder_input_ids = torch.full((1, n_tokens), 1)\n",
        "    with torch.no_grad():\n",
        "        outputs = model(input_ids=inputs['input_ids'].cuda(), decoder_input_ids=decoder_input_ids.cuda())\n",
        "    logits = outputs['logits'][:, -1, 3:6]\n",
        "    pred = logits.argmax(-1).detach().cpu().numpy()[0]\n",
        "    # print(logits)\n",
        "    return pred"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ohlT29oGueUD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ef705633-ccb3-4a66-d732-3e272834f6f2"
      },
      "source": [
        "train_rets = []\n",
        "for i in tqdm(range(len(title_train))):\n",
        "    pred = predict(title_train[i])\n",
        "    train_rets.append((label_train[i], pred, title_train[i]))"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 5355/5355 [04:39<00:00, 19.19it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NE84AeC9BC8U"
      },
      "source": [
        "rets = []\n",
        "for i in tqdm(range(len(title_test))):\n",
        "    pred = predict(title_test[i])\n",
        "    rets.append((label_test[i], pred, title_test[i]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DZOnkUtKuiwa",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "83bda13e-9dc1-41c3-ff0e-6b696dc65ccc"
      },
      "source": [
        "print(\n",
        "    accuracy_score(\n",
        "        [x[0] for x in train_rets],\n",
        "        [x[1] for x in train_rets],\n",
        "    )\n",
        ")"
      ],
      "execution_count": 23,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.861624649859944\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "74CWYzXXuW-W",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c1a7e002-279d-4602-9ebf-19df96135ea1"
      },
      "source": [
        "print(\n",
        "    accuracy_score(\n",
        "        [x[0] for x in rets],\n",
        "        [x[1] for x in rets],\n",
        "    )\n",
        ")"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.7447447447447447\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D8vIKsvAuCpX",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "73f9bd61-f349-41af-b08b-96c20a5ea304"
      },
      "source": [
        "print(\n",
        "    accuracy_score(\n",
        "        [x[0] for x in rets],\n",
        "        [0] * len(rets),\n",
        "    ),\n",
        "    accuracy_score(\n",
        "        [x[0] for x in rets],\n",
        "        [1] * len(rets),\n",
        "    ),\n",
        "    accuracy_score(\n",
        "        [x[0] for x in rets],\n",
        "        [2] * len(rets),\n",
        "    )\n",
        ")"
      ],
      "execution_count": 25,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0.0990990990990991 0.4944944944944945 0.4064064064064064\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1O1uwg8irIdc"
      },
      "source": [
        ""
      ],
      "execution_count": 22,
      "outputs": []
    }
  ]
}