{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.srgnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SRGNN\n",
    "> [Wu et. al. Session-based Recommendation with Graph Neural Networks. AAAI, 2019.](https://arxiv.org/abs/1811.00855)\n",
    "\n",
    "SR-GNN first models all session sequences as session graphs. Then, each session graph is proceeded one by one and the resulting node vectors can be obtained through a gated graph neural network. After that, each session is represented as the combination of the global preference and current interests of this session using an attention net. Finally, it predict the probability of each item that will appear to be the next-click one for each session.\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/RecoHut-Projects/sessrec-gnn/main/report/images/img2.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 math\n",
    "import datetime\n",
    "import numpy as np\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from recohut.datasets.session import SampleSessionDataset, GraphData"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class GNN(nn.Module):\n",
    "    def __init__(self, hidden_size, step=1):\n",
    "        super(GNN, self).__init__()\n",
    "        self.step = step\n",
    "        self.hidden_size = hidden_size\n",
    "        self.input_size = hidden_size * 2\n",
    "        self.gate_size = 3 * hidden_size\n",
    "        self.w_ih = nn.Parameter(torch.Tensor(self.gate_size, self.input_size))\n",
    "        self.w_hh = nn.Parameter(torch.Tensor(self.gate_size, self.hidden_size))\n",
    "        self.b_ih = nn.Parameter(torch.Tensor(self.gate_size))\n",
    "        self.b_hh = nn.Parameter(torch.Tensor(self.gate_size))\n",
    "        self.b_iah = nn.Parameter(torch.Tensor(self.hidden_size))\n",
    "        self.b_oah = nn.Parameter(torch.Tensor(self.hidden_size))\n",
    "\n",
    "        self.linear_edge_in = nn.Linear(self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_edge_out = nn.Linear(self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_edge_f = nn.Linear(self.hidden_size, self.hidden_size, bias=True)\n",
    "\n",
    "    def GNNCell(self, A, hidden):\n",
    "        input_in = torch.matmul(A[:, :, :A.shape[1]], self.linear_edge_in(hidden)) + self.b_iah\n",
    "        input_out = torch.matmul(A[:, :, A.shape[1]: 2 * A.shape[1]], self.linear_edge_out(hidden)) + self.b_oah\n",
    "        inputs = torch.cat([input_in, input_out], 2)\n",
    "        gi = F.linear(inputs, self.w_ih, self.b_ih)\n",
    "        gh = F.linear(hidden, self.w_hh, self.b_hh)\n",
    "        i_r, i_i, i_n = gi.chunk(3, 2)\n",
    "        h_r, h_i, h_n = gh.chunk(3, 2)\n",
    "        resetgate = torch.sigmoid(i_r + h_r)\n",
    "        inputgate = torch.sigmoid(i_i + h_i)\n",
    "        newgate = torch.tanh(i_n + resetgate * h_n)\n",
    "        hy = newgate + inputgate * (hidden - newgate)\n",
    "        return hy\n",
    "\n",
    "    def forward(self, A, hidden):\n",
    "        for i in range(self.step):\n",
    "            hidden = self.GNNCell(A, hidden)\n",
    "        return hidden"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SRGNN(nn.Module):\n",
    "    def __init__(self, opt):\n",
    "        super(SRGNN, self).__init__()\n",
    "        self.hidden_size = opt.hiddenSize\n",
    "        self.n_node = opt.n_node\n",
    "        self.batch_size = opt.batchSize\n",
    "        self.nonhybrid = opt.nonhybrid\n",
    "        self.embedding = nn.Embedding(self.n_node, self.hidden_size)\n",
    "        self.gnn = GNN(self.hidden_size, step=opt.step)\n",
    "        self.linear_one = nn.Linear(self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_two = nn.Linear(self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_three = nn.Linear(self.hidden_size, 1, bias=False)\n",
    "        self.linear_transform = nn.Linear(self.hidden_size * 2, self.hidden_size, bias=True)\n",
    "        self.loss_function = nn.CrossEntropyLoss()\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=opt.lr, weight_decay=opt.l2)\n",
    "        self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=opt.lr_dc_step, gamma=opt.lr_dc)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1.0 / math.sqrt(self.hidden_size)\n",
    "        for weight in self.parameters():\n",
    "            weight.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def compute_scores(self, hidden, mask):\n",
    "        ht = hidden[torch.arange(mask.shape[0]).long(), torch.sum(mask, 1) - 1]  # batch_size x latent_size\n",
    "        q1 = self.linear_one(ht).view(ht.shape[0], 1, ht.shape[1])  # batch_size x 1 x latent_size\n",
    "        q2 = self.linear_two(hidden)  # batch_size x seq_length x latent_size\n",
    "        alpha = self.linear_three(torch.sigmoid(q1 + q2))\n",
    "        a = torch.sum(alpha * hidden * mask.view(mask.shape[0], -1, 1).float(), 1)\n",
    "        if not self.nonhybrid:\n",
    "            a = self.linear_transform(torch.cat([a, ht], 1))\n",
    "        b = self.embedding.weight[1:]  # n_nodes x latent_size\n",
    "        scores = torch.matmul(a, b.transpose(1, 0))\n",
    "        return scores\n",
    "\n",
    "    def forward(self, inputs, A):\n",
    "        hidden = self.embedding(inputs)\n",
    "        hidden = self.gnn(A, hidden)\n",
    "        return hidden"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training a session-based recommender using SR-GNN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args():\n",
    "    dataset = 'sample'\n",
    "    batchSize = 100 # input batch size\n",
    "    hiddenSize = 100 # hidden state size\n",
    "    epoch = 30 # the number of epochs to train for\n",
    "    lr = 0.001 # learning rate')  # [0.001, 0.0005, 0.000\n",
    "    lr_dc = 0.1 # learning rate decay rate\n",
    "    lr_dc_step = 3 # the number of steps after which the learning rate decay\n",
    "    l2 = 1e-5 # l2 penalty')  # [0.001, 0.0005, 0.0001, 0.00005, 0.0000\n",
    "    step = 1 # gnn propogation steps\n",
    "    patience = 10 # the number of epoch to wait before early stop \n",
    "    nonhybrid = True # only use the global preference to predict\n",
    "    validation = True # validation\n",
    "    valid_portion = 0.1 # split the portion of training set as validation set\n",
    "    n_node = 310\n",
    "\n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trans_to_cuda(variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return variable.cuda()\n",
    "    else:\n",
    "        return variable\n",
    "\n",
    "\n",
    "def trans_to_cpu(variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return variable.cpu()\n",
    "    else:\n",
    "        return variable\n",
    "\n",
    "\n",
    "def forward(model, i, data):\n",
    "    alias_inputs, A, items, mask, targets = data.get_slice(i)\n",
    "    alias_inputs = trans_to_cuda(torch.Tensor(alias_inputs).long())\n",
    "    items = trans_to_cuda(torch.Tensor(items).long())\n",
    "    A = trans_to_cuda(torch.Tensor(A).float())\n",
    "    mask = trans_to_cuda(torch.Tensor(mask).long())\n",
    "    hidden = model(items, A)\n",
    "    get = lambda i: hidden[i][alias_inputs[i]]\n",
    "    seq_hidden = torch.stack([get(i) for i in torch.arange(len(alias_inputs)).long()])\n",
    "    return targets, model.compute_scores(seq_hidden, mask)\n",
    "\n",
    "\n",
    "def train_test(model, train_data, test_data):\n",
    "    model.scheduler.step()\n",
    "    print('start training: ', datetime.datetime.now())\n",
    "    model.train()\n",
    "    total_loss = 0.0\n",
    "    slices = train_data.generate_batch(model.batch_size)\n",
    "    for i, j in zip(slices, np.arange(len(slices))):\n",
    "        model.optimizer.zero_grad()\n",
    "        targets, scores = forward(model, i, train_data)\n",
    "        targets = trans_to_cuda(torch.Tensor(targets).long())\n",
    "        loss = model.loss_function(scores, targets - 1)\n",
    "        loss.backward()\n",
    "        model.optimizer.step()\n",
    "        total_loss += loss\n",
    "        if j % int(len(slices) / 5 + 1) == 0:\n",
    "            print('[%d/%d] Loss: %.4f' % (j, len(slices), loss.item()))\n",
    "    print('\\tLoss:\\t%.3f' % total_loss)\n",
    "\n",
    "    print('start predicting: ', datetime.datetime.now())\n",
    "    model.eval()\n",
    "    hit, mrr = [], []\n",
    "    slices = test_data.generate_batch(model.batch_size)\n",
    "    for i in slices:\n",
    "        targets, scores = forward(model, i, test_data)\n",
    "        sub_scores = scores.topk(20)[1]\n",
    "        sub_scores = trans_to_cpu(sub_scores).detach().numpy()\n",
    "        for score, target, mask in zip(sub_scores, targets, test_data.mask):\n",
    "            hit.append(np.isin(target - 1, score))\n",
    "            if len(np.where(score == target - 1)[0]) == 0:\n",
    "                mrr.append(0)\n",
    "            else:\n",
    "                mrr.append(1 / (np.where(score == target - 1)[0][0] + 1))\n",
    "    hit = np.mean(hit) * 100\n",
    "    mrr = np.mean(mrr) * 100\n",
    "    return hit, mrr\n",
    "\n",
    "\n",
    "def split_validation(train_set, valid_portion):\n",
    "    train_set_x, train_set_y = train_set\n",
    "    n_samples = len(train_set_x)\n",
    "    sidx = np.arange(n_samples, dtype='int32')\n",
    "    np.random.shuffle(sidx)\n",
    "    n_train = int(np.round(n_samples * (1. - valid_portion)))\n",
    "    valid_set_x = [train_set_x[s] for s in sidx[n_train:]]\n",
    "    valid_set_y = [train_set_y[s] for s in sidx[n_train:]]\n",
    "    train_set_x = [train_set_x[s] for s in sidx[:n_train]]\n",
    "    train_set_y = [train_set_y[s] for s in sidx[:n_train]]\n",
    "\n",
    "    return (train_set_x, train_set_y), (valid_set_x, valid_set_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/sample_session/raw/v1/sample_train-item-views.csv\n",
      "Processing...\n",
      "Done!\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- Reading data\n",
      "Splitting date 1464134400.0\n",
      "469\n",
      "47\n",
      "[('2671', 1451952000.0), ('1211', 1452384000.0), ('3780', 1452384000.0)]\n",
      "[('1864', 1464220800.0), ('1867', 1464220800.0), ('1868', 1464220800.0)]\n",
      "-- Splitting train set and test set\n",
      "310\n",
      "1205\n",
      "99\n",
      "[[1, 2], [1], [4]] [1451952000.0, 1451952000.0, 1452384000.0] [3, 2, 5]\n",
      "[[282], [281, 308], [281]] [1464220800.0, 1464220800.0, 1464220800.0] [282, 281, 308]\n",
      "avg length:  3.5669291338582676\n",
      "-------------------------------------------------------\n",
      "epoch:  0\n",
      "start training:  2021-12-23 10:32:44.027013\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/torch/optim/lr_scheduler.py:134: UserWarning: Detected call of `lr_scheduler.step()` before `optimizer.step()`. In PyTorch 1.1.0 and later, you should call them in the opposite order: `optimizer.step()` before `lr_scheduler.step()`.  Failure to do this will result in PyTorch skipping the first value of the learning rate schedule. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\n",
      "  \"https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate\", UserWarning)\n",
      "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:19: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at  ../torch/csrc/utils/tensor_new.cpp:201.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.6835\n",
      "[3/11] Loss: 5.6416\n",
      "[6/11] Loss: 5.5608\n",
      "[9/11] Loss: 5.4904\n",
      "\tLoss:\t61.538\n",
      "start predicting:  2021-12-23 10:32:44.908047\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  1\n",
      "start training:  2021-12-23 10:32:44.973611\n",
      "[0/11] Loss: 5.4299\n",
      "[3/11] Loss: 5.2619\n",
      "[6/11] Loss: 5.2992\n",
      "[9/11] Loss: 4.9818\n",
      "\tLoss:\t56.973\n",
      "start predicting:  2021-12-23 10:32:45.623729\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  2\n",
      "start training:  2021-12-23 10:32:45.670660\n",
      "[0/11] Loss: 4.8519\n",
      "[3/11] Loss: 4.8910\n",
      "[6/11] Loss: 4.7712\n",
      "[9/11] Loss: 4.8673\n",
      "\tLoss:\t53.259\n",
      "start predicting:  2021-12-23 10:32:46.292004\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  3\n",
      "start training:  2021-12-23 10:32:46.344736\n",
      "[0/11] Loss: 4.7358\n",
      "[3/11] Loss: 4.6773\n",
      "[6/11] Loss: 4.9566\n",
      "[9/11] Loss: 4.7277\n",
      "\tLoss:\t52.641\n",
      "start predicting:  2021-12-23 10:32:46.968312\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  4\n",
      "start training:  2021-12-23 10:32:47.017120\n",
      "[0/11] Loss: 4.7608\n",
      "[3/11] Loss: 4.6921\n",
      "[6/11] Loss: 4.6916\n",
      "[9/11] Loss: 4.6267\n",
      "\tLoss:\t51.958\n",
      "start predicting:  2021-12-23 10:32:47.636812\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  5\n",
      "start training:  2021-12-23 10:32:47.683054\n",
      "[0/11] Loss: 4.9229\n",
      "[3/11] Loss: 4.6586\n",
      "[6/11] Loss: 4.6485\n",
      "[9/11] Loss: 4.5362\n",
      "\tLoss:\t51.537\n",
      "start predicting:  2021-12-23 10:32:48.308910\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  6\n",
      "start training:  2021-12-23 10:32:48.368369\n",
      "[0/11] Loss: 4.7077\n",
      "[3/11] Loss: 4.6776\n",
      "[6/11] Loss: 4.5178\n",
      "[9/11] Loss: 4.7291\n",
      "\tLoss:\t51.484\n",
      "start predicting:  2021-12-23 10:32:48.987707\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  7\n",
      "start training:  2021-12-23 10:32:49.035722\n",
      "[0/11] Loss: 4.7499\n",
      "[3/11] Loss: 4.4730\n",
      "[6/11] Loss: 4.5056\n",
      "[9/11] Loss: 4.7769\n",
      "\tLoss:\t51.375\n",
      "start predicting:  2021-12-23 10:32:49.683579\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  8\n",
      "start training:  2021-12-23 10:32:49.737845\n",
      "[0/11] Loss: 4.8842\n",
      "[3/11] Loss: 4.7976\n",
      "[6/11] Loss: 4.6972\n",
      "[9/11] Loss: 4.6249\n",
      "\tLoss:\t51.350\n",
      "start predicting:  2021-12-23 10:32:50.347680\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  9\n",
      "start training:  2021-12-23 10:32:50.394962\n",
      "[0/11] Loss: 4.6556\n",
      "[3/11] Loss: 4.9243\n",
      "[6/11] Loss: 4.4315\n",
      "[9/11] Loss: 4.7741\n",
      "\tLoss:\t51.400\n",
      "start predicting:  2021-12-23 10:32:50.934105\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  10\n",
      "start training:  2021-12-23 10:32:50.982599\n",
      "[0/11] Loss: 4.4425\n",
      "[3/11] Loss: 4.6440\n",
      "[6/11] Loss: 4.6355\n",
      "[9/11] Loss: 4.5264\n",
      "\tLoss:\t51.341\n",
      "start predicting:  2021-12-23 10:32:51.514488\n",
      "Best Result:\n",
      "\tRecall@20:\t56.1983\tMMR@20:\t40.6045\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "Run time: 7.538349 s\n"
     ]
    }
   ],
   "source": [
    "import pickle\n",
    "import time\n",
    "\n",
    "_ = SampleSessionDataset('./session_ds')\n",
    "train_data = pickle.load(open('./session_ds/processed/train.txt', 'rb'))\n",
    "\n",
    "if args.validation:\n",
    "    train_data, valid_data = split_validation(train_data, args.valid_portion)\n",
    "    test_data = valid_data\n",
    "else:\n",
    "    test_data = pickle.load(open('test.txt', 'rb'))\n",
    "\n",
    "train_data = GraphData(train_data, shuffle=True)\n",
    "test_data = GraphData(test_data, shuffle=False)\n",
    "\n",
    "model = trans_to_cuda(SRGNN(args))\n",
    "\n",
    "start = time.time()\n",
    "best_result = [0, 0]\n",
    "best_epoch = [0, 0]\n",
    "bad_counter = 0\n",
    "\n",
    "for epoch in range(args.epoch):\n",
    "    print('-------------------------------------------------------')\n",
    "    print('epoch: ', epoch)\n",
    "    hit, mrr = train_test(model, train_data, test_data)\n",
    "    flag = 0\n",
    "    if hit >= best_result[0]:\n",
    "        best_result[0] = hit\n",
    "        best_epoch[0] = epoch\n",
    "        flag = 1\n",
    "    if mrr >= best_result[1]:\n",
    "        best_result[1] = mrr\n",
    "        best_epoch[1] = epoch\n",
    "        flag = 1\n",
    "    print('Best Result:')\n",
    "    print('\\tRecall@20:\\t%.4f\\tMMR@20:\\t%.4f\\tEpoch:\\t%d,\\t%d'% (best_result[0], best_result[1], best_epoch[0], best_epoch[1]))\n",
    "    bad_counter += 1 - flag\n",
    "    if bad_counter >= args.patience:\n",
    "        break\n",
    "print('-------------------------------------------------------')\n",
    "end = time.time()\n",
    "print(\"Run time: %f s\" % (end - start))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-23 10:35:27\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",
      "torch  : 1.10.0+cu111\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
}
