{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.cotrec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# COTREC\n",
    "> Self-Supervised Graph Co-Training for Session-based Recommendation.\n",
    "\n",
    "Session-based recommendation targets next-item prediction by exploiting user behaviors within a short time period. Compared with other recommendation paradigms, session-based recommendation suffers more from the problem of data sparsity due to the very limited short-term interactions. Self-supervised learning, which can discover ground-truth samples from the raw data, holds vast potentials to tackle this problem. However, existing self-supervised recommendation models mainly rely on item/segment dropout to augment data, which are not fit for session-based recommendation because the dropout leads to sparser data, creating unserviceable self-supervision signals.\n",
    "\n",
    "For informative session-based data augmentation, COTREC combine self-supervised learning with co-training, and then develop a framework to enhance session-based recommendation. Technically, we first exploit the session-based graph to augment two views that exhibit the internal and external connectivities of sessions, and then we build two distinct graph encoders over the two views, which recursively leverage the different connectivity information to generate ground-truth samples to supervise each other by contrastive learning. In contrast to the dropout strategy, the proposed self-supervised graph co-training preserves the complete session information and fulfills genuine data augmentation.\n",
    "\n",
    "<img src='https://github.com/RecoHut-Stanzas/S969915/raw/main/images/cotrec_model.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 numpy as np\n",
    "import math\n",
    "\n",
    "import torch\n",
    "from torch import nn, backends\n",
    "import torch.nn.functional as F\n",
    "import torch.sparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def trans_to_cuda(variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return variable.cuda()\n",
    "    else:\n",
    "        return variable\n",
    "\n",
    "def trans_to_cpu(variable):\n",
    "    if torch.cuda.is_available():\n",
    "        return variable.cpu()\n",
    "    else:\n",
    "        return variable"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class ItemConv(nn.Module):\n",
    "    def __init__(self, layers, emb_size=100):\n",
    "        super(ItemConv, self).__init__()\n",
    "        self.emb_size = emb_size\n",
    "        self.layers = layers[0]\n",
    "        self.w_item = {}\n",
    "        for i in range(self.layers):\n",
    "            self.w_item['weight_item%d' % (i)] = nn.Linear(self.emb_size, self.emb_size, bias=False)\n",
    "\n",
    "    def forward(self, adjacency, embedding):\n",
    "        values = adjacency.data\n",
    "        indices = np.vstack((adjacency.row, adjacency.col))\n",
    "        i = torch.LongTensor(indices)\n",
    "        v = torch.FloatTensor(values)\n",
    "        shape = adjacency.shape\n",
    "        adjacency = torch.sparse.FloatTensor(i, v, torch.Size(shape))\n",
    "        item_embeddings = embedding\n",
    "        item_embedding_layer0 = item_embeddings\n",
    "        final = [item_embedding_layer0]\n",
    "        for i in range(self.layers):\n",
    "            item_embeddings = trans_to_cuda(self.w_item['weight_item%d' % (i)])(item_embeddings)\n",
    "            item_embeddings = torch.sparse.mm(trans_to_cuda(adjacency), item_embeddings)\n",
    "            final.append(F.normalize(item_embeddings, dim=-1, p=2))\n",
    "        item_embeddings = np.sum(final, 0)/(self.layers+1)\n",
    "        return item_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class SessConv(nn.Module):\n",
    "    def __init__(self, layers, batch_size, emb_size=100):\n",
    "        super(SessConv, self).__init__()\n",
    "        self.emb_size = emb_size\n",
    "        self.batch_size = batch_size\n",
    "        self.layers = layers[0]\n",
    "        self.w_sess = {}\n",
    "        for i in range(self.layers):\n",
    "            self.w_sess['weight_sess%d' % (i)] = nn.Linear(self.emb_size, self.emb_size, bias=False)\n",
    "\n",
    "    def forward(self, item_embedding, D, A, session_item, session_len):\n",
    "        # zeros = torch.cuda.FloatTensor(1, self.emb_size).fill_(0)\n",
    "        zeros = trans_to_cuda(torch.FloatTensor(1, self.emb_size).fill_(0))\n",
    "        # zeros = torch.zeros([1,self.emb_size])\n",
    "        item_embedding = torch.cat([zeros, item_embedding], 0)\n",
    "        seq_h = []\n",
    "        for i in torch.arange(len(session_item)):\n",
    "            seq_h.append(torch.index_select(item_embedding, 0, session_item[i]))\n",
    "        seq_h1 = trans_to_cuda(torch.tensor([item.cpu().detach().numpy() for item in seq_h]))\n",
    "        session_emb = torch.div(torch.sum(seq_h1, 1), session_len)\n",
    "        session = [session_emb]\n",
    "        DA = torch.mm(D, A).float()\n",
    "        for i in range(self.layers):\n",
    "            session_emb = trans_to_cuda(self.w_sess['weight_sess%d' % (i)])(session_emb)\n",
    "            session_emb = torch.mm(DA, session_emb)\n",
    "            session.append(F.normalize(session_emb, p=2, dim=-1))\n",
    "        sess = trans_to_cuda(torch.tensor([item.cpu().detach().numpy() for item in session]))\n",
    "        session_emb = torch.sum(sess, 0)/(self.layers+1)\n",
    "        return session_emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class COTREC(nn.Module):\n",
    "    def __init__(self, adjacency, n_node, lr, layers, l2, beta,lam,eps, dataset, emb_size=100, batch_size=100):\n",
    "        super(COTREC, self).__init__()\n",
    "        self.emb_size = emb_size\n",
    "        self.batch_size = batch_size\n",
    "        self.n_node = n_node\n",
    "        self.dataset = dataset\n",
    "        self.L2 = l2\n",
    "        self.lr = lr\n",
    "        self.layers = layers\n",
    "        self.beta = beta\n",
    "        self.lam = lam\n",
    "        self.eps = eps\n",
    "        self.K = 10\n",
    "        self.w_k = 10\n",
    "        self.num = 5000\n",
    "        self.adjacency = adjacency\n",
    "        self.embedding = nn.Embedding(self.n_node, self.emb_size)\n",
    "        self.pos_len = 200\n",
    "        if self.dataset == 'retailrocket':\n",
    "            self.pos_len = 300\n",
    "        self.pos_embedding = nn.Embedding(self.pos_len, self.emb_size)\n",
    "        self.ItemGraph = ItemConv(self.layers)\n",
    "        self.SessGraph = SessConv(self.layers, self.batch_size)\n",
    "        self.w_1 = nn.Parameter(torch.Tensor(2 * self.emb_size, self.emb_size))\n",
    "        self.w_2 = nn.Parameter(torch.Tensor(self.emb_size, 1))\n",
    "        self.w_i = nn.Linear(self.emb_size, self.emb_size)\n",
    "        self.w_s = nn.Linear(self.emb_size, self.emb_size)\n",
    "        self.glu1 = nn.Linear(self.emb_size, self.emb_size)\n",
    "        self.glu2 = nn.Linear(self.emb_size, self.emb_size, bias=False)\n",
    "\n",
    "        self.adv_item = trans_to_cuda(torch.FloatTensor(self.n_node, self.emb_size).fill_(0).requires_grad_(True))\n",
    "        self.adv_sess = trans_to_cuda(torch.FloatTensor(self.n_node, self.emb_size).fill_(0).requires_grad_(True))\n",
    "        # self.adv_item = torch.zeros(self.n_node, self.emb_size).requires_grad_(True)\n",
    "        # self.adv_sess = torch.zeros(self.n_node, self.emb_size).requires_grad_(True)\n",
    "        self.loss_function = nn.CrossEntropyLoss()\n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=self.lr)\n",
    "        self.init_parameters()\n",
    "\n",
    "    def init_parameters(self):\n",
    "        stdv = 1.0 / math.sqrt(self.emb_size)\n",
    "        for weight in self.parameters():\n",
    "            weight.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def generate_sess_emb(self, item_embedding, session_item, session_len, reversed_sess_item, mask):\n",
    "        zeros = trans_to_cuda(torch.FloatTensor(1, self.emb_size).fill_(0))\n",
    "        # zeros = torch.zeros(1, self.emb_size)\n",
    "        item_embedding = torch.cat([zeros, item_embedding], 0)\n",
    "        get = lambda i: item_embedding[reversed_sess_item[i]]\n",
    "        seq_h = trans_to_cuda(torch.FloatTensor(self.batch_size, list(reversed_sess_item.shape)[1], self.emb_size).fill_(0))\n",
    "        # seq_h = torch.zeros(self.batch_size, list(reversed_sess_item.shape)[1], self.emb_size)\n",
    "        for i in torch.arange(session_item.shape[0]):\n",
    "            seq_h[i] = get(i)\n",
    "        hs = torch.div(torch.sum(seq_h, 1), session_len)\n",
    "        mask = mask.float().unsqueeze(-1)\n",
    "        len = seq_h.shape[1]\n",
    "        pos_emb = self.pos_embedding.weight[:len]\n",
    "        pos_emb = pos_emb.unsqueeze(0).repeat(self.batch_size, 1, 1)\n",
    "\n",
    "        hs = hs.unsqueeze(-2).repeat(1, len, 1)\n",
    "        nh = torch.matmul(torch.cat([pos_emb, seq_h], -1), self.w_1)\n",
    "        nh = torch.tanh(nh)\n",
    "        nh = torch.sigmoid(self.glu1(nh) + self.glu2(hs))\n",
    "        beta = torch.matmul(nh, self.w_2)\n",
    "        beta = beta * mask\n",
    "        select = torch.sum(beta * seq_h, 1)\n",
    "        return select\n",
    "\n",
    "    def generate_sess_emb_npos(self, item_embedding, session_item, session_len, reversed_sess_item, mask):\n",
    "        zeros = trans_to_cuda(torch.FloatTensor(1, self.emb_size).fill_(0))\n",
    "        # zeros = torch.zeros(1, self.emb_size)\n",
    "        item_embedding = torch.cat([zeros, item_embedding], 0)\n",
    "        get = lambda i: item_embedding[reversed_sess_item[i]]\n",
    "        seq_h = trans_to_cuda(torch.FloatTensor(self.batch_size, list(reversed_sess_item.shape)[1], self.emb_size).fill_(0))\n",
    "        # seq_h = torch.zeros(self.batch_size, list(reversed_sess_item.shape)[1], self.emb_size)\n",
    "        for i in torch.arange(session_item.shape[0]):\n",
    "            seq_h[i] = get(i)\n",
    "        hs = torch.div(torch.sum(seq_h, 1), session_len)\n",
    "        mask = mask.float().unsqueeze(-1)\n",
    "        len = seq_h.shape[1]\n",
    "\n",
    "        hs = hs.unsqueeze(-2).repeat(1, len, 1)\n",
    "        nh = torch.sigmoid(self.glu1(seq_h) + self.glu2(hs))\n",
    "        beta = torch.matmul(nh, self.w_2)\n",
    "        beta = beta * mask\n",
    "        select = torch.sum(beta * seq_h, 1)\n",
    "        return select\n",
    "\n",
    "    def example_predicting(self, item_emb, sess_emb):\n",
    "        x_u = torch.matmul(item_emb, sess_emb)\n",
    "        pos = torch.softmax(x_u, 0)\n",
    "        return pos\n",
    "\n",
    "    def adversarial_item(self, item_emb, tar,sess_emb):\n",
    "        adv_item_emb = item_emb + self.adv_item\n",
    "        score = torch.mm(sess_emb, torch.transpose(adv_item_emb, 1, 0))\n",
    "        loss = self.loss_function(score, tar)\n",
    "        grad = torch.autograd.grad(loss, self.adv_item,retain_graph=True)[0]\n",
    "        adv = grad.detach()\n",
    "        self.adv_item = (F.normalize(adv, p=2,dim=1) * self.eps).requires_grad_(True)\n",
    "\n",
    "    def adversarial_sess(self, item_emb, tar,sess_emb):\n",
    "        adv_item_emb = item_emb + self.adv_sess\n",
    "        score = torch.mm(sess_emb, torch.transpose(adv_item_emb, 1, 0))\n",
    "        loss = self.loss_function(score, tar)\n",
    "        grad = torch.autograd.grad(loss, self.adv_sess,retain_graph=True)[0]\n",
    "        adv = grad.detach()\n",
    "        self.adv_sess = (F.normalize(adv, p=2,dim=1) * self.eps).requires_grad_(True)\n",
    "\n",
    "    def diff(self, score_item, score_sess, score_adv2, score_adv1, diff_mask):\n",
    "        # compute KL(score_item, score_adv2), KL(score_sess, score_adv1)\n",
    "        score_item = F.softmax(score_item, dim=1)\n",
    "        score_sess = F.softmax(score_sess, dim=1)\n",
    "        score_adv2 = F.softmax(score_adv2, dim=1)\n",
    "        score_adv1 = F.softmax(score_adv1, dim=1)\n",
    "        score_item = torch.mul(score_item, diff_mask)\n",
    "        score_sess = torch.mul(score_sess, diff_mask)\n",
    "        score_adv1 = torch.mul(score_adv1, diff_mask)\n",
    "        score_adv2 = torch.mul(score_adv2, diff_mask)\n",
    "\n",
    "        h1 = torch.sum(torch.mul(score_item, torch.log(1e-8 + ((score_item + 1e-8)/(score_adv2 + 1e-8)))))\n",
    "        h2 = torch.sum(torch.mul(score_sess, torch.log(1e-8 + ((score_sess + 1e-8)/(score_adv1 + 1e-8)))))\n",
    "\n",
    "        return h1+h2\n",
    "\n",
    "    def SSL_topk(self, anchor, sess_emb, pos, neg):\n",
    "        def score(x1, x2):\n",
    "            return torch.sum(torch.mul(x1, x2), 2)\n",
    "\n",
    "        anchor = F.normalize(anchor + sess_emb, p=2, dim=-1)\n",
    "        pos = torch.reshape(pos, (self.batch_size, self.K, self.emb_size)) + sess_emb.unsqueeze(1).repeat(1, self.K, 1)\n",
    "        neg = torch.reshape(neg, (self.batch_size, self.K, self.emb_size)) + sess_emb.unsqueeze(1).repeat(1, self.K, 1)\n",
    "        pos_score = score(anchor.unsqueeze(1).repeat(1, self.K, 1), F.normalize(pos, p=2, dim=-1))\n",
    "        neg_score = score(anchor.unsqueeze(1).repeat(1, self.K, 1), F.normalize(neg, p=2, dim=-1))\n",
    "        pos_score = torch.sum(torch.exp(pos_score / 0.2), 1)\n",
    "        neg_score = torch.sum(torch.exp(neg_score / 0.2), 1)\n",
    "        con_loss = -torch.sum(torch.log(pos_score / (pos_score + neg_score)))\n",
    "        return con_loss\n",
    "\n",
    "    def topk_func_random(self, score1,score2, item_emb_I, item_emb_S):\n",
    "        values, pos_ind_I = score1.topk(self.num, dim=0, largest=True, sorted=True)\n",
    "        values, pos_ind_S = score2.topk(self.num, dim=0, largest=True, sorted=True)\n",
    "        pos_emb_I = trans_to_cuda(torch.FloatTensor(self.K, self.batch_size, self.emb_size).fill_(0))\n",
    "        pos_emb_S = trans_to_cuda(torch.FloatTensor(self.K, self.batch_size, self.emb_size).fill_(0))\n",
    "        neg_emb_I = trans_to_cuda(torch.FloatTensor(self.K, self.batch_size, self.emb_size).fill_(0))\n",
    "        neg_emb_S = trans_to_cuda(torch.FloatTensor(self.K, self.batch_size, self.emb_size).fill_(0))\n",
    "        for i in torch.arange(self.K):\n",
    "            pos_emb_S[i] = item_emb_S[pos_ind_I[i]]\n",
    "            pos_emb_I[i] = item_emb_I[pos_ind_S[i]]\n",
    "        random_slices = torch.randint(self.K, self.num, (self.K,))  # choose negative items\n",
    "        for i in torch.arange(self.K):\n",
    "            neg_emb_S[i] = item_emb_S[pos_ind_I[random_slices[i]]]\n",
    "            neg_emb_I[i] = item_emb_I[pos_ind_S[random_slices[i]]]\n",
    "        return pos_emb_I, neg_emb_I, pos_emb_S, neg_emb_S\n",
    "\n",
    "    def forward(self, session_item, session_len, D, A, reversed_sess_item, mask, epoch, tar, train, diff_mask):\n",
    "        if train:\n",
    "            item_embeddings_i = self.ItemGraph(self.adjacency, self.embedding.weight)\n",
    "            if self.dataset == 'Tmall':\n",
    "                # for Tmall dataset, we do not use position embedding to learn temporal order\n",
    "                sess_emb_i = self.generate_sess_emb_npos(item_embeddings_i, session_item, session_len,reversed_sess_item, mask)\n",
    "            else:\n",
    "                sess_emb_i = self.generate_sess_emb(item_embeddings_i, session_item, session_len, reversed_sess_item, mask)\n",
    "            sess_emb_i = self.w_k * F.normalize(sess_emb_i, dim=-1, p=2)\n",
    "            item_embeddings_i = F.normalize(item_embeddings_i, dim=-1, p=2)\n",
    "            scores_item = torch.mm(sess_emb_i, torch.transpose(item_embeddings_i, 1, 0))\n",
    "            loss_item = self.loss_function(scores_item, tar)\n",
    "\n",
    "            sess_emb_s = self.SessGraph(self.embedding.weight, D, A, session_item, session_len)\n",
    "            scores_sess = torch.mm(sess_emb_s, torch.transpose(item_embeddings_i, 1, 0))\n",
    "            # compute probability of items to be positive examples\n",
    "            pos_prob_I = self.example_predicting(item_embeddings_i, sess_emb_i)\n",
    "            pos_prob_S = self.example_predicting(self.embedding.weight, sess_emb_s)\n",
    "\n",
    "            # choose top-10 items as positive samples and randomly choose 10 items as negative and get their embedding\n",
    "            pos_emb_I, neg_emb_I, pos_emb_S, neg_emb_S = self.topk_func_random(pos_prob_I,pos_prob_S, item_embeddings_i, self.embedding.weight)\n",
    "\n",
    "            last_item = torch.squeeze(reversed_sess_item[:, 0])\n",
    "            last_item = last_item - 1\n",
    "            last = item_embeddings_i.index_select(0, last_item)\n",
    "            con_loss = self.SSL_topk(last, sess_emb_i, pos_emb_I, neg_emb_I)\n",
    "            last = self.embedding(last_item)\n",
    "            con_loss += self.SSL_topk(last, sess_emb_s, pos_emb_S, neg_emb_S)\n",
    "\n",
    "            # compute and update adversarial examples\n",
    "            self.adversarial_item(item_embeddings_i, tar, sess_emb_i)\n",
    "            self.adversarial_sess(item_embeddings_i, tar, sess_emb_s)\n",
    "\n",
    "            adv_emb_item = item_embeddings_i + self.adv_item\n",
    "            adv_emb_sess = item_embeddings_i + self.adv_sess\n",
    "\n",
    "            score_adv1 = torch.mm(sess_emb_s, torch.transpose(adv_emb_item, 1, 0))\n",
    "            score_adv2 = torch.mm(sess_emb_i, torch.transpose(adv_emb_sess, 1, 0))\n",
    "            # add difference constraint\n",
    "            loss_diff = self.diff(scores_item, scores_sess, score_adv2, score_adv1, diff_mask)\n",
    "        else:\n",
    "            item_embeddings_i = self.ItemGraph(self.adjacency, self.embedding.weight)\n",
    "            if self.dataset == 'Tmall':\n",
    "                sess_emb_i = self.generate_sess_emb_npos(item_embeddings_i, session_item, session_len, reversed_sess_item, mask)\n",
    "            else:\n",
    "                sess_emb_i = self.generate_sess_emb(item_embeddings_i, session_item, session_len, reversed_sess_item, mask)\n",
    "            sess_emb_i = self.w_k * F.normalize(sess_emb_i, dim=-1, p=2)\n",
    "            item_embeddings_i = F.normalize(item_embeddings_i, dim=-1, p=2)\n",
    "            scores_item = torch.mm(sess_emb_i, torch.transpose(item_embeddings_i, 1, 0))\n",
    "            loss_item = self.loss_function(scores_item, tar)\n",
    "            loss_diff = 0\n",
    "            con_loss = 0\n",
    "        return self.beta * con_loss, loss_item, scores_item, loss_diff*self.lam"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://arxiv.org/pdf/2108.10560v1.pdf\n",
    "2. https://github.com/RecoHut-Stanzas/S969915/blob/main/reports/S969915_report.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-30 07:12:08\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",
      "sys    : 3.7.12 (default, Sep 10 2021, 00:21:48) \n",
      "[GCC 7.5.0]\n",
      "numpy  : 1.19.5\n",
      "torch  : 1.10.0+cu111\n",
      "IPython: 5.5.0\n",
      "csv    : 1.0\n",
      "recohut: 0.0.8\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
}
