{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.narm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NARM\n",
    "> Neural Attentive Session-based Recommendation."
   ]
  },
  {
   "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",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class NARMEmbedding(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        vocab_size = args.num_items + 1\n",
    "        embed_size = args.bert_hidden_units\n",
    "        \n",
    "        self.token = nn.Embedding(vocab_size, embed_size)\n",
    "        self.embed_dropout = nn.Dropout(args.bert_dropout)\n",
    "\n",
    "    def get_mask(self, x, lengths):\n",
    "        if len(x.shape) > 2:\n",
    "            return torch.ones(x.shape[:2])[:, :max(lengths)].to(x.device)\n",
    "        else:\n",
    "            return ((x > 0) * 1)[:, :max(lengths)]\n",
    "\n",
    "    def forward(self, x, lengths):\n",
    "        mask = self.get_mask(x, lengths)\n",
    "        if len(x.shape) > 2:\n",
    "            x = torch.matmul(x, self.token.weight)\n",
    "        else:\n",
    "            x = self.token(x)\n",
    "\n",
    "        return self.embed_dropout(x), mask\n",
    "\n",
    "\n",
    "class NARMModel(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super().__init__()\n",
    "        embed_size = args.bert_hidden_units\n",
    "        hidden_size = 2 * args.bert_hidden_units\n",
    "\n",
    "        self.gru = nn.GRU(embed_size, hidden_size, num_layers=1, batch_first=True)\n",
    "        self.a_global = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.a_local = nn.Linear(hidden_size, hidden_size, bias=False)\n",
    "        self.act = HardSigmoid()\n",
    "        self.v_vector = nn.Linear(hidden_size, 1, bias=False)\n",
    "        self.proj_dropout = nn.Dropout(args.bert_attn_dropout)\n",
    "        self.b_vetor = nn.Linear(embed_size, 2 * hidden_size, bias=False)\n",
    "\n",
    "    def forward(self, x, embedding_weight, lengths, mask):\n",
    "        x = pack_padded_sequence(x, lengths, batch_first=True, enforce_sorted=False)\n",
    "        gru_out, hidden = self.gru(x)\n",
    "        gru_out, _ = pad_packed_sequence(gru_out, batch_first=True)\n",
    "        c_global = hidden[-1]\n",
    "\n",
    "        state2 = self.a_local(gru_out)\n",
    "        state1 = self.a_global(c_global).unsqueeze(1).expand_as(state2)\n",
    "        state1 = mask.unsqueeze(2).expand_as(state2) * state1\n",
    "        alpha = self.act(state1 + state2).view(-1, state1.size(-1))\n",
    "        attn = self.v_vector(alpha).view(mask.size())\n",
    "        attn = F.softmax(attn.masked_fill(mask == 0, -1e9), dim=-1)\n",
    "        c_local = torch.sum(attn.unsqueeze(2).expand_as(gru_out) * gru_out, 1)\n",
    "\n",
    "        proj = self.proj_dropout(torch.cat([c_global, c_local], 1))\n",
    "        scores = torch.matmul(proj, self.b_vetor(embedding_weight).permute(1, 0))\n",
    "        return scores\n",
    "\n",
    "\n",
    "class HardSigmoid(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return torch.clamp((x / 6 + 0.5), min=0., max=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NARM(nn.Module):\n",
    "    def __init__(self, args):\n",
    "        super(NARM, self).__init__()\n",
    "        self.args = args\n",
    "        self.embedding = NARMEmbedding(self.args)\n",
    "        self.model = NARMModel(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 p in self.parameters():\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, lengths):\n",
    "        x, mask = self.embedding(x, lengths)\n",
    "        scores = self.model(x, self.embedding.token.weight, lengths, mask)\n",
    "        return scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Module.parameters of NARM(\n",
       "  (embedding): NARMEmbedding(\n",
       "    (token): Embedding(11, 4)\n",
       "    (embed_dropout): Dropout(p=0.2, inplace=False)\n",
       "  )\n",
       "  (model): NARMModel(\n",
       "    (gru): GRU(4, 8, batch_first=True)\n",
       "    (a_global): Linear(in_features=8, out_features=8, bias=False)\n",
       "    (a_local): Linear(in_features=8, out_features=8, bias=False)\n",
       "    (act): HardSigmoid()\n",
       "    (v_vector): Linear(in_features=8, out_features=1, bias=False)\n",
       "    (proj_dropout): Dropout(p=0.2, inplace=False)\n",
       "    (b_vetor): Linear(in_features=4, out_features=16, bias=False)\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 = NARM(args)\n",
    "model.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://arxiv.org/abs/1711.04725\n",
    "2. https://github.com/Yueeeeeeee/RecSys-Extraction-Attack/blob/main/model/narm.py\n",
    "3. https://recbole.io/docs/recbole/recbole.model.sequential_recommender.narm.html\n",
    "4. https://github.com/Wang-Shuo/Neural-Attentive-Session-Based-Recommendation-PyTorch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-31 07:01:15\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",
      "numpy     : 1.19.5\n",
      "matplotlib: 3.2.2\n",
      "torch     : 1.10.0+cu111\n",
      "PIL       : 7.1.2\n",
      "IPython   : 5.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\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
}
