examples/notebook_examples/sqaudv1.1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/airaria/TextBrewer/blob/master/examples/notebook_examples/sqaudv1.1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rOvyG9zChDJu"
      },
      "source": [
        "This notebook shows how to fine-tune a model on squadv1.1 dataset and how to distill the model with TextBrewer.\n",
        "\n",
        "Detailed Docs can be find here: https://github.com/airaria/TextBrewer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PqJEpBdjiMKQ"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "device='cuda' if torch.cuda.is_available() else 'cpu'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sgDADHlOiSw4"
      },
      "outputs": [],
      "source": [
        "!pip install transformers\n",
        "!pip install textbrewer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UWyCSU6u0ido"
      },
      "source": [
        "### Prepare dataset to train teacher model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nYjP-ZDOv8zR"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import random\n",
        "import timeit\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
        "from tqdm import tqdm_notebook\n",
        "from transformers import AdamW,BertConfig,BertForQuestionAnswering,BertTokenizer,get_linear_schedule_with_warmup,squad_convert_examples_to_features\n",
        "\n",
        "from transformers.data.metrics.squad_metrics import compute_predictions_logits,squad_evaluate\n",
        "from transformers.data.processors.squad import SquadResult, SquadV1Processor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cir5Dfr8ie0Z",
        "outputId": "536e8619-a4dc-413b-c8de-1cbc091ceda8"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "--2021-07-08 04:22:17--  https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\n",
            "Resolving rajpurkar.github.io (rajpurkar.github.io)... 185.199.108.153, 185.199.109.153, 185.199.110.153, ...\n",
            "Connecting to rajpurkar.github.io (rajpurkar.github.io)|185.199.108.153|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 30288272 (29M) [application/json]\n",
            "Saving to: ‘train-v1.1.json’\n",
            "\n",
            "train-v1.1.json     100%[===================>]  28.88M  78.2MB/s    in 0.4s    \n",
            "\n",
            "2021-07-08 04:22:18 (78.2 MB/s) - ‘train-v1.1.json’ saved [30288272/30288272]\n",
            "\n",
            "--2021-07-08 04:22:18--  https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\n",
            "Reusing existing connection to rajpurkar.github.io:443.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 4854279 (4.6M) [application/json]\n",
            "Saving to: ‘dev-v1.1.json’\n",
            "\n",
            "dev-v1.1.json       100%[===================>]   4.63M  --.-KB/s    in 0.02s   \n",
            "\n",
            "2021-07-08 04:22:19 (215 MB/s) - ‘dev-v1.1.json’ saved [4854279/4854279]\n",
            "\n",
            "FINISHED --2021-07-08 04:22:19--\n",
            "Total wall clock time: 1.7s\n",
            "Downloaded: 2 files, 34M in 0.4s (85.7 MB/s)\n"
          ]
        }
      ],
      "source": [
        "!wget https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0NR7hLfqvyxZ",
        "outputId": "01fb674d-fb22-41e0-db98-8f44873bbb84"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 442/442 [00:33<00:00, 13.29it/s]\n"
          ]
        }
      ],
      "source": [
        "processor = SquadV1Processor()\n",
        "examples = processor.get_train_examples('/content/')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JOjXEHwQwBRZ",
        "outputId": "2ed2e42c-95cc-4c4b-a7b4-667c7642277c"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at bert-base-cased were not used when initializing BertForQuestionAnswering: ['cls.predictions.decoder.weight', 'cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.transform.dense.bias', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
            "- This IS expected if you are initializing BertForQuestionAnswering from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertForQuestionAnswering from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForQuestionAnswering were not initialized from the model checkpoint at bert-base-cased and are newly initialized: ['qa_outputs.bias', 'qa_outputs.weight']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "BertForQuestionAnswering(\n",
              "  (bert): BertModel(\n",
              "    (embeddings): BertEmbeddings(\n",
              "      (word_embeddings): Embedding(28996, 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): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (1): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (2): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (3): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (4): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (5): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (6): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (7): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (8): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (9): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (10): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "        (11): BertLayer(\n",
              "          (attention): BertAttention(\n",
              "            (self): BertSelfAttention(\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",
              "          )\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",
              "  )\n",
              "  (qa_outputs): Linear(in_features=768, out_features=2, bias=True)\n",
              ")"
            ]
          },
          "execution_count": 17,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tokenizer = BertTokenizer.from_pretrained('bert-base-cased')\n",
        "model = BertForQuestionAnswering.from_pretrained('bert-base-cased')\n",
        "\n",
        "model.to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GEcJPauczrjK",
        "outputId": "49111455-999e-4455-faf7-ce6e88b0369e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "convert squad examples to features: 100%|██████████| 87599/87599 [11:27<00:00, 127.42it/s]\n",
            "add example index and unique id: 100%|██████████| 87599/87599 [00:00<00:00, 761102.36it/s]\n"
          ]
        }
      ],
      "source": [
        "features,train_dataset = squad_convert_examples_to_features(\n",
        "    examples=examples,\n",
        "    tokenizer=tokenizer,\n",
        "    max_seq_length=384,\n",
        "    doc_stride=128,\n",
        "    max_query_length=64,\n",
        "    is_training=True,\n",
        "    return_dataset=\"pt\"\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xNJhv3Y_wJxU"
      },
      "outputs": [],
      "source": [
        "train_sampler = RandomSampler(train_dataset)\n",
        "train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=8)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "sH8Ivk5E2Cuf",
        "outputId": "710fd6d0-1119-44b3-a86d-b07d17f33eb4"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[tensor([[ 101, 1332, 1108,  ...,    0,    0,    0],\n",
              "         [ 101, 1327, 2523,  ...,    0,    0,    0],\n",
              "         [ 101, 1130, 1184,  ...,    0,    0,    0],\n",
              "         ...,\n",
              "         [ 101, 1731, 1242,  ...,    0,    0,    0],\n",
              "         [ 101, 1327, 1110,  ...,    0,    0,    0],\n",
              "         [ 101, 1327, 1160,  ...,    0,    0,    0]]),\n",
              " tensor([[1, 1, 1,  ..., 0, 0, 0],\n",
              "         [1, 1, 1,  ..., 0, 0, 0],\n",
              "         [1, 1, 1,  ..., 0, 0, 0],\n",
              "         ...,\n",
              "         [1, 1, 1,  ..., 0, 0, 0],\n",
              "         [1, 1, 1,  ..., 0, 0, 0],\n",
              "         [1, 1, 1,  ..., 0, 0, 0]]),\n",
              " tensor([[0, 0, 0,  ..., 0, 0, 0],\n",
              "         [0, 0, 0,  ..., 0, 0, 0],\n",
              "         [0, 0, 0,  ..., 0, 0, 0],\n",
              "         ...,\n",
              "         [0, 0, 0,  ..., 0, 0, 0],\n",
              "         [0, 0, 0,  ..., 0, 0, 0],\n",
              "         [0, 0, 0,  ..., 0, 0, 0]]),\n",
              " tensor([ 81,  99,  21, 146, 122, 130, 191,  77]),\n",
              " tensor([ 82, 100,  21, 147, 125, 130, 192,  80]),\n",
              " tensor([0, 0, 0, 0, 0, 0, 0, 0]),\n",
              " tensor([[0., 1., 1.,  ..., 1., 1., 1.],\n",
              "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
              "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
              "         ...,\n",
              "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
              "         [0., 1., 1.,  ..., 1., 1., 1.],\n",
              "         [0., 1., 1.,  ..., 1., 1., 1.]]),\n",
              " tensor([0., 0., 0., 0., 0., 0., 0., 0.])]"
            ]
          },
          "execution_count": 9,
          "metadata": {
            "tags": []
          },
          "output_type": "execute_result"
        }
      ],
      "source": [
        "next(iter(train_dataloader))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rXkT9uN79JLi"
      },
      "outputs": [],
      "source": [
        "#Start training\n",
        "import tqdm\n",
        "epochs = 2\n",
        "t_total = len(train_dataloader) * epochs\n",
        "\n",
        "# Prepare optimizer and schedule \n",
        "optimizer = AdamW(model.parameters(), lr=3e-5, eps=1e-8)\n",
        "scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=t_total)\n",
        "\n",
        "\n",
        "print(\"  Num examples = \", len(train_dataset))\n",
        "print(\"  Total optimization steps = \", t_total)\n",
        "\n",
        "steps = 1\n",
        "tr_loss = 0.0\n",
        "model.zero_grad()\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    print('Epoch:{}'.format(epoch+1))\n",
        "    epoch_iterator = tqdm.notebook.tqdm(train_dataloader, desc=\"Iteration\", disable=False)\n",
        "    for step, batch in enumerate(epoch_iterator):\n",
        "\n",
        "        model.train()\n",
        "        batch = tuple(t.to(device) for t in batch)\n",
        "\n",
        "        inputs = {\n",
        "            \"input_ids\": batch[0],\n",
        "            \"attention_mask\": batch[1],\n",
        "            \"token_type_ids\": batch[2],\n",
        "            \"start_positions\": batch[3],\n",
        "            \"end_positions\": batch[4],\n",
        "        }\n",
        "\n",
        "        outputs = model(**inputs)\n",
        "        loss = outputs.loss\n",
        "\n",
        "        loss.backward()\n",
        "\n",
        "        tr_loss += loss.item()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
        "\n",
        "        optimizer.step()\n",
        "        scheduler.step()  \n",
        "        model.zero_grad()\n",
        "        steps += 1\n",
        "        \n",
        "        # Log \n",
        "        if steps % 500 == 0:\n",
        "            print('steps = {}, logging_loss = {}'.format(steps,tr_loss))\n",
        "\n",
        "print(\" steps = %s, average loss = %s\", steps, tr_loss / steps)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k-QLh-3IACr1"
      },
      "outputs": [],
      "source": [
        "torch.save(model.state_dict(), '/content/drive/MyDrive/squad_teacher_model.pt') #save the teacher model weights to distill"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SQHyopKl0qB4"
      },
      "source": [
        "### Start distillation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a8AY5MZMBqpv"
      },
      "outputs": [],
      "source": [
        "import textbrewer\n",
        "from textbrewer import GeneralDistiller\n",
        "from textbrewer import TrainingConfig, DistillationConfig\n",
        "from transformers import BertForSequenceClassification, BertConfig, AdamW,BertTokenizer\n",
        "from transformers import get_linear_schedule_with_warmup\n",
        "import torch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g0SWHH47jaz6"
      },
      "outputs": [],
      "source": [
        "from transformers import BertForQuestionAnswering,BertConfig\n",
        "\n",
        "\n",
        "bert_config = BertConfig.from_json_file('/content/drive/MyDrive/TextBrewer-master/examples/student_config/bert_base_cased_config/bert_config.json')\n",
        "bert_config_T3 = BertConfig.from_json_file('/content/drive/MyDrive/TextBrewer-master/examples/student_config/bert_base_cased_config/bert_config_L3.json')\n",
        "\n",
        "bert_config.output_hidden_states = True\n",
        "bert_config_T3.output_hidden_states = True\n",
        "\n",
        "\n",
        "teacher_model = BertForQuestionAnswering(bert_config) \n",
        "student_model = BertForQuestionAnswering(bert_config_T3) \n",
        "\n",
        "teacher_model.load_state_dict(torch.load('/content/drive/MyDrive/squad_teacher_model.pt'))\n",
        "\n",
        "teacher_model.to(device=device)\n",
        "student_model.to(device=device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KD_vql_pje_J"
      },
      "outputs": [],
      "source": [
        "num_epochs = 20\n",
        "num_training_steps = len(train_dataloader) * num_epochs\n",
        "# Optimizer and learning rate scheduler\n",
        "optimizer = AdamW(student_model.parameters(), lr=1e-4)\n",
        "\n",
        "scheduler_class = get_linear_schedule_with_warmup\n",
        "# arguments dict except 'optimizer'\n",
        "scheduler_args = {'num_warmup_steps':int(0.1*num_training_steps), 'num_training_steps':num_training_steps}\n",
        "\n",
        "\n",
        "def simple_adaptor(batch, model_outputs):\n",
        "    return {'logits': (model_outputs.start_logits,model_outputs.end_logits), \n",
        "        'hidden': model_outputs.hidden_states, \n",
        "        'attention': model_outputs.attentions}\n",
        "\n",
        "distill_config = DistillationConfig(\n",
        "    temperature = 1,\n",
        "    intermediate_matches=[{\"layer_T\":[0,0],  \"layer_S\":[0,0], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
        "               {\"layer_T\":[4,4],  \"layer_S\":[1,1], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
        "               {\"layer_T\":[8,8],  \"layer_S\":[2,2], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1},\n",
        "               {\"layer_T\":[12,12], \"layer_S\":[3,3], \"feature\":\"hidden\", \"loss\":\"mmd\", \"weight\":1}])\n",
        "train_config = TrainingConfig()\n",
        "\n",
        "distiller = GeneralDistiller(\n",
        "    train_config=train_config, distill_config=distill_config,\n",
        "    model_T=teacher_model, model_S=student_model, \n",
        "    adaptor_T=simple_adaptor, adaptor_S=simple_adaptor)\n",
        "\n",
        "def batch_postprocessor(batch):\n",
        "  return {\"input_ids\": batch[0],\n",
        "       \"attention_mask\": batch[1],\n",
        "       \"token_type_ids\": batch[2],\n",
        "       \"start_positions\": batch[3],\n",
        "       \"end_positions\": batch[4]}\n",
        "\n",
        "with distiller:\n",
        "    distiller.train(optimizer, train_dataloader, num_epochs, scheduler_class=scheduler_class, scheduler_args = scheduler_args, batch_postprocessor=batch_postprocessor, callback=None)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5i0ucOPulw4T"
      },
      "outputs": [],
      "source": [
        "bert_config_T3 = BertConfig.from_json_file('/content/drive/MyDrive/data/bert_config/bert_config_L3.json')\n",
        "test_model = BertForQuestionAnswering(bert_config_T3)\n",
        "\n",
        "test_model.load_state_dict(torch.load('/content/saved_models/gs2813.pkl'))\n",
        "test_model.to(device)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ExZU8DTXMAkl",
        "outputId": "a2926358-328a-40d0-aab0-cb5b41f7116e"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 48/48 [00:06<00:00,  7.61it/s]\n",
            "convert squad examples to features: 100%|██████████| 10570/10570 [01:24<00:00, 125.69it/s]\n",
            "add example index and unique id: 100%|██████████| 10570/10570 [00:00<00:00, 685666.02it/s]\n"
          ]
        }
      ],
      "source": [
        "processor = SquadV1Processor()\n",
        "examples = processor.get_dev_examples('/content/')\n",
        "\n",
        "features,eval_dataset = squad_convert_examples_to_features(\n",
        "    examples=examples,\n",
        "    tokenizer=tokenizer,\n",
        "    max_seq_length=384,\n",
        "    doc_stride=128,\n",
        "    max_query_length=64,\n",
        "    is_training=False,\n",
        "    return_dataset=\"pt\"\n",
        ")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VZScjwwolz1V"
      },
      "outputs": [],
      "source": [
        "# Evaluate\n",
        "import tqdm\n",
        "from typing import Any, BinaryIO, Dict, List, Optional, Tuple, Union\n",
        "eval_sampler = SequentialSampler(eval_dataset)\n",
        "eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=8)\n",
        "\n",
        "print(\"  Num examples = \", len(eval_dataset))\n",
        "\n",
        "all_results = []\n",
        "\n",
        "def to_tuple(self) -> Tuple[Any]:\n",
        "    \"\"\"\n",
        "    Convert self to a tuple containing all the attributes/keys that are not ``None``.\n",
        "    \"\"\"\n",
        "    return tuple(self[k] for k in self.keys())\n",
        "\n",
        "for batch in tqdm.notebook.tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
        "    test_model.eval()\n",
        "    batch = tuple(t.to(device) for t in batch)\n",
        "\n",
        "    with torch.no_grad():\n",
        "        inputs = {\n",
        "            \"input_ids\": batch[0],\n",
        "            \"attention_mask\": batch[1],\n",
        "            \"token_type_ids\": batch[2],\n",
        "        }\n",
        "\n",
        "        feature_indices = batch[3]\n",
        "\n",
        "        outputs = model(**inputs)\n",
        "\n",
        "    for i, feature_index in enumerate(feature_indices):\n",
        "        eval_feature = features[feature_index.item()]\n",
        "        unique_id = int(eval_feature.unique_id)\n",
        "        #output = [output[i].detach().cpu().tolist() for output in outputs]\n",
        "        #start_logits, end_logits = output\n",
        "        output = [(output[i]).detach().cpu().tolist() for output in outputs.to_tuple()]\n",
        "        start_logits, end_logits = output\n",
        "        result = SquadResult(unique_id, start_logits, end_logits)\n",
        "        all_results.append(result)\n",
        "\n",
        "\n",
        "predictions = compute_predictions_logits(\n",
        "    examples,\n",
        "    features,\n",
        "    all_results,\n",
        "    n_best_size=20,\n",
        "    max_answer_length=30,\n",
        "    do_lower_case=False,\n",
        "    output_prediction_file=\"predictions.json\",\n",
        "    output_nbest_file=\"nbest_predictions.json\",\n",
        "    output_null_log_odds_file=None,\n",
        "    verbose_logging=False,\n",
        "    version_2_with_negative=False,\n",
        "    null_score_diff_threshold=0.0,\n",
        "    tokenizer=tokenizer,)\n",
        "\n",
        "# Compute the F1 and exact scores.\n",
        "results = squad_evaluate(examples, predictions)\n",
        "\n",
        "print(\"Results: {}\".format(results))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "squadv1.1.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}