{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.bases.interactions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Interactions Dataset\n",
    "> Implementation of base modules for interactions dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from typing import Any, Iterable, List, Optional, Tuple, Union, Callable\n",
    "\n",
    "import sys\n",
    "import os\n",
    "from os import path as osp\n",
    "import collections\n",
    "\n",
    "import random\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.sparse import coo_matrix\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "from pytorch_lightning import LightningDataModule\n",
    "\n",
    "from recohut.utils.common_utils import *\n",
    "from recohut.datasets import base\n",
    "from recohut.transforms.splitting import random_split, stratified_split\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## InteractionsDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class InteractionsDataset(torch.utils.data.Dataset, base.Dataset):\n",
    "\n",
    "    def __init__(self,\n",
    "                 data_dir,\n",
    "                 min_rating=None,\n",
    "                 min_uc=5,\n",
    "                 min_sc=5,\n",
    "                 num_negative_samples=100,\n",
    "                 max_samples=200,\n",
    "                 data_type=None,\n",
    "                 split_type='random',\n",
    "                 val_p=0.2,\n",
    "                 test_p=0.2,\n",
    "                 seed=42,\n",
    "                 *args,\n",
    "                 **kwargs):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            data_dir: Where to save/load the data\n",
    "            min_uc: minimum user count to keep in the data\n",
    "            min_sc: minimum item count to keep in the data\n",
    "            min_rating: minimum rating threshold to convert explicit feedback into implicit\n",
    "            num_negative_samples: number of negative samples for each positive one\n",
    "            max_samples: max samples limit\n",
    "            data_type: train/valid/test\n",
    "            split_type: data split method - stratified/random\n",
    "            val_p: Percent (float) or number (int) of samples to use for the validation split\n",
    "            test_p: Percent (float) or number (int) of samples to use for the test split\n",
    "            seed: Random seed to be used for train/val/test splits\n",
    "        \"\"\"\n",
    "        self.min_rating = min_rating\n",
    "        self.min_uc = min_uc\n",
    "        self.min_sc = min_sc\n",
    "        self.num_negative_samples = num_negative_samples\n",
    "        self.max_samples = 200\n",
    "        self.data_type = data_type\n",
    "        self.val_p = val_p if val_p is not None else 0.2\n",
    "        self.test_p = test_p if test_p is not None else 0.2\n",
    "        self.seed = seed\n",
    "        self.split_type = split_type\n",
    "\n",
    "        super().__init__(data_dir)\n",
    "\n",
    "        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'\n",
    "\n",
    "        self._process()\n",
    "\n",
    "        if self.data_type is not None:\n",
    "            self.load()\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return ['data_train.pt',\n",
    "                'data_valid_pos.pt',\n",
    "                'data_valid_neg.pt',\n",
    "                'data_test_pos.pt',\n",
    "                'data_test_neg.pt']\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def make_implicit(self, df):\n",
    "        \"convert the explicit data to implicit by only keeping interactions with a rating >= min_rating\"\n",
    "        print('Turning into implicit ratings')\n",
    "        df = df[df['rating'] >= self.min_rating].reset_index(drop=True)\n",
    "        df['rating'] = 1\n",
    "        return df\n",
    "\n",
    "    def filter_triplets(self, df):\n",
    "        print('Filtering triplets')\n",
    "        if self.min_sc > 0 or self.min_uc > 0:\n",
    "            item_sizes = df.groupby('sid').size()\n",
    "            good_items = item_sizes.index[item_sizes >= self.min_sc]\n",
    "            user_sizes = df.groupby('uid').size()\n",
    "            good_users = user_sizes.index[user_sizes >= self.min_uc]\n",
    "            while len(good_items) < len(item_sizes) or len(good_users) < len(user_sizes):\n",
    "                if self.min_sc > 0:\n",
    "                    item_sizes = df.groupby('sid').size()\n",
    "                    good_items = item_sizes.index[item_sizes >= self.min_sc]\n",
    "                    df = df[df['sid'].isin(good_items)]\n",
    "\n",
    "                if self.min_uc > 0:\n",
    "                    user_sizes = df.groupby('uid').size()\n",
    "                    good_users = user_sizes.index[user_sizes >= self.min_uc]\n",
    "                    df = df[df['uid'].isin(good_users)]\n",
    "\n",
    "                item_sizes = df.groupby('sid').size()\n",
    "                good_items = item_sizes.index[item_sizes >= self.min_sc]\n",
    "                user_sizes = df.groupby('uid').size()\n",
    "                good_users = user_sizes.index[user_sizes >= self.min_uc]\n",
    "        return df\n",
    "\n",
    "    def densify_index(self, df):\n",
    "        print('Densifying index')\n",
    "        umap = {u: i for i, u in enumerate(set(df['uid']))}\n",
    "        smap = {s: i for i, s in enumerate(set(df['sid']))}\n",
    "        df['uid'] = df['uid'].map(umap)\n",
    "        df['sid'] = df['sid'].map(smap)\n",
    "        return df, umap, smap\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @staticmethod\n",
    "    def _convert_to_torch_sparse(mat):\n",
    "        values = mat.data\n",
    "        indices = np.vstack((mat.row, mat.col))\n",
    "\n",
    "        i = torch.LongTensor(indices)\n",
    "        v = torch.FloatTensor(values)\n",
    "        shape = mat.shape\n",
    "\n",
    "        return torch.sparse.FloatTensor(i, v, torch.Size(shape))\n",
    "\n",
    "    def process(self):\n",
    "        df = self.load_ratings_df()\n",
    "        if self.min_rating:\n",
    "            df = self.make_implicit(df)\n",
    "        df = self.filter_triplets(df)\n",
    "        df, umap, smap = self.densify_index(df)\n",
    "        self.num_users = max(df.uid) + 1 # df.uid.nunique()\n",
    "        self.num_items = max(df.sid) + 1 # df.sid.nunique()\n",
    "        mat = coo_matrix((np.array(df.rating),\n",
    "                          (np.array(df.uid), np.array(df.sid))),\n",
    "                         shape=(self.num_users, self.num_items))\n",
    "        \n",
    "        if self.split_type == 'random':\n",
    "            mat_train, mat_valid, mat_test = random_split(mat = mat,\n",
    "                                                        val_p = self.val_p,\n",
    "                                                        test_p = self.test_p,\n",
    "                                                        seed = self.seed)\n",
    "        elif self.split_type == 'stratified':\n",
    "            mat_train, mat_valid, mat_test = stratified_split(mat = mat,\n",
    "                                                        val_p = self.val_p,\n",
    "                                                        test_p = self.test_p,\n",
    "                                                        seed = self.seed)\n",
    "\n",
    "        mat_train = self._convert_to_torch_sparse(mat_train)\n",
    "        torch.save(mat_train, self.processed_paths[0])\n",
    "\n",
    "        mat_valid_pos = self._convert_to_torch_sparse(mat_valid)._indices().T \n",
    "        _, indices = np.unique(mat_valid_pos[:, 0], return_index=True)\n",
    "        mat_valid_pos = mat_valid_pos[indices, :]\n",
    "        torch.save(mat_valid_pos, self.processed_paths[1])\n",
    "\n",
    "        pos_items = set(zip(mat_valid.row, mat_valid.col))\n",
    "        mat_valid_neg = self._negative_sample(np.arange(mat_valid.shape[0]), pos_items)\n",
    "        mat_valid_neg = torch.tensor(mat_valid_neg, dtype=torch.int)\n",
    "        torch.save(mat_valid_neg, self.processed_paths[2])\n",
    "\n",
    "        mat_test_pos = self._convert_to_torch_sparse(mat_test)._indices().T \n",
    "        _, indices = np.unique(mat_test_pos[:, 0], return_index=True)\n",
    "        mat_test_pos = mat_test_pos[indices, :]\n",
    "        torch.save(mat_test_pos, self.processed_paths[3])\n",
    "\n",
    "        pos_items = set(zip(mat_test.row, mat_test.col))\n",
    "        mat_test_neg = self._negative_sample(np.arange(mat_test.shape[0]), pos_items)\n",
    "        mat_test_neg = torch.tensor(mat_test_neg, dtype=torch.int)\n",
    "        torch.save(mat_test_neg, self.processed_paths[4])\n",
    "\n",
    "        return mat\n",
    "        \n",
    "    def todense(self) -> np.matrix:\n",
    "        \"\"\"Transforms sparse matrix to np.matrix, 2-d.\"\"\"\n",
    "        return self.mat.todense()\n",
    "\n",
    "    def toarray(self) -> np.array:\n",
    "        \"\"\"Transforms sparse matrix to np.array, 2-d.\"\"\"\n",
    "        return self.mat.toarray()\n",
    "\n",
    "    def head(self, n: int = 5) -> np.array:\n",
    "        \"\"\"Return the first ``n`` rows of the dense matrix as a np.array, 2-d.\"\"\"\n",
    "        n = self._prep_head_tail_n(n=n)\n",
    "        return self.mat.tocsr()[range(n), :].toarray()\n",
    "\n",
    "    def tail(self, n: int = 5) -> np.array:\n",
    "        \"\"\"Return the last ``n`` rows of the dense matrix as a np.array, 2-d.\"\"\"\n",
    "        n = self._prep_head_tail_n(n=n)\n",
    "        return self.mat.tocsr()[range(-n, 0), :].toarray()\n",
    "\n",
    "    def _prep_head_tail_n(self, n: int) -> int:\n",
    "        \"\"\"Ensure we don't run into an ``IndexError`` when using ``head`` or ``tail`` methods.\"\"\"\n",
    "        if n < 0:\n",
    "            n = self.num_users + n\n",
    "        if n > self.num_users:\n",
    "            n = self.num_users\n",
    "        return n\n",
    "\n",
    "    def _negative_sample(self, user_id: Union[int, np.array], positive_items) -> np.array:\n",
    "        \"\"\"Generate negative samples for a ``user_id``.\"\"\"\n",
    "        if self.max_samples > 0:\n",
    "            # if we are here, we are doing true negative sampling\n",
    "            negative_item_ids_list = list()\n",
    "\n",
    "            if not isinstance(user_id, collections.abc.Iterable):\n",
    "                user_id = [user_id]\n",
    "\n",
    "            for specific_user_id in user_id:\n",
    "                # generate true negative samples for the ``user_id``\n",
    "                samples_checked = 0\n",
    "                temp_negative_item_ids_list = list()\n",
    "\n",
    "                while len(temp_negative_item_ids_list) < self.num_negative_samples:\n",
    "                    negative_item_id = random.choice(range(self.num_items))\n",
    "                    # we have a negative sample, make sure the user has not interacted with it\n",
    "                    # before, else we resample and try again\n",
    "                    while (\n",
    "                        (specific_user_id, negative_item_id) in positive_items\n",
    "                        or negative_item_id in temp_negative_item_ids_list\n",
    "                    ):\n",
    "                        if samples_checked >= self.max_samples:\n",
    "                            num_samples_left_to_generate = (\n",
    "                                self.num_negative_samples - len(temp_negative_item_ids_list) - 1\n",
    "                            )\n",
    "                            temp_negative_item_ids_list += random.choices(\n",
    "                                range(self.num_items), k=num_samples_left_to_generate\n",
    "                            )\n",
    "                            break\n",
    "\n",
    "                        negative_item_id = random.choice(range(self.num_items))\n",
    "                        samples_checked += 1\n",
    "\n",
    "                    temp_negative_item_ids_list.append(negative_item_id)\n",
    "\n",
    "                negative_item_ids_list += [np.array(temp_negative_item_ids_list)]\n",
    "\n",
    "            if len(user_id) > 1:\n",
    "                negative_item_ids_array = np.stack(negative_item_ids_list)\n",
    "            else:\n",
    "                negative_item_ids_array = negative_item_ids_list[0]\n",
    "        else:\n",
    "            # if we are here, we are doing approximate negative sampling\n",
    "            if isinstance(user_id, collections.abc.Iterable):\n",
    "                size = (len(user_id), self.num_negative_samples)\n",
    "            else:\n",
    "                size = (self.num_negative_samples,)\n",
    "\n",
    "            negative_item_ids_array = np.random.randint(\n",
    "                low=0,\n",
    "                high=self.num_items,\n",
    "                size=size,\n",
    "            )\n",
    "\n",
    "        return negative_item_ids_array\n",
    "\n",
    "    def load(self):\n",
    "        if self.data_type=='train':\n",
    "            self.train = torch.load(self.processed_paths[0])\n",
    "            self.train_pos = self.train._indices().T\n",
    "            self.n_users, self.n_items = self.train.size()\n",
    "\n",
    "            self.train_score = torch.sparse.sum(self.train, dim=0).to_dense().repeat((self.n_users, 1))\n",
    "            self.train_score[self.train_pos[:, 0], self.train_pos[:, 1]] = 0\n",
    "        elif self.data_type=='valid':\n",
    "            self.valid_pos = torch.load(self.processed_paths[1])\n",
    "            self.valid_neg = torch.load(self.processed_paths[2])\n",
    "            self.n_users = self.valid_pos.shape[0]\n",
    "\n",
    "            valid_items = []\n",
    "            for u in range(self.n_users):\n",
    "                items = torch.cat((self.valid_pos[u, 1].view(1), self.valid_neg[u]))\n",
    "                valid_items.append(items)\n",
    "\n",
    "            self.valid_items = torch.vstack(valid_items)\n",
    "            self.valid_labels = torch.zeros(self.valid_items.shape)\n",
    "            self.valid_labels[:, 0] += 1\n",
    "        else:\n",
    "            self.test_pos = torch.load(self.processed_paths[3])\n",
    "            self.test_neg = torch.load(self.processed_paths[4])\n",
    "            self.n_users = self.test_pos.shape[0]\n",
    "\n",
    "            test_items = []\n",
    "            for u in range(self.n_users):\n",
    "                items = torch.cat((self.test_pos[u, 1].view(1), self.test_neg[u]))\n",
    "                test_items.append(items)\n",
    "\n",
    "            self.test_items = torch.vstack(test_items)\n",
    "            self.test_labels = torch.zeros(self.test_items.shape)\n",
    "            self.test_labels[:, 0] += 1\n",
    "\n",
    "    def __len__(self):\n",
    "            return self.n_users\n",
    "\n",
    "    def __train__(self, index):\n",
    "            return self.train_pos[index], self.train_score[self.train_pos[index][0]]\n",
    "\n",
    "    def __valid__(self, index):\n",
    "            return self.valid_pos[index], self.valid_items[index], self.valid_labels[index]\n",
    "\n",
    "    def __test__(self, index):\n",
    "            return self.test_pos[index], self.test_items[index], self.test_labels[index]\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        if self.data_type=='train':\n",
    "            return self.__train__(index)\n",
    "        elif self.data_type=='valid':\n",
    "            return self.__valid__(index)\n",
    "        else:\n",
    "            return self.__test__(index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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(osp.join(self.raw_dir, 'ml-1m', self.raw_file_names), self.raw_dir)\n",
    "        rmtree(osp.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"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## InteractionsDataModule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class InteractionsDataModule(LightningDataModule):\n",
    "\n",
    "    dataset_cls: str = \"\"\n",
    "\n",
    "    def __init__(self,\n",
    "                 data_dir: Optional[str] = None,\n",
    "                 num_workers: int = 0,\n",
    "                 normalize: bool = False,\n",
    "                 batch_size: int = 32,\n",
    "                 shuffle: bool = True,\n",
    "                 pin_memory: bool = True,\n",
    "                 drop_last: bool = False,\n",
    "                 *args, \n",
    "                 **kwargs) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            data_dir: Where to save/load the data\n",
    "            num_workers: How many workers to use for loading data\n",
    "            normalize: If true applies rating normalize\n",
    "            batch_size: How many samples per batch to load\n",
    "            shuffle: If true shuffles the train data every epoch\n",
    "            pin_memory: If true, the data loader will copy Tensors into CUDA pinned memory before\n",
    "                        returning them\n",
    "            drop_last: If true drops the last incomplete batch\n",
    "        \"\"\"\n",
    "        super().__init__(data_dir)\n",
    "\n",
    "        self.data_dir = data_dir if data_dir is not None else os.getcwd()\n",
    "        self.num_workers = num_workers\n",
    "        self.normalize = normalize\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle = shuffle\n",
    "        self.pin_memory = pin_memory\n",
    "        self.drop_last = drop_last\n",
    "        self.kwargs = kwargs\n",
    "\n",
    "    def prepare_data(self, *args: Any, **kwargs: Any) -> None:\n",
    "        \"\"\"Saves files to data_dir.\"\"\"\n",
    "        self.data = self.dataset_cls(self.data_dir, **self.kwargs)\n",
    "\n",
    "    def setup(self, stage: Optional[str] = None) -> None:\n",
    "        \"\"\"Creates train, val, and test dataset.\"\"\"\n",
    "        if stage == \"fit\" or stage is None:\n",
    "            self.dataset_train = self.dataset_cls(self.data_dir, **self.kwargs, data_type='train')\n",
    "            self.dataset_val = self.dataset_cls(self.data_dir, **self.kwargs, data_type='valid')\n",
    "        if stage == \"test\" or stage is None:\n",
    "            self.dataset_test = self.dataset_cls(self.data_dir, **self.kwargs, data_type='test')\n",
    "\n",
    "    def train_dataloader(self, *args: Any, **kwargs: Any) -> DataLoader:\n",
    "        \"\"\"The train dataloader.\"\"\"\n",
    "        return self._data_loader(self.dataset_train, shuffle=self.shuffle)\n",
    "\n",
    "    def val_dataloader(self, *args: Any, **kwargs: Any) -> Union[DataLoader, List[DataLoader]]:\n",
    "        \"\"\"The val dataloader.\"\"\"\n",
    "        return self._data_loader(self.dataset_val)\n",
    "\n",
    "    def test_dataloader(self, *args: Any, **kwargs: Any) -> Union[DataLoader, List[DataLoader]]:\n",
    "        \"\"\"The test dataloader.\"\"\"\n",
    "        return self._data_loader(self.dataset_test)\n",
    "\n",
    "    def _data_loader(self, dataset: Dataset, shuffle: bool = False) -> DataLoader:\n",
    "        return DataLoader(\n",
    "            dataset,\n",
    "            batch_size=self.batch_size,\n",
    "            shuffle=shuffle,\n",
    "            num_workers=self.num_workers,\n",
    "            drop_last=self.drop_last,\n",
    "            pin_memory=self.pin_memory,\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ML1mDataModule(InteractionsDataModule):\n",
    "    dataset_cls = ML1mDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2022-01-09 19:37:01\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",
      "recohut: 0.0.10\n",
      "torch  : 1.10.0+cu111\n",
      "pandas : 1.1.5\n",
      "sys    : 3.7.12 (default, Sep 10 2021, 00:21:48) \n",
      "[GCC 7.5.0]\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
}
