{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.hofm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# HOFM\n",
    "> A pytorch implementation of Higher-Order Factorization Machines."
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\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\n",
    "\n",
    "class AnovaKernel(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, order, reduce_sum=True):\n",
    "        super().__init__()\n",
    "        self.order = order\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",
    "        batch_size, num_fields, embed_dim = x.shape\n",
    "        a_prev = torch.ones((batch_size, num_fields + 1, embed_dim), dtype=torch.float).to(x.device)\n",
    "        for t in range(self.order):\n",
    "            a = torch.zeros((batch_size, num_fields + 1, embed_dim), dtype=torch.float).to(x.device)\n",
    "            a[:, t+1:, :] += x[:, t:, :] * a_prev[:, t:-1, :]\n",
    "            a = torch.cumsum(a, dim=1)\n",
    "            a_prev = a\n",
    "        if self.reduce_sum:\n",
    "            return torch.sum(a[:, -1, :], dim=-1, keepdim=True)\n",
    "        else:\n",
    "            return a[:, -1, :]\n",
    "\n",
    "class HOFM(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    A pytorch implementation of Higher-Order Factorization Machines.\n",
    "    Reference:\n",
    "        M Blondel, et al. Higher-Order Factorization Machines, 2016.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, field_dims, order, embed_dim):\n",
    "        super().__init__()\n",
    "        if order < 1:\n",
    "            raise ValueError(f'invalid order: {order}')\n",
    "        self.order = order\n",
    "        self.embed_dim = embed_dim\n",
    "        self.linear = FeaturesLinear(field_dims)\n",
    "        if order >= 2:\n",
    "            self.embedding = FeaturesEmbedding(field_dims, embed_dim * (order - 1))\n",
    "            self.fm = FactorizationMachine(reduce_sum=True)\n",
    "        if order >= 3:\n",
    "            self.kernels = torch.nn.ModuleList([\n",
    "                AnovaKernel(order=i, reduce_sum=True) for i in range(3, order + 1)\n",
    "            ])\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        :param x: Long tensor of size ``(batch_size, num_fields)``\n",
    "        \"\"\"\n",
    "        y = self.linear(x).squeeze(1)\n",
    "        if self.order >= 2:\n",
    "            x = self.embedding(x)\n",
    "            x_part = x[:, :, :self.embed_dim]\n",
    "            y += self.fm(x_part).squeeze(1)\n",
    "            for i in range(self.order - 2):\n",
    "                x_part = x[:, :, (i + 1) * self.embed_dim: (i + 2) * self.embed_dim]\n",
    "                y += self.kernels[i](x_part).squeeze(1)\n",
    "        return torch.sigmoid(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- M Blondel, et al. Higher-Order Factorization Machines, 2016.\n",
    "- https://github.com/rixwew/pytorch-fm/blob/master/torchfm/model/hofm.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
}
