{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.autoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AutoInt\n",
    "> A pytorch implementation of AutoInt."
   ]
  },
  {
   "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",
    "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 AutoInt(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of AutoInt.\n",
    "    Reference:\n",
    "        W Song, et al. AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks, 2018.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, embed_dim, atten_embed_dim, num_heads, num_layers, mlp_dims, dropouts, has_residual=True):\n",
    "        super().__init__()\n",
    "        self.num_fields = len(field_dims)\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        self.embedding = FeaturesEmbedding(field_dims, embed_dim)\n",
    "        self.atten_embedding = torch.nn.Linear(embed_dim, atten_embed_dim)\n",
    "        self.embed_output_dim = len(field_dims) * embed_dim\n",
    "        self.atten_output_dim = len(field_dims) * atten_embed_dim\n",
    "        self.has_residual = has_residual\n",
    "        self.mlp = MultiLayerPerceptron(self.embed_output_dim, mlp_dims, dropouts[1])\n",
    "        self.self_attns = torch.nn.ModuleList([\n",
    "            torch.nn.MultiheadAttention(atten_embed_dim, num_heads, dropout=dropouts[0]) for _ in range(num_layers)\n",
    "        ])\n",
    "        self.attn_fc = torch.nn.Linear(self.atten_output_dim, 1)\n",
    "        if self.has_residual:\n",
    "            self.V_res_embedding = torch.nn.Linear(embed_dim, atten_embed_dim)\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",
    "        atten_x = self.atten_embedding(embed_x)\n",
    "        cross_term = atten_x.transpose(0, 1)\n",
    "        for self_attn in self.self_attns:\n",
    "            cross_term, _ = self_attn(cross_term, cross_term, cross_term)\n",
    "        cross_term = cross_term.transpose(0, 1)\n",
    "        if self.has_residual:\n",
    "            V_res = self.V_res_embedding(embed_x)\n",
    "            cross_term += V_res\n",
    "        cross_term = F.relu(cross_term).contiguous().view(-1, self.atten_output_dim)\n",
    "        x = self.linear(x) + self.attn_fc(cross_term) + self.mlp(embed_x.view(-1, self.embed_output_dim))\n",
    "        return torch.sigmoid(x.squeeze(1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- W Song, et al. AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks, 2018. https://arxiv.org/abs/1810.11921.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/afi.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
}
