chatbot/48.transformer-xl-lstm.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, maxlen):\n",
    "    padded_seqs = []\n",
    "    seq_lens = []\n",
    "    max_sentence_len = maxlen\n",
    "    for sentence in sentence_batch:\n",
    "        padded_seqs.append(sentence + [pad_int] * (max_sentence_len - len(sentence)))\n",
    "        seq_lens.append(maxlen)\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": [
    "maxlen_question = max([len(x) for x in X]) * 2\n",
    "maxlen_answer = max([len(y) for y in Y]) * 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "learning_rate = 1e-3\n",
    "batch_size = 16\n",
    "epoch = 20\n",
    "n_layer = 3\n",
    "d_model = 256\n",
    "d_embed = 256\n",
    "n_head = 10\n",
    "d_head = 50\n",
    "d_inner = 512"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def positional_embedding(pos_seq, inv_freq, bsz = None):\n",
    "    sinusoid_inp = tf.einsum('i,j->ij', pos_seq, inv_freq)\n",
    "    pos_emb = tf.concat([tf.sin(sinusoid_inp), tf.cos(sinusoid_inp)], -1)\n",
    "    if bsz is not None:\n",
    "        return tf.tile(pos_emb[:, None, :], [1, bsz, 1])\n",
    "    else:\n",
    "        return pos_emb[:, None, :]\n",
    "\n",
    "\n",
    "def positionwise_FF(inp, d_model, d_inner, kernel_initializer, scope = 'ff'):\n",
    "    output = inp\n",
    "    with tf.variable_scope(scope):\n",
    "        output = tf.layers.dense(\n",
    "            inp,\n",
    "            d_inner,\n",
    "            activation = tf.nn.relu,\n",
    "            kernel_initializer = kernel_initializer,\n",
    "            name = 'layer_1',\n",
    "        )\n",
    "        output = tf.layers.dense(\n",
    "            output,\n",
    "            d_model,\n",
    "            kernel_initializer = kernel_initializer,\n",
    "            name = 'layer_2',\n",
    "        )\n",
    "        output = tf.contrib.layers.layer_norm(\n",
    "            output + inp, begin_norm_axis = -1\n",
    "        )\n",
    "    return output\n",
    "\n",
    "\n",
    "def rel_shift(x):\n",
    "    x_size = tf.shape(x)\n",
    "\n",
    "    x = tf.pad(x, [[0, 0], [1, 0], [0, 0], [0, 0]])\n",
    "    x = tf.reshape(x, [x_size[1] + 1, x_size[0], x_size[2], x_size[3]])\n",
    "    x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])\n",
    "    x = tf.reshape(x, x_size)\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "def rel_multihead_attn(\n",
    "    w,\n",
    "    r,\n",
    "    r_w_bias,\n",
    "    r_r_bias,\n",
    "    attn_mask,\n",
    "    mems,\n",
    "    d_model,\n",
    "    n_head,\n",
    "    d_head,\n",
    "    kernel_initializer,\n",
    "    scope = 'rel_attn',\n",
    "):\n",
    "    scale = 1 / (d_head ** 0.5)\n",
    "    with tf.variable_scope(scope):\n",
    "        qlen = tf.shape(w)[0]\n",
    "        rlen = tf.shape(r)[0]\n",
    "        bsz = tf.shape(w)[1]\n",
    "\n",
    "        cat = (\n",
    "            tf.concat([mems, w], 0)\n",
    "            if mems is not None and mems.shape.ndims > 1\n",
    "            else w\n",
    "        )\n",
    "        w_heads = tf.layers.dense(\n",
    "            cat,\n",
    "            3 * n_head * d_head,\n",
    "            use_bias = False,\n",
    "            kernel_initializer = kernel_initializer,\n",
    "            name = 'qkv',\n",
    "        )\n",
    "        r_head_k = tf.layers.dense(\n",
    "            r,\n",
    "            n_head * d_head,\n",
    "            use_bias = False,\n",
    "            kernel_initializer = kernel_initializer,\n",
    "            name = 'r',\n",
    "        )\n",
    "\n",
    "        w_head_q, w_head_k, w_head_v = tf.split(w_heads, 3, -1)\n",
    "        w_head_q = w_head_q[-qlen:]\n",
    "\n",
    "        klen = tf.shape(w_head_k)[0]\n",
    "\n",
    "        w_head_q = tf.reshape(w_head_q, [qlen, bsz, n_head, d_head])\n",
    "        w_head_k = tf.reshape(w_head_k, [klen, bsz, n_head, d_head])\n",
    "        w_head_v = tf.reshape(w_head_v, [klen, bsz, n_head, d_head])\n",
    "\n",
    "        r_head_k = tf.reshape(r_head_k, [rlen, n_head, d_head])\n",
    "\n",
    "        rw_head_q = w_head_q + r_w_bias\n",
    "        rr_head_q = w_head_q + r_r_bias\n",
    "\n",
    "        AC = tf.einsum('ibnd,jbnd->ijbn', rw_head_q, w_head_k)\n",
    "        BD = tf.einsum('ibnd,jnd->ijbn', rr_head_q, r_head_k)\n",
    "        BD = rel_shift(BD)\n",
    "        \n",
    "        paddings = tf.fill(tf.shape(BD), float('-inf'))\n",
    "\n",
    "        attn_score = (AC + BD) * scale\n",
    "        attn_mask_t = attn_mask[:, :, None, None]\n",
    "        attn_score = attn_score * (1 - attn_mask_t) - 1e30 * attn_mask_t\n",
    "\n",
    "        attn_prob = tf.nn.softmax(attn_score, 1)\n",
    "        attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, w_head_v)\n",
    "        size_t = tf.shape(attn_vec)\n",
    "        attn_vec = tf.reshape(attn_vec, [size_t[0], size_t[1], n_head * d_head])\n",
    "\n",
    "        attn_out = tf.layers.dense(\n",
    "            attn_vec,\n",
    "            d_model,\n",
    "            use_bias = False,\n",
    "            kernel_initializer = kernel_initializer,\n",
    "            name = 'o',\n",
    "        )\n",
    "\n",
    "        output = tf.contrib.layers.layer_norm(\n",
    "            attn_out + w, begin_norm_axis = -1\n",
    "        )\n",
    "    return output\n",
    "\n",
    "\n",
    "def embedding_lookup(lookup_table, x):\n",
    "    return tf.nn.embedding_lookup(lookup_table, x)\n",
    "\n",
    "\n",
    "def mask_adaptive_embedding_lookup(\n",
    "    x,\n",
    "    n_token,\n",
    "    d_embed,\n",
    "    d_proj,\n",
    "    cutoffs,\n",
    "    initializer,\n",
    "    proj_initializer,\n",
    "    div_val = 1,\n",
    "    proj_same_dim = True,\n",
    "    scope = 'adaptive_embed',\n",
    "    **kwargs\n",
    "):\n",
    "    emb_scale = d_proj ** 0.5\n",
    "    with tf.variable_scope(scope):\n",
    "        if div_val == 1:\n",
    "            lookup_table = tf.get_variable(\n",
    "                'lookup_table', [n_token, d_embed], initializer = initializer\n",
    "            )\n",
    "            y = embedding_lookup(lookup_table, x)\n",
    "            if d_proj != d_embed:\n",
    "                proj_W = tf.get_variable(\n",
    "                    'proj_W', [d_embed, d_proj], initializer = proj_initializer\n",
    "                )\n",
    "                y = tf.einsum('ibe,ed->ibd', y, proj_W)\n",
    "            else:\n",
    "                proj_W = None\n",
    "            ret_params = [lookup_table, proj_W]\n",
    "        else:\n",
    "            tables, projs = [], []\n",
    "            cutoff_ends = [0] + cutoffs + [n_token]\n",
    "            x_size = tf.shape(x)\n",
    "            y = tf.zeros([x_size[0], x_size[1], d_proj])\n",
    "            for i in range(len(cutoff_ends) - 1):\n",
    "                with tf.variable_scope('cutoff_{}'.format(i)):\n",
    "                    l_idx, r_idx = cutoff_ends[i], cutoff_ends[i + 1]\n",
    "                    mask = (x >= l_idx) & (x < r_idx)\n",
    "                    cur_x = tf.boolean_mask(x, mask) - l_idx\n",
    "                    cur_d_embed = d_embed // (div_val ** i)\n",
    "                    lookup_table = tf.get_variable(\n",
    "                        'lookup_table',\n",
    "                        [r_idx - l_idx, cur_d_embed],\n",
    "                        initializer = initializer,\n",
    "                    )\n",
    "                    cur_y = embedding_lookup(lookup_table, cur_x)\n",
    "                    if d_proj == cur_d_embed and not proj_same_dim:\n",
    "                        proj_W = None\n",
    "                    else:\n",
    "                        proj_W = tf.get_variable(\n",
    "                            'proj_W',\n",
    "                            [cur_d_embed, d_proj],\n",
    "                            initializer = proj_initializer,\n",
    "                        )\n",
    "                        cur_y = tf.einsum('id,de->ie', cur_y, proj_W)\n",
    "                    mask_idx = tf.to_int64(tf.where(mask))\n",
    "                    y += tf.scatter_nd(\n",
    "                        mask_idx, cur_y, tf.to_int64(tf.shape(y))\n",
    "                    )\n",
    "                    tables.append(lookup_table)\n",
    "                    projs.append(proj_W)\n",
    "            ret_params = [tables, projs]\n",
    "\n",
    "    y *= emb_scale\n",
    "    return y, ret_params\n",
    "\n",
    "\n",
    "def _create_mask(qlen, mlen, same_length = False):\n",
    "    attn_mask = tf.ones([qlen, qlen])\n",
    "    mask_u = tf.matrix_band_part(attn_mask, 0, -1)\n",
    "    mask_dia = tf.matrix_band_part(attn_mask, 0, 0)\n",
    "    attn_mask_pad = tf.zeros([qlen, mlen])\n",
    "    ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)\n",
    "    if same_length:\n",
    "        mask_l = tf.matrix_band_part(attn_mask, -1, 0)\n",
    "        ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)\n",
    "    return ret\n",
    "\n",
    "\n",
    "def _cache_mem(curr_out, prev_mem, mem_len = None):\n",
    "    if mem_len is None or prev_mem is None:\n",
    "        new_mem = curr_out\n",
    "    elif mem_len == 0:\n",
    "        return prev_mem\n",
    "    else:\n",
    "        new_mem = tf.concat([prev_mem, curr_out], 0)[-mem_len:]\n",
    "\n",
    "    return tf.stop_gradient(new_mem)\n",
    "\n",
    "\n",
    "def transformer(\n",
    "    dec_inp,\n",
    "    mems,\n",
    "    n_token,\n",
    "    n_layer,\n",
    "    d_model,\n",
    "    d_embed,\n",
    "    n_head,\n",
    "    d_head,\n",
    "    d_inner,\n",
    "    initializer,\n",
    "    proj_initializer = None,\n",
    "    mem_len = None,\n",
    "    cutoffs = [],\n",
    "    div_val = 1,\n",
    "    tie_projs = [],\n",
    "    same_length = False,\n",
    "    clamp_len = -1,\n",
    "    untie_r = False,\n",
    "    proj_same_dim = True,\n",
    "    scope = 'transformer',\n",
    "    reuse = tf.AUTO_REUSE\n",
    "):\n",
    "    \"\"\"\n",
    "  cutoffs: a list of python int. Cutoffs for adaptive softmax.\n",
    "  tie_projs: a list of python bools. Whether to tie the projections.\n",
    "  perms: a list of tensors. Each tensor should of size [len, bsz, bin_size].\n",
    "        Only used in the adaptive setting.\n",
    "  \"\"\"\n",
    "    new_mems = []\n",
    "    with tf.variable_scope(scope,reuse=reuse):\n",
    "        if untie_r:\n",
    "            r_w_bias = tf.get_variable(\n",
    "                'r_w_bias', [n_layer, n_head, d_head], initializer = initializer\n",
    "            )\n",
    "            r_r_bias = tf.get_variable(\n",
    "                'r_r_bias', [n_layer, n_head, d_head], initializer = initializer\n",
    "            )\n",
    "        else:\n",
    "            r_w_bias = tf.get_variable(\n",
    "                'r_w_bias', [n_head, d_head], initializer = initializer\n",
    "            )\n",
    "            r_r_bias = tf.get_variable(\n",
    "                'r_r_bias', [n_head, d_head], initializer = initializer\n",
    "            )\n",
    "\n",
    "        qlen = tf.shape(dec_inp)[0]\n",
    "        mlen = tf.shape(mems[0])[0] if mems is not None else 0\n",
    "        klen = mlen + qlen\n",
    "\n",
    "        if proj_initializer is None:\n",
    "            proj_initializer = initializer\n",
    "        lookup_fn = mask_adaptive_embedding_lookup\n",
    "        embeddings, shared_params = lookup_fn(\n",
    "            x = dec_inp,\n",
    "            n_token = n_token,\n",
    "            d_embed = d_embed,\n",
    "            d_proj = d_model,\n",
    "            cutoffs = cutoffs,\n",
    "            initializer = initializer,\n",
    "            proj_initializer = proj_initializer,\n",
    "            div_val = div_val,\n",
    "            proj_same_dim = proj_same_dim,\n",
    "        )\n",
    "\n",
    "        attn_mask = _create_mask(qlen, mlen, same_length)\n",
    "\n",
    "        pos_seq = tf.range(klen - 1, -1, -1.0)\n",
    "        if clamp_len > 0:\n",
    "            pos_seq = tf.minimum(pos_seq, clamp_len)\n",
    "        inv_freq = 1 / (10000 ** (tf.range(0, d_model, 2.0) / d_model))\n",
    "        pos_emb = positional_embedding(pos_seq, inv_freq)\n",
    "\n",
    "        if mems is None:\n",
    "            mems = [None] * n_layer\n",
    "        output = embeddings\n",
    "        for i in range(n_layer):\n",
    "            # cache new mems\n",
    "            new_mems.append(_cache_mem(output, mems[i], mem_len))\n",
    "\n",
    "            with tf.variable_scope('layer_{}'.format(i)):\n",
    "                output = rel_multihead_attn(\n",
    "                    w = output,\n",
    "                    r = pos_emb,\n",
    "                    r_w_bias = r_w_bias if not untie_r else r_w_bias[i],\n",
    "                    r_r_bias = r_r_bias if not untie_r else r_r_bias[i],\n",
    "                    attn_mask = attn_mask,\n",
    "                    mems = mems[i],\n",
    "                    d_model = d_model,\n",
    "                    n_head = n_head,\n",
    "                    d_head = d_head,\n",
    "                    kernel_initializer = initializer,\n",
    "                )\n",
    "                output = positionwise_FF(\n",
    "                    inp = output,\n",
    "                    d_model = d_model,\n",
    "                    d_inner = d_inner,\n",
    "                    kernel_initializer = initializer,\n",
    "                )\n",
    "\n",
    "        return output, new_mems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Chatbot:\n",
    "    def __init__(self):\n",
    "\n",
    "        self.X = tf.placeholder(tf.int32, [None, None])\n",
    "        self.Y = 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, 1, dtype = tf.int32)\n",
    "        batch_size = tf.shape(self.X)[0]\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",
    "        initializer = tf.initializers.random_normal(stddev = 0.1)\n",
    "        memory = tf.fill(\n",
    "                [n_layer, tf.shape(self.X)[0], tf.shape(self.X)[1], d_model], PAD\n",
    "            )\n",
    "        memory = tf.cast(memory, tf.float32)\n",
    "        logits, next_memory = transformer(\n",
    "                self.X,\n",
    "                memory,\n",
    "                len(dictionary_from),\n",
    "                n_layer,\n",
    "                d_model,\n",
    "                d_embed,\n",
    "                n_head,\n",
    "                d_head,\n",
    "                d_inner,\n",
    "                initializer,\n",
    "                scope='encoder',\n",
    "                reuse=False\n",
    "            )\n",
    "        logits, next_memory = transformer(\n",
    "                self.X,\n",
    "                next_memory,\n",
    "                len(dictionary_from),\n",
    "                n_layer,\n",
    "                d_model,\n",
    "                d_embed,\n",
    "                n_head,\n",
    "                d_head,\n",
    "                d_inner,\n",
    "                initializer,\n",
    "                scope='decoder',\n",
    "                reuse=False\n",
    "            )\n",
    "        embedding = tf.Variable(tf.random_uniform([len(dictionary_to), d_model], -1, 1))\n",
    "        init_state = tf.reduce_mean(tf.reduce_mean(next_memory,axis=0),axis=1)\n",
    "        cell = tf.nn.rnn_cell.LSTMCell(d_model)\n",
    "        vocab_proj = tf.layers.Dense(len(dictionary_to))\n",
    "        helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "            inputs = tf.nn.embedding_lookup(embedding, decoder_input),\n",
    "            sequence_length = tf.to_int32(self.Y_seq_len))\n",
    "        encoder_state = tf.nn.rnn_cell.LSTMStateTuple(c=init_state, h=init_state)\n",
    "        decoder = tf.contrib.seq2seq.BasicDecoder(cell = cell,\n",
    "                                                  helper = helper,\n",
    "                                                  initial_state = encoder_state,\n",
    "                                                  output_layer = vocab_proj)\n",
    "        decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(decoder = decoder,\n",
    "                                                                maximum_iterations = tf.reduce_max(self.Y_seq_len))\n",
    "        \n",
    "        helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(embedding = embedding,\n",
    "                                                          start_tokens = tf.tile(\n",
    "                                                              tf.constant([GO], \n",
    "                                                                          dtype=tf.int32), \n",
    "                                                              [tf.shape(init_state)[0]]),\n",
    "                                                          end_token = EOS)\n",
    "        decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "            cell = cell,\n",
    "            helper = helper,\n",
    "            initial_state = encoder_state,\n",
    "            output_layer = vocab_proj)\n",
    "        predicting_decoder_output, _, _ = tf.contrib.seq2seq.dynamic_decode(\n",
    "            decoder = decoder,\n",
    "            maximum_iterations = 2 * tf.reduce_max(self.X_seq_len))\n",
    "        self.training_logits = decoder_output.rnn_output\n",
    "        self.predicting_ids = predicting_decoder_output.sample_id\n",
    "        self.logits = decoder_output.sample_id\n",
    "        masks = tf.sequence_mask(self.Y_seq_len, tf.reduce_max(self.Y_seq_len), dtype=tf.float32)\n",
    "        self.cost = tf.contrib.seq2seq.sequence_loss(logits = self.training_logits,\n",
    "                                                     targets = self.Y,\n",
    "                                                     weights = masks)\n",
    "        self.optimizer = tf.train.AdamOptimizer(learning_rate).minimize(self.cost)\n",
    "        y_t = tf.argmax(self.training_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": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Chatbot()\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad_sentence_batch_dynamic(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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, avg loss: 5.496688, avg accuracy: 0.227262\n",
      "epoch: 2, avg loss: 4.586239, avg accuracy: 0.272152\n",
      "epoch: 3, avg loss: 4.316728, avg accuracy: 0.294077\n",
      "epoch: 4, avg loss: 4.028886, avg accuracy: 0.307932\n",
      "epoch: 5, avg loss: 3.730364, avg accuracy: 0.326791\n",
      "epoch: 6, avg loss: 3.432248, avg accuracy: 0.346253\n",
      "epoch: 7, avg loss: 3.162382, avg accuracy: 0.371888\n",
      "epoch: 8, avg loss: 2.906495, avg accuracy: 0.401091\n",
      "epoch: 9, avg loss: 2.673143, avg accuracy: 0.423848\n",
      "epoch: 10, avg loss: 2.442581, avg accuracy: 0.484632\n",
      "epoch: 11, avg loss: 2.218746, avg accuracy: 0.518694\n",
      "epoch: 12, avg loss: 1.971839, avg accuracy: 0.594742\n",
      "epoch: 13, avg loss: 1.753959, avg accuracy: 0.655731\n",
      "epoch: 14, avg loss: 1.587028, avg accuracy: 0.698523\n",
      "epoch: 15, avg loss: 1.451729, avg accuracy: 0.727547\n",
      "epoch: 16, avg loss: 1.331589, avg accuracy: 0.751697\n",
      "epoch: 17, avg loss: 1.238719, avg accuracy: 0.770444\n",
      "epoch: 18, avg loss: 1.153919, avg accuracy: 0.785509\n",
      "epoch: 19, avg loss: 1.070685, avg accuracy: 0.806163\n",
      "epoch: 20, avg loss: 0.999631, avg accuracy: 0.826571\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_dynamic(X[k: index], PAD)\n",
    "        batch_y, seq_y = pad_sentence_batch_dynamic(Y[k: index], PAD)\n",
    "        predicted, accuracy,loss, _ = sess.run([model.predicting_ids,model.accuracy, model.cost, model.optimizer], \n",
    "                                      feed_dict={model.X:batch_x,\n",
    "                                                model.Y:batch_y})\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": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "row 1\n",
      "QUESTION: i am a werewolf\n",
      "REAL ANSWER: a werewolf\n",
      "PREDICTED ANSWER: a werewolf \n",
      "\n",
      "row 2\n",
      "QUESTION: i was dreaming again\n",
      "REAL ANSWER: i would think so\n",
      "PREDICTED ANSWER: i would think so \n",
      "\n",
      "row 3\n",
      "QUESTION: the kitchen\n",
      "REAL ANSWER: very nice\n",
      "PREDICTED ANSWER: very nice \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": 18,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "row 1\n",
      "QUESTION: but david\n",
      "REAL ANSWER: is here that\n",
      "PREDICTED ANSWER: you have a good time \n",
      "\n",
      "row 2\n",
      "QUESTION: hopeless it is hopeless\n",
      "REAL ANSWER: tell ballet then back\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 3\n",
      "QUESTION: miss price\n",
      "REAL ANSWER: yes learning\n",
      "PREDICTED ANSWER: you have a good time \n",
      "\n",
      "row 4\n",
      "QUESTION: mr kessler wake up please\n",
      "REAL ANSWER: is here are\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 5\n",
      "QUESTION: there were witnesses\n",
      "REAL ANSWER: why she out\n",
      "PREDICTED ANSWER: i am on \n",
      "\n",
      "row 6\n",
      "QUESTION: what about it\n",
      "REAL ANSWER: not you are\n",
      "PREDICTED ANSWER: i will be right \n",
      "\n",
      "row 7\n",
      "QUESTION: go on ask them\n",
      "REAL ANSWER: i just home\n",
      "PREDICTED ANSWER: you have a good \n",
      "\n",
      "row 8\n",
      "QUESTION: beware the moon\n",
      "REAL ANSWER: seen hi is he\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 9\n",
      "QUESTION: did you hear that\n",
      "REAL ANSWER: is down what\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 10\n",
      "QUESTION: i heard that\n",
      "REAL ANSWER: it here not\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 11\n",
      "QUESTION: the hound of the baskervilles\n",
      "REAL ANSWER: heard\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 12\n",
      "QUESTION: it is moving\n",
      "REAL ANSWER: not you hear\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 13\n",
      "QUESTION: nice doggie good boy\n",
      "REAL ANSWER: bill stupid\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 14\n",
      "QUESTION: it sounds far away\n",
      "REAL ANSWER: that pecos baby seen hi\n",
      "PREDICTED ANSWER: i will be another \n",
      "\n",
      "row 15\n",
      "QUESTION: debbie klein cried a lot\n",
      "REAL ANSWER: is will srai not\n",
      "PREDICTED ANSWER: i am so \n",
      "\n",
      "row 16\n",
      "QUESTION: what are you doing here\n",
      "REAL ANSWER: is know look i\n",
      "PREDICTED ANSWER: i am so \n",
      "\n"
     ]
    }
   ],
   "source": [
    "batch_x, seq_x = pad_sentence_batch(X_test[:batch_size], PAD, maxlen_question)\n",
    "batch_y, seq_y = pad_sentence_batch(Y_test[:batch_size], PAD, maxlen_answer)\n",
    "predicted = sess.run(model.predicting_ids, feed_dict={model.X:batch_x})\n",
    "\n",
    "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
}