{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.sasrec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SASRec\n",
    "> Self-Attentive Sequential Recommendation Model.\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",
    "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 *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "from torch import nn\n",
    "import math\n",
    "\n",
    "from recohut.models.layers.attention import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class SASEmbedding(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        vocab_size = args.num_items + 1\n",
    "        hidden = args.bert_hidden_units\n",
    "        max_len = args.bert_max_len\n",
    "        dropout = args.bert_dropout\n",
    "\n",
    "        self.token = TokenEmbedding(\n",
    "            vocab_size=vocab_size, embed_size=hidden)\n",
    "        self.position = PositionalEmbedding(\n",
    "            max_len=max_len, d_model=hidden)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def get_mask(self, x):\n",
    "        if len(x.shape) > 2:\n",
    "            x = torch.ones(x.shape[:2]).to(x.device)\n",
    "        return (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        mask = self.get_mask(x)\n",
    "        if len(x.shape) > 2:\n",
    "            pos = self.position(torch.ones(x.shape[:2]).to(x.device))\n",
    "            x = torch.matmul(x, self.token.weight) + pos\n",
    "        else:\n",
    "            x = self.token(x) + self.position(x)\n",
    "        return self.dropout(x), mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class SASModel(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        hidden = args.bert_hidden_units\n",
    "        heads = args.bert_num_heads\n",
    "        head_size = args.bert_head_size\n",
    "        dropout = args.bert_dropout\n",
    "        attn_dropout = args.bert_attn_dropout\n",
    "        layers = args.bert_num_blocks\n",
    "\n",
    "        self.transformer_blocks = nn.ModuleList([SASTransformerBlock(\n",
    "            hidden, heads, head_size, hidden * 4, dropout, attn_dropout) for _ in range(layers)])\n",
    "\n",
    "    def forward(self, x, embedding_weight, mask):\n",
    "        for transformer in self.transformer_blocks:\n",
    "            x = transformer.forward(x, mask)\n",
    "        scores = torch.matmul(x, embedding_weight.permute(1, 0))\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SASRec(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        self.args = args\n",
    "        self.embedding = SASEmbedding(self.args)\n",
    "        self.model = SASModel(self.args)\n",
    "        self.truncated_normal_init()\n",
    "\n",
    "    def truncated_normal_init(self, mean=0, std=0.02, lower=-0.04, upper=0.04):\n",
    "        with torch.no_grad():\n",
    "            l = (1. + math.erf(((lower - mean) / std) / math.sqrt(2.))) / 2.\n",
    "            u = (1. + math.erf(((upper - mean) / std) / math.sqrt(2.))) / 2.\n",
    "\n",
    "            for n, p in self.model.named_parameters():\n",
    "                if not 'layer_norm' in n:\n",
    "                    p.uniform_(2 * l - 1, 2 * u - 1)\n",
    "                    p.erfinv_()\n",
    "                    p.mul_(std * math.sqrt(2.))\n",
    "                    p.add_(mean)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x, mask = self.embedding(x)\n",
    "        scores = self.model(x, self.embedding.token.weight, mask)\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Module.parameters of SASRec(\n",
       "  (embedding): SASEmbedding(\n",
       "    (token): TokenEmbedding(11, 4, padding_idx=0)\n",
       "    (position): PositionalEmbedding(\n",
       "      (pe): Embedding(9, 4)\n",
       "    )\n",
       "    (dropout): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       "  (model): SASModel(\n",
       "    (transformer_blocks): ModuleList(\n",
       "      (0): SASTransformerBlock(\n",
       "        (layer_norm): LayerNorm()\n",
       "        (attention): SASMultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "        (feed_forward): SASPositionwiseFeedForward(\n",
       "          (conv1): Conv1d(4, 16, kernel_size=(1,), stride=(1,))\n",
       "          (activation): ReLU()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (conv2): Conv1d(16, 4, kernel_size=(1,), stride=(1,))\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "      )\n",
       "      (1): SASTransformerBlock(\n",
       "        (layer_norm): LayerNorm()\n",
       "        (attention): SASMultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "        (feed_forward): SASPositionwiseFeedForward(\n",
       "          (conv1): Conv1d(4, 16, kernel_size=(1,), stride=(1,))\n",
       "          (activation): ReLU()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (conv2): Conv1d(16, 4, kernel_size=(1,), stride=(1,))\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "      )\n",
       "      (2): SASTransformerBlock(\n",
       "        (layer_norm): LayerNorm()\n",
       "        (attention): SASMultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "        (feed_forward): SASPositionwiseFeedForward(\n",
       "          (conv1): Conv1d(4, 16, kernel_size=(1,), stride=(1,))\n",
       "          (activation): ReLU()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (conv2): Conv1d(16, 4, kernel_size=(1,), stride=(1,))\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "      )\n",
       "      (3): SASTransformerBlock(\n",
       "        (layer_norm): LayerNorm()\n",
       "        (attention): SASMultiHeadedAttention(\n",
       "          (linear_layers): ModuleList(\n",
       "            (0): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (1): Linear(in_features=4, out_features=8, bias=True)\n",
       "            (2): Linear(in_features=4, out_features=8, bias=True)\n",
       "          )\n",
       "          (attention): Attention()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "        (feed_forward): SASPositionwiseFeedForward(\n",
       "          (conv1): Conv1d(4, 16, kernel_size=(1,), stride=(1,))\n",
       "          (activation): ReLU()\n",
       "          (dropout): Dropout(p=0.2, inplace=False)\n",
       "          (conv2): Conv1d(16, 4, kernel_size=(1,), stride=(1,))\n",
       "          (layer_norm): LayerNorm()\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Args:\n",
    "    bert_hidden_units = 4\n",
    "    bert_num_heads = 2\n",
    "    bert_head_size = 4\n",
    "    bert_dropout = 0.2\n",
    "    bert_attn_dropout = 0.2\n",
    "    bert_num_blocks = 4\n",
    "    num_items = 10\n",
    "    bert_hidden_units = 4\n",
    "    bert_max_len = 8\n",
    "    bert_dropout = 0.2\n",
    "\n",
    "args = Args()\n",
    "model = SASRec(args)\n",
    "model.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://cseweb.ucsd.edu/~jmcauley/pdfs/icdm18.pdf\n",
    "2. https://github.com/Yueeeeeeee/RecSys-Extraction-Attack/blob/main/model/sasrec.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-31 06:51:33\n",
      "\n",
      "recohut: 0.0.8\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "PIL       : 7.1.2\n",
      "torch     : 1.10.0+cu111\n",
      "matplotlib: 3.2.2\n",
      "numpy     : 1.19.5\n",
      "IPython   : 5.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p recohut"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
