{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.fnfm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FNFM\n",
    "> A pytorch implementation of Field-aware Neural Factorization Machine."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "from recohut.models.layers.common import FeaturesLinear, MultiLayerPerceptron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class FieldAwareFactorizationMachine(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim):\n",
    "        super().__init__()\n",
    "        self.num_fields = len(field_dims)\n",
    "        self.embeddings = torch.nn.ModuleList([\n",
    "            torch.nn.Embedding(sum(field_dims), embed_dim) for _ in range(self.num_fields)\n",
    "        ])\n",
    "        self.offsets = np.array((0, *np.cumsum(field_dims)[:-1]), dtype=np.long)\n",
    "        for embedding in self.embeddings:\n",
    "            torch.nn.init.xavier_uniform_(embedding.weight.data)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        x = x + x.new_tensor(self.offsets).unsqueeze(0)\n",
    "        xs = [self.embeddings[i](x) for i in range(self.num_fields)]\n",
    "        ix = list()\n",
    "        for i in range(self.num_fields - 1):\n",
    "            for j in range(i + 1, self.num_fields):\n",
    "                ix.append(xs[j][:, i] * xs[i][:, j])\n",
    "        ix = torch.stack(ix, dim=1)\n",
    "        return ix\n",
    "\n",
    "class FNFM(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Field-aware Neural Factorization Machine.\n",
    "    Reference:\n",
    "        L Zhang, et al. Field-aware Neural Factorization Machine for Click-Through Rate Prediction, 2019.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, mlp_dims, dropouts):\n",
    "        super().__init__()\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        self.ffm = FieldAwareFactorizationMachine(field_dims, embed_dim)\n",
    "        self.ffm_output_dim = len(field_dims) * (len(field_dims) - 1) // 2 * embed_dim\n",
    "        self.bn = torch.nn.BatchNorm1d(self.ffm_output_dim)\n",
    "        self.dropout = torch.nn.Dropout(dropouts[0])\n",
    "        self.mlp = MultiLayerPerceptron(self.ffm_output_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.ffm(x).view(-1, self.ffm_output_dim)\n",
    "        cross_term = self.bn(cross_term)\n",
    "        cross_term = self.dropout(cross_term)\n",
    "        x = self.linear(x) + self.mlp(cross_term)\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- L Zhang, et al. Field-aware Neural Factorization Machine for Click-Through Rate Prediction, 2019.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/fnfm.py"
   ]
  },
  {
   "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
}
