{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.nfm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NFM\n",
    "> A pytorch implementation of Neural Factorization Machine.\n",
    "\n",
    "NFM seamlessly combines the linearity of FM in modelling second-order feature interactions and the non-linearity of neural network in modelling higher-order feature interactions. Conceptually, NFM is more expressive than FM since FM can be seen as a special case of NFM without hidden layers.\n",
    "\n",
    "![Untitled](https://github.com/RecoHut-Stanzas/S021355/raw/main/images/img10.png)"
   ]
  },
  {
   "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 NFM(PointModel):\n",
    "\n",
    "    def __init__(self, n_users, n_items, embedding_dim, batch_norm=True, dropout=0.1, num_layers=3, act_function='relu'):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            n_users : int, the number of users\n",
    "            n_items : int, the number of items\n",
    "            embedding_dim : int, the number of latent factor\n",
    "            act_function : str, activation function for hidden layer\n",
    "            num_layers : int, number of hidden layers\n",
    "            batch_norm : bool, whether to normalize a batch of data\n",
    "            dropout : float, dropout rate\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "\n",
    "        self.num_layers = num_layers\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",
    "        mlp_modules = []\n",
    "        in_dim = embedding_dim\n",
    "        for _ in range(num_layers):  # dim\n",
    "            out_dim = in_dim # dim\n",
    "            mlp_modules.append(nn.Linear(in_dim, out_dim))\n",
    "            in_dim = out_dim\n",
    "            if batch_norm:\n",
    "                mlp_modules.append(nn.BatchNorm1d(out_dim))\n",
    "            if act_function == 'relu':\n",
    "                mlp_modules.append(nn.ReLU())\n",
    "            elif act_function == 'sigmoid':\n",
    "                mlp_modules.append(nn.Sigmoid())\n",
    "            elif act_function == 'tanh':\n",
    "                mlp_modules.append(nn.Tanh())\n",
    "            mlp_modules.append(nn.Dropout(dropout))\n",
    "        self.deep_layers = nn.Sequential(*mlp_modules)\n",
    "        predict_size = embedding_dim  # layers[-1] if layers else embedding_dim\n",
    "\n",
    "        self.pred = nn.Linear(predict_size, 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.constant_(self.user_bias.weight, 0.0)\n",
    "        nn.init.constant_(self.item_bias.weight, 0.0)\n",
    "\n",
    "        # for deep layers\n",
    "        if self.num_layers > 0:  # len(self.layers)\n",
    "            for m in self.deep_layers:\n",
    "                if isinstance(m, nn.Linear):\n",
    "                    nn.init.xavier_normal_(m.weight)\n",
    "            nn.init.xavier_normal_(self.pred.weight)\n",
    "        else:\n",
    "            nn.init.constant_(self.pred.weight, 1.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",
    "        if self.num_layers:\n",
    "            fm = self.deep_layers(fm)\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": "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 | deep_layers    | Sequential | 1.4 K \n",
      "6 | pred           | Linear     | 20    \n",
      "----------------------------------------------\n",
      "193 K     Trainable params\n",
      "0         Non-trainable params\n",
      "193 K     Total params\n",
      "0.775     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a6f0a13fabd540dcb5a13d18b3b13086",
       "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": "61c8b7f4d710446ca973f6d95a7d0ce0",
       "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": "02f5d0db3388436ca608ed57fb161c06",
       "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.0468),\n",
      "                  'hr': tensor(0.1622),\n",
      "                  'loss': tensor(321.0711),\n",
      "                  'ncdg': tensor(0.0732)}}\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'Test Metrics': {'apak': tensor(0.0468),\n",
       "   'hr': tensor(0.1622),\n",
       "   'loss': tensor(321.0711),\n",
       "   'ncdg': tensor(0.0732)}}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = NFM(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",
    "- X He and TS Chua, Neural Factorization Machines for Sparse Predictive Analytics, 2017.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/nfm.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 FactorizationMachine(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, reduce_sum=True):\n",
    "        super().__init__()\n",
    "        self.reduce_sum = reduce_sum\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Float tensor of size ``(batch_size, num_fields, embed_dim)``\n",
    "        \"\"\"\n",
    "        square_of_sum = torch.sum(x, dim=1) ** 2\n",
    "        sum_of_square = torch.sum(x ** 2, dim=1)\n",
    "        ix = square_of_sum - sum_of_square\n",
    "        if self.reduce_sum:\n",
    "            ix = torch.sum(ix, dim=1, keepdim=True)\n",
    "        return 0.5 * ix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NFMv2(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Neural Factorization Machine.\n",
    "    Reference:\n",
    "        X He and TS Chua, Neural Factorization Machines for Sparse Predictive Analytics, 2017.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, mlp_dims, dropouts):\n",
    "        super().__init__()\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        self.fm = torch.nn.Sequential(\n",
    "            FactorizationMachine(reduce_sum=False),\n",
    "            torch.nn.BatchNorm1d(embed_dim),\n",
    "            torch.nn.Dropout(dropouts[0])\n",
    "        )\n",
    "        self.mlp = MultiLayerPerceptron(embed_dim, mlp_dims, dropouts[1])\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        cross_term = self.fm(self.embedding(x))\n",
    "        x = self.linear(x) + self.mlp(cross_term)\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
}
