{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hackathon 9\n",
    "\n",
    "Topics:\n",
    "- Neural Machine Translation (NTM)\n",
    "- Sequence to Sequence (Seq2Seq)\n",
    "- Attention in NMT\n",
    "\n",
    "In today's demo, we'll teach an model how to translate from English to French. Significant portions sourced from [TensorFlow's documentation](https://www.tensorflow.org/tutorials/seq2seq).\n",
    "\n",
    "This is all setup in a IPython notebook so you can run any code you want to experiment with. Feel free to edit any cell, or add some to run your own code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# we'll start with our library imports...\n",
    "# tensorflow to specify and run computation graphs\n",
    "# numpy to run any numerical operations that need to take place outside of the TF graph\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "# these ones let us draw images in our notebook\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Machine Translation\n",
    "\n",
    "Researchers at Google have pioneered Neural Machine Translation (NMT), a method of translating between natural languages using RNNs-based Seq2Seq models. Previous models relied on memorizing phrase translations and combining those without regard for context. This led, famously, to translations from the English \"The spirit is willing, but the flesh is weak\" to Russian, back to English as \"The vodka is good, but the meat is rotten\". NTM uses context and learned experience to translate using deep networks.\n",
    "\n",
    "NTM relies on a model called Sequence to Sequence (Seq2Seq). Basically, it reads the input sentence one word at a time with an encoder, and outputs a fixed-length \"thought vector\". Then, a decoder unrolls the thought vector into a sentence in the target language.\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/encdec.jpg\">\n",
    "\n",
    "Specifically, the English encoder will take one input word per timestep, but not be expected to produce any output while it's reading the input sentence. After reading in the full sentence, we'll use the cell state as the thought vector that we'll use as the initial state of the French decoder. The decoder will then process the input sentence as it predicts the words of the output sentence.\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/seq2seq.jpg\">\n",
    "Here, \"(s)\" is a special character that marks the start of the decoding process while \"(/s)\" is a special character that tells the decoder to stop."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as in last week's hackathon, the first thing we'll do is a word embedding by mapping from integers to learned vectors. This week, we'll be working with time-major data. This simply switches the `batch` and `time` dimensions, but we need to carefully track this in the code to ensure correctness. We have some code here to load vocabularies and data from source and target languages. The data is an English-Vietnamese corpus of TED talks (133K sentence pairs) provided by the IWSLT Evaluation Campaign."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import codecs\n",
    "\n",
    "def load_data(inference_input_file, hparams=None):\n",
    "  \"\"\"Load inference data.\"\"\"\n",
    "  with codecs.getreader(\"utf-8\")(\n",
    "      tf.gfile.GFile(inference_input_file, mode=\"rb\")) as f:\n",
    "    inference_data = f.read().splitlines()\n",
    "\n",
    "  if hparams and hparams.inference_indices:\n",
    "    inference_data = [inference_data[i] for i in hparams.inference_indices]\n",
    "\n",
    "  return inference_data\n",
    "\n",
    "\n",
    "def get_sequences(src_data, dst_data, src_encode_fn, dst_encode_fn, num, go_symbol, stop_symbol):\n",
    "    # pull from source data\n",
    "    decode_ids = np.random.randint(0, len(src_data) - 1, size=num)\n",
    "    sentences = [src_data[x].split() for x in decode_ids]\n",
    "    input_lengths = [len(s) for s in sentences]\n",
    "    \n",
    "    # pull from target data\n",
    "    target_inputs = [[go_symbol] + dst_data[x].split() for x in decode_ids]\n",
    "    target_outputs = [dst_data[x].split() + [stop_symbol] for x in decode_ids]\n",
    "    target_lengths = [len(s) for s in target_inputs]\n",
    "\n",
    "    # pad sequences to uniform lengths and be time-major\n",
    "    def zero_pad(sequences, lengths):\n",
    "        padded = np.zeros([np.max(lengths), len(sequences)], dtype=np.int32)\n",
    "        for (i, s) in enumerate(sequences):\n",
    "            padded[0:len(s), i] = s\n",
    "        return padded\n",
    "    \n",
    "    input_sentences = zero_pad(src_encode_fn(sentences), input_lengths)\n",
    "    target_inputs = zero_pad(dst_encode_fn(target_inputs), target_lengths)\n",
    "    target_outputs = zero_pad(dst_encode_fn(target_outputs), target_lengths)\n",
    "    \n",
    "    return input_sentences, input_lengths, target_inputs, target_outputs, target_lengths\n",
    "\n",
    "\n",
    "def code(vocab_dict, sequences):\n",
    "    \"\"\"Use vocab for int -> word and inv_vocab for word -> int\"\"\"\n",
    "    return [[vocab_dict.get(x, list(vocab_dict.values())[0]) for x in s] for s in sequences]\n",
    "\n",
    "\n",
    "list_to_dict = lambda l: {k: v for k, v in zip(l, range(len(l)))}\n",
    "invert_dict = lambda d: {v: k for k, v in d.items()}\n",
    "\n",
    "# load source vocab and data\n",
    "base_path = '/work/cse496dl/shared/hackathon/09/'\n",
    "src_vocab = list_to_dict(load_data(base_path + 'vocab.en'))\n",
    "encode_src = lambda s: code(src_vocab, s)\n",
    "decode_src = lambda s: code(invert_dict(src_vocab), s)\n",
    "src_data = load_data(base_path + 'train.en')\n",
    "\n",
    "# load target vocab and data\n",
    "dst_vocab = list_to_dict(load_data(base_path + 'vocab.vi'))\n",
    "encode_dst = lambda s: code(dst_vocab, s)\n",
    "decode_dst = lambda s: code(invert_dict(dst_vocab), s)\n",
    "dst_data = load_data(base_path + 'train.vi')\n",
    "\n",
    "# set constants for later\n",
    "SRC_VOCAB_SIZE = len(src_vocab)\n",
    "DST_VOCAB_SIZE = len(dst_vocab)\n",
    "GO_SYMBOL = '<s>'\n",
    "END_SYMBOL = '</s>'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Just as in the last hackathon, we'll use learned word embeddings for both source and target languages. Different from the last hackathon, we're not using anyone else's code for loading the data, so we'll use placeholders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "EMBEDDING_SIZE = 100\n",
    "MAX_TIME = 20\n",
    "\n",
    "tf.reset_default_graph()\n",
    "encoder_inputs = tf.placeholder(tf.int32, shape=[None, None])\n",
    "source_sequence_length = tf.placeholder(tf.int32, shape=[None])\n",
    "\n",
    "# Embedding\n",
    "src_embedding_matrix = tf.get_variable('src_embedding_matrix', dtype=tf.float32,\n",
    "                                   shape=[SRC_VOCAB_SIZE, EMBEDDING_SIZE], trainable=True)\n",
    "\n",
    "# Look up embedding:\n",
    "#   encoder_inputs: [max_time, batch_size]\n",
    "#   encoder_emb_inp: [max_time, batch_size, embedding_size]\n",
    "encoder_emb_inp = tf.nn.embedding_lookup(src_embedding_matrix, encoder_inputs)\n",
    "\n",
    "# decoder placeholders\n",
    "decoder_inputs = tf.placeholder(tf.int32, shape=[None, None])\n",
    "decoder_lengths = tf.placeholder(tf.int32, shape=[None])\n",
    "\n",
    "# Embed decoder input\n",
    "dst_embedding_matrix = tf.get_variable('dst_embedding_matrix', dtype=tf.float32,\n",
    "                                   shape=[DST_VOCAB_SIZE, EMBEDDING_SIZE], trainable=True)\n",
    "decoder_emb_inp = tf.nn.embedding_lookup(dst_embedding_matrix, decoder_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll declare the `LSTMCell` just as in the last hackathon, but this time we'll make sure to pass the `time_major` and `sequence_length` arguments to `dynamic_rnn`. We can use the default zero state to start the LSTM by passing the data type argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "NUM_UNITS = 200\n",
    "\n",
    "# Build RNN cell\n",
    "encoder_cell = tf.nn.rnn_cell.BasicLSTMCell(NUM_UNITS)\n",
    "\n",
    "# Run Dynamic RNN\n",
    "#   encoder_outputs: [max_time, batch_size, num_units]\n",
    "#   encoder_state: [batch_size, num_units]\n",
    "encoder_outputs, encoder_state = tf.nn.dynamic_rnn(\n",
    "    encoder_cell, encoder_emb_inp,\n",
    "    sequence_length=source_sequence_length, time_major=True, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the same seq2seq `sequence_loss` again, but this time we dynamically set the batch size and time steps to allow variable shaped inputs.\n",
    "\n",
    "Google Brain's seq2seq code uses `Helper` objects to do dynamic decoding. Factoring out the decoding code from the rest allows us to easy switch between regular sampling for training and greedy sampling when doing inference. We use a flag to set which one to use.\n",
    "\n",
    "Then, the seq2seq code uses `BasicDecoder`, which takes all the inputs the decoding process needs, and `dynamic_decode` to run the decoder rnn."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "MODE = \"train\"\n",
    "\n",
    "# Helper\n",
    "if MODE == \"train\":\n",
    "  helper = tf.contrib.seq2seq.TrainingHelper(\n",
    "    inputs=decoder_emb_inp,\n",
    "    sequence_length=decoder_lengths,\n",
    "    time_major=True)\n",
    "elif MODE == \"infer\":\n",
    "  helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(\n",
    "      embedding=lambda params: tf.nn.embedding_lookup(embedding_matrix, params),\n",
    "      start_tokens=tf.tile([GO_SYMBOL], [batch_size]),\n",
    "      end_token=END_SYMBOL,\n",
    "      time_major=True)\n",
    "\n",
    "# Build RNN cell and projection layer\n",
    "decoder_cell = tf.nn.rnn_cell.BasicLSTMCell(NUM_UNITS)\n",
    "projection_layer = tf.layers.Dense(DST_VOCAB_SIZE, use_bias=False, name=\"output_projection\")\n",
    "\n",
    "# Decoder\n",
    "decoder = tf.contrib.seq2seq.BasicDecoder(\n",
    "    decoder_cell, helper, encoder_state,\n",
    "    output_layer=projection_layer)\n",
    "# Dynamic decoding\n",
    "(final_outputs, _, _) = tf.contrib.seq2seq.dynamic_decode(decoder, output_time_major=True)\n",
    "logits = final_outputs.rnn_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the important steps in training RNNs is gradient clipping. Here, we clip by the global norm which clips the values of multiple tensors by the ratio of the sum of their norms. The max value, `max_gradient_norm`, is often set to a value like 5 or 1.\n",
    "\n",
    "Google Brain uses standard SGD (tf.train.GradientDescentOptimizer) with a decreasing learning rate schedule, which yields better performance, but that takes a lot of tuning to work, so we'll stick with the `RMSProp` optimizer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "MAX_GRADIENT_NORM = 1.0\n",
    "LEARNING_RATE = 1e-4\n",
    "\n",
    "decoder_outputs = tf.placeholder(tf.int32, shape=[None, None])\n",
    "\n",
    "dynamic_time_steps = tf.shape(logits)[0]\n",
    "dynamic_batch_size = tf.shape(logits)[1]\n",
    "train_loss = tf.contrib.seq2seq.sequence_loss(\n",
    "    logits,\n",
    "    decoder_outputs,\n",
    "    tf.ones([dynamic_batch_size, dynamic_time_steps], dtype=tf.float32),\n",
    "    average_across_timesteps=True,\n",
    "    average_across_batch=True)\n",
    "\n",
    "# Calculate and clip gradients\n",
    "params = tf.trainable_variables()\n",
    "gradients = tf.gradients(train_loss, params)\n",
    "clipped_gradients, _ = tf.clip_by_global_norm(\n",
    "    gradients, MAX_GRADIENT_NORM)\n",
    "\n",
    "# Optimization\n",
    "optimizer = tf.train.RMSPropOptimizer(LEARNING_RATE)\n",
    "update_step = optimizer.apply_gradients(\n",
    "    zip(clipped_gradients, params))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we use the `get_sequences` functions we declared so long ago to get a minibatch from the corpus and run the update step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 10\n",
    "\n",
    "session = tf.Session()\n",
    "session.run(tf.global_variables_initializer())\n",
    "\n",
    "src_sequences, src_lens, dst_inputs, dst_outputs, dst_lens = get_sequences(src_data, dst_data, encode_src, encode_dst, BATCH_SIZE, GO_SYMBOL, END_SYMBOL)\n",
    "feed_dict = {encoder_inputs: src_sequences, source_sequence_length: src_lens,\n",
    "             decoder_inputs: dst_inputs, decoder_lengths: dst_lens,\n",
    "             decoder_outputs: dst_outputs}\n",
    "\n",
    "# we can run the train op as usual\n",
    "_ = session.run(update_step, feed_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## Attention in NMT\n",
    "\n",
    "Now we'll look at the details of an attention system described in ([Luong et al., 2015](https://arxiv.org/abs/1508.04025)) which is commonly used in NMT systems.\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_mechanism.jpg\">\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_vis.jpg\" width=75%>\n",
    "\n",
    "1. The current target hidden state is compared with all source states to derive **attention weights** (can be visualized as in the figure immediately above).\n",
    "2. Based on the attention weights we compute a **context vector** as the weighted average of the source states.\n",
    "3. Combine the context vector with the current target hidden state to yield the **final attention vector**\n",
    "4. The attention vector is fed as an input to the next time step (input feeding). The first three steps can be summarized by the equations below:\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_0.jpg\" width=80%>\n",
    "\n",
    "Once computed, the attention vector $a$ is used to derive the softmax logit and loss. This is similar to the target hidden state at the top layer of a vanilla seq2seq model. The score function and the function $f$ can take other forms:\n",
    "\n",
    "<img src=\"https://www.tensorflow.org/images/seq2seq/attention_equation_1.jpg\" width=80%>\n",
    "\n",
    "Practically, the code for using this attention mechanism in the decoder is below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# attention_states: [batch_size, max_time, num_units], transposing because we used time major above\n",
    "attention_states = tf.transpose(encoder_outputs, [1, 0, 2])\n",
    "\n",
    "# Create an attention mechanism\n",
    "attention_mechanism = tf.contrib.seq2seq.LuongAttention(\n",
    "    NUM_UNITS, attention_states,\n",
    "    memory_sequence_length=source_sequence_length)\n",
    "\n",
    "# Wrap the decoder cell\n",
    "decoder_cell = tf.contrib.seq2seq.AttentionWrapper(\n",
    "    decoder_cell, attention_mechanism,\n",
    "    attention_layer_size=NUM_UNITS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# No exercise this week!\n",
    "\n",
    "Instead, finish strong on homework 3 and make some progress on your projects. The next hackathon will be the first project help session. You should show up to work with your group, ask questions, and get my help with any roadblocks you've encountered.\n",
    "\n",
    "Have a good spring break!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "TensorFlow 1.4 (py36)",
   "language": "python",
   "name": "tensorflow-1.4.0-py36"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
