notebooks/model_implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing a new model with Jack \n",
    "\n",
    "In this tutorial, we focus on the minimal steps required to implement a new model from scratch using Jack.\n",
    "Please note that this tutorial has a lot of detail. It is aimed at developers who want to understand the internals of Jack. \n",
    "\n",
    "In order to implement a Jack Reader, we define three modules:\n",
    "- **Input Module**: Responsible for mapping `QASetting`s to numpy arrays assoicated with `TensorPort`s\n",
    "- **Model Module**: Defines the differentiable model architecture graph (_TensorFlow_ or _PyTorch_)\n",
    "- **Output Module**: Converting the network output to human-readable overall system output. \n",
    "\n",
    "Jack is modular, in the sense that any particular input/model/output module can be exchanged with another one.\n",
    "To illustrate this, we will implement an entire reader, and will then go on to implement another reader, but reusing the _Input_ and _Output_ module of the first.\n",
    "\n",
    "The first reader will have a _Model Module_ based on *TensorFlow*, the second will have a _Model Module_ based on *PyTorch*.\n",
    "\n",
    "### Model Overview\n",
    "As example, we will implement a simple Bi-LSTM baseline for extractive question answering, which involves extracting the answer to a question from a given text. On a high level, the architecture looks as follows:\n",
    "- Words of question and support are embedded using random embeddings (not trained)\n",
    "- Both word and question are encoded using a bi-directional LSTM\n",
    "- The question is summarized by a weighted token representation\n",
    "- A feedforward NN scores each of the support tokens to be the _start_ of the answer\n",
    "- A feedforward NN scores each of the support tokens to be the _end_ of the answer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First change dir to jack parent\n",
    "import os\n",
    "os.chdir('..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from jack.core import *\n",
    "from jack.core.tensorflow import TFReader, TFModelModule\n",
    "from jack.io.embeddings import Embeddings\n",
    "from jack.util.hooks import LossHook\n",
    "from jack.util.vocab import *\n",
    "from jack.readers.extractive_qa.shared import XQAPorts, XQAOutputModule\n",
    "from jack.readers.extractive_qa.util import prepare_data\n",
    "from jack.readers.extractive_qa.util import tokenize\n",
    "from jack import tfutil\n",
    "from jack.tfutil import sequence_encoder\n",
    "from jack.tfutil.misc import mask_for_lengths\n",
    "from jack.util.map import numpify\n",
    "from jack.util.preprocessing import stack_and_pad\n",
    "import tensorflow as tf\n",
    "_tokenize_pattern = re.compile('\\w+|[^\\w\\s]')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ports\n",
    "\n",
    "All communication between _Input_, _Model_ and _Output_ modules happens via `TensorPort`s (see `jack/core/tensorport.py`). Tensorports can be understood as placeholders for tensors, and define the ways in which information is communicated between the differentiable model architecture (_Model_ module), and the _Input_ and _Output_ modules.\n",
    "\n",
    "This is useful when implementing new models: often there already exists a model for the same task, and you can re-use existing _Input_ or _Output_ modules. You can re-use existing modules by making sure that your new module is compatible to the  ports specified in the already existing modules.\n",
    "\n",
    "In case you can reuse existing _Input_ or _Output_ modules, it is then enough to simply\n",
    "implement a new _Model_ Module (see below) that adheres to the same Tensorport interface.\n",
    "See `jack/readers/implementations.py` to see how different readers re-use the same modules.\n",
    "\n",
    "If you need a new port, however, it is also straight-forward to define one.\n",
    "For this tutorial, we will define most ports here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyPorts:\n",
    "\n",
    "    embedded_question = TensorPort(np.float32, [None, None, None],\n",
    "                                   \"embedded_question\",\n",
    "                                   \"Represents the embedded question\",\n",
    "                                   \"[B, max_num_question_tokens, N]\")\n",
    "    # or reuse Ports.Misc.embedded_question\n",
    "\n",
    "    question_length = TensorPort(np.int32, [None],\n",
    "                                 \"question_length\",\n",
    "                                 \"Represents length of questions in batch\",\n",
    "                                 \"[B]\")\n",
    "    # or reuse Ports.Input.question_length\n",
    "\n",
    "    embedded_support = TensorPort(np.float32, [None, None, None],\n",
    "                                  \"embedded_support\",\n",
    "                                  \"Represents the embedded support\",\n",
    "                                  \"[B, max_num_tokens, N]\")\n",
    "    # or reuse Ports.Misc.embedded_support\n",
    "\n",
    "    support_length = TensorPort(np.int32, [None],\n",
    "                                \"support_length\",\n",
    "                                \"Represents length of support in batch\",\n",
    "                                \"[B]\")\n",
    "    # or reuse Ports.Input.support_length\n",
    "\n",
    "    start_scores = TensorPort(np.float32, [None, None],\n",
    "                              \"start_scores\",\n",
    "                              \"Represents start scores for each support sequence\",\n",
    "                              \"[B, max_num_tokens]\")\n",
    "    # or reuse Ports.Prediction.start_scores\n",
    "\n",
    "    end_scores = TensorPort(np.float32, [None, None],\n",
    "                            \"end_scores\",\n",
    "                            \"Represents end scores for each support sequence\",\n",
    "                            \"[B, max_num_tokens]\")\n",
    "    # or reuse Ports.Prediction.end_scores\n",
    "\n",
    "    span_prediction = TensorPort(np.int32, [None, 2],\n",
    "                                 \"span_prediction\",\n",
    "                                 \"Represents predicted answer as a (start, end) span\",\n",
    "                                 \"[B, 2]\")\n",
    "    # or reuse Ports.Prediction.span_prediction\n",
    "\n",
    "    answer_span = TensorPort(np.int32, [None, 2],\n",
    "                             \"answer_span_target\",\n",
    "                             \"Represents target answer as a (start, end) span\",\n",
    "                             \"[B, 2]\")\n",
    "    # or reuse Ports.Target.answer_span\n",
    "\n",
    "    token_offsets = TensorPort(np.int32, [None, None],\n",
    "                               \"token_offsets\",\n",
    "                               \"Character index of tokens in support.\",\n",
    "                               \"[B, support_length]\")\n",
    "    # or reuse XQAPorts.token_offsets\n",
    "    \n",
    "    loss = Ports.loss  # this port must be used"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing an Input Module\n",
    "\n",
    "The _Input_ module is responsible for converting `QASetting` instances (the inputs to the reader) into numpy\n",
    "arrays, which are mapped to `TensorPort`s and passed on to the _Model_ module.\n",
    "Effectively, we are building the tensorflow _feed dictionary_ used during training and inference. \n",
    "There are _Input_ modules for\n",
    "several readers that can easily be reused when your model requires the same\n",
    "pre-processing and input as another model. \n",
    "**Note**: Similarly, this is also true for the _Output_ Module. \n",
    "\n",
    "To implement a new _Input_ module, you could implement the `InputModule` interface, but in many cases it'll be\n",
    "easier to inherit from `OnlineInputModule`, which already comes with useful functionality. In our implementation we will do the latter. We will need to:\n",
    "- Define the output `TensorPort`s of our input module. These will be used to communicate with the _Model_ module\n",
    "- Implement the actual preprocessing (e.g. tokenization, mapping to embedding vectors, ...). The result of this step is one *annotation* per instance; this annotation is a `dict` with values for every Tensorport to pass on to the _Model_ module (see `_preprocess_instance()` below).\n",
    "- Implement batching."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyInputModule(OnlineInputModule):\n",
    "    \n",
    "    def setup(self):\n",
    "        self.vocab = self.shared_resources.vocab\n",
    "        self.embeddings = self.shared_resources.embeddings\n",
    "        self.__default_embedding = np.zeros([self.embeddings.shape[-1]])\n",
    "\n",
    "    # We will now define the input and output TensorPorts of our model.\n",
    "\n",
    "    @property\n",
    "    def output_ports(self):\n",
    "        return [MyPorts.embedded_question,           # Question embeddings\n",
    "                MyPorts.question_length,             # Lengths of the questions\n",
    "                MyPorts.embedded_support,            # Support embeddings\n",
    "                MyPorts.support_length,              # Lengths of the supports\n",
    "                MyPorts.token_offsets  # Character offsets of tokens in support, used for in ouput module\n",
    "               ]\n",
    "\n",
    "    @property\n",
    "    def training_ports(self):\n",
    "        return [MyPorts.answer_span]                 # Answer span, one for each question\n",
    "\n",
    "    # Now, we implement our preprocessing. This involves tokenization,\n",
    "    # mapping to token IDs, mapping to to token embeddings,\n",
    "    # and computing the answer spans.\n",
    "\n",
    "\n",
    "    def preprocess(self, questions, answers=None, is_eval=False):\n",
    "        \"\"\"Maps a list of instances to a list of annotations.\n",
    "\n",
    "        Since in our case, all instances can be preprocessed independently, we'll\n",
    "        delegate the preprocessing to a `_preprocess_instance()` method.\n",
    "        \"\"\"\n",
    "\n",
    "        if answers is None:\n",
    "            answers = [None] * len(questions)\n",
    "\n",
    "        return [self._preprocess_instance(q, a)\n",
    "                for q, a in zip(questions, answers)]\n",
    "\n",
    "    def _preprocess_instance(self, question, answers=None):\n",
    "        \"\"\"Maps an instance to an annotation.\n",
    "\n",
    "        An annotation contains the embeddings and length of question and support,\n",
    "        token offsets, and optionally answer spans.\n",
    "        \"\"\"\n",
    "\n",
    "        has_answers = answers is not None\n",
    "\n",
    "        # `prepare_data()` handles most of the computation in our case, but\n",
    "        # you could implement your own preprocessing here.\n",
    "        q_tokenized, q_ids, _, q_length, s_tokenized, s_ids, _, s_length, \\\n",
    "        word_in_question, offsets, answer_spans = \\\n",
    "            prepare_data(question, answers, self.vocab,\n",
    "                         with_answers=has_answers,\n",
    "                         max_support_length=100)\n",
    "        # there is only 1 support\n",
    "        s_tokenized, s_ids, s_length, offsets = s_tokenized[0], s_ids[0], s_length[0], offsets[0]\n",
    "\n",
    "        # For both question and support, we'll fill an embedding tensor\n",
    "        emb_support = np.zeros([s_length, self.embeddings.shape[-1]])\n",
    "        emb_question = np.zeros([q_length, self.embeddings.shape[-1]])\n",
    "        for k in range(len(s_tokenized)):\n",
    "            emb_support[k] = self.embeddings.get(s_tokenized[k], self.__default_embedding)\n",
    "        for k in range(len(q_tokenized)):\n",
    "            emb_question[k] = self.embeddings.get(q_tokenized[k], self.__default_embedding)\n",
    "\n",
    "        # Now, we build the annotation for the question instance. We'll use a\n",
    "        # dict that maps from `TensorPort` to numpy array, but this could be\n",
    "        # any data type, like a custom class, or a named tuple.\n",
    "\n",
    "        annotation = {\n",
    "            MyPorts.question_length: q_length,\n",
    "            MyPorts.embedded_question: emb_question,\n",
    "            MyPorts.support_length: s_length,\n",
    "            MyPorts.embedded_support: emb_support,\n",
    "            MyPorts.token_offsets: offsets\n",
    "        }\n",
    "\n",
    "        if has_answers:\n",
    "            # For the purpose of this tutorial, we'll only use the first answer, such\n",
    "            # that we will have exactly as many answers as questions.\n",
    "            annotation[MyPorts.answer_span] = answer_spans[0][0]\n",
    "\n",
    "        return numpify(annotation, keys=annotation.keys())\n",
    "\n",
    "    def create_batch(self, annotations, is_eval, with_answers):\n",
    "        \"\"\"Now, we need to implement the mapping of a list of annotations to a feed dict.\n",
    "        \n",
    "        Because our annotations already are dicts mapping TensorPorts to numpy\n",
    "        arrays, we only need to do padding here.\n",
    "        \"\"\"\n",
    "\n",
    "        return {key: stack_and_pad([a[key] for a in annotations])\n",
    "                for key in annotations[0].keys()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing a Model Module\n",
    "\n",
    "The _Model_ module defines the differentiable computation graph.\n",
    "It takes _Input_ module outputs as inputs, and produces outputs (such as the loss, or logits)\n",
    "that match the inputs to the _Output_ module.\n",
    "\n",
    "We first look at a _TensorFlow_ implementation of the _Model_ module; futher below you can find an implementation using _PyTorch_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModelModule(TFModelModule):\n",
    "\n",
    "    @property\n",
    "    def input_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.embedded_question,\n",
    "                MyPorts.question_length,\n",
    "                MyPorts.embedded_support,\n",
    "                MyPorts.support_length]\n",
    "\n",
    "    @property\n",
    "    def output_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.start_scores,\n",
    "                MyPorts.end_scores,\n",
    "                MyPorts.span_prediction]\n",
    "\n",
    "    @property\n",
    "    def training_input_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.start_scores,\n",
    "                MyPorts.end_scores,\n",
    "                MyPorts.answer_span]\n",
    "\n",
    "    @property\n",
    "    def training_output_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.loss]\n",
    "\n",
    "    def create_output(self, shared_resources, input_tensors):\n",
    "        \"\"\"\n",
    "        Implements the \"core\" model: The TensorFlow subgraph which computes the\n",
    "        answer span from the embedded question and support.\n",
    "        Args:\n",
    "            emb_question: [Q, L_q, N]\n",
    "            question_length: [Q]\n",
    "            emb_support: [Q, L_s, N]\n",
    "            support_length: [Q]\n",
    "\n",
    "        Returns:\n",
    "            start_scores [B, L_s, N], end_scores [B, L_s, N], span_prediction [B, 2]\n",
    "        \"\"\"\n",
    "        tensors = TensorPortTensors(input_tensors)\n",
    "        with tf.variable_scope(\"fast_qa\", initializer=tf.contrib.layers.xavier_initializer()):\n",
    "            dim = shared_resources.config['repr_dim']\n",
    "            # set shapes for inputs\n",
    "            tensors.embedded_question.set_shape([None, None, dim])\n",
    "            tensors.embedded_support.set_shape([None, None, dim])\n",
    "\n",
    "            # encode question and support\n",
    "            rnn = tf.contrib.rnn.LSTMBlockFusedCell\n",
    "            encoded_question = sequence_encoder.bi_lstm(dim, tensors.embedded_question,\n",
    "                                                        tensors.question_length, name='bilstm',\n",
    "                                                        with_projection=True)\n",
    "\n",
    "            encoded_support = sequence_encoder.bi_lstm(dim, tensors.embedded_support,\n",
    "                                                       tensors.support_length, name='bilstm',\n",
    "                                                       reuse=True, with_projection=True)\n",
    "\n",
    "            start_scores, end_scores, predicted_start_pointer, predicted_end_pointer = \\\n",
    "                self._output_layer(dim, encoded_question, tensors.question_length,\n",
    "                                   encoded_support, tensors.support_length)\n",
    "\n",
    "            span = tf.concat([predicted_start_pointer, predicted_end_pointer], 1)\n",
    "\n",
    "            return TensorPort.to_mapping(self.output_ports, (start_scores, end_scores, span))\n",
    "\n",
    "    def _output_layer(self,\n",
    "                      dim,\n",
    "                      encoded_question,\n",
    "                      question_length,\n",
    "                      encoded_support,\n",
    "                      support_length):\n",
    "        \"\"\"Simple span prediction layer of our network\"\"\"\n",
    "        batch_size = tf.shape(question_length)[0]\n",
    "\n",
    "        # Computing weighted question state\n",
    "        attention_scores = tf.contrib.layers.fully_connected(encoded_question, 1,\n",
    "                                                             scope=\"question_attention\")\n",
    "        q_mask = mask_for_lengths(question_length, batch_size)\n",
    "        attention_scores = attention_scores + tf.expand_dims(q_mask, 2)\n",
    "        question_attention_weights = tf.nn.softmax(attention_scores, 1,\n",
    "                                                   name=\"question_attention_weights\")\n",
    "        question_state = tf.reduce_sum(question_attention_weights * encoded_question, [1])\n",
    "\n",
    "        # Prediction\n",
    "        support_mask = mask_for_lengths(support_length, batch_size)\n",
    "        interaction = tf.expand_dims(question_state, 1) * encoded_support\n",
    "        \n",
    "        def predict():\n",
    "            scores = tf.layers.dense(tf.concat([interaction, encoded_support], axis=2), 1)\n",
    "            scores = tf.squeeze(scores, [2])\n",
    "            scores = scores + support_mask\n",
    "            _, predicted = tf.nn.top_k(scores, 1)\n",
    "            return scores, predicted\n",
    "\n",
    "        start_scores, predicted_start_pointer = predict()\n",
    "        end_scores, predicted_end_pointer = predict()\n",
    "\n",
    "        return start_scores, end_scores, predicted_start_pointer, predicted_end_pointer\n",
    "\n",
    "    def create_training_output(self,\n",
    "                               shared_resources,\n",
    "                               input_tensors) -> Sequence[TensorPort]:\n",
    "        \"\"\"Compute loss from start & end scores and the gold-standard `answer_span`.\"\"\"\n",
    "        tensors = TensorPortTensors(input_tensors)\n",
    "        start, end = [tf.squeeze(t, 1) for t in tf.split(tensors.answer_span_target, 2, 1)]\n",
    "\n",
    "        start_score_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=tensors.start_scores,\n",
    "                                                                          labels=start)\n",
    "        end_score_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=tensors.end_scores,\n",
    "                                                                        labels=end)\n",
    "        loss = start_score_loss + end_score_loss\n",
    "        return TensorPort.to_mapping(self.training_output_ports, [tf.reduce_mean(loss)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing an Output Module\n",
    "\n",
    "The _Output_ module converts model predictions from the differentiable computation graph into `Answer` instances.\n",
    "Since our model is a standard extractive QA model, we could reuse the existing `XQAOutputModule`, rather than implementing our own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyOutputModule(OutputModule):\n",
    "    @property\n",
    "    def input_ports(self) -> List[TensorPort]:\n",
    "        return [MyPorts.span_prediction,\n",
    "                MyPorts.token_offsets,\n",
    "                MyPorts.start_scores,\n",
    "                MyPorts.end_scores]\n",
    "    \n",
    "    def __call__(self,\n",
    "                 questions,\n",
    "                 input_tensors) -> Sequence[Answer]:\n",
    "        \"\"\"Produces best answer for each question.\"\"\"\n",
    "        answers = []\n",
    "        tensors = TensorPortTensors(input_tensors)\n",
    "        for i, question in enumerate(questions):\n",
    "            offsets = tensors.token_offsets[i]\n",
    "            start, end = tensors.span_prediction[i]\n",
    "            score = tensors.start_scores[i, start] + tensors.end_scores[i, end]\n",
    "            # map token to char span\n",
    "            char_start = offsets[start]\n",
    "            char_end = offsets[end + 1] if end < len(offsets) - 1 else len(question.support[0])\n",
    "            answer = question.support[0][char_start: char_end]\n",
    "            answer = answer.rstrip()\n",
    "            char_end = char_start + len(answer)\n",
    "            \n",
    "            answers.append(Answer(answer, span=(char_start, char_start), score=score))\n",
    "\n",
    "        return answers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Putting Together all Modules\n",
    "\n",
    "We are now ready to put together the above defined _Input_, _Model_, and _Output_ modules into one _Reader_.\n",
    "\n",
    "For illustration purposes, we will use a toy data example with just one example question:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_set = [\n",
    "    (QASetting(\n",
    "        question=\"Which is it?\",\n",
    "        support=[\"While b seems plausible, answer a is correct.\"],\n",
    "        id=\"1\"),\n",
    "     [Answer(text=\"a\", span=(32, 33))])\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before assembling the parts of our newly defined reader, we will need to define some shared resources, which all of the modules can depend on. This includes a vocabulary `Vocab`, and a configuration hyperparameter dictionary `config`.\n",
    "\n",
    "We build the vocabulary directly from the above data set using the function `build_vocab()`, which also associates each word with random embedding vectors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_dim = 10\n",
    "\n",
    "def build_vocab(questions):\n",
    "    \"\"\"Build a vocabulary of random vectors.\"\"\"\n",
    "\n",
    "    vocab_dict = dict()\n",
    "    for question in questions:\n",
    "        for t in tokenize(question.question):\n",
    "            if t not in vocab_dict:\n",
    "                vocab_dict[t] = len(vocab_dict)\n",
    "    embeddings = Embeddings(vocab_dict, \n",
    "                            np.random.random([len(vocab_dict),\n",
    "                                              embedding_dim]))\n",
    "\n",
    "    vocab = Vocab(vocab=vocab_dict)\n",
    "    return vocab, embeddings\n",
    "\n",
    "questions = [q for q, _ in data_set]\n",
    "# build some vocab and embeddings\n",
    "vocab, embeddings = build_vocab(questions)\n",
    "shared_resources = SharedResources(vocab,\n",
    "                                   config={'repr_dim': 10},\n",
    "                                   embeddings=embeddings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then instantiate our above defined modules with these `shared_resources` as input parameter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "input_module = MyInputModule(shared_resources)\n",
    "model_module = MyModelModule(shared_resources)\n",
    "output_module = MyOutputModule()\n",
    "\n",
    "reader = TFReader(shared_resources, input_module, model_module, output_module)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point, the Reader is complete! It is composed of the three modules and shared resources, and is ready to generate predictions, or to train it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:jack.core.reader:Number of parameters: 1943\n",
      "INFO:jack.core.reader:Start training...\n",
      "INFO:jack.util.hooks:Epoch 1\tIter 1\ttrain loss 4.607385635375977\n",
      "INFO:jack.util.hooks:Epoch 2\tIter 2\ttrain loss 3.731391429901123\n",
      "INFO:jack.util.hooks:Epoch 3\tIter 3\ttrain loss 1.7199833393096924\n",
      "INFO:jack.util.hooks:Epoch 4\tIter 4\ttrain loss 0.1438484787940979\n",
      "INFO:jack.util.hooks:Epoch 5\tIter 5\ttrain loss 0.001254733419045806\n",
      "INFO:jack.util.hooks:Epoch 6\tIter 6\ttrain loss 9.536739753457368e-07\n",
      "INFO:jack.util.hooks:Epoch 7\tIter 7\ttrain loss 0.0\n",
      "INFO:jack.util.hooks:Epoch 8\tIter 8\ttrain loss 0.0\n",
      "INFO:jack.util.hooks:Epoch 9\tIter 9\ttrain loss 0.0\n",
      "INFO:jack.util.hooks:Epoch 10\tIter 10\ttrain loss 0.0\n",
      "\n",
      "Which is it? While b seems plausible, answer a is correct.\n",
      "-96.17295837402344, (32, 32), a\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "\n",
    "hooks = [LossHook(reader, iter_interval=1)]\n",
    "optimizer = tf.train.AdamOptimizer(learning_rate=0.1)\n",
    "reader.train(optimizer, data_set, batch_size, max_epochs=10, hooks=hooks)\n",
    "\n",
    "print()\n",
    "print(questions[0].question, questions[0].support[0])\n",
    "answers = reader(questions)\n",
    "print(\"{}, {}, {}\".format(answers[0].score, answers[0].span, answers[0].text))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Note:** If you want to train your newly implemented model using the main training script `jack/train_reader.py`, you first have to register a name for your new model in `jack.core.implementations`.\n",
    "\n",
    "### Hooks\n",
    "In the above example, we are making use of a _hook_. Hooks are used to monitor progress throughout training. For example, the `LossHook` monitors the loss throughout training, but other hooks can measure validation accuracy, time elapsed, etc. \n",
    "Jack comes with several hooks predefined (see `jack.util.hooks`), but you can always extend them or add your own.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing a QA model in PyTorch\n",
    "\n",
    "Above, we have implemented a complete reader from scratch, using _TensorFlow_ to define the differentiable computation graph in the _Model_ module. Let's now implement another reader, reusing as much as possible, but change frameworks from _TensorFlow_ to _PyTorch_.\n",
    "\n",
    "All we need to do to accomplish this, is to write another _ModelModule_.\n",
    "\n",
    "**Note:** the following code requires that you to have installed [PyTorch](http://pytorch.org/).\n",
    "\n",
    "### Differentiable Model Architecture (PyTorch)\n",
    "\n",
    "Let's first define _PyTorch_ modules that define the differentiable model architecture. This is independent of Jack, but we do offer some convenience functions, similar to TF. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from jack.torch_util import embedding, misc, xqa\n",
    "from jack.torch_util.highway import Highway\n",
    "from jack.torch_util.rnn import BiLSTM\n",
    "\n",
    "class MyPredictionTorchModule(nn.Module):\n",
    "    def __init__(self, shared_resources):\n",
    "        super(MyPredictionTorchModule, self).__init__()\n",
    "        self._shared_resources = shared_resources\n",
    "        repr_dim_input = shared_resources.embeddings.shape[-1]\n",
    "        repr_dim = shared_resources.config[\"repr_dim\"]\n",
    "        \n",
    "        # nn child modules\n",
    "        self._bilstm = BiLSTM(repr_dim_input, repr_dim)\n",
    "        self._linear_question_attention = nn.Linear(2 * repr_dim, 1, bias=False)\n",
    "        self._linear_start_scores = nn.Linear(2 * repr_dim, 1, bias=False)\n",
    "        self._linear_end_scores = nn.Linear(2 * repr_dim, 1, bias=False)\n",
    "\n",
    "\n",
    "    def forward(self, emb_question, question_length, emb_support, support_length):\n",
    "        # encode\n",
    "        encoded_question = self._bilstm(emb_question)[0]\n",
    "        encoded_support = self._bilstm(emb_support)[0]\n",
    "\n",
    "        # answer\n",
    "        # computing attention over question\n",
    "        attention_scores = self._linear_question_attention(encoded_question)\n",
    "        q_mask = misc.mask_for_lengths(question_length)\n",
    "        attention_scores = attention_scores.squeeze(2) + q_mask\n",
    "        question_attention_weights = F.softmax(attention_scores, dim=1)\n",
    "        question_state = torch.matmul(question_attention_weights.unsqueeze(1),\n",
    "                                      encoded_question).squeeze(1)\n",
    "        \n",
    "        interaction = question_state * encoded_support\n",
    "        # Prediction\n",
    "        start_scores = self._linear_start_scores(interaction).squeeze(2)\n",
    "        end_scores = self._linear_start_scores(interaction).squeeze(2)\n",
    "        # Mask\n",
    "        support_mask = misc.mask_for_lengths(support_length)\n",
    "        start_scores += support_mask\n",
    "        end_scores += support_mask\n",
    "\n",
    "        _, predicted_start_pointer = start_scores.max(1)\n",
    "        _, predicted_end_pointer = end_scores.max(1)\n",
    "        \n",
    "        # end pointer cannot come before start\n",
    "        predicted_end_pointer = torch.max(predicted_end_pointer, predicted_start_pointer)\n",
    "\n",
    "        span = torch.stack([predicted_start_pointer, predicted_end_pointer], 1)\n",
    "        return start_scores, end_scores, span\n",
    "    \n",
    "class MyLossTorchModule(nn.Module):\n",
    "    def forward(self, start_scores, end_scores, answer_span):\n",
    "        start, end = answer_span[:, 0], answer_span[:, 1]\n",
    "        \n",
    "        # start prediction loss\n",
    "        loss = -torch.index_select(F.log_softmax(start_scores, dim=1), dim=1, index=start.long())\n",
    "        # end prediction loss\n",
    "        loss -= torch.index_select(F.log_softmax(end_scores, dim=1), dim=1, index=end.long())\n",
    "        \n",
    "        # mean loss over the current batch\n",
    "        return loss.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementing the Jack _Model_ Module with PyTorch \n",
    "\n",
    "After defining our `torch nn.Module` classes, we can use them in a Jack `ModelModule`. Note that the signature of the `nn.Module` torch implementations above must match the tensorport signature of the `ModelModule`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jack.core.torch import PyTorchModelModule, PyTorchReader\n",
    "\n",
    "\n",
    "class MyTorchModelModule(PyTorchModelModule):\n",
    "\n",
    "    @property\n",
    "    def input_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.embedded_question,\n",
    "                MyPorts.question_length,\n",
    "                MyPorts.embedded_support,\n",
    "                MyPorts.support_length]\n",
    "\n",
    "    @property\n",
    "    def output_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.start_scores,\n",
    "                MyPorts.end_scores,\n",
    "                MyPorts.span_prediction]\n",
    "\n",
    "    @property\n",
    "    def training_input_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.start_scores,\n",
    "                MyPorts.end_scores,\n",
    "                MyPorts.answer_span]\n",
    "\n",
    "    @property\n",
    "    def training_output_ports(self) -> Sequence[TensorPort]:\n",
    "        return [MyPorts.loss]\n",
    "    \n",
    "    \n",
    "    def create_loss_module(self, shared_resources: SharedResources):\n",
    "        return MyLossTorchModule()\n",
    "\n",
    "    def create_prediction_module(self, shared_resources: SharedResources):\n",
    "        return MyPredictionTorchModule(shared_resources)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After defining our new PyTorchModelModule we can create our JackReader similar as before, by instantiating a `PyTorchReader`, rather than a `TFReader`, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_module = MyInputModule(shared_resources)\n",
    "model_module = MyTorchModelModule(shared_resources)  # was MyModelModule\n",
    "output_module = MyOutputModule()\n",
    "\n",
    "reader = PyTorchReader(shared_resources,\n",
    "                       input_module,\n",
    "                       model_module,\n",
    "                       output_module)  # was TFReader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Interacting with the instantiated readers is transparent. For the user it doesn't matter whether it is a `TFReader` or a `PyTorchReader`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:jack.core.reader:Setting up data and model...\n",
      "INFO:jack.core.input_module:OnlineInputModule pre-processes data on-the-fly in first epoch and caches results for subsequent epochs! That means, first epoch might be slower.\n",
      "INFO:jack.core.reader:Start training...\n",
      "INFO:jack.util.hooks:Epoch 1\tIter 1\ttrain loss 4.5954461097717285\n",
      "INFO:jack.util.hooks:Epoch 2\tIter 2\ttrain loss 4.420642852783203\n",
      "INFO:jack.util.hooks:Epoch 3\tIter 3\ttrain loss 3.8911027908325195\n",
      "INFO:jack.util.hooks:Epoch 4\tIter 4\ttrain loss 3.1102333068847656\n",
      "INFO:jack.util.hooks:Epoch 5\tIter 5\ttrain loss 1.9925825595855713\n",
      "INFO:jack.util.hooks:Epoch 6\tIter 6\ttrain loss 1.2458298206329346\n",
      "INFO:jack.util.hooks:Epoch 7\tIter 7\ttrain loss 0.747711181640625\n",
      "INFO:jack.util.hooks:Epoch 8\tIter 8\ttrain loss 0.31256622076034546\n",
      "INFO:jack.util.hooks:Epoch 9\tIter 9\ttrain loss 0.1637781262397766\n",
      "INFO:jack.util.hooks:Epoch 10\tIter 10\ttrain loss 0.03822925314307213\n",
      "\n",
      "Which is it? While b seems plausible, answer a is correct.\n",
      "16.32866668701172, (32, 32), a\n"
     ]
    }
   ],
   "source": [
    "batch_size = 1\n",
    "\n",
    "# torch needs to be setup already at this point, to get the parameters\n",
    "reader.setup_from_data(data_set, is_training=True)\n",
    "optimizer = torch.optim.Adam(reader.model_module.prediction_module.parameters(), lr=0.1)\n",
    "hooks = [LossHook(reader, iter_interval=1)]\n",
    "\n",
    "reader.train(optimizer,\n",
    "             data_set,\n",
    "             batch_size,\n",
    "             max_epochs=10,\n",
    "             hooks=hooks)\n",
    "\n",
    "print()\n",
    "print(questions[0].question, questions[0].support[0])\n",
    "answers = reader(questions)\n",
    "print(\"{}, {}, {}\".format(answers[0].score, answers[0].span, answers[0].text))"
   ]
  }
 ],
 "metadata": {
  "hide_input": false,
  "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}