Self-Attention-Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/rahulchandran/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
      "  from ._conv import register_converters as _register_converters\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "from transformer.hparams import Hparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data_load utils:\n",
    "\n",
    "#1. tokenize using google sentence-piece split vocab \n",
    "def load_vocab(vocab_fpath):\n",
    "    vocab = [line.split()[0] for line in open(vocab_fpath,'r').read().splitlines()]\n",
    "    token2idx = {token:idx for idx,token in  enumerate(vocab)}\n",
    "    idx2token = {idx:token for idx,token in enumerate(vocab)}\n",
    "    return token2idx,idx2token\n",
    "\n",
    "# load length controlled sentences\n",
    "\n",
    "def load_data(fpath1,fpath2,maxlen1,maxlen2):\n",
    "    sents1,sents2 = [],[]\n",
    "    with open(fpath1,'r') as f1, open(fpath2,'r') as f2:\n",
    "        for sent1,sent2 in zip(f1,f2):\n",
    "            if len(sent1.split())+1 > maxlen1:\n",
    "                continue\n",
    "            if eln(sent2.split())+1 > maxlen2:\n",
    "                continue\n",
    "            sents1.append(sent1.strip())\n",
    "            sents2.append(sent2.strip())\n",
    "    return sents1,sents2\n",
    "\n",
    "def calc_num_batches(total_num,batch_size):\n",
    "    return total_num // batch_size + int(total_num % batch_size != 0)\n",
    "\n",
    "def sent2num_encode(inp_sent,type_inp,t2idx_dict):\n",
    "    inp_str = inp_sent.decode(\"utf-8\")\n",
    "    if type_inp == \"x\":\n",
    "        tokens = inp_str.split() + [\"</s>\"]\n",
    "    else:\n",
    "        tokens= [\"<s>\"] + inp_str.split() + [\"</s>\"]\n",
    "    enc_list = [t2idx_dict.get(t,t2idx_dict[\"<unk>\"]) for t in tokens]\n",
    "    return enc_list\n",
    "\n",
    "\n",
    "def gen_fn(sents1,sents2,vocab_fpath):\n",
    "    token2idx, _ = load_vocab(vocab_fpath) \n",
    "    #Now the generator part using yield\n",
    "    for sent1,sent2 in zip(sents1,sents2):\n",
    "        x = sent2num_encode(sents1,\"x\",token2idx)\n",
    "        y = sent2num_encode(sents2,\"y\",token2idx)\n",
    "        \n",
    "        decoder_inp,y = y[:-1], y[1:]\n",
    "        x_seq_len, y_seq_len =  len(x),len(y)\n",
    "        yield (x,x_seq_len,sent1) , (decoder_inp,y,y_seq_len,sent2)\n",
    "\n",
    "\n",
    "def data_feed_fn(sents1,sents2,vocab_fpath,batch_size,shuffle=False):\n",
    "    # xs = (x[int32] , x_seq_len[int32], sents1[string])\n",
    "    # ys = (decoder_inp[int32], y[int32], y_seq_len[int32], sents2[string])\n",
    "    shapes=(([None],(), ()),([None],[None],(),()))  # represents xs, ys\n",
    "    \n",
    "    types = (\n",
    "    (tf.int32,tf.int32,tf.string),\n",
    "    (tf.int32,tf.int32,tf.int32,tf.string)\n",
    "    )\n",
    "    \n",
    "    paddings = (\n",
    "    (0,0,''),\n",
    "    (0,0,0,'')\n",
    "    )\n",
    "    \n",
    "    dataset = tf.data.Dataset.from_generator(\n",
    "        gen_fn,\n",
    "        output_shapes=shapes,\n",
    "        output_types = types,\n",
    "        args = (sents1,sents2,vocab_fpath))\n",
    "    if shuffle:\n",
    "        dataset = dataset.shuffle(128 * batch_size)\n",
    "    \n",
    "    dataset = dataset.repeat() # infinite fetch\n",
    "    dataset  = dataset.padded_batch(batch_size, shapes, paddings).prefetch(1)\n",
    "    return dataset\n",
    "\n",
    "def get_batch(fpath1,fpath2,maxlen1,maxlen2,vocab_fpath,batch_size,shuffle=False):\n",
    "    sents1,sents2 = load_data(fpath1,fpath2,maxlen1,maxlen2)\n",
    "    batches = data_feed_fn(sents1,sents2,vocab_fpath,batch_size,shuffle)\n",
    "    num_batches = calc_num_batches(len(sents1), batch_size)\n",
    "    return batches, num_batches, len(sents1)\n",
    "    \n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Useful module functions..\n",
    "#1. Token_embedding layer\n",
    "def get_token_embeddings(vocab_sz,d_embed,zero_pad=True):\n",
    "    with tf.variable_scope(\"shared_weight_matrix\"):\n",
    "        embeddings = tf.get_variable(\"weight_mat\",\n",
    "                                     dtype = tf.float32,\n",
    "                                     shape= (vocab_sz,d_embed),\n",
    "                                     initializer=tf.contrib.layers.xavier_initializer())\n",
    "        if zero_pad:\n",
    "            embeddings = tf.concat((tf.zeros(shape=[1,d_embed]) ,embeddings[1:,:]),0)\n",
    "    return embeddings\n",
    "\n",
    "\n",
    "#2. positional encoding as given in the paper \n",
    "\n",
    "def positional_encoding(inputs,maxlen,masking=True,scope=\"positional_encoding\"):\n",
    "    # Get embedding_dims into E : this is a fixed value\n",
    "    E = inputs.get_shape().as_list()[-1]\n",
    "    \n",
    "    # Now get the batch_len [here 1] N  and  seq size T <= maxlen .. Both of these are dynamic , so the tf.shape here\n",
    "    N, T = tf.shape(inputs)[0],tf.shape(inputs)[1]\n",
    "    \n",
    "    with tf.variable_scope(scope,reuse=tf.AUTO_REUSE):\n",
    "        #Now get positional indices. This is implemented by creating a range(T) ie. [0..T]  , which can then be \n",
    "        # used to look for index in embedding. Tile it to create a stack of rows, which will be easy for look up\n",
    "        \n",
    "        positional_indices = tf.tile(tf.expand_dims(tf.range(T),0),[N,1]) # Tile N,1 times or [1,1] / 1 times here..\n",
    "        \n",
    "        # And that's super smart code.. B-)  \n",
    "        position_enc = np.array([[pos/ np.power(10000,(i-i%2)/E) for i in range(E)]\n",
    "                            for pos in range(maxlen)])\n",
    "        \n",
    "        #Apply sin to odd positions and cos to even positions in embeddings.. \n",
    "        position[:,::2] = np.sin(position_enc[:,::2]) # 2i\n",
    "        position[:,1::2] = np.cos(position_enc[:,1::2]) # 2i+1\n",
    "        \n",
    "        #Convert to tensor to use it in embedding \n",
    "        position_enc = tf.convert_to_tensor(position_enc, tf.float32) # (maxlen, E)\n",
    "        \n",
    "        \n",
    "        # lookup\n",
    "        outputs = tf.nn.embedding_lookup(position_enc, position_ind)\n",
    "        \n",
    "        #Masking logic : if inputs contains 0 at that position, make it zero or use outputs\n",
    "        \n",
    "        if masking:\n",
    "            outputs = tf.where(tff.equal(inputs,0) , inputs, outputs)\n",
    "        \n",
    "    return tf.to_float(outputs)\n",
    "    \n",
    "    \n",
    "def mask(inputs,queries = None,keys = None,mask_type=None):\n",
    "    outputs = None\n",
    "    # define negative infinity\n",
    "    padding_neg_infinity = - 2**32 + 1\n",
    "    \n",
    "    if mask_type == \"key\": # This is to mask padded value sgetting attention .... w.r.t key \n",
    "        # get paddings part, which typically contains 0's in embeddings part\n",
    "        # MAtch the mask to inputs  which is Q.K having dims (N,T_Q,T_K)\n",
    "        \n",
    "        masks =  tf.sign(tf.reduce_sum(tf.abs(keys),axis = -1)) # (N, T_k)\n",
    "        masks  = tf.expand_dims(masks,axis=1) # (N, 1, T_k)\n",
    "        masks = tf.tile(masks, [1,tf.shape(queries)[1],1]) # Tile to (N,T_q,T_k)\n",
    "        \n",
    "        \n",
    "        # Now apply masks to inputs with negaive infinity \n",
    "        paddings = tf.ones_like(inputs) * padding_neg_infinity\n",
    "        outputs = tf.where(tf.equal(masks,0),paddings, inputs)\n",
    "        \n",
    "    elif mask_type==\"query\": # This is to mask padded value sgetting attention .... w.r.t query \n",
    "        masks = tf.sign(tf,reduce_sum(tf.abs(queries),axis = -1)) #(N, T_q)\n",
    "        masks = tf.expand_dims(masks,axis=-1) # (N, T_q,1) <- Note the change in logic...So awesome ha.. :P\n",
    "        masks  = tf.tile(masks,  [1,1,tf.shape(keys)[1]]) # Tile to (N,T_q,T_k)\n",
    "        outputs = inputs * masks\n",
    "        \n",
    "    elif mask_type==\"future\": # This is to mask the other side [next prediction] getting eyed by the decoder..\n",
    "        #This is done via lower unit triangle\n",
    "        outer_dim_set = tf.ones_like(inputs[0,:,:]) # Get (T_q, T_k) values\n",
    "        low_triangle_diag = tf.linalg.LinearOperatorLowerTriangular(outer_dim_set).to_dense() # (T_q, T_k)\n",
    "        masks = tf.tile(tf.expand_dims(low_triangle_diag,0),[tf.shape(inputs)[0], 1, 1]) # (N,T_q,T_k)\n",
    "        paddings = tf.ones_like(inputs) * padding_neg_infinity\n",
    "        outputs = tf.where(tf.equal(masks, 0), paddings, inputs)\n",
    "    else:\n",
    "        print(\"Check if you entered mask type correctly!\")\n",
    "    return outputs\n",
    "\n",
    "\n",
    "def scaled_dot_product_attention(Q,K,V,causality=False,\n",
    "                                 dropout_rate=0.0,training=True,scope=\"scaled_dot_produt_attention\"):\n",
    "    with tf.variable_scope(scope,reuse=tf.AUTO_REUSE):\n",
    "        # get embedding dims\n",
    "        d_k = Q.get_shape().as_list()[-1]\n",
    "        \n",
    "        # dot product (Q dot K)\n",
    "        # Q is (N,T_q,d_k) and K is (N,T_k,d_k)\n",
    "        outputs = tf.matmul(Q,(tf.transpose(k,[0,2,1])))\n",
    "        \n",
    "        # Now scale with sqrt(d_k)\n",
    "        outputs /=  d_k ** 0.5\n",
    "        \n",
    "        # Now do the key masking to remove attention from padded values w.r.t keys\n",
    "        outputs = mask(outputs, Q,K, mask_type=\"key\")\n",
    "        \n",
    "        # Now do it for decoder right side blind masking case\n",
    "        if causality:\n",
    "            outputs = mask(outputs, mask_type=\"future\")\n",
    "            \n",
    "        # Now find the softmax\n",
    "        outputs = tf.nn.softmax(outputs)\n",
    "        attention  = tf.transpose(outputs,[0,2,1])\n",
    "        tf.summary.image(\"attention\",tf.expand_dims(attention[:1],-1))\n",
    "        \n",
    "        \n",
    "        #query masking\n",
    "        outputs = mask(outputs,Q,K,mask_type=\"query\")\n",
    "        \n",
    "        #Add dropout\n",
    "        outputs = tf.layers.dropout(outputs, rate=dropout_rate, training=training)\n",
    "        \n",
    "        # Now the last remaining output x value part\n",
    "        outputs = tf.matmul(outputs, V) #(N,T_q,d_v)\n",
    "        \n",
    "        \n",
    "    return outputs\n",
    "    \n",
    "def layer_normalize(inputs,epsilon=1e-10, scope=\"layer_normalize\"):\n",
    "    with tf.variable_scope(scope,reuse=tf.AUTO_REUSE):\n",
    "        vec_dims = inputs.get_shape().as_list[-1:] # : is used to return as array with the last dim\n",
    "        mean_vec_axis, variance_vec_axis =  tf.nn.moments(inputs,[-1],keep_dims=True)\n",
    "        beta = tf.get_variable(\"beta\",vec_dims,initializer=tf.zeros_initializer())\n",
    "        gamma = tf.get_variable(\"gamma\", vec_dims, initializer=tf.ones_initializer())\n",
    "        normalized = (inputs-mean_vec_axis) / ((variance_vec_axis+epsilon) ** (0.5))\n",
    "        outputs = gamma * normalized+ beta \n",
    "    return outputs\n",
    "\n",
    "\n",
    "        \n",
    "def multihead_attention(queries,keys,values,num_heads,dropout_rate=0,\n",
    "                        training=True,causality=False,scope=\"multihead_attention\"):\n",
    "    # Take embedding dim\n",
    "    d_model = queries.get_shape().as_list()[-1]\n",
    "    with tf.variable_scope(scope,reuse=tf.AUTO_REUSE):\n",
    "        \n",
    "        #Setting W_q, W_k, W_v\n",
    "        Q = tf.layers.dense(queries,d_model) # N x T x d_model\n",
    "        K = tf.layers.dense(keys,d_model) # N x T x d_model \n",
    "        V = tf.layers.dense(values,d_model) # N x T x d_model\n",
    "        \n",
    "        Q_ = tf.split(Q,num_heads,axis=2)\n",
    "        Q_ = tf.concat(Q_,axis=0) # [(num_heads * N) x T x (d_model/num_heads) ]\n",
    "        \n",
    "        K_ = tf.split(K, num_heads,axis=2)\n",
    "        K_ = tf.concat(K_,axis=0)  # [(num_heads * N) x T x (d_model/num_heads) ]\n",
    "        \n",
    "        V_ = tf.split(V, num_heads,aaxis=2)\n",
    "        V_ = tf.concat(V_, axis=0)\n",
    "        \n",
    "        outputs = scaled_dot_product_attention(Q_, K_, V_, causality, dropout_rate, training)\n",
    "\n",
    "        #Reshape to default\n",
    "        \n",
    "        outputs = tf.concat(tf.split(outputs, num_heads, axis=0), axis=2 ) #(N, T_q, d_model)\n",
    "        \n",
    "        # Residual connection\n",
    "        outputs += queries\n",
    "        \n",
    "        # Now layer normalization : Effective for recurrent nets than batch normalization \n",
    "        # https://arxiv.org/abs/1607.06450\n",
    "        outputs += layer_normalize(outputs)\n",
    "        \n",
    "    return outputs\n",
    "\n",
    "def ff(inputs,num_units,scope=\"block_level_feedforward\"):\n",
    "    \n",
    "    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):\n",
    "        # inner layer\n",
    "        outputs = tf.nn.dense(inputs,num_units[0],activation=tf.nn.relu)\n",
    "\n",
    "        #Outer layer\n",
    "        outputs = tf.nn.dense(outputs,num_units[1])\n",
    "\n",
    "        #Residual connections\n",
    "        outputs += inputs\n",
    "\n",
    "        #Layer normalize\n",
    "        outputs += layer_normalize(outputs)\n",
    "    \n",
    "    return outputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now comes the Transformer class\n",
    "\n",
    "class Transformer:\n",
    "    \n",
    "    def __init__(self,hp):\n",
    "        self.hp = hp\n",
    "        self.token2idx, self.idx2token = load_vocab(hp.vocab)\n",
    "        self.embeddings = get_token_embeddings(self.hp.vocab_size,self.hp.d_model,zero_pad=True)\n",
    "        \n",
    "    \n",
    "    def encode(self,xs, training=True):\n",
    "        with tf.variable_scope(\"encoder\",reuse=tf.AUTO_REUSE):\n",
    "            x, seqlens, sents1 = xs\n",
    "            \n",
    "            # Now the embedding part\n",
    "            enc = tf.nn.embedding_lookup(self.embeddings,x)\n",
    "            # scale the embedding with sqrt of model_dim\n",
    "            enc *= self.hp.d_model ** 0.5\n",
    "            \n",
    "            # Now add positional encoding using the formulas given in the paper.\n",
    "            enc += positional_encoding(enc,self.hp.maxlen1)\n",
    "            \n",
    "            #Add dropout for embedding layer..\n",
    "            enc = tf.layers.dropout(enc, self.hp.dropout_rate, training=training)\n",
    "            \n",
    "            \n",
    "            #Now add the encoder blocks\n",
    "            \n",
    "            for i in range(self.hp.num_blocks):\n",
    "                with tf.variable_scope(\"num_blocks_{}\".format(i),reuse=tf.AUTO_REUSE):\n",
    "                    # Self- attention with residual sum and layer normalization\n",
    "                    enc = multihead_attention(queries=enc,\n",
    "                                              keys=enc,\n",
    "                                              values=enc,\n",
    "                                              num_heads = self.hp.num_heads,\n",
    "                                              dropout_rate=self.hp.dropout_rate,\n",
    "                                              training=training,\n",
    "                                              causality=False\n",
    "                                             )\n",
    "                    \n",
    "                    #ff layer\n",
    "                    enc = ff(enc,[self.hp.d_ff,self.hp.d_model])\n",
    "        memory = enc\n",
    "        return memory,sents1\n",
    "   \n",
    "    def decode(self,ys,memory,training=True):\n",
    "        with tf.variable_scope(\"decoder\",reuse=tf.AUTO_REUSE):\n",
    "            decoder_inp,y, seqlens,sents2 = ys\n",
    "            \n",
    "            #embedding\n",
    "            dec = tf.nn.embedding_lookup(self.embeddings,decoder_inp)\n",
    "            #scale\n",
    "            dec *= self.hp.d_model ** 0.5 \n",
    "            \n",
    "            #positional encoding \n",
    "            dec += positional_encoding(dec,self.hp.maxlen2)\n",
    "            \n",
    "            #Add dropout\n",
    "            dec = tf.layers.dropout(dec,self.hp.dropout_rate,training=training)\n",
    "            \n",
    "            # Now decoder blocks \n",
    "            \n",
    "            for i in range(self.hp.num_blocks):\n",
    "                with tf.variable_scope(\"num_blocks_{}\".format(i), reuse=tf.AUTO_REUSE):\n",
    "                    # 1. masked_attention with residual sum and layer norm\n",
    "                    dec = multihead_attention(\n",
    "                        queries=dec,\n",
    "                        keys=dec,\n",
    "                        values = dec,\n",
    "                        num_heads = self.hp.num_heads,\n",
    "                        training  = training,\n",
    "                        causality = True,\n",
    "                        dropout_rate = self.hp.dropout_rate,\n",
    "                        scope=\"self_attention\"\n",
    "                    )\n",
    "                    \n",
    "                    # 2.  Vanilla normalization  with queries from dec and key and query {K_mem,V_mem} from encoder layer\n",
    "                    \n",
    "                    dec = multihead_attention(\n",
    "                        queries = dec,\n",
    "                        keys = memory,\n",
    "                        values = memory,\n",
    "                        num_heads = self.hp.num_heads,\n",
    "                        training = training,\n",
    "                        causality = False , # no we don't need future masking\n",
    "                        dropout_rate = self.hp.dropout_rate,\n",
    "                        scope=\"vanilla_attention\"\n",
    "                    )\n",
    "                    \n",
    "                    \n",
    "                    dec = ff(dec, num_units = [self.hp.d_ff,self.d_model])\n",
    "        weights = tf.transpose(self.embeddings) # (d_model,vocab_size)\n",
    "        # Love the einstein summation method.. :) \n",
    "        logits  = tf.einsum(\"ntd,dv->ntv\",dec,weights) #(N,T,vocab_sz)\n",
    "        #Predicted output\n",
    "        y_hat = tf.to_int32(tf.argmax(logits, axis=-1))\n",
    "        \n",
    "        return logits,y_hat,y , sents2\n",
    "    \n",
    "    def train():\n",
    "        \n",
    "        pass\n",
    "    def eval():\n",
    "        pass\n",
    "        \n",
    "                    "
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}