{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.tf.attrec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AttRec\n",
    "> Self-Attentive Sequential Recommendation\n",
    "\n",
    "The loss function can either be pointwise or pairwise. Due to the non-convexity of the objective function of NeuMF, gradient-based optimization methods only find locally optimal solutions. It is reported that initialization plays an important role in the convergence and performance of deep learning models. Since NeuMF is an ensemble of GMF and MLP, we usually initialize NeuMF using the pre-trained models of GMF and MLP.\n",
    "\n",
    "![Untitled](https://github.com/RecoHut-Stanzas/S021355/raw/main/images/img3.png)"
   ]
  },
  {
   "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 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.layers import Layer\n",
    "from tensorflow.keras.layers import Embedding, Input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SelfAttention_Layer(Layer):\n",
    "    def __init__(self):\n",
    "        super(SelfAttention_Layer, self).__init__()\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        self.dim = input_shape[0][-1]\n",
    "        self.W = self.add_weight(shape=[self.dim, self.dim], name='weight', \n",
    "            initializer='random_uniform')\n",
    "\n",
    "    def call(self, inputs, **kwargs):\n",
    "        q, k, v, mask = inputs\n",
    "        # pos encoding\n",
    "        k += self.positional_encoding(k)\n",
    "        q += self.positional_encoding(q)\n",
    "        # Nonlinear transformation\n",
    "        q = tf.nn.relu(tf.matmul(q, self.W))  # (None, seq_len, dim)\n",
    "        k = tf.nn.relu(tf.matmul(k, self.W))  # (None, seq_len, dim)\n",
    "        mat_qk = tf.matmul(q, k, transpose_b=True)  # (None, seq_len, seq_len)\n",
    "        dk = tf.cast(self.dim, dtype=tf.float32)\n",
    "        # Scaled\n",
    "        scaled_att_logits = mat_qk / tf.sqrt(dk)\n",
    "        # Mask\n",
    "        mask = tf.tile(tf.expand_dims(mask, 1), [1, q.shape[1], 1])  # (None, seq_len, seq_len)\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",
    "        # softmax\n",
    "        outputs = tf.nn.softmax(logits=outputs, axis=-1)  # (None, seq_len, seq_len)\n",
    "        # output\n",
    "        outputs = tf.matmul(outputs, v)  # (None, seq_len, dim)\n",
    "        outputs = tf.reduce_mean(outputs, axis=1)  # (None, dim)\n",
    "        return outputs\n",
    "\n",
    "    @staticmethod\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",
    "    def positional_encoding(self, QK_input):\n",
    "        angle_rads = self.get_angles(np.arange(QK_input.shape[1])[:, np.newaxis],\n",
    "                                np.arange(self.dim)[np.newaxis, :], self.dim)\n",
    "        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n",
    "        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n",
    "        pos_encoding = angle_rads[np.newaxis, ...]\n",
    "\n",
    "        return tf.cast(pos_encoding, dtype=tf.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AttRec(Model):\n",
    "    def __init__(self, feature_columns, maxlen=40, mode='inner', gamma=0.5, w=0.5, embed_reg=1e-6, **kwargs):\n",
    "        \"\"\"\n",
    "        AttRec\n",
    "        :param feature_columns: A feature columns list. user + seq\n",
    "        :param maxlen: A scalar. In the paper, maxlen is L, the number of latest items.\n",
    "        :param gamma: A scalar. if mode == 'dist', gamma is the margin.\n",
    "        :param mode: A string. inner or dist.\n",
    "        :param w: A scalar. The weight of short interest.\n",
    "        :param embed_reg: A scalar. The regularizer of embedding.\n",
    "        \"\"\"\n",
    "        super(AttRec, self).__init__(**kwargs)\n",
    "        # maxlen\n",
    "        self.maxlen = maxlen\n",
    "        # w\n",
    "        self.w = w\n",
    "        self.gamma = gamma\n",
    "        self.mode = mode\n",
    "        # feature columns\n",
    "        self.user_fea_col, self.item_fea_col = feature_columns\n",
    "        # embed_dim\n",
    "        self.embed_dim = self.item_fea_col['embed_dim']\n",
    "        # user embedding\n",
    "        self.user_embedding = Embedding(input_dim=self.user_fea_col['feat_num'],\n",
    "                                        input_length=1,\n",
    "                                        output_dim=self.user_fea_col['embed_dim'],\n",
    "                                        mask_zero=False,\n",
    "                                        embeddings_initializer='random_normal',\n",
    "                                        embeddings_regularizer=l2(embed_reg))\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_normal',\n",
    "                                        embeddings_regularizer=l2(embed_reg))\n",
    "        # item2 embedding, not share embedding\n",
    "        self.item2_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_normal',\n",
    "                                        embeddings_regularizer=l2(embed_reg))\n",
    "        # self-attention\n",
    "        self.self_attention = SelfAttention_Layer()\n",
    "\n",
    "    def call(self, inputs, **kwargs):\n",
    "        # input\n",
    "        user_inputs, seq_inputs, pos_inputs, neg_inputs = inputs\n",
    "        # mask\n",
    "        # mask = self.item_embedding.compute_mask(seq_inputs)\n",
    "        mask = tf.cast(tf.not_equal(seq_inputs, 0), dtype=tf.float32)  # (None, maxlen)\n",
    "        # user info\n",
    "        user_embed = self.user_embedding(tf.squeeze(user_inputs, axis=-1))  # (None, dim)\n",
    "        # seq info\n",
    "        seq_embed = self.item_embedding(seq_inputs)  # (None, maxlen, dim)\n",
    "        # item\n",
    "        pos_embed = self.item_embedding(tf.squeeze(pos_inputs, axis=-1))  # (None, dim)\n",
    "        neg_embed = self.item_embedding(tf.squeeze(neg_inputs, axis=-1))  # (None, dim)\n",
    "        # item2 embed\n",
    "        pos_embed2 = self.item2_embedding(tf.squeeze(pos_inputs, axis=-1))  # (None, dim)\n",
    "        neg_embed2 = self.item2_embedding(tf.squeeze(neg_inputs, axis=-1))  # (None, dim)\n",
    "\n",
    "        # short-term interest\n",
    "        short_interest = self.self_attention([seq_embed, seq_embed, seq_embed, mask])  # (None, dim)\n",
    "\n",
    "        # mode\n",
    "        if self.mode == 'inner':\n",
    "            # long-term interest, pos and neg\n",
    "            pos_long_interest = tf.multiply(user_embed, pos_embed2)\n",
    "            neg_long_interest = tf.multiply(user_embed, neg_embed2)\n",
    "            # combine\n",
    "            pos_scores = self.w * tf.reduce_sum(pos_long_interest, axis=-1, keepdims=True) \\\n",
    "                         + (1 - self.w) * tf.reduce_sum(tf.multiply(short_interest, pos_embed), axis=-1, keepdims=True)\n",
    "            neg_scores = self.w * tf.reduce_sum(neg_long_interest, axis=-1, keepdims=True) \\\n",
    "                         + (1 - self.w) * tf.reduce_sum(tf.multiply(short_interest, neg_embed), axis=-1, keepdims=True)\n",
    "            self.add_loss(tf.reduce_mean(-tf.math.log(tf.nn.sigmoid(pos_scores - neg_scores))))\n",
    "        else:\n",
    "            # clip by norm\n",
    "            user_embed = tf.clip_by_norm(user_embed, 1, -1)\n",
    "            pos_embed = tf.clip_by_norm(pos_embed, 1, -1)\n",
    "            neg_embed = tf.clip_by_norm(neg_embed, 1, -1)\n",
    "            pos_embed2 = tf.clip_by_norm(pos_embed2, 1, -1)\n",
    "            neg_embed2 = tf.clip_by_norm(neg_embed2, 1, -1)\n",
    "            # distance\n",
    "            # long-term interest, pos and neg\n",
    "            pos_long_interest = tf.square(user_embed - pos_embed2)  # (None, dim)\n",
    "            neg_long_interest = tf.square(user_embed - neg_embed2)  # (None, dim)\n",
    "            # combine. Here is a difference from the original paper.\n",
    "            pos_scores = self.w * tf.reduce_sum(pos_long_interest, axis=-1, keepdims=True) + \\\n",
    "                         (1 - self.w) * tf.reduce_sum(tf.square(short_interest - pos_embed), axis=-1, keepdims=True)\n",
    "            neg_scores = self.w * tf.reduce_sum(neg_long_interest, axis=-1, keepdims=True) + \\\n",
    "                         (1 - self.w) * tf.reduce_sum(tf.square(short_interest - neg_embed), axis=-1, keepdims=True)\n",
    "            # minimize loss\n",
    "            # self.add_loss(tf.reduce_sum(tf.maximum(pos_scores - neg_scores + self.gamma, 0)))\n",
    "            self.add_loss(tf.reduce_sum(tf.nn.relu(pos_scores - neg_scores + self.gamma)))\n",
    "        return pos_scores, neg_scores\n",
    "\n",
    "    def summary(self):\n",
    "        seq_inputs = Input(shape=(self.maxlen,), dtype=tf.int32)\n",
    "        user_inputs = Input(shape=(1, ), dtype=tf.int32)\n",
    "        pos_inputs = Input(shape=(1, ), dtype=tf.int32)\n",
    "        neg_inputs = Input(shape=(1, ), dtype=tf.int32)\n",
    "        Model(inputs=[user_inputs, seq_inputs, pos_inputs, neg_inputs], \n",
    "            outputs=self.call([user_inputs, seq_inputs, pos_inputs, neg_inputs])).summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model():\n",
    "    user_features = {'feat': 'user_id', 'feat_num': 100, 'embed_dim': 8}\n",
    "    seq_features = {'feat': 'item_id', 'feat_num': 100, 'embed_dim': 8}\n",
    "    features = [user_features, seq_features]\n",
    "    model = AttRec(features, mode='dist')\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_2 (InputLayer)           [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " input_3 (InputLayer)           [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " input_1 (InputLayer)           [(None, 40)]         0           []                               \n",
      "                                                                                                  \n",
      " input_4 (InputLayer)           [(None, 1)]          0           []                               \n",
      "                                                                                                  \n",
      " tf.compat.v1.squeeze (TFOpLamb  (None,)             0           ['input_2[0][0]']                \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.compat.v1.squeeze_3 (TFOpLa  (None,)             0           ['input_3[0][0]']                \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.math.not_equal (TFOpLambda)  (None, 40)          0           ['input_1[0][0]']                \n",
      "                                                                                                  \n",
      " tf.compat.v1.squeeze_1 (TFOpLa  (None,)             0           ['input_3[0][0]']                \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.compat.v1.squeeze_4 (TFOpLa  (None,)             0           ['input_4[0][0]']                \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " tf.compat.v1.squeeze_2 (TFOpLa  (None,)             0           ['input_4[0][0]']                \n",
      " mbda)                                                                                            \n",
      "                                                                                                  \n",
      " embedding (Embedding)          (None, 8)            800         ['tf.compat.v1.squeeze[0][0]']   \n",
      "                                                                                                  \n",
      " embedding_2 (Embedding)        (None, 8)            800         ['tf.compat.v1.squeeze_3[0][0]', \n",
      "                                                                  'tf.compat.v1.squeeze_4[0][0]'] \n",
      "                                                                                                  \n",
      " embedding_1 (Embedding)        multiple             800         ['input_1[0][0]',                \n",
      "                                                                  'tf.compat.v1.squeeze_1[0][0]', \n",
      "                                                                  'tf.compat.v1.squeeze_2[0][0]'] \n",
      "                                                                                                  \n",
      " tf.cast (TFOpLambda)           (None, 40)           0           ['tf.math.not_equal[0][0]']      \n",
      "                                                                                                  \n",
      " tf.clip_by_norm (TFOpLambda)   (None, 8)            0           ['embedding[0][0]']              \n",
      "                                                                                                  \n",
      " tf.clip_by_norm_3 (TFOpLambda)  (None, 8)           0           ['embedding_2[0][0]']            \n",
      "                                                                                                  \n",
      " self_attention__layer (SelfAtt  (None, 8)           64          ['embedding_1[0][0]',            \n",
      " ention_Layer)                                                    'embedding_1[0][0]',            \n",
      "                                                                  'embedding_1[0][0]',            \n",
      "                                                                  'tf.cast[0][0]']                \n",
      "                                                                                                  \n",
      " tf.clip_by_norm_1 (TFOpLambda)  (None, 8)           0           ['embedding_1[1][0]']            \n",
      "                                                                                                  \n",
      " tf.clip_by_norm_4 (TFOpLambda)  (None, 8)           0           ['embedding_2[1][0]']            \n",
      "                                                                                                  \n",
      " tf.clip_by_norm_2 (TFOpLambda)  (None, 8)           0           ['embedding_1[2][0]']            \n",
      "                                                                                                  \n",
      " tf.math.subtract (TFOpLambda)  (None, 8)            0           ['tf.clip_by_norm[0][0]',        \n",
      "                                                                  'tf.clip_by_norm_3[0][0]']      \n",
      "                                                                                                  \n",
      " tf.math.subtract_2 (TFOpLambda  (None, 8)           0           ['self_attention__layer[0][0]',  \n",
      " )                                                                'tf.clip_by_norm_1[0][0]']      \n",
      "                                                                                                  \n",
      " tf.math.subtract_1 (TFOpLambda  (None, 8)           0           ['tf.clip_by_norm[0][0]',        \n",
      " )                                                                'tf.clip_by_norm_4[0][0]']      \n",
      "                                                                                                  \n",
      " tf.math.subtract_3 (TFOpLambda  (None, 8)           0           ['self_attention__layer[0][0]',  \n",
      " )                                                                'tf.clip_by_norm_2[0][0]']      \n",
      "                                                                                                  \n",
      " tf.math.square (TFOpLambda)    (None, 8)            0           ['tf.math.subtract[0][0]']       \n",
      "                                                                                                  \n",
      " tf.math.square_2 (TFOpLambda)  (None, 8)            0           ['tf.math.subtract_2[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.square_1 (TFOpLambda)  (None, 8)            0           ['tf.math.subtract_1[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.square_3 (TFOpLambda)  (None, 8)            0           ['tf.math.subtract_3[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum (TFOpLambda  (None, 1)           0           ['tf.math.square[0][0]']         \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum_1 (TFOpLamb  (None, 1)           0           ['tf.math.square_2[0][0]']       \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum_2 (TFOpLamb  (None, 1)           0           ['tf.math.square_1[0][0]']       \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.math.reduce_sum_3 (TFOpLamb  (None, 1)           0           ['tf.math.square_3[0][0]']       \n",
      " da)                                                                                              \n",
      "                                                                                                  \n",
      " tf.math.multiply (TFOpLambda)  (None, 1)            0           ['tf.math.reduce_sum[0][0]']     \n",
      "                                                                                                  \n",
      " tf.math.multiply_1 (TFOpLambda  (None, 1)           0           ['tf.math.reduce_sum_1[0][0]']   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.multiply_2 (TFOpLambda  (None, 1)           0           ['tf.math.reduce_sum_2[0][0]']   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.math.multiply_3 (TFOpLambda  (None, 1)           0           ['tf.math.reduce_sum_3[0][0]']   \n",
      " )                                                                                                \n",
      "                                                                                                  \n",
      " tf.__operators__.add (TFOpLamb  (None, 1)           0           ['tf.math.multiply[0][0]',       \n",
      " da)                                                              'tf.math.multiply_1[0][0]']     \n",
      "                                                                                                  \n",
      " tf.__operators__.add_1 (TFOpLa  (None, 1)           0           ['tf.math.multiply_2[0][0]',     \n",
      " mbda)                                                            'tf.math.multiply_3[0][0]']     \n",
      "                                                                                                  \n",
      "==================================================================================================\n",
      "Total params: 2,464\n",
      "Trainable params: 2,464\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:36:36\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
}
