JaQuAD.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CEvslqjDCiUM"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "awj7ZLT7rOGQ"
      },
      "outputs": [],
      "source": [
        "# We assume that you are working on CoLab. This code works on Python 3.7.\n",
        "# You may change the version of PyTorch/XLA and other packages depending on your\n",
        "# python version.\n",
        "!pip install -q torch==1.9.0 torchtext torchvision pytorch_lightning\n",
        "!pip install -q cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.9-cp37-cp37m-linux_x86_64.whl\n",
        "\n",
        "!pip install -q transformers==4.15 transformers[\"ja\"] datasets==1.18\n",
        "\n",
        "# You can ignore pip's dependency error on torchaudio and earthengine-api. They\n",
        "# won't be used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "533xQ1HdZBhv"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "if not os.environ.get('COLAB_TPU_ADDR'):\n",
        "    print('Make sure to select TPU from Edit > Notebook settings > Hardware accelerator')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65cwlByIqbVC"
      },
      "source": [
        "## Define model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "GVNY-Mt3rSYn"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import string\n",
        "import unicodedata\n",
        "from typing import Any, Dict, Iterator, List, Tuple, Union\n",
        "\n",
        "import datasets\n",
        "import numpy as np\n",
        "from pytorch_lightning import LightningModule, Trainer, seed_everything\n",
        "import torch\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.optim import lr_scheduler\n",
        "from transformers import AutoModelForQuestionAnswering, AutoTokenizer, AdamW"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cmT8NWJbrVFa",
        "outputId": "989e2fd5-8b6c-4cb8-cb76-d602e1a156c6"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'batch_size': 32,\n",
              " 'cpu_workers': 2,\n",
              " 'dataset': 'SkelterLabsInc/JaQuAD',\n",
              " 'huggingface_auth_token': None,\n",
              " 'doc_stride': 128,\n",
              " 'epochs': 4,\n",
              " 'fp16': False,\n",
              " 'lr': 2e-05,\n",
              " 'lr_scheduler': 'warmup_lin',\n",
              " 'max_length': 384,\n",
              " 'norm_form': 'NFKC',\n",
              " 'optimizer': 'AdamW',\n",
              " 'pretrained_model': 'cl-tohoku/bert-base-japanese-whole-word-masking',\n",
              " 'pretrained_tokenizer': '',\n",
              " 'random_seed': 42,\n",
              " 'test_mode': False,\n",
              " 'tpu_cores': 8,\n",
              " 'warmup_ratio': 0.1,\n",
              " 'weight_decay': 0.01}"
            ]
          },
          "execution_count": 6,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "args = {\n",
        "    'random_seed': 42,  # Random Seed\n",
        "    # Transformers PLM name.\n",
        "    'pretrained_model': 'cl-tohoku/bert-base-japanese-whole-word-masking',\n",
        "    # Optional, Transformers Tokenizer name. Overrides `pretrained_model`\n",
        "    'pretrained_tokenizer': '',\n",
        "    'norm_form': 'NFKC',\n",
        "    'batch_size': 32,  # <=32 for TPUv2-8\n",
        "    'lr': 2e-5,  # Learning Rate\n",
        "    'max_length': 384,  # Max Length input size\n",
        "    'doc_stride': 128,  # The interval of the context when splitting is needed\n",
        "    'epochs': 4,  # Max Epochs\n",
        "    'dataset': 'SkelterLabsInc/JaQuAD',\n",
        "    'huggingface_auth_token': None,\n",
        "    'test_mode': False,  # Test Mode enables `fast_dev_run`\n",
        "    'optimizer': 'AdamW',\n",
        "    'weight_decay': 0.01,  # Weight decaying parameter for AdamW\n",
        "    'lr_scheduler': 'warmup_lin',\n",
        "    'warmup_ratio': 0.1,\n",
        "    'fp16': False,  # Enable train on FP16 (if GPU)\n",
        "    'tpu_cores': 8,  # Enable TPU with 1 core or 8 cores\n",
        "    'cpu_workers': os.cpu_count(),\n",
        "}\n",
        "\n",
        "args"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YYsBE7rf0Fm6"
      },
      "outputs": [],
      "source": [
        "class WarmupLinearLR(lr_scheduler.LambdaLR):\n",
        "    '''The learning rate is linearly increased for the first `warmup_steps`\n",
        "    and linearly decreased to zero afterward.\n",
        "    '''\n",
        "\n",
        "    def __init__(self, optimizer, warmup_steps, max_steps, last_epoch=-1):\n",
        "\n",
        "        def lr_lambda(step):\n",
        "            if step < warmup_steps:\n",
        "                return float(step) / float(max(1.0, warmup_steps))\n",
        "            ratio = 1 - float(step - warmup_steps) / float(max_steps -\n",
        "                                                           warmup_steps)\n",
        "            return max(0.0, min(1.0, ratio))\n",
        "\n",
        "        super().__init__(optimizer, lr_lambda, last_epoch=last_epoch)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ke6TdX1EX26I"
      },
      "outputs": [],
      "source": [
        "def make_spans(\n",
        "    inputs: Dict[str, Union[int, List[int]]],\n",
        "    question_len: int,\n",
        "    max_seq_len: int,\n",
        "    stride: int,\n",
        "    answer_start_position: int = -1,\n",
        "    answer_end_position: int = -1\n",
        ") -> Iterator[Tuple[Dict[str, List[int]], Tuple[int, int]]]:\n",
        "    input_len = len(inputs['input_ids'])\n",
        "    context_len = input_len - question_len\n",
        "\n",
        "    def make_value(input_list, i, padding=0):\n",
        "        context_end = min(max_seq_len - question_len, context_len - i)\n",
        "        pad_len = max_seq_len - question_len - context_end\n",
        "        val = input_list[:question_len]\n",
        "        val += input_list[question_len + i:question_len + i + context_end]\n",
        "        val[-1] = input_list[-1]\n",
        "        val += [padding] * pad_len\n",
        "        return val\n",
        "\n",
        "    for i in range(0, input_len - max_seq_len + stride, stride):\n",
        "        span = {key: make_value(val, i) for key, val in inputs.items()}\n",
        "        answer_start = answer_start_position - i\n",
        "        answer_end = answer_end_position - i\n",
        "        if answer_start < question_len or answer_end >= max_seq_len - 1:\n",
        "            answer_start = answer_end = 0\n",
        "        yield span, (answer_start, answer_end)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "def get_offsets(input_ids: List[int],\n",
        "                context: str,\n",
        "                tokenizer: AutoTokenizer,\n",
        "                norm_form='NFKC') -> List[Tuple[int, int]]:\n",
        "    '''The character-level start/end offsets of a token within a context.\n",
        "    Algorithm:\n",
        "    1. Make offsets of normalized context within the original context.\n",
        "    2. Make offsets of tokens (input_ids) within the normalized context.\n",
        "\n",
        "    Arguments:\n",
        "    input_ids -- Token ids of tokenized context (by tokenizer).\n",
        "    context -- String of context\n",
        "    tokenizer\n",
        "    norm_form\n",
        "\n",
        "    Return:\n",
        "        List[Tuple[int, int]]: Offsets of tokens within the input context.\n",
        "        For each token, the offsets are presented as a tuple of (start\n",
        "        position index, end position index). Both indices are inclusive.\n",
        "    '''\n",
        "    cxt_start = input_ids.index(tokenizer.sep_token_id) + 1\n",
        "    cxt_end = cxt_start + input_ids[cxt_start:].index(tokenizer.sep_token_id)\n",
        "    tokens = tokenizer.convert_ids_to_tokens(input_ids[cxt_start:cxt_end])\n",
        "    tokens = [tok[2:] if tok.startswith('##') else tok for tok in tokens]\n",
        "    whitespace = string.whitespace + '\\u3000'\n",
        "\n",
        "    # 1. Make offsets of normalized context within the original context.\n",
        "    offsets_norm_context = []\n",
        "    norm_context = ''\n",
        "    for idx, char in enumerate(context):\n",
        "        norm_char = unicodedata.normalize(norm_form, char)\n",
        "        norm_context += norm_char\n",
        "        offsets_norm_context.extend([idx] * len(norm_char))\n",
        "    norm_context_org = unicodedata.normalize(norm_form, context)\n",
        "    assert norm_context == norm_context_org, \\\n",
        "        'Normalized contexts are not the same: ' \\\n",
        "        + f'{norm_context} != {norm_context_org}'\n",
        "    assert len(norm_context) == len(offsets_norm_context), \\\n",
        "        'Normalized contexts have different numbers of tokens: ' \\\n",
        "        + f'{len(norm_context)} != {len(offsets_norm_context)}'\n",
        "\n",
        "    # 2. Make offsets of tokens (input_ids) within the normalized context.\n",
        "    offsets_token = []\n",
        "    unk_pointer = None\n",
        "    cid = 0\n",
        "    tid = 0\n",
        "    while tid < len(tokens):\n",
        "        cur_token = tokens[tid]\n",
        "        if cur_token == tokenizer.unk_token:\n",
        "            unk_pointer = tid\n",
        "            offsets_token.append([cid, cid])\n",
        "            cid += 1\n",
        "        elif norm_context[cid:cid + len(cur_token)] != cur_token:\n",
        "            # Wrong offsets of the previous UNK token\n",
        "            assert unk_pointer is not None, \\\n",
        "                'Normalized context and tokens are not matched'\n",
        "            prev_unk_expected = offsets_token[unk_pointer]\n",
        "            prev_unk_expected[1] += norm_context[prev_unk_expected[1] + 2:]\\\n",
        "                .index(tokens[unk_pointer + 1]) + 1\n",
        "            tid = unk_pointer\n",
        "            offsets_token = offsets_token[:tid] + [prev_unk_expected]\n",
        "            cid = prev_unk_expected[1] + 1\n",
        "        else:\n",
        "            start_pos = norm_context[cid:].index(cur_token)\n",
        "            if start_pos > 0 and tokens[tid - 1] == tokenizer.unk_token:\n",
        "                offsets_token[-1][1] += start_pos\n",
        "                cid += start_pos\n",
        "                start_pos = 0\n",
        "            assert start_pos == 0, f'{start_pos} != 0 (cur: {cur_token}'\n",
        "            offsets_token.append([cid, cid + len(cur_token) - 1])\n",
        "            cid += len(cur_token)\n",
        "            while cid < len(norm_context) and norm_context[cid] in whitespace:\n",
        "                offsets_token[-1][1] += 1\n",
        "                cid += 1\n",
        "        tid += 1\n",
        "    if tokens[-1] == tokenizer.unk_token:\n",
        "        offsets_token[-1][1] = len(norm_context) - 1\n",
        "    else:\n",
        "        assert cid == len(norm_context) == offsets_token[-1][1] + 1, \\\n",
        "            'Offsets do not include all characters'\n",
        "    assert len(offsets_token) == len(tokens), \\\n",
        "        'The numbers of tokens and offsets are different'\n",
        "\n",
        "    offsets_mapping = [(offsets_norm_context[start], offsets_norm_context[end])\n",
        "                       for start, end in offsets_token]\n",
        "    return [(-1, -1)] * cxt_start + offsets_mapping"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "erLcpPrsYJTN"
      },
      "outputs": [],
      "source": [
        "class QAModel(LightningModule):\n",
        "    def __init__(self, **kwargs):\n",
        "        super().__init__()\n",
        "        self.save_hyperparameters()  # kwargs are saved in self.hparams\n",
        "        self.automatic_optimization = False\n",
        "\n",
        "        self.question_answerer = AutoModelForQuestionAnswering.from_pretrained(\n",
        "            self.hparams.pretrained_model,\n",
        "            use_auth_token=self.hparams.huggingface_auth_token)\n",
        "        self.tokenizer = AutoTokenizer.from_pretrained(\n",
        "            self.hparams.pretrained_tokenizer\n",
        "            if self.hparams.pretrained_tokenizer else\n",
        "            self.hparams.pretrained_model,\n",
        "            use_auth_token=self.hparams.huggingface_auth_token)\n",
        "\n",
        "    def forward(self, **kwargs):\n",
        "        return self.question_answerer(**kwargs)\n",
        "\n",
        "    def step(self, batch, batch_idx):\n",
        "        outputs = self(**batch)\n",
        "\n",
        "        loss = outputs.loss\n",
        "        start_preds = outputs.start_logits.argmax(dim=-1).cpu().detach()\n",
        "        end_preds = outputs.end_logits.argmax(dim=-1).cpu().detach()\n",
        "        start_positions = batch['start_positions'].cpu().detach()\n",
        "        end_positions = batch['end_positions'].cpu().detach()\n",
        "\n",
        "        return {\n",
        "            'loss': loss,\n",
        "            'start_preds': start_preds,\n",
        "            'end_preds': end_preds,\n",
        "            'start_positions': start_positions,\n",
        "            'end_positions': end_positions,\n",
        "        }\n",
        "\n",
        "    def training_step(self, batch, batch_idx):\n",
        "        opt = self.optimizers()\n",
        "        opt.zero_grad()\n",
        "        outputs = self.step(batch, batch_idx)\n",
        "        self.manual_backward(outputs['loss'])\n",
        "        opt.step()\n",
        "\n",
        "        # single scheduler\n",
        "        sch = self.lr_schedulers()\n",
        "        sch.step()\n",
        "        return outputs\n",
        "\n",
        "    def validation_step(self, batch, batch_idx):\n",
        "        return self.step(batch, batch_idx)\n",
        "\n",
        "    @staticmethod\n",
        "    def calculate_f1_score(start_positions, end_positions, start_preds,\n",
        "                           end_preds):\n",
        "        start_overlap = np.maximum(start_positions, start_preds)\n",
        "        start_overlap = np.maximum(start_positions, start_preds)\n",
        "        end_overlap = np.minimum(end_positions, end_preds)\n",
        "        overlap = np.maximum(end_overlap - start_overlap + 1, 0)\n",
        "\n",
        "        pred_token_count = np.maximum(end_preds - start_preds + 1, 0)\n",
        "        ground_token_count = np.maximum(end_positions - start_positions + 1, 0)\n",
        "\n",
        "        precision = torch.nan_to_num(overlap / pred_token_count, nan=0.)\n",
        "        recall = torch.nan_to_num(overlap / ground_token_count, nan=0.)\n",
        "        f1 = torch.nan_to_num(\n",
        "            2 * precision * recall / (precision + recall), nan=0.)\n",
        "        return {\n",
        "            'precision': precision,\n",
        "            'recall': recall,\n",
        "            'f1': f1,\n",
        "        }\n",
        "\n",
        "    @staticmethod\n",
        "    def calculate_exact_match(start_positions, end_positions, start_preds,\n",
        "                              end_preds):\n",
        "        equal_start = (start_preds == start_positions)\n",
        "        equal_end = (end_preds == end_positions)\n",
        "        return (equal_start * equal_end).type(torch.float)\n",
        "\n",
        "    def epoch_end(self, outputs, state='train'):\n",
        "        loss = torch.tensor(0, dtype=torch.float)\n",
        "        precision = torch.tensor(0, dtype=torch.float)\n",
        "        recall = torch.tensor(0, dtype=torch.float)\n",
        "        f1 = torch.tensor(0, dtype=torch.float)\n",
        "        em = torch.tensor(0, dtype=torch.float)\n",
        "\n",
        "        for i in outputs:\n",
        "            loss += i['loss'].cpu().detach()\n",
        "            f1_metrics = self.calculate_f1_score(i['start_positions'],\n",
        "                                                 i['end_positions'],\n",
        "                                                 i['start_preds'],\n",
        "                                                 i['end_preds'])\n",
        "            precision += f1_metrics['precision'].mean()\n",
        "            recall += f1_metrics['recall'].mean()\n",
        "            f1 += f1_metrics['f1'].mean()\n",
        "            em += self.calculate_exact_match(i['start_positions'],\n",
        "                                             i['end_positions'],\n",
        "                                             i['start_preds'],\n",
        "                                             i['end_preds']).mean()\n",
        "        loss = loss / len(outputs)\n",
        "        precision = precision / len(outputs)\n",
        "        recall = recall / len(outputs)\n",
        "        f1 = f1 / len(outputs)\n",
        "        em = em / len(outputs)\n",
        "        metrics = {\n",
        "            state + '_loss': float(loss),\n",
        "            state + '_precision': precision,\n",
        "            state + '_recall': recall,\n",
        "            state + '_f1': f1,\n",
        "            state + '_em': em,\n",
        "        }\n",
        "\n",
        "        self.log_dict(metrics, on_epoch=True)\n",
        "\n",
        "        return metrics\n",
        "\n",
        "    def training_epoch_end(self, outputs):\n",
        "        self.epoch_end(outputs, state='train')\n",
        "\n",
        "    def validation_epoch_end(self, outputs):\n",
        "        self.epoch_end(outputs, state='val')\n",
        "\n",
        "    def configure_optimizers(self):\n",
        "        if self.hparams.optimizer == 'AdamW':\n",
        "            optimizer = AdamW(\n",
        "                self.parameters(),\n",
        "                lr=self.hparams.lr,\n",
        "                weight_decay=self.hparams.weight_decay)\n",
        "        else:\n",
        "            raise NotImplementedError('Only AdamW is Supported!')\n",
        "\n",
        "        if self.hparams.lr_scheduler == 'cos':\n",
        "            scheduler = lr_scheduler.CosineAnnealingWarmRestarts(\n",
        "                optimizer, T_0=1, T_mult=2)\n",
        "        elif self.hparams.lr_scheduler == 'exp':\n",
        "            scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.5)\n",
        "        elif self.hparams.lr_scheduler == 'warmup_lin':\n",
        "            steps_per_epoch = len(self.train_dataloader())\n",
        "            if os.environ.get('COLAB_TPU_ADDR'):\n",
        "                steps_per_epoch = steps_per_epoch // self.hparams.tpu_cores\n",
        "            total_steps = steps_per_epoch * self.hparams.epochs\n",
        "            warmup_steps = int(total_steps * self.hparams.warmup_ratio)\n",
        "            scheduler = WarmupLinearLR(\n",
        "                optimizer, warmup_steps=warmup_steps, max_steps=total_steps)\n",
        "        else:\n",
        "            raise NotImplementedError(\n",
        "                'Only cos, exp, and warmup_lin lr scheduler is Supported!')\n",
        "\n",
        "        return [optimizer], [scheduler]\n",
        "\n",
        "    def preprocess_function(self, examples):\n",
        "        tokenized_examples = self.tokenizer(\n",
        "            examples['question'],\n",
        "            examples['context'],\n",
        "        )\n",
        "\n",
        "        inputs = {\n",
        "            'input_ids': [],\n",
        "            'attention_mask': [],\n",
        "            'token_type_ids': [],\n",
        "            'start_positions': [],\n",
        "            'end_positions': [],\n",
        "        }\n",
        "        for tokens, att_mask, type_ids, context, answer, start_char \\\n",
        "                in zip(tokenized_examples['input_ids'],\n",
        "                       tokenized_examples['attention_mask'],\n",
        "                       tokenized_examples['token_type_ids'],\n",
        "                       examples['context'],\n",
        "                       examples['answer'],\n",
        "                       examples['answer_start']):\n",
        "            answer = answer[0]\n",
        "            start_char = start_char[0]\n",
        "            offsets = get_offsets(tokens, context, self.tokenizer,\n",
        "                                  self.hparams.norm_form)\n",
        "\n",
        "            ctx_start = tokens.index(self.tokenizer.sep_token_id) + 1\n",
        "            answer_start_index = ctx_start\n",
        "            answer_end_index = len(offsets) - 1\n",
        "            while offsets[answer_start_index][0] < start_char:\n",
        "                answer_start_index += 1\n",
        "            while offsets[answer_end_index][1] > start_char + len(answer):\n",
        "                answer_end_index -= 1\n",
        "\n",
        "            span_inputs = {\n",
        "                'input_ids': tokens,\n",
        "                'attention_mask': att_mask,\n",
        "                'token_type_ids': type_ids,\n",
        "            }\n",
        "            for span, answer_idx in make_spans(\n",
        "                    span_inputs,\n",
        "                    question_len=ctx_start,\n",
        "                    max_seq_len=self.hparams.max_length,\n",
        "                    stride=self.hparams.doc_stride,\n",
        "                    answer_start_position=answer_start_index,\n",
        "                    answer_end_position=answer_end_index):\n",
        "                inputs['input_ids'].append(span['input_ids'])\n",
        "                inputs['attention_mask'].append(span['attention_mask'])\n",
        "                inputs['token_type_ids'].append(span['token_type_ids'])\n",
        "                inputs['start_positions'].append(answer_idx[0])\n",
        "                inputs['end_positions'].append(answer_idx[1])\n",
        "        return inputs\n",
        "\n",
        "    def prepare_data(self):\n",
        "        datasetdict = datasets.load_dataset(\n",
        "            self.hparams.dataset,\n",
        "            use_auth_token=self.hparams.huggingface_auth_token)\n",
        "        datasetdict = datasetdict.flatten()\\\n",
        "            .rename_column('answers.text', 'answer')\\\n",
        "            .rename_column('answers.answer_start', 'answer_start')\\\n",
        "            .rename_column('answers.answer_type', 'answer_type')\n",
        "\n",
        "        self.tokenized_dataset = datasetdict.map(\n",
        "            self.preprocess_function,\n",
        "            batched=True,\n",
        "            remove_columns=datasetdict['train'].column_names)\n",
        "\n",
        "    def dataloader(self, dataset, shuffle=False):\n",
        "        dataset.set_format(type='torch')\n",
        "\n",
        "        return DataLoader(\n",
        "            dataset,\n",
        "            batch_size=self.hparams.batch_size,\n",
        "            shuffle=shuffle,\n",
        "            num_workers=self.hparams.cpu_workers,\n",
        "        )\n",
        "\n",
        "    def train_dataloader(self):\n",
        "        return self.dataloader(self.tokenized_dataset['train'], shuffle=True)\n",
        "\n",
        "    def val_dataloader(self):\n",
        "        return self.dataloader(\n",
        "            self.tokenized_dataset['validation'], shuffle=False)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nN41xPPunWQJ"
      },
      "source": [
        "## Start training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZzpIpf64rZsA"
      },
      "outputs": [],
      "source": [
        "from pytorch_lightning import callbacks\n",
        "from pytorch_lightning import loggers\n",
        "\n",
        "checkpoint_callback = callbacks.ModelCheckpoint(\n",
        "    filename='val_loss{val_loss:.4f}-val_f1{val_f1:.4f}-epoch{epoch}',\n",
        "    monitor='val_loss',\n",
        "    mode='min',\n",
        "    save_top_k=5,\n",
        "    auto_insert_metric_name=False,\n",
        ")\n",
        "lr_callback = callbacks.LearningRateMonitor(logging_interval='step')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_aPXGrWBnWQL",
        "outputId": "dc9a0eee-2b45-46c9-e980-2c2d8ab3859e"
      },
      "outputs": [],
      "source": [
        "print('Using PyTorch Ver', torch.__version__)\n",
        "print('Fix Seed:', args['random_seed'])\n",
        "seed_everything(args['random_seed'])\n",
        "\n",
        "model = QAModel(**args)\n",
        "\n",
        "trainer = Trainer(\n",
        "    callbacks=[lr_callback,\n",
        "               checkpoint_callback],\n",
        "    log_every_n_steps=16,  # Logging frequency of **learning rate**\n",
        "    max_epochs=args['epochs'],\n",
        "    fast_dev_run=args['test_mode'],\n",
        "    num_sanity_val_steps=None if args['test_mode'] else 0,\n",
        "    # For GPU Setup\n",
        "    # deterministic=torch.cuda.is_available(),\n",
        "    # gpus=[0] if torch.cuda.is_available() else None,  # Use one GPU (idx 0)\n",
        "    # precision=16 if args['fp16'] and torch.cuda.is_available() else 32,\n",
        "    # For TPU Setup\n",
        "    tpu_cores=args['tpu_cores'] if args['tpu_cores'] else None,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Es7nKQ3w26vq"
      },
      "outputs": [],
      "source": [
        "print(':: Start Training ::')\n",
        "trainer.fit(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0XbLPtuPY-hJ"
      },
      "outputs": [],
      "source": [
        "%reload_ext tensorboard\n",
        "%tensorboard --logdir lightning_logs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gqAgK6JVnWQM"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UTNeTGApNlWf"
      },
      "outputs": [],
      "source": [
        "# Load fine-tuned model from huggingface (fine-tuned by Skelter Labs)\n",
        "# If you are trying to use your model, skip this cell.\n",
        "!apt-get install git-lfs\n",
        "\n",
        "args['pretrained_model'] = 'SkelterLabsInc/bert-base-japanese-jaquad'\n",
        "model = QAModel(**args)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GTpDsD8-bapF"
      },
      "outputs": [],
      "source": [
        "def get_answers(model: AutoModelForQuestionAnswering,\n",
        "                context: str,\n",
        "                question: str,\n",
        "                n_best_size: int = 5,\n",
        "                max_seq_len: int = 384,\n",
        "                doc_stride: int = 128) -> List[Dict[str, Any]]:\n",
        "    valid_answers = []\n",
        "    inputs = model.tokenizer(question, context)\n",
        "    offsets = get_offsets(inputs['input_ids'], context, model.tokenizer,\n",
        "                          model.hparams.norm_form)\n",
        "    question_len = inputs['input_ids'].index(model.tokenizer.sep_token_id) + 1\n",
        "    i = 0\n",
        "    for span, _ in make_spans(\n",
        "            inputs,\n",
        "            question_len=question_len,\n",
        "            max_seq_len=max_seq_len,\n",
        "            stride=doc_stride):\n",
        "        for key, val in span.items():\n",
        "            span[key] = torch.Tensor([val]).type(torch.long)\n",
        "        output = model(**span)\n",
        "        start_logits = output.start_logits[0].cpu().detach().numpy()\n",
        "        end_logits = output.end_logits[0].cpu().detach().numpy()\n",
        "        start_indexes = np.argsort(start_logits)[-1:-n_best_size -\n",
        "                                                 1:-1].tolist()\n",
        "        end_indexes = np.argsort(end_logits)[-1:-n_best_size - 1:-1].tolist()\n",
        "        cur_offsets = offsets[i:]\n",
        "        i += doc_stride\n",
        "        for start_index in start_indexes:\n",
        "            for end_index in end_indexes:\n",
        "                if 0 < start_index <= end_index < len(cur_offsets):\n",
        "                    # We need to refine that test to check the answer is inside\n",
        "                    # the context\n",
        "                    valid_answers.append({\n",
        "                        'score':\n",
        "                            start_logits[start_index] + end_logits[end_index],\n",
        "                        'position': (start_index, end_index),\n",
        "                        'text':\n",
        "                            context[cur_offsets[start_index][0]:\n",
        "                                    cur_offsets[end_index - 1][1] + 1],\n",
        "                    })\n",
        "    if not valid_answers:\n",
        "        return [{\n",
        "            'score': -float('inf'),\n",
        "            'position': (-1, -1),\n",
        "            'text': '',\n",
        "        }]\n",
        "    valid_answers.sort(key=lambda x: x['score'], reverse=True)\n",
        "    return valid_answers[:n_best_size]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0k1DIdDYg5_N"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "\n",
        "def compute_f1_score(ground_truth_values: List[str],\n",
        "                     prediction_values: List[str]) -> float:\n",
        "    '''Compute f1 score comparing two list of values.'''\n",
        "    common = (\n",
        "        collections.Counter(prediction_values) &\n",
        "        collections.Counter(ground_truth_values))\n",
        "    num_same = sum(common.values())\n",
        "\n",
        "    # No answer case.\n",
        "    if not ground_truth_values or not prediction_values:\n",
        "        return int(ground_truth_values == prediction_values)\n",
        "\n",
        "    if num_same == 0:\n",
        "        return 0.\n",
        "\n",
        "    precision = 1.0 * num_same / len(prediction_values)\n",
        "    recall = 1.0 * num_same / len(ground_truth_values)\n",
        "    f1_score = (2 * precision * recall) / (precision + recall)\n",
        "    return f1_score\n",
        "\n",
        "\n",
        "def char_f1_score(prediction: str, ground_truth: str) -> float:\n",
        "    '''Character F1 score.'''\n",
        "    prediction_tokens = prediction.split()\n",
        "    ground_truth_tokens = ground_truth.split()\n",
        "\n",
        "    # F1 by character\n",
        "    prediction_char = []\n",
        "    for tok in prediction_tokens:\n",
        "        prediction_char.extend(list(tok))\n",
        "\n",
        "    ground_truth_char = []\n",
        "    for tok in ground_truth_tokens:\n",
        "        ground_truth_char.extend(list(tok))\n",
        "    return compute_f1_score(ground_truth_char, prediction_char)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uKGSIdi0skje"
      },
      "outputs": [],
      "source": [
        "jaquad_dataset = datasets.load_dataset(\n",
        "    args['dataset'], use_auth_token=args['huggingface_auth_token']).flatten()\n",
        "jaquad_dataset = jaquad_dataset.rename_column('answers.text', 'answer')\n",
        "\n",
        "infer_data = jaquad_dataset['validation'][:]\n",
        "\n",
        "cnt = len(infer_data['question'])\n",
        "print(f'Eval {cnt} data')\n",
        "\n",
        "f1_scores = []\n",
        "em = []\n",
        "for i, (context, question, answer) in enumerate(\n",
        "        zip(infer_data['context'], infer_data['question'],\n",
        "            infer_data['answer'])):\n",
        "    answer = answer[0]\n",
        "    predictions = get_answers(model, context=context, question=question)\n",
        "    pred_text = predictions[0]['text']\n",
        "    f1 = char_f1_score(pred_text, answer)\n",
        "    f1_scores.append(f1)\n",
        "    em.append(f1 == 1.)\n",
        "    if i % 200 == 0:\n",
        "        print(f'  {i+1}/{cnt} | EM: {sum(em) / (i+1):.4f}, '\\\n",
        "              + f'F1: {sum(f1_scores) / (i+1):.4f}')\n",
        "        print(f'        (Sample) pred: \"{pred_text}\", answer: \"{answer}\"')\n",
        "print(f'F1 score: {sum(f1_scores) / cnt}')\n",
        "print(f'Exact Match: {sum(em) / cnt}')\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "khT4Jiln92Tb"
      },
      "source": [
        "## Analysis\n",
        "\n",
        "Following cells generate graphs used in the paper \"JaQuAD: Japanese Question Answering Dataset for Machine Reading Comprehension\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pzzMnITpcAfO"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "\n",
        "import datasets\n",
        "from transformers import AutoTokenizer\n",
        "\n",
        "\n",
        "jaquad = datasets.load_dataset(\n",
        "    args['dataset'],\n",
        "    use_auth_token=args['huggingface_auth_token']).flatten()\n",
        "\n",
        "tokenizer = model.tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tdkqrdc694OI"
      },
      "source": [
        "### Context/Answer lengths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ubY_ftnV916m",
        "outputId": "57efdc26-c7ee-43cc-99b5-ea1e5c6c21e7"
      },
      "outputs": [],
      "source": [
        "context_lengths = []\n",
        "question_lengths = []\n",
        "answer_lengths = []\n",
        "\n",
        "for dset in jaquad.values():\n",
        "    for batch in dset:\n",
        "        context_tokens = tokenizer(batch['context'])['input_ids']\n",
        "        context_lengths.append(len(context_tokens) - 2)\n",
        "        question_tokens = tokenizer(batch['question'])['input_ids']\n",
        "        question_lengths.append(len(question_tokens) - 2)\n",
        "        answer_tokens = tokenizer(batch['answers.text'][0])['input_ids']\n",
        "        answer_lengths.append(len(answer_tokens) - 2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "15dPWXpLgM00",
        "outputId": "86dbbb3c-3d13-4740-b89f-cbf137cfc9c5"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAD4CAYAAAAO9oqkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAARJElEQVR4nO3df6zddX3H8edrrYDipK3cNLVt1hobl2qywRooYTGEulKYsSxBU2JGZd2aTNzULXF0/tH6g0U2I8oyUSJ1xTB+DNlomI50BbLsDypFHPLDrlcYtE2hVwu4SfxRfe+P87l46O4F7j333nMv9/lITs73+/5+vuf7+fR7e173++Ocm6pCkjS7/Uq/OyBJ6j/DQJJkGEiSDANJEoaBJAmY2+8OjNepp55ay5Yt63c3JGnGuP/++79fVQMjLZuxYbBs2TL27t3b725I0oyR5InRlnmaSJJkGEiSDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJzOBPIKtj2z3bJn8b50z+NiT1l0cGkiTDQJJkGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkXkEYJNme5EiSh7pqC5LsSrK/Pc9v9SS5OslgkgeTnN61zsbWfn+SjV3130rynbbO1Uky0YOUJL20V3Jk8PfAuuNqlwO7q2oFsLvNA5wPrGiPzcA10AkPYCtwJnAGsHU4QFqbP+pa7/htSZIm2cuGQVX9O3D0uPJ6YEeb3gFc2FW/vjruBeYlWQScB+yqqqNV9QywC1jXlr2hqu6tqgKu73otSdIUGe81g4VVdbhNPwUsbNOLgQNd7Q622kvVD45QH1GSzUn2Jtk7NDQ0zq5Lko7X8wXk9ht9TUBfXsm2rq2qVVW1amBgYCo2KUmzwnjD4Ol2iof2fKTVDwFLu9otabWXqi8ZoS5JmkLj/RvIO4GNwKfb8+1d9Q8muYnOxeLnqupwkjuBv+q6aLwW2FJVR5P8MMlqYA9wCfC34+zTtDIVf5tYkibKy4ZBkhuBc4BTkxykc1fQp4FbkmwCngDe25p/HbgAGASeBy4FaG/6nwTua+0+UVXDF6U/QOeOpdcC32gPSdIUetkwqKqLR1m0ZoS2BVw2yutsB7aPUN8LvP3l+iFJmjx+AlmSZBhIkgwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CSRI9hkOQjSR5O8lCSG5OclGR5kj1JBpPcnOSE1vbENj/Yli/rep0trb4vyXm9DUmSNFbjDoMki4E/BVZV1duBOcAG4Ergqqp6C/AMsKmtsgl4ptWvau1IsrKt9zZgHfCFJHPG2y9J0tj1eppoLvDaJHOB1wGHgXOBW9vyHcCFbXp9m6ctX5MkrX5TVf2kqh4HBoEzeuyXJGkMxh0GVXUI+AzwJJ0QeA64H3i2qo61ZgeBxW16MXCgrXustX9jd32EdV4kyeYke5PsHRoaGm/XJUnH6eU00Xw6v9UvB94EnEznNM+kqaprq2pVVa0aGBiYzE1J0qzSy2midwKPV9VQVf0MuA04G5jXThsBLAEOtelDwFKAtvwU4Afd9RHWkSRNgV7C4ElgdZLXtXP/a4BHgLuBi1qbjcDtbXpnm6ctv6uqqtU3tLuNlgMrgG/20C9J0hjNffkmI6uqPUluBb4FHAMeAK4F/gW4KcmnWu26tsp1wFeTDAJH6dxBRFU9nOQWOkFyDLisqn4+3n5JksZu3GEAUFVbga3HlR9jhLuBqurHwHtGeZ0rgCt66Yskafz8BLIkyTCQJBkGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJose/Z6DZYds92yZ/G+dM/jYkjc4jA0mSYSBJMgwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEj2GQZJ5SW5N8t0kjyY5K8mCJLuS7G/P81vbJLk6yWCSB5Oc3vU6G1v7/Uk29jooSdLY9Hpk8HngX6vq14HfAB4FLgd2V9UKYHebBzgfWNEem4FrAJIsALYCZwJnAFuHA0SSNDXGHQZJTgHeAVwHUFU/rapngfXAjtZsB3Bhm14PXF8d9wLzkiwCzgN2VdXRqnoG2AWsG2+/JElj18uRwXJgCPhKkgeSfDnJycDCqjrc2jwFLGzTi4EDXesfbLXR6v9Pks1J9ibZOzQ01EPXJUndegmDucDpwDVVdRrwI355SgiAqiqgetjGi1TVtVW1qqpWDQwMTNTLStKs10sYHAQOVtWeNn8rnXB4up3+oT0facsPAUu71l/SaqPVJUlTZNxhUFVPAQeSvLWV1gCPADuB4TuCNgK3t+mdwCXtrqLVwHPtdNKdwNok89uF47WtJkmaInN7XP9PgBuSnAA8BlxKJ2BuSbIJeAJ4b2v7deACYBB4vrWlqo4m+SRwX2v3iao62mO/JElj0FMYVNW3gVUjLFozQtsCLhvldbYD23vpiyRp/PwEsiTJMJAkGQaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSSJCQiDJHOSPJDkjja/PMmeJINJbk5yQquf2OYH2/JlXa+xpdX3JTmv1z5JksZmIo4MPgQ82jV/JXBVVb0FeAbY1OqbgGda/arWjiQrgQ3A24B1wBeSzJmAfkmSXqGewiDJEuB3gS+3+QDnAre2JjuAC9v0+jZPW76mtV8P3FRVP6mqx4FB4Ixe+iVJGptejww+B3wU+EWbfyPwbFUda/MHgcVtejFwAKAtf661f6E+wjovkmRzkr1J9g4NDfXYdUnSsHGHQZJ3AUeq6v4J7M9Lqqprq2pVVa0aGBiYqs1K0qve3B7WPRt4d5ILgJOANwCfB+Ylmdt++18CHGrtDwFLgYNJ5gKnAD/oqg/rXkeSNAXGfWRQVVuqaklVLaNzAfiuqnofcDdwUWu2Ebi9Te9s87Tld1VVtfqGdrfRcmAF8M3x9kuSNHa9HBmM5i+Am5J8CngAuK7VrwO+mmQQOEonQKiqh5PcAjwCHAMuq6qfT0K/JEmjmJAwqKp7gHva9GOMcDdQVf0YeM8o618BXDERfZEkjZ2fQJYkGQaSJMNAksTkXECWxmzbPdumZjvnTM12pJnGIwNJkmEgSTIMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiRgbr870A9T9cfXJWmm8MhAkmQYSJIMA0kShoEkiR7CIMnSJHcneSTJw0k+1OoLkuxKsr89z2/1JLk6yWCSB5Oc3vVaG1v7/Uk29j4sSdJY9HJkcAz486paCawGLkuyErgc2F1VK4DdbR7gfGBFe2wGroFOeABbgTOBM4CtwwEiSZoa4w6DqjpcVd9q0/8DPAosBtYDO1qzHcCFbXo9cH113AvMS7IIOA/YVVVHq+oZYBewbrz9kiSN3YRcM0iyDDgN2AMsrKrDbdFTwMI2vRg40LXawVYbrT7SdjYn2Ztk79DQ0ER0XZLEBIRBktcDXwM+XFU/7F5WVQVUr9voer1rq2pVVa0aGBiYqJeVpFmvpzBI8ho6QXBDVd3Wyk+30z+05yOtfghY2rX6klYbrS5JmiK93E0U4Drg0ar6bNeincDwHUEbgdu76pe0u4pWA8+100l3AmuTzG8Xjte2miRpivTy3URnA78PfCfJt1vtL4FPA7ck2QQ8Aby3Lfs6cAEwCDwPXApQVUeTfBK4r7X7RFUd7aFfkqQxGncYVNV/ABll8ZoR2hdw2SivtR3YPt6+SJJ64yeQJUmGgSTJMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJLo7VtLpRln2z3bJn8b50z+NqSJ5pGBJMkwkCQZBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJ+HUU0oTzKy80E3lkIEkyDCRJhoEkCa8ZSDPSVFyXmCpe/5gePDKQJE2fMEiyLsm+JINJLu93fyRpNpkWYZBkDvB3wPnASuDiJCv72ytJmj2mRRgAZwCDVfVYVf0UuAlY3+c+SdKsMV0uIC8GDnTNHwTOPL5Rks3A5jb7v0n2jXN7pwLfH+e608mrZRzgWKajKRnHx/n4ZG8CXj37BHoby6+NtmC6hMErUlXXAtf2+jpJ9lbVqgnoUl+9WsYBjmU6erWMAxzLKzFdThMdApZ2zS9pNUnSFJguYXAfsCLJ8iQnABuAnX3ukyTNGtPiNFFVHUvyQeBOYA6wvaoensRN9nyqaZp4tYwDHMt09GoZBziWl5WqmozXlSTNINPlNJEkqY8MA0nS7AqDmfaVF0mWJrk7ySNJHk7yoVZfkGRXkv3teX6rJ8nVbXwPJjm9vyN4sSRzkjyQ5I42vzzJntbfm9vNAyQ5sc0PtuXL+tnv4yWZl+TWJN9N8miSs2bwPvlI+9l6KMmNSU6aKfslyfYkR5I81FUb835IsrG1359k4zQZx9+0n68Hk/xTknldy7a0cexLcl5Xvbf3t6qaFQ86F6a/B7wZOAH4T2Blv/v1Mn1eBJzepn8V+C86X9fx18DlrX45cGWbvgD4BhBgNbCn32M4bjx/BvwDcEebvwXY0Ka/CPxxm/4A8MU2vQG4ud99P24cO4A/bNMnAPNm4j6h82HPx4HXdu2P98+U/QK8AzgdeKirNqb9ACwAHmvP89v0/GkwjrXA3DZ9Zdc4Vrb3rhOB5e09bc5EvL/1/QdyCv/BzwLu7JrfAmzpd7/GOIbbgd8B9gGLWm0RsK9Nfwm4uKv9C+36/aDz2ZHdwLnAHe0/5fe7fuBf2D907io7q03Pbe3S7zG0/pzS3kBzXH0m7pPhT/4vaP/OdwDnzaT9Aiw77k10TPsBuBj4Ulf9Re36NY7jlv0ecEObftH71vA+mYj3t9l0mmikr7xY3Ke+jFk7JD8N2AMsrKrDbdFTwMI2PZ3H+Dngo8Av2vwbgWer6lib7+7rC+Noy59r7aeD5cAQ8JV2yuvLSU5mBu6TqjoEfAZ4EjhM59/5fmbmfhk21v0wbfdPlz+gc1QDkziO2RQGM1aS1wNfAz5cVT/sXladXwOm9f3BSd4FHKmq+/vdlwkwl84h/TVVdRrwIzqnI14wE/YJQDufvp5OwL0JOBlY19dOTaCZsh9eSpKPAceAGyZ7W7MpDGbkV14keQ2dILihqm5r5aeTLGrLFwFHWn26jvFs4N1J/pvON9KeC3wemJdk+IOP3X19YRxt+SnAD6aywy/hIHCwqva0+VvphMNM2ycA7wQer6qhqvoZcBudfTUT98uwse6Habt/krwfeBfwvhZsMInjmE1hMOO+8iJJgOuAR6vqs12LdgLDdz1spHMtYbh+SbtzYjXwXNchc99U1ZaqWlJVy+j8u99VVe8D7gYuas2OH8fw+C5q7afFb3hV9RRwIMlbW2kN8AgzbJ80TwKrk7yu/awNj2XG7ZcuY90PdwJrk8xvR0prW62vkqyjc1r13VX1fNeincCGdmfXcmAF8E0m4v2tnxd/+nCR5gI6d+R8D/hYv/vzCvr723QOcx8Evt0eF9A5T7sb2A/8G7CgtQ+dPxL0PeA7wKp+j2GEMZ3DL+8menP7QR4E/hE4sdVPavODbfmb+93v48bwm8Detl/+mc5dKDNynwAfB74LPAR8lc5dKjNivwA30rnW8TM6R2ybxrMf6JyTH2yPS6fJOAbpXAMY/n//xa72H2vj2Aec31Xv6f3Nr6OQJM2q00SSpFEYBpIkw0CSZBhIkjAMJEkYBpIkDANJEvB/nCpk72cO5OUAAAAASUVORK5CYII=",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "### Context lengths\n",
        "\n",
        "bin_size = 100\n",
        "bins = range(1, max(context_lengths) + bin_size, bin_size)\n",
        "plt.hist(context_lengths,\n",
        "         bins,\n",
        "         density=False,\n",
        "         facecolor='g',\n",
        "         alpha=0.5)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MpvnzjoIx1lt"
      },
      "outputs": [],
      "source": [
        "bins = range(0, max(context_lengths) + bin_size, bin_size)\n",
        "df = pd.DataFrame({'context_lengths': context_lengths})\n",
        "df['bin'] = pd.cut(df.context_lengths, bins)\n",
        "df.bin.value_counts().sort_index()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "YP-rRk6FmrGo",
        "outputId": "f4ed8639-e07c-4adc-d5ce-96c959947611"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAD4CAYAAAAO9oqkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAT1UlEQVR4nO3df4xdZ33n8fdn7SZLabNxyKxl7HRtqGEVotYBK2TVgrKkJE4W4VBVrK2qMTTCIBItbCt1k+WPGLqR6BbKbiRqZIg3zgpi0oQ0FjINxosarbQOHoPl2PlRT341Yzn2FAPplipg+O4f95ntiTNjj+eO54f9fklX95zv+fU8OuP5+Dzn3DupKiRJ57Z/NtMNkCTNPMNAkmQYSJIMA0kShoEkCZg/0w2YrIsvvriWLl06082QpDllz549f1dVAyfW52wYLF26lMHBwZluhiTNKUmeH6vuMJEkyTCQJBkGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkpjDn0Ce9TZsOLuOI+ms5pWBJMkwkCRNIAySbE5yNMn+Tu2rSfa213NJ9rb60iT/2Fn2hc42b0vyWJKhJHcmSatflGRHkoPtfcGZ6KgkaXwTuTK4G1jVLVTVv6+qFVW1AngA+Fpn8dOjy6rqI536RuBDwPL2Gt3nrcDOqloO7GzzkqRpdMowqKpHgGNjLWv/u38/cO/J9pFkEXBBVe2qqgLuAW5oi1cDW9r0lk5dkjRN+r1n8A7gSFUd7NSWJflekr9O8o5WWwwMd9YZbjWAhVV1uE2/CCwc72BJ1icZTDI4MjLSZ9MlSaP6DYO1vPKq4DDwK1V1OfAHwFeSXDDRnbWrhjrJ8k1VtbKqVg4MvOoP9UiSJmnSnzNIMh/4beBto7Wqehl4uU3vSfI08CbgELCks/mSVgM4kmRRVR1uw0lHJ9smSdLk9HNl8FvAk1X1/4d/kgwkmdem30DvRvEzbRjopSRXtvsMNwIPtc22Aeva9LpOXZI0TSbyaOm9wP8B3pxkOMlNbdEaXn3j+J3Avvao6f3AR6pq9ObzR4EvAUPA08A3Wv3TwLuTHKQXMJ/uoz+SpEk45TBRVa0dp/6BMWoP0HvUdKz1B4HLxqh/H7j6VO2QJJ05fgJZkmQYSJIMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRITCIMkm5McTbK/U9uQ5FCSve11fWfZbUmGkjyV5NpOfVWrDSW5tVNfluTRVv9qkvOmsoOSpFObyJXB3cCqMeqfq6oV7bUdIMmlwBrgLW2bP08yL8k84PPAdcClwNq2LsCftH39KvAD4KZ+OiRJOn2nDIOqegQ4NsH9rQa2VtXLVfUsMARc0V5DVfVMVf0E2AqsThLgXcD9bfstwA2n2QdJUp/6uWdwS5J9bRhpQastBl7orDPcauPVXwf8sKqOn1CXJE2jyYbBRuCNwArgMPDZKWvRSSRZn2QwyeDIyMh0HFKSzgmTCoOqOlJVP6uqnwNfpDcMBHAIuKSz6pJWG6/+feDCJPNPqI933E1VtbKqVg4MDEym6ZKkMUwqDJIs6sy+Dxh90mgbsCbJ+UmWAcuB7wC7geXtyaHz6N1k3lZVBXwb+J22/Trgocm0SZI0efNPtUKSe4GrgIuTDAO3A1clWQEU8BzwYYCqOpDkPuBx4Dhwc1X9rO3nFuBhYB6wuaoOtEP8J2Brkv8CfA+4a8p6J0makFOGQVWtHaM87i/sqroDuGOM+nZg+xj1Z/inYSZJ0gzwE8iSJMNAkmQYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJElMIAySbE5yNMn+Tu1PkzyZZF+SB5Nc2OpLk/xjkr3t9YXONm9L8liSoSR3JkmrX5RkR5KD7X3BmeioJGl8E7kyuBtYdUJtB3BZVf0a8DfAbZ1lT1fVivb6SKe+EfgQsLy9Rvd5K7CzqpYDO9u8JGkanTIMquoR4NgJtW9W1fE2uwtYcrJ9JFkEXFBVu6qqgHuAG9ri1cCWNr2lU5ckTZOpuGfw+8A3OvPLknwvyV8neUerLQaGO+sMtxrAwqo63KZfBBaOd6Ak65MMJhkcGRmZgqZLkqDPMEjyCeA48OVWOgz8SlVdDvwB8JUkF0x0f+2qoU6yfFNVrayqlQMDA320XJLUNX+yGyb5APAe4Or2S5yqehl4uU3vSfI08CbgEK8cSlrSagBHkiyqqsNtOOnoZNskSZqcSV0ZJFkF/BHw3qr6cac+kGRem34DvRvFz7RhoJeSXNmeIroReKhttg1Y16bXdeqSpGlyyiuDJPcCVwEXJxkGbqf39ND5wI72hOiu9uTQO4FPJfkp8HPgI1U1evP5o/SeTHoNvXsMo/cZPg3cl+Qm4Hng/VPSM0nShJ0yDKpq7Rjlu8ZZ9wHggXGWDQKXjVH/PnD1qdohSTpz/ASyJMkwkCQZBpIk+ni0VLPEhg1nxzEkzSivDCRJhoEkyTCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkiQmGQZLNSY4m2d+pXZRkR5KD7X1BqyfJnUmGkuxL8tbONuva+geTrOvU35bksbbNnUkylZ2UJJ3cRK8M7gZWnVC7FdhZVcuBnW0e4DpgeXutBzZCLzyA24G3A1cAt48GSFvnQ53tTjyWJOkMmlAYVNUjwLETyquBLW16C3BDp35P9ewCLkyyCLgW2FFVx6rqB8AOYFVbdkFV7aqqAu7p7EuSNA36uWewsKoOt+kXgYVtejHwQme94VY7WX14jPqrJFmfZDDJ4MjISB9NlyR1TckN5PY/+pqKfZ3iOJuqamVVrRwYGDjTh5Okc0Y/YXCkDfHQ3o+2+iHgks56S1rtZPUlY9QlSdOknzDYBow+EbQOeKhTv7E9VXQl8KM2nPQwcE2SBe3G8TXAw23ZS0mubE8R3djZlyRpGsyfyEpJ7gWuAi5OMkzvqaBPA/cluQl4Hnh/W307cD0wBPwY+CBAVR1L8sfA7rbep6pq9Kb0R+k9sfQa4BvtJUmaJhMKg6paO86iq8dYt4Cbx9nPZmDzGPVB4LKJtEWSNPX8BLIkyTCQJBkGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRJ9hEGSNyfZ23m9lOTjSTYkOdSpX9/Z5rYkQ0meSnJtp76q1YaS3NpvpyRJp2f+ZDesqqeAFQBJ5gGHgAeBDwKfq6rPdNdPcimwBngL8HrgW0ne1BZ/Hng3MAzsTrKtqh6fbNskSadn0mFwgquBp6vq+STjrbMa2FpVLwPPJhkCrmjLhqrqGYAkW9u6hoEkTZOpumewBri3M39Lkn1JNidZ0GqLgRc66wy32nj1V0myPslgksGRkZEparokqe8wSHIe8F7gL1ppI/BGekNIh4HP9nuMUVW1qapWVtXKgYGBqdqtJJ3zpmKY6Drgu1V1BGD0HSDJF4Gvt9lDwCWd7Za0GiepS5KmwVQME62lM0SUZFFn2fuA/W16G7AmyflJlgHLge8Au4HlSZa1q4w1bV1J0jTp68ogyWvpPQX04U75vyZZARTw3OiyqjqQ5D56N4aPAzdX1c/afm4BHgbmAZur6kA/7ZIknZ6+wqCq/gF43Qm13zvJ+ncAd4xR3w5s76ctkqTJ8xPIkiTDQJJkGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJKYgjBI8lySx5LsTTLYahcl2ZHkYHtf0OpJcmeSoST7kry1s591bf2DSdb12y5J0sRN1ZXBv62qFVW1ss3fCuysquXAzjYPcB2wvL3WAxuhFx7A7cDbgSuA20cDRJJ05p2pYaLVwJY2vQW4oVO/p3p2ARcmWQRcC+yoqmNV9QNgB7DqDLVNknSCqQiDAr6ZZE+S9a22sKoOt+kXgYVtejHwQmfb4VYbr/4KSdYnGUwyODIyMgVNlyQBzJ+CffxmVR1K8i+BHUme7C6sqkpSU3AcqmoTsAlg5cqVU7JPSdIUXBlU1aH2fhR4kN6Y/5E2/EN7P9pWPwRc0tl8SauNV5ckTYO+wiDJa5P88ug0cA2wH9gGjD4RtA54qE1vA25sTxVdCfyoDSc9DFyTZEG7cXxNq0mSpkG/w0QLgQeTjO7rK1X1V0l2A/cluQl4Hnh/W387cD0wBPwY+CBAVR1L8sfA7rbep6rqWJ9tkyRNUF9hUFXPAL8+Rv37wNVj1Au4eZx9bQY299MeSdLk+AlkSZJhIEkyDCRJGAaSJAwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAksTU/A1kne02bDg7jiFpXF4ZSJIMA0mSYSBJwjCQJNFHGCS5JMm3kzye5ECSj7X6hiSHkuxtr+s729yWZCjJU0mu7dRXtdpQklv765Ik6XT18zTRceAPq+q7SX4Z2JNkR1v2uar6THflJJcCa4C3AK8HvpXkTW3x54F3A8PA7iTbqurxPtomSToNkw6DqjoMHG7Tf5/kCWDxSTZZDWytqpeBZ5MMAVe0ZUNV9QxAkq1tXcNAkqbJlNwzSLIUuBx4tJVuSbIvyeYkC1ptMfBCZ7PhVhuvPtZx1icZTDI4MjIyFU2XJDEFYZDkl4AHgI9X1UvARuCNwAp6Vw6f7fcYo6pqU1WtrKqVAwMDU7VbSTrn9fUJ5CS/QC8IvlxVXwOoqiOd5V8Evt5mDwGXdDZf0mqcpC5Jmgb9PE0U4C7giar6s059UWe19wH72/Q2YE2S85MsA5YD3wF2A8uTLEtyHr2bzNsm2y5J0unr58rgN4DfAx5LsrfV/jOwNskKoIDngA8DVNWBJPfRuzF8HLi5qn4GkOQW4GFgHrC5qg700S5J0mnq52mi/w1kjEXbT7LNHcAdY9S3n2w7SdKZ5SeQJUmGgSTJMJAkYRhIkjhX/9KZf1VLkl7BKwNJkmEgSTIMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJHGufgJZs890fSrcT59LY/LKQJJkGEiSDANJEoaBJAnDQJLELAqDJKuSPJVkKMmtM90eSTqXzIpHS5PMAz4PvBsYBnYn2VZVj89sy3TWmY5HS318VXPQrAgD4ApgqKqeAUiyFVgNGAaae/zMhOag2RIGi4EXOvPDwNtPXCnJemB9m/2/SZ46zeNcDPzdpFo4u5wt/QD7Mnmf/OSZ2rPnZHaaqr78q7GKsyUMJqSqNgGbJrt9ksGqWjmFTZoRZ0s/wL7MRmdLP8C+nI7ZcgP5EHBJZ35Jq0mSpsFsCYPdwPIky5KcB6wBts1wmyTpnDErhomq6niSW4CHgXnA5qo6cAYONekhplnmbOkH2JfZ6GzpB9iXCUtVncn9S5LmgNkyTCRJmkGGgSTp3AiDufxVF0kuSfLtJI8nOZDkY61+UZIdSQ629wUz3daJSDIvyfeSfL3NL0vyaDs3X20PEMx6SS5Mcn+SJ5M8keTfzOFz8h/bz9b+JPcm+edz5bwk2ZzkaJL9ndqY5yE9d7Y+7Uvy1plr+SuN048/bT9f+5I8mOTCzrLbWj+eSnLtVLThrA+DzlddXAdcCqxNcunMtuq0HAf+sKouBa4Ebm7tvxXYWVXLgZ1tfi74GPBEZ/5PgM9V1a8CPwBumpFWnb7/DvxVVf1r4Nfp9WnOnZMki4H/AKysqsvoPcCxhrlzXu4GVp1QG+88XAcsb6/1wMZpauNE3M2r+7EDuKyqfg34G+A2gPbvfw3wlrbNn7ffc30568OAzlddVNVPgNGvupgTqupwVX23Tf89vV86i+n1YUtbbQtww8y0cOKSLAH+HfClNh/gXcD9bZW50o9/AbwTuAugqn5SVT9kDp6TZj7wmiTzgV8EDjNHzktVPQIcO6E83nlYDdxTPbuAC5Msmp6WntxY/aiqb1bV8Ta7i97nr6DXj61V9XJVPQsM0fs915dzIQzG+qqLxTPUlr4kWQpcDjwKLKyqw23Ri8DCGWrW6fhvwB8BP2/zrwN+2PmBnyvnZhkwAvyPNuT1pSSvZQ6ek6o6BHwG+Ft6IfAjYA9z87yMGu88zOXfBb8PfKNNn5F+nAthcFZI8kvAA8DHq+ql7rLqPR88q58RTvIe4GhV7ZnptkyB+cBbgY1VdTnwD5wwJDQXzglAG09fTS/gXg+8llcPV8xZc+U8nEyST9AbLv7ymTzOuRAGc/6rLpL8Ar0g+HJVfa2Vj4xe4rb3ozPVvgn6DeC9SZ6jN1T3Lnrj7he24QmYO+dmGBiuqkfb/P30wmGunROA3wKeraqRqvop8DV652ounpdR452HOfe7IMkHgPcAv1v/9KGwM9KPcyEM5vRXXbRx9buAJ6rqzzqLtgHr2vQ64KHpbtvpqKrbqmpJVS2ldw7+V1X9LvBt4HfaarO+HwBV9SLwQpI3t9LV9L5ufU6dk+ZvgSuT/GL7WRvty5w7Lx3jnYdtwI3tqaIrgR91hpNmnSSr6A2rvreqftxZtA1Yk+T8JMvo3RD/Tt8HrKqz/gVcT+9u/NPAJ2a6PafZ9t+kd5m7D9jbXtfTG2/fCRwEvgVcNNNtPY0+XQV8vU2/of0gDwF/AZw/0+2bYB9WAIPtvPwlsGCunhPgk8CTwH7gfwLnz5XzAtxL717HT+ldsd003nkAQu/JwqeBx+g9QTXjfThJP4bo3RsY/Xf/hc76n2j9eAq4bira4NdRSJLOiWEiSdIpGAaSJMNAkmQYSJIwDCRJGAaSJAwDSRLw/wCd7ty9wdaP0AAAAABJRU5ErkJggg==",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "### Question lengths\n",
        "\n",
        "bin_size = 10\n",
        "bins = range(1, max(question_lengths) + bin_size, bin_size)\n",
        "plt.hist(question_lengths,\n",
        "        bins,\n",
        "        density=False,\n",
        "        facecolor='r',\n",
        "        alpha=0.5)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KTgMc7xXbC04"
      },
      "outputs": [],
      "source": [
        "bins = range(0, max(question_lengths) + bin_size, bin_size)\n",
        "df = pd.DataFrame({'question_lengths': question_lengths})\n",
        "df['bin'] = pd.cut(df.question_lengths, bins)\n",
        "df.bin.value_counts().sort_index()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "id": "EQl_S6-VEtDX",
        "outputId": "598e87d1-c17c-4358-d68b-c984da1d0d2c"
      },
      "outputs": [
        {
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYMAAAD4CAYAAAAO9oqkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAU+klEQVR4nO3df5Bd9Xnf8ffHkrGNbSxAG0okEalBpgGmdvAG1HHtOiYRgmYs2toeGLeojiZqa5E4bTo2JJ0gGzNj2rQ0tDYexVIsMgRQiV3UBBurGId/jGAxGBA/wpof1moAyZaApkwggqd/3K/qa3lX0t67u3dB79fMnT3nOd9zznO1d/XZe865e1JVSJKObG8YdAOSpMEzDCRJhoEkyTCQJGEYSJKAuYNuoFfz58+vxYsXD7oNSXpNueeee35YVUMH1l+zYbB48WJGRkYG3YYkvaYkeWq8uoeJJEmGgSTJMJAkYRhIkjiMMEiyMcmuJA8eUP/NJI8k2Z7kP3bVL00ymuTRJOd01Ve02miSS7rqS5Jsa/Ubkxw1VU9OknR4DuedwVeAFd2FJL8MrATeVVWnAX/Q6qcCFwCntXW+mGROkjnAF4BzgVOBC9tYgCuBq6rqZGAvsLrfJyVJmpxDhkFV3QHsOaD8b4DPV9VLbcyuVl8J3FBVL1XVE8AocGZ7jFbV41X1MnADsDJJgA8CN7X1NwHn9/mcJEmT1Os5g3cC72uHd/4yyS+1+gJgR9e4sVabqH488FxV7TugLkmaQb1+6GwucBywDPglYHOSvztlXU0gyRpgDcBJJ5003buTpCNGr2EwBny1OnfGuSvJq8B8YCewqGvcwlZjgvqPgHlJ5rZ3B93jf0pVrQfWAwwPD/d8V55163pdc2rNlj4kqdfDRP8T+GWAJO8EjgJ+CGwBLkjypiRLgKXAXcDdwNJ25dBRdE4yb2lhcjvw4bbdVcDNvT4ZSVJvDvnOIMn1wAeA+UnGgMuAjcDGdrnpy8Cq9h/79iSbgYeAfcDaqnqlbedi4FZgDrCxqra3XXwauCHJ54B7gQ1T+PwkSYfhkGFQVRdOsOifTzD+CuCKceq3ALeMU3+cztVGkqQB8RPIkiTDQJJkGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJ4jDCIMnGJLvaLS4PXPY7SSrJ/DafJFcnGU1yf5IzusauSvJYe6zqqr8nyQNtnauTZKqenCTp8BzOO4OvACsOLCZZBCwHftBVPhdY2h5rgGva2OPo3Dv5LDq3uLwsybFtnWuA3+ha76f2JUmaXocMg6q6A9gzzqKrgE8B1VVbCVxbHXcC85KcCJwDbK2qPVW1F9gKrGjLjqmqO6uqgGuB8/t7SpKkyerpnEGSlcDOqvreAYsWADu65sda7WD1sXHqkqQZNHeyKyQ5GvhdOoeIZlSSNXQOP3HSSSfN9O4l6XWrl3cGPw8sAb6X5ElgIfDdJH8H2Aks6hq7sNUOVl84Tn1cVbW+qoaranhoaKiH1iVJ45l0GFTVA1X1M1W1uKoW0zm0c0ZVPQNsAS5qVxUtA56vqqeBW4HlSY5tJ46XA7e2ZS8kWdauIroIuHmKnpsk6TAdzqWl1wPfAU5JMpZk9UGG3wI8DowCfwR8AqCq9gCXA3e3x2dbjTbmy22d7wNf7+2pSJJ6dchzBlV14SGWL+6aLmDtBOM2AhvHqY8Apx+qD0nS9PETyJIkw0CSZBhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJAwDSRKGgSQJw0CSxGHc3EbTZ926QXfQMVv6kDQ4h3Pby41JdiV5sKv2n5I8kuT+JF9LMq9r2aVJRpM8muScrvqKVhtNcklXfUmSba1+Y5KjpvIJSpIO7XAOE30FWHFAbStwelX9feCvgEsBkpwKXACc1tb5YpI5SeYAXwDOBU4FLmxjAa4Erqqqk4G9wMHusSxJmgaHDIOqugPYc0Dtm1W1r83eCSxs0yuBG6rqpap6gs5N7s9sj9GqeryqXgZuAFYmCfBB4Ka2/ibg/D6fkyRpkqbiBPKvA19v0wuAHV3LxlptovrxwHNdwbK/LkmaQX2FQZLfA/YB101NO4fc35okI0lGdu/ePRO7lKQjQs9hkORfAr8GfKyqqpV3Aou6hi1stYnqPwLmJZl7QH1cVbW+qoaranhoaKjX1iVJB+gpDJKsAD4FfKiqXuxatAW4IMmbkiwBlgJ3AXcDS9uVQ0fROcm8pYXI7cCH2/qrgJt7eyqSpF4dzqWl1wPfAU5JMpZkNfDfgbcDW5Pcl+RLAFW1HdgMPAR8A1hbVa+0cwIXA7cCDwOb21iATwP/LskonXMIG6b0GUqSDumQHzqrqgvHKU/4H3ZVXQFcMU79FuCWceqP07naSJI0IP45CkmSYSBJMgwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJHF4t73cmGRXkge7ascl2Zrksfb12FZPkquTjCa5P8kZXeusauMfS7Kqq/6eJA+0da5Okql+kpKkgzucdwZfAVYcULsEuK2qlgK3tXmAc4Gl7bEGuAY64QFcBpxF5xaXl+0PkDbmN7rWO3BfkqRpdsgwqKo7gD0HlFcCm9r0JuD8rvq11XEnMC/JicA5wNaq2lNVe4GtwIq27JiqurOqCri2a1uSpBnS6zmDE6rq6Tb9DHBCm14A7OgaN9ZqB6uPjVOXJM2gvk8gt9/oawp6OaQka5KMJBnZvXv3TOxSko4IvYbBs+0QD+3rrlbfCSzqGrew1Q5WXzhOfVxVtb6qhqtqeGhoqMfWJUkH6jUMtgD7rwhaBdzcVb+oXVW0DHi+HU66FVie5Nh24ng5cGtb9kKSZe0qoou6tiVJmiFzDzUgyfXAB4D5ScboXBX0eWBzktXAU8BH2/BbgPOAUeBF4OMAVbUnyeXA3W3cZ6tq/0npT9C5YuktwNfbQ5I0gw4ZBlV14QSLzh5nbAFrJ9jORmDjOPUR4PRD9SFJmj5+AlmSZBhIkgwDSRKGgSQJw0CShGEgScIwkCRhGEiSMAwkSRgGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEn0GQZJ/m2S7UkeTHJ9kjcnWZJkW5LRJDcmOaqNfVObH23LF3dt59JWfzTJOf09JUnSZPUcBkkWAL8FDFfV6cAc4ALgSuCqqjoZ2AusbqusBva2+lVtHElObeudBqwAvphkTq99SZImr9/DRHOBtySZCxwNPA18ELipLd8EnN+mV7Z52vKzk6TVb6iql6rqCWAUOLPPviRJk9BzGFTVTuAPgB/QCYHngXuA56pqXxs2Bixo0wuAHW3dfW388d31cdaRJM2Afg4THUvnt/olwM8Cb6VzmGfaJFmTZCTJyO7du6dzV5J0ROnnMNGvAE9U1e6q+lvgq8B7gXntsBHAQmBnm94JLAJoy98B/Ki7Ps46P6Gq1lfVcFUNDw0N9dG6JKlbP2HwA2BZkqPbsf+zgYeA24EPtzGrgJvb9JY2T1v+raqqVr+gXW20BFgK3NVHX5KkSZp76CHjq6ptSW4CvgvsA+4F1gN/AdyQ5HOttqGtsgH4kySjwB46VxBRVduTbKYTJPuAtVX1Sq99SZImr+cwAKiqy4DLDig/zjhXA1XV3wAfmWA7VwBX9NOLJKl3fgJZkmQYSJIMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJPoMgyTzktyU5JEkDyf5B0mOS7I1yWPt67FtbJJcnWQ0yf1Jzujazqo2/rEkq/p9UpKkyen3ncEfAt+oqr8HvAt4GLgEuK2qlgK3tXmAc4Gl7bEGuAYgyXF07qN8Fp17J1+2P0AkSTOj5zBI8g7g/cAGgKp6uaqeA1YCm9qwTcD5bXolcG113AnMS3IicA6wtar2VNVeYCuwote+JEmT1887gyXAbuCPk9yb5MtJ3gqcUFVPtzHPACe06QXAjq71x1ptovpPSbImyUiSkd27d/fRuiSpWz9hMBc4A7imqn4R+L/8+JAQAFVVQPWxj59QVeurariqhoeGhqZqs5J0xOsnDMaAsara1uZvohMOz7bDP7Svu9ryncCirvUXttpEdUnSDOk5DKrqGWBHklNa6WzgIWALsP+KoFXAzW16C3BRu6poGfB8O5x0K7A8ybHtxPHyVpMkzZC5fa7/m8B1SY4CHgc+TidgNidZDTwFfLSNvQU4DxgFXmxjqao9SS4H7m7jPltVe/rsS5I0CX2FQVXdBwyPs+jsccYWsHaC7WwENvbTiySpd34CWZJkGEiSDANJEoaBJAnDQJKEYSBJwjCQJGEYSJLo/xPIeh1Yt27QHcyOHqQjme8MJEmGgSTJMJAkYRhIkjAMJEkYBpIkDANJEoaBJIkpCIMkc5Lcm+TP2/ySJNuSjCa5sd0SkyRvavOjbfnirm1c2uqPJjmn354kSZMzFe8MPgk83DV/JXBVVZ0M7AVWt/pqYG+rX9XGkeRU4ALgNGAF8MUkc6agL0nSYeorDJIsBP4x8OU2H+CDwE1tyCbg/Da9ss3Tlp/dxq8Ebqiql6rqCWAUOLOfviRJk9PvO4P/CnwKeLXNHw88V1X72vwYsKBNLwB2ALTlz7fx/78+zjo/IcmaJCNJRnbv3t1n65Kk/XoOgyS/BuyqqnumsJ+Dqqr1VTVcVcNDQ0MztVtJet3r56+Wvhf4UJLzgDcDxwB/CMxLMrf99r8Q2NnG7wQWAWNJ5gLvAH7UVd+vex1J0gzo+Z1BVV1aVQurajGdE8DfqqqPAbcDH27DVgE3t+ktbZ62/FtVVa1+QbvaaAmwFLir174kSZM3Hfcz+DRwQ5LPAfcCG1p9A/AnSUaBPXQChKranmQz8BCwD1hbVa9MQ1+SpAlMSRhU1beBb7fpxxnnaqCq+hvgIxOsfwVwxVT0IkmaPD+BLEkyDCRJhoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJAnDQJKEYSBJwjCQJGEYSJIwDCRJGAaSJPoIgySLktye5KEk25N8stWPS7I1yWPt67GtniRXJxlNcn+SM7q2taqNfyzJqon2KUmaHv28M9gH/E5VnQosA9YmORW4BLitqpYCt7V5gHPp3Ox+KbAGuAY64QFcBpxF53aZl+0PEEnSzOg5DKrq6ar6bpv+P8DDwAJgJbCpDdsEnN+mVwLXVsedwLwkJwLnAFurak9V7QW2Ait67UuSNHlTcs4gyWLgF4FtwAlV9XRb9AxwQpteAOzoWm2s1Saqj7efNUlGkozs3r17KlqXJDEFYZDkbcCfAb9dVS90L6uqAqrffXRtb31VDVfV8NDQ0FRtVpKOeH2FQZI30gmC66rqq638bDv8Q/u6q9V3Aou6Vl/YahPVJUkzpJ+riQJsAB6uqv/StWgLsP+KoFXAzV31i9pVRcuA59vhpFuB5UmObSeOl7eaJGmGzO1j3fcC/wJ4IMl9rfa7wOeBzUlWA08BH23LbgHOA0aBF4GPA1TVniSXA3e3cZ+tqj199CVJmqR0Duu/9gwPD9fIyEhP665bN7W96PXD14Ze75LcU1XDB9b9BLIkyTCQJBkGkiQMA0kShoEkCcNAkoRhIEnCMJAkYRhIkjAMJEkYBpIkDANJEoaBJIn+/oS19LozW/5q6WzpQ0cO3xlIkgwDSZJhIEliFoVBkhVJHk0ymuSSQfcjSUeSWXECOckc4AvArwJjwN1JtlTVQ4PtTBqM2XICebb0oek3K8IAOBMYrarHAZLcAKwEDANpgAyD2We6viezJQwWADu65seAsw4clGQNsKbN/nWSR3vc33zghz2uO5XsY3b1APZxIPuYXT3wmc/03cfPjVecLWFwWKpqPbC+3+0kGamq4SloyT5eRz3Yh33M9h6ms4/ZcgJ5J7Coa35hq0mSZsBsCYO7gaVJliQ5CrgA2DLgniTpiDErDhNV1b4kFwO3AnOAjVW1fRp32fehpiliHz82G3oA+ziQffzYbOgBpqmPVNV0bFeS9BoyWw4TSZIGyDCQJB1ZYZBkY5JdSR4cYA9vTnJXku8l2Z7kMwPs5ckkDyS5L8nIgHo4pe1//+OFJL89oF4+meTB9n2ZsR7Ge10m+Ujr49UkM3I54wR9XJ7k/va9+WaSnx1AD+uS7Ox6jZw3nT0cpI8bu3p4Msl9A+rjXUm+0352/1eSY6ZkZ1V1xDyA9wNnAA8OsIcAb2vTbwS2AcsG1MuTwPxBf1+6+pkDPAP83AD2fTrwIHA0nQsr/jdw8gzt+6del8AvAKcA3waGB9jHMV3TvwV8aQA9rAP+/Qy/Hg76fwXwn4HfH9D35G7gH7XpXwcun4p9HVHvDKrqDmDPgHuoqvrrNvvG9vAsfsfZwPer6qkB7PsXgG1V9WJV7QP+EvinM7Hj8V6XVfVwVfX6Cfup7OOFrtm3Ms2v1dnwM3qoPpIE+Chw/YD6eCdwR5veCvyzqdjXERUGs0WSOe0t5i5ga1VtG1ArBXwzyT3tT30M2gXMwA/YBB4E3pfk+CRHA+fxkx+EPGIluSLJDuBjwO8PqI2L2+GqjUmOHVAP+70PeLaqHhvQ/rfT+dttAB9hil6nhsEAVNUrVfVuOp+0PjPJ6QNq5R9W1RnAucDaJO8fUB+0Dxt+CPgfg9h/VT0MXAl8E/gGcB/wyiB6mW2q6veqahFwHXDxAFq4Bvh54N3A03QO0QzShQzulxboHBr6RJJ7gLcDL0/FRg2DAaqq54DbgRUD2v/O9nUX8DU6fz12UM4FvltVzw6qgaraUFXvqar3A3uBvxpUL7PUdUzRIYnJqKpn2y9QrwJ/xABfp0nm0jl8eOOgeqiqR6pqeVW9h04ofX8qtmsYzLAkQ0nmtem30LmHwyMD6OOtSd6+fxpYTudQyaAM+rctkvxM+3oSnR/4Px1kP7NBkqVdsysZzGv1xK7Zf8JgX6e/AjxSVWODaqDrdfoG4D8AX5qK7c6KP0cxU5JcD3wAmJ9kDLisqjbMcBsnApvaDX3eAGyuqj+f4R4ATgC+1jkXxlzgT6vqGwPoY38Y/Srwrwax/y5/luR44G+Bte2d27Qb73VJ56ThfwOGgL9Icl9VnTOAPs5LcgrwKvAU8K8H0MMHkrybzjmuJ5mB18lB/q+Y0fNaE/x7vC3J2jbkq8AfT8m+2uVJkqQjmIeJJEmGgSTJMJAkYRhIkjAMJEkYBpIkDANJEvD/AG73aYP0YSYyAAAAAElFTkSuQmCC",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "needs_background": "light"
          },
          "output_type": "display_data"
        }
      ],
      "source": [
        "### Answer lengths\n",
        "\n",
        "bin_size = 2\n",
        "bins = range(1, max(answer_lengths) + bin_size, bin_size)\n",
        "plt.hist([a for a in answer_lengths if a <= 20],\n",
        "         bins[:10],\n",
        "         density=False,\n",
        "         facecolor='b',\n",
        "         alpha=0.5)\n",
        "plt.xticks(bins[:10])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "psu7iHwf0N38"
      },
      "outputs": [],
      "source": [
        "bins = range(0, max(answer_lengths) + bin_size, bin_size)\n",
        "\n",
        "df = pd.DataFrame({'answer_lengths': answer_lengths})\n",
        "df['bin'] = pd.cut(df.answer_lengths, bins)\n",
        "df.bin.value_counts()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jcOaM1HQHo4e"
      },
      "source": [
        "### Performance by question/answer types and answer_lengths"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lp9TACgPMYZG"
      },
      "outputs": [],
      "source": [
        "dev_answer_lengths = []\n",
        "for datum in jaquad['validation']:\n",
        "    answer = tokenizer(datum['answers.text'])['input_ids']\n",
        "    dev_answer_lengths.append(len(answer) - 2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LczWxMSnH02f"
      },
      "outputs": [],
      "source": [
        "df = pd.DataFrame({\n",
        "    'question_type': jaquad['validation']['question_type'],\n",
        "    'answer_type': [typ[0] for typ in jaquad['validation']['answers.answer_type']],\n",
        "    'answer_len': dev_answer_lengths,\n",
        "    'f1': f1_scores,\n",
        "    'em': em,\n",
        "})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HsJ2gKRoH6Bq"
      },
      "outputs": [],
      "source": [
        "df[['question_type', 'f1', 'em']].groupby(by='question_type').mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5_nM9S4TLt3K"
      },
      "outputs": [],
      "source": [
        "df[['answer_type', 'f1', 'em']].groupby(by='answer_type').mean()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZxlfAcEGH8eq"
      },
      "outputs": [],
      "source": [
        "bins = 2\n",
        "\n",
        "df['answer_len_bin'] = (df['answer_len']-1) // bins * bins\n",
        "df['answer_len_bin'][df.answer_len_bin >= 8] = 8\n",
        "df['answer_len_bin'] = df['answer_len_bin'].apply(lambda x: f'{x+1:02d}-{x+bins:02d}')\n",
        "df[['answer_len_bin', 'f1', 'em']].groupby(by='answer_len_bin').mean()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "JaQuAD.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}