{
 "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 *\n",
    "from fastcore.nb_imports import *\n",
    "from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Any, Iterable, List, Optional, Tuple, Union, Callable\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "from recohut.models.bases.common import PointModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NMF(PairModel):\n",
    "    def __init__(self, n_users, n_items, embedding_dim, dropout=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        self.user_embedding = nn.Embedding(\n",
    "            num_embeddings=n_users, embedding_dim=embedding_dim\n",
    "        )\n",
    "        self.item_embedding = nn.Embedding(\n",
    "            num_embeddings=n_items, embedding_dim=embedding_dim\n",
    "        )\n",
    "\n",
    "        self.user_embedding_gmf = nn.Embedding(\n",
    "            num_embeddings=n_users, embedding_dim=embedding_dim\n",
    "        )\n",
    "        self.item_embedding_gmf = nn.Embedding(\n",
    "            num_embeddings=n_items, embedding_dim=embedding_dim\n",
    "        )\n",
    "\n",
    "        self.gmf = nn.Linear(embedding_dim, int(embedding_dim / 2))\n",
    "\n",
    "        self.fc1 = nn.Linear(embedding_dim * 2, embedding_dim)\n",
    "        self.fc2 = nn.Linear(embedding_dim, embedding_dim)\n",
    "        self.fc3 = nn.Linear(embedding_dim, int(embedding_dim / 2))\n",
    "\n",
    "        self.fc_final = nn.Linear(embedding_dim, 1)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, users, items):\n",
    "        user_embeddings = self.user_embedding(users)\n",
    "        item_embeddings = self.item_embedding(items)\n",
    "        embeddings = torch.cat([user_embeddings, item_embeddings], dim=1)\n",
    "\n",
    "        user_embeddings_gmf = self.user_embedding_gmf(users)\n",
    "        item_embeddings_gmf = self.item_embedding_gmf(items)\n",
    "        embeddings_gmf = user_embeddings_gmf.mul(item_embeddings_gmf)\n",
    "\n",
    "        output_gmf = self.gmf(embeddings_gmf)\n",
    "        output = nn.ReLU()(self.fc1(embeddings))\n",
    "        output = self.dropout(output)\n",
    "        output = nn.ReLU()(self.fc2(output))\n",
    "        output = self.dropout(output)\n",
    "        output = self.fc3(output)\n",
    "\n",
    "        output = torch.cat([output, output_gmf], dim=1)\n",
    "        output = self.fc_final(output)\n",
    "\n",
    "        return output.squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4212, -0.6745], grad_fn=<SqueezeBackward0>)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = NMF(n_users=5, n_items=5, embedding_dim=4)\n",
    "model.forward(users=torch.tensor([0,1]), items=torch.tensor([1,3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class NMFv2(nn.Module):\n",
    "    def __init__(self, args, num_users, num_items):\n",
    "        super(NMFv2, 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.5000, 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 = NMFv2(args, num_users=5, num_items=5)\n",
    "model.forward(torch.tensor([0,1]), torch.tensor([1,3]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from recohut.utils.common_utils import *\n",
    "from recohut.datasets.bases.commons.interactions import InteractionsDataset, InteractionsDataModule\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "\n",
    "class ML1mDataset(InteractionsDataset):\n",
    "    url = \"http://files.grouplens.org/datasets/movielens/ml-1m.zip\"\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return 'ratings.dat'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        from shutil import move, rmtree\n",
    "        move(os.path.join(self.raw_dir, 'ml-1m', self.raw_file_names), self.raw_dir)\n",
    "        rmtree(os.path.join(self.raw_dir, 'ml-1m'))\n",
    "        os.unlink(path)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        df = pd.read_csv(self.raw_paths[0], sep='::', header=None, engine='python')\n",
    "        df.columns = ['uid', 'sid', 'rating', 'timestamp']\n",
    "        # drop duplicate user-item pair records, keeping recent ratings only\n",
    "        df.drop_duplicates(subset=['uid', 'sid'], keep='last', inplace=True)\n",
    "        return df\n",
    "\n",
    "\n",
    "class ML1mDataModule(InteractionsDataModule):\n",
    "    dataset_cls = ML1mDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self):\n",
    "        self.data_dir = '/content/data'\n",
    "        self.min_rating = 4\n",
    "        self.num_negative_samples = 99\n",
    "        self.min_uc = 5\n",
    "        self.min_sc = 5\n",
    "\n",
    "        self.log_dir = '/content/logs'\n",
    "        self.model_dir = '/content/models'\n",
    "\n",
    "        self.val_p = 0.2\n",
    "        self.test_p = 0.2\n",
    "        self.num_workers = 2\n",
    "        self.normalize = False\n",
    "        self.batch_size = 32\n",
    "        self.seed = 42\n",
    "        self.shuffle = True\n",
    "        self.pin_memory = True\n",
    "        self.drop_last = False\n",
    "        self.split_type = 'stratified'\n",
    "\n",
    "        self.embedding_dim = 20\n",
    "        self.max_epochs = 5\n",
    "\n",
    "args = Args()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_lightning import Trainer\n",
    "from pytorch_lightning.callbacks import ModelCheckpoint\n",
    "from pytorch_lightning.loggers import TensorBoardLogger\n",
    "\n",
    "ds = ML1mDataModule(**args.__dict__)\n",
    "\n",
    "logger = TensorBoardLogger(\n",
    "    save_dir=args.log_dir,\n",
    ")\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    monitor=\"valid_loss\",\n",
    "    mode=\"min\",\n",
    "    dirpath=args.model_dir,\n",
    "    filename=\"recommender\",\n",
    ")\n",
    "\n",
    "def pl_trainer(model, datamodule):\n",
    "\n",
    "    trainer = Trainer(\n",
    "    max_epochs=args.max_epochs,\n",
    "    logger=logger,\n",
    "    check_val_every_n_epoch=10,\n",
    "    callbacks=[checkpoint_callback],\n",
    "    gpus=None\n",
    "    )\n",
    "\n",
    "    trainer.fit(model, datamodule=datamodule)\n",
    "    test_result = trainer.test(model, datamodule=datamodule)\n",
    "    return test_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Turning into implicit ratings\n",
      "Filtering triplets\n",
      "Densifying index\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds.prepare_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: False, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "\n",
      "  | Name               | Type      | Params\n",
      "-------------------------------------------------\n",
      "0 | user_embedding     | Embedding | 120 K \n",
      "1 | item_embedding     | Embedding | 62.5 K\n",
      "2 | user_embedding_gmf | Embedding | 120 K \n",
      "3 | item_embedding_gmf | Embedding | 62.5 K\n",
      "4 | gmf                | Linear    | 210   \n",
      "5 | fc1                | Linear    | 820   \n",
      "6 | fc2                | Linear    | 420   \n",
      "7 | fc3                | Linear    | 210   \n",
      "8 | fc_final           | Linear    | 21    \n",
      "9 | dropout            | Dropout   | 0     \n",
      "-------------------------------------------------\n",
      "368 K     Trainable params\n",
      "0         Non-trainable params\n",
      "368 K     Total params\n",
      "1.472     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c32cd3c9a61e4f9d8cd1e2ae5329238c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation sanity check: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7984926a8f3e4d33afb9277303aa8cdf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "932b7d90cfc84baa9e01f498ee28426d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "DATALOADER:0 TEST RESULTS\n",
      "{'Test Metrics': {'apak': tensor(0.0392),\n",
      "                  'hr': tensor(0.1308),\n",
      "                  'loss': tensor(0.2274),\n",
      "                  'ncdg': tensor(0.0602)}}\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'Test Metrics': {'apak': tensor(0.0392),\n",
       "   'hr': tensor(0.1308),\n",
       "   'loss': tensor(0.2274),\n",
       "   'ncdg': tensor(0.0602)}}]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = NMF(n_items=ds.data.num_items, n_users=ds.data.num_users, embedding_dim=args.embedding_dim)\n",
    "\n",
    "pl_trainer(model, ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m/content/data\u001b[00m\n",
      "├── [ 11M]  \u001b[01;34mprocessed\u001b[00m\n",
      "│   ├── [2.3M]  data_test_neg.pt\n",
      "│   ├── [ 95K]  data_test_pos.pt\n",
      "│   ├── [6.5M]  data_train.pt\n",
      "│   ├── [2.3M]  data_valid_neg.pt\n",
      "│   └── [ 95K]  data_valid_pos.pt\n",
      "└── [ 23M]  \u001b[01;34mraw\u001b[00m\n",
      "    └── [ 23M]  ratings.dat\n",
      "\n",
      "  35M used in 2 directories, 6 files\n"
     ]
    }
   ],
   "source": [
    "!tree -h --du -C \"{args.data_dir}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2022-01-10 09:09:05\n",
      "\n",
      "recohut          : 0.0.10\n",
      "pytorch_lightning: 1.5.8\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",
      "numpy  : 1.19.5\n",
      "pandas : 1.1.5\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,pytorch_lightning"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
