{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.groupim"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# GroupIM\n",
    "> Implementation of GroupIM recommendation model - A Mutual Information Maximization Framework for Neural Group Recommendation.\n",
    "\n",
    "Group interactions are sparse in nature which makes it difficult to provide relevant recommendation to the group. GroupIM regularize the user-group latent space to overcome group interaction sparsity by: maximizing mutual information between representations of groups and group members; and dynamically prioritizing the preferences of highly informative members through contextual preference weighting."
   ]
  },
  {
   "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 torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class Encoder(nn.Module):\n",
    "    \"\"\" User Preference Encoder implemented as fully connected layers over binary bag-of-words vector\n",
    "    (over item set) per user \"\"\"\n",
    "\n",
    "    def __init__(self, n_items, user_layers, embedding_dim, drop_ratio):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.n_items = n_items\n",
    "        self.embedding_dim = embedding_dim\n",
    "        self.drop = nn.Dropout(drop_ratio)\n",
    "        self.user_preference_encoder = torch.nn.ModuleList()  # user individual preference encoder layers.\n",
    "\n",
    "        for idx, (in_size, out_size) in enumerate(zip([self.n_items] + user_layers[:-1], user_layers)):\n",
    "            layer = torch.nn.Linear(in_size, out_size, bias=True)\n",
    "            nn.init.xavier_uniform_(layer.weight)\n",
    "            nn.init.zeros_(layer.bias)\n",
    "            self.user_preference_encoder.append(layer)\n",
    "\n",
    "        self.transform_layer = nn.Linear(self.embedding_dim, self.embedding_dim)\n",
    "        nn.init.xavier_uniform_(self.transform_layer.weight)\n",
    "        nn.init.zeros_(self.transform_layer.bias)\n",
    "\n",
    "        self.user_predictor = nn.Linear(self.embedding_dim, self.n_items, bias=False)  # item embedding for pre-training\n",
    "        nn.init.xavier_uniform_(self.user_predictor.weight)\n",
    "\n",
    "    def pre_train_forward(self, user_items):\n",
    "        \"\"\" user individual preference encoder (excluding final layer) for user-item pre-training\n",
    "            :param user_items: [B, G, I] or [B, I]\n",
    "        \"\"\"\n",
    "        user_items_norm = F.normalize(user_items)  # [B, G, I] or [B, I]\n",
    "        user_pref_embedding = self.drop(user_items_norm)\n",
    "        for idx, _ in enumerate(range(len(self.user_preference_encoder))):\n",
    "            user_pref_embedding = self.user_preference_encoder[idx](user_pref_embedding)  # [B, G, D] or [B, D]\n",
    "            user_pref_embedding = torch.tanh(user_pref_embedding)  # [B, G, D] or [B, D]\n",
    "\n",
    "        logits = self.user_predictor(user_pref_embedding)  # [B, G, D] or [B, D]\n",
    "        return logits, user_pref_embedding\n",
    "\n",
    "    def forward(self, user_items):\n",
    "        \"\"\" user individual preference encoder\n",
    "            :param user_items: [B, G, I]\n",
    "        \"\"\"\n",
    "        _, user_embeds = self.pre_train_forward(user_items)  # [B, G, D]\n",
    "        user_embeds = torch.tanh(self.transform_layer(user_embeds))  # [B, G, D]\n",
    "        return user_embeds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Aggregator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class MaxPoolAggregator(nn.Module):\n",
    "    \"\"\" Group Preference Aggregator implemented as max pooling over group member embeddings \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim, output_dim, drop_ratio=0):\n",
    "        super(MaxPoolAggregator, self).__init__()\n",
    "\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(input_dim, output_dim, bias=True),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(drop_ratio)\n",
    "        )\n",
    "        nn.init.xavier_uniform_(self.mlp[0].weight)\n",
    "        if self.mlp[0].bias is not None:\n",
    "            self.mlp[0].bias.data.fill_(0.0)\n",
    "\n",
    "    def forward(self, x, mask, mlp=False):\n",
    "        \"\"\" max pooling aggregator:\n",
    "            :param x: [B, G, D]  group member embeddings\n",
    "            :param mask: [B, G]  -inf/0 for absent/present\n",
    "            :param mlp: flag to add a linear layer before max pooling\n",
    "        \"\"\"\n",
    "        if mlp:\n",
    "            h = torch.tanh(self.mlp(x))\n",
    "        else:\n",
    "            h = x\n",
    "\n",
    "        if mask is None:\n",
    "            return torch.max(h, dim=1)\n",
    "        else:\n",
    "            res = torch.max(h + mask.unsqueeze(2), dim=1)\n",
    "            return res.values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "# mask:  -inf/0 for absent/present.\n",
    "class MeanPoolAggregator(nn.Module):\n",
    "    \"\"\" Group Preference Aggregator implemented as mean pooling over group member embeddings \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim, output_dim, drop_ratio=0):\n",
    "        super(MeanPoolAggregator, self).__init__()\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(input_dim, output_dim, bias=True),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(drop_ratio)\n",
    "        )\n",
    "        nn.init.xavier_uniform_(self.mlp[0].weight)\n",
    "        if self.mlp[0].bias is not None:\n",
    "            self.mlp[0].bias.data.fill_(0.0)\n",
    "\n",
    "    def forward(self, x, mask, mlp=False):\n",
    "        \"\"\" mean pooling aggregator:\n",
    "            :param x: [B, G, D]  group member embeddings\n",
    "            :param mask: [B, G]  -inf/0 for absent/present\n",
    "            :param mlp: flag to add a linear layer before mean pooling\n",
    "        \"\"\"\n",
    "        if mlp:\n",
    "            h = torch.tanh(self.mlp(x))\n",
    "        else:\n",
    "            h = x\n",
    "        if mask is None:\n",
    "            return torch.mean(h, dim=1)\n",
    "        else:\n",
    "            mask = torch.exp(mask)\n",
    "            res = torch.sum(h * mask.unsqueeze(2), dim=1) / mask.sum(1).unsqueeze(1)\n",
    "            return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class AttentionAggregator(nn.Module):\n",
    "    \"\"\" Group Preference Aggregator implemented as attention over group member embeddings \"\"\"\n",
    "\n",
    "    def __init__(self, input_dim, output_dim, drop_ratio=0):\n",
    "        super(AttentionAggregator, self).__init__()\n",
    "        self.mlp = nn.Sequential(\n",
    "            nn.Linear(input_dim, output_dim, bias=True),\n",
    "            nn.ReLU(),\n",
    "            nn.Dropout(drop_ratio)\n",
    "        )\n",
    "\n",
    "        self.attention = nn.Linear(output_dim, 1)\n",
    "        self.drop = nn.Dropout(drop_ratio)\n",
    "        nn.init.xavier_uniform_(self.mlp[0].weight)\n",
    "        if self.mlp[0].bias is not None:\n",
    "            self.mlp[0].bias.data.fill_(0.0)\n",
    "\n",
    "    def forward(self, x, mask, mlp=False):\n",
    "        \"\"\" attentive aggregator:\n",
    "            :param x: [B, G, D]  group member embeddings\n",
    "            :param mask: [B, G]  -inf/0 for absent/present\n",
    "            :param mlp: flag to add a linear layer before attention\n",
    "        \"\"\"\n",
    "        if mlp:\n",
    "            h = torch.tanh(self.mlp(x))\n",
    "        else:\n",
    "            h = x\n",
    "\n",
    "        attention_out = torch.tanh(self.attention(h))\n",
    "        if mask is None:\n",
    "            weight = torch.softmax(attention_out, dim=1)\n",
    "        else:\n",
    "            weight = torch.softmax(attention_out + mask.unsqueeze(2), dim=1)\n",
    "        ret = torch.matmul(h.transpose(2, 1), weight).squeeze(2)\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class Discriminator(nn.Module):\n",
    "    \"\"\" Discriminator for Mutual Information Estimation and Maximization, implemented with bilinear layers and\n",
    "    binary cross-entropy loss training \"\"\"\n",
    "\n",
    "    def __init__(self, embedding_dim=64):\n",
    "        super(Discriminator, self).__init__()\n",
    "        self.embedding_dim = embedding_dim\n",
    "\n",
    "        self.fc_layer = torch.nn.Linear(self.embedding_dim, self.embedding_dim, bias=True)\n",
    "        nn.init.xavier_uniform_(self.fc_layer.weight)\n",
    "        nn.init.zeros_(self.fc_layer.bias)\n",
    "\n",
    "        self.bilinear_layer = nn.Bilinear(self.embedding_dim, self.embedding_dim, 1)  # output_dim = 1 => single score.\n",
    "        nn.init.zeros_(self.bilinear_layer.weight)\n",
    "        nn.init.zeros_(self.bilinear_layer.bias)\n",
    "\n",
    "        self.bce_loss = nn.BCEWithLogitsLoss()\n",
    "\n",
    "    def forward(self, group_inputs, user_inputs, group_mask):\n",
    "        \"\"\" bilinear discriminator:\n",
    "            :param group_inputs: [B, I]\n",
    "            :param user_inputs: [B, n_samples, I] where n_samples is either G or # negs\n",
    "            :param group_mask: [B, G]\n",
    "        \"\"\"\n",
    "        # FC + activation.\n",
    "        group_encoded = self.fc_layer(group_inputs)  # [B, D]\n",
    "        group_embed = torch.tanh(group_encoded)  # [B, D]\n",
    "\n",
    "        # FC + activation.\n",
    "        user_pref_embedding = self.fc_layer(user_inputs)\n",
    "        user_embed = torch.tanh(user_pref_embedding)  # [B, n_samples, D]\n",
    "\n",
    "        return self.bilinear_layer(user_embed, group_embed.unsqueeze(1).repeat(1, user_inputs.shape[1], 1))\n",
    "\n",
    "    def mi_loss(self, scores_group, group_mask, scores_corrupted, device='cpu'):\n",
    "        \"\"\" binary cross-entropy loss over (group, user) pairs for discriminator training\n",
    "            :param scores_group: [B, G]\n",
    "            :param group_mask: [B, G]\n",
    "            :param scores_corrupted: [B, N]\n",
    "            :param device (cpu/gpu)\n",
    "         \"\"\"\n",
    "        batch_size = scores_group.shape[0]\n",
    "        pos_size, neg_size = scores_group.shape[1], scores_corrupted.shape[1]\n",
    "\n",
    "        one_labels = torch.ones(batch_size, pos_size).to(device)  # [B, G]\n",
    "        zero_labels = torch.zeros(batch_size, neg_size).to(device)  # [B, N]\n",
    "\n",
    "        labels = torch.cat((one_labels, zero_labels), 1)  # [B, G+N]\n",
    "        logits = torch.cat((scores_group, scores_corrupted), 1).squeeze(2)  # [B, G + N]\n",
    "\n",
    "        mask = torch.cat((torch.exp(group_mask), torch.ones([batch_size, neg_size]).to(device)),\n",
    "                         1)  # torch.exp(.) to binarize since original mask has -inf.\n",
    "\n",
    "        mi_loss = self.bce_loss(logits * mask, labels * mask) * (batch_size * (pos_size + neg_size)) \\\n",
    "                  / (torch.exp(group_mask).sum() + batch_size * neg_size)\n",
    "\n",
    "        return mi_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GroupIM Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class GroupIM(nn.Module):\n",
    "    \"\"\"\n",
    "    GroupIM framework for Group Recommendation:\n",
    "    (a) User Preference encoding: user_preference_encoder\n",
    "    (b) Group Aggregator: preference_aggregator\n",
    "    (c) InfoMax Discriminator: discriminator\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, n_items, user_layers, lambda_mi=0.1, drop_ratio=0.4, aggregator_type='attention'):\n",
    "        super(GroupIM, self).__init__()\n",
    "        self.n_items = n_items\n",
    "        self.lambda_mi = lambda_mi\n",
    "        self.drop = nn.Dropout(drop_ratio)\n",
    "        self.embedding_dim = user_layers[-1]\n",
    "        self.aggregator_type = aggregator_type\n",
    "\n",
    "        self.user_preference_encoder = Encoder(self.n_items, user_layers, self.embedding_dim, drop_ratio)\n",
    "\n",
    "        if self.aggregator_type == 'maxpool':\n",
    "            self.preference_aggregator = MaxPoolAggregator(self.embedding_dim, self.embedding_dim)\n",
    "        elif self.aggregator_type == 'meanpool':\n",
    "            self.preference_aggregator = MeanPoolAggregator(self.embedding_dim, self.embedding_dim)\n",
    "        elif self.aggregator_type == 'attention':\n",
    "            self.preference_aggregator = AttentionAggregator(self.embedding_dim, self.embedding_dim)\n",
    "        else:\n",
    "            raise NotImplementedError(\"Aggregator type {} not implemented \".format(self.aggregator_type))\n",
    "\n",
    "        self.group_predictor = nn.Linear(self.embedding_dim, self.n_items, bias=False)\n",
    "        nn.init.xavier_uniform_(self.group_predictor.weight)\n",
    "\n",
    "        self.discriminator = Discriminator(embedding_dim=self.embedding_dim)\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            if isinstance(m, nn.Embedding):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "\n",
    "    def forward(self, group, group_users, group_mask, user_items):\n",
    "        \"\"\" compute group embeddings and item recommendations by user preference encoding, group aggregation and\n",
    "        item prediction\n",
    "        :param group: [B] group id\n",
    "        :param group_users: [B, G] group user ids with padding\n",
    "        :param group_mask: [B, G] -inf/0 for absent/present user\n",
    "        :param user_items: [B, G, I] individual item interactions of group members\n",
    "        \"\"\"\n",
    "        user_pref_embeds = self.user_preference_encoder(user_items)\n",
    "        group_embed = self.preference_aggregator(user_pref_embeds, group_mask, mlp=False)  # [B, D]\n",
    "        group_logits = self.group_predictor(group_embed)  # [B, I]\n",
    "\n",
    "        if self.train:\n",
    "            obs_user_embeds = self.user_preference_encoder(user_items)  # [B, G, D]\n",
    "            scores_ug = self.discriminator(group_embed, obs_user_embeds, group_mask).detach()  # [B, G]\n",
    "            return group_logits, group_embed, scores_ug\n",
    "        else:\n",
    "            return group_logits, group_embed\n",
    "\n",
    "    def multinomial_loss(self, logits, items):\n",
    "        \"\"\" multinomial likelihood with softmax over item set \"\"\"\n",
    "        return -torch.mean(torch.sum(F.log_softmax(logits, 1) * items, -1))\n",
    "\n",
    "    def user_loss(self, user_logits, user_items):\n",
    "        return self.multinomial_loss(user_logits, user_items)\n",
    "\n",
    "    def infomax_group_loss(self, group_logits, group_embeds, scores_ug, group_mask, group_items, user_items,\n",
    "                           corrupted_user_items, device='cpu'):\n",
    "        \"\"\" loss function with three terms: L_G, L_UG, L_MI\n",
    "            :param group_logits: [B, G, I] group item predictions\n",
    "            :param group_embeds: [B, D] group embedding\n",
    "            :param scores_ug: [B, G] discriminator scores for group members\n",
    "            :param group_mask: [B, G] -inf/0 for absent/present user\n",
    "            :param group_items: [B, I] item interactions of group\n",
    "            :param user_items: [B, G, I] individual item interactions of group members\n",
    "            :param corrupted_user_items: [B, N, I] individual item interactions of negative user samples\n",
    "            :param device: cpu/gpu\n",
    "        \"\"\"\n",
    "\n",
    "        group_user_embeds = self.user_preference_encoder(user_items)  # [B, G, D]\n",
    "        corrupt_user_embeds = self.user_preference_encoder(corrupted_user_items)  # [B, N, D]\n",
    "\n",
    "        scores_observed = self.discriminator(group_embeds, group_user_embeds, group_mask)  # [B, G]\n",
    "        scores_corrupted = self.discriminator(group_embeds, corrupt_user_embeds, group_mask)  # [B, N]\n",
    "\n",
    "        mi_loss = self.discriminator.mi_loss(scores_observed, group_mask, scores_corrupted, device=device)\n",
    "\n",
    "        ui_sum = user_items.sum(2, keepdim=True)  # [B, G]\n",
    "        user_items_norm = user_items / torch.max(torch.ones_like(ui_sum), ui_sum)  # [B, G, I]\n",
    "        gi_sum = group_items.sum(1, keepdim=True)\n",
    "        group_items_norm = group_items / torch.max(torch.ones_like(gi_sum), gi_sum)  # [B, I]\n",
    "        assert scores_ug.requires_grad is False\n",
    "\n",
    "        group_mask_zeros = torch.exp(group_mask).unsqueeze(2)  # [B, G, 1]\n",
    "        scores_ug = torch.sigmoid(scores_ug)  # [B, G, 1]\n",
    "\n",
    "        user_items_norm = torch.sum(user_items_norm * scores_ug * group_mask_zeros, dim=1) / group_mask_zeros.sum(1)\n",
    "        user_group_loss = self.multinomial_loss(group_logits, user_items_norm)\n",
    "        group_loss = self.multinomial_loss(group_logits, group_items_norm)\n",
    "\n",
    "        return mi_loss, user_group_loss, group_loss\n",
    "\n",
    "    def loss(self, group_logits, summary_embeds, scores_ug, group_mask, group_items, user_items, corrupted_user_items,\n",
    "             device='cpu'):\n",
    "        \"\"\" L_G + lambda L_UG + L_MI \"\"\"\n",
    "        mi_loss, user_group_loss, group_loss = self.infomax_group_loss(group_logits, summary_embeds, scores_ug,\n",
    "                                                                       group_mask, group_items, user_items,\n",
    "                                                                       corrupted_user_items, device)\n",
    "\n",
    "        return group_loss + mi_loss + self.lambda_mi * user_group_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<bound method Module.parameters of GroupIM(\n",
       "  (drop): Dropout(p=0.4, inplace=False)\n",
       "  (user_preference_encoder): Encoder(\n",
       "    (drop): Dropout(p=0.4, inplace=False)\n",
       "    (user_preference_encoder): ModuleList(\n",
       "      (0): Linear(in_features=10, out_features=5, bias=True)\n",
       "    )\n",
       "    (transform_layer): Linear(in_features=5, out_features=5, bias=True)\n",
       "    (user_predictor): Linear(in_features=5, out_features=10, bias=False)\n",
       "  )\n",
       "  (preference_aggregator): AttentionAggregator(\n",
       "    (mlp): Sequential(\n",
       "      (0): Linear(in_features=5, out_features=5, bias=True)\n",
       "      (1): ReLU()\n",
       "      (2): Dropout(p=0, inplace=False)\n",
       "    )\n",
       "    (attention): Linear(in_features=5, out_features=1, bias=True)\n",
       "    (drop): Dropout(p=0, inplace=False)\n",
       "  )\n",
       "  (group_predictor): Linear(in_features=5, out_features=10, bias=False)\n",
       "  (discriminator): Discriminator(\n",
       "    (fc_layer): Linear(in_features=5, out_features=5, bias=True)\n",
       "    (bilinear_layer): Bilinear(in1_features=5, in2_features=5, out_features=1, bias=True)\n",
       "    (bce_loss): BCEWithLogitsLoss()\n",
       "  )\n",
       ")>"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = GroupIM(n_items=10, user_layers=[5])\n",
    "model.parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> References\n",
    "1. https://github.com/RecoHut-Stanzas/S168471/blob/main/reports/S168471_report.ipynb\n",
    "2. https://arxiv.org/abs/2006.03736"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-29 14:27:33\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",
      "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
}
