{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp models.nmf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NMF\n",
    "> Neural Matrix Factorization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import torch\n",
    "from torch import nn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NMF(nn.Module):\n",
    "    def __init__(self, args, num_users, num_items):\n",
    "        super(NMF, self).__init__()\n",
    "        self.num_users = num_users\n",
    "        self.num_items = num_items\n",
    "        self.factor_num_mf = args.factor_num\n",
    "        self.factor_num_mlp =  int(args.layers[0]/2)\n",
    "        self.layers = args.layers\n",
    "        self.dropout = args.dropout\n",
    "\n",
    "        self.embedding_user_mlp = nn.Embedding(num_embeddings=self.num_users, embedding_dim=self.factor_num_mlp)\n",
    "        self.embedding_item_mlp = nn.Embedding(num_embeddings=self.num_items, embedding_dim=self.factor_num_mlp)\n",
    "\n",
    "        self.embedding_user_mf = nn.Embedding(num_embeddings=self.num_users, embedding_dim=self.factor_num_mf)\n",
    "        self.embedding_item_mf = nn.Embedding(num_embeddings=self.num_items, embedding_dim=self.factor_num_mf)\n",
    "\n",
    "        self.fc_layers = nn.ModuleList()\n",
    "        for idx, (in_size, out_size) in enumerate(zip(args.layers[:-1], args.layers[1:])):\n",
    "            self.fc_layers.append(torch.nn.Linear(in_size, out_size))\n",
    "            self.fc_layers.append(nn.ReLU())\n",
    "\n",
    "        self.affine_output = nn.Linear(in_features=args.layers[-1] + self.factor_num_mf, out_features=1)\n",
    "        self.logistic = nn.Sigmoid()\n",
    "        self.init_weight()\n",
    "\n",
    "    def init_weight(self):\n",
    "        nn.init.normal_(self.embedding_user_mlp.weight, std=0.01)\n",
    "        nn.init.normal_(self.embedding_item_mlp.weight, std=0.01)\n",
    "        nn.init.normal_(self.embedding_user_mf.weight, std=0.01)\n",
    "        nn.init.normal_(self.embedding_item_mf.weight, std=0.01)\n",
    "        \n",
    "        for m in self.fc_layers:\n",
    "            if isinstance(m, nn.Linear):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "                \n",
    "        nn.init.xavier_uniform_(self.affine_output.weight)\n",
    "\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear) and m.bias is not None:\n",
    "                m.bias.data.zero_()\n",
    "\n",
    "    def forward(self, user_indices, item_indices):\n",
    "        user_embedding_mlp = self.embedding_user_mlp(user_indices)\n",
    "        item_embedding_mlp = self.embedding_item_mlp(item_indices)\n",
    "\n",
    "        user_embedding_mf = self.embedding_user_mf(user_indices)\n",
    "        item_embedding_mf = self.embedding_item_mf(item_indices)\n",
    "\n",
    "        mlp_vector = torch.cat([user_embedding_mlp, item_embedding_mlp], dim=-1)\n",
    "        mf_vector =torch.mul(user_embedding_mf, item_embedding_mf)\n",
    "\n",
    "        for idx, _ in enumerate(range(len(self.fc_layers))):\n",
    "            mlp_vector = self.fc_layers[idx](mlp_vector)\n",
    "\n",
    "        vector = torch.cat([mlp_vector, mf_vector], dim=-1)\n",
    "        logits = self.affine_output(vector)\n",
    "        rating = self.logistic(logits)\n",
    "        return rating.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.5041, 0.5000], grad_fn=<SqueezeBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Args:\n",
    "    dropout = 0.2\n",
    "    factor_num = 4\n",
    "    layers = [8,4,2]\n",
    "args = Args()\n",
    "\n",
    "model = NMF(args, num_users=5, num_items=5)\n",
    "model.forward(torch.tensor([0,1]), torch.tensor([1,3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-26 04:53:14\n",
      "\n",
      "recohut: 0.0.7\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.144+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "torch  : 1.10.0+cu111\n",
      "IPython: 5.5.0\n",
      "\n"
     ]
    }
   ],
   "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
}
