notebooks/Question-answering-with-BERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Question-answering with BERT\n",
    "\n",
    "In Question-Answering tasks, the model receives a text and a question regarding to the text, and it should mark the beginning and end of the answer in the text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare and import modules\n",
    "\n",
    "We need to install a few modules:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install torchvision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip uninstall -y transformers\n",
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And, then, we can import all the necessary modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import string\n",
    "import numpy as np\n",
    "import pickle\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras as K\n",
    "\n",
    "from tokenizers import BertWordPieceTokenizer\n",
    "from transformers import BertTokenizer, TFBertModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of replicas: 1\n"
     ]
    }
   ],
   "source": [
    "# Detect hardware, return appropriate distribution strategy.\n",
    "# You can see that it is pretty easy to set up.\n",
    "try:\n",
    "    # TPU detection: no parameters necessary if TPU_NAME environment\n",
    "    # variable is set (always set in Kaggle)\n",
    "    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
    "    tf.config.experimental_connect_to_cluster(tpu)\n",
    "    tf.tpu.experimental.initialize_tpu_system(tpu)\n",
    "    strategy = tf.distribute.experimental.TPUStrategy(tpu)\n",
    "    print('Running on TPU ', tpu.master())\n",
    "except ValueError:\n",
    "    # Default distribution strategy in Tensorflow. Works on CPU and single GPU.\n",
    "    strategy = tf.distribute.get_strategy()\n",
    "\n",
    "print('Number of replicas:', strategy.num_replicas_in_sync)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Constants\n",
    "\n",
    "**Dataset.** We will use XQuAD (Cross-lingual Question Answering Dataset), which consists of a subset of 240 paragraphs and 1190 question-answer pairs from the development set of SQuAD v1.1 together with their professional translations into 10 languages. See: https://github.com/deepmind/xquad\n",
    "\n",
    "**Model name.** This is the model name that we need to provide to HuggingFace to import the pretrained BERT model and tokenizer. We will import BERT Multilingual Base model which has been pretrained on the top 104 languages with the largest Wikipedia using a masked language modeling (MLM) objective. See: https://huggingface.co/bert-base-multilingual-cased"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATASET_URL = 'https://raw.githubusercontent.com/deepmind/xquad/99910ec0f10151652f6726282ca922dd8eb0207a/xquad.es.json'\n",
    "MODEL_NAME = 'bert-base-multilingual-cased'\n",
    "MAX_LEN = 384\n",
    "\n",
    "# Set language name to save model\n",
    "LANGUAGE = 'spanish'\n",
    "\n",
    "# Depends on whether we are using TPUs or not, increase BATCH_SIZE\n",
    "BATCH_SIZE = 8 * strategy.num_replicas_in_sync\n",
    "\n",
    "# Detect environment\n",
    "if os.environ.get('KAGGLE_KERNEL_RUN_TYPE',''):\n",
    "    print('Detected Kaggle environment')\n",
    "    ARTIFACTS_PATH = 'artifacts/'\n",
    "else:\n",
    "    ARTIFACTS_PATH = '../artifacts/'\n",
    "    \n",
    "if not os.path.exists(ARTIFACTS_PATH):\n",
    "    os.makedirs(ARTIFACTS_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import pretrained tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import tokenizer from HuggingFace\n",
    "slow_tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)\n",
    "\n",
    "save_path = '%s%s-%s/' % (ARTIFACTS_PATH, LANGUAGE, MODEL_NAME)\n",
    "if not os.path.exists(save_path):\n",
    "    os.makedirs(save_path)\n",
    "\n",
    "slow_tokenizer.save_pretrained(save_path)\n",
    "\n",
    "# You can already use the Slow Tokenizer, but its implementation\n",
    "# in Rust is much faster.\n",
    "tokenizer = BertWordPieceTokenizer('%s/vocab.txt' % save_path, lowercase=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize & encode dataset\n",
    "\n",
    "<u>Skip this step if you already have a fine-tunned model.</u>\n",
    "\n",
    "BERT expects 3 different inputs:\n",
    "- Token IDs: the tokens transformed into numbers.\n",
    "- Mask: sequence of `0` (if there are PAD tokens in that position) and `1` (otherwise).\n",
    "- Segments or Type IDs: sequence of `0` and `1` to distinguish between the first and the second sentence in NSP tasks. In this notebook, we do not need this input, so it will be always `0`.\n",
    "\n",
    "For example:\n",
    "```\n",
    "Text:       Is this jacksonville? Yes, it is.\n",
    "---------------------------------------------------------------------------------\n",
    "Tokens:     [CLS] Is    this  ja    ##cks ##on  ##ville ?   [SEP] Yes   ,    it   is   .    [SEP] [PAD] [PAD] ...\n",
    "Token IDs:  101   12034 10531 10201 18676 10263 12043   136 102   2160  117  1122 1110 119  102   100   100   ...\n",
    "Mask:       1     1     1     1     1     1     1       1   1     1     1    1    1    1    1     0     0     ...\n",
    "Type IDs:   0     0     0     0     0     0     0       0   0     1     1    1    1    1    1     0     0     ...\n",
    "```\n",
    "\n",
    "Notes:\n",
    "- Token IDs may be different depending on the tokenizer.\n",
    "- I have manually introduced `[SEP]` between both sentences.\n",
    "\n",
    "For further details, see BERT paper: https://arxiv.org/pdf/1810.04805.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This code is a modified version from https://keras.io/examples/nlp/text_extraction_with_bert/\n",
    "class SquadExample:\n",
    "    def __init__(\n",
    "        self,\n",
    "        question,\n",
    "        context,\n",
    "        start_char_idx,\n",
    "        answer_text,\n",
    "        all_answers,\n",
    "        tokenizer\n",
    "    ):\n",
    "        self.question = question\n",
    "        self.context = context\n",
    "        self.start_char_idx = start_char_idx\n",
    "        self.answer_text = answer_text\n",
    "        self.all_answers = all_answers\n",
    "        self.tokenizer = tokenizer\n",
    "        self.skip = False\n",
    "\n",
    "    def preprocess(self):\n",
    "        context = self.context\n",
    "        question = self.question\n",
    "        answer_text = self.answer_text\n",
    "        start_char_idx = self.start_char_idx\n",
    "\n",
    "        # Fix white spaces\n",
    "        context = re.sub(r\"\\s+\", ' ', context).strip()\n",
    "        question = re.sub(r\"\\s+\", ' ', question).strip()\n",
    "        answer = re.sub(r\"\\s+\", ' ', answer_text).strip()\n",
    "\n",
    "        # Find end token index of answer in context\n",
    "        end_char_idx = start_char_idx + len(answer)\n",
    "        if end_char_idx >= len(context):\n",
    "            self.skip = True\n",
    "            return\n",
    "\n",
    "        # Mark the character indexes in context that are in answer\n",
    "        is_char_in_ans = [0] * len(context)\n",
    "        for idx in range(start_char_idx, end_char_idx):\n",
    "            is_char_in_ans[idx] = 1\n",
    "\n",
    "        # Encode context (token IDs, mask and token types)\n",
    "        tokenized_context = tokenizer.encode(context)\n",
    "\n",
    "        # Find tokens that were created from answer characters\n",
    "        ans_token_idx = []\n",
    "        for idx, (start, end) in enumerate(tokenized_context.offsets):\n",
    "            if sum(is_char_in_ans[start:end]) > 0:\n",
    "                ans_token_idx.append(idx)\n",
    "\n",
    "        if len(ans_token_idx) == 0:\n",
    "            self.skip = True\n",
    "            return\n",
    "\n",
    "        # Find start and end token index for tokens from answer\n",
    "        start_token_idx = ans_token_idx[0]\n",
    "        end_token_idx = ans_token_idx[-1]\n",
    "\n",
    "        # Encode question (token IDs, mask and token types)\n",
    "        tokenized_question = tokenizer.encode(question)\n",
    "\n",
    "        # Create inputs\n",
    "        input_ids = tokenized_context.ids + tokenized_question.ids[1:]\n",
    "        token_type_ids = [0] * len(tokenized_context.ids) + [1] * len(\n",
    "            tokenized_question.ids[1:]\n",
    "        )\n",
    "        attention_mask = [1] * len(input_ids)\n",
    "\n",
    "        # Pad and create attention masks.\n",
    "        # Skip if truncation is needed\n",
    "        padding_length = MAX_LEN - len(input_ids)\n",
    "        if padding_length > 0:  # pad\n",
    "            input_ids = input_ids + ([0] * padding_length)\n",
    "            attention_mask = attention_mask + ([0] * padding_length)\n",
    "            token_type_ids = token_type_ids + ([0] * padding_length)\n",
    "        elif padding_length < 0:  # skip\n",
    "            self.skip = True\n",
    "            return\n",
    "\n",
    "        self.input_ids = input_ids\n",
    "        self.token_type_ids = token_type_ids\n",
    "        self.attention_mask = attention_mask\n",
    "        self.start_token_idx = start_token_idx\n",
    "        self.end_token_idx = end_token_idx\n",
    "        self.context_token_to_char = tokenized_context.offsets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_squad_examples(raw_data, tokenizer):\n",
    "    squad_examples = []\n",
    "    for item in raw_data[\"data\"]:\n",
    "        for para in item[\"paragraphs\"]:\n",
    "            context = para[\"context\"]\n",
    "            for qa in para[\"qas\"]:\n",
    "                question = qa[\"question\"]\n",
    "                answer_text = qa[\"answers\"][0][\"text\"]\n",
    "                all_answers = [_[\"text\"] for _ in qa[\"answers\"]]\n",
    "                start_char_idx = qa[\"answers\"][0][\"answer_start\"]\n",
    "                squad_eg = SquadExample(\n",
    "                    question,\n",
    "                    context,\n",
    "                    start_char_idx,\n",
    "                    answer_text,\n",
    "                    all_answers,\n",
    "                    tokenizer\n",
    "                )\n",
    "                squad_eg.preprocess()\n",
    "                squad_examples.append(squad_eg)\n",
    "    return squad_examples\n",
    "\n",
    "\n",
    "def create_inputs_targets(squad_examples):\n",
    "    dataset_dict = {\n",
    "        \"input_ids\": [],\n",
    "        \"token_type_ids\": [],\n",
    "        \"attention_mask\": [],\n",
    "        \"start_token_idx\": [],\n",
    "        \"end_token_idx\": [],\n",
    "    }\n",
    "    for item in squad_examples:\n",
    "        if item.skip == False:\n",
    "            for key in dataset_dict:\n",
    "                dataset_dict[key].append(getattr(item, key))\n",
    "    for key in dataset_dict:\n",
    "        dataset_dict[key] = np.array(dataset_dict[key])\n",
    "\n",
    "    x = [\n",
    "        dataset_dict[\"input_ids\"],\n",
    "        dataset_dict[\"token_type_ids\"],\n",
    "        dataset_dict[\"attention_mask\"],\n",
    "    ]\n",
    "    y = [dataset_dict[\"start_token_idx\"], dataset_dict[\"end_token_idx\"]]\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = K.utils.get_file('dataset.json', DATASET_URL)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(dataset_path) as fp:\n",
    "    raw_data = json.load(fp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split into train and test sets\n",
    "raw_train_data = {}\n",
    "raw_eval_data = {}\n",
    "raw_train_data['data'], raw_eval_data['data'] = np.split(np.asarray(raw_data['data']), [int(.8*len(raw_data['data']))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_squad_examples = create_squad_examples(raw_train_data, tokenizer)\n",
    "x_train, y_train = create_inputs_targets(train_squad_examples)\n",
    "print(f\"{len(train_squad_examples)} training points created.\")\n",
    "\n",
    "eval_squad_examples = create_squad_examples(raw_eval_data, tokenizer)\n",
    "x_eval, y_eval = create_inputs_targets(eval_squad_examples)\n",
    "print(f\"{len(eval_squad_examples)} evaluation points created.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build model\n",
    "\n",
    "As we explained before, we import a pretrained BERT model from HuggingFace and attach the necessary inputs to it. Then, we add two fully-connected layer to obtain two outputs: the initial token and the ending token.\n",
    "\n",
    "![images/BERT-QA-arch.png](images/BERT-QA-arch.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_model():\n",
    "    input_ids = tf.keras.layers.Input(shape=(MAX_LEN,), name='input_ids', dtype=tf.int32)\n",
    "    token_type_ids = tf.keras.layers.Input(shape=(MAX_LEN,), name='token_type_ids', dtype=tf.int32)\n",
    "    attention_mask = tf.keras.layers.Input(shape=(MAX_LEN,), name='attention_mask', dtype=tf.int32)\n",
    "    \n",
    "    encoder = TFBertModel.from_pretrained(MODEL_NAME)\n",
    "    x = encoder(input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)\n",
    "    \n",
    "    # Huggingface transformers have multiple outputs, embeddings are the first one,\n",
    "    # so let's slice out the first position.\n",
    "    x = x[0]\n",
    "\n",
    "    # Define two outputs\n",
    "    start_logits = tf.keras.layers.Dense(1, name='start_logit', use_bias=False)(x)\n",
    "    start_logits = tf.keras.layers.Flatten()(start_logits)\n",
    "\n",
    "    end_logits = tf.keras.layers.Dense(1, name='end_logit', use_bias=False)(x)\n",
    "    end_logits = tf.keras.layers.Flatten()(end_logits)\n",
    "\n",
    "    # Normalize outputs with softmax\n",
    "    start_probs = tf.keras.layers.Activation(tf.keras.activations.softmax, name='start_probs')(start_logits)\n",
    "    end_probs = tf.keras.layers.Activation(tf.keras.activations.softmax, name='end_probs')(end_logits)\n",
    "\n",
    "    model = tf.keras.Model(\n",
    "        inputs=[input_ids, token_type_ids, attention_mask],\n",
    "        outputs=[start_probs, end_probs],\n",
    "    )\n",
    "    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False)\n",
    "    optimizer = tf.keras.optimizers.Adam(lr=2e-5)\n",
    "    model.compile(optimizer=optimizer, loss=[loss, loss])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing TFBertModel: ['nsp___cls', 'mlm___cls']\n",
      "- This IS expected if you are initializing TFBertModel 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 TFBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "All the weights of TFBertModel were initialized from the model checkpoint at bert-base-multilingual-cased.\n",
      "If your task is similar to the task the model of the ckeckpoint was trained on, you can already use TFBertModel for predictions without further training.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"functional_1\"\n",
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_ids (InputLayer)          [(None, 384)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "attention_mask (InputLayer)     [(None, 384)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "token_type_ids (InputLayer)     [(None, 384)]        0                                            \n",
      "__________________________________________________________________________________________________\n",
      "tf_bert_model (TFBertModel)     ((None, 384, 768), ( 177853440   input_ids[0][0]                  \n",
      "                                                                 attention_mask[0][0]             \n",
      "                                                                 token_type_ids[0][0]             \n",
      "__________________________________________________________________________________________________\n",
      "start_logit (Dense)             (None, 384, 1)       768         tf_bert_model[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "end_logit (Dense)               (None, 384, 1)       768         tf_bert_model[0][0]              \n",
      "__________________________________________________________________________________________________\n",
      "flatten (Flatten)               (None, 384)          0           start_logit[0][0]                \n",
      "__________________________________________________________________________________________________\n",
      "flatten_1 (Flatten)             (None, 384)          0           end_logit[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "start_probs (Activation)        (None, 384)          0           flatten[0][0]                    \n",
      "__________________________________________________________________________________________________\n",
      "end_probs (Activation)          (None, 384)          0           flatten_1[0][0]                  \n",
      "==================================================================================================\n",
      "Total params: 177,854,976\n",
      "Trainable params: 177,854,976\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "with strategy.scope():\n",
    "    model = create_model()\n",
    "    model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation callback\n",
    "\n",
    "<u>Skip this step if you already have a fine-tunned model.</u>\n",
    "\n",
    "Each `SquadExample` object contains the character level offsets for each token in its input paragraph. We use them to get back the span of text corresponding to the tokens between our predicted start and end tokens. All the ground-truth answers are also present in each `SquadExample` object. We calculate the percentage of data points where the span of text obtained from model predictions matches one of the ground-truth answers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source: https://keras.io/examples/nlp/text_extraction_with_bert/\n",
    "class ExactMatch(tf.keras.callbacks.Callback):\n",
    "    def __init__(self, x_eval, y_eval):\n",
    "        self.x_eval = x_eval\n",
    "        self.y_eval = y_eval\n",
    "\n",
    "    def on_epoch_end(self, epoch, logs=None):\n",
    "        pred_start, pred_end = self.model.predict(self.x_eval)\n",
    "        count = 0\n",
    "        eval_examples_no_skip = [_ for _ in eval_squad_examples if _.skip == False]\n",
    "        for idx, (start, end) in enumerate(zip(pred_start, pred_end)):\n",
    "            squad_eg = eval_examples_no_skip[idx]\n",
    "            offsets = squad_eg.context_token_to_char\n",
    "            start = np.argmax(start)\n",
    "            end = np.argmax(end)\n",
    "            if start >= len(offsets):\n",
    "                continue\n",
    "            \n",
    "            # Get answer from context text\n",
    "            pred_char_start = offsets[start][0]\n",
    "            if end < len(offsets):\n",
    "                pred_char_end = offsets[end][1]\n",
    "                pred_ans = squad_eg.context[pred_char_start:pred_char_end]\n",
    "            else:\n",
    "                pred_ans = squad_eg.context[pred_char_start:]\n",
    "\n",
    "            # Normalize answers before comparing prediction and true answers\n",
    "            normalized_pred_ans = self._normalize_text(pred_ans)\n",
    "            normalized_true_ans = [self._normalize_text(_) for _ in squad_eg.all_answers]\n",
    "            \n",
    "            # If the prediction is contained in the true answer, it counts as a hit\n",
    "            if normalized_pred_ans in normalized_true_ans:\n",
    "                count += 1\n",
    "\n",
    "        acc = count / len(self.y_eval[0])\n",
    "        print(f\"\\nepoch={epoch+1}, exact match score={acc:.2f}\")\n",
    "    \n",
    "    def _normalize_text(self, text):\n",
    "        text = text.lower()\n",
    "\n",
    "        # Remove punctuations\n",
    "        exclude = set(string.punctuation)\n",
    "        text = ''.join(ch for ch in text if ch not in exclude)\n",
    "\n",
    "        # Remove articles\n",
    "        regex = re.compile(r\"\\b(a|an|the)\\b\", re.UNICODE)\n",
    "        text = re.sub(regex, ' ', text)\n",
    "\n",
    "        # Remove extra white spaces\n",
    "        text = re.sub(r\"\\s+\", ' ', text).strip()\n",
    "\n",
    "        return text"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train & save model weights\n",
    "\n",
    "<u>Skip this step if you already have a fine-tunned model.</u>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 8\n",
    "\n",
    "with strategy.scope():\n",
    "    exact_match_callback = ExactMatch(x_eval, y_eval)\n",
    "    model.fit(\n",
    "        x_train,\n",
    "        y_train,\n",
    "        epochs=EPOCHS,\n",
    "        verbose=1,\n",
    "        batch_size=BATCH_SIZE,\n",
    "        callbacks=[exact_match_callback],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "weigh = model.get_weights()\n",
    "pklfile = '%s%s-%s.pickle' % (ARTIFACTS_PATH, LANGUAGE, MODEL_NAME)\n",
    "\n",
    "with open(pklfile, 'wb') as fp:\n",
    "    pickle.dump(weigh, fp, protocol= pickle.HIGHEST_PROTOCOL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Answering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load model weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "pklfile = '%s%s-%s.pickle' % (ARTIFACTS_PATH, LANGUAGE, MODEL_NAME)\n",
    "with open(pklfile, 'rb') as fp:\n",
    "    data = pickle.load(fp)\n",
    "    model.set_weights(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Infer answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_answer_question(question, context, model, tokenizer):\n",
    "    # Fix white spaces\n",
    "    context = re.sub(r\"\\s+\", ' ', context).strip()\n",
    "    question = re.sub(r\"\\s+\", ' ', question).strip()\n",
    "\n",
    "    # Encode context (token IDs, mask and token types)\n",
    "    tokenized_context = tokenizer.encode(context)\n",
    "    tokenized_question = tokenizer.encode(question)\n",
    "\n",
    "    # Create inputs\n",
    "    input_ids = tokenized_context.ids + tokenized_question.ids[1:]\n",
    "    token_type_ids = [0] * len(tokenized_context.ids) + [1] * len(\n",
    "        tokenized_question.ids[1:]\n",
    "    )\n",
    "    attention_mask = [1] * len(input_ids)\n",
    "\n",
    "    # Pad and create attention masks.\n",
    "    padding_length = MAX_LEN - len(input_ids)\n",
    "    if padding_length > 0:\n",
    "        input_ids = input_ids + ([0] * padding_length)\n",
    "        attention_mask = attention_mask + ([0] * padding_length)\n",
    "        token_type_ids = token_type_ids + ([0] * padding_length)\n",
    "    elif padding_length < 0:\n",
    "        raise Exception('Too long!')\n",
    "\n",
    "    input_ids = np.asarray(input_ids, dtype='int32')\n",
    "    token_type_ids = np.asarray(token_type_ids, dtype='int32')\n",
    "    attention_mask = np.asarray(attention_mask, dtype='int32')\n",
    "        \n",
    "    encoded_input = [\n",
    "        np.asarray([input_ids]),\n",
    "        np.asarray([token_type_ids]),\n",
    "        np.asarray([attention_mask])\n",
    "    ]\n",
    "    \n",
    "    # Get prediction of answer for the given question and context.\n",
    "    pred_start, pred_end = model.predict(encoded_input)\n",
    "    \n",
    "    start = np.argmax(pred_start[0])\n",
    "    end = np.argmax(pred_end[0])\n",
    "    \n",
    "    offsets = tokenized_context.offsets\n",
    "    if start >= len(offsets):\n",
    "        print('Cannot capture answer.')\n",
    "\n",
    "    pred_char_start = offsets[start][0]\n",
    "    if end < len(offsets):\n",
    "        pred_char_end = offsets[end][1]\n",
    "        pred_ans = context[pred_char_start:pred_char_end]\n",
    "    else:\n",
    "        pred_ans = context[pred_char_start:]\n",
    "\n",
    "    # Remove extra white spaces\n",
    "    normalized_pred_ans = re.sub(r\"\\s+\", ' ', pred_ans).strip()\n",
    "    \n",
    "    return normalized_pred_ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "======  TEXT  ======\n",
      "El Tempo Expiatorio de la Sagrada Familia (en catalán, Temple Expiatori de la\n",
      "Sagrada Família), conocido simplemente como la Sagrada Familia, es una basílica\n",
      "católica de Barcelona (España), diseñada por el arquitecto Antoni Gaudí.\n",
      "Iniciada en 1882, todavía está en construcción. Es la obra maestra de Gaudí, y\n",
      "el máximo exponente de la arquitectura modernista catalana. Es uno de los\n",
      "monumentos más visitados de España, junto al Museo del Prado y la Alhambra de\n",
      "Granada, y es la iglesia más visitada de Europa tras la basílica de San Pedro\n",
      "del Vaticano. Cuando esté finalizada será la iglesia cristiana más alta del\n",
      "mundo.\n",
      "=====================\n",
      "===== QUESTION  =====\n",
      "¿Quién diseñó la Sagrada Familia?\n",
      "=====================\n",
      "=====  ANSWER  ======\n",
      "Antoni Gaudí\n"
     ]
    }
   ],
   "source": [
    "# Write a context and a question about the context (in the same language as the XQuAD dataset)\n",
    "context = \"El Tempo Expiatorio de la Sagrada Familia (en catalán, Temple Expiatori de la Sagrada Família), conocido simplemente como la Sagrada Familia, es una basílica católica de Barcelona (España), diseñada por el arquitecto Antoni Gaudí. Iniciada en 1882, todavía está en construcción. Es la obra maestra de Gaudí, y el máximo exponente de la arquitectura modernista catalana. Es uno de los monumentos más visitados de España, junto al Museo del Prado y la Alhambra de Granada, y es la iglesia más visitada de Europa tras la basílica de San Pedro del Vaticano. Cuando esté finalizada será la iglesia cristiana más alta del mundo.\"\n",
    "question = '¿Quién diseñó la Sagrada Familia?'\n",
    "\n",
    "# Import textwrap library to display context\n",
    "import textwrap\n",
    "wrapper = textwrap.TextWrapper(width=80) \n",
    "\n",
    "# Display\n",
    "print('='*6, ' TEXT ', '='*6)\n",
    "print(wrapper.fill(context))\n",
    "print('='*21)\n",
    "\n",
    "print('='*5, 'QUESTION ', '='*5)\n",
    "print(question)\n",
    "print('='*21)\n",
    "\n",
    "# Infer answer\n",
    "answer = get_answer_question(question, context, model, tokenizer)\n",
    "print('='*5, ' ANSWER ', '='*6)\n",
    "print(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}