{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.tagnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TAGNN\n",
    "> [Yu et. al. Target Attentive Graph Neural Networks for Session-based Recommendation. SIGIR, 2020.](https://arxiv.org/abs/2005.02844)\n",
    "\n",
    "TAGNN first models all session sequences as session graphs. Then, graph neural networks capture rich item transitions in sessions. Lastly, from one session embedding vector, target-aware attention adaptively activates different user interests concerning varied target items to be predicted.\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/RecoHut-Projects/sessrec-gnn/main/report/images/img6.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 TAGNN(nn.Module):\n",
    "    def __init__(self, opt):\n",
    "        super(TAGNN, 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.linear_t = nn.Linear(self.hidden_size, self.hidden_size, bias=False)  #target attention\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))  # (b,s,1)\n",
    "        # alpha = torch.sigmoid(alpha) # B,S,1\n",
    "        alpha = F.softmax(alpha, 1) # B,S,1\n",
    "        a = torch.sum(alpha * hidden * mask.view(mask.shape[0], -1, 1).float(), 1)  # (b,d)\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",
    "        # target attention: sigmoid(hidden M b)\n",
    "        # mask  # batch_size x seq_length\n",
    "        hidden = hidden * mask.view(mask.shape[0], -1, 1).float()  # batch_size x seq_length x latent_size\n",
    "        qt = self.linear_t(hidden)  # batch_size x seq_length x latent_size\n",
    "        # beta = torch.sigmoid(b @ qt.transpose(1,2))  # batch_size x n_nodes x seq_length\n",
    "        beta = F.softmax(b @ qt.transpose(1,2), -1)  # batch_size x n_nodes x seq_length\n",
    "        target = beta @ hidden  # batch_size x n_nodes x latent_size\n",
    "        a = a.view(ht.shape[0], 1, ht.shape[1])  # b,1,d\n",
    "        a = a + target  # b,n,d\n",
    "        scores = torch.sum(a * b, -1)  # b,n\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 TAGNN"
   ]
  },
  {
   "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.item()\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": "stdout",
     "output_type": "stream",
     "text": [
      "-------------------------------------------------------\n",
      "epoch:  0\n",
      "start training:  2021-12-23 11:17:17.358225\n",
      "[0/11] Loss: 5.7136\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": [
      "[3/11] Loss: 5.7028\n",
      "[6/11] Loss: 5.7000\n",
      "[9/11] Loss: 5.6981\n",
      "\tLoss:\t62.736\n",
      "start predicting:  2021-12-23 11:17:18.800668\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  1\n",
      "start training:  2021-12-23 11:17:18.881197\n",
      "[0/11] Loss: 5.6988\n",
      "[3/11] Loss: 5.6828\n",
      "[6/11] Loss: 5.6606\n",
      "[9/11] Loss: 5.6704\n",
      "\tLoss:\t62.421\n",
      "start predicting:  2021-12-23 11:17:20.583203\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  2\n",
      "start training:  2021-12-23 11:17:20.668376\n",
      "[0/11] Loss: 5.6598\n",
      "[3/11] Loss: 5.6544\n",
      "[6/11] Loss: 5.6495\n",
      "[9/11] Loss: 5.6483\n",
      "\tLoss:\t62.155\n",
      "start predicting:  2021-12-23 11:17:22.019096\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  3\n",
      "start training:  2021-12-23 11:17:22.106588\n",
      "[0/11] Loss: 5.6384\n",
      "[3/11] Loss: 5.6427\n",
      "[6/11] Loss: 5.6393\n",
      "[9/11] Loss: 5.6504\n",
      "\tLoss:\t62.093\n",
      "start predicting:  2021-12-23 11:17:23.456983\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  4\n",
      "start training:  2021-12-23 11:17:23.542167\n",
      "[0/11] Loss: 5.6427\n",
      "[3/11] Loss: 5.6565\n",
      "[6/11] Loss: 5.6579\n",
      "[9/11] Loss: 5.6316\n",
      "\tLoss:\t62.017\n",
      "start predicting:  2021-12-23 11:17:24.940504\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  5\n",
      "start training:  2021-12-23 11:17:25.023468\n",
      "[0/11] Loss: 5.6347\n",
      "[3/11] Loss: 5.6364\n",
      "[6/11] Loss: 5.6495\n",
      "[9/11] Loss: 5.6352\n",
      "\tLoss:\t61.969\n",
      "start predicting:  2021-12-23 11:17:26.257310\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  6\n",
      "start training:  2021-12-23 11:17:26.341687\n",
      "[0/11] Loss: 5.6456\n",
      "[3/11] Loss: 5.6245\n",
      "[6/11] Loss: 5.6465\n",
      "[9/11] Loss: 5.6343\n",
      "\tLoss:\t61.965\n",
      "start predicting:  2021-12-23 11:17:27.782870\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  7\n",
      "start training:  2021-12-23 11:17:27.867611\n",
      "[0/11] Loss: 5.6178\n",
      "[3/11] Loss: 5.6378\n",
      "[6/11] Loss: 5.6158\n",
      "[9/11] Loss: 5.6217\n",
      "\tLoss:\t61.952\n",
      "start predicting:  2021-12-23 11:17:29.276856\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  8\n",
      "start training:  2021-12-23 11:17:29.355513\n",
      "[0/11] Loss: 5.6418\n",
      "[3/11] Loss: 5.6355\n",
      "[6/11] Loss: 5.6157\n",
      "[9/11] Loss: 5.6426\n",
      "\tLoss:\t61.950\n",
      "start predicting:  2021-12-23 11:17:30.763400\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  9\n",
      "start training:  2021-12-23 11:17:30.851242\n",
      "[0/11] Loss: 5.6316\n",
      "[3/11] Loss: 5.6091\n",
      "[6/11] Loss: 5.6344\n",
      "[9/11] Loss: 5.6349\n",
      "\tLoss:\t61.947\n",
      "start predicting:  2021-12-23 11:17:32.230495\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "epoch:  10\n",
      "start training:  2021-12-23 11:17:32.310292\n",
      "[0/11] Loss: 5.6269\n",
      "[3/11] Loss: 5.6355\n",
      "[6/11] Loss: 5.6396\n",
      "[9/11] Loss: 5.6452\n",
      "\tLoss:\t61.945\n",
      "start predicting:  2021-12-23 11:17:33.564034\n",
      "Best Result:\n",
      "\tRecall@20:\t62.8099\tMMR@20:\t46.2166\tEpoch:\t0,\t0\n",
      "-------------------------------------------------------\n",
      "Run time: 16.294159 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('./session_ds/processed/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(TAGNN(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 11:17:51\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",
      "IPython: 5.5.0\n",
      "numpy  : 1.19.5\n",
      "torch  : 1.10.0+cu111\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
}
