{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.afn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AFN\n",
    "> A pytorch implementation of AFN."
   ]
  },
  {
   "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 math\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from recohut.models.layers.common import FeaturesEmbedding, FeaturesLinear, MultiLayerPerceptron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class LNN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of LNN layer\n",
    "    Input shape\n",
    "        - A 3D tensor with shape: ``(batch_size,field_size,embedding_size)``.\n",
    "    Output shape\n",
    "        - 2D tensor with shape:``(batch_size,LNN_dim*embedding_size)``.\n",
    "    Arguments\n",
    "        - **in_features** : Embedding of feature.\n",
    "        - **num_fields**: int.The field size of feature.\n",
    "        - **LNN_dim**: int.The number of Logarithmic neuron.\n",
    "        - **bias**: bool.Whether or not use bias in LNN.\n",
    "    \"\"\"\n",
    "    def __init__(self, num_fields, embed_dim, LNN_dim, bias=False):\n",
    "        super(LNN, self).__init__()\n",
    "        self.num_fields = num_fields\n",
    "        self.embed_dim = embed_dim\n",
    "        self.LNN_dim = LNN_dim\n",
    "        self.lnn_output_dim = LNN_dim * embed_dim\n",
    "        self.weight = torch.nn.Parameter(torch.Tensor(LNN_dim, num_fields))\n",
    "        if bias:\n",
    "            self.bias = torch.nn.Parameter(torch.Tensor(LNN_dim, embed_dim))\n",
    "        else:\n",
    "            self.register_parameter('bias', None)\n",
    "        self.reset_parameters()\n",
    "    \n",
    "    def reset_parameters(self):\n",
    "        stdv = 1. / math.sqrt(self.weight.size(1))\n",
    "        self.weight.data.uniform_(-stdv, stdv)\n",
    "        if self.bias is not None:\n",
    "            self.bias.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields, embedding_size)``\n",
    "        \"\"\"\n",
    "        embed_x_abs = torch.abs(x) # Computes the element-wise absolute value of the given input tensor.\n",
    "        embed_x_afn = torch.add(embed_x_abs, 1e-7)\n",
    "        # Logarithmic Transformation\n",
    "        embed_x_log = torch.log1p(embed_x_afn) # torch.log1p and torch.expm1\n",
    "        lnn_out = torch.matmul(self.weight, embed_x_log)\n",
    "        if self.bias is not None:\n",
    "            lnn_out += self.bias\n",
    "        lnn_exp = torch.expm1(lnn_out)\n",
    "        output = F.relu(lnn_exp).contiguous().view(-1, self.lnn_output_dim)\n",
    "        return output\n",
    "\n",
    "\n",
    "class AFN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of AFN.\n",
    "    Reference:\n",
    "        Cheng W, et al. Adaptive Factorization Network: Learning Adaptive-Order Feature Interactions, 2019.\n",
    "    \"\"\"\n",
    "    def __init__(self, field_dims, embed_dim, LNN_dim, mlp_dims, dropouts):\n",
    "        super().__init__()\n",
    "        self.num_fields = len(field_dims)\n",
    "        self.linear = FeaturesLinear(field_dims)    # Linear\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)   # Embedding\n",
    "        self.LNN_dim = LNN_dim\n",
    "        self.LNN_output_dim = self.LNN_dim * embed_dim\n",
    "        self.LNN = LNN(self.num_fields, embed_dim, LNN_dim)\n",
    "        self.mlp = MultiLayerPerceptron(self.LNN_output_dim, mlp_dims, dropouts[0])\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",
    "        lnn_out = self.LNN(embed_x)\n",
    "        x = self.linear(x) + self.mlp(lnn_out)\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- Cheng W, et al. Adaptive Factorization Network: Learning Adaptive-Order Feature Interactions, 2019.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/afn.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
}
