{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.afm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AFM\n",
    "> A pytorch implementation of Attentional Factorization Machines (AFM).\n",
    "\n",
    "Improves FM by discriminating the importance of different feature interactions. It learns the importance of each feature interaction from data via a neural attention network. Empirically, it is shown on regression task AFM betters FM with a 8.6% relative improvement, and consistently outperforms the state-of-the-art deep learning methods Wide&Deep and DeepCross with a much simpler structure and fewer model parameters.\n",
    "\n",
    "![Untitled](https://github.com/RecoHut-Stanzas/S021355/raw/main/images/img11.png)\n",
    "\n",
    "Formally, the AFM model can be defined as:\n",
    "\n",
    "$$\\hat{y}_{AFM} (x) = w_0 + \\sum_{i=1}^nw_ix_i + p^T\\sum_{i=1}^n\\sum_{j=i+1}^na_{ij}(v_i\\odot v_j)x_ix_j$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## V1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Any, Iterable, List, Optional, Tuple, Union, Callable\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from recohut.models.bases.common import PointModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AFM(PointModel):\n",
    "\n",
    "    def __init__(self, n_users, n_items, embedding_dim, batch_norm=True, dropout=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        self.user_embedding = nn.Embedding(\n",
    "            num_embeddings=n_users, embedding_dim=embedding_dim\n",
    "        )\n",
    "        self.item_embedding = nn.Embedding(\n",
    "            num_embeddings=n_items, embedding_dim=embedding_dim\n",
    "        )\n",
    "        self.user_bias = nn.Embedding(n_users, 1)\n",
    "        self.item_bias = nn.Embedding(n_items, 1)\n",
    "        self.bias_ = nn.Parameter(torch.tensor([0.0]))\n",
    "\n",
    "        fm_modules = []\n",
    "        if batch_norm:\n",
    "            fm_modules.append(nn.BatchNorm1d(embedding_dim))\n",
    "        fm_modules.append(nn.Dropout(dropout))\n",
    "        self.fm_layers = nn.Sequential(*fm_modules)\n",
    "\n",
    "        # consider attention score layer dimension should be (embedding_dim, num_features)\n",
    "        # here we only consider 2 features, user & item, then K=2\n",
    "        K = 2   # num_features\n",
    "        self.lin = nn.Linear(embedding_dim, K)\n",
    "        self.h = nn.Parameter(torch.rand(K, 1))\n",
    "\n",
    "        # final prediction for reducer sum\n",
    "        self.pred = nn.Linear(embedding_dim, 1, bias=False)\n",
    "\n",
    "        self._init_weights()\n",
    "\n",
    "    def _init_weights(self):\n",
    "        nn.init.normal_(self.item_embedding.weight, std=0.01)\n",
    "        nn.init.normal_(self.user_embedding.weight, std=0.01)\n",
    "        nn.init.normal_(self.lin.weight, std=0.01)\n",
    "        nn.init.xavier_normal_(self.pred.weight)\n",
    "\n",
    "        nn.init.constant_(self.user_bias.weight, 0.0)\n",
    "        nn.init.constant_(self.item_bias.weight, 0.0)\n",
    "\n",
    "    def forward(self, users, items):\n",
    "        embed_user = self.user_embedding(users)\n",
    "        embed_item = self.item_embedding(items)\n",
    "\n",
    "        fm = embed_user * embed_item\n",
    "        fm = self.fm_layers(fm)\n",
    "\n",
    "        ''' attention part '''\n",
    "        att = F.relu(self.lin(fm)).mm(self.h)\n",
    "        fm = fm * att\n",
    "\n",
    "        fm = fm + self.user_bias(users) + self.item_bias(items) + self.bias_\n",
    "        pred = self.pred(fm)\n",
    "\n",
    "        return pred.view(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from recohut.datasets.movielens import ML1mDataModule\n",
    "from recohut.trainers.pl_trainer import pl_trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self):\n",
    "        self.data_dir = '/content/data'\n",
    "        self.min_rating = 4\n",
    "        self.num_negative_samples = 99\n",
    "        self.min_uc = 5\n",
    "        self.min_sc = 5\n",
    "\n",
    "        self.log_dir = '/content/logs'\n",
    "        self.model_dir = '/content/models'\n",
    "\n",
    "        self.val_p = 0.2\n",
    "        self.test_p = 0.2\n",
    "        self.num_workers = 2\n",
    "        self.normalize = False\n",
    "        self.batch_size = 32\n",
    "        self.seed = 42\n",
    "        self.shuffle = True\n",
    "        self.pin_memory = True\n",
    "        self.drop_last = False\n",
    "        self.split_type = 'stratified'\n",
    "\n",
    "        self.embedding_dim = 20\n",
    "        self.max_epochs = 5\n",
    "\n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = ML1mDataModule(**args.__dict__)\n",
    "\n",
    "ds.prepare_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "\n",
      "  | Name           | Type       | Params\n",
      "----------------------------------------------\n",
      "0 | user_embedding | Embedding  | 120 K \n",
      "1 | item_embedding | Embedding  | 62.5 K\n",
      "2 | user_bias      | Embedding  | 6.0 K \n",
      "3 | item_bias      | Embedding  | 3.1 K \n",
      "4 | fm_layers      | Sequential | 40    \n",
      "5 | lin            | Linear     | 42    \n",
      "6 | pred           | Linear     | 20    \n",
      "----------------------------------------------\n",
      "192 K     Trainable params\n",
      "0         Non-trainable params\n",
      "192 K     Total params\n",
      "0.770     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "921f52bac98c4ffcbb5678032494e5ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "980933afbdf84a1d8dc57266d958ea06",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validating: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ed6103a15c3746239d1c402d4ac9c2a0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{'Test Metrics': {'apak': tensor(0.1085),\n",
      "                  'hr': tensor(0.2565),\n",
      "                  'loss': tensor(214.0854),\n",
      "                  'ncdg': tensor(0.1429)}}\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'Test Metrics': {'apak': tensor(0.1085),\n",
       "   'hr': tensor(0.2565),\n",
       "   'loss': tensor(214.0854),\n",
       "   'ncdg': tensor(0.1429)}}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = AFM(n_items=ds.data.num_items, n_users=ds.data.num_users, embedding_dim=args.embedding_dim)\n",
    "\n",
    "pl_trainer(model, ds, max_epochs=args.max_epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## V2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- J Xiao, et al. Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks, 2017.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/afm.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "\n",
    "from recohut.layers.common import FeaturesEmbedding, FeaturesLinear, MultiLayerPerceptron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class AttentionalFactorizationMachine(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, attn_size, dropouts):\n",
    "        super().__init__()\n",
    "        self.attention = torch.nn.Linear(embed_dim, attn_size)\n",
    "        self.projection = torch.nn.Linear(attn_size, 1)\n",
    "        self.fc = torch.nn.Linear(embed_dim, 1)\n",
    "        self.dropouts = dropouts\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``\n",
    "        \"\"\"\n",
    "        num_fields = x.shape[1]\n",
    "        row, col = list(), list()\n",
    "        for i in range(num_fields - 1):\n",
    "            for j in range(i + 1, num_fields):\n",
    "                row.append(i), col.append(j)\n",
    "        p, q = x[:, row], x[:, col]\n",
    "        inner_product = p * q\n",
    "        attn_scores = F.relu(self.attention(inner_product))\n",
    "        attn_scores = F.softmax(self.projection(attn_scores), dim=1)\n",
    "        attn_scores = F.dropout(attn_scores, p=self.dropouts[0], training=self.training)\n",
    "        attn_output = torch.sum(attn_scores * inner_product, dim=1)\n",
    "        attn_output = F.dropout(attn_output, p=self.dropouts[1], training=self.training)\n",
    "        return self.fc(attn_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class AFMv2(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Attentional Factorization Machine.\n",
    "    Reference:\n",
    "        J Xiao, et al. Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks, 2017.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, attn_size, dropouts):\n",
    "        super().__init__()\n",
    "        self.num_fields = len(field_dims)\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        self.afm = AttentionalFactorizationMachine(embed_dim, attn_size, dropouts)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        x = self.linear(x) + self.afm(self.embedding(x))\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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
}
