{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.tagnn_pp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TAGNN++\n",
    "> [Mitheran et. al. Improved Representation Learning for Session-based Recommendation. arXiv, 2021.](https://arxiv.org/abs/2107.01516v2)\n",
    "\n",
    "TAGNN models item interactions with GNN, and both local and global user interactions with  a Transformer.\n",
    "\n",
    "<img src='https://raw.githubusercontent.com/RecoHut-Projects/sessrec-gnn/main/report/images/img7.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",
    "from collections import Iterable\n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def unitwise_norm(x: torch.Tensor):\n",
    "    if x.ndim <= 1:\n",
    "        dim = 0\n",
    "        keepdim = False\n",
    "    elif x.ndim in [2, 3]:\n",
    "        dim = 0\n",
    "        keepdim = True\n",
    "    elif x.ndim == 4:\n",
    "        dim = [1, 2, 3]\n",
    "        keepdim = True\n",
    "    else:\n",
    "        raise ValueError('Wrong dimensions of x')\n",
    "\n",
    "    return torch.sum(x**2, dim=dim, keepdim=keepdim) ** 0.5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class AGC(optim.Optimizer):\n",
    "    \n",
    "    \"\"\"Generic implementation of the Adaptive Gradient Clipping\n",
    "    Args:\n",
    "      params (iterable): iterable of parameters to optimize or dicts defining\n",
    "            parameter groups\n",
    "      optim (torch.optim.Optimizer): Optimizer with base class optim.Optimizer\n",
    "      clipping (float, optional): clipping value (default: 1e-3)\n",
    "      eps (float, optional): eps (default: 1e-3)\n",
    "      model (torch.nn.Module, optional): The original model\n",
    "      ignore_agc (str, Iterable, optional): Layers for AGC to ignore\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, params, optim: optim.Optimizer, clipping: float = 1e-2, eps: float = 1e-3, model=None, ignore_agc= ['']):\n",
    "        if clipping < 0.0:\n",
    "            raise ValueError(\"Invalid clipping value: {}\".format(clipping))\n",
    "        if eps < 0.0:\n",
    "            raise ValueError(\"Invalid eps value: {}\".format(eps))\n",
    "\n",
    "        self.optim = optim\n",
    "\n",
    "        defaults = dict(clipping=clipping, eps=eps)\n",
    "        defaults = {**defaults, **optim.defaults}\n",
    "\n",
    "        if not isinstance(ignore_agc, Iterable):\n",
    "            ignore_agc = [ignore_agc]\n",
    "\n",
    "        if model is not None:\n",
    "            assert ignore_agc not in [\n",
    "                None, []], \"Specify args ignore_agc to ignore fc-like(or other) layers\"\n",
    "            names = [name for name, module in model.named_modules()]\n",
    "\n",
    "            for module_name in ignore_agc:\n",
    "                if module_name not in names:\n",
    "                    raise ModuleNotFoundError(\n",
    "                        \"Module name {} not found in the model\".format(module_name))\n",
    "            parameters = [{\"params\": module.parameters()} for name,\n",
    "                          module in model.named_modules() if name not in ignore_agc]\n",
    "\n",
    "        super(AGC, self).__init__(params, defaults)\n",
    "\n",
    "    @torch.no_grad()\n",
    "    def step(self, closure=None):\n",
    "        \"\"\"Performs a single optimization step.\n",
    "        Arguments:\n",
    "            closure (callable, optional): A closure that reevaluates the model\n",
    "                and returns the loss.\n",
    "        \"\"\"\n",
    "        loss = None\n",
    "        if closure is not None:\n",
    "            with torch.enable_grad():\n",
    "                loss = closure()\n",
    "\n",
    "        for group in self.param_groups:\n",
    "            for p in group['params']:\n",
    "                if p.grad is None:\n",
    "                    continue\n",
    "                param_norm = torch.max(unitwise_norm(\n",
    "                    p.detach()), torch.tensor(group['eps']).to(p.device))\n",
    "                grad_norm = unitwise_norm(p.grad.detach())\n",
    "                max_norm = param_norm * group['clipping']\n",
    "\n",
    "                trigger = grad_norm < max_norm\n",
    "\n",
    "                clipped_grad = p.grad * \\\n",
    "                    (max_norm / torch.max(grad_norm,\n",
    "                                          torch.tensor(1e-6).to(grad_norm.device)))\n",
    "                p.grad.data.copy_(torch.where(trigger, clipped_grad, p.grad))\n",
    "\n",
    "        return self.optim.step(closure)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class Attention_GNN(nn.Module):\n",
    "    def __init__(self, hidden_size, step=1):\n",
    "        super(Attention_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(\n",
    "            self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_edge_out = nn.Linear(\n",
    "            self.hidden_size, self.hidden_size, bias=True)\n",
    "        self.linear_edge_f = nn.Linear(\n",
    "            self.hidden_size, self.hidden_size, bias=True)\n",
    "\n",
    "    def GNNCell(self, A, hidden):\n",
    "        input_in = torch.matmul(A[:, :, :A.shape[1]],\n",
    "                                self.linear_edge_in(hidden)) + self.b_iah\n",
    "\n",
    "        input_out = torch.matmul(\n",
    "            A[:, :, A.shape[1]: 2 * A.shape[1]], self.linear_edge_out(hidden)) + self.b_oah\n",
    "\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_PP(nn.Module):\n",
    "    def __init__(self, opt):\n",
    "        super(TAGNN_PP, self).__init__()\n",
    "        self.hidden_size = args.hiddenSize\n",
    "        self.n_node = opt.n_node\n",
    "        self.batch_size = args.batchSize\n",
    "        self.nonhybrid = args.nonhybrid\n",
    "        self.embedding = nn.Embedding(self.n_node, self.hidden_size)\n",
    "        self.tagnn = Attention_GNN(self.hidden_size, step=args.step)\n",
    "\n",
    "        self.layer_norm1 = nn.LayerNorm(self.hidden_size)\n",
    "        self.attn = nn.MultiheadAttention(\n",
    "            embed_dim=self.hidden_size, num_heads=2, dropout=0.1)\n",
    "\n",
    "        self.linear_one = nn.Linear(\n",
    "            self.hidden_size, self.hidden_size, bias=True)\n",
    "\n",
    "        self.linear_two = nn.Linear(\n",
    "            self.hidden_size, self.hidden_size, bias=True)\n",
    "\n",
    "        self.linear_three = nn.Linear(self.hidden_size, 1, bias=False)\n",
    "        self.linear_transform = nn.Linear(\n",
    "            self.hidden_size * 2, self.hidden_size, bias=True)\n",
    "        self.linear_t = nn.Linear(\n",
    "            self.hidden_size, self.hidden_size, bias=False)  # target attention\n",
    "        self.loss_function = nn.CrossEntropyLoss()\n",
    "        self.optimizer = torch.optim.Adam(\n",
    "            self.parameters(), lr=args.lr, weight_decay=args.l2)\n",
    "        self.agc_optimizer = AGC(self.parameters(), self.optimizer, model=self)\n",
    "        self.scheduler = torch.optim.lr_scheduler.StepLR(\n",
    "            self.optimizer, step_size=args.lr_dc_step, gamma=args.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(\n",
    "            mask, 1) - 1]  # batch_size x latent_size\n",
    "        # batch_size x 1 x latent_size\n",
    "        q1 = self.linear_one(ht).view(ht.shape[0], 1, ht.shape[1])\n",
    "        q2 = self.linear_two(hidden)  # batch_size x seq_length x latent_size\n",
    "        # batch_size x seq_length x 1\n",
    "        alpha = self.linear_three(torch.sigmoid(q1 + q2))\n",
    "        alpha = F.softmax(alpha, 1)  # batch_size x seq_length x 1\n",
    "        # batch_size x latent_size\n",
    "        a = torch.sum(alpha * hidden *\n",
    "                      mask.view(mask.shape[0], -1, 1).float(), 1)\n",
    "\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",
    "\n",
    "        # batch_size x seq_length x latent_size\n",
    "        hidden = hidden * mask.view(mask.shape[0], -1, 1).float()\n",
    "        qt = self.linear_t(hidden)  # batch_size x seq_length x latent_size\n",
    "        # batch_size x n_nodes x seq_length\n",
    "        beta = F.softmax(b @ qt.transpose(1, 2), -1)\n",
    "        target = beta @ hidden  # batch_size x n_nodes x latent_size\n",
    "        a = a.view(ht.shape[0], 1, ht.shape[1])  # batch_size x 1 x latent_size\n",
    "        a = a + target  # batch_size x n_nodes x latent_size\n",
    "        scores = torch.sum(a * b, -1)  # batch_size x n_nodes\n",
    "        return scores\n",
    "\n",
    "    @staticmethod\n",
    "    def get_mask(seq_len):\n",
    "        return torch.from_numpy(np.triu(np.ones((seq_len, seq_len)), k=1).astype('bool'))#.to('cuda')\n",
    "\n",
    "    def forward(self, inputs, A):\n",
    "        hidden = self.embedding(inputs)\n",
    "        hidden = self.tagnn(A, hidden)\n",
    "        hidden = hidden.permute(1, 0, 2)\n",
    "\n",
    "        skip = self.layer_norm1(hidden)\n",
    "        hidden, attn_w = self.attn(\n",
    "            hidden, hidden, hidden, attn_mask=self.get_mask(hidden.shape[0]))\n",
    "        hidden = hidden+skip\n",
    "        hidden = hidden.permute(1, 0, 2)\n",
    "\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 = 5 # 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 to_cuda(input_variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return input_variable.cuda()\n",
    "    else:\n",
    "        return input_variable\n",
    "\n",
    "\n",
    "def to_cpu(input_variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return input_variable.cpu()\n",
    "    else:\n",
    "        return input_variable\n",
    "\n",
    "\n",
    "def forward(model, i, data):\n",
    "    alias_inputs, A, items, mask, targets = data.get_slice(i)\n",
    "    alias_inputs = to_cuda(torch.Tensor(alias_inputs).long())\n",
    "    items = to_cuda(torch.Tensor(items).long())\n",
    "    A = to_cuda(torch.Tensor(A).float())\n",
    "    mask = to_cuda(torch.Tensor(mask).long())\n",
    "    hidden = model(items, A)\n",
    "\n",
    "    def get(i): return hidden[i][alias_inputs[i]]\n",
    "    seq_hidden = torch.stack([get(i)\n",
    "                             for i in torch.arange(len(alias_inputs)).long()])\n",
    "\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",
    "\n",
    "    from tqdm.notebook import tqdm\n",
    "    for i, j in tqdm(zip(slices, np.arange(len(slices))), total=len(slices)):\n",
    "        model.optimizer.zero_grad()\n",
    "        targets, scores = forward(model, i, train_data)\n",
    "        targets = 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",
    "\n",
    "        if j % int(len(slices) / 5 + 1) == 0:\n",
    "            print('[%d/%d] Loss: %.4f' % (j, len(slices), loss.item()))\n",
    "\n",
    "    print('\\tLoss Value:\\t%.3f' % total_loss)\n",
    "    print('Start Prediction: ', datetime.datetime.now())\n",
    "\n",
    "    model.eval()\n",
    "    hit, mrr = [], []\n",
    "    slices = test_data.generate_batch(model.batch_size)\n",
    "\n",
    "    for i in slices:\n",
    "        targets, scores = forward(model, i, test_data)\n",
    "        sub_scores = scores.topk(20)[1]\n",
    "        sub_scores = to_cpu(sub_scores).detach().numpy()\n",
    "\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",
    "\n",
    "    hit = np.mean(hit) * 100\n",
    "    mrr = np.mean(mrr) * 100\n",
    "    return hit, mrr\n",
    "\n",
    "\n",
    "def get_pos(seq_len):\n",
    "    return torch.arange(seq_len).unsqueeze(0)\n",
    "\n",
    "\n",
    "def str2bool(v):\n",
    "    return v.lower() in ('true')\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:37:30.078500\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3437a677f8834075912b519d1ab5b02e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.7335\n",
      "[3/11] Loss: 5.7263\n",
      "[6/11] Loss: 5.7291\n",
      "[9/11] Loss: 5.7135\n",
      "\tLoss Value:\t62.968\n",
      "Start Prediction:  2021-12-23 11:37:31.807639\n",
      "Best Result:\n",
      "\tRecall@20:\t17.3554\tMRR@20:\t5.2499\tEpoch:\t0,\t0\n",
      "--------------------------------------------------\n",
      "Epoch:  1\n",
      "Start training:  2021-12-23 11:37:31.935884\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "91c69fba49d6493390082a81987e9f1d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.7028\n",
      "[3/11] Loss: 5.6784\n",
      "[6/11] Loss: 5.6666\n",
      "[9/11] Loss: 5.6629\n",
      "\tLoss Value:\t62.498\n",
      "Start Prediction:  2021-12-23 11:37:33.677782\n",
      "Best Result:\n",
      "\tRecall@20:\t19.8347\tMRR@20:\t5.2499\tEpoch:\t1,\t0\n",
      "--------------------------------------------------\n",
      "Epoch:  2\n",
      "Start training:  2021-12-23 11:37:33.802865\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "38cacdddf31a4d1ba099b031b217ed42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.6609\n",
      "[3/11] Loss: 5.6498\n",
      "[6/11] Loss: 5.6462\n",
      "[9/11] Loss: 5.6439\n",
      "\tLoss Value:\t62.090\n",
      "Start Prediction:  2021-12-23 11:37:35.502085\n",
      "Best Result:\n",
      "\tRecall@20:\t20.6612\tMRR@20:\t5.2499\tEpoch:\t2,\t0\n",
      "--------------------------------------------------\n",
      "Epoch:  3\n",
      "Start training:  2021-12-23 11:37:35.604826\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "17bf5edbf9624ec19a26bbc1b20191db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.6461\n",
      "[3/11] Loss: 5.6324\n",
      "[6/11] Loss: 5.6430\n",
      "[9/11] Loss: 5.6417\n",
      "\tLoss Value:\t61.997\n",
      "Start Prediction:  2021-12-23 11:37:37.213415\n",
      "Best Result:\n",
      "\tRecall@20:\t21.4876\tMRR@20:\t5.2499\tEpoch:\t3,\t0\n",
      "--------------------------------------------------\n",
      "Epoch:  4\n",
      "Start training:  2021-12-23 11:37:37.333273\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4fe37570e1bc40cf9161a06e0bf4f4af",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/11 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0/11] Loss: 5.6551\n",
      "[3/11] Loss: 5.6529\n",
      "[6/11] Loss: 5.5994\n",
      "[9/11] Loss: 5.6067\n",
      "\tLoss Value:\t61.898\n",
      "Start Prediction:  2021-12-23 11:37:38.855207\n",
      "Best Result:\n",
      "\tRecall@20:\t21.4876\tMRR@20:\t5.2499\tEpoch:\t4,\t0\n",
      "--------------------------------------------------\n",
      "Running time: 8.897622 seconds\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pickle\n",
    "import time\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "from recohut.datasets.session import SampleSessionDataset, GraphData\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "model_save_dir = 'saved/'\n",
    "log_dir='saved/logs'\n",
    "writer = SummaryWriter(log_dir=log_dir)\n",
    "os.makedirs(log_dir, exist_ok=True)\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 = to_cuda(TAGNN_PP(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('-' * 50)\n",
    "    print('Epoch: ', epoch)\n",
    "    hit, mrr = train_test(model, train_data, test_data)\n",
    "    flag = 0\n",
    "\n",
    "    # Logging\n",
    "    writer.add_scalar('epoch/recall', hit, epoch)\n",
    "    writer.add_scalar('epoch/mrr', mrr, epoch)\n",
    "\n",
    "    flag = 0\n",
    "\n",
    "    if hit >= best_result[0]:\n",
    "        best_result[0] = hit\n",
    "        best_epoch[0] = epoch\n",
    "        flag = 1\n",
    "        torch.save(model, model_save_dir + 'epoch_' +\n",
    "                    str(epoch) + '_recall_' + str(hit) + '_.pt')\n",
    "    if mrr >= best_result[1]:\n",
    "        best_result[1] = mrr\n",
    "        best_epoch[1] = epoch\n",
    "        flag = 1\n",
    "        torch.save(model, model_save_dir + 'epoch_' +\n",
    "                    str(epoch) + '_mrr_' + str(mrr) + '_.pt')\n",
    "\n",
    "    print('Best Result:')\n",
    "    print('\\tRecall@20:\\t%.4f\\tMRR@20:\\t%.4f\\tEpoch:\\t%d,\\t%d' %\n",
    "            (best_result[0], best_result[1], best_epoch[0], best_epoch[1]))\n",
    "\n",
    "    bad_counter += 1 - flag\n",
    "\n",
    "    if bad_counter >= args.patience:\n",
    "        break\n",
    "\n",
    "print('-' * 50)\n",
    "end = time.time()\n",
    "print(\"Running time: %f seconds\" % (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:37:42\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
}
