{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp trainers.trainer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trainer\n",
    "> Implementation of PyTorch model trainer."
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "## v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dataset\n",
    "from recohut.datasets.movielens import ML1mRatingDataset\n",
    "\n",
    "# models\n",
    "from recohut.models.afm import AFM\n",
    "from recohut.models.afn import AFN\n",
    "from recohut.models.autoint import AutoInt\n",
    "from recohut.models.dcn import DCN\n",
    "from recohut.models.deepfm import DeepFM\n",
    "from recohut.models.ffm import FFM\n",
    "from recohut.models.fm import FM\n",
    "from recohut.models.fnfm import FNFM\n",
    "from recohut.models.fnn import FNN\n",
    "from recohut.models.hofm import HOFM\n",
    "from recohut.models.lr import LR\n",
    "from recohut.models.ncf import NCF\n",
    "from recohut.models.nfm import NFM\n",
    "from recohut.models.ncf import NCF\n",
    "from recohut.models.pnn import PNN\n",
    "from recohut.models.wide_and_deep import WideAndDeep\n",
    "from recohut.models.xdeepfm import xDeepFM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading http://files.grouplens.org/datasets/movielens/ml-1m.zip\n",
      "Extracting /content/ML1m/raw/ml-1m.zip\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = ML1mRatingDataset(root='/content/ML1m', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import os\n",
    "import tqdm\n",
    "from sklearn.metrics import roc_auc_score\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    def __init__(self,\n",
    "                 dataset='ml_1m',\n",
    "                 model='wide_and_deep'\n",
    "                 ):\n",
    "        self.dataset = dataset\n",
    "        self.model = model\n",
    "        # dataset\n",
    "        if dataset == 'ml_1m':\n",
    "            self.dataset_root = '/content/ML1m'\n",
    "            self.min_uc = 20\n",
    "            self.min_sc = 20\n",
    "\n",
    "        # model training\n",
    "        self.device = 'cpu' # 'cuda:0'\n",
    "        self.num_workers = 2\n",
    "        self.batch_size = 256\n",
    "        self.lr = 0.001\n",
    "        self.weight_decay = 1e-6\n",
    "        self.save_dir = '/content/chkpt'\n",
    "        self.n_epochs = 2\n",
    "        self.dropout = 0.2\n",
    "        self.log_interval = 100\n",
    "\n",
    "        # model architecture\n",
    "        if model == 'wide_and_deep':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'fm':\n",
    "            self.embed_dim = 16\n",
    "        elif model == 'ffm':\n",
    "            self.embed_dim = 4\n",
    "        elif model == 'hofm':\n",
    "            self.embed_dim = 16\n",
    "            self.order = 3\n",
    "        elif model == 'fnn':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'ipnn':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16,)\n",
    "            self.method = 'inner'\n",
    "        elif model == 'opnn':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16,)\n",
    "            self.method = 'outer'\n",
    "        elif model == 'dcn':\n",
    "            self.embed_dim = 16\n",
    "            self.num_layers = 3\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'nfm':\n",
    "            self.embed_dim = 64\n",
    "            self.mlp_dims = (64,)\n",
    "            self.dropouts = (0.2, 0.2)\n",
    "        elif model == 'ncf':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'fnfm':\n",
    "            self.embed_dim = 4\n",
    "            self.mlp_dims = (64,)\n",
    "            self.dropouts = (0.2, 0.2)\n",
    "        elif model == 'deep_fm':\n",
    "            self.embed_dim = 16\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'xdeep_fm':\n",
    "            self.embed_dim = 16\n",
    "            self.cross_layer_sizes = (16, 16)\n",
    "            self.split_half = False\n",
    "            self.mlp_dims = (16, 16)\n",
    "        elif model == 'afm':\n",
    "            self.embed_dim = 16\n",
    "            self.attn_size = 16\n",
    "            self.dropouts = (0.2, 0.2)\n",
    "        elif model == 'autoint':\n",
    "            self.embed_dim = 16\n",
    "            self.atten_embed_dim = 64\n",
    "            self.num_heads = 2\n",
    "            self.num_layers = 3\n",
    "            self.mlp_dims = (400, 400)\n",
    "            self.dropouts = (0, 0, 0)\n",
    "        elif model == 'afn':\n",
    "            self.embed_dim = 16\n",
    "            self.LNN_dim = 1500\n",
    "            self.mlp_dims = (400, 400, 400)\n",
    "            self.dropouts = (0, 0, 0)\n",
    "\n",
    "    def get_dataset(self):\n",
    "        if self.dataset == 'ml_1m':\n",
    "            return ML1mRatingDataset(root = self.dataset_root,\n",
    "                                     min_uc = self.min_uc,\n",
    "                                     min_sc = self.min_sc\n",
    "                                     )\n",
    "    \n",
    "    def get_model(self, field_dims, user_field_idx=None, item_field_idx=None):\n",
    "        if self.model == 'wide_and_deep':\n",
    "            return WideAndDeep(field_dims,\n",
    "                               embed_dim=self.embed_dim,\n",
    "                               mlp_dims = self.mlp_dims,\n",
    "                               dropout = self.dropout\n",
    "                               )\n",
    "        elif self.model == 'fm':\n",
    "            return FM(field_dims,\n",
    "                      embed_dim = self.embed_dim\n",
    "                      )\n",
    "        elif self.model == 'lr':\n",
    "            return LR(field_dims\n",
    "                      )\n",
    "        elif self.model == 'ffm':\n",
    "            return FFM(field_dims,\n",
    "                       embed_dim = self.embed_dim\n",
    "                      )\n",
    "        elif self.model == 'hofm':\n",
    "            return HOFM(field_dims,\n",
    "                        embed_dim = self.embed_dim,\n",
    "                        order = self.order\n",
    "                      )\n",
    "        elif self.model == 'fnn':\n",
    "            return FNN(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropout = self.dropout\n",
    "                      )\n",
    "        elif self.model == 'ipnn':\n",
    "            return PNN(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       method = self.method,\n",
    "                       dropout = self.dropout\n",
    "                      )\n",
    "        elif self.model == 'opnn':\n",
    "            return PNN(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       method = self.method,\n",
    "                       dropout = self.dropout\n",
    "                      )\n",
    "        elif self.model == 'dcn':\n",
    "            return DCN(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       num_layers = self.num_layers,\n",
    "                       dropout = self.dropout,\n",
    "                      )\n",
    "        elif self.model == 'nfm':\n",
    "            return NFM(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropouts = self.dropouts,\n",
    "                      )\n",
    "        elif self.model == 'ncf':\n",
    "            return NCF(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropout = self.dropout,\n",
    "                       user_field_idx=user_field_idx,\n",
    "                       item_field_idx=item_field_idx\n",
    "                      )\n",
    "        elif self.model == 'fnfm':\n",
    "            return FNFM(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropouts = self.dropouts,\n",
    "                      )\n",
    "        elif self.model == 'deep_fm':\n",
    "            return DeepFM(field_dims,\n",
    "                          embed_dim = self.embed_dim,\n",
    "                          mlp_dims = self.mlp_dims,\n",
    "                          dropout = self.dropout,\n",
    "                      )\n",
    "        elif self.model == 'xdeep_fm':\n",
    "            return xDeepFM(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropout = self.dropout,\n",
    "                       cross_layer_sizes = self.cross_layer_sizes,\n",
    "                       split_half = self.split_half,\n",
    "                      )\n",
    "        elif self.model == 'afm':\n",
    "            return AFM(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       dropouts = self.dropouts,\n",
    "                       attn_size = self.attn_size,\n",
    "                      )\n",
    "        elif self.model == 'autoint':\n",
    "            return AutoInt(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropouts = self.dropouts,\n",
    "                       atten_embed_dim = self.atten_embed_dim,\n",
    "                       num_heads = self.num_heads,\n",
    "                       num_layers = self.num_layers,\n",
    "                      )\n",
    "        elif self.model == 'afn':\n",
    "            return AFN(field_dims,\n",
    "                       embed_dim = self.embed_dim,\n",
    "                       mlp_dims = self.mlp_dims,\n",
    "                       dropouts = self.dropouts,\n",
    "                       LNN_dim = self.LNN_dim,\n",
    "                      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EarlyStopper(object):\n",
    "\n",
    "    def __init__(self, num_trials, save_path):\n",
    "        self.num_trials = num_trials\n",
    "        self.trial_counter = 0\n",
    "        self.best_accuracy = 0\n",
    "        self.save_path = save_path\n",
    "\n",
    "    def is_continuable(self, model, accuracy):\n",
    "        if accuracy > self.best_accuracy:\n",
    "            self.best_accuracy = accuracy\n",
    "            self.trial_counter = 0\n",
    "            torch.save(model, self.save_path)\n",
    "            return True\n",
    "        elif self.trial_counter + 1 < self.num_trials:\n",
    "            self.trial_counter += 1\n",
    "            return True\n",
    "        else:\n",
    "            return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trainer:\n",
    "    def __init__(self, args):\n",
    "        device = torch.device(args.device)\n",
    "        # dataset\n",
    "        dataset = args.get_dataset()\n",
    "        # model\n",
    "        model = args.get_model(dataset.field_dims,\n",
    "                               user_field_idx = dataset.user_field_idx,\n",
    "                               item_field_idx = dataset.item_field_idx)\n",
    "        model = model.to(device)\n",
    "        model_name = type(model).__name__\n",
    "        # data split\n",
    "        train_length = int(len(dataset) * 0.8)\n",
    "        valid_length = int(len(dataset) * 0.1)\n",
    "        test_length = len(dataset) - train_length - valid_length\n",
    "        # data loader\n",
    "        train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(\n",
    "            dataset, (train_length, valid_length, test_length))\n",
    "        train_data_loader = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\n",
    "        valid_data_loader = DataLoader(valid_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\n",
    "        test_data_loader = DataLoader(test_dataset, batch_size=args.batch_size, num_workers=args.num_workers)\n",
    "        # handlers\n",
    "        criterion = torch.nn.BCELoss()\n",
    "        optimizer = torch.optim.Adam(params=model.parameters(), lr=args.lr, weight_decay=args.weight_decay)\n",
    "        os.makedirs(args.save_dir, exist_ok=True)\n",
    "        early_stopper = EarlyStopper(num_trials=2, save_path=f'{args.save_dir}/{model_name}.pt')\n",
    "        # # scheduler\n",
    "        # # ref - https://github.com/sparsh-ai/stanza/blob/7961a0a00dc06b9b28b71954b38181d6a87aa803/trainer/bert.py#L36\n",
    "        # import torch.optim as optim\n",
    "        # if args.enable_lr_schedule:\n",
    "        #     if args.enable_lr_warmup:\n",
    "        #         self.lr_scheduler = self.get_linear_schedule_with_warmup(\n",
    "        #             optimizer, args.warmup_steps, len(train_data_loader) * self.n_epochs)\n",
    "        #     else:\n",
    "        #         self.lr_scheduler = optim.lr_scheduler.StepLR(\n",
    "        #             optimizer, step_size=args.decay_step, gamma=args.gamma)\n",
    "        # training\n",
    "        for epoch_i in range(args.n_epochs):\n",
    "            self._train(model, optimizer, train_data_loader, criterion, device)\n",
    "            auc = self._test(model, valid_data_loader, device)\n",
    "            print('epoch:', epoch_i, 'validation: auc:', auc)\n",
    "            if not early_stopper.is_continuable(model, auc):\n",
    "                print(f'validation: best auc: {early_stopper.best_accuracy}')\n",
    "                break\n",
    "        auc = self._test(model, test_data_loader, device)\n",
    "        print(f'test auc: {auc}')\n",
    "\n",
    "    @staticmethod\n",
    "    def _train(model, optimizer, data_loader, criterion, device, log_interval=100):\n",
    "        model.train()\n",
    "        total_loss = 0\n",
    "        tk0 = tqdm.tqdm(data_loader, smoothing=0, mininterval=1.0)\n",
    "        for i, (fields, target) in enumerate(tk0):\n",
    "            fields, target = fields.to(device), target.to(device)\n",
    "            y = model(fields)\n",
    "            loss = criterion(y, target.float())\n",
    "            model.zero_grad()\n",
    "            loss.backward()\n",
    "            # self.clip_gradients(5)\n",
    "            optimizer.step()\n",
    "            # if self.args.enable_lr_schedule:\n",
    "            #     self.lr_scheduler.step()\n",
    "            total_loss += loss.item()\n",
    "            if (i + 1) % log_interval == 0:\n",
    "                tk0.set_postfix(loss=total_loss / log_interval)\n",
    "                total_loss = 0\n",
    "    \n",
    "    @staticmethod\n",
    "    def _test(model, data_loader, device):\n",
    "        model.eval()\n",
    "        targets, predicts = list(), list()\n",
    "        with torch.no_grad():\n",
    "            for fields, target in tqdm.tqdm(data_loader, smoothing=0, mininterval=1.0):\n",
    "                fields, target = fields.to(device), target.to(device)\n",
    "                y = model(fields)\n",
    "                targets.extend(target.tolist())\n",
    "                predicts.extend(y.tolist())\n",
    "        return roc_auc_score(targets, predicts)\n",
    "\n",
    "    # def clip_gradients(self, limit=5):\n",
    "    #     \"\"\"\n",
    "    #     Reference:\n",
    "    #         1. https://github.com/sparsh-ai/stanza/blob/7961a0a00dc06b9b28b71954b38181d6a87aa803/trainer/bert.py#L175\n",
    "    #     \"\"\"\n",
    "    #     for p in self.model.parameters():\n",
    "    #         nn.utils.clip_grad_norm_(p, 5)\n",
    "\n",
    "    # def _create_optimizer(self):\n",
    "    #     args = self.args\n",
    "    #     param_optimizer = list(self.model.named_parameters())\n",
    "    #     no_decay = ['bias', 'layer_norm']\n",
    "    #     optimizer_grouped_parameters = [\n",
    "    #         {\n",
    "    #             'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],\n",
    "    #             'weight_decay': args.weight_decay,\n",
    "    #         },\n",
    "    #         {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0},\n",
    "    #     ]\n",
    "    #     if args.optimizer.lower() == 'adamw':\n",
    "    #         return optim.AdamW(optimizer_grouped_parameters, lr=args.lr, eps=args.adam_epsilon)\n",
    "    #     elif args.optimizer.lower() == 'adam':\n",
    "    #         return optim.Adam(optimizer_grouped_parameters, lr=args.lr, weight_decay=args.weight_decay)\n",
    "    #     elif args.optimizer.lower() == 'sgd':\n",
    "    #         return optim.SGD(optimizer_grouped_parameters, lr=args.lr, weight_decay=args.weight_decay, momentum=args.momentum)\n",
    "    #     else:\n",
    "    #         raise ValueError\n",
    "\n",
    "    # def get_linear_schedule_with_warmup(self, optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):\n",
    "    #     # based on hugging face get_linear_schedule_with_warmup\n",
    "    #     def lr_lambda(current_step: int):\n",
    "    #         if current_step < num_warmup_steps:\n",
    "    #             return float(current_step) / float(max(1, num_warmup_steps))\n",
    "    #         return max(\n",
    "    #             0.0, float(num_training_steps - current_step) / float(max(1, num_training_steps - num_warmup_steps))\n",
    "    #         )\n",
    "\n",
    "    #     return LambdaLR(optimizer, lr_lambda, last_epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n",
      "100%|██████████| 3126/3126 [00:23<00:00, 135.91it/s, loss=0.57]\n",
      "100%|██████████| 391/391 [00:01<00:00, 252.62it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7781601005135064\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:22<00:00, 137.03it/s, loss=0.557]\n",
      "100%|██████████| 391/391 [00:01<00:00, 259.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7842454181872189\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 261.87it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.783773499847308\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:15<00:00, 200.85it/s, loss=0.587]\n",
      "100%|██████████| 391/391 [00:01<00:00, 277.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7511323978391329\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:15<00:00, 203.24it/s, loss=0.542]\n",
      "100%|██████████| 391/391 [00:01<00:00, 286.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7852232398637453\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 286.62it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7851983970544512\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:12<00:00, 243.01it/s, loss=0.713]\n",
      "100%|██████████| 391/391 [00:01<00:00, 290.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.606845663039941\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:12<00:00, 243.68it/s, loss=0.625]\n",
      "100%|██████████| 391/391 [00:01<00:00, 290.60it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.6962495583229628\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 280.21it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.6917994954031111\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:16<00:00, 189.89it/s, loss=0.639]\n",
      "100%|██████████| 391/391 [00:01<00:00, 275.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.6956660360854087\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:16<00:00, 190.43it/s, loss=0.559]\n",
      "100%|██████████| 391/391 [00:01<00:00, 279.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.769259926201433\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 275.84it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7694256825177728\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:23<00:00, 135.66it/s, loss=0.585]\n",
      "100%|██████████| 391/391 [00:01<00:00, 229.47it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7508361070441243\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:23<00:00, 135.52it/s, loss=0.538]\n",
      "100%|██████████| 391/391 [00:01<00:00, 229.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7867336507526798\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 226.59it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7849653473859624\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:21<00:00, 146.97it/s, loss=0.554]\n",
      "100%|██████████| 391/391 [00:01<00:00, 272.83it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7899586086314532\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:21<00:00, 148.00it/s, loss=0.544]\n",
      "100%|██████████| 391/391 [00:01<00:00, 268.07it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7938707366151592\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 267.69it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7935777015287597\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:20<00:00, 151.01it/s, loss=0.55]\n",
      "100%|██████████| 391/391 [00:01<00:00, 253.01it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7901787607777198\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:20<00:00, 151.50it/s, loss=0.536]\n",
      "100%|██████████| 391/391 [00:01<00:00, 258.53it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7958062417181883\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 265.78it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7959379435427811\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:21<00:00, 144.98it/s, loss=0.548]\n",
      "100%|██████████| 391/391 [00:01<00:00, 256.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7943316704845618\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:21<00:00, 145.26it/s, loss=0.53]\n",
      "100%|██████████| 391/391 [00:01<00:00, 252.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.8027591784990165\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 259.79it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.8016146552653354\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:26<00:00, 116.37it/s, loss=0.537]\n",
      "100%|██████████| 391/391 [00:01<00:00, 240.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7898151214837668\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:26<00:00, 116.92it/s, loss=0.527]\n",
      "100%|██████████| 391/391 [00:01<00:00, 239.57it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7955138244674892\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 240.84it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7964998271099959\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:22<00:00, 138.66it/s, loss=0.586]\n",
      "100%|██████████| 391/391 [00:01<00:00, 252.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7631548463451637\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:22<00:00, 137.08it/s, loss=0.551]\n",
      "100%|██████████| 391/391 [00:01<00:00, 251.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7752154803420491\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 252.42it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7727792981788815\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:23<00:00, 132.24it/s, loss=0.554]\n",
      "100%|██████████| 391/391 [00:01<00:00, 248.61it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7876433331502086\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:23<00:00, 132.00it/s, loss=0.543]\n",
      "100%|██████████| 391/391 [00:01<00:00, 249.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7923030405914255\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 257.83it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7930787548185895\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:23<00:00, 133.99it/s, loss=0.61]\n",
      "100%|██████████| 391/391 [00:01<00:00, 250.20it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7376150945998978\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:23<00:00, 135.18it/s, loss=0.583]\n",
      "100%|██████████| 391/391 [00:01<00:00, 246.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7583206065924306\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 245.77it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7594084947700983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:24<00:00, 127.45it/s, loss=0.569]\n",
      "100%|██████████| 391/391 [00:01<00:00, 244.49it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7806048647711028\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:24<00:00, 128.70it/s, loss=0.554]\n",
      "100%|██████████| 391/391 [00:01<00:00, 246.00it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7857091265544482\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 245.62it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7857843263334994\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:30<00:00, 103.68it/s, loss=0.558]\n",
      "100%|██████████| 391/391 [00:01<00:00, 219.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7814674364890849\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:29<00:00, 104.41it/s, loss=0.539]\n",
      "100%|██████████| 391/391 [00:01<00:00, 214.94it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7899837530572655\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 216.12it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7863345464272122\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [00:23<00:00, 133.32it/s, loss=0.606]\n",
      "100%|██████████| 391/391 [00:01<00:00, 244.59it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7590887701790624\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:23<00:00, 134.06it/s, loss=0.576]\n",
      "100%|██████████| 391/391 [00:01<00:00, 247.06it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7820711568875622\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:01<00:00, 247.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7835448236219698\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "models = [\n",
    "          'wide_and_deep',\n",
    "          'fm',\n",
    "          'lr',\n",
    "          'ffm',\n",
    "          'hofm',\n",
    "          'fnn',\n",
    "          'ipnn',\n",
    "          'opnn',\n",
    "          'dcn',\n",
    "          'nfm',\n",
    "          'ncf',\n",
    "          'fnfm',\n",
    "          'deep_fm',\n",
    "          'xdeep_fm',\n",
    "          'afm',\n",
    "        #   'autoint',\n",
    "        #   'afn'\n",
    "          ]\n",
    "\n",
    "for model in models:\n",
    "    args = Args(model=model)\n",
    "    trainer = Trainer(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "Done!\n",
      "100%|██████████| 3126/3126 [00:43<00:00, 72.44it/s, loss=0.551]\n",
      "100%|██████████| 391/391 [00:02<00:00, 171.82it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7838440329134869\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [00:42<00:00, 73.24it/s, loss=0.532]\n",
      "100%|██████████| 391/391 [00:02<00:00, 172.49it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7924653055551055\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:02<00:00, 169.07it/s]\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7935854845577758\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "100%|██████████| 3126/3126 [01:24<00:00, 37.12it/s, loss=0.564]\n",
      "100%|██████████| 391/391 [00:03<00:00, 107.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0 validation: auc: 0.7796980126749351\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3126/3126 [01:23<00:00, 37.50it/s, loss=0.547]\n",
      "100%|██████████| 391/391 [00:03<00:00, 108.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 1 validation: auc: 0.7879478169612124\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [00:03<00:00, 108.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test auc: 0.7893059350190452\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "models = [\n",
    "          'autoint',\n",
    "          'afn'\n",
    "          ]\n",
    "\n",
    "for model in models:\n",
    "    args = Args(model=model)\n",
    "    trainer = Trainer(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m/content/chkpt\u001b[00m\n",
      "├── [669K]  AFM.pt\n",
      "├── [ 39M]  AFN.pt\n",
      "├── [1.5M]  AutoInt.pt\n",
      "├── [640K]  DCN.pt\n",
      "├── [676K]  DeepFM.pt\n",
      "├── [355K]  FFM.pt\n",
      "├── [666K]  FM.pt\n",
      "├── [363K]  FNFM.pt\n",
      "├── [636K]  FNN.pt\n",
      "├── [1.3M]  HOFM.pt\n",
      "├── [ 41K]  LR.pt\n",
      "├── [636K]  NCF.pt\n",
      "├── [2.5M]  NFM.pt\n",
      "├── [1.2M]  PNN.pt\n",
      "├── [676K]  WideAndDeep.pt\n",
      "└── [682K]  xDeepFM.pt\n",
      "\n",
      "  51M used in 0 directories, 16 files\n"
     ]
    }
   ],
   "source": [
    "!tree --du -h -C /content/chkpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## v2\n",
    "\n",
    "**References:-**\n",
    "1. https://nbviewer.org/github/CS-512-Recsys/Recsys/blob/main/nbs/basic_implementation.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "\n",
    "import os\n",
    "import copy\n",
    "import random\n",
    "from pathlib import Path\n",
    "from collections import defaultdict\n",
    "\n",
    "from argparse import Namespace\n",
    "from joblib import dump, load\n",
    "from tqdm import tqdm\n",
    "\n",
    "import wandb\n",
    "from torch.utils.data import DataLoader as dl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RecsysDataset(torch.utils.data.Dataset):\n",
    "\n",
    "    def __init__(self,df,usr_dict=None,mov_dict=None):\n",
    "        self.df = df\n",
    "        self.usr_dict = usr_dict\n",
    "        self.mov_dict = mov_dict\n",
    "\n",
    "    def __getitem__(self,index):\n",
    "        if self.usr_dict and self.mov_dict:\n",
    "            return [self.usr_dict[int(self.df.iloc[index]['user_id'])],self.mov_dict[int(self.df.iloc[index]['movie_id'])]],self.df.iloc[index]['rating']\n",
    "        else:\n",
    "            return [int(self.df.iloc[index]['user_id']-1),int(self.df.iloc[index]['movie_id']-1)],self.df.iloc[index]['rating']\n",
    "            \n",
    "    def __len__(self):\n",
    "        return len(self.df)\n",
    "\n",
    "\n",
    "sample = pd.DataFrame({'user_id':[1,2,3,2,2,3,2,2],'movie_id':[1,2,3,3,3,2,1,1],'rating':[2.0,1.0,4.0,5.0,1.3,3.5,3.0,4.5]})\n",
    "trn_ids = random.sample(range(8),4,)\n",
    "valid_ids = [i for i in range(8) if i not in trn_ids]\n",
    "\n",
    "sample_trn,sample_vld = copy.deepcopy(sample.iloc[trn_ids].reset_index()),copy.deepcopy(sample.iloc[valid_ids].reset_index())\n",
    "\n",
    "sample_vld = RecsysDataset(sample_vld)\n",
    "sample_trn = RecsysDataset(sample_trn)\n",
    "\n",
    "train_loader = dl(sample_trn, batch_size=2, shuffle=True)\n",
    "valid_loader = dl(sample_vld, batch_size=2, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NCF(nn.Module):\n",
    "    \n",
    "    def __init__(self,user_sz,item_sz,embd_sz,dropout_fac,min_r=0.0,max_r=5.0,alpha=0.5,with_variable_alpha=False):\n",
    "        super().__init__()\n",
    "        self.dropout_fac = dropout_fac\n",
    "        self.user_embd_mtrx = nn.Embedding(user_sz,embd_sz)\n",
    "        self.item_embd_mtrx = nn.Embedding(item_sz,embd_sz)\n",
    "        #bias = torch.zeros(size=(user_sz, 1), requires_grad=True)\n",
    "        self.h =  nn.Linear(embd_sz,1)\n",
    "        self.fst_lyr = nn.Linear(embd_sz*2,embd_sz)\n",
    "        self.snd_lyr = nn.Linear(embd_sz,embd_sz//2)\n",
    "        self.thrd_lyr = nn.Linear(embd_sz//2,embd_sz//4)\n",
    "        self.out_lyr = nn.Linear(embd_sz//4,1)\n",
    "        self.alpha = torch.tensor(alpha)\n",
    "        self.min_r,self.max_r = min_r,max_r\n",
    "        if with_variable_alpha:\n",
    "            self.alpha = torch.tensor(alpha,requires_grad=True)\n",
    "        \n",
    "    def forward(self,x):\n",
    "        user_emd = self.user_embd_mtrx(x[0])\n",
    "        item_emd = self.item_embd_mtrx(x[-1])\n",
    "        #hadamard-product\n",
    "        gmf = user_emd*item_emd\n",
    "        gmf = self.h(gmf)\n",
    "        \n",
    "        \n",
    "        mlp = torch.cat([user_emd,item_emd],dim=-1)\n",
    "        mlp = self.out_lyr(F.relu(self.thrd_lyr(F.relu(self.snd_lyr(F.dropout(F.relu(self.fst_lyr(mlp)),p=self.dropout_fac))))))\n",
    "        fac = torch.clip(self.alpha,min=0.0,max=1.0)\n",
    "        out = fac*gmf+ (1-fac)*mlp\n",
    "        out = torch.clip(out,min=self.min_r,max=self.max_r)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "user:tensor([2, 1]),item:tensor([2, 0]) and rating:tensor([4.0000, 4.5000], dtype=torch.float64)\n",
      "output of the network=> out:tensor([[0.4322],\n",
      "        [0.5724]], grad_fn=<ClampBackward1>),shape:torch.Size([2, 1])\n"
     ]
    }
   ],
   "source": [
    "#does it work\n",
    "model = NCF(3,3,4,0.5)\n",
    "for u,r in train_loader:\n",
    "    #user,item = u\n",
    "    print(f'user:{u[0]},item:{u[-1]} and rating:{r}')\n",
    "    #print(u)\n",
    "    out = model(u)\n",
    "    print(f'output of the network=> out:{out},shape:{out.shape}')\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Trainer(object):\n",
    "    def __init__(self, model, device,loss_fn=None, optimizer=None, scheduler=None,artifacts_loc=None,exp_tracker=None):\n",
    "\n",
    "        # Set params\n",
    "        self.model = model\n",
    "        self.device = device\n",
    "        self.loss_fn = loss_fn\n",
    "        self.optimizer = optimizer\n",
    "        self.scheduler = scheduler\n",
    "        self.store_loc = artifacts_loc\n",
    "        self.exp_tracker = exp_tracker\n",
    "\n",
    "    def train_step(self, dataloader):\n",
    "        \"\"\"Train step.\"\"\"\n",
    "        # Set model to train mode\n",
    "        self.model.train()\n",
    "        loss = 0.0\n",
    "\n",
    "        # Iterate over train batches\n",
    "        for i, batch in enumerate(dataloader):\n",
    "            #batch = [item.to(self.device) for item in batch]  # Set device\n",
    "            inputs,targets = batch\n",
    "            inputs = [item.to(self.device) for item in inputs]\n",
    "            targets = targets.to(self.device)\n",
    "            #inputs, targets = batch[:-1], batch[-1]\n",
    "            #import pdb;pdb.set_trace()\n",
    "            self.optimizer.zero_grad()  # Reset gradients\n",
    "            z = self.model(inputs)  # Forward pass\n",
    "            targets = targets.reshape(z.shape)\n",
    "            J = self.loss_fn(z.float(), targets.float())  # Define loss\n",
    "            J.backward()  # Backward pass\n",
    "            self.optimizer.step()  # Update weights\n",
    "\n",
    "            # Cumulative Metrics\n",
    "            loss += (J.detach().item() - loss) / (i + 1)\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def eval_step(self, dataloader):\n",
    "        \"\"\"Validation or test step.\"\"\"\n",
    "        # Set model to eval mode\n",
    "        self.model.eval()\n",
    "        loss = 0.0\n",
    "        y_trues, y_probs = [], []\n",
    "\n",
    "        # Iterate over val batches\n",
    "        with torch.inference_mode():\n",
    "            for i, batch in enumerate(dataloader):\n",
    "                inputs,y_true = batch\n",
    "                inputs = [item.to(self.device) for item in inputs]\n",
    "                y_true = y_true.to(self.device).float()\n",
    "\n",
    "                # Step\n",
    "                z = self.model(inputs).float()  # Forward pass\n",
    "                y_true = y_true.reshape(z.shape)\n",
    "                J = self.loss_fn(z, y_true).item()\n",
    "\n",
    "                # Cumulative Metrics\n",
    "                loss += (J - loss) / (i + 1)\n",
    "\n",
    "                # Store outputs\n",
    "                y_prob = z.cpu().numpy()\n",
    "                y_probs.extend(y_prob)\n",
    "                y_trues.extend(y_true.cpu().numpy())\n",
    "\n",
    "        return loss, np.vstack(y_trues), np.vstack(y_probs)\n",
    "\n",
    "    def predict_step(self, dataloader):\n",
    "        \"\"\"Prediction step.\"\"\"\n",
    "        # Set model to eval mode\n",
    "        self.model.eval()\n",
    "        y_probs = []\n",
    "\n",
    "        # Iterate over val batches\n",
    "        with torch.inference_mode():\n",
    "            for i, batch in enumerate(dataloader):\n",
    "\n",
    "                # Forward pass w/ inputs\n",
    "                inputs, targets = batch\n",
    "                z = self.model(inputs).float()\n",
    "\n",
    "                # Store outputs\n",
    "                y_prob = z.cpu().numpy()\n",
    "                y_probs.extend(y_prob)\n",
    "\n",
    "        return np.vstack(y_probs)\n",
    "    \n",
    "    def train(self, num_epochs, patience, train_dataloader, val_dataloader, \n",
    "              tolerance=1e-5):\n",
    "        best_val_loss = np.inf\n",
    "        training_stats = defaultdict(list)\n",
    "        for epoch in tqdm(range(num_epochs)):\n",
    "            # Steps\n",
    "            train_loss = self.train_step(dataloader=train_dataloader)\n",
    "            val_loss, _, _ = self.eval_step(dataloader=val_dataloader)\n",
    "            #store stats\n",
    "            training_stats['epoch'].append(epoch)\n",
    "            training_stats['train_loss'].append(train_loss)\n",
    "            training_stats['val_loss'].append(val_loss)\n",
    "            #log-stats\n",
    "            # wandb.init(project=f\"{args.trail_id}_{args.dataset}_{args.data_type}\",config=config_dict)\n",
    "            if self.exp_tracker == 'wandb':\n",
    "                log_metrics = {'epoch':epoch,'train_loss':train_loss,'val_loss':val_loss}\n",
    "                wandb.log(log_metrics,step=epoch)\n",
    "            \n",
    "            self.scheduler.step(val_loss)\n",
    "\n",
    "            # Early stopping\n",
    "            if val_loss < best_val_loss - tolerance:\n",
    "                best_val_loss = val_loss\n",
    "                best_model = self.model\n",
    "                _patience = patience  # reset _patience\n",
    "            else:\n",
    "                _patience -= 1\n",
    "            if not _patience:  # 0\n",
    "                print(\"Stopping early!\")\n",
    "                break\n",
    "\n",
    "            # Tracking\n",
    "            #mlflow.log_metrics({\"train_loss\": train_loss, \"val_loss\": val_loss}, step=epoch)\n",
    "\n",
    "            # Logging\n",
    "            if epoch%5 == 0:\n",
    "                print(\n",
    "                    f\"Epoch: {epoch+1} | \"\n",
    "                    f\"train_loss: {train_loss:.5f}, \"\n",
    "                    f\"val_loss: {val_loss:.5f}, \"\n",
    "                    f\"lr: {self.optimizer.param_groups[0]['lr']:.2E}, \"\n",
    "                    f\"_patience: {_patience}\"\n",
    "                )\n",
    "        if self.store_loc:\n",
    "            pd.DataFrame(training_stats).to_csv(self.store_loc/'training_stats.csv',index=False)\n",
    "        return best_model, best_val_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_fn = nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer, mode=\"min\", factor=0.1, patience=5)\n",
    "\n",
    "trainer = Trainer(model,'cpu',loss_fn,optimizer,scheduler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 11%|█         | 11/100 [00:00<00:01, 50.86it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1 | train_loss: 8.55765, val_loss: 8.12049, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 6 | train_loss: 8.41908, val_loss: 8.07732, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 11 | train_loss: 8.28326, val_loss: 8.00117, lr: 1.00E-03, _patience: 8\n",
      "Epoch: 16 | train_loss: 8.15473, val_loss: 7.88881, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 21 | train_loss: 8.01803, val_loss: 7.85543, lr: 1.00E-03, _patience: 8\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 45%|████▌     | 45/100 [00:00<00:00, 93.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 26 | train_loss: 7.88548, val_loss: 7.76981, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 31 | train_loss: 7.75244, val_loss: 7.69328, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 36 | train_loss: 7.61687, val_loss: 7.61532, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 41 | train_loss: 7.48551, val_loss: 7.53750, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 46 | train_loss: 7.35090, val_loss: 7.42478, lr: 1.00E-03, _patience: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 65%|██████▌   | 65/100 [00:00<00:00, 96.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 51 | train_loss: 7.20997, val_loss: 7.34263, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 56 | train_loss: 7.07457, val_loss: 7.30977, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 61 | train_loss: 6.93769, val_loss: 7.23912, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 66 | train_loss: 6.80419, val_loss: 7.16288, lr: 1.00E-03, _patience: 9\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▌ | 85/100 [00:00<00:00, 95.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 71 | train_loss: 6.66622, val_loss: 7.08776, lr: 1.00E-03, _patience: 8\n",
      "Epoch: 76 | train_loss: 6.52944, val_loss: 6.96159, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 81 | train_loss: 6.38914, val_loss: 6.89294, lr: 1.00E-03, _patience: 10\n",
      "Epoch: 86 | train_loss: 6.24970, val_loss: 6.86126, lr: 1.00E-03, _patience: 9\n",
      "Epoch: 91 | train_loss: 6.10893, val_loss: 6.74132, lr: 1.00E-03, _patience: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:01<00:00, 87.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 96 | train_loss: 5.96172, val_loss: 6.66766, lr: 1.00E-03, _patience: 10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(NCF(\n",
       "   (user_embd_mtrx): Embedding(3, 4)\n",
       "   (item_embd_mtrx): Embedding(3, 4)\n",
       "   (h): Linear(in_features=4, out_features=1, bias=True)\n",
       "   (fst_lyr): Linear(in_features=8, out_features=4, bias=True)\n",
       "   (snd_lyr): Linear(in_features=4, out_features=2, bias=True)\n",
       "   (thrd_lyr): Linear(in_features=2, out_features=1, bias=True)\n",
       "   (out_lyr): Linear(in_features=1, out_features=1, bias=True)\n",
       " ), 6.608727216720581)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train(100,10,train_loader,valid_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
