abstractive-summarization/7.xueyouluo-pointer-generator-luong.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/sklearn/cross_validation.py:41: DeprecationWarning: This module was deprecated in version 0.18 in favor of the model_selection module into which all the refactored classes and functions are moved. Also note that the interface of the new CV iterators are different from that of this module. This module will be removed in 0.20.\n",
      "  \"This module will be removed in 0.20.\", DeprecationWarning)\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import collections\n",
    "from sklearn.cross_validation import train_test_split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('ctexts.json','r') as fopen:\n",
    "    ctexts = json.load(fopen)\n",
    "    \n",
    "with open('headlines.json','r') as fopen:\n",
    "    headlines = json.load(fopen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from unidecode import unidecode\n",
    "import random\n",
    "\n",
    "max_len = 500\n",
    "\n",
    "def textcleaning(string):\n",
    "    string = unidecode(string).replace('.', '. ').replace(',', ' , ')\n",
    "    string = re.sub('[^\\'\"A-Za-z\\- ]+', ' ', string)\n",
    "    string = re.sub(r'[ ]+', ' ', string.lower()).strip().split()\n",
    "    if len(string) > max_len:\n",
    "        string = random.sample(string, max_len)\n",
    "    return ' '.join(string)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 880 ms, sys: 4 ms, total: 884 ms\n",
      "Wall time: 885 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "h, c = [], []\n",
    "for i in range(len(ctexts)):\n",
    "    try:\n",
    "        c.append(textcleaning(ctexts[i]))\n",
    "        h.append(textcleaning(headlines[i]))\n",
    "    except:\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4396"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(h)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(words, n_words):\n",
    "    count = [['PAD', 0], ['GO', 1], ['EOS', 2], ['UNK', 3]]\n",
    "    count.extend(collections.Counter(words).most_common(n_words))\n",
    "    dictionary = dict()\n",
    "    for word, _ in count:\n",
    "        dictionary[word] = len(dictionary)\n",
    "    data = list()\n",
    "    unk_count = 0\n",
    "    for word in words:\n",
    "        index = dictionary.get(word, 0)\n",
    "        if index == 0:\n",
    "            unk_count += 1\n",
    "        data.append(index)\n",
    "    count[0][1] = unk_count\n",
    "    reversed_dictionary = dict(zip(dictionary.values(), dictionary.keys()))\n",
    "    return data, count, dictionary, reversed_dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab from size: 45165\n",
      "Most common words [('the', 81889), ('dot', 62474), ('comma', 61046), ('and', 34452), ('to', 33701), ('of', 31796)]\n",
      "Sample data [4, 9995, 7, 11743, 643, 12, 468, 7805, 11, 3523] ['the', 'daman', 'and', 'diu', 'administration', 'on', 'wednesday', 'withdrew', 'a', 'circular']\n",
      "filtered vocab size: 45169\n",
      "% of vocab used: 100.01%\n"
     ]
    }
   ],
   "source": [
    "concat = ' '.join(c).split()\n",
    "vocabulary_size = len(list(set(concat)))\n",
    "data, count, dictionary, rev_dictionary = build_dataset(concat, vocabulary_size)\n",
    "print('vocab from size: %d'%(vocabulary_size))\n",
    "print('Most common words', count[4:10])\n",
    "print('Sample data', data[:10], [rev_dictionary[i] for i in data[:10]])\n",
    "print('filtered vocab size:',len(dictionary))\n",
    "print(\"% of vocab used: {}%\".format(round(len(dictionary)/vocabulary_size,4)*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'daman and diu revokes mandatory rakshabandhan in offices order EOS'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "for i in range(len(h)):\n",
    "    h[i] = h[i] + ' EOS'\n",
    "h[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "GO = dictionary['GO']\n",
    "PAD = dictionary['PAD']\n",
    "EOS = dictionary['EOS']\n",
    "UNK = dictionary['UNK']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_X, test_X, train_Y, test_Y = train_test_split(c, h, test_size = 0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sent2idx(sent, vocab, UNK=UNK):\n",
    "    tokens = sent.split()\n",
    "    oovs = []\n",
    "    extend_tokens = []\n",
    "    tokenized = []\n",
    "    for token in tokens:\n",
    "        if token not in vocab:\n",
    "            tokenized.append(UNK)\n",
    "            if token not in oovs:\n",
    "                oovs.append(token)\n",
    "            extend_tokens.append(len(vocab) + oovs.index(token))\n",
    "        else:\n",
    "            extend_tokens.append(vocab[token])\n",
    "            tokenized.append(vocab[token])\n",
    "    return tokenized, extend_tokens, oovs\n",
    "\n",
    "def target2idx(sent, oovs, vocab,UNK=UNK):\n",
    "    tokens = sent.split()\n",
    "    tokenized = []\n",
    "    for token in tokens:\n",
    "        if token not in vocab:\n",
    "            if token not in oovs:\n",
    "                tokenized.append(UNK)\n",
    "            else:\n",
    "                tokenized.append(len(vocab) + oovs.index(token))\n",
    "        else:\n",
    "            tokenized.append(vocab[token])\n",
    "    return tokenized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.python.framework import ops, tensor_shape\n",
    "from tensorflow.contrib.rnn.python.ops.core_rnn_cell import _linear\n",
    "from tensorflow.python.ops.rnn_cell_impl import _zero_state_tensors\n",
    "from tensorflow.python.util import nest\n",
    "from tensorflow.python.ops import init_ops, array_ops, variable_scope\n",
    "from tensorflow.contrib.seq2seq.python.ops.attention_wrapper import (\n",
    "    _compute_attention,\n",
    ")\n",
    "\n",
    "UNK_ID = 3\n",
    "\n",
    "\n",
    "class PointerGeneratorGreedyEmbeddingHelper(\n",
    "    tf.contrib.seq2seq.GreedyEmbeddingHelper\n",
    "):\n",
    "    def __init__(self, embedding, start_tokens, end_token):\n",
    "        self.vocab_size = tf.shape(embedding)[-1]\n",
    "        super(PointerGeneratorGreedyEmbeddingHelper, self).__init__(\n",
    "            embedding, start_tokens, end_token\n",
    "        )\n",
    "\n",
    "    def sample(self, time, outputs, state, name = None):\n",
    "        \"\"\"sample for PointerGeneratorGreedyEmbeddingHelper.\"\"\"\n",
    "        del time, state\n",
    "        if not isinstance(outputs, ops.Tensor):\n",
    "            raise TypeError(\n",
    "                'Expected outputs to be a single Tensor, got: %s'\n",
    "                % type(outputs)\n",
    "            )\n",
    "        sample_ids = tf.argmax(outputs, axis = -1, output_type = tf.int32)\n",
    "        return sample_ids\n",
    "\n",
    "    def next_inputs(self, time, outputs, state, sample_ids, name = None):\n",
    "        \"\"\"next_inputs_fn for GreedyEmbeddingHelper.\"\"\"\n",
    "        del time, outputs  # unused by next_inputs_fn\n",
    "        finished = tf.equal(sample_ids, self._end_token)\n",
    "        all_finished = tf.reduce_all(finished)\n",
    "\n",
    "        condition = tf.less(sample_ids, self.vocab_size)\n",
    "        sample_ids = tf.where(\n",
    "            condition, sample_ids, tf.ones_like(sample_ids) * UNK_ID\n",
    "        )\n",
    "\n",
    "        next_inputs = tf.cond(\n",
    "            all_finished,\n",
    "            lambda: self._start_inputs,\n",
    "            lambda: self._embedding_fn(sample_ids),\n",
    "        )\n",
    "        return (finished, next_inputs, state)\n",
    "\n",
    "\n",
    "class PointerGeneratorDecoder(tf.contrib.seq2seq.BasicDecoder):\n",
    "    \"\"\"Pointer Generator sampling decoder.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        source_extend_tokens,\n",
    "        source_oov_words,\n",
    "        coverage,\n",
    "        cell,\n",
    "        helper,\n",
    "        initial_state,\n",
    "        output_layer = None,\n",
    "    ):\n",
    "        self.source_oov_words = source_oov_words\n",
    "        self.source_extend_tokens = source_extend_tokens\n",
    "        self.coverage = coverage\n",
    "        super(PointerGeneratorDecoder, self).__init__(\n",
    "            cell, helper, initial_state, output_layer\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def output_size(self):\n",
    "        # Return the cell output and the id\n",
    "        return tf.contrib.seq2seq.BasicDecoderOutput(\n",
    "            rnn_output = self._rnn_output_size() + self.source_oov_words,\n",
    "            sample_id = self._helper.sample_ids_shape,\n",
    "        )\n",
    "\n",
    "    @property\n",
    "    def output_dtype(self):\n",
    "        # Assume the dtype of the cell is the output_size structure\n",
    "        # containing the input_state's first component's dtype.\n",
    "        # Return that structure and the sample_ids_dtype from the helper.\n",
    "        dtype = nest.flatten(self._initial_state)[0].dtype\n",
    "        return tf.contrib.seq2seq.BasicDecoderOutput(\n",
    "            nest.map_structure(\n",
    "                lambda _: dtype, self._rnn_output_size() + self.source_oov_words\n",
    "            ),\n",
    "            self._helper.sample_ids_dtype,\n",
    "        )\n",
    "\n",
    "    def step(self, time, inputs, state, name = None):\n",
    "        with ops.name_scope(name, 'PGDecoderStep', (time, inputs, state)):\n",
    "            cell_outputs, cell_state = self._cell(inputs, state)\n",
    "            attention = cell_state[0].attention\n",
    "            att_cell_state = cell_state[0].cell_state\n",
    "            alignments = cell_state[0].alignments\n",
    "\n",
    "            with tf.variable_scope('calculate_pgen'):\n",
    "                p_gen = _linear([attention, inputs, att_cell_state], 1, True)\n",
    "                p_gen = tf.sigmoid(p_gen)\n",
    "\n",
    "            if self._output_layer is not None:\n",
    "                cell_outputs = self._output_layer(cell_outputs)\n",
    "\n",
    "            vocab_dist = tf.nn.softmax(cell_outputs) * p_gen\n",
    "            alignments = alignments * (1 - p_gen)\n",
    "            vocab_size = tf.shape(vocab_dist)[-1]\n",
    "            extended_vsize = vocab_size + self.source_oov_words\n",
    "            batch_size = tf.shape(vocab_dist)[0]\n",
    "            extra_zeros = tf.zeros((batch_size, self.source_oov_words))\n",
    "            vocab_dists_extended = tf.concat(\n",
    "                axis = -1, values = [vocab_dist, extra_zeros]\n",
    "            )\n",
    "\n",
    "            batch_nums = tf.range(0, limit = batch_size)\n",
    "            batch_nums = tf.expand_dims(batch_nums, 1)\n",
    "            attn_len = tf.shape(self.source_extend_tokens)[1]\n",
    "            batch_nums = tf.tile(\n",
    "                batch_nums, [1, attn_len]\n",
    "            )\n",
    "            indices = tf.stack(\n",
    "                (batch_nums, self.source_extend_tokens), axis = 2\n",
    "            )\n",
    "            shape = [batch_size, extended_vsize]\n",
    "            attn_dists_projected = tf.scatter_nd(indices, alignments, shape)\n",
    "\n",
    "            final_dists = attn_dists_projected + vocab_dists_extended\n",
    "            sample_ids = self._helper.sample(\n",
    "                time = time, outputs = final_dists, state = cell_state\n",
    "            )\n",
    "\n",
    "            (finished, next_inputs, next_state) = self._helper.next_inputs(\n",
    "                time = time,\n",
    "                outputs = cell_outputs,\n",
    "                state = cell_state,\n",
    "                sample_ids = sample_ids,\n",
    "            )\n",
    "\n",
    "            outputs = tf.contrib.seq2seq.BasicDecoderOutput(\n",
    "                final_dists, sample_ids\n",
    "            )\n",
    "            return (outputs, next_state, next_inputs, finished)\n",
    "\n",
    "\n",
    "class PointerGeneratorAttentionWrapper(tf.contrib.seq2seq.AttentionWrapper):\n",
    "    def __init__(\n",
    "        self,\n",
    "        cell,\n",
    "        attention_mechanism,\n",
    "        attention_layer_size = None,\n",
    "        alignment_history = False,\n",
    "        cell_input_fn = None,\n",
    "        output_attention = True,\n",
    "        initial_cell_state = None,\n",
    "        name = None,\n",
    "        coverage = False,\n",
    "    ):\n",
    "        super(PointerGeneratorAttentionWrapper, self).__init__(\n",
    "            cell,\n",
    "            attention_mechanism,\n",
    "            attention_layer_size,\n",
    "            alignment_history,\n",
    "            cell_input_fn,\n",
    "            output_attention,\n",
    "            initial_cell_state,\n",
    "            name,\n",
    "        )\n",
    "        self.coverage = coverage\n",
    "\n",
    "    def zero_state(self, batch_size, dtype):\n",
    "        with ops.name_scope(\n",
    "            type(self).__name__ + 'ZeroState', values = [batch_size]\n",
    "        ):\n",
    "            if self._initial_cell_state is not None:\n",
    "                cell_state = self._initial_cell_state\n",
    "            else:\n",
    "                cell_state = self._cell.zero_state(batch_size, dtype)\n",
    "            error_message = (\n",
    "                'When calling zero_state of AttentionWrapper %s: '\n",
    "                % self._base_name\n",
    "                + 'Non-matching batch sizes between the memory '\n",
    "                '(encoder output) and the requested batch size.  Are you using '\n",
    "                'the BeamSearchDecoder?  If so, make sure your encoder output has '\n",
    "                'been tiled to beam_width via tf.contrib.seq2seq.tile_batch, and '\n",
    "                'the batch_size= argument passed to zero_state is '\n",
    "                'batch_size * beam_width.'\n",
    "            )\n",
    "            with tf.control_dependencies(\n",
    "                self._batch_size_checks(batch_size, error_message)\n",
    "            ):\n",
    "                cell_state = nest.map_structure(\n",
    "                    lambda s: tf.identity(s, name = 'checked_cell_state'),\n",
    "                    cell_state,\n",
    "                )\n",
    "            return tf.contrib.seq2seq.AttentionWrapperState(\n",
    "                cell_state = cell_state,\n",
    "                time = tf.zeros([], dtype = tf.int32),\n",
    "                attention = _zero_state_tensors(\n",
    "                    self._attention_layer_size, batch_size, dtype\n",
    "                ),\n",
    "                alignments = self._item_or_tuple(\n",
    "                    attention_mechanism.initial_alignments(batch_size, dtype)\n",
    "                    for attention_mechanism in self._attention_mechanisms\n",
    "                ),\n",
    "                attention_state = self._item_or_tuple(\n",
    "                    attention_mechanism.initial_state(batch_size, dtype)\n",
    "                    for attention_mechanism in self._attention_mechanisms\n",
    "                ),\n",
    "                # since we need to read the alignment history several times, so we need set clear_after_read to False\n",
    "                alignment_history = self._item_or_tuple(\n",
    "                    tf.TensorArray(\n",
    "                        dtype = dtype,\n",
    "                        size = 0,\n",
    "                        clear_after_read = False,\n",
    "                        dynamic_size = True,\n",
    "                    )\n",
    "                    if self._alignment_history\n",
    "                    else ()\n",
    "                    for _ in self._attention_mechanisms\n",
    "                ),\n",
    "            )\n",
    "\n",
    "    def call(self, inputs, state):\n",
    "        if not isinstance(state, tf.contrib.seq2seq.AttentionWrapperState):\n",
    "            raise TypeError(\n",
    "                'Expected state to be instance of AttentionWrapperState. '\n",
    "                'Received type %s instead.' % type(state)\n",
    "            )\n",
    "\n",
    "        # Step 1: Calculate the true inputs to the cell based on the\n",
    "        # previous attention value.\n",
    "        cell_inputs = self._cell_input_fn(inputs, state.attention)\n",
    "        cell_state = state.cell_state\n",
    "        cell_output, next_cell_state = self._cell(cell_inputs, cell_state)\n",
    "\n",
    "        cell_batch_size = cell_output.shape[0].value or tf.shape(cell_output)[0]\n",
    "        error_message = (\n",
    "            'When applying AttentionWrapper %s: ' % self.name\n",
    "            + 'Non-matching batch sizes between the memory '\n",
    "            '(encoder output) and the query (decoder output).  Are you using '\n",
    "            'the BeamSearchDecoder?  You may need to tile your memory input via '\n",
    "            'the tf.contrib.seq2seq.tile_batch function with argument '\n",
    "            'multiple=beam_width.'\n",
    "        )\n",
    "        with tf.control_dependencies(\n",
    "            self._batch_size_checks(cell_batch_size, error_message)\n",
    "        ):\n",
    "            cell_output = tf.identity(cell_output, name = 'checked_cell_output')\n",
    "\n",
    "        if self._is_multi:\n",
    "            previous_alignments = state.alignments\n",
    "            previous_alignment_history = state.alignment_history\n",
    "        else:\n",
    "            previous_alignments = [state.alignments]\n",
    "            previous_alignment_history = [state.alignment_history]\n",
    "\n",
    "        all_alignments = []\n",
    "        all_attentions = []\n",
    "        all_histories = []\n",
    "\n",
    "        for i, attention_mechanism in enumerate(self._attention_mechanisms):\n",
    "            print(attention_mechanism)\n",
    "            if self.coverage:\n",
    "                # if we use coverage mode, previous alignments is coverage vector\n",
    "                # alignment history stack has shape:  decoder time * batch * atten_len\n",
    "                # convert it to coverage vector\n",
    "                previous_alignments[i] = tf.cond(\n",
    "                    previous_alignment_history[i].size() > 0,\n",
    "                    lambda: tf.reduce_sum(\n",
    "                        tf.transpose(\n",
    "                            previous_alignment_history[i].stack(), [1, 2, 0]\n",
    "                        ),\n",
    "                        axis = 2,\n",
    "                    ),\n",
    "                    lambda: tf.zeros_like(previous_alignments[i]),\n",
    "                )\n",
    "            # debug\n",
    "            # previous_alignments[i] = tf.Print(previous_alignments[i],[previous_alignment_history[i].size(), tf.shape(previous_alignments[i]),previous_alignments[i]],message=\"atten wrapper:\")\n",
    "            attention, alignments, next_attention_state = _compute_attention(\n",
    "                attention_mechanism,\n",
    "                cell_output,\n",
    "                previous_alignments[i],\n",
    "                self._attention_layers[i] if self._attention_layers else None,\n",
    "            )\n",
    "            alignment_history = (\n",
    "                previous_alignment_history[i].write(state.time, alignments)\n",
    "                if self._alignment_history\n",
    "                else ()\n",
    "            )\n",
    "\n",
    "            all_alignments.append(alignments)\n",
    "            all_histories.append(alignment_history)\n",
    "            all_attentions.append(attention)\n",
    "\n",
    "        attention = tf.concat(all_attentions, 1)\n",
    "        next_state = tf.contrib.seq2seq.AttentionWrapperState(\n",
    "            time = state.time + 1,\n",
    "            cell_state = next_cell_state,\n",
    "            attention = attention,\n",
    "            alignments = self._item_or_tuple(all_alignments),\n",
    "            attention_state = self._item_or_tuple(all_alignments),\n",
    "            alignment_history = self._item_or_tuple(all_histories),\n",
    "        )\n",
    "\n",
    "        if self._output_attention:\n",
    "            return attention, next_state\n",
    "        else:\n",
    "            return cell_output, next_state\n",
    "\n",
    "\n",
    "def _pg_luong_score(query, keys, scale, coverage, coverage_vector):\n",
    "    depth = query.get_shape()[-1]\n",
    "    key_units = keys.get_shape()[-1]\n",
    "    num_units = keys.shape[2].value or tf.shape(keys)[2]\n",
    "    if depth != key_units:\n",
    "        raise ValueError(\n",
    "            'Incompatible or unknown inner dimensions between query and keys.  '\n",
    "            'Query (%s) has units: %s.  Keys (%s) have units: %s.  '\n",
    "            \"Perhaps you need to set num_units to the keys' dimension (%s)?\"\n",
    "            % (query, depth, keys, key_units, key_units)\n",
    "        )\n",
    "    dtype = query.dtype\n",
    "    query = array_ops.expand_dims(query, 1)\n",
    "    score = tf.matmul(query, keys, transpose_b = True)\n",
    "    score = array_ops.squeeze(score, [1])\n",
    "    if scale:\n",
    "        g = variable_scope.get_variable(\n",
    "            'attention_g',\n",
    "            dtype = dtype,\n",
    "            initializer = init_ops.ones_initializer,\n",
    "            shape = (),\n",
    "        )\n",
    "        score = g * score\n",
    "    if coverage:\n",
    "        w_c = tf.get_variable('coverage_w', [num_units], dtype = dtype)\n",
    "        coverage_vector = tf.expand_dims(coverage_vector, -1)\n",
    "        score = tf.expand_dims(score, -1)\n",
    "        return tf.reduce_sum(score + coverage_vector, [2])\n",
    "\n",
    "\n",
    "class PointerGeneratorLuongAttention(tf.contrib.seq2seq.LuongAttention):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_units,\n",
    "        memory,\n",
    "        memory_sequence_length = None,\n",
    "        probability_fn = None,\n",
    "        score_mask_value = float('-inf'),\n",
    "        name = 'PointerGeneratorLuongAttention',\n",
    "        coverage = False,\n",
    "        scale = True,\n",
    "    ):\n",
    "        super(PointerGeneratorLuongAttention, self).__init__(\n",
    "            num_units = num_units,\n",
    "            memory = memory,\n",
    "            scale = scale,\n",
    "            memory_sequence_length = memory_sequence_length,\n",
    "            probability_fn = probability_fn,\n",
    "            score_mask_value = score_mask_value,\n",
    "            name = name,\n",
    "        )\n",
    "        self.coverage = coverage\n",
    "\n",
    "    def __call__(self, query, state):\n",
    "        with tf.variable_scope(\n",
    "            None, 'pointer_generator_luong_attention', [query]\n",
    "        ):\n",
    "\n",
    "            score = _pg_luong_score(\n",
    "                query, self._keys, self._scale, self.coverage, state\n",
    "            )\n",
    "        alignments = self._probability_fn(score, state)\n",
    "        next_state = alignments\n",
    "        return alignments, next_state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Summarization:\n",
    "    def __init__(self, size_layer, num_layers, embedded_size, dict_size):\n",
    "        \n",
    "        def lstm_cell(reuse=False):\n",
    "            return tf.nn.rnn_cell.GRUCell(size_layer, reuse=reuse)\n",
    "        \n",
    "        self.X = tf.placeholder(tf.int32, [None, None])\n",
    "        self.Y = tf.placeholder(tf.int32, [None, None])\n",
    "        self.X_seq_len = tf.count_nonzero(self.X, 1, dtype=tf.int32)\n",
    "        self.Y_seq_len = tf.count_nonzero(self.Y, 1, dtype=tf.int32)\n",
    "        batch_size = tf.shape(self.X)[0]\n",
    "        self.source_oov_words = tf.placeholder(tf.int32, shape=[])\n",
    "        self.source_extend_tokens = tf.placeholder(tf.int32, shape=[None, None])\n",
    "        main = tf.strided_slice(self.Y, [0, 0], [batch_size, -1], [1, 1])\n",
    "        decoder_input = tf.concat([tf.fill([batch_size, 1], GO), main], 1)\n",
    "        \n",
    "        condition = tf.less(decoder_input, dict_size)\n",
    "        self.decoder_input = decoder_input\n",
    "        self.decoder_input_length = self.Y_seq_len\n",
    "        self.predict_count = tf.reduce_sum(self.decoder_input_length)\n",
    "        \n",
    "        embeddings = tf.Variable(tf.random_uniform([dict_size, embedded_size], -1, 1))\n",
    "        \n",
    "        encoder_embedded = tf.nn.embedding_lookup(embeddings, self.X)\n",
    "        encoder_cells = tf.nn.rnn_cell.MultiRNNCell([lstm_cell() for _ in range(num_layers)])\n",
    "        self.encoder_out, self.encoder_state = tf.nn.dynamic_rnn(cell = encoder_cells, \n",
    "                                                                 inputs = encoder_embedded, \n",
    "                                                                 sequence_length = self.X_seq_len,\n",
    "                                                                 dtype = tf.float32)\n",
    "        self.decode_initial_state = self.encoder_state[-1]\n",
    "        print(self.decode_initial_state)\n",
    "        \n",
    "        atten_mech = PointerGeneratorLuongAttention(\n",
    "                size_layer, self.encoder_out, memory_sequence_length=self.X_seq_len,\n",
    "        coverage = True)\n",
    "        decoder_cells = [lstm_cell() for _ in range(num_layers)]\n",
    "        decoder_cells[0] = PointerGeneratorAttentionWrapper(\n",
    "                cell=decoder_cells[0],\n",
    "                attention_mechanism=atten_mech,\n",
    "                attention_layer_size=size_layer,\n",
    "                alignment_history = True,\n",
    "                coverage = True\n",
    "            )\n",
    "        initial_state = [self.decode_initial_state for i in range(num_layers)]\n",
    "        attention_cell_state = decoder_cells[0].zero_state(\n",
    "                dtype=tf.float32, batch_size=batch_size)\n",
    "        initial_state[0] = attention_cell_state.clone(\n",
    "                cell_state=initial_state[0])\n",
    "        self.initial_state = tuple(initial_state)\n",
    "        decoder_cells = tf.contrib.rnn.MultiRNNCell(decoder_cells)\n",
    "        \n",
    "        decoded = tf.nn.embedding_lookup(embeddings, self.decoder_input)\n",
    "        \n",
    "        training_helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "            decoded,\n",
    "            self.decoder_input_length\n",
    "        )\n",
    "        dense_layer = tf.layers.Dense(dict_size)\n",
    "        \n",
    "        training_decoder = PointerGeneratorDecoder(\n",
    "            source_extend_tokens = self.source_extend_tokens,\n",
    "            source_oov_words = self.source_oov_words,\n",
    "            coverage = True,\n",
    "            cell=decoder_cells,\n",
    "            helper=training_helper,\n",
    "            initial_state=self.initial_state,\n",
    "            output_layer=dense_layer\n",
    "        )\n",
    "        \n",
    "        maxlen = tf.reduce_max(self.decoder_input_length)\n",
    "        train_dec_outputs, train_dec_last_state, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "            training_decoder,\n",
    "            output_time_major=False,\n",
    "            impute_finished=True,\n",
    "            maximum_iterations=maxlen,\n",
    "            swap_memory=True)\n",
    "        logits = train_dec_outputs.rnn_output\n",
    "        \n",
    "        masks = tf.sequence_mask(\n",
    "            self.decoder_input_length, maxlen, \n",
    "            dtype=tf.float32)\n",
    "        \n",
    "        targets = tf.slice(self.Y, [0, 0], [-1, maxlen])\n",
    "        i1, i2 = tf.meshgrid(tf.range(batch_size),\n",
    "                     tf.range(maxlen), indexing=\"ij\")\n",
    "        indices = tf.stack((i1,i2,targets),axis=2)\n",
    "        probs = tf.gather_nd(logits, indices)\n",
    "        probs = tf.where(tf.less_equal(probs,0),tf.ones_like(probs)*1e-10,probs)\n",
    "        crossent = -tf.log(probs)\n",
    "        self.cost = tf.reduce_sum(crossent * masks) / tf.to_float(batch_size)\n",
    "        alignment_history = train_dec_last_state[0].alignment_history.stack()\n",
    "        alignment_history = tf.transpose(alignment_history,[1,2,0])\n",
    "        coverage_loss = tf.minimum(alignment_history,tf.cumsum(alignment_history, axis=2, exclusive=True))\n",
    "        self.coverage_loss = tf.reduce_sum(coverage_loss / tf.to_float(batch_size))\n",
    "        self.cost = self.cost + self.coverage_loss\n",
    "        \n",
    "        self.optimizer = tf.train.AdamOptimizer().minimize(self.cost)\n",
    "        \n",
    "        helper = PointerGeneratorGreedyEmbeddingHelper(\n",
    "            embedding=embeddings,\n",
    "            start_tokens=tf.tile(tf.constant([GO], dtype=tf.int32), [batch_size]),\n",
    "            end_token=EOS\n",
    "        )\n",
    "        \n",
    "        inference_decoder = PointerGeneratorDecoder(\n",
    "            source_extend_tokens = self.source_extend_tokens,\n",
    "            source_oov_words = self.source_oov_words,\n",
    "            coverage = True,\n",
    "            cell=decoder_cells,\n",
    "            helper=helper,\n",
    "            initial_state=self.initial_state,\n",
    "            output_layer=dense_layer\n",
    "        )\n",
    "        \n",
    "        dec_outputs, dec_last_state, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "            inference_decoder,\n",
    "            output_time_major=False,\n",
    "            maximum_iterations=tf.reduce_max(self.X_seq_len),\n",
    "            swap_memory=True)\n",
    "        \n",
    "        self.beam_predictions = dec_outputs.sample_id\n",
    "        print(self.beam_predictions)\n",
    "        \n",
    "        masks = tf.sequence_mask(self.Y_seq_len, tf.reduce_max(self.Y_seq_len), dtype=tf.float32)\n",
    "        y_t = tf.argmax(logits,axis=2)\n",
    "        y_t = tf.cast(y_t, tf.int32)\n",
    "        self.prediction = tf.boolean_mask(y_t, masks)\n",
    "        mask_label = tf.boolean_mask(self.Y, masks)\n",
    "        correct_pred = tf.equal(self.prediction, mask_label)\n",
    "        correct_index = tf.cast(correct_pred, tf.float32)\n",
    "        self.accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "size_layer = 128\n",
    "num_layers = 2\n",
    "embedded_size = 128\n",
    "batch_size = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor(\"rnn/while/Exit_4:0\", shape=(?, 128), dtype=float32)\n",
      "<__main__.PointerGeneratorLuongAttention object at 0x7fd0bc606ba8>\n",
      "<__main__.PointerGeneratorLuongAttention object at 0x7fd0bc606ba8>\n",
      "Tensor(\"decoder_1/transpose_1:0\", shape=(?, ?), dtype=int32)\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Summarization(size_layer, num_layers, embedded_size, len(dictionary))\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "from keras.preprocessing.sequence import pad_sequences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batching(X, Y):\n",
    "    s_, es_, oovs_, target_ = [], [], [], []\n",
    "    for x, y in zip(X, Y):\n",
    "        s,es,oovs = sent2idx(x, dictionary)\n",
    "        target = target2idx(y, oovs,dictionary)\n",
    "        s_.append(s)\n",
    "        es_.append(es)\n",
    "        oovs_.append(oovs)\n",
    "        target_.append(target)\n",
    "    s_ = pad_sequences(s_,padding='post')\n",
    "    es_ = pad_sequences(es_,padding='post')\n",
    "    target_ = pad_sequences(target_,padding='post')\n",
    "    maxlen = max([len(o) for o in oovs_])\n",
    "    return s_, es_, target_, maxlen"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:45<00:00,  1.95it/s, accuracy=0.151, cost=78.9] \n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  6.71it/s, accuracy=0.174, cost=72.6]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, avg loss: 73.637576, avg accuracy: 0.117944\n",
      "epoch: 0, avg loss test: 68.686492, avg accuracy test: 0.154677\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:44<00:00,  1.96it/s, accuracy=0.233, cost=55.8] \n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  6.55it/s, accuracy=0.135, cost=73.9]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, avg loss: 61.329693, avg accuracy: 0.188649\n",
      "epoch: 1, avg loss test: 64.766316, avg accuracy test: 0.186535\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:41<00:00,  2.31it/s, accuracy=0.282, cost=39.8]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.02it/s, accuracy=0.136, cost=68.2]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 2, avg loss: 53.691719, avg accuracy: 0.256157\n",
      "epoch: 2, avg loss test: 64.007127, avg accuracy test: 0.198297\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:36<00:00,  2.22it/s, accuracy=0.286, cost=45.9]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.14it/s, accuracy=0.2, cost=61.4]  \n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 3, avg loss: 46.313700, avg accuracy: 0.331684\n",
      "epoch: 3, avg loss test: 65.356595, avg accuracy test: 0.188881\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:36<00:00,  2.04it/s, accuracy=0.429, cost=34.1]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.08it/s, accuracy=0.157, cost=62.7]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 4, avg loss: 39.288931, avg accuracy: 0.423175\n",
      "epoch: 4, avg loss test: 68.959602, avg accuracy test: 0.179051\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:37<00:00,  2.01it/s, accuracy=0.386, cost=39.5]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  6.90it/s, accuracy=0.214, cost=68.6]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 5, avg loss: 33.383819, avg accuracy: 0.508171\n",
      "epoch: 5, avg loss test: 72.910164, avg accuracy test: 0.176751\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:37<00:00,  2.02it/s, accuracy=0.452, cost=30.7]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  6.98it/s, accuracy=0.114, cost=85.8] \n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 6, avg loss: 28.742008, avg accuracy: 0.572048\n",
      "epoch: 6, avg loss test: 76.335111, avg accuracy test: 0.167679\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:37<00:00,  2.02it/s, accuracy=0.646, cost=25.5]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.25it/s, accuracy=0.226, cost=66.8]\n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 7, avg loss: 25.249980, avg accuracy: 0.611214\n",
      "epoch: 7, avg loss test: 81.496095, avg accuracy test: 0.168926\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:37<00:00,  2.14it/s, accuracy=0.558, cost=26.1]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.08it/s, accuracy=0.176, cost=76.1] \n",
      "train minibatch loop:   0%|          | 0/440 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 8, avg loss: 22.495210, avg accuracy: 0.636025\n",
      "epoch: 8, avg loss test: 83.243105, avg accuracy test: 0.173153\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "train minibatch loop: 100%|██████████| 440/440 [03:36<00:00,  2.05it/s, accuracy=0.617, cost=22.1]\n",
      "test minibatch loop: 100%|██████████| 110/110 [00:15<00:00,  7.07it/s, accuracy=0.17, cost=92.3] "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 9, avg loss: 19.937910, avg accuracy: 0.661205\n",
      "epoch: 9, avg loss test: 88.395862, avg accuracy test: 0.165144\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "from sklearn.utils import shuffle\n",
    "import time\n",
    "\n",
    "for EPOCH in range(10):\n",
    "    lasttime = time.time()\n",
    "    total_loss, total_accuracy, total_loss_test, total_accuracy_test = 0, 0, 0, 0\n",
    "    train_X, train_Y = shuffle(train_X, train_Y)\n",
    "    test_X, test_Y = shuffle(test_X, test_Y)\n",
    "    pbar = tqdm(range(0, len(train_X), batch_size), desc='train minibatch loop')\n",
    "    for k in pbar:\n",
    "        index = min(k+batch_size,len(train_X))\n",
    "        batch_x, batch_es, batch_y, maxlen = batching(train_X[k: index],\n",
    "                                                     train_Y[k: index])\n",
    "        acc, loss, _ = sess.run([model.accuracy, model.cost, model.optimizer], \n",
    "                                      feed_dict={model.X:batch_x,\n",
    "                                                 model.source_extend_tokens:batch_es,\n",
    "                                                model.Y:batch_y,\n",
    "                                                model.source_oov_words:maxlen})\n",
    "        total_loss += loss\n",
    "        total_accuracy += acc\n",
    "        pbar.set_postfix(cost=loss, accuracy = acc)\n",
    "        \n",
    "    pbar = tqdm(range(0, len(test_X), batch_size), desc='test minibatch loop')\n",
    "    for k in pbar:\n",
    "        index = min(k+batch_size,len(test_X))\n",
    "        batch_x, batch_es, batch_y, maxlen = batching(test_X[k: index],\n",
    "                                                     test_Y[k: index])\n",
    "        acc, loss = sess.run([model.accuracy, model.cost], \n",
    "                                      feed_dict={model.X:batch_x,\n",
    "                                                 model.source_extend_tokens:batch_es,\n",
    "                                                model.Y:batch_y,\n",
    "                                                model.source_oov_words:maxlen})\n",
    "        total_loss_test += loss\n",
    "        total_accuracy_test += acc\n",
    "        pbar.set_postfix(cost=loss, accuracy = acc)\n",
    "        \n",
    "    total_loss /= (len(train_X) / batch_size)\n",
    "    total_accuracy /= (len(train_X) / batch_size)\n",
    "    total_loss_test /= (len(test_X) / batch_size)\n",
    "    total_accuracy_test /= (len(test_X) / batch_size)\n",
    "        \n",
    "    print('epoch: %d, avg loss: %f, avg accuracy: %f'%(EPOCH, total_loss, total_accuracy))\n",
    "    print('epoch: %d, avg loss test: %f, avg accuracy test: %f'%(EPOCH, total_loss_test, total_accuracy_test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_x, batch_es, batch_y, maxlen = batching(test_X[:1], test_Y[:1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "out = [rev_dictionary[i] for i in sess.run(model.beam_predictions, feed_dict = {model.X: batch_x,\n",
    "                                             model.Y: batch_y,\n",
    "                                             model.source_extend_tokens:batch_es,\n",
    "                                             model.source_oov_words:maxlen})[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "def f7(seq):\n",
    "    seen = set()\n",
    "    seen_add = seen.add\n",
    "    return ' '.join([x for x in seq if not (x in seen or seen_add(x))])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'baby birth boy hospital nursing EOS'"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f7(out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['gave birth to boy comma hospital handed us girl delhi couple EOS']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_Y[:1]"
   ]
  },
  {
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}