{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.movielens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MovieLens Dataset\n",
    "> Implementation of MovieLens datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "import os\n",
    "import tempfile\n",
    "from pathlib import Path\n",
    "from collections import defaultdict\n",
    "import shutil\n",
    "from datetime import date"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def download(url, savepath):\n",
    "    import wget\n",
    "    wget.download(url, str(savepath))\n",
    "\n",
    "\n",
    "def unzip(zippath, savepath):\n",
    "    import zipfile\n",
    "    zip = zipfile.ZipFile(zippath)\n",
    "    zip.extractall(savepath)\n",
    "    zip.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ML1MDataset:\n",
    "    def __init__(self, args):\n",
    "        self.args = args\n",
    "        self.min_rating = args.min_rating\n",
    "        self.min_uc = args.min_uc\n",
    "        self.min_sc = args.min_sc\n",
    "        self.split = args.split\n",
    "\n",
    "        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'\n",
    "\n",
    "    @classmethod\n",
    "    def code(cls):\n",
    "        return 'ml-1m'\n",
    "\n",
    "    @classmethod\n",
    "    def raw_code(cls):\n",
    "        return cls.code()\n",
    "\n",
    "    @classmethod\n",
    "    def url(cls):\n",
    "        return 'http://files.grouplens.org/datasets/movielens/ml-1m.zip'\n",
    "\n",
    "    @classmethod\n",
    "    def is_zipfile(cls):\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def zip_file_content_is_folder(cls):\n",
    "        return True\n",
    "\n",
    "    @classmethod\n",
    "    def all_raw_file_names(cls):\n",
    "        return ['README',\n",
    "                'movies.dat',\n",
    "                'ratings.dat',\n",
    "                'users.dat']\n",
    "\n",
    "    def load_dataset(self):\n",
    "        self.preprocess()\n",
    "        dataset_path = self._get_preprocessed_dataset_path()\n",
    "        dataset = pickle.load(dataset_path.open('rb'))\n",
    "        return dataset\n",
    "\n",
    "    def preprocess(self):\n",
    "        dataset_path = self._get_preprocessed_dataset_path()\n",
    "        if dataset_path.is_file():\n",
    "            print('Already preprocessed. Skip preprocessing')\n",
    "            return\n",
    "        if not dataset_path.parent.is_dir():\n",
    "            dataset_path.parent.mkdir(parents=True)\n",
    "        self.maybe_download_raw_dataset()\n",
    "        df = self.load_ratings_df()\n",
    "        df = self.make_implicit(df)\n",
    "        df = self.filter_triplets(df)\n",
    "        df, umap, smap = self.densify_index(df)\n",
    "        train, val, test = self.split_df(df, len(umap))\n",
    "        dataset = {'train': train,\n",
    "                   'val': val,\n",
    "                   'test': test,\n",
    "                   'umap': umap,\n",
    "                   'smap': smap}\n",
    "        with dataset_path.open('wb') as f:\n",
    "            pickle.dump(dataset, f)\n",
    "\n",
    "    def maybe_download_raw_dataset(self):\n",
    "        folder_path = self._get_rawdata_folder_path()\n",
    "        if folder_path.is_dir() and\\\n",
    "           all(folder_path.joinpath(filename).is_file() for filename in self.all_raw_file_names()):\n",
    "            print('Raw data already exists. Skip downloading')\n",
    "            return\n",
    "        print(\"Raw file doesn't exist. Downloading...\")\n",
    "        if self.is_zipfile():\n",
    "            tmproot = Path(tempfile.mkdtemp())\n",
    "            tmpzip = tmproot.joinpath('file.zip')\n",
    "            tmpfolder = tmproot.joinpath('folder')\n",
    "            download(self.url(), tmpzip)\n",
    "            unzip(tmpzip, tmpfolder)\n",
    "            if self.zip_file_content_is_folder():\n",
    "                tmpfolder = tmpfolder.joinpath(os.listdir(tmpfolder)[0])\n",
    "            shutil.move(tmpfolder, folder_path)\n",
    "            shutil.rmtree(tmproot)\n",
    "            print()\n",
    "        else:\n",
    "            tmproot = Path(tempfile.mkdtemp())\n",
    "            tmpfile = tmproot.joinpath('file')\n",
    "            download(self.url(), tmpfile)\n",
    "            folder_path.mkdir(parents=True)\n",
    "            shutil.move(tmpfile, folder_path.joinpath('ratings.csv'))\n",
    "            shutil.rmtree(tmproot)\n",
    "            print()\n",
    "\n",
    "    def make_implicit(self, df):\n",
    "        print('Turning into implicit ratings')\n",
    "        df = df[df['rating'] >= self.min_rating]\n",
    "        # return df[['uid', 'sid', 'timestamp']]\n",
    "        return df\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        folder_path = self._get_rawdata_folder_path()\n",
    "        file_path = folder_path.joinpath('ratings.dat')\n",
    "        df = pd.read_csv(file_path, sep='::', header=None, engine='python')\n",
    "        df.columns = ['uid', 'sid', 'rating', 'timestamp']\n",
    "        return df\n",
    "\n",
    "    def filter_triplets(self, df):\n",
    "        print('Filtering triplets')\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",
    "        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 split_df(self, df, user_count):\n",
    "        if self.args.split == 'leave_one_out':\n",
    "            print('Splitting')\n",
    "            user_group = df.groupby('uid')\n",
    "            user2items = user_group.progress_apply(lambda d: list(d.sort_values(by='timestamp')['sid']))\n",
    "            train, val, test = {}, {}, {}\n",
    "            for user in range(user_count):\n",
    "                items = user2items[user]\n",
    "                train[user], val[user], test[user] = items[:-2], items[-2:-1], items[-1:]\n",
    "            return train, val, test\n",
    "        elif self.args.split == 'holdout':\n",
    "            print('Splitting')\n",
    "            np.random.seed(self.args.dataset_split_seed)\n",
    "            eval_set_size = self.args.eval_set_size\n",
    "\n",
    "            # Generate user indices\n",
    "            permuted_index = np.random.permutation(user_count)\n",
    "            train_user_index = permuted_index[                :-2*eval_set_size]\n",
    "            val_user_index   = permuted_index[-2*eval_set_size:  -eval_set_size]\n",
    "            test_user_index  = permuted_index[  -eval_set_size:                ]\n",
    "\n",
    "            # Split DataFrames\n",
    "            train_df = df.loc[df['uid'].isin(train_user_index)]\n",
    "            val_df   = df.loc[df['uid'].isin(val_user_index)]\n",
    "            test_df  = df.loc[df['uid'].isin(test_user_index)]\n",
    "\n",
    "            # DataFrame to dict => {uid : list of sid's}\n",
    "            train = dict(train_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            val   = dict(val_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            test  = dict(test_df.groupby('uid').progress_apply(lambda d: list(d['sid'])))\n",
    "            return train, val, test\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "\n",
    "    def _get_rawdata_root_path(self):\n",
    "        return Path(self.args.data_raw_path)\n",
    "\n",
    "    def _get_rawdata_folder_path(self):\n",
    "        root = self._get_rawdata_root_path()\n",
    "        return root.joinpath(self.raw_code())\n",
    "\n",
    "    def _get_preprocessed_root_path(self):\n",
    "        return Path(self.args.data_save_path)\n",
    "\n",
    "    def _get_preprocessed_folder_path(self):\n",
    "        preprocessed_root = self._get_preprocessed_root_path()\n",
    "        folder_name = '{}_min_rating{}-min_uc{}-min_sc{}-split{}' \\\n",
    "            .format(self.code(), self.min_rating, self.min_uc, self.min_sc, self.split)\n",
    "        return preprocessed_root.joinpath(folder_name)\n",
    "\n",
    "    def _get_preprocessed_dataset_path(self):\n",
    "        folder = self._get_preprocessed_folder_path()\n",
    "        return folder.joinpath('dataset.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Raw file doesn't exist. Downloading...\n",
      "\n",
      "Turning into implicit ratings\n",
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6040/6040 [00:02<00:00, 2707.83it/s]\n"
     ]
    }
   ],
   "source": [
    "class Args:\n",
    "    min_rating = 0\n",
    "    min_uc = 5\n",
    "    min_sc = 5\n",
    "    split = 'leave_one_out'\n",
    "    data_raw_path = 'data/bronze'\n",
    "    data_save_path = 'data/silver'\n",
    "\n",
    "args = Args()\n",
    "\n",
    "os.makedirs(args.data_raw_path, exist_ok=True)\n",
    "os.makedirs(args.data_save_path, exist_ok=True)\n",
    "\n",
    "dataset = ML1MDataset(args)\n",
    "dataset.preprocess()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m./data\u001b[00m\n",
      "├── [ 24M]  \u001b[01;34mbronze\u001b[00m\n",
      "│   └── [ 24M]  \u001b[01;34mml-1m\u001b[00m\n",
      "│       ├── [167K]  movies.dat\n",
      "│       ├── [ 23M]  ratings.dat\n",
      "│       ├── [5.4K]  README\n",
      "│       └── [131K]  users.dat\n",
      "└── [3.0M]  \u001b[01;34msilver\u001b[00m\n",
      "    └── [3.0M]  \u001b[01;34mml-1m_min_rating0-min_uc5-min_sc5-splitleave_one_out\u001b[00m\n",
      "        └── [3.0M]  dataset.pkl\n",
      "\n",
      "  27M used in 4 directories, 5 files\n"
     ]
    }
   ],
   "source": [
    "!tree --du -h -C ./data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-21 09:41:25\n",
      "\n",
      "Compiler    : GCC 7.5.0\n",
      "OS          : Linux\n",
      "Release     : 5.4.104+\n",
      "Machine     : x86_64\n",
      "Processor   : x86_64\n",
      "CPU cores   : 2\n",
      "Architecture: 64bit\n",
      "\n",
      "google : 2.0.3\n",
      "numpy  : 1.19.5\n",
      "pandas : 1.1.5\n",
      "IPython: 5.5.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
