#1 - Pointer Network.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "import numpy as np\n",
    "from rouge import Rouge "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.enable_eager_execution()\n",
    "\n",
    "def mmap(fn, elem):\n",
    "    return list(map(fn, elem))\n",
    "\n",
    "def getel(n, lst): \n",
    "    return mmap(lambda x: x[n], lst)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_pretrained_embeddings(embed_matrix, trainable=False):\n",
    "    return keras.layers.Embedding(embed_matrix.shape[0],\n",
    "                                  embed_matrix.shape[1],\n",
    "                                  weights=[embed_matrix],\n",
    "                                  trainable=trainable)\n",
    "\n",
    "def get_new_embeddings(voc_len, embedding_dim):\n",
    "    return keras.layers.Embedding(voc_len, embedding_dim)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(keras.Model):\n",
    "    def __init__(self, units):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.units = units\n",
    "        self.lstm = keras.layers.LSTM(self.units,\n",
    "                                      return_sequences=True, \n",
    "                                      return_state=True)\n",
    "        \n",
    "    def call(self, x):\n",
    "        sequences, state1, state2 = self.lstm(x)\n",
    "        return sequences, state1, state2\n",
    "    \n",
    "class Attention(keras.Model):\n",
    "    def __init__(self, w_units):\n",
    "        super().__init__()\n",
    "        self.W = keras.layers.Dense(w_units, use_bias=False)\n",
    "        self.W1 = keras.layers.Dense(w_units, use_bias=False)\n",
    "        self.v = keras.layers.Dense(1)\n",
    "        self.tanh = keras.activations.tanh\n",
    "        self.softmax = keras.activations.softmax\n",
    "        \n",
    "    def call(self, enc_hidden, dec_hidden):\n",
    "        dec_hidden = tf.expand_dims(dec_hidden, 1)\n",
    "        unnorm = self.v(\n",
    "            self.tanh(self.W(enc_hidden) + self.W1(dec_hidden))\n",
    "        )\n",
    "\n",
    "        attention_weights = self.softmax(unnorm, axis=1)\n",
    "        \n",
    "        # Compute the context vector used to generate the decoder state\n",
    "        c_vec = tf.reduce_sum(attention_weights * enc_hidden, axis=1)\n",
    "        \n",
    "        # Return the context vector and the pointer logits and the pointer probs\n",
    "        return c_vec, tf.squeeze(attention_weights, axis=2)  \n",
    "    \n",
    "class Decoder(keras.Model):\n",
    "    def __init__(self, units, output_size):\n",
    "        super().__init__()\n",
    "        self.units = units\n",
    "        self.lstm = keras.layers.LSTM(self.units, return_state=True)\n",
    "        self.output_layer = keras.layers.Dense(output_size, activation='softmax')\n",
    "        \n",
    "    def call(self, x, enc_out, prev_state):\n",
    "        # Concatenate encoder output (or context vector) and the target/predicted embedding\n",
    "        concatenated_inp = tf.concat([x, enc_out], axis=1)\n",
    "        concatenated_inp = tf.expand_dims(concatenated_inp, 1)\n",
    "        # Compute the hidden state h^d\n",
    "        d, dec_h, dec_c =  self.lstm(concatenated_inp, initial_state=prev_state)\n",
    "        \n",
    "        # Decode using vocabulary\n",
    "        flattened = tf.layers.flatten(d)\n",
    "        decoded_probs = self.output_layer(flattened)\n",
    "        \n",
    "        # Return Decode hidden states and vocabulary logits over the vocabulary\n",
    "        return d, dec_h, dec_c, decoded_probs\n",
    "\n",
    "class PointerSwitch(keras.Model):\n",
    "    def __init__(self, units):\n",
    "        super().__init__()\n",
    "        self.W1 = keras.layers.Dense(units,  use_bias=False)\n",
    "        self.W2 = keras.layers.Dense(units, use_bias=False)\n",
    "        self.v = keras.layers.Dense(1)\n",
    "    \n",
    "    def call(self, enc, c_vec):\n",
    "        '''\n",
    "            Compute switch probabilities from the context vector and\n",
    "            the encoder last output state\n",
    "        '''\n",
    "        switch_prob = tf.keras.activations.sigmoid(\n",
    "            self.v(\n",
    "                self.W1(enc) + self.W2(c_vec)\n",
    "            )\n",
    "        )\n",
    "        \n",
    "        return switch_prob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = Rouge()\n",
    "\n",
    "def rouge_score(y, y_):\n",
    "    true_tok = [' '.join(mmap(str,sent)) for sent in y]\n",
    "    pred_tok = [' '.join(mmap(str,sent)) for sent in y_]\n",
    "    scores = r.get_scores(true_tok, pred_tok)\n",
    "\n",
    "    values = []\n",
    "    for key in scores[0].keys():\n",
    "        for sub_metric in ['f', 'p', 'r']:\n",
    "            mean_score = np.mean(mmap(lambda x: x[key][sub_metric], scores))\n",
    "            values.append(mean_score)\n",
    "    return values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PointerNetwork(keras.Model):\n",
    "    def __init__(self,\n",
    "                 enc_units,\n",
    "                 dec_units, \n",
    "                 voc_size,\n",
    "                 att_units, \n",
    "                 switch_units,\n",
    "                 max_len, \n",
    "                 start_token,\n",
    "                 end_token,\n",
    "                 padding_char):\n",
    "        super().__init__()\n",
    "        self.encoder = Encoder(enc_units)\n",
    "        self.decoder = Decoder(dec_units, voc_size)\n",
    "        self.attention = Attention(att_units)\n",
    "        self.pointer_switch = PointerSwitch(switch_units)\n",
    "        self.embeddings = False\n",
    "        self.max_len = max_len\n",
    "        self.start_token = start_token\n",
    "        self.end_token = end_token\n",
    "        self.voc_size = voc_size\n",
    "        self.padding_char = padding_char\n",
    "        \n",
    "        self.optimizer = tf.train.AdamOptimizer()\n",
    "    \n",
    "    def set_embeddings_layer(self, embeddings_layer):\n",
    "        self.embeddings = embeddings_layer\n",
    "    \n",
    "    def predict_batch(self, X):\n",
    "        assert self.embeddings, \"Call self.set_embeddings_layer first\"\n",
    "        X = tf.convert_to_tensor(X)\n",
    "        \n",
    "        embed = self.embeddings(X)\n",
    "        enc_states, h1, h2 = self.encoder(embed)\n",
    "        input_tokens = tf.convert_to_tensor([self.start_token] * embed.shape[0])\n",
    "        # put last encoder state as attention vec at start\n",
    "        c_vec = h1\n",
    "        outputs = []\n",
    "        \n",
    "        for _ in range(self.max_len):\n",
    "            dec_input = self.embeddings(input_tokens)\n",
    "            decoded_state, h1, h2, decoded_probs = self.decoder(dec_input, \n",
    "                                                                c_vec, \n",
    "                                                                [h1, h2])\n",
    "            c_vec, pointer_probs = self.attention(enc_states, \n",
    "                                                     decoded_state)\n",
    "            \n",
    "            # Compute switch probability to decide where to extract the next\n",
    "            # word token\n",
    "            switch_probs = self.pointer_switch(h1, c_vec)\n",
    "            # Decode based on switch probs            \n",
    "            input_tokens = self.decode_next_word(switch_probs, \n",
    "                                                 decoded_probs,\n",
    "                                                 X,\n",
    "                                                 pointer_probs)\n",
    "            outputs.append(input_tokens)\n",
    "            \n",
    "        return tf.transpose(tf.convert_to_tensor(outputs))\n",
    "    \n",
    "    def decode_next_word(self, switch_probs, decoded_probs, inputs, att_probs):\n",
    "        sampled_probs = tf.random.uniform(switch_probs.shape, 0 , 1)\n",
    "        tokens = []\n",
    "        token = None\n",
    "\n",
    "        for prob, sampled, decoded, inp, att_p in zip(switch_probs,\n",
    "                                                      sampled_probs, \n",
    "                                                      decoded_probs,\n",
    "                                                      inputs,\n",
    "                                                      att_probs):\n",
    "            if prob.numpy()[0] >= sampled.numpy()[0]:\n",
    "                token = self.fixed_vocab_decode(decoded)\n",
    "            else:\n",
    "                token = self.pointer_greedy_search(att_p, inp)\n",
    "\n",
    "            tokens.append(token)\n",
    "            \n",
    "        return tf.convert_to_tensor(tokens, dtype=tf.float32)\n",
    "    \n",
    "    def pointer_greedy_search(self, probs, inputs):\n",
    "        return inputs[tf.argmax(probs)]\n",
    "    \n",
    "    def fixed_vocab_decode(self, decoded_probs):\n",
    "        return tf.argmax(decoded_probs)\n",
    "    \n",
    "    def pointer_batch_loss(self, gen, y, d_prob, p_prob, s_prob):\n",
    "        # Compute the mask to ignore the padding in the loss\n",
    "        mask = 1-tf.cast(tf.equal(gen[:,None],\n",
    "                                  tf.ones(gen[:,None].shape) * self.padding_char\n",
    "                                 ), tf.float32)\n",
    "        \n",
    "        # Compute pointer loss across all values of y for the pointer and generated probs\n",
    "        pointer_mat = (p_prob + (1 - s_prob)) * mask\n",
    "        generator_mat = (d_prob + s_prob) * mask\n",
    "        \n",
    "        # Add the expected loss in terms of likelihood\n",
    "        batch_loss = 0\n",
    "        for i, g in enumerate(gen):\n",
    "            # Add if the word was taken from the input\n",
    "            if g == 0:\n",
    "                batch_loss += pointer_mat[i, y[i]]\n",
    "            # Add if the word was generated by the network\n",
    "            else:\n",
    "                batch_loss += generator_mat[i, y[i]]\n",
    "            \n",
    "        # Reduce to scalar, dont forget to include minus sign (its a loss not a likelihood)\n",
    "        return -batch_loss\n",
    "    \n",
    "    def __train_batch(self, X, y, gen):\n",
    "        assert self.embeddings, \"Call self.load_embeddings first\"\n",
    "\n",
    "        X = tf.convert_to_tensor(X)\n",
    "        y = tf.convert_to_tensor(y, dtype='int32')\n",
    "        gen = tf.convert_to_tensor(gen, dtype='float32')\n",
    "\n",
    "        enc_inp = self.embeddings(X)\n",
    "        enc_states, h1, h2 = self.encoder(enc_inp)\n",
    "        c_vec = h1\n",
    "        input_tokens = y[:,0]\n",
    "        loss = 0\n",
    "        for t in range(1, y.shape[1]):\n",
    "            # Get embeddings\n",
    "            dec_input = self.embeddings(input_tokens)\n",
    "            \n",
    "            # Get decoder output\n",
    "            decoded_state, h1, h2, decoded_probs = self.decoder(dec_input, c_vec, [h1, h2])\n",
    "            \n",
    "            # Get context vector for the next step, and pointer probabilities\n",
    "            c_vec, pointer_probs = self.attention(enc_states, decoded_state)\n",
    "            \n",
    "            # Get switch probability (BS*1)\n",
    "            switch_probs = self.pointer_switch(h1, c_vec)\n",
    "                        \n",
    "            # Is target generated or extracted from the input (BS * 1)\n",
    "            batch_gen = tf.convert_to_tensor(gen[:, t])\n",
    "            \n",
    "            # Compute Pointer Network batch loss at timestep t\n",
    "            loss += self.pointer_batch_loss(batch_gen, y[:, t], decoded_probs,\n",
    "                                       pointer_probs, switch_probs)\n",
    "\n",
    "            # Get next decoder input tokens\n",
    "            input_tokens = y[:, t]\n",
    "        \n",
    "        # Dont forget to divide by summary lenght N, since we lose the /N component n by calling\n",
    "        # N times softmax cross entropy\n",
    "        loss = loss / int(y.shape[1]-1)\n",
    "        self._loss = loss\n",
    "        return loss\n",
    "    \n",
    "    def train_batch(self, X, y, gen):\n",
    "        self.optimizer.minimize(lambda: self.__train_batch(X, y, gen))   \n",
    "        return [self._loss]\n",
    "    \n",
    "    def evaluate(self, X, y, verbose=0):\n",
    "        y_ = self.predict_batch(X)\n",
    "        return rouge_score(y, y_)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_units = 128\n",
    "dec_units = 128\n",
    "voc_size = 300\n",
    "att_units = 128 \n",
    "switch_units = 128\n",
    "max_len = 200\n",
    "start_index_token = 0\n",
    "end_index_token = 1\n",
    "padding_char = -1\n",
    "ptr = PointerNetwork(enc_units, \n",
    "                     dec_units, \n",
    "                     voc_size, \n",
    "                     att_units, \n",
    "                     switch_units, \n",
    "                     max_len, \n",
    "                     start_index_token, \n",
    "                     end_index_token,\n",
    "                     padding_char)\n",
    "\n",
    "ptr.set_embeddings_layer(load_pretrained_embeddings(np.ones((300,300))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=176852, shape=(2, 200), dtype=float32, numpy=\n",
       "array([[  1.,   1.,   1.,  55., 190.,  19.,  19.,   1.,  19.,   1.,   1.,\n",
       "          1.,  55.,  19.,  19.,   1.,   1.,   1.,  19.,   1.,  19.,  19.,\n",
       "         19.,   1.,  19.,  19.,   1.,   1.,  19.,   1.,   1.,  19.,  19.,\n",
       "         19.,   1.,  19.,   1.,   1.,  19.,   1.,  19.,   1.,  19.,   1.,\n",
       "          1.,  19.,  19.,   1.,  19.,   1.,  19.,   1.,   1.,   1.,   1.,\n",
       "         55.,  19.,   1.,  19.,  19.,   1.,  19.,   1.,   1.,  19.,   1.,\n",
       "         19.,  19.,  19.,   1.,  19.,   1.,  19.,   1.,   1.,   1.,  19.,\n",
       "         19.,   1.,  19.,   1.,   1.,   1.,   1.,   1.,  55.,   1., 190.,\n",
       "         19.,  19.,  19.,  19.,  19.,   1.,  19.,   1.,  19.,   1.,   1.,\n",
       "         19.,  19.,   1.,   1.,  19.,   1.,   1.,  19.,  19.,  19.,  19.,\n",
       "          1.,  19.,  19.,   1.,  19.,  19.,  19.,   1.,   1.,  19.,   1.,\n",
       "         19.,  19.,  19.,   1.,  19.,  19.,   1.,   1.,   1.,   1.,   1.,\n",
       "         55.,   1., 190.,  19.,   1.,  19.,  19.,   1.,  19.,   1.,  19.,\n",
       "         19.,   1.,   1.,  19.,   1.,  19.,   1.,   1.,   1.,   1.,  55.,\n",
       "          1.,   1., 190.,   1.,  55., 190.,   1., 190.,  55.,   1.,   1.,\n",
       "          1.,   1.,   1.,   1.,   1.,   1.,  55., 190.,  34.,   1.,   1.,\n",
       "         55.,   1.,   1.,   1., 190.,   1.,  55.,   1., 190.,  55.,   1.,\n",
       "          1., 190.,   1.,  55., 190., 167.,   1.,   1.,  55.,   1., 190.,\n",
       "          1.,   1.],\n",
       "       [144.,  12.,   1.,   1.,  55., 205.,  19.,  19.,  19.,  19.,   1.,\n",
       "         19.,   1.,  19.,  19.,   1.,  19.,   1.,   1.,  19.,   1.,   1.,\n",
       "          1.,  55.,  19.,   1.,  19.,   1.,  19.,   1.,   1.,  19.,   1.,\n",
       "         19.,   1.,   1.,   1.,  55.,   1.,   1., 190.,   1.,  55.,   1.,\n",
       "          1.,   1.,   1.,   1., 190.,  55., 190., 167.,   1.,  55.,   1.,\n",
       "        190.,   1.,  55., 190., 167.,  55.,   1., 190.,   1.,   1.,  55.,\n",
       "          1.,   1.,   1., 190.,   1.,  55.,   1.,   1.,   1., 190.,   1.,\n",
       "          1.,   1.,   1.,  55.,   1.,   1., 190.,   1.,  55., 190., 167.,\n",
       "          1.,  55.,   1.,   1.,   1., 190.,  55., 190., 167.,  55., 167.,\n",
       "          1.,   1., 190.,  55., 167.,   1.,   1.,   1.,   1.,  55.,   1.,\n",
       "        190.,   1., 190.,   1.,   1.,  55.,   1.,   1., 190.,   1.,   1.,\n",
       "          1.,  55., 190., 190.,   1.,   1.,  55., 190.,   1.,   1.,  55.,\n",
       "        190.,   1.,   1., 190.,   1.,  55., 190., 167.,  55.,   1.,   1.,\n",
       "          1.,   1.,   1., 190.,   1.,  55.,   1., 190.,  55., 167.,   1.,\n",
       "        190.,  55.,   1., 190.,   1.,  55., 190., 167.,  55., 167., 167.,\n",
       "        210.,   1., 179.,   1.,   1.,  55., 190.,  19.,  19.,  19.,  19.,\n",
       "         19.,  19.,   1.,   1.,  19.,  19.,  19.,  19.,  19.,  19.,  19.,\n",
       "         19.,  19.,  19.,   1.,   1.,   1.,   1.,  19.,  19.,   1.,   1.,\n",
       "          1.,   1.]], dtype=float32)>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ptr.predict_batch(tf.convert_to_tensor(np.ones((2,10))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "X = [np.random.randint(0,300, size=20)] * 32\n",
    "y = [[0, 1, 22, 44, 87,1]] * 32\n",
    "gen = np.array([np.random.randint(0,2, size=6) for _ in range(32)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<tf.Tensor: id=580100, shape=(), dtype=float32, numpy=-31.768637>]\n",
      "[<tf.Tensor: id=593259, shape=(), dtype=float32, numpy=-32.1079>]\n",
      "[<tf.Tensor: id=606418, shape=(), dtype=float32, numpy=-32.175247>]\n",
      "[<tf.Tensor: id=619577, shape=(), dtype=float32, numpy=-32.23389>]\n",
      "[<tf.Tensor: id=632736, shape=(), dtype=float32, numpy=-32.31731>]\n",
      "[<tf.Tensor: id=645895, shape=(), dtype=float32, numpy=-32.43171>]\n",
      "[<tf.Tensor: id=659054, shape=(), dtype=float32, numpy=-32.58477>]\n",
      "[<tf.Tensor: id=672213, shape=(), dtype=float32, numpy=-32.791733>]\n",
      "[<tf.Tensor: id=685372, shape=(), dtype=float32, numpy=-33.083458>]\n",
      "[<tf.Tensor: id=698531, shape=(), dtype=float32, numpy=-33.48859>]\n"
     ]
    }
   ],
   "source": [
    "for _ in range(10):\n",
    "    print(ptr.train_batch(X, y, gen))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: id=5131056, shape=(200,), dtype=float32, numpy=\n",
       "array([149., 192., 227., 243., 243., 168., 274., 274., 274., 274., 274.,\n",
       "       168., 274., 233., 233., 160., 274., 274., 233., 160., 274., 233.,\n",
       "       160., 233., 274., 274., 233., 274., 274., 168., 233., 233., 160.,\n",
       "       160., 274., 233., 274., 233., 160., 274., 233., 160., 233., 160.,\n",
       "       274., 274., 274., 274., 274., 168., 274., 233., 274., 274., 274.,\n",
       "       233., 233., 160., 233., 274., 274., 274., 168., 233., 274., 233.,\n",
       "       274., 233., 160., 274., 274., 274., 274., 168., 274., 233., 233.,\n",
       "       160., 233., 274., 274., 233., 160., 274., 274., 233., 233., 160.,\n",
       "       160.,  46., 233., 274., 274., 168., 274., 274., 233., 233., 160.,\n",
       "       274., 274., 233., 274., 274., 274., 168., 274., 274., 274., 274.,\n",
       "       274., 233., 274., 274., 233., 274., 274., 233.,  80.,  78., 274.,\n",
       "        78., 274., 168., 233., 233., 160., 274., 233., 274., 233., 160.,\n",
       "       233., 160., 232., 233., 233., 274., 274., 274., 168., 233., 274.,\n",
       "       233., 274., 233., 160., 160., 274., 274., 274., 233., 274., 274.,\n",
       "       233., 274., 233., 274., 274., 233., 160., 274., 233., 274., 233.,\n",
       "       160., 274., 233., 160., 274., 233., 160., 233., 160.,  46., 274.,\n",
       "       168., 274., 233., 233., 160., 274., 233., 274., 274., 168., 274.,\n",
       "       233., 274., 274., 233., 274., 274., 274., 274., 274., 233.,  80.,\n",
       "        78.,  78.], dtype=float32)>"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ptr.predict_batch(X)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ptr.evaluate(X, y)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Tensorflow",
   "language": "python",
   "name": "tensorflow"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}