Keras_nlp_QnA.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jv3iFl2fg2OP"
      },
      "outputs": [],
      "source": [
        "!pip install -q --upgrade keras-nlp"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import keras\n",
        "import keras_nlp\n",
        "import tensorflow_datasets as tfds\n",
        "from tensorflow.keras import layers\n",
        "import os\n",
        "import re\n",
        "import json\n",
        "import string\n",
        "import numpy as np\n",
        "\n",
        "print(tf.__version__)"
      ],
      "metadata": {
        "id": "SOwEZm8XhDqD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tokenizers import BertWordPieceTokenizer\n",
        "from transformers import BertTokenizer\n",
        "\n",
        "# Save the slow pretrained tokenizer\n",
        "slow_tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n",
        "save_path = \"bert_base_uncased/\"\n",
        "if not os.path.exists(save_path):\n",
        "    os.makedirs(save_path)\n",
        "slow_tokenizer.save_pretrained(save_path)\n",
        "\n",
        "# Load the fast tokenizer from saved file\n",
        "tokenizer = BertWordPieceTokenizer(\"bert_base_uncased/vocab.txt\", lowercase=True)"
      ],
      "metadata": {
        "id": "7_9eMkwliqz6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data_url = \"https://rajpurkar.github.io/SQuAD-explorer/dataset/train-v1.1.json\"\n",
        "train_path = keras.utils.get_file(\"train.json\", train_data_url)\n",
        "eval_data_url = \"https://rajpurkar.github.io/SQuAD-explorer/dataset/dev-v1.1.json\"\n",
        "eval_path = keras.utils.get_file(\"eval.json\", eval_data_url)"
      ],
      "metadata": {
        "id": "zDJGBEj0iDI4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "max_len = 512\n",
        "class SquadExample:\n",
        "    def __init__(self, question, context, start_char_idx, answer_text, all_answers):\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.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",
        "        # Clean context, answer and question\n",
        "        context = \" \".join(str(context).split())\n",
        "        question = \" \".join(str(question).split())\n",
        "        answer = \" \".join(str(answer_text).split())\n",
        "\n",
        "        # Find end character 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",
        "        # Tokenize context\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",
        "        # Tokenize question\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\n",
        "\n",
        "\n",
        "with open(train_path) as f:\n",
        "    raw_train_data = json.load(f)\n",
        "\n",
        "with open(eval_path) as f:\n",
        "    raw_eval_data = json.load(f)\n",
        "\n",
        "\n",
        "def create_squad_examples(raw_data):\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, context, start_char_idx, answer_text, all_answers\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\n",
        "\n",
        "\n",
        "train_squad_examples = create_squad_examples(raw_train_data)\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)\n",
        "x_eval, y_eval = create_inputs_targets(eval_squad_examples)\n",
        "print(f\"{len(eval_squad_examples)} evaluation points created.\")"
      ],
      "metadata": {
        "id": "O0OXuOdgiKdk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def create_model():\n",
        "    ## BERT encoder\n",
        "    #Select variants from https://keras.io/api/keras_nlp/models/bert/bert_backbone/\n",
        "    encoder = keras_nlp.models.BertBackbone.from_preset(\"bert_medium_en_uncased\")\n",
        "\n",
        "    #There is an option to train only the last layers. This does not give a well trained model though.\n",
        "    #encoder.trainable = False\n",
        "\n",
        "    ## QA Model\n",
        "    token_ids = layers.Input(shape=(max_len,), dtype=tf.int32)\n",
        "    segment_ids = layers.Input(shape=(max_len,), dtype=tf.int32)\n",
        "    padding_mask = layers.Input(shape=(max_len,), dtype=tf.int32)\n",
        "\n",
        "    embedding = encoder(inputs={'token_ids': token_ids, 'segment_ids': segment_ids,'padding_mask': padding_mask})\n",
        "\n",
        "    start_logits = layers.Dense(1, name=\"start_logit\", use_bias=False)(embedding['sequence_output'])\n",
        "    start_logits = layers.Flatten()(start_logits)\n",
        "\n",
        "    end_logits = layers.Dense(1, name=\"end_logit\", use_bias=False)(embedding['sequence_output'])\n",
        "    end_logits = layers.Flatten()(end_logits)\n",
        "\n",
        "    start_probs = layers.Activation(keras.activations.softmax)(start_logits)\n",
        "    end_probs = layers.Activation(keras.activations.softmax)(end_logits)\n",
        "\n",
        "    model = keras.Model(\n",
        "        inputs=[token_ids, segment_ids, padding_mask],\n",
        "        outputs=[start_probs, end_probs],\n",
        "    )\n",
        "    loss = keras.losses.SparseCategoricalCrossentropy(from_logits=False)\n",
        "    optimizer = keras.optimizers.Adam(learning_rate=5e-5)\n",
        "    model.compile(optimizer=optimizer, loss=[loss, loss])\n",
        "    return model"
      ],
      "metadata": {
        "id": "rjmWyK8Mi7b7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "model = create_model()\n",
        "model.summary()"
      ],
      "metadata": {
        "id": "A5pFbZVmjG3G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "def normalize_text(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 space\n",
        "    text = \" \".join(text.split())\n",
        "    return text\n",
        "\n",
        "\n",
        "class ExactMatch(keras.callbacks.Callback):\n",
        "    \"\"\"\n",
        "    Each `SquadExample` object contains the character level offsets for each token\n",
        "    in its input paragraph. We use them to get back the span of text corresponding\n",
        "    to the tokens between our predicted start and end tokens.\n",
        "    All the ground-truth answers are also present in each `SquadExample` object.\n",
        "    We calculate the percentage of data points where the span of text obtained\n",
        "    from model predictions matches one of the ground-truth answers.\n",
        "    \"\"\"\n",
        "\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",
        "            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",
        "            normalized_pred_ans = normalize_text(pred_ans)\n",
        "            normalized_true_ans = [normalize_text(_) for _ in squad_eg.all_answers]\n",
        "            if normalized_pred_ans in normalized_true_ans:\n",
        "                count += 1\n",
        "        acc = count / len(self.y_eval[0])\n",
        "        print(f\"\\nepoch={epoch+1}, exact match score={acc:.2f}\")\n"
      ],
      "metadata": {
        "id": "PJ79MuzqjS_f"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "exact_match_callback = ExactMatch(x_eval, y_eval)\n",
        "model.fit(\n",
        "    x_train,\n",
        "    y_train,\n",
        "    epochs=10,  # 10 epochs are recommended\n",
        "    verbose=1,\n",
        "    batch_size=16,  # change based on your GPU's memory size\n",
        "    callbacks=[exact_match_callback],\n",
        ")"
      ],
      "metadata": {
        "id": "WzxZwMcGjXBv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!mkdir -p saved_model\n",
        "model.save('saved_model/my_model')"
      ],
      "metadata": {
        "id": "Ve4OidwBBS63"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!zip -r /content/trained_model.zip /content/saved_model/my_model"
      ],
      "metadata": {
        "id": "XZj_ixfiBrtC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "converter = tf.lite.TFLiteConverter.from_keras_model(model)\n",
        "converter.optimizations = [tf.lite.Optimize.DEFAULT]\n",
        "tflite_model = converter.convert()\n",
        "\n",
        "# Save the model.\n",
        "with open('model_bert_medium_512_quant.tflite', 'wb') as f:\n",
        "  f.write(tflite_model)"
      ],
      "metadata": {
        "id": "aYH1fWNNjacQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "fgSZHL75N7I5"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Check QnA with Interpreter"
      ],
      "metadata": {
        "id": "y7ew9sDvW8xJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "max_len = 512\n",
        "context = \"Nikola Tesla (Serbian Cyrillic: 10 July 1856 - 7 January 1943) was a Serbian American inventor, electrical engineer, mechanical engineer, physicist, and futurist best known for his contributions to the design of the modern alternating current (AC) electricity supply system.\"\n",
        "question = \"In what year did Tesla die?\"\n",
        "\n",
        "tokenized_context = tokenizer.encode(context)\n",
        "print(tokenized_context)\n",
        "# Tokenize question\n",
        "tokenized_question = tokenizer.encode(question)\n",
        "print(tokenized_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"
      ],
      "metadata": {
        "id": "Lag7uBTJN6XH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "\n",
        "# Load the TFLite model and allocate tensors.\n",
        "interpreter = tf.lite.Interpreter(model_path=\"/content/model_bert_medium_512_quant.tflite\")\n",
        "\n",
        "interpreter.allocate_tensors()\n",
        "\n",
        "# Get input and output tensors.\n",
        "input_details = interpreter.get_input_details()\n",
        "output_details = interpreter.get_output_details()\n",
        "#print(input_details)\n",
        "#print(output_details)\n",
        "\n",
        "input_ids = np.array(input_ids, dtype=np.int32)\n",
        "input_ids = np.reshape(input_ids,(1, input_ids.size))\n",
        "\n",
        "token_type_ids = np.array(token_type_ids, dtype=np.int32)\n",
        "token_type_ids = np.reshape(token_type_ids,(1, token_type_ids.size))\n",
        "\n",
        "attention_mask = np.array(attention_mask, dtype=np.int32)\n",
        "attention_mask = np.reshape(attention_mask,(1, attention_mask.size))\n",
        "\n",
        "interpreter.set_tensor(input_details[0]['index'], input_ids)\n",
        "interpreter.set_tensor(input_details[1]['index'], attention_mask)\n",
        "interpreter.set_tensor(input_details[2]['index'], token_type_ids)\n",
        "\n",
        "\n",
        "interpreter.invoke()\n"
      ],
      "metadata": {
        "id": "ObAkg9CtBsIR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "output_data_0 = interpreter.get_tensor(output_details[0]['index'])\n",
        "print(output_data_0)\n",
        "output_data_1 = interpreter.get_tensor(output_details[1]['index'])\n",
        "print(output_data_1)"
      ],
      "metadata": {
        "id": "kmTodVNBTBhL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "start_position = tf.argmax(output_data_1, axis=1)\n",
        "print(start_position)\n",
        "end_position = tf.argmax(output_data_0, axis=1)\n",
        "print(end_position)\n",
        "print(int(start_position), int(end_position))"
      ],
      "metadata": {
        "id": "qmxOaeLiZlBP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(input_ids.tolist())\n",
        "answer = input_ids[0, int(start_position) : int(end_position) + 1]"
      ],
      "metadata": {
        "id": "WU9h_ks7bGBR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "logit = tokenizer.decode(answer.tolist())\n",
        "print(logit)"
      ],
      "metadata": {
        "id": "7Rp401KNW4g2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "UEkOTJ7tpqAH"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}