{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DCN\n",
    "\n",
    "1. Sentinel vector $d_{\\emptyset}$, is not included\n",
    "2. The graph isn't tested with real data\n",
    "3. Dropout wrapper and Multi-cell are not added\n",
    "4. I may have fucked up the affinity matrix calculation - need your input on it\n",
    "5. The code is extremely non-pythonic - calculated individual equations in DCN paper in a single line\n",
    "6. **Decoder** is not added yet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dummy Params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab_size = 100\n",
    "dlen = 30 # m\n",
    "qlen = 15 # n\n",
    "B = 8 # batch size\n",
    "hdim = 32 # l\n",
    "emb_dim = 32"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def shapes(l):\n",
    "    if type(l) == type((23, 69)) or type(l) == type([]):\n",
    "        return [ shapes(t) for t in l ]\n",
    "    elif type(l) == np.ndarray:\n",
    "        return l.shape\n",
    "    \n",
    "def execute_graph(t, feed_dict=None):\n",
    "    feed_dict = {\n",
    "        _document : np.random.randint(0, vocab_size, [B, dlen]),\n",
    "        _query    : np.random.randint(0, vocab_size, [B, qlen]),\n",
    "        _answer   : np.random.randint(0, dlen, [B, ]),\n",
    "    } if type(feed_dict) == type(None) else feed_dict\n",
    "    with tf.Session() as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        return sess.run(t, feed_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "_document = tf.placeholder(tf.int32, [None, None], 'document')\n",
    "_query    = tf.placeholder(tf.int32, [None, None], 'query')\n",
    "_answer   = tf.placeholder(tf.int32, [None, ], 'answer')\n",
    "batch_size_, dlen_ = tf.unstack(tf.shape(_document))\n",
    "qlen_ = tf.shape(_query)[1]\n",
    "E = tf.get_variable('embedding', [vocab_size, emb_dim], tf.float32, \n",
    "                    tf.random_uniform_initializer(-0.01, 0.01))\n",
    "enc_cell = tf.nn.rnn_cell.LSTMCell(hdim)\n",
    "with tf.variable_scope('encoder') as scope:\n",
    "    document, _ = tf.nn.dynamic_rnn(enc_cell, \n",
    "                      tf.nn.embedding_lookup(E, _document),\n",
    "                      tf.count_nonzero(_document, axis=1), # sequence lengths\n",
    "                      dtype=tf.float32)\n",
    "    scope.reuse_variables()\n",
    "    query, _ = tf.nn.dynamic_rnn(enc_cell, \n",
    "                  tf.nn.embedding_lookup(E, _query),\n",
    "                  tf.count_nonzero(_query, axis=1), # sequence lengths\n",
    "                  dtype=tf.float32)\n",
    "with tf.variable_scope('query_projection'):\n",
    "    query = tf.contrib.layers.fully_connected(query, hdim, \n",
    "                                                  activation_fn=tf.nn.tanh)\n",
    "with tf.variable_scope('affinity_matrix'):\n",
    "    affinity = tf.matmul(document, #[(B,Ld,d]x[B,d,Lq]=[B,Ld,Lq]\n",
    "                         tf.transpose(query, [0, 2, 1]))\n",
    "with tf.variable_scope('attention_weights'):\n",
    "    Ad = tf.nn.softmax(affinity) # normalize along Lq\n",
    "    Aq = tf.nn.softmax(tf.transpose(affinity, [0, 2, 1])) # normalize along Ld\n",
    "with tf.variable_scope('summary'):\n",
    "    Cq = tf.matmul(Aq, document) # [B,Lq,Ld] x [B,Ld,d] = [B,Lq,d]\n",
    "    Cd = tf.transpose(tf.matmul( # [B,Lq,d*2] x [B,Ld,Lq] = [B,Ld,2*d]\n",
    "        tf.transpose(tf.concat([query, Cq], axis=-1), [0, 2, 1]),\n",
    "        tf.transpose(Ad, [0, 2, 1])\n",
    "    ), [0, 2, 1])\n",
    "\n",
    "with tf.variable_scope('temporal_fusion'):\n",
    "    context_states, _ = tf.nn.bidirectional_dynamic_rnn(\n",
    "        tf.nn.rnn_cell.LSTMCell(hdim),\n",
    "        tf.nn.rnn_cell.LSTMCell(hdim),\n",
    "        tf.concat([document, Cd], axis=-1), # [B, Ld, 2*d + d]\n",
    "        tf.count_nonzero(_document, axis=1), # sequence lengths\n",
    "        dtype=tf.float32)\n",
    "    context = tf.concat(context_states, axis=-1) # final representation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Execute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(8, 30, 64)]"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "shapes(execute_graph([context]))"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}