{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.pnn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Product-based Neural Network (PNN)\n",
    "> A pytorch implementation of inner/outer Product Neural Network.\n",
    "\n",
    "PNN uses an embedding layer to learn a distributed representation of the categorical data, a product layer to capture interactive patterns between inter-field categories, and further fully connected layers to explore high-order feature interactions.\n",
    "\n",
    "<img src='https://github.com/RecoHut-Stanzas/S021355/raw/main/images/img8.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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "\n",
    "from recohut.models.layers.common import FeaturesEmbedding, FeaturesLinear, MultiLayerPerceptron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class InnerProductNetwork(torch.nn.Module):\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",
    "        return torch.sum(x[:, row] * x[:, col], dim=2)\n",
    "\n",
    "\n",
    "class OuterProductNetwork(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, num_fields, embed_dim, kernel_type='mat'):\n",
    "        super().__init__()\n",
    "        num_ix = num_fields * (num_fields - 1) // 2\n",
    "        if kernel_type == 'mat':\n",
    "            kernel_shape = embed_dim, num_ix, embed_dim\n",
    "        elif kernel_type == 'vec':\n",
    "            kernel_shape = num_ix, embed_dim\n",
    "        elif kernel_type == 'num':\n",
    "            kernel_shape = num_ix, 1\n",
    "        else:\n",
    "            raise ValueError('unknown kernel type: ' + kernel_type)\n",
    "        self.kernel_type = kernel_type\n",
    "        self.kernel = torch.nn.Parameter(torch.zeros(kernel_shape))\n",
    "        torch.nn.init.xavier_uniform_(self.kernel.data)\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",
    "        if self.kernel_type == 'mat':\n",
    "            kp = torch.sum(p.unsqueeze(1) * self.kernel, dim=-1).permute(0, 2, 1)\n",
    "            return torch.sum(kp * q, -1)\n",
    "        else:\n",
    "            return torch.sum(p * q * self.kernel.unsqueeze(0), -1)\n",
    "\n",
    "class PNN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of inner/outer Product Neural Network.\n",
    "    Reference:\n",
    "        Y Qu, et al. Product-based Neural Networks for User Response Prediction, 2016.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, mlp_dims, dropout, method='inner'):\n",
    "        super().__init__()\n",
    "        num_fields = len(field_dims)\n",
    "        if method == 'inner':\n",
    "            self.pn = InnerProductNetwork()\n",
    "        elif method == 'outer':\n",
    "            self.pn = OuterProductNetwork(num_fields, embed_dim)\n",
    "        else:\n",
    "            raise ValueError('unknown product type: ' + method)\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)\n",
    "        self.linear = FeaturesLinear(field_dims, embed_dim)\n",
    "        self.embed_output_dim = num_fields * embed_dim\n",
    "        self.mlp = MultiLayerPerceptron(num_fields * (num_fields - 1) // 2 + self.embed_output_dim, mlp_dims, dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        embed_x = self.embedding(x)\n",
    "        cross_term = self.pn(embed_x)\n",
    "        x = torch.cat([embed_x.view(-1, self.embed_output_dim), cross_term], dim=1)\n",
    "        x = self.mlp(x)\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- Y Qu, et al. Product-based Neural Networks for User Response Prediction, 2016. https://arxiv.org/abs/1611.00144.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/pnn.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
}
