neural-machine-translation/45.attention-is-all-you-need.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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len from: 500, len to: 500\n"
     ]
    }
   ],
   "source": [
    "with open('english-train', 'r') as fopen:\n",
    "    text_from = fopen.read().lower().split('\\n')[:-1]\n",
    "with open('vietnam-train', 'r') as fopen:\n",
    "    text_to = fopen.read().lower().split('\\n')[:-1]\n",
    "print('len from: %d, len to: %d'%(len(text_from), len(text_to)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab from size: 1935\n",
      "Most common words [(',', 564), ('.', 477), ('the', 368), ('and', 286), ('to', 242), ('of', 220)]\n",
      "Sample data [482, 483, 78, 6, 137, 484, 10, 226, 787, 14] ['rachel', 'pike', ':', 'the', 'science', 'behind', 'a', 'climate', 'headline', 'in']\n"
     ]
    }
   ],
   "source": [
    "concat_from = ' '.join(text_from).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]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "vocab to size: 1461\n",
      "Most common words [(',', 472), ('.', 430), ('tôi', 283), ('và', 230), ('có', 199), ('chúng', 196)]\n",
      "Sample data [84, 22, 668, 73, 10, 389, 110, 34, 81, 299] ['khoa', 'học', 'đằng', 'sau', 'một', 'tiêu', 'đề', 'về', 'khí', 'hậu']\n"
     ]
    }
   ],
   "source": [
    "concat_to = ' '.join(text_to).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 to 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]])"
   ]
  },
  {
   "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(text_to)):\n",
    "    text_to[i] += ' EOS'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def layer_norm(inputs, epsilon=1e-8):\n",
    "    mean, variance = tf.nn.moments(inputs, [-1], keep_dims=True)\n",
    "    normalized = (inputs - mean) / (tf.sqrt(variance + epsilon))\n",
    "\n",
    "    params_shape = inputs.get_shape()[-1:]\n",
    "    gamma = tf.get_variable('gamma', params_shape, tf.float32, tf.ones_initializer())\n",
    "    beta = tf.get_variable('beta', params_shape, tf.float32, tf.zeros_initializer())\n",
    "    \n",
    "    outputs = gamma * normalized + beta\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def multihead_attn(queries, keys, q_masks, k_masks, future_binding, num_units, num_heads):\n",
    "    \n",
    "    T_q = tf.shape(queries)[1]                                      \n",
    "    T_k = tf.shape(keys)[1]                  \n",
    "\n",
    "    Q = tf.layers.dense(queries, num_units, name='Q')                              \n",
    "    K_V = tf.layers.dense(keys, 2*num_units, name='K_V')    \n",
    "    K, V = tf.split(K_V, 2, -1)        \n",
    "\n",
    "    Q_ = tf.concat(tf.split(Q, num_heads, axis=2), axis=0)                         \n",
    "    K_ = tf.concat(tf.split(K, num_heads, axis=2), axis=0)                    \n",
    "    V_ = tf.concat(tf.split(V, num_heads, axis=2), axis=0)                      \n",
    "\n",
    "    align = tf.matmul(Q_, tf.transpose(K_, [0,2,1]))                      \n",
    "    align = align / np.sqrt(K_.get_shape().as_list()[-1])                 \n",
    "\n",
    "    paddings = tf.fill(tf.shape(align), float('-inf'))                   \n",
    "\n",
    "    key_masks = k_masks                                                 \n",
    "    key_masks = tf.tile(key_masks, [num_heads, 1])                       \n",
    "    key_masks = tf.tile(tf.expand_dims(key_masks, 1), [1, T_q, 1])            \n",
    "    align = tf.where(tf.equal(key_masks, 0), paddings, align)       \n",
    "\n",
    "    if future_binding:\n",
    "        lower_tri = tf.ones([T_q, T_k])                                          \n",
    "        lower_tri = tf.linalg.LinearOperatorLowerTriangular(lower_tri).to_dense()  \n",
    "        masks = tf.tile(tf.expand_dims(lower_tri,0), [tf.shape(align)[0], 1, 1]) \n",
    "        align = tf.where(tf.equal(masks, 0), paddings, align)                      \n",
    "    \n",
    "    align = tf.nn.softmax(align)                                            \n",
    "    query_masks = tf.to_float(q_masks)                                             \n",
    "    query_masks = tf.tile(query_masks, [num_heads, 1])                             \n",
    "    query_masks = tf.tile(tf.expand_dims(query_masks, -1), [1, 1, T_k])            \n",
    "    align *= query_masks                                                           \n",
    "          \n",
    "    outputs = tf.matmul(align, V_)                                                 \n",
    "    outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2)             \n",
    "    outputs += queries                                                             \n",
    "    outputs = layer_norm(outputs)                                                 \n",
    "    return outputs\n",
    "\n",
    "\n",
    "def pointwise_feedforward(inputs, hidden_units, activation=None):\n",
    "    outputs = tf.layers.dense(inputs, 4*hidden_units, activation=activation)\n",
    "    outputs = tf.layers.dense(outputs, hidden_units, activation=None)\n",
    "    outputs += inputs\n",
    "    outputs = layer_norm(outputs)\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def learned_position_encoding(inputs, mask, embed_dim):\n",
    "    T = tf.shape(inputs)[1]\n",
    "    outputs = tf.range(tf.shape(inputs)[1])                # (T_q)\n",
    "    outputs = tf.expand_dims(outputs, 0)                   # (1, T_q)\n",
    "    outputs = tf.tile(outputs, [tf.shape(inputs)[0], 1])   # (N, T_q)\n",
    "    outputs = embed_seq(outputs, T, embed_dim, zero_pad=False, scale=False)\n",
    "    return tf.expand_dims(tf.to_float(mask), -1) * outputs\n",
    "\n",
    "\n",
    "def sinusoidal_position_encoding(inputs, mask, repr_dim):\n",
    "    T = tf.shape(inputs)[1]\n",
    "    pos = tf.reshape(tf.range(0.0, tf.to_float(T), dtype=tf.float32), [-1, 1])\n",
    "    i = np.arange(0, repr_dim, 2, np.float32)\n",
    "    denom = np.reshape(np.power(10000.0, i / repr_dim), [1, -1])\n",
    "    enc = tf.expand_dims(tf.concat([tf.sin(pos / denom), tf.cos(pos / denom)], 1), 0)\n",
    "    return tf.tile(enc, [tf.shape(inputs)[0], 1, 1]) * tf.expand_dims(tf.to_float(mask), -1)\n",
    "\n",
    "\n",
    "def label_smoothing(inputs, epsilon=0.1):\n",
    "    C = inputs.get_shape().as_list()[-1]\n",
    "    return ((1 - epsilon) * inputs) + (epsilon / C)\n",
    "\n",
    "\n",
    "class Chatbot:\n",
    "    def __init__(self, size_layer, embedded_size, from_dict_size, to_dict_size, learning_rate,\n",
    "                 num_blocks = 2,\n",
    "                 num_heads = 8,\n",
    "                 min_freq = 50):\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",
    "        \n",
    "        encoder_embedding = tf.Variable(tf.random_uniform([from_dict_size, embedded_size], -1, 1))\n",
    "        decoder_embedding = tf.Variable(tf.random_uniform([to_dict_size, embedded_size], -1, 1))\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",
    "        def forward(x, y):\n",
    "            encoder_embedded = tf.nn.embedding_lookup(encoder_embedding, x)\n",
    "            en_masks = tf.sign(x)\n",
    "            encoder_embedded += sinusoidal_position_encoding(x, en_masks, embedded_size)\n",
    "        \n",
    "            for i in range(num_blocks):\n",
    "                with tf.variable_scope('encoder_self_attn_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                    encoder_embedded = multihead_attn(queries = encoder_embedded,\n",
    "                                             keys = encoder_embedded,\n",
    "                                             q_masks = en_masks,\n",
    "                                             k_masks = en_masks,\n",
    "                                             future_binding = False,\n",
    "                                             num_units = size_layer,\n",
    "                                             num_heads = num_heads)\n",
    "\n",
    "                with tf.variable_scope('encoder_feedforward_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                    encoder_embedded = pointwise_feedforward(encoder_embedded,\n",
    "                                                    embedded_size,\n",
    "                                                    activation = tf.nn.relu)\n",
    "            \n",
    "            decoder_embedded = tf.nn.embedding_lookup(decoder_embedding, y)\n",
    "            de_masks = tf.sign(y)\n",
    "            decoder_embedded += sinusoidal_position_encoding(y, de_masks, embedded_size)\n",
    "            \n",
    "            for i in range(num_blocks):\n",
    "                with tf.variable_scope('decoder_self_attn_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                    decoder_embedded = multihead_attn(queries = decoder_embedded,\n",
    "                                         keys = decoder_embedded,\n",
    "                                         q_masks = de_masks,\n",
    "                                         k_masks = de_masks,\n",
    "                                         future_binding = True,\n",
    "                                         num_units = size_layer,\n",
    "                                         num_heads = num_heads)\n",
    "                \n",
    "                with tf.variable_scope('decoder_attn_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                    decoder_embedded = multihead_attn(queries = decoder_embedded,\n",
    "                                         keys = encoder_embedded,\n",
    "                                         q_masks = de_masks,\n",
    "                                         k_masks = en_masks,\n",
    "                                         future_binding = False,\n",
    "                                         num_units = size_layer,\n",
    "                                         num_heads = num_heads)\n",
    "                \n",
    "                with tf.variable_scope('decoder_feedforward_%d'%i,reuse=tf.AUTO_REUSE):\n",
    "                    decoder_embedded = pointwise_feedforward(decoder_embedded,\n",
    "                                                    embedded_size,\n",
    "                                            activation = tf.nn.relu)\n",
    "            \n",
    "            return tf.layers.dense(decoder_embedded, to_dict_size, reuse=tf.AUTO_REUSE)\n",
    "        \n",
    "        self.training_logits = forward(self.X, decoder_input)\n",
    "        \n",
    "        def cond(i, y, temp):\n",
    "            return i < 2 * tf.reduce_max(self.X_seq_len)\n",
    "        \n",
    "        def body(i, y, temp):\n",
    "            logits = forward(self.X, y)\n",
    "            ids = tf.argmax(logits, -1)[:, i]\n",
    "            ids = tf.expand_dims(ids, -1)\n",
    "            temp = tf.concat([temp[:, 1:], ids], -1)\n",
    "            y = tf.concat([temp[:, -(i+1):], temp[:, :-(i+1)]], -1)\n",
    "            y = tf.reshape(y, [tf.shape(temp)[0], 2 * tf.reduce_max(self.X_seq_len)])\n",
    "            i += 1\n",
    "            return i, y, temp\n",
    "        \n",
    "        target = tf.fill([batch_size, 2 * tf.reduce_max(self.X_seq_len)], GO)\n",
    "        target = tf.cast(target, tf.int64)\n",
    "        self.target = target\n",
    "        \n",
    "        _, self.predicting_ids, _ = tf.while_loop(cond, body, \n",
    "                                                  [tf.constant(0), target, target])\n",
    "        \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 = 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": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedded_size = 256\n",
    "learning_rate = 0.001\n",
    "batch_size = 16\n",
    "epoch = 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "sess = tf.InteractiveSession()\n",
    "model = Chatbot(embedded_size, embedded_size, len(dictionary_from), \n",
    "                len(dictionary_to), learning_rate)\n",
    "sess.run(tf.global_variables_initializer())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = str_idx(text_from, dictionary_from)\n",
    "Y = str_idx(text_to, dictionary_to)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "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": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1, avg loss: 6.446041, avg accuracy: 0.063624\n",
      "epoch: 2, avg loss: 6.000207, avg accuracy: 0.087263\n",
      "epoch: 3, avg loss: 5.758100, avg accuracy: 0.105811\n",
      "epoch: 4, avg loss: 5.600954, avg accuracy: 0.123585\n",
      "epoch: 5, avg loss: 5.564998, avg accuracy: 0.113341\n",
      "epoch: 6, avg loss: 5.493285, avg accuracy: 0.120724\n",
      "epoch: 7, avg loss: 5.529355, avg accuracy: 0.117925\n",
      "epoch: 8, avg loss: 5.522790, avg accuracy: 0.119511\n",
      "epoch: 9, avg loss: 5.488758, avg accuracy: 0.124838\n",
      "epoch: 10, avg loss: 5.481946, avg accuracy: 0.124881\n",
      "epoch: 11, avg loss: 5.495231, avg accuracy: 0.115534\n",
      "epoch: 12, avg loss: 5.358830, avg accuracy: 0.120977\n",
      "epoch: 13, avg loss: 5.233122, avg accuracy: 0.136445\n",
      "epoch: 14, avg loss: 5.153393, avg accuracy: 0.143617\n",
      "epoch: 15, avg loss: 5.057733, avg accuracy: 0.152985\n",
      "epoch: 16, avg loss: 4.906600, avg accuracy: 0.162316\n",
      "epoch: 17, avg loss: 4.826649, avg accuracy: 0.164304\n",
      "epoch: 18, avg loss: 4.789770, avg accuracy: 0.162977\n",
      "epoch: 19, avg loss: 4.777955, avg accuracy: 0.166525\n",
      "epoch: 20, avg loss: 4.712728, avg accuracy: 0.170279\n"
     ]
    }
   ],
   "source": [
    "for i in range(epoch):\n",
    "    total_loss, total_accuracy = 0, 0\n",
    "    for k in range(0, len(text_to), batch_size):\n",
    "        index = min(k+batch_size, len(text_to))\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.predicting_ids, \n",
    "                                                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(text_to) / batch_size)\n",
    "    total_accuracy /= (len(text_to) / batch_size)\n",
    "    print('epoch: %d, avg loss: %f, avg accuracy: %f'%(i+1, total_loss, total_accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "row 1\n",
      "QUESTION: or , if you had to choose between the last two , which one would you choose ?\n",
      "REAL ANSWER: sau khi thôi không làm đau mọi người nữa , tôi sẽ hỏi họ bạn có đau không ? đau như thế nào ?\n",
      "PREDICTED ANSWER: và đây là một một đây một một một lớp lớp một một một một một một một những những cách cách những những những những lớp cách cách cách cách cách cách cách cách cách những những một một một những cách cách những những những những những cách \n",
      "\n",
      "row 2\n",
      "QUESTION: i kept on doing this for a while .\n",
      "REAL ANSWER: hoặc nếu được chọn giữa 2 kiểu đau cuối , bạn sẽ chọn cái nào ?\n",
      "PREDICTED ANSWER: và đây là một một đây một một một lớp lớp một một một một một một một những những cách cách những những những những lớp cách cách cách cách cách cách cách cách cách những những một một một những cách cách những những những những những cách \n",
      "\n",
      "row 3\n",
      "QUESTION: and then , like all good academic projects , i got more funding .\n",
      "REAL ANSWER: tôi tiếp tục làm thí nghiệm này 1 thời gian\n",
      "PREDICTED ANSWER: và đây là một một đây một một một lớp lớp một một một một một một một những những cách cách những những những những lớp cách cách cách cách cách cách cách cách cách những những một một một những cách cách những những những những những cách \n",
      "\n",
      "row 4\n",
      "QUESTION: i moved to sounds , electrical shocks -- i even had a pain suit that i could get people to feel much more pain .\n",
      "REAL ANSWER: và sau đó , giống các đề tài nghiên cứu hay khác , tôi nhận thêm nguồn tài trợ .\n",
      "PREDICTED ANSWER: và đây là một một đây một một một lớp lớp một một một một một một một những những cách cách những những những những lớp cách cách cách cách cách cách cách cách cách những những một một một những cách cách những những những những những cách \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
}