text-augmentation/4.gru-vae-beam-bahdanau-tfprob.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxlen = 20\n",
    "max_vocab = 20000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "word2idx = tf.keras.datasets.imdb.get_word_index()\n",
    "word2idx = {k: (v + 4) for k, v in word2idx.items()}\n",
    "word2idx['<PAD>'] = 0\n",
    "word2idx['<START>'] = 1\n",
    "word2idx['<UNK>'] = 2\n",
    "word2idx['<END>'] = 3\n",
    "idx2word = {i: w for w, i in word2idx.items()}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "(train_X, _), (test_X, _) = tf.contrib.keras.datasets.imdb.load_data(num_words = max_vocab, index_from= 4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.concatenate([train_X, test_X])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.concatenate((tf.keras.preprocessing.sequence.pad_sequences(\n",
    "                            X, maxlen, truncating='post', padding='post'),\n",
    "                        tf.keras.preprocessing.sequence.pad_sequences(\n",
    "                            X, maxlen, truncating='pre', padding='post')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "Y_input = X[:]\n",
    "Y_output = np.concatenate([X[:, 1:], np.full([X.shape[0], 1], word2idx['<END>'])], 1)\n",
    "X = X[:, 1:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((100000, 19), (100000, 20), (100000, 20))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape, Y_input.shape, Y_output.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.python.util import nest\n",
    "from tensorflow.contrib.seq2seq.python.ops.beam_search_decoder import _beam_search_step\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class ModifiedBasicDecoder(tf.contrib.seq2seq.BasicDecoder):\n",
    "    def __init__(self, cell, helper, initial_state, concat_z, output_layer=None):\n",
    "        super().__init__(cell, helper, initial_state, output_layer = output_layer)\n",
    "        self.z = concat_z\n",
    "\n",
    "    def initialize(self, name=None):\n",
    "        (finished, first_inputs, initial_state) =  self._helper.initialize() + (self._initial_state,)\n",
    "        first_inputs = tf.concat([first_inputs, self.z], -1)\n",
    "        return (finished, first_inputs, initial_state)\n",
    "\n",
    "    def step(self, time, inputs, state, name=None):\n",
    "        with tf.name_scope(name, \"BasicDecoderStep\", (time, inputs, state)):\n",
    "            cell_outputs, cell_state = self._cell(inputs, state)\n",
    "        print(self._output_layer)\n",
    "        if self._output_layer is not None:\n",
    "            cell_outputs = self._output_layer(cell_outputs)\n",
    "        print(cell_outputs)\n",
    "        sample_ids = self._helper.sample(\n",
    "            time=time, outputs=cell_outputs, state=cell_state)\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",
    "        outputs = tf.contrib.seq2seq.BasicDecoderOutput(cell_outputs, sample_ids)\n",
    "        next_inputs = tf.concat([next_inputs, self.z], -1)\n",
    "        return (outputs, next_state, next_inputs, finished)\n",
    "\n",
    "\n",
    "class ModifiedBeamSearchDecoder(tf.contrib.seq2seq.BeamSearchDecoder):\n",
    "    def __init__(self,\n",
    "                 cell,\n",
    "                 embedding,\n",
    "                 start_tokens,\n",
    "                 end_token,\n",
    "                 initial_state,\n",
    "                 beam_width,\n",
    "                 concat_z,\n",
    "                 output_layer=None,\n",
    "                 length_penalty_weight=0.0):\n",
    "        super().__init__(cell, embedding, start_tokens, end_token, initial_state, beam_width, output_layer, length_penalty_weight)\n",
    "        self.z = concat_z\n",
    "\n",
    "    def initialize(self, name=None):\n",
    "        finished, start_inputs = self._finished, self._start_inputs\n",
    "\n",
    "        start_inputs = tf.concat([start_inputs, self.z], -1)\n",
    "\n",
    "        log_probs = tf.one_hot(  # shape(batch_sz, beam_sz)\n",
    "            tf.zeros([self._batch_size], dtype=tf.int32),\n",
    "            depth=self._beam_width,\n",
    "            on_value=0.0,\n",
    "            off_value=-np.Inf,\n",
    "            dtype=nest.flatten(self._initial_cell_state)[0].dtype)\n",
    "\n",
    "        initial_state = tf.contrib.seq2seq.BeamSearchDecoderState(\n",
    "            cell_state=self._initial_cell_state,\n",
    "            log_probs=log_probs,\n",
    "            finished=finished,\n",
    "            lengths=tf.zeros(\n",
    "                [self._batch_size, self._beam_width], dtype=tf.int64),\n",
    "            accumulated_attention_probs=())\n",
    "\n",
    "        return (finished, start_inputs, initial_state)\n",
    "\n",
    "    def step(self, time, inputs, state, name=None):\n",
    "        batch_size = self._batch_size\n",
    "        beam_width = self._beam_width\n",
    "        end_token = self._end_token\n",
    "        length_penalty_weight = self._length_penalty_weight\n",
    "\n",
    "        with tf.name_scope(name, \"BeamSearchDecoderStep\", (time, inputs, state)):\n",
    "            cell_state = state.cell_state\n",
    "            inputs = nest.map_structure(\n",
    "                lambda inp: self._merge_batch_beams(inp, s=inp.shape[2:]), inputs)\n",
    "            cell_state = nest.map_structure(self._maybe_merge_batch_beams, cell_state,\n",
    "                                            self._cell.state_size)\n",
    "            cell_outputs, next_cell_state = self._cell(inputs, cell_state)\n",
    "            cell_outputs = nest.map_structure(\n",
    "                lambda out: self._split_batch_beams(out, out.shape[1:]), cell_outputs)\n",
    "            next_cell_state = nest.map_structure(\n",
    "                self._maybe_split_batch_beams, next_cell_state, self._cell.state_size)\n",
    "            print(self._output_layer)\n",
    "            if self._output_layer is not None:\n",
    "                cell_outputs = self._output_layer(cell_outputs)\n",
    "\n",
    "            beam_search_output, beam_search_state = _beam_search_step(\n",
    "                time=time,\n",
    "                logits=cell_outputs,\n",
    "                next_cell_state=next_cell_state,\n",
    "                beam_state=state,\n",
    "                batch_size=batch_size,\n",
    "                beam_width=beam_width,\n",
    "                end_token=end_token,\n",
    "                length_penalty_weight=length_penalty_weight,\n",
    "                coverage_penalty_weight = 0.0)\n",
    "\n",
    "            finished = beam_search_state.finished\n",
    "            sample_ids = beam_search_output.predicted_ids\n",
    "            next_inputs = tf.cond(\n",
    "                tf.reduce_all(finished), lambda: self._start_inputs,\n",
    "                lambda: self._embedding_fn(sample_ids))\n",
    "\n",
    "            next_inputs = tf.concat([next_inputs, self.z], -1)\n",
    "\n",
    "        return (beam_search_output, beam_search_state, next_inputs, finished)\n",
    "    \n",
    "from tensorflow.python.util import nest\n",
    "from tensorflow.contrib.seq2seq.python.ops.beam_search_decoder import (\n",
    "    _beam_search_step,\n",
    ")\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "\n",
    "\n",
    "class ModifiedBasicDecoder(tf.contrib.seq2seq.BasicDecoder):\n",
    "    def __init__(\n",
    "        self, cell, helper, initial_state, concat_z, output_layer = None\n",
    "    ):\n",
    "        super().__init__(\n",
    "            cell, helper, initial_state, output_layer = output_layer\n",
    "        )\n",
    "        self.z = concat_z\n",
    "\n",
    "    def initialize(self, name = None):\n",
    "        (finished, first_inputs, initial_state) = self._helper.initialize() + (\n",
    "            self._initial_state,\n",
    "        )\n",
    "        first_inputs = tf.concat([first_inputs, self.z], -1)\n",
    "        return (finished, first_inputs, initial_state)\n",
    "\n",
    "    def step(self, time, inputs, state, name = None):\n",
    "        with tf.name_scope(name, 'BasicDecoderStep', (time, inputs, state)):\n",
    "            cell_outputs, cell_state = self._cell(inputs, state)\n",
    "        if self._output_layer is not None:\n",
    "            cell_outputs = self._output_layer(cell_outputs)\n",
    "        sample_ids = self._helper.sample(\n",
    "            time = time, outputs = cell_outputs, state = cell_state\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",
    "        outputs = tf.contrib.seq2seq.BasicDecoderOutput(\n",
    "            cell_outputs, sample_ids\n",
    "        )\n",
    "        next_inputs = tf.concat([next_inputs, self.z], -1)\n",
    "        return (outputs, next_state, next_inputs, finished)\n",
    "\n",
    "\n",
    "class ModifiedBeamSearchDecoder(tf.contrib.seq2seq.BeamSearchDecoder):\n",
    "    def __init__(\n",
    "        self,\n",
    "        cell,\n",
    "        embedding,\n",
    "        start_tokens,\n",
    "        end_token,\n",
    "        initial_state,\n",
    "        beam_width,\n",
    "        concat_z,\n",
    "        output_layer = None,\n",
    "        length_penalty_weight = 0.0,\n",
    "    ):\n",
    "        super().__init__(\n",
    "            cell,\n",
    "            embedding,\n",
    "            start_tokens,\n",
    "            end_token,\n",
    "            initial_state,\n",
    "            beam_width,\n",
    "            output_layer,\n",
    "            length_penalty_weight,\n",
    "        )\n",
    "        self.z = concat_z\n",
    "\n",
    "    def initialize(self, name = None):\n",
    "        finished, start_inputs = self._finished, self._start_inputs\n",
    "\n",
    "        start_inputs = tf.concat([start_inputs, self.z], -1)\n",
    "\n",
    "        log_probs = tf.one_hot(  # shape(batch_sz, beam_sz)\n",
    "            tf.zeros([self._batch_size], dtype = tf.int32),\n",
    "            depth = self._beam_width,\n",
    "            on_value = 0.0,\n",
    "            off_value = -np.Inf,\n",
    "            dtype = nest.flatten(self._initial_cell_state)[0].dtype,\n",
    "        )\n",
    "\n",
    "        initial_state = tf.contrib.seq2seq.BeamSearchDecoderState(\n",
    "            cell_state = self._initial_cell_state,\n",
    "            log_probs = log_probs,\n",
    "            finished = finished,\n",
    "            lengths = tf.zeros(\n",
    "                [self._batch_size, self._beam_width], dtype = tf.int64\n",
    "            ),\n",
    "            accumulated_attention_probs = (),\n",
    "        )\n",
    "\n",
    "        return (finished, start_inputs, initial_state)\n",
    "\n",
    "    def step(self, time, inputs, state, name = None):\n",
    "        batch_size = self._batch_size\n",
    "        beam_width = self._beam_width\n",
    "        end_token = self._end_token\n",
    "        length_penalty_weight = self._length_penalty_weight\n",
    "\n",
    "        with tf.name_scope(\n",
    "            name, 'BeamSearchDecoderStep', (time, inputs, state)\n",
    "        ):\n",
    "            cell_state = state.cell_state\n",
    "            inputs = nest.map_structure(\n",
    "                lambda inp: self._merge_batch_beams(inp, s = inp.shape[2:]),\n",
    "                inputs,\n",
    "            )\n",
    "            cell_state = nest.map_structure(\n",
    "                self._maybe_merge_batch_beams, cell_state, self._cell.state_size\n",
    "            )\n",
    "            cell_outputs, next_cell_state = self._cell(inputs, cell_state)\n",
    "            cell_outputs = nest.map_structure(\n",
    "                lambda out: self._split_batch_beams(out, out.shape[1:]),\n",
    "                cell_outputs,\n",
    "            )\n",
    "            next_cell_state = nest.map_structure(\n",
    "                self._maybe_split_batch_beams,\n",
    "                next_cell_state,\n",
    "                self._cell.state_size,\n",
    "            )\n",
    "            if self._output_layer is not None:\n",
    "                cell_outputs = self._output_layer(cell_outputs)\n",
    "\n",
    "            beam_search_output, beam_search_state = _beam_search_step(\n",
    "                time = time,\n",
    "                logits = cell_outputs,\n",
    "                next_cell_state = next_cell_state,\n",
    "                beam_state = state,\n",
    "                batch_size = batch_size,\n",
    "                beam_width = beam_width,\n",
    "                end_token = end_token,\n",
    "                length_penalty_weight = length_penalty_weight,\n",
    "                coverage_penalty_weight = 0.0,\n",
    "            )\n",
    "\n",
    "            finished = beam_search_state.finished\n",
    "            sample_ids = beam_search_output.predicted_ids\n",
    "            next_inputs = tf.cond(\n",
    "                tf.reduce_all(finished),\n",
    "                lambda: self._start_inputs,\n",
    "                lambda: self._embedding_fn(sample_ids),\n",
    "            )\n",
    "\n",
    "            next_inputs = tf.concat([next_inputs, self.z], -1)\n",
    "\n",
    "        return (beam_search_output, beam_search_state, next_inputs, finished)\n",
    "\n",
    "\n",
    "def _bahdanau_score(processed_query, keys, normalize):\n",
    "    dtype = processed_query.dtype\n",
    "    num_units = keys.shape[2].value or tf.shape(keys)[2]\n",
    "    processed_query = tf.expand_dims(processed_query, 1)\n",
    "    v = tf.get_variable('attention_v', [num_units], dtype = dtype)\n",
    "    if normalize:\n",
    "        # Scalar used in weight normalization\n",
    "        g = variable_scope.get_variable(\n",
    "            'attention_g',\n",
    "            dtype = dtype,\n",
    "            initializer = init_ops.constant_initializer(\n",
    "                math.sqrt((1.0 / num_units))\n",
    "            ),\n",
    "            shape = (),\n",
    "        )\n",
    "        # Bias added prior to the nonlinearity\n",
    "        b = variable_scope.get_variable(\n",
    "            'attention_b',\n",
    "            [num_units],\n",
    "            dtype = dtype,\n",
    "            initializer = init_ops.zeros_initializer(),\n",
    "        )\n",
    "        # normed_v = g * v / ||v||\n",
    "        normed_v = (\n",
    "            g * v * math_ops.rsqrt(math_ops.reduce_sum(math_ops.square(v)))\n",
    "        )\n",
    "        return math_ops.reduce_sum(\n",
    "            normed_v * math_ops.tanh(keys + processed_query + b), [2]\n",
    "        )\n",
    "    else:\n",
    "        return tf.reduce_sum(\n",
    "            v * tf.tanh(keys + processed_query), [2]\n",
    "        )\n",
    "\n",
    "\n",
    "class VAEBahdanauAttention(tf.contrib.seq2seq.BahdanauAttention):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_units,\n",
    "        memory,\n",
    "        concat_z,\n",
    "        memory_sequence_length = None,\n",
    "        normalize = False,\n",
    "        probability_fn = None,\n",
    "        score_mask_value = float('-inf'),\n",
    "        name = 'VAEBahdanauAttention',\n",
    "    ):\n",
    "        super(VAEBahdanauAttention, self).__init__(\n",
    "            num_units = num_units,\n",
    "            memory = memory,\n",
    "            memory_sequence_length = memory_sequence_length,\n",
    "            normalize = normalize,\n",
    "            probability_fn = probability_fn,\n",
    "            score_mask_value = score_mask_value,\n",
    "            name = name,\n",
    "        )\n",
    "        self.z = concat_z\n",
    "\n",
    "    def __call__(self, query, state):\n",
    "        with tf.variable_scope(\n",
    "            None, 'vae_bahdanau_attention', [query]\n",
    "        ):\n",
    "            processed_query = (\n",
    "                self.query_layer(query) if self.query_layer else query\n",
    "            )\n",
    "            score = _bahdanau_score(\n",
    "                processed_query, self._keys, self._normalize\n",
    "            )\n",
    "        alignments = self._probability_fn(score, state)\n",
    "        next_state = alignments\n",
    "        next_state = tf.concat([next_state, self.z], -1)\n",
    "        return alignments, next_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE:\n",
    "    def __init__(self, size_layer, num_layers, embedded_size, dict_size, learning_rate,\n",
    "                beam_size = 15, latent_size = 16, anneal_max = 1.0, anneal_bias = 6000):\n",
    "        \n",
    "        def cells(reuse=False):\n",
    "            return tf.nn.rnn_cell.GRUCell(size_layer, reuse=reuse)\n",
    "        \n",
    "        def attention(encoder_out, z, seq_len, reuse=False):\n",
    "            attention_mechanism = VAEBahdanauAttention(num_units = size_layer, \n",
    "                                                       memory = encoder_out,\n",
    "                                                       concat_z = z,\n",
    "                                                       memory_sequence_length = seq_len)\n",
    "            return tf.contrib.seq2seq.AttentionWrapper(\n",
    "            cell = cells(reuse), \n",
    "                attention_mechanism = attention_mechanism,\n",
    "                attention_layer_size = size_layer)\n",
    "        \n",
    "        def kl_w_fn(global_step):\n",
    "            return anneal_max * tf.sigmoid((10 / anneal_bias) * \\\n",
    "                                           (tf.to_float(global_step) - tf.constant(anneal_bias / 2)))\n",
    "        \n",
    "        self.X = tf.placeholder(tf.int32, [None, None])\n",
    "        self.Y_input = tf.placeholder(tf.int32, [None, None])\n",
    "        self.Y_output = tf.placeholder(tf.int32, [None, None])\n",
    "        \n",
    "        self.X_seq_len = tf.count_nonzero(self.X, 1, dtype=tf.int32)\n",
    "        self.Y_seq_len = tf.count_nonzero(self.Y_input, 1, dtype=tf.int32)\n",
    "        batch_size = tf.shape(self.X)[0]\n",
    "        main = tf.strided_slice(self.Y_input, [0, 0], [batch_size, -1], [1, 1])\n",
    "        decoder_input = tf.concat([tf.fill([batch_size, 1], word2idx['<START>']), main], 1)\n",
    "        \n",
    "        embeddings = tf.Variable(tf.random_uniform([dict_size, embedded_size], -1, 1))\n",
    "        x = tf.nn.embedding_lookup(embeddings, self.X)\n",
    "        \n",
    "        encoder_out, encoder_state = tf.nn.dynamic_rnn(\n",
    "            cell = tf.nn.rnn_cell.MultiRNNCell([cells() for _ in range(num_layers)]), \n",
    "            inputs = x,\n",
    "            sequence_length = self.X_seq_len,\n",
    "            dtype = tf.float32)\n",
    "        encoder_state = encoder_state[-1]\n",
    "        \n",
    "        z_mean = tf.layers.dense(encoder_state, latent_size)\n",
    "        z_var = tf.layers.dense(encoder_state, latent_size)\n",
    "        \n",
    "        posterior = tf.contrib.distributions.MultivariateNormalDiag(z_mean, z_var)\n",
    "        prior = tf.contrib.distributions.MultivariateNormalDiag(tf.zeros_like(z_mean),\n",
    "                                                            tf.ones_like(z_var))\n",
    "        z = posterior.sample()\n",
    "        z_reshape = tf.reshape(z, (batch_size, 1, latent_size))\n",
    "        z_timeseries = tf.tile(z_reshape, [1, tf.shape(encoder_out)[1], 1])\n",
    "        print(z_timeseries, encoder_out)\n",
    "        encoder_out = tf.multiply(encoder_out, z_timeseries)\n",
    "        print(encoder_out)\n",
    "        init_state = tf.layers.dense(z, size_layer, tf.nn.elu)\n",
    "        \n",
    "        output_proj = tf.layers.Dense(dict_size)\n",
    "        \n",
    "        with tf.variable_scope('decode'):\n",
    "            decoder_cell = attention(encoder_out, z, self.X_seq_len)\n",
    "\n",
    "            helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "                    inputs = tf.nn.embedding_lookup(embeddings, decoder_input),\n",
    "                    sequence_length = self.Y_seq_len)\n",
    "\n",
    "            decoder = ModifiedBasicDecoder(\n",
    "                    cell = decoder_cell,\n",
    "                    helper = helper,\n",
    "                    initial_state = decoder_cell.zero_state(batch_size, tf.float32).clone(cell_state=init_state),\n",
    "                    output_layer = output_proj,\n",
    "                    concat_z = z)\n",
    "\n",
    "            decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                    decoder = decoder,\n",
    "                    impute_finished = True,\n",
    "                    maximum_iterations = tf.reduce_max(self.Y_seq_len))\n",
    "        \n",
    "        self.training_logits = decoder_output.rnn_output\n",
    "        out_dist = tf.distributions.Categorical(self.training_logits)\n",
    "        global_step = tf.Variable(0, trainable=False)\n",
    "        self.out_dist = out_dist.log_prob(self.Y_output)\n",
    "        nll_loss = -tf.reduce_sum(self.out_dist)\n",
    "        self.nll_loss = nll_loss\n",
    "        kl_w = kl_w_fn(global_step)\n",
    "        self.kl_w = kl_w\n",
    "        kl_loss = tf.reduce_sum(tf.distributions.kl_divergence(posterior, prior))\n",
    "        self.kl_loss = kl_loss\n",
    "        self.cost = nll_loss + kl_w * kl_loss\n",
    "        self.optimizer = tf.train.AdamOptimizer(learning_rate).minimize(self.cost, \n",
    "                                                                        global_step = global_step)\n",
    "        \n",
    "        with tf.variable_scope('decode', reuse=True):\n",
    "            encoder_out_tiled = tf.contrib.seq2seq.tile_batch(encoder_out, beam_size)\n",
    "            X_seq_len_tiled = tf.contrib.seq2seq.tile_batch(self.X_seq_len, beam_size)\n",
    "            z_tiled = tf.contrib.seq2seq.tile_batch(z, beam_size)\n",
    "            decoder_cell = attention(encoder_out_tiled, z_tiled, X_seq_len_tiled)\n",
    "            tiled_z = tf.tile(tf.expand_dims(z, 1), [1, beam_size, 1])\n",
    "            encoder_state_tiled = tf.contrib.seq2seq.tile_batch(init_state, beam_size)\n",
    "            init_state = decoder_cell.zero_state(batch_size * beam_size, tf.float32).\\\n",
    "            clone(cell_state=encoder_state_tiled)\n",
    "            decoder = ModifiedBeamSearchDecoder(\n",
    "                    cell = decoder_cell,\n",
    "                    embedding = embeddings,\n",
    "                    start_tokens = tf.tile(tf.constant([word2idx['<START>']], tf.int32),\n",
    "                                           [batch_size]),\n",
    "                    end_token = word2idx['<END>'],\n",
    "                    initial_state = init_state,\n",
    "                    beam_width = beam_size,\n",
    "                    output_layer = output_proj,\n",
    "                    concat_z = tiled_z)\n",
    "            decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "                    maximum_iterations = tf.reduce_max(self.X_seq_len),\n",
    "                    decoder = decoder)\n",
    "            self.predict_ids = decoder_output.predicted_ids[:, :, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "size_layer = 128\n",
    "num_layers = 2\n",
    "embedded_size = 128\n",
    "learning_rate = 1e-3\n",
    "batch_size = 8\n",
    "epoch = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From <ipython-input-10-c06859524df3>:45: MultivariateNormalDiag.__init__ (from tensorflow.contrib.distributions.python.ops.mvn_diag) is deprecated and will be removed after 2018-10-01.\n",
      "Instructions for updating:\n",
      "The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.contrib.distributions`.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/contrib/distributions/python/ops/mvn_diag.py:224: MultivariateNormalLinearOperator.__init__ (from tensorflow.contrib.distributions.python.ops.mvn_linear_operator) is deprecated and will be removed after 2018-10-01.\n",
      "Instructions for updating:\n",
      "The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.contrib.distributions`.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/contrib/distributions/python/ops/mvn_linear_operator.py:201: AffineLinearOperator.__init__ (from tensorflow.contrib.distributions.python.ops.bijectors.affine_linear_operator) is deprecated and will be removed after 2018-10-01.\n",
      "Instructions for updating:\n",
      "The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.contrib.distributions`.\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/contrib/distributions/python/ops/bijectors/affine_linear_operator.py:158: _DistributionShape.__init__ (from tensorflow.contrib.distributions.python.ops.shape) is deprecated and will be removed after 2018-10-01.\n",
      "Instructions for updating:\n",
      "The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.contrib.distributions`.\n",
      "Tensor(\"Tile:0\", shape=(?, ?, 128), dtype=float32) Tensor(\"rnn/transpose_1:0\", shape=(?, ?, 128), dtype=float32)\n",
      "Tensor(\"Mul:0\", shape=(?, ?, 128), dtype=float32)\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/distributions/kullback_leibler.py:98: _kl_brute_force (from tensorflow.contrib.distributions.python.ops.mvn_linear_operator) is deprecated and will be removed after 2018-10-01.\n",
      "Instructions for updating:\n",
      "The TensorFlow Distributions library has moved to TensorFlow Probability (https://github.com/tensorflow/probability). You should update all references to use `tfp.distributions` instead of `tf.contrib.distributions`.\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = VAE(size_layer, num_layers, embedded_size, len(word2idx), learning_rate,\n",
    "           latent_size = size_layer)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word_dropout(x):\n",
    "    is_dropped = np.random.binomial(1, 0.5, x.shape)\n",
    "    fn = np.vectorize(lambda x, k: word2idx['<UNK>'] if (\n",
    "                      k and (x not in range(4))) else x)\n",
    "    return fn(x, is_dropped)\n",
    "\n",
    "def inf_inp(test_strs):\n",
    "    x = [[word2idx.get(w, 2) for w in s.split()] for s in test_strs]\n",
    "    x = tf.keras.preprocessing.sequence.pad_sequences(\n",
    "        x, maxlen, truncating='post', padding='post')\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_strings = ['i love this film and i think it is one of the best films',\n",
    "             'this movie is a waste of time and there is no point to watch it']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 14, 120,  15,  23,   6,  14, 105,  13,  10,  32,   8,   5, 119,\n",
       "        109,   0,   0,   0,   0,   0,   0],\n",
       "       [ 15,  21,  10,   7, 438,   8,  59,   6,  51,  10,  58, 214,   9,\n",
       "        107,  13,   0,   0,   0,   0,   0]], dtype=int32)"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "inf_inp(test_strings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_input = word_dropout(Y_input[:2])\n",
    "y_output = Y_output[:2]\n",
    "x = X[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[459.06894,\n",
       " array([[-11.385559 , -11.385246 , -11.393755 , -11.390775 , -11.39856  ,\n",
       "         -11.399674 , -11.389531 , -11.390428 , -11.375457 , -11.397874 ,\n",
       "         -11.380082 , -11.387401 , -11.3699045, -11.399622 , -11.370014 ,\n",
       "         -11.398319 , -11.3987465, -11.403546 , -11.405607 , -11.382941 ],\n",
       "        [-11.386787 , -11.3887615, -11.391078 , -11.393381 , -11.383347 ,\n",
       "         -11.388348 , -11.406068 , -11.381132 , -11.385797 , -11.386624 ,\n",
       "         -11.379927 , -11.386466 , -11.398993 , -11.398088 , -11.40023  ,\n",
       "         -11.405902 , -11.401047 , -11.406738 , -11.385779 , -11.386157 ]],\n",
       "       dtype=float32),\n",
       " 455.64368,\n",
       " 511.7788,\n",
       " array([[[ 0.00605713,  0.00100482, -0.00248459, ..., -0.00274286,\n",
       "          -0.00042495,  0.00513934],\n",
       "         [ 0.00753104, -0.0006326 ,  0.00039052, ...,  0.00077708,\n",
       "           0.00275912,  0.00477254],\n",
       "         [ 0.00336134,  0.00094279,  0.00207165, ...,  0.0014382 ,\n",
       "           0.00137247,  0.00993976],\n",
       "         ...,\n",
       "         [-0.00545961,  0.00610403, -0.00068409, ...,  0.01330453,\n",
       "           0.00010723,  0.03450532],\n",
       "         [-0.00629055,  0.00834379, -0.00227644, ...,  0.01385184,\n",
       "          -0.00027043,  0.03543622],\n",
       "         [-0.00257956,  0.01771955, -0.00433735, ...,  0.01278491,\n",
       "           0.00662502,  0.02987854]],\n",
       " \n",
       "        [[ 0.00143907,  0.00075642,  0.01064086, ...,  0.01056396,\n",
       "           0.01109017,  0.00227523],\n",
       "         [ 0.00381149,  0.00035135,  0.0094788 , ...,  0.01450245,\n",
       "           0.01234856, -0.00036613],\n",
       "         [ 0.00070621,  0.00066832,  0.008125  , ...,  0.01468847,\n",
       "           0.00935656,  0.00438248],\n",
       "         ...,\n",
       "         [-0.00210279,  0.00427182,  0.00383839, ...,  0.00690867,\n",
       "           0.00212132,  0.02549111],\n",
       "         [-0.0034757 ,  0.00567319,  0.00166805, ...,  0.00799169,\n",
       "           0.00010306,  0.02731323],\n",
       "         [-0.00696126,  0.00991146, -0.00273268, ...,  0.0063918 ,\n",
       "           0.00073863,  0.02924314]]], dtype=float32)]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sess.run([model.cost, model.out_dist, model.nll_loss, model.kl_loss, model.training_logits],\n",
    "         feed_dict = {model.X: x, model.Y_input: y_input,\n",
    "                      model.Y_output: y_output})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"bewilderingly 24' 24' 24' 24' 24' drat melford's melford's melford's melford's melford's bigg's 50usd 50usd\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "r_aug = sess.run(model.predict_ids, feed_dict = {model.X: inf_inp(test_strings)})[0]\n",
    "' '.join([idx2word[r] for r in r_aug])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "epoch = 10\n",
    "batch_size = 32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:08<00:00,  3.98it/s, cost=1.44e+3]\n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, average loss 2655.045520\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i i this film would i it it is one of the best films films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:07<00:00,  3.97it/s, cost=1.22e+3]\n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 2, average loss 1499.365106\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: to love a and and would think it this one than the movie part word\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.98it/s, cost=830]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 3, average loss 1109.726648\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.98it/s, cost=699]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 4, average loss 995.606202\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END> <END>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.98it/s, cost=648]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 5, average loss 842.265805\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this film it's i i it is one of the best films films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.97it/s, cost=2e+3]   \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 6, average loss 1219.888623\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this film and i think it is one of the the of films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:07<00:00,  3.97it/s, cost=717]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 7, average loss 1091.723570\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this film and i think it is one of the best films films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:07<00:00,  3.95it/s, cost=670]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 8, average loss 930.946201\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this film i i think it is is of of the the best\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.97it/s, cost=698]    \n",
      "minibatch loop:   0%|          | 0/3125 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 9, average loss 731.723512\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this this and i think it is one one of the best films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "minibatch loop: 100%|██████████| 3125/3125 [13:06<00:00,  3.97it/s, cost=898]    "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 10, average loss 1197.459520\n",
      "real string: i love this film and i think it is one of the best films\n",
      "augmented string: i love this film and i think it is one of of the best films\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "for e in range(epoch):\n",
    "    pbar = tqdm(\n",
    "        range(0, len(X), batch_size), desc = 'minibatch loop')\n",
    "    cost = 0\n",
    "    for i in pbar:\n",
    "        index = min(i + batch_size, len(X))\n",
    "        y_input = word_dropout(Y_input[i: index])\n",
    "        y_output = Y_output[i: index]\n",
    "        x = X[i: index]\n",
    "        c, _ = sess.run([model.cost, model.optimizer],\n",
    "         feed_dict = {model.X: x, model.Y_input: y_input,\n",
    "                      model.Y_output: y_output})\n",
    "        cost += c\n",
    "        pbar.set_postfix(cost = c)\n",
    "    cost /= (len(X) / batch_size)\n",
    "    r_aug = sess.run(model.predict_ids, feed_dict = {model.X: inf_inp(test_strings)})[0]\n",
    "    print('epoch %d, average loss %f'%(e + 1, cost))\n",
    "    print('real string: %s'%(test_strings[0]))\n",
    "    print('augmented string: %s'%(' '.join([idx2word[r] for r in r_aug])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "augmented string: i love this film and i think it is one of of the films films\n"
     ]
    }
   ],
   "source": [
    "r_aug = sess.run(model.predict_ids, feed_dict = {model.X: inf_inp(test_strings)})[0]\n",
    "print('augmented string: %s'%(' '.join([idx2word[r] for r in r_aug])))"
   ]
  },
  {
   "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
}