chatbot/53.tacotron-beam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from sklearn.utils import shuffle\n",
    "import re\n",
    "import time\n",
    "import collections\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(words, n_words, atleast=1):\n",
    "    count = [['PAD', 0], ['GO', 1], ['EOS', 2], ['UNK', 3]]\n",
    "    counter = collections.Counter(words).most_common(n_words)\n",
    "    counter = [i for i in counter if i[1] >= atleast]\n",
    "    count.extend(counter)\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": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "lines = open('movie_lines.txt', encoding='utf-8', errors='ignore').read().split('\\n')\n",
    "conv_lines = open('movie_conversations.txt', encoding='utf-8', errors='ignore').read().split('\\n')\n",
    "\n",
    "id2line = {}\n",
    "for line in lines:\n",
    "    _line = line.split(' +++$+++ ')\n",
    "    if len(_line) == 5:\n",
    "        id2line[_line[0]] = _line[4]\n",
    "        \n",
    "convs = [ ]\n",
    "for line in conv_lines[:-1]:\n",
    "    _line = line.split(' +++$+++ ')[-1][1:-1].replace(\"'\",\"\").replace(\" \",\"\")\n",
    "    convs.append(_line.split(','))\n",
    "    \n",
    "questions = []\n",
    "answers = []\n",
    "\n",
    "for conv in convs:\n",
    "    for i in range(len(conv)-1):\n",
    "        questions.append(id2line[conv[i]])\n",
    "        answers.append(id2line[conv[i+1]])\n",
    "        \n",
    "def clean_text(text):\n",
    "    text = text.lower()\n",
    "    text = re.sub(r\"i'm\", \"i am\", text)\n",
    "    text = re.sub(r\"he's\", \"he is\", text)\n",
    "    text = re.sub(r\"she's\", \"she is\", text)\n",
    "    text = re.sub(r\"it's\", \"it is\", text)\n",
    "    text = re.sub(r\"that's\", \"that is\", text)\n",
    "    text = re.sub(r\"what's\", \"that is\", text)\n",
    "    text = re.sub(r\"where's\", \"where is\", text)\n",
    "    text = re.sub(r\"how's\", \"how is\", text)\n",
    "    text = re.sub(r\"\\'ll\", \" will\", text)\n",
    "    text = re.sub(r\"\\'ve\", \" have\", text)\n",
    "    text = re.sub(r\"\\'re\", \" are\", text)\n",
    "    text = re.sub(r\"\\'d\", \" would\", text)\n",
    "    text = re.sub(r\"\\'re\", \" are\", text)\n",
    "    text = re.sub(r\"won't\", \"will not\", text)\n",
    "    text = re.sub(r\"can't\", \"cannot\", text)\n",
    "    text = re.sub(r\"n't\", \" not\", text)\n",
    "    text = re.sub(r\"n'\", \"ng\", text)\n",
    "    text = re.sub(r\"'bout\", \"about\", text)\n",
    "    text = re.sub(r\"'til\", \"until\", text)\n",
    "    text = re.sub(r\"[-()\\\"#/@;:<>{}`+=~|.!?,]\", \"\", text)\n",
    "    return ' '.join([i.strip() for i in filter(None, text.split())])\n",
    "\n",
    "clean_questions = []\n",
    "for question in questions:\n",
    "    clean_questions.append(clean_text(question))\n",
    "    \n",
    "clean_answers = []    \n",
    "for answer in answers:\n",
    "    clean_answers.append(clean_text(answer))\n",
    "    \n",
    "min_line_length = 2\n",
    "max_line_length = 5\n",
    "short_questions_temp = []\n",
    "short_answers_temp = []\n",
    "\n",
    "i = 0\n",
    "for question in clean_questions:\n",
    "    if len(question.split()) >= min_line_length and len(question.split()) <= max_line_length:\n",
    "        short_questions_temp.append(question)\n",
    "        short_answers_temp.append(clean_answers[i])\n",
    "    i += 1\n",
    "\n",
    "short_questions = []\n",
    "short_answers = []\n",
    "\n",
    "i = 0\n",
    "for answer in short_answers_temp:\n",
    "    if len(answer.split()) >= min_line_length and len(answer.split()) <= max_line_length:\n",
    "        short_answers.append(answer)\n",
    "        short_questions.append(short_questions_temp[i])\n",
    "    i += 1\n",
    "    \n",
    "question_test = short_questions[500:550]\n",
    "answer_test = short_answers[500:550]\n",
    "short_questions = short_questions[:500]\n",
    "short_answers = short_answers[:500]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab from size: 657\n",
      "Most common words [('you', 132), ('is', 78), ('i', 68), ('what', 51), ('it', 50), ('that', 49)]\n",
      "Sample data [7, 28, 129, 35, 61, 42, 12, 22, 82, 225] ['what', 'good', 'stuff', 'she', 'okay', 'they', 'do', 'to', 'hey', 'sweet']\n",
      "filtered vocab size: 661\n",
      "% of vocab used: 100.61%\n"
     ]
    }
   ],
   "source": [
    "concat_from = ' '.join(short_questions+question_test).split()\n",
    "vocabulary_size_from = len(list(set(concat_from)))\n",
    "data_from, count_from, dictionary_from, rev_dictionary_from = build_dataset(concat_from, vocabulary_size_from)\n",
    "print('vocab from size: %d'%(vocabulary_size_from))\n",
    "print('Most common words', count_from[4:10])\n",
    "print('Sample data', data_from[:10], [rev_dictionary_from[i] for i in data_from[:10]])\n",
    "print('filtered vocab size:',len(dictionary_from))\n",
    "print(\"% of vocab used: {}%\".format(round(len(dictionary_from)/vocabulary_size_from,4)*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab from size: 660\n",
      "Most common words [('i', 97), ('you', 91), ('is', 62), ('it', 58), ('not', 47), ('what', 39)]\n",
      "Sample data [12, 216, 5, 4, 94, 25, 59, 10, 8, 79] ['the', 'real', 'you', 'i', 'hope', 'so', 'they', 'do', 'not', 'hi']\n",
      "filtered vocab size: 664\n",
      "% of vocab used: 100.61%\n"
     ]
    }
   ],
   "source": [
    "concat_to = ' '.join(short_answers+answer_test).split()\n",
    "vocabulary_size_to = len(list(set(concat_to)))\n",
    "data_to, count_to, dictionary_to, rev_dictionary_to = build_dataset(concat_to, vocabulary_size_to)\n",
    "print('vocab from size: %d'%(vocabulary_size_to))\n",
    "print('Most common words', count_to[4:10])\n",
    "print('Sample data', data_to[:10], [rev_dictionary_to[i] for i in data_to[:10]])\n",
    "print('filtered vocab size:',len(dictionary_to))\n",
    "print(\"% of vocab used: {}%\".format(round(len(dictionary_to)/vocabulary_size_to,4)*100))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "GO = dictionary_from['GO']\n",
    "PAD = dictionary_from['PAD']\n",
    "EOS = dictionary_from['EOS']\n",
    "UNK = dictionary_from['UNK']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(len(short_answers)):\n",
    "    short_answers[i] += ' EOS'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def str_idx(corpus, dic):\n",
    "    X = []\n",
    "    for i in corpus:\n",
    "        ints = []\n",
    "        for k in i.split():\n",
    "            ints.append(dic.get(k,UNK))\n",
    "        X.append(ints)\n",
    "    return X\n",
    "\n",
    "def pad_sentence_batch(sentence_batch, pad_int):\n",
    "    padded_seqs = []\n",
    "    seq_lens = []\n",
    "    max_sentence_len = max([len(sentence) for sentence in sentence_batch])\n",
    "    for sentence in sentence_batch:\n",
    "        padded_seqs.append(sentence + [pad_int] * (max_sentence_len - len(sentence)))\n",
    "        seq_lens.append(len(sentence))\n",
    "    return padded_seqs, seq_lens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = str_idx(short_questions, dictionary_from)\n",
    "Y = str_idx(short_answers, dictionary_to)\n",
    "X_test = str_idx(question_test, dictionary_from)\n",
    "Y_test = str_idx(answer_test, dictionary_from)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_size = 256\n",
    "encoder_num_banks = 16\n",
    "decoder_num_banks = 8\n",
    "num_highway_blocks = 4\n",
    "reduction_factor = 5\n",
    "learning_rate = 1e-4\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from setting import embed_size\n",
    "import tensorflow as tf\n",
    "\n",
    "\n",
    "def embed(inputs, vocab_size, dimension, scope = 'embedding', reuse = None):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        lookup_table = tf.get_variable(\n",
    "            'lookup_table',\n",
    "            dtype = tf.float32,\n",
    "            shape = [vocab_size, dimension],\n",
    "            initializer = tf.truncated_normal_initializer(\n",
    "                mean = 0.0, stddev = 0.01\n",
    "            ),\n",
    "        )\n",
    "        lookup_table = tf.concat(\n",
    "            (tf.zeros(shape = [1, dimension]), lookup_table[1:, :]), 0\n",
    "        )\n",
    "    return tf.nn.embedding_lookup(lookup_table, inputs)\n",
    "\n",
    "\n",
    "def normalize_bn(\n",
    "    inputs,\n",
    "    decay = 0.99,\n",
    "    is_training = True,\n",
    "    activation_fn = None,\n",
    "    scope = 'normalize_bn',\n",
    "):\n",
    "    inputs_shape = inputs.get_shape()\n",
    "    inputs_rank = inputs_shape.ndims\n",
    "    if inputs_rank in [2, 3, 4]:\n",
    "        if inputs_rank == 2:\n",
    "            inputs = tf.expand_dims(inputs, axis = 1)\n",
    "            inputs = tf.expand_dims(inputs, axis = 2)\n",
    "        elif inputs_rank == 3:\n",
    "            inputs = tf.expand_dims(inputs, axis = 1)\n",
    "        outputs = tf.contrib.layers.batch_norm(\n",
    "            inputs = inputs,\n",
    "            decay = decay,\n",
    "            center = True,\n",
    "            scale = True,\n",
    "            activation_fn = activation_fn,\n",
    "            updates_collections = None,\n",
    "            is_training = is_training,\n",
    "            scope = scope,\n",
    "            zero_debias_moving_mean = True,\n",
    "            fused = True,\n",
    "        )\n",
    "        if inputs_rank == 2:\n",
    "            outputs = tf.squeeze(outputs, axis = [1, 2])\n",
    "        elif inputs_rank == 3:\n",
    "            outputs = tf.squeeze(outputs, axis = 1)\n",
    "    else:\n",
    "        outputs = tf.contrib.layers.batch_norm(\n",
    "            inputs = inputs,\n",
    "            decay = decay,\n",
    "            center = True,\n",
    "            scale = True,\n",
    "            activation_fn = activation_fn,\n",
    "            updates_collections = None,\n",
    "            is_training = is_training,\n",
    "            scope = scope,\n",
    "            fused = False,\n",
    "        )\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def normalize_layer_norm(\n",
    "    inputs, activation_fn = None, scope = 'normalize_layer_norm'\n",
    "):\n",
    "    return tf.contrib.layers.layer_norm(\n",
    "        inputs = inputs,\n",
    "        center = True,\n",
    "        scale = True,\n",
    "        activation_fn = activation_fn,\n",
    "        scope = scope,\n",
    "    )\n",
    "\n",
    "\n",
    "def normalize_in(inputs, activation_fn = None, scope = 'normalize_in', reuse = False):\n",
    "    with tf.variable_scope(scope + '_normalize_in', reuse = reuse):\n",
    "        batch, steps, channels = inputs.get_shape().as_list()\n",
    "        var_shape = [channels]\n",
    "        mu, sigma_sq = tf.nn.moments(inputs, [1], keep_dims = True)\n",
    "        shift = tf.get_variable(\n",
    "            'shift',\n",
    "            dtype = tf.float32,\n",
    "            shape = var_shape,\n",
    "            initializer = tf.initializers.zeros(),\n",
    "        )\n",
    "        scale = tf.get_variable(\n",
    "            'scale',\n",
    "            dtype = tf.float32,\n",
    "            shape = var_shape,\n",
    "            initializer = tf.initializers.ones(),\n",
    "        )\n",
    "        epsilon = 1e-8\n",
    "        normalized = (inputs - mu) / (sigma_sq + epsilon) ** (0.5)\n",
    "        outputs = scale * normalized + shift\n",
    "        if activation_fn:\n",
    "            outputs = activation_fn(outputs)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def conv1d(\n",
    "    inputs,\n",
    "    filters = None,\n",
    "    size = 1,\n",
    "    rate = 1,\n",
    "    padding = 'SAME',\n",
    "    use_bias = False,\n",
    "    activation_fn = None,\n",
    "    scope = 'conv1d',\n",
    "    reuse = None,\n",
    "):\n",
    "    with tf.variable_scope(scope):\n",
    "        if padding.lower() == 'causal':\n",
    "            pad_len = (size - 1) * rate\n",
    "            inputs = tf.pad(inputs, [[0, 0], [pad_len, 0], [0, 0]])\n",
    "            padding = 'valid'\n",
    "        if filters is None:\n",
    "            filters = inputs.get_shape().as_list()[-1]\n",
    "        params = {\n",
    "            'inputs': inputs,\n",
    "            'filters': filters,\n",
    "            'kernel_size': size,\n",
    "            'dilation_rate': rate,\n",
    "            'padding': padding,\n",
    "            'activation': activation_fn,\n",
    "            'use_bias': use_bias,\n",
    "            'reuse': reuse,\n",
    "        }\n",
    "        outputs = tf.layers.conv1d(**params)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def conv1d_banks(\n",
    "    inputs, K = 16, is_training = True, scope = 'conv1d_banks', reuse = None\n",
    "):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        outputs = conv1d(inputs, embed_size // 2, 1)\n",
    "        outputs = normalize_in(outputs, tf.nn.relu, reuse = reuse)\n",
    "        for k in range(2, K + 1):\n",
    "            with tf.variable_scope('num_%d' % (k)):\n",
    "                output = conv1d(inputs, embed_size // 2, k)\n",
    "                output = normalize_in(output, tf.nn.relu, reuse = reuse)\n",
    "                outputs = tf.concat((outputs, output), -1)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def gru(inputs, units = None, bidirection = False, scope = 'gru', reuse = None):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        if units is None:\n",
    "            units = inputs.get_shape().as_list()[-1]\n",
    "        cell = tf.contrib.rnn.GRUCell(units)\n",
    "        if bidirection:\n",
    "            cell_bw = tf.contrib.rnn.GRUCell(units)\n",
    "            outputs, _ = tf.nn.bidirectional_dynamic_rnn(\n",
    "                cell, cell_bw, inputs, dtype = tf.float32\n",
    "            )\n",
    "            return tf.concat(outputs, 2)\n",
    "        else:\n",
    "            outputs, _ = tf.nn.dynamic_rnn(cell, inputs, dtype = tf.float32)\n",
    "            return outputs\n",
    "\n",
    "\n",
    "def attention_decoder(\n",
    "    inputs, memory, units = None, scope = 'attention_decoder', reuse = None\n",
    "):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        if units is None:\n",
    "            units = inputs.get_shape().as_list()[-1]\n",
    "        attention_mechanism = tf.contrib.seq2seq.BahdanauAttention(\n",
    "            units, memory\n",
    "        )\n",
    "        decoder_cell = tf.contrib.rnn.GRUCell(units)\n",
    "        cell_with_attention = tf.contrib.seq2seq.AttentionWrapper(\n",
    "            decoder_cell, attention_mechanism, units\n",
    "        )\n",
    "        outputs, _ = tf.nn.dynamic_rnn(\n",
    "            cell_with_attention, inputs, dtype = tf.float32\n",
    "        )\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def prenet(inputs, is_training = True, scope = 'prenet', reuse = None):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        outputs = tf.layers.dense(\n",
    "            inputs, units = embed_size, activation = tf.nn.relu, name = 'dense1'\n",
    "        )\n",
    "        outputs = tf.nn.dropout(\n",
    "            outputs,\n",
    "            keep_prob = 0.5 if is_training == True else 1.0,\n",
    "            name = 'dropout1',\n",
    "        )\n",
    "        outputs = tf.layers.dense(\n",
    "            outputs,\n",
    "            units = embed_size // 2,\n",
    "            activation = tf.nn.relu,\n",
    "            name = 'dense2',\n",
    "        )\n",
    "        outputs = tf.nn.dropout(\n",
    "            outputs,\n",
    "            keep_prob = 0.5 if is_training == True else 1.0,\n",
    "            name = 'dropout2',\n",
    "        )\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def highwaynet(inputs, units = None, scope = 'highwaynet', reuse = None):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        if units is None:\n",
    "            units = inputs.get_shape().as_list()[-1]\n",
    "        H = tf.layers.dense(\n",
    "            inputs, units = units, activation = tf.nn.relu, name = 'dense1'\n",
    "        )\n",
    "        T = tf.layers.dense(\n",
    "            inputs, units = units, activation = tf.nn.sigmoid, name = 'dense2'\n",
    "        )\n",
    "        C = 1.0 - T\n",
    "        return H * T + inputs * C\n",
    "\n",
    "\n",
    "def shift_by_one(inputs):\n",
    "    return tf.concat((tf.zeros_like(inputs[:, :1]), inputs[:, :-1]), 1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(inputs, is_training = True, scope = 'encoder', reuse = None):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        prenet_out = prenet(inputs, scope = 'prenet', is_training = is_training)\n",
    "        enc = conv1d_banks(\n",
    "            prenet_out, K = encoder_num_banks, is_training = is_training\n",
    "        )\n",
    "        enc = tf.layers.max_pooling1d(enc, 2, 1, padding = 'same')\n",
    "        enc = conv1d(enc, embed_size // 2, 3, scope = 'conv1d_1')\n",
    "        enc = normalize_in(enc, activation_fn = tf.nn.relu, scope = 'conv1d_1', reuse = reuse)\n",
    "        enc = conv1d(enc, embed_size // 2, 3, scope = 'conv1d_2')\n",
    "        enc = normalize_in(enc, activation_fn = tf.nn.relu, scope = 'conv1d_2', reuse = reuse)\n",
    "        enc += prenet_out\n",
    "        for i in range(num_highway_blocks):\n",
    "            enc = highwaynet(\n",
    "                enc, units = embed_size // 2, scope = 'highwaynet_%d' % (i)\n",
    "            )\n",
    "        memory = gru(enc, embed_size // 2, True)\n",
    "    return memory\n",
    "\n",
    "def decode(\n",
    "    inputs, memory, output_size, is_training = True, scope = 'decoder_layers', reuse = None\n",
    "):\n",
    "    with tf.variable_scope(scope, reuse = reuse):\n",
    "        dec = prenet(inputs, is_training = is_training)\n",
    "        dec = attention_decoder(dec, memory, embed_size * 2)\n",
    "        dec += gru(dec, embed_size, True, scope = 'gru1')\n",
    "        dec += gru(dec, embed_size, True, scope = 'gru2')\n",
    "        return tf.layers.dense(dec, output_size)\n",
    "\n",
    "class Model:\n",
    "    def __init__(self):\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",
    "        self.is_training = tf.placeholder(tf.bool, None)\n",
    "        batch_size = tf.shape(self.X)[0]\n",
    "        \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",
    "        encoder_embeddings = tf.Variable(tf.random_uniform([len(dictionary_from), embed_size], -1, 1))\n",
    "        decoder_embeddings = tf.Variable(tf.random_uniform([len(dictionary_to), embed_size], -1, 1))\n",
    "        \n",
    "        \n",
    "        def forward(x, y, reuse = False):\n",
    "            encoder_embedded = tf.nn.embedding_lookup(encoder_embeddings, x)\n",
    "            decoder_inputs = tf.nn.embedding_lookup(decoder_embeddings, y)\n",
    "            with tf.variable_scope('net', reuse = reuse):\n",
    "                memory = encode(encoder_embedded, is_training = self.is_training)\n",
    "                outputs = decode(\n",
    "                    decoder_inputs, memory, len(dictionary_to), is_training = self.is_training\n",
    "                )\n",
    "            return outputs\n",
    "        \n",
    "        self.outputs = forward(self.X, decoder_input)\n",
    "        \n",
    "        self.preds = tf.argmax(self.outputs, axis = -1)\n",
    "        masks = tf.sequence_mask(\n",
    "                self.Y_seq_len,\n",
    "                tf.reduce_max(self.Y_seq_len),\n",
    "                dtype = tf.float32,\n",
    "        )\n",
    "        self.cost = tf.contrib.seq2seq.sequence_loss(\n",
    "                logits = self.outputs, targets = self.Y, weights = masks\n",
    "        )\n",
    "        self.optimizer = tf.train.AdamOptimizer(learning_rate).minimize(\n",
    "                self.cost\n",
    "        )\n",
    "        y_t = tf.cast(self.preds, 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))\n",
    "        \n",
    "        self.logits = forward(self.X, self.Y, reuse = True)\n",
    "        self.k = tf.placeholder(dtype = tf.int32)\n",
    "        p = tf.nn.softmax(self.logits)\n",
    "        self.topk_logprobs, self.topk_ids = tf.nn.top_k(tf.log(p), self.k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/client/session.py:1702: UserWarning: An interactive session is already active. This can cause out-of-memory errors in some cases. You must explicitly call `InteractiveSession.close()` to release resources held by the other session(s).\n",
      "  warnings.warn('An interactive session is already active. This can '\n"
     ]
    }
   ],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Model()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 16\n",
    "epoch = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Hypothesis:\n",
    "    def __init__(self, log_prob, seq):\n",
    "        self.log_prob = log_prob\n",
    "        self.seq = seq\n",
    "\n",
    "    @property\n",
    "    def step(self):\n",
    "        return len(self.seq) - 1\n",
    "\n",
    "\n",
    "def beam_search(\n",
    "    batch_x,\n",
    "    beam_size,\n",
    "    num_ans = 5,\n",
    "    normalize_by_len = 1.0,\n",
    "):\n",
    "    assert 0 <= normalize_by_len <= 1\n",
    "    batch_size = len(batch_x)\n",
    "    max_len = len(batch_x[0]) * 2\n",
    "    dec_inputs = np.ones((batch_size, 2), dtype=np.int32)\n",
    "    answers = [[] for i in range(batch_size)]\n",
    "    H = [[] for i in range(batch_size)]\n",
    "    \n",
    "    tkl, tkid = sess.run([model.topk_logprobs, \n",
    "                          model.topk_ids],\n",
    "                         feed_dict = {model.X: batch_x,\n",
    "                                     model.Y: dec_inputs,\n",
    "                                     model.k: beam_size,\n",
    "                                     model.is_training: False})\n",
    "    for i in range(batch_size):\n",
    "        for j, log_prob in enumerate(tkl[i, 0]):\n",
    "            if tkid[i, 0, j] != EOS:\n",
    "                h = Hypothesis(log_prob, [1, tkid[i, 0, j]])\n",
    "                H[i].append(h)\n",
    "        H[i].sort(key=lambda h: h.log_prob)\n",
    "    \n",
    "    done = [False] * batch_size\n",
    "    while not all(done):\n",
    "        tkl_beam = []\n",
    "        tkid_beam = []\n",
    "        dec_inputs_beam = []\n",
    "        steps_beam = []\n",
    "        for i in range(beam_size):\n",
    "            steps = [1] * batch_size\n",
    "            prev_log_probs = np.zeros(batch_size, dtype=np.float32)\n",
    "            dec_inputs = np.ones((batch_size, max_len), dtype=np.int32)\n",
    "            for j, h in enumerate(H):\n",
    "                while h:\n",
    "                    hi = h.pop()\n",
    "                    lp, step, candidate_seq = hi.log_prob, hi.step, hi.seq\n",
    "                    if candidate_seq[-1] != EOS:\n",
    "                        dec_inputs[j, :len(candidate_seq)] = candidate_seq\n",
    "                        steps[j] = step\n",
    "                        prev_log_probs[j] = lp\n",
    "                        break\n",
    "                    else:\n",
    "                        answers[j].append((lp, candidate_seq))\n",
    "                        \n",
    "            max_step = max(steps)\n",
    "            dec_inputs = dec_inputs[:, :max_step + 2]\n",
    "            tkl, tkid = sess.run([model.topk_logprobs, \n",
    "                          model.topk_ids],\n",
    "                         feed_dict = {model.X: batch_x,\n",
    "                                     model.Y: dec_inputs,\n",
    "                                     model.k: beam_size,\n",
    "                                     model.is_training: False})\n",
    "            \n",
    "            tkl_beam.append(tkl + prev_log_probs[:, None, None])\n",
    "            tkid_beam.append(tkid)\n",
    "            dec_inputs_beam.append(dec_inputs.copy())\n",
    "            steps_beam.append(steps)\n",
    "            \n",
    "        for i in range(beam_size):\n",
    "            tkl = tkl_beam[i]\n",
    "            tkid = tkid_beam[i]\n",
    "            dec_inputs = dec_inputs_beam[i]\n",
    "            steps = steps_beam[i]\n",
    "            for j in range(batch_size):\n",
    "                step = steps[j]\n",
    "                for k in range(tkid.shape[2]):\n",
    "                    extended_seq = np.hstack((dec_inputs[j, :step+1], [tkid[j, step, k]]))\n",
    "                    log_prob = tkl[j, step, k]\n",
    "                    if len(extended_seq) <= max_len and log_prob > -10:\n",
    "                        h = Hypothesis(log_prob, extended_seq)\n",
    "                        H[j].append(h)\n",
    "                H[j].sort(key=lambda h: h.log_prob / (h.step**normalize_by_len))\n",
    "            \n",
    "        for i in range(batch_size):\n",
    "            done[i] = (len(answers[i]) >= num_ans) or (not H[i]) or (len(H[i]) > 100)\n",
    "            \n",
    "    return answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, avg loss: 5.611297, avg accuracy: 0.221634\n",
      "epoch: 2, avg loss: 4.900788, avg accuracy: 0.263580\n",
      "epoch: 3, avg loss: 4.714738, avg accuracy: 0.267999\n",
      "epoch: 4, avg loss: 4.544914, avg accuracy: 0.271799\n",
      "epoch: 5, avg loss: 4.281002, avg accuracy: 0.286341\n",
      "epoch: 6, avg loss: 4.030073, avg accuracy: 0.320546\n",
      "epoch: 7, avg loss: 3.780790, avg accuracy: 0.356800\n",
      "epoch: 8, avg loss: 3.517768, avg accuracy: 0.392613\n",
      "epoch: 9, avg loss: 3.249952, avg accuracy: 0.429736\n",
      "epoch: 10, avg loss: 2.981893, avg accuracy: 0.462068\n",
      "epoch: 11, avg loss: 2.720933, avg accuracy: 0.488644\n",
      "epoch: 12, avg loss: 2.500093, avg accuracy: 0.514025\n",
      "epoch: 13, avg loss: 2.384659, avg accuracy: 0.540364\n",
      "epoch: 14, avg loss: 2.213157, avg accuracy: 0.564702\n",
      "epoch: 15, avg loss: 1.936738, avg accuracy: 0.622486\n",
      "epoch: 16, avg loss: 1.717999, avg accuracy: 0.674084\n",
      "epoch: 17, avg loss: 1.516243, avg accuracy: 0.726900\n",
      "epoch: 18, avg loss: 1.343975, avg accuracy: 0.776601\n",
      "epoch: 19, avg loss: 1.173745, avg accuracy: 0.816437\n",
      "epoch: 20, avg loss: 1.020963, avg accuracy: 0.855140\n"
     ]
    }
   ],
   "source": [
    "for i in range(epoch):\n",
    "    total_loss, total_accuracy = 0, 0\n",
    "    for k in range(0, len(short_questions), batch_size):\n",
    "        index = min(k+batch_size, len(short_questions))\n",
    "        batch_x, seq_x = pad_sentence_batch(X[k: index], PAD)\n",
    "        batch_y, seq_y = pad_sentence_batch(Y[k: index], PAD)\n",
    "        predicted, accuracy,loss, _ = sess.run([model.outputs,model.accuracy, model.cost, model.optimizer], \n",
    "                                      feed_dict={model.X:batch_x,\n",
    "                                                model.Y:batch_y,\n",
    "                                                model.is_training: True})\n",
    "        total_loss += loss\n",
    "        total_accuracy += accuracy\n",
    "    total_loss /= (len(short_questions) / batch_size)\n",
    "    total_accuracy /= (len(short_questions) / batch_size)\n",
    "    print('epoch: %d, avg loss: %f, avg accuracy: %f'%(i+1, total_loss, total_accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_x, seq_x = pad_sentence_batch(X[k: index], PAD)\n",
    "batch_y, _ = pad_sentence_batch(Y[k: index], PAD)\n",
    "beamed = beam_search(batch_x, 5)\n",
    "predicted = [max(b, key = lambda t: t[0])[1] for b in beamed]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "row 1\n",
      "QUESTION: i am a werewolf\n",
      "REAL ANSWER: a werewolf\n",
      "PREDICTED ANSWER: no thank \n",
      "\n",
      "row 2\n",
      "QUESTION: i was dreaming again\n",
      "REAL ANSWER: i would think so\n",
      "PREDICTED ANSWER: and would think so so \n",
      "\n",
      "row 3\n",
      "QUESTION: the kitchen\n",
      "REAL ANSWER: very nice\n",
      "PREDICTED ANSWER: yeah one one room \n",
      "\n",
      "row 4\n",
      "QUESTION: the bedroom\n",
      "REAL ANSWER: there is only one bed\n",
      "PREDICTED ANSWER: there is only one bed \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(len(batch_x)):\n",
    "    print('row %d'%(i+1))\n",
    "    print('QUESTION:',' '.join([rev_dictionary_from[n] for n in batch_x[i] if n not in [0,1,2,3]]))\n",
    "    print('REAL ANSWER:',' '.join([rev_dictionary_to[n] for n in batch_y[i] if n not in[0,1,2,3]]))\n",
    "    print('PREDICTED ANSWER:',' '.join([rev_dictionary_to[n] for n in predicted[i] if n not in[0,1,2,3]]),'\\n')"
   ]
  },
  {
   "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
}