{
 "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 *\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 pandas as pd\n",
    "\n",
    "from recohut.utils.common_utils import *\n",
    "from recohut.datasets.bases.common import Dataset\n",
    "from recohut.datasets.bases.interactions import InteractionsDataset, InteractionsDataModule"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ML1m Rating Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ML1mDataModule(InteractionsDataModule):\n",
    "    dataset_cls = ML1mDataset"
   ]
  },
  {
   "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",
    "        self.val_p = 0.2\n",
    "        self.test_p = 0.2\n",
    "        self.seed = 42\n",
    "        self.split_type = 'stratified'\n",
    "\n",
    "args = Args()"
   ]
  },
  {
   "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 = ML1mDataModule(**args.__dict__)\n",
    "ds.prepare_data()"
   ]
  },
  {
   "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": "markdown",
   "metadata": {},
   "source": [
    "## ML100k Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class ML100kDataset(Dataset):\n",
    "    url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'\n",
    "    \n",
    "    def __init__(self, root):\n",
    "        super().__init__(root)\n",
    "    \n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        return ['u1.base', 'u1.test', 'u4.test', 'allbut.pl', 'u.item', \n",
    "                'ua.test', 'u.occupation', 'u3.test', 'u5.base', 'ub.test', \n",
    "                'u2.test', 'u3.base', 'u.genre', 'u.data', 'u4.base', \n",
    "                'u5.test', 'u.info', 'README', 'ub.base', 'mku.sh', 'u2.base', \n",
    "                'u.user', 'ua.base']\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        raise NotImplementedError\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",
    "        file_names = os.listdir(osp.join(self.raw_dir, 'ml-100k'))   \n",
    "        for file_name in file_names:\n",
    "            move(osp.join(self.raw_dir, 'ml-100k', file_name), self.raw_dir)\n",
    "        rmtree(osp.join(self.raw_dir, 'ml-100k'))\n",
    "        os.unlink(path)\n",
    "\n",
    "    def process(self):\n",
    "        raise NotImplementedError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2022-01-10 10:04:23\n",
      "\n",
      "recohut: 0.0.10\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",
      "IPython: 5.5.0\n",
      "pandas : 1.1.5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "#hide\n",
    "%reload_ext watermark\n",
    "%watermark -a \"Sparsh A.\" -m -iv -u -t -d -p recohut"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
