{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.tf.sasrec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SASRec\n",
    "> Self-Attentive Sequential Recommendation (SASRec)\n",
    "\n",
    "Sequential dynamics are a key feature of many modern recommender systems, which seek to capture the ‘context’ of users’ activities on the basis of actions they have performed recently. To capture such patterns, two approaches have proliferated: Markov Chains (MCs) and Recurrent Neural Networks (RNNs). Markov Chains assume that a user’s next action can be predicted on the basis of just their last (or last few) actions, while RNNs in principle allow for longer-term semantics to be uncovered. Generally speaking, MC-based methods perform best in extremely sparse datasets, where model parsimony is critical, while RNNs perform better in denser datasets where higher model complexity is affordable. SASRec captures the long-term semantics (like an RNN), but, using an attention mechanism, makes its predictions based on relatively few actions (like an MC).\n",
    "\n",
    "![US512148 _ General Recommenders-L186674 _ SASRec Model.drawio.png](https://github.com/RecoHut-Stanzas/S021355/raw/main/images/sasrec.png)\n",
    "\n",
    "At each time step, SASRec seeks to identify which items are ‘relevant’ from a user’s action history, and use them to predict the next item. Extensive empirical studies show that this method outperforms various state-of-the-art sequential models (including MC/CNN/RNN-based approaches) on both sparse and dense datasets. Moreover, the model is an order of magnitude more efficient than comparable CNN/RNN-based models.\n",
    "\n",
    "We adopt the binary cross entropy loss as the objective function:\n",
    "\n",
    "$$-\\sum_{S^u\\in S} \\sum_{t \\in [1,2,\\dots,n]}\\left[ log(\\sigma(r_{o_t,t})) + \\sum_{j \\notin S^u} log(1-\\sigma(r_{j,t})) \\right]$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import os\n",
    "import numpy as np\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import Model\n",
    "from tensorflow.keras.regularizers import l2\n",
    "from tensorflow.keras.losses import Loss\n",
    "from tensorflow.keras.layers import Layer, Dense, LayerNormalization\n",
    "from tensorflow.keras.layers import Dropout, Embedding, Input, Conv1D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "def get_angles(pos, i, d_model):\n",
    "    angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))\n",
    "    return pos * angle_rates\n",
    "\n",
    "\n",
    "def positional_encoding(seq_inputs, embed_dim):\n",
    "    angle_rads = get_angles(np.arange(seq_inputs.shape[-1])[:, np.newaxis],\n",
    "                            np.arange(embed_dim)[np.newaxis, :], embed_dim)\n",
    "    angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
    "    angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
    "\n",
    "    pos_encoding = angle_rads[np.newaxis, ...]\n",
    "\n",
    "    return tf.cast(pos_encoding, dtype=tf.float32)\n",
    "\n",
    "\n",
    "def scaled_dot_product_attention(q, k, v, mask, causality=True):\n",
    "    \"\"\"\n",
    "    Attention Mechanism\n",
    "    :param q: A 3d tensor with shape of (None, seq_len, depth), depth = d_model // num_heads\n",
    "    :param k: A 3d tensor with shape of (None, seq_len, depth)\n",
    "    :param v: A 3d tensor with shape of (None, seq_len, depth)\n",
    "    :param mask:\n",
    "    :param causality: Boolean. If True, using causality, default True\n",
    "    :return:\n",
    "    \"\"\"\n",
    "    mat_qk = tf.matmul(q, k, transpose_b=True)  # (None, seq_len, seq_len)\n",
    "    dk = tf.cast(k.shape[-1], dtype=tf.float32)\n",
    "    # Scaled\n",
    "    scaled_att_logits = mat_qk / tf.sqrt(dk)\n",
    "\n",
    "    paddings = tf.ones_like(scaled_att_logits) * (-2 ** 32 + 1)\n",
    "    outputs = tf.where(tf.equal(mask, 0), paddings, scaled_att_logits)  # (None, seq_len, seq_len)\n",
    "\n",
    "    # Causality\n",
    "    if causality:\n",
    "        diag_vals = tf.ones_like(outputs)  # (None, seq_len, seq_len)\n",
    "        masks = tf.linalg.LinearOperatorLowerTriangular(diag_vals).to_dense()  # (None, seq_len, seq_len)\n",
    "\n",
    "        paddings = tf.ones_like(masks) * (-2 ** 32 + 1)\n",
    "        outputs = tf.where(tf.equal(masks, 0), paddings, outputs)  # (None, seq_len, seq_len)\n",
    "\n",
    "    # softmax\n",
    "    outputs = tf.nn.softmax(logits=outputs)  # (None, seq_len, seq_len)\n",
    "    outputs = tf.matmul(outputs, v)  # (None, seq_len, depth)\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class MultiHeadAttention(Layer):\n",
    "    def __init__(self, d_model, num_heads, causality=True):\n",
    "        \"\"\"\n",
    "        Multi Head Attention Mechanism\n",
    "        :param d_model: A scalar. The self-attention hidden size.\n",
    "        :param num_heads: A scalar. Number of heads. If num_heads == 1, the layer is a single self-attention layer.\n",
    "        :param causality: Boolean. If True, using causality, default True\n",
    "        \"\"\"\n",
    "        super(MultiHeadAttention, self).__init__()\n",
    "        self.d_model = d_model\n",
    "        self.num_heads = num_heads\n",
    "        self.causality = causality\n",
    "\n",
    "        self.wq = Dense(d_model, activation=None)\n",
    "        self.wk = Dense(d_model, activation=None)\n",
    "        self.wv = Dense(d_model, activation=None)\n",
    "\n",
    "    def call(self, q, k, v, mask):\n",
    "        q = self.wq(q)  # (None, seq_len, d_model)\n",
    "        k = self.wk(k)  # (None, seq_len, d_model)\n",
    "        v = self.wv(v)  # (None, seq_len, d_model)\n",
    "\n",
    "        # split d_model into num_heads * depth, and concatenate\n",
    "        q = tf.reshape(tf.concat([tf.split(q, self.num_heads, axis=2)], axis=0),\n",
    "                       (-1, q.shape[1], q.shape[2] // self.num_heads))  # (None * num_heads, seq_len, d_model // num_heads)\n",
    "        k = tf.reshape(tf.concat([tf.split(k, self.num_heads, axis=2)], axis=0),\n",
    "                       (-1, k.shape[1], k.shape[2] // self.num_heads))  # (None * num_heads, seq_len, d_model // num_heads)\n",
    "        v = tf.reshape(tf.concat([tf.split(v, self.num_heads, axis=2)], axis=0),\n",
    "                       (-1, v.shape[1], v.shape[2] // self.num_heads))  # (None * num_heads, seq_len, d_model // num_heads)\n",
    "\n",
    "        # attention\n",
    "        scaled_attention = scaled_dot_product_attention(q, k, v, mask, self.causality)  # (None * num_heads, seq_len, d_model // num_heads)\n",
    "\n",
    "        # Reshape\n",
    "        outputs = tf.concat(tf.split(scaled_attention, self.num_heads, axis=0), axis=2)  # (N, seq_len, d_model)\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FFN(Layer):\n",
    "    def __init__(self, hidden_unit, d_model):\n",
    "        \"\"\"\n",
    "        Feed Forward Network\n",
    "        :param hidden_unit: A scalar. W1\n",
    "        :param d_model: A scalar. W2\n",
    "        \"\"\"\n",
    "        super(FFN, self).__init__()\n",
    "        self.conv1 = Conv1D(filters=hidden_unit, kernel_size=1, activation='relu', use_bias=True)\n",
    "        self.conv2 = Conv1D(filters=d_model, kernel_size=1, activation=None, use_bias=True)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x = self.conv1(inputs)\n",
    "        output = self.conv2(x)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class EncoderLayer(Layer):\n",
    "    def __init__(self, d_model, num_heads=1, ffn_hidden_unit=128, dropout=0., norm_training=True, causality=True):\n",
    "        \"\"\"\n",
    "        Encoder Layer\n",
    "        :param d_model: A scalar. The self-attention hidden size.\n",
    "        :param num_heads: A scalar. Number of heads.\n",
    "        :param ffn_hidden_unit: A scalar. Number of hidden unit in FFN\n",
    "        :param dropout: A scalar. Number of dropout.\n",
    "        :param norm_training: Boolean. If True, using layer normalization, default True\n",
    "        :param causality: Boolean. If True, using causality, default True\n",
    "        \"\"\"\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.mha = MultiHeadAttention(d_model, num_heads, causality)\n",
    "        self.ffn = FFN(ffn_hidden_unit, d_model)\n",
    "\n",
    "        self.layernorm1 = LayerNormalization(epsilon=1e-6, trainable=norm_training)\n",
    "        self.layernorm2 = LayerNormalization(epsilon=1e-6, trainable=norm_training)\n",
    "\n",
    "        self.dropout1 = Dropout(dropout)\n",
    "        self.dropout2 = Dropout(dropout)\n",
    "\n",
    "    def call(self, inputs):\n",
    "        x, mask = inputs\n",
    "        # self-attention\n",
    "        att_out = self.mha(x, x, x, mask)  # （None, seq_len, d_model)\n",
    "        att_out = self.dropout1(att_out)\n",
    "        # residual add\n",
    "        out1 = self.layernorm1(x + att_out)\n",
    "        # ffn\n",
    "        ffn_out = self.ffn(out1)\n",
    "        ffn_out = self.dropout2(ffn_out)\n",
    "        # residual add\n",
    "        out2 = self.layernorm2(out1 + ffn_out)  # (None, seq_len, d_model)\n",
    "        return out2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SASRec(tf.keras.Model):\n",
    "    def __init__(self, item_fea_col, blocks=1, num_heads=1, ffn_hidden_unit=128,\n",
    "                 dropout=0., maxlen=40, norm_training=True, causality=False, embed_reg=1e-6):\n",
    "        \"\"\"\n",
    "        SASRec model\n",
    "        :param item_fea_col: A dict contains 'feat_name', 'feat_num' and 'embed_dim'.\n",
    "        :param blocks: A scalar. The Number of blocks.\n",
    "        :param num_heads: A scalar. Number of heads.\n",
    "        :param ffn_hidden_unit: A scalar. Number of hidden unit in FFN\n",
    "        :param dropout: A scalar. Number of dropout.\n",
    "        :param maxlen: A scalar. Number of length of sequence\n",
    "        :param norm_training: Boolean. If True, using layer normalization, default True\n",
    "        :param causality: Boolean. If True, using causality, default True\n",
    "        :param embed_reg: A scalar. The regularizer of embedding\n",
    "        \"\"\"\n",
    "        super(SASRec, self).__init__()\n",
    "        # sequence length\n",
    "        self.maxlen = maxlen\n",
    "        # item feature columns\n",
    "        self.item_fea_col = item_fea_col\n",
    "        # embed_dim\n",
    "        self.embed_dim = self.item_fea_col['embed_dim']\n",
    "        # d_model must be the same as embedding_dim, because of residual connection\n",
    "        self.d_model = self.embed_dim\n",
    "        # item embedding\n",
    "        self.item_embedding = Embedding(input_dim=self.item_fea_col['feat_num'],\n",
    "                                        input_length=1,\n",
    "                                        output_dim=self.item_fea_col['embed_dim'],\n",
    "                                        mask_zero=True,\n",
    "                                        embeddings_initializer='random_uniform',\n",
    "                                        embeddings_regularizer=l2(embed_reg))\n",
    "        self.pos_embedding = Embedding(input_dim=self.maxlen,\n",
    "                                       input_length=1,\n",
    "                                       output_dim=self.embed_dim,\n",
    "                                       mask_zero=False,\n",
    "                                       embeddings_initializer='random_uniform',\n",
    "                                       embeddings_regularizer=l2(embed_reg))\n",
    "        self.dropout = Dropout(dropout)\n",
    "        # attention block\n",
    "        self.encoder_layer = [EncoderLayer(self.d_model, num_heads, ffn_hidden_unit,\n",
    "                                           dropout, norm_training, causality) for b in range(blocks)]\n",
    "\n",
    "    def call(self, inputs, training=None):\n",
    "        # inputs\n",
    "        seq_inputs, pos_inputs, neg_inputs = inputs  # (None, maxlen), (None, 1), (None, 1)\n",
    "        # mask\n",
    "        mask = tf.expand_dims(tf.cast(tf.not_equal(seq_inputs, 0), dtype=tf.float32), axis=-1)  # (None, maxlen, 1)\n",
    "        # seq info\n",
    "        seq_embed = self.item_embedding(seq_inputs)  # (None, maxlen, dim)\n",
    "        # pos encoding\n",
    "        # pos_encoding = positional_encoding(seq_inputs, self.embed_dim)\n",
    "        pos_encoding = tf.expand_dims(self.pos_embedding(tf.range(self.maxlen)), axis=0)\n",
    "        seq_embed += pos_encoding\n",
    "        seq_embed = self.dropout(seq_embed)\n",
    "        att_outputs = seq_embed  # (None, maxlen, dim)\n",
    "        att_outputs *= mask\n",
    "\n",
    "        # self-attention\n",
    "        for block in self.encoder_layer:\n",
    "            att_outputs = block([att_outputs, mask])  # (None, seq_len, dim)\n",
    "            att_outputs *= mask\n",
    "\n",
    "        # user_info = tf.reduce_mean(att_outputs, axis=1)  # (None, dim)\n",
    "        user_info = tf.expand_dims(att_outputs[:, -1], axis=1)  # (None, 1, dim)\n",
    "        # item info\n",
    "        pos_info = self.item_embedding(pos_inputs)  # (None, 1, dim)\n",
    "        neg_info = self.item_embedding(neg_inputs)  # (None, 1/100, dim)\n",
    "        pos_logits = tf.reduce_sum(user_info * pos_info, axis=-1)  # (None, 1)\n",
    "        neg_logits = tf.reduce_sum(user_info * neg_info, axis=-1)  # (None, 1)\n",
    "        # loss\n",
    "        losses = tf.reduce_mean(- tf.math.log(tf.nn.sigmoid(pos_logits)) -\n",
    "                                tf.math.log(1 - tf.nn.sigmoid(neg_logits))) / 2\n",
    "        self.add_loss(losses)\n",
    "        logits = tf.concat([pos_logits, neg_logits], axis=-1)\n",
    "        return logits\n",
    "\n",
    "    def summary(self):\n",
    "        seq_inputs = Input(shape=(self.maxlen,), dtype=tf.int32)\n",
    "        pos_inputs = Input(shape=(1,), dtype=tf.int32)\n",
    "        neg_inputs = Input(shape=(1,), dtype=tf.int32)\n",
    "        tf.keras.Model(inputs=[seq_inputs, pos_inputs, neg_inputs],\n",
    "                       outputs=self.call([seq_inputs, pos_inputs, neg_inputs])).summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model():\n",
    "    item_fea_col = {'feat': 'item_id', 'feat_num': 100, 'embed_dim': 8}\n",
    "    model = SASRec(item_fea_col, num_heads=8)\n",
    "    model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "__________________________________________________________________________________________________\n",
      " Layer (type)                   Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      " input_1 (InputLayer)           [(None, 40)]         0           []                               \n",
      "                                                                                                  \n",
      " embedding (Embedding)          multiple             800         ['input_1[0][0]',                \n",
      "                                                                  'input_2[0][0]',                \n",
      "                                                                  'input_3[0][0]']                \n",
      "                                                                                                  \n",
      " tf.math.not_equal (TFOpLambda)  (None, 40)          0           ['input_1[0][0]']                \n",
      "                                                                                                  \n",
      " tf.__operators__.add (TFOpLamb  (None, 40, 8)       0           ['embedding[0][0]']              \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.cast (TFOpLambda)           (None, 40)           0           ['tf.math.not_equal[0][0]']      \n",
      "                                                                                                  \n",
      " dropout (Dropout)              (None, 40, 8)        0           ['tf.__operators__.add[0][0]']   \n",
      "                                                                                                  \n",
      " tf.expand_dims (TFOpLambda)    (None, 40, 1)        0           ['tf.cast[0][0]']                \n",
      "                                                                                                  \n",
      " tf.math.multiply (TFOpLambda)  (None, 40, 8)        0           ['dropout[0][0]',                \n",
      "                                                                  'tf.expand_dims[0][0]']         \n",
      "                                                                                                  \n",
      " encoder_layer (EncoderLayer)   (None, 40, 8)        2432        ['tf.math.multiply[0][0]',       \n",
      "                                                                  'tf.expand_dims[0][0]']         \n",
      "                                                                                                  \n",
      " tf.math.multiply_1 (TFOpLambda  (None, 40, 8)       0           ['encoder_layer[0][0]',          \n",
      " )                                                                'tf.expand_dims[0][0]']         \n",
      "                                                                                                  \n",
      " tf.__operators__.getitem (Slic  (None, 8)           0           ['tf.math.multiply_1[0][0]']     \n",
      " ingOpLambda)                                                                                     \n",
      "                                                                                                  \n",
      " input_2 (InputLayer)           [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " input_3 (InputLayer)           [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " tf.expand_dims_1 (TFOpLambda)  (None, 1, 8)         0           ['tf.__operators__.getitem[0][0]'\n",
      "                                                                 ]                                \n",
      "                                                                                                  \n",
      " tf.math.multiply_2 (TFOpLambda  (None, 1, 8)        0           ['tf.expand_dims_1[0][0]',       \n",
      " )                                                                'embedding[1][0]']              \n",
      "                                                                                                  \n",
      " tf.math.multiply_3 (TFOpLambda  (None, 1, 8)        0           ['tf.expand_dims_1[0][0]',       \n",
      " )                                                                'embedding[2][0]']              \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum (TFOpLambda  (None, 1)           0           ['tf.math.multiply_2[0][0]']     \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum_1 (TFOpLamb  (None, 1)           0           ['tf.math.multiply_3[0][0]']     \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.concat (TFOpLambda)         (None, 2)            0           ['tf.math.reduce_sum[0][0]',     \n",
      "                                                                  'tf.math.reduce_sum_1[0][0]']   \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 3,232\n",
      "Trainable params: 3,232\n",
      "Non-trainable params: 0\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "test_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-20 09:07:30\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "tensorflow: 2.7.0\n",
      "IPython   : 5.5.0\n",
      "numpy     : 1.19.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "!pip install -q watermark\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
