{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp datasets.bases.common"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Base dataset\n",
    "> Base class for dataset module."
   ]
  },
  {
   "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",
    "from typing import List, Optional, Callable, Union, Any, Tuple\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "tqdm.pandas()\n",
    "\n",
    "from abc import *\n",
    "from pathlib import Path\n",
    "import os\n",
    "import os.path as osp\n",
    "from collections.abc import Sequence\n",
    "import sys\n",
    "import tempfile\n",
    "import shutil\n",
    "import pickle\n",
    "import time\n",
    "import csv\n",
    "import math\n",
    "import operator\n",
    "import itertools\n",
    "from collections import defaultdict\n",
    "import datetime\n",
    "from datetime import date, timezone, timedelta\n",
    "from pandas import Timedelta\n",
    "\n",
    "import torch\n",
    "\n",
    "from recohut.utils.common_utils import download_url, extract_zip, extract_gz, makedirs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "def to_list(value: Any) -> Sequence:\n",
    "    if isinstance(value, Sequence) and not isinstance(value, str):\n",
    "        return value\n",
    "    else:\n",
    "        return [value]\n",
    "\n",
    "def files_exist(files: List[str]) -> bool:\n",
    "    # NOTE: We return `False` in case `files` is empty, leading to a\n",
    "    # re-processing of files on every instantiation.\n",
    "    return len(files) != 0 and all([osp.exists(f) for f in files])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#exporti\n",
    "class Dataset:\n",
    "    \"\"\"Dataset base class\n",
    "    \"\"\"\n",
    "    @property\n",
    "    def raw_file_names(self) -> Union[str, List[str], Tuple]:\n",
    "        r\"\"\"The name of the files in the :obj:`self.raw_dir` folder that must\n",
    "        be present in order to skip downloading.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> Union[str, List[str], Tuple]:\n",
    "        r\"\"\"The name of the files in the :obj:`self.processed_dir` folder that\n",
    "        must be present in order to skip processing.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def download(self):\n",
    "        r\"\"\"Downloads the dataset to the :obj:`self.raw_dir` folder.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def process(self):\n",
    "        r\"\"\"Processes the dataset to the :obj:`self.processed_dir` folder.\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def __init__(self, root=None):\n",
    "        self.root = root\n",
    "\n",
    "        if 'download' in self.__class__.__dict__:\n",
    "            self._download()\n",
    "\n",
    "        # if 'process' in self.__class__.__dict__:\n",
    "        #     self._process()\n",
    "\n",
    "    @property\n",
    "    def raw_dir(self) -> str:\n",
    "        return osp.join(self.root, 'raw')\n",
    "\n",
    "    @property\n",
    "    def processed_dir(self) -> str:\n",
    "        return osp.join(self.root, 'processed')\n",
    "\n",
    "    @property\n",
    "    def raw_paths(self) -> List[str]:\n",
    "        r\"\"\"The absolute filepaths that must be present in order to skip\n",
    "        downloading.\"\"\"\n",
    "        files = to_list(self.raw_file_names)\n",
    "        return [osp.join(self.raw_dir, f) for f in files]\n",
    "\n",
    "    @property\n",
    "    def processed_paths(self) -> List[str]:\n",
    "        r\"\"\"The absolute filepaths that must be present in order to skip\n",
    "        processing.\"\"\"\n",
    "        files = to_list(self.processed_file_names)\n",
    "        return [osp.join(self.processed_dir, f) for f in files]\n",
    "\n",
    "    def _download(self):\n",
    "        if files_exist(self.raw_paths):  # pragma: no cover\n",
    "            return\n",
    "\n",
    "        makedirs(self.raw_dir)\n",
    "        self.download()\n",
    "\n",
    "    def _process(self):\n",
    "        if files_exist(self.processed_paths):  # pragma: no cover\n",
    "            return\n",
    "\n",
    "        print('Processing...', file=sys.stderr)\n",
    "\n",
    "        makedirs(self.processed_dir)\n",
    "        self.process()\n",
    "\n",
    "        print('Done!', file=sys.stderr)\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        arg_repr = str(len(self)) if len(self) > 1 else ''\n",
    "        return f'{self.__class__.__name__}({arg_repr})'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SessionDataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDataset(Dataset):\n",
    "    r\"\"\"Session data base class.\n",
    "\n",
    "    Args:\n",
    "        min_session_length (int): Minimum number of items for a session to be valid\n",
    "        min_item_support (int): Minimum number of interactions for an item to be valid\n",
    "        eval_sec (int): these many seconds from the end will be taken as validation data\n",
    "\n",
    "    References:\n",
    "        1. https://github.com/Ethan-Yys/GRU4REC-pytorch-master/blob/master/preprocessing.py\n",
    "    \"\"\"\n",
    "    def __init__(self,\n",
    "                 root,\n",
    "                 min_session_length: int = None,\n",
    "                 min_item_support: int = None,\n",
    "                 eval_sec: int = 86400,  # valid/test data is one day = 86400 seconds\n",
    "                 ):\n",
    "        super().__init__(root)\n",
    "        self.min_session_length = min_session_length\n",
    "        self.min_item_support = min_item_support\n",
    "        self.eval_sec = eval_sec\n",
    "\n",
    "        self._process()\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        r\"\"\"load raw dataset into pandas dataframe\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def remove_short_sessions(self, df):\n",
    "        r\"\"\"delete sessions where session length is less than `min_session_length`\"\"\"\n",
    "        # groupby session id and get size of each session\n",
    "        session_len = df.groupby('uid').size()\n",
    "        df = df[np.in1d(df.uid, session_len[session_len >= self.min_session_length].index)]\n",
    "        return df\n",
    "\n",
    "    def remove_sparse_items(self, df):\n",
    "        r\"\"\"delete records of items which appeared less than `min_item_support` times\"\"\"\n",
    "        # groupby itemID and get size of each item\n",
    "        item_len = df.groupby('sid').size()\n",
    "        df = df[np.in1d(df.sid, item_len[item_len >= self.min_item_support].index)]\n",
    "        # remove sessions of less than n interactions again\n",
    "        df = self.remove_short_sessions(df)\n",
    "        return df\n",
    "\n",
    "    def split_df(self, df):\n",
    "        timeMax = df.timestamp.max() #maximum time in all records\n",
    "        sessionMaxTime = df.groupby('uid').timestamp.max() #group by uid and get the maximum time of each session\n",
    "        sessionTrain = sessionMaxTime[sessionMaxTime < (timeMax - self.eval_sec)].index #training split is all sessions that ended before the `eval_sec` seconds\n",
    "        sessionTest  = sessionMaxTime[sessionMaxTime >= (timeMax - self.eval_sec)].index #testing split is all sessions has records in the `eval_sec` seconds\n",
    "        train = df[np.in1d(df.uid, sessionTrain)]\n",
    "        test = df[np.in1d(df.uid, sessionTest)]\n",
    "        # Delete records in testing split where items are not in training split\n",
    "        test = test[np.in1d(test.sid, train.sid)]\n",
    "        # Delete Sessions in testing split\n",
    "        test = self.remove_short_sessions(test)\n",
    "        print('Training Set has {} Events, {} Sessions, and {} Items\\n\\n'.format(\n",
    "            len(train), train.uid.nunique(), train.sid.nunique()))\n",
    "        print('Validation Set has {} Events, {} Sessions, and {} Items\\n\\n'.format(\n",
    "            len(test), test.uid.nunique(), test.sid.nunique()))\n",
    "        return train, test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class YoochooseDataset(SessionDataset):\n",
    "    data_id = '1UEcKC4EfgMVD2n_zBvAyp0vRNyv7ndSF'\n",
    "\n",
    "    def __init__(self,\n",
    "                 root,\n",
    "                 min_session_length: int = 2,\n",
    "                 min_item_support: int = 5,\n",
    "                 eval_sec: int = 86400,\n",
    "                 ):\n",
    "        super().__init__(root, min_session_length, min_item_support, eval_sec)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        return 'rsc15-clicks.dat'\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['yoochoose_train.txt','yoochoose_valid.txt']\n",
    "\n",
    "    def download(self):\n",
    "        from google_drive_downloader import GoogleDriveDownloader as gdd\n",
    "        from shutil import move, rmtree\n",
    "\n",
    "        path = osp.join(self.raw_dir, 'rsc15.zip')\n",
    "        gdd.download_file_from_google_drive(self.data_id, path)\n",
    "        extract_zip(path, self.raw_dir)\n",
    "        move(osp.join(self.raw_dir, 'rsc15', 'raw', ),\n",
    "             osp.join(self.raw_dir, self.raw_file_names))\n",
    "        rmtree(osp.join(self.raw_dir, 'rsc15'))\n",
    "        os.unlink(path)\n",
    "\n",
    "    def process(self):\n",
    "        df = self.load_ratings_df()\n",
    "        if self.min_session_length is not None:\n",
    "            df = self.remove_short_sessions(df)\n",
    "        if self.min_item_support is not None:\n",
    "            df = self.remove_sparse_items(df)\n",
    "        train, test = self.split_df(df)\n",
    "        train.to_csv(self.processed_paths[0], sep=',', index=False)\n",
    "        test.to_csv(self.processed_paths[1], sep=',', index=False)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        df = pd.read_csv(self.raw_paths[0], header=None, usecols=[0, 1, 2],\n",
    "                         dtype={0: np.int32, 1: str, 2: np.int64})\n",
    "        df.columns = ['uid', 'timestamp', 'sid']\n",
    "        df['timestamp'] = df['timestamp'].apply(lambda x: datetime.datetime.strptime(\n",
    "            x, '%Y-%m-%dT%H:%M:%S.%fZ').timestamp())\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Training Set has 31637239 Events, 7966257 Sessions, and 37483 Items\n",
      "\n",
      "\n",
      "Validation Set has 71222 Events, 15324 Sessions, and 6751 Items\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = YoochooseDataset(root='/content/yoochoose')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[01;34m/content/yoochoose\u001b[00m\n",
      "├── [995M]  \u001b[01;34mprocessed\u001b[00m\n",
      "│   ├── [993M]  yoochoose_train.txt\n",
      "│   └── [2.3M]  yoochoose_valid.txt\n",
      "└── [1.4G]  \u001b[01;34mraw\u001b[00m\n",
      "    └── [1.4G]  rsc15-clicks.dat\n",
      "\n",
      " 2.4G used in 2 directories, 3 files\n"
     ]
    }
   ],
   "source": [
    "!tree --du -h -C /content/yoochoose"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDatasetv2(Dataset):\n",
    "    def __init__(self, root, column_names):\n",
    "        super().__init__(root)\n",
    "        self.SESSION_ID = column_names['SESSION_ID']\n",
    "        self.ITEM_ID = column_names['ITEM_ID']\n",
    "        self.TIMEFRAME = column_names['TIMEFRAME']\n",
    "        self.EVENT_DATE = column_names['EVENT_DATE']\n",
    "\n",
    "        self._process()\n",
    "    \n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['train.txt','test.txt','all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def process(self):\n",
    "        with open(self.raw_paths[0], \"r\") as f:\n",
    "            reader = csv.DictReader(f, delimiter=';')\n",
    "            sess_clicks = {}\n",
    "            sess_date = {}\n",
    "            ctr = 0\n",
    "            curid = -1\n",
    "            curdate = None\n",
    "            for data in reader:\n",
    "                sessid = data[self.SESSION_ID]\n",
    "                if curdate and not curid == sessid:\n",
    "                    date = ''\n",
    "                    date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))\n",
    "                    sess_date[curid] = date\n",
    "                curid = sessid\n",
    "                item = data[self.ITEM_ID], int(data[self.TIMEFRAME])\n",
    "                curdate = ''\n",
    "                curdate = data[self.EVENT_DATE]\n",
    "\n",
    "                if sessid in sess_clicks:\n",
    "                    sess_clicks[sessid] += [item]\n",
    "                else:\n",
    "                    sess_clicks[sessid] = [item]\n",
    "                ctr += 1\n",
    "            date = ''\n",
    "            date = time.mktime(time.strptime(curdate, '%Y-%m-%d'))\n",
    "            for i in list(sess_clicks):\n",
    "                sorted_clicks = sorted(sess_clicks[i], key=operator.itemgetter(1))\n",
    "                sess_clicks[i] = [c[0] for c in sorted_clicks]\n",
    "            sess_date[curid] = date\n",
    "\n",
    "        print(\"-- Reading data\")\n",
    "\n",
    "        # Filter out length 1 sessions\n",
    "        for s in list(sess_clicks):\n",
    "            if len(sess_clicks[s]) == 1:\n",
    "                del sess_clicks[s]\n",
    "                del sess_date[s]\n",
    "\n",
    "        # Count number of times each item appears\n",
    "        iid_counts = {}\n",
    "        for s in sess_clicks:\n",
    "            seq = sess_clicks[s]\n",
    "            for iid in seq:\n",
    "                if iid in iid_counts:\n",
    "                    iid_counts[iid] += 1\n",
    "                else:\n",
    "                    iid_counts[iid] = 1\n",
    "\n",
    "        sorted_counts = sorted(iid_counts.items(), key=operator.itemgetter(1))\n",
    "\n",
    "        length = len(sess_clicks)\n",
    "        for s in list(sess_clicks):\n",
    "            curseq = sess_clicks[s]\n",
    "            filseq = list(filter(lambda i: iid_counts[i] >= 5, curseq))\n",
    "            if len(filseq) < 2:\n",
    "                del sess_clicks[s]\n",
    "                del sess_date[s]\n",
    "            else:\n",
    "                sess_clicks[s] = filseq\n",
    "\n",
    "        # Split out test set based on dates\n",
    "        dates = list(sess_date.items())\n",
    "        maxdate = dates[0][1]\n",
    "\n",
    "        for _, date in dates:\n",
    "            if maxdate < date:\n",
    "                maxdate = date\n",
    "\n",
    "        # 7 days for test\n",
    "        splitdate = 0\n",
    "        splitdate = maxdate - 86400 * 7\n",
    "\n",
    "        print('Splitting date', splitdate)\n",
    "        tra_sess = filter(lambda x: x[1] < splitdate, dates)\n",
    "        tes_sess = filter(lambda x: x[1] > splitdate, dates)\n",
    "\n",
    "        # Sort sessions by date\n",
    "        tra_sess = sorted(tra_sess, key=operator.itemgetter(1))\n",
    "        tes_sess = sorted(tes_sess, key=operator.itemgetter(1))\n",
    "\n",
    "        print(len(tra_sess))\n",
    "        print(len(tes_sess))\n",
    "        print(tra_sess[:3])\n",
    "        print(tes_sess[:3])\n",
    "\n",
    "        print(\"-- Splitting train set and test set\")\n",
    "\n",
    "        item_dict = {}\n",
    "        # Convert training sessions to sequences and renumber items to start from 1\n",
    "        def obtian_tra():\n",
    "            train_ids = []\n",
    "            train_seqs = []\n",
    "            train_dates = []\n",
    "            item_ctr = 1\n",
    "            for s, date in tra_sess:\n",
    "                seq = sess_clicks[s]\n",
    "                outseq = []\n",
    "                for i in seq:\n",
    "                    if i in item_dict:\n",
    "                        outseq += [item_dict[i]]\n",
    "                    else:\n",
    "                        outseq += [item_ctr]\n",
    "                        item_dict[i] = item_ctr\n",
    "                        item_ctr += 1\n",
    "                if len(outseq) < 2:  # Doesn't occur\n",
    "                    continue\n",
    "                train_ids += [s]\n",
    "                train_dates += [date]\n",
    "                train_seqs += [outseq]\n",
    "            print(item_ctr)     # 43098, 37484\n",
    "            return train_ids, train_dates, train_seqs\n",
    "\n",
    "        # Convert test sessions to sequences, ignoring items that do not appear in training set\n",
    "        def obtian_tes():\n",
    "            test_ids = []\n",
    "            test_seqs = []\n",
    "            test_dates = []\n",
    "            for s, date in tes_sess:\n",
    "                seq = sess_clicks[s]\n",
    "                outseq = []\n",
    "                for i in seq:\n",
    "                    if i in item_dict:\n",
    "                        outseq += [item_dict[i]]\n",
    "                if len(outseq) < 2:\n",
    "                    continue\n",
    "                test_ids += [s]\n",
    "                test_dates += [date]\n",
    "                test_seqs += [outseq]\n",
    "            return test_ids, test_dates, test_seqs\n",
    "\n",
    "        tra_ids, tra_dates, tra_seqs = obtian_tra()\n",
    "        tes_ids, tes_dates, tes_seqs = obtian_tes()\n",
    "\n",
    "        def process_seqs(iseqs, idates):\n",
    "            out_seqs = []\n",
    "            out_dates = []\n",
    "            labs = []\n",
    "            ids = []\n",
    "            for id, seq, date in zip(range(len(iseqs)), iseqs, idates):\n",
    "                for i in range(1, len(seq)):\n",
    "                    tar = seq[-i]\n",
    "                    labs += [tar]\n",
    "                    out_seqs += [seq[:-i]]\n",
    "                    out_dates += [date]\n",
    "                    ids += [id]\n",
    "            return out_seqs, out_dates, labs, ids\n",
    "\n",
    "        tr_seqs, tr_dates, tr_labs, tr_ids = process_seqs(tra_seqs, tra_dates)\n",
    "        te_seqs, te_dates, te_labs, te_ids = process_seqs(tes_seqs, tes_dates)\n",
    "        tra = (tr_seqs, tr_labs)\n",
    "        tes = (te_seqs, te_labs)\n",
    "\n",
    "        print(len(tr_seqs))\n",
    "        print(len(te_seqs))\n",
    "        print(tr_seqs[:3], tr_dates[:3], tr_labs[:3])\n",
    "        print(te_seqs[:3], te_dates[:3], te_labs[:3])\n",
    "\n",
    "        all = 0\n",
    "\n",
    "        for seq in tra_seqs:\n",
    "            all += len(seq)\n",
    "        for seq in tes_seqs:\n",
    "            all += len(seq)\n",
    "        print('avg length: ', all/(len(tra_seqs) + len(tes_seqs) * 1.0))\n",
    "\n",
    "        pickle.dump(tra, open(self.processed_paths[0], 'wb'))\n",
    "        pickle.dump(tes, open(self.processed_paths[1], 'wb'))\n",
    "        pickle.dump(tra_seqs, open(self.processed_paths[2], 'wb'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDatasetv3(Dataset):\n",
    "    def __init__(self, root):\n",
    "        super().__init__(root)\n",
    "\n",
    "        self._process()\n",
    "    \n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        return ['train.txt','test.txt','all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def process(self):\n",
    "        def get_session_id(df, interval):\n",
    "            df_prev = df.shift()\n",
    "            is_new_session = (df.userId != df_prev.userId) | (\n",
    "                df.timestamp - df_prev.timestamp > interval\n",
    "            )\n",
    "            session_id = is_new_session.cumsum() - 1\n",
    "            return session_id\n",
    "\n",
    "        def group_sessions(df, interval):\n",
    "            sessionId = get_session_id(df, interval)\n",
    "            df = df.assign(sessionId=sessionId)\n",
    "            return df\n",
    "\n",
    "        def filter_short_sessions(df, min_len=2):\n",
    "            session_len = df.groupby('sessionId', sort=False).size()\n",
    "            long_sessions = session_len[session_len >= min_len].index\n",
    "            df_long = df[df.sessionId.isin(long_sessions)]\n",
    "            return df_long\n",
    "\n",
    "        def filter_infreq_items(df, min_support=5):\n",
    "            item_support = df.groupby('itemId', sort=False).size()\n",
    "            freq_items = item_support[item_support >= min_support].index\n",
    "            df_freq = df[df.itemId.isin(freq_items)]\n",
    "            return df_freq\n",
    "\n",
    "        def filter_until_all_long_and_freq(df, min_len=2, min_support=5):\n",
    "            while True:\n",
    "                df_long = filter_short_sessions(df, min_len)\n",
    "                df_freq = filter_infreq_items(df_long, min_support)\n",
    "                if len(df_freq) == len(df):\n",
    "                    break\n",
    "                df = df_freq\n",
    "            return df\n",
    "\n",
    "        def truncate_long_sessions(df, max_len=20, is_sorted=False):\n",
    "            if not is_sorted:\n",
    "                df = df.sort_values(['sessionId', 'timestamp'])\n",
    "            itemIdx = df.groupby('sessionId').cumcount()\n",
    "            df_t = df[itemIdx < max_len]\n",
    "            return df_t\n",
    "\n",
    "        def update_id(df, field):\n",
    "            labels = pd.factorize(df[field])[0]\n",
    "            kwargs = {field: labels}\n",
    "            df = df.assign(**kwargs)\n",
    "            return df\n",
    "\n",
    "        def remove_immediate_repeats(df):\n",
    "            df_prev = df.shift()\n",
    "            is_not_repeat = (df.sessionId != df_prev.sessionId) | (df.itemId != df_prev.itemId)\n",
    "            df_no_repeat = df[is_not_repeat]\n",
    "            return df_no_repeat\n",
    "\n",
    "        def reorder_sessions_by_endtime(df):\n",
    "            endtime = df.groupby('sessionId', sort=False).timestamp.max()\n",
    "            df_endtime = endtime.sort_values().reset_index()\n",
    "            oid2nid = dict(zip(df_endtime.sessionId, df_endtime.index))\n",
    "            sessionId_new = df.sessionId.map(oid2nid)\n",
    "            df = df.assign(sessionId=sessionId_new)\n",
    "            df = df.sort_values(['sessionId', 'timestamp'])\n",
    "            return df\n",
    "\n",
    "        def keep_top_n_items(df, n):\n",
    "            item_support = df.groupby('itemId', sort=False).size()\n",
    "            top_items = item_support.nlargest(n).index\n",
    "            df_top = df[df.itemId.isin(top_items)]\n",
    "            return df_top\n",
    "\n",
    "        def split_by_time(df, timedelta):\n",
    "            max_time = df.timestamp.max()\n",
    "            end_time = df.groupby('sessionId').timestamp.max()\n",
    "            split_time = max_time - timedelta\n",
    "            train_sids = end_time[end_time < split_time].index\n",
    "            df_train = df[df.sessionId.isin(train_sids)]\n",
    "            df_test = df[~df.sessionId.isin(train_sids)]\n",
    "            return df_train, df_test\n",
    "\n",
    "        def train_test_split(df, test_split=0.2):\n",
    "            endtime = df.groupby('sessionId', sort=False).timestamp.max()\n",
    "            endtime = endtime.sort_values()\n",
    "            num_tests = int(len(endtime) * test_split)\n",
    "            test_session_ids = endtime.index[-num_tests:]\n",
    "            df_train = df[~df.sessionId.isin(test_session_ids)]\n",
    "            df_test = df[df.sessionId.isin(test_session_ids)]\n",
    "            return df_train, df_test\n",
    "\n",
    "        def save_sessions(df, filepath):\n",
    "            df = reorder_sessions_by_endtime(df)\n",
    "            sessions = df.groupby('sessionId').itemId.apply(lambda x: ','.join(map(str, x)))\n",
    "            sessions.to_csv(filepath, sep='\\t', header=False, index=False)\n",
    "\n",
    "        def save_dataset(df_train, df_test):\n",
    "            # filter items in test but not in train\n",
    "            df_test = df_test[df_test.itemId.isin(df_train.itemId.unique())]\n",
    "            df_test = filter_short_sessions(df_test)\n",
    "\n",
    "            print(f'No. of Clicks: {len(df_train) + len(df_test)}')\n",
    "            print(f'No. of Items: {df_train.itemId.nunique()}')\n",
    "\n",
    "            # update itemId\n",
    "            train_itemId_new, uniques = pd.factorize(df_train.itemId)\n",
    "            df_train = df_train.assign(itemId=train_itemId_new)\n",
    "            oid2nid = {oid: i for i, oid in enumerate(uniques)}\n",
    "            test_itemId_new = df_test.itemId.map(oid2nid)\n",
    "            df_test = df_test.assign(itemId=test_itemId_new)\n",
    "\n",
    "            print('saving dataset')\n",
    "            save_sessions(df_train, self.processed_paths[0])\n",
    "            save_sessions(df_test, self.processed_paths[1])\n",
    "            num_items = len(uniques)\n",
    "            with open(self.processed_paths[2], 'w') as f:\n",
    "                f.write(str(num_items))\n",
    "\n",
    "        df = pd.read_csv(\n",
    "            osp.join(self.raw_dir,self.raw_file_names),\n",
    "            sep='\\t',\n",
    "            header=None,\n",
    "            names=['userId', 'timestamp', 'itemId'],\n",
    "            usecols=[0, 1, 4],\n",
    "            parse_dates=['timestamp'],\n",
    "            infer_datetime_format=True,\n",
    "        )\n",
    "\n",
    "        print('start preprocessing')\n",
    "        df = df.dropna()\n",
    "        df = update_id(df, 'userId')\n",
    "        df = update_id(df, 'itemId')\n",
    "        df = df.sort_values(['userId', 'timestamp'])\n",
    "\n",
    "        df = group_sessions(df, Timedelta(days=1))\n",
    "        df = remove_immediate_repeats(df)\n",
    "        df = truncate_long_sessions(df, is_sorted=True)\n",
    "        df = keep_top_n_items(df, n=30000)\n",
    "        df = filter_until_all_long_and_freq(df)\n",
    "        df_train, df_test = train_test_split(df, test_split=0.2)\n",
    "        save_dataset(df_train, df_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### v4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionDatasetv4(Dataset):\n",
    "    r\"\"\"Session dataset base class.\n",
    "\n",
    "    Args:\n",
    "        root (string): Root directory where the dataset should be saved.\n",
    "        process_method (string):\n",
    "            last: last day => test set\n",
    "            last_min_date: last day => test set, but from a minimal date onwards\n",
    "            days_test: last N days => test set\n",
    "            slice: create multiple train-test-combinations with a sliding window approach\n",
    "        min_date (string): Minimum date\n",
    "        session_length (int): Session time length :default = 30 * 60 #30 minutes\n",
    "        min_session_length (int): Minimum number of items for a session to be valid\n",
    "        min_item_support (int): Minimum number of interactions for an item to be valid\n",
    "        num_slices (int): Offset in days from the first date in the data set\n",
    "        days_offset (int): Number of days the training start date is shifted after creating one slice\n",
    "        days_shift (int): Days shift\n",
    "        days_train (int): Days in train set in each slice\n",
    "        days_test (int): Days in test set in each slice\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, root, process_method, min_date=None,\n",
    "                 session_length=None, min_session_length=None, min_item_support=None,\n",
    "                 num_slices=None, days_offset=None, days_shift=None, days_train=None,\n",
    "                 days_test=None, data=None):\n",
    "        super().__init__(root)\n",
    "        self.process_method = process_method\n",
    "        self.min_date = min_date\n",
    "        self.session_length = session_length\n",
    "        self.min_session_length = min_session_length\n",
    "        self.min_item_support = min_item_support\n",
    "        self.num_slices = num_slices\n",
    "        self.days_offset = days_offset\n",
    "        self.days_shift = days_shift\n",
    "        self.days_train = days_train\n",
    "        self.days_test = days_test\n",
    "        self.data = None\n",
    "\n",
    "        self._process()\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def load(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def filter_data(self): \n",
    "        data = self.data\n",
    "\n",
    "        #filter session length\n",
    "        session_lengths = data.groupby('SessionId').size()\n",
    "        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>1].index)]\n",
    "        \n",
    "        #filter item support\n",
    "        item_supports = data.groupby('ItemId').size()\n",
    "        data = data[np.in1d(data.ItemId, item_supports[item_supports>= self.min_item_support].index)]\n",
    "        \n",
    "        #filter session length\n",
    "        session_lengths = data.groupby('SessionId').size()\n",
    "        data = data[np.in1d(data.SessionId, session_lengths[session_lengths>= self.min_session_length].index)]\n",
    "        \n",
    "        #output\n",
    "        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Filtered data set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}\\n\\n'.\n",
    "              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))\n",
    "    \n",
    "        self.data = data\n",
    "        \n",
    "    def filter_min_date(self):\n",
    "        data = self.data\n",
    "\n",
    "        min_datetime = datetime.datetime.strptime(self.min_date + ' 00:00:00', '%Y-%m-%d %H:%M:%S')\n",
    "        \n",
    "        #filter\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_keep = session_max_times[session_max_times > min_datetime.datetime.timestamp()].index\n",
    "        \n",
    "        data = data[np.in1d(data.SessionId, session_keep)]\n",
    "        \n",
    "        #output\n",
    "        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Filtered data set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}\\n\\n'.\n",
    "              format(len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.date().isoformat(), data_end.date().isoformat()))\n",
    "        \n",
    "        self.data = data\n",
    "\n",
    "    def split_data_org(self):\n",
    "        data = self.data\n",
    "        tmax = data.Time.max()\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < tmax-86400].index\n",
    "        session_test = session_max_times[session_max_times >= tmax-86400].index\n",
    "        train = data[np.in1d(data.SessionId, session_train)]\n",
    "        test = data[np.in1d(data.SessionId, session_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Full train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))\n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\\t', index=False)\n",
    "        print('Test set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))\n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\\t', index=False)\n",
    "        \n",
    "        tmax = train.Time.max()\n",
    "        session_max_times = train.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < tmax-86400].index\n",
    "        session_valid = session_max_times[session_max_times >= tmax-86400].index\n",
    "        train_tr = train[np.in1d(train.SessionId, session_train)]\n",
    "        valid = train[np.in1d(train.SessionId, session_valid)]\n",
    "        valid = valid[np.in1d(valid.ItemId, train_tr.ItemId)]\n",
    "        tslength = valid.groupby('SessionId').size()\n",
    "        valid = valid[np.in1d(valid.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train_tr), train_tr.SessionId.nunique(), train_tr.ItemId.nunique()))\n",
    "        train_tr.to_csv(osp.join(self.processed_dir,'events_train_tr.txt'), sep='\\t', index=False)\n",
    "        print('Validation set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(valid), valid.SessionId.nunique(), valid.ItemId.nunique()))\n",
    "        valid.to_csv(osp.join(self.processed_dir,'events_train_valid.txt'), sep='\\t', index=False)\n",
    "\n",
    "    def split_data(self):\n",
    "        data = self.data\n",
    "        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        test_from = data_end - timedelta(self.days_test)\n",
    "        \n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        session_train = session_max_times[session_max_times < test_from.timestamp()].index\n",
    "        session_test = session_max_times[session_max_times >= test_from.timestamp()].index\n",
    "        train = data[np.in1d(data.SessionId, session_train)]\n",
    "        test = data[np.in1d(data.SessionId, session_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        print('Full train set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(train), train.SessionId.nunique(), train.ItemId.nunique()))\n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.txt'), sep='\\t', index=False)\n",
    "        print('Test set\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}'.format(len(test), test.SessionId.nunique(), test.ItemId.nunique()))\n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.txt'), sep='\\t', index=False)\n",
    "\n",
    "    def slice_data(self):\n",
    "        for slice_id in range(0, self.num_slices):\n",
    "            self.split_data_slice(slice_id, self.days_offset+(slice_id*self.days_shift))\n",
    "\n",
    "    def split_data_slice(self, slice_id, days_offset):\n",
    "        data = self.data\n",
    "        data_start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc)\n",
    "        data_end = datetime.datetime.fromtimestamp(data.Time.max(), timezone.utc)\n",
    "        \n",
    "        print('Full data set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}'.\n",
    "            format(slice_id, len(data), data.SessionId.nunique(), data.ItemId.nunique(), data_start.isoformat(), data_end.isoformat()))    \n",
    "        \n",
    "        start = datetime.datetime.fromtimestamp(data.Time.min(), timezone.utc ) + timedelta(days_offset) \n",
    "        middle =  start + timedelta(self.days_train)\n",
    "        end =  middle + timedelta(self.days_test)\n",
    "        \n",
    "        #prefilter the timespan\n",
    "        session_max_times = data.groupby('SessionId').Time.max()\n",
    "        greater_start = session_max_times[session_max_times >= start.timestamp()].index\n",
    "        lower_end = session_max_times[session_max_times <= end.timestamp()].index\n",
    "        data_filtered = data[np.in1d(data.SessionId, greater_start.intersection(lower_end))]\n",
    "        \n",
    "        print('Slice data set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {} / {}'.\n",
    "            format( slice_id, len(data_filtered), data_filtered.SessionId.nunique(), data_filtered.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat(), end.date().isoformat() ) )\n",
    "        \n",
    "        #split to train and test\n",
    "        session_max_times = data_filtered.groupby('SessionId').Time.max()\n",
    "        sessions_train = session_max_times[session_max_times < middle.timestamp()].index\n",
    "        sessions_test = session_max_times[session_max_times >= middle.timestamp()].index\n",
    "        \n",
    "        train = data[np.in1d(data.SessionId, sessions_train)]\n",
    "        \n",
    "        print('Train set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {}'.\n",
    "            format( slice_id, len(train), train.SessionId.nunique(), train.ItemId.nunique(), start.date().isoformat(), middle.date().isoformat() ) )\n",
    "        \n",
    "        train.to_csv(osp.join(self.processed_dir,'events_train_full.'+str(slice_id)+'.txt'), sep='\\t', index=False)\n",
    "        \n",
    "        test = data[np.in1d(data.SessionId, sessions_test)]\n",
    "        test = test[np.in1d(test.ItemId, train.ItemId)]\n",
    "        \n",
    "        tslength = test.groupby('SessionId').size()\n",
    "        test = test[np.in1d(test.SessionId, tslength[tslength>=2].index)]\n",
    "        \n",
    "        print('Test set {}\\n\\tEvents: {}\\n\\tSessions: {}\\n\\tItems: {}\\n\\tSpan: {} / {} \\n\\n'.\n",
    "            format( slice_id, len(test), test.SessionId.nunique(), test.ItemId.nunique(), middle.date().isoformat(), end.date().isoformat() ) )\n",
    "        \n",
    "        test.to_csv(osp.join(self.processed_dir,'events_test.'+str(slice_id)+'.txt'), sep='\\t', index=False)\n",
    "        \n",
    "    def process(self):\n",
    "        self.load()\n",
    "        self.filter_data()\n",
    "        if self.process_method == 'last':\n",
    "            self.split_data_org()\n",
    "        elif self.process_method == 'last_min_date':\n",
    "            self.filter_min_date()\n",
    "            self.split_data_org()\n",
    "        elif self.process_method == 'days_test':\n",
    "            self.split_data()\n",
    "        elif self.process_method == 'slice':\n",
    "            self.slice_data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Graph Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class GraphDataset():\n",
    "    def __init__(self, data, shuffle=False, graph=None):\n",
    "        inputs = data[0]\n",
    "        inputs, mask, len_max = self.data_masks(inputs, [0])\n",
    "        self.inputs = np.asarray(inputs)\n",
    "        self.mask = np.asarray(mask)\n",
    "        self.len_max = len_max\n",
    "        self.targets = np.asarray(data[1])\n",
    "        self.length = len(inputs)\n",
    "        self.shuffle = shuffle\n",
    "        self.graph = graph\n",
    "\n",
    "    def generate_batch(self, batch_size):\n",
    "        if self.shuffle:\n",
    "            shuffled_arg = np.arange(self.length)\n",
    "            np.random.shuffle(shuffled_arg)\n",
    "            self.inputs = self.inputs[shuffled_arg]\n",
    "            self.mask = self.mask[shuffled_arg]\n",
    "            self.targets = self.targets[shuffled_arg]\n",
    "        n_batch = int(self.length / batch_size)\n",
    "        if self.length % batch_size != 0:\n",
    "            n_batch += 1\n",
    "        slices = np.split(np.arange(n_batch * batch_size), n_batch)\n",
    "        slices[-1] = slices[-1][:(self.length - batch_size * (n_batch - 1))]\n",
    "        return slices\n",
    "\n",
    "    def get_slice(self, i):\n",
    "        inputs, mask, targets = self.inputs[i], self.mask[i], self.targets[i]\n",
    "        items, n_node, A, alias_inputs = [], [], [], []\n",
    "        for u_input in inputs:\n",
    "            n_node.append(len(np.unique(u_input)))\n",
    "        max_n_node = np.max(n_node)\n",
    "        for u_input in inputs:\n",
    "            node = np.unique(u_input)\n",
    "            items.append(node.tolist() + (max_n_node - len(node)) * [0])\n",
    "            u_A = np.zeros((max_n_node, max_n_node))\n",
    "            for i in np.arange(len(u_input) - 1):\n",
    "                if u_input[i + 1] == 0:\n",
    "                    break\n",
    "                u = np.where(node == u_input[i])[0][0]\n",
    "                v = np.where(node == u_input[i + 1])[0][0]\n",
    "                u_A[u][v] = 1\n",
    "            u_sum_in = np.sum(u_A, 0)\n",
    "            u_sum_in[np.where(u_sum_in == 0)] = 1\n",
    "            u_A_in = np.divide(u_A, u_sum_in)\n",
    "            u_sum_out = np.sum(u_A, 1)\n",
    "            u_sum_out[np.where(u_sum_out == 0)] = 1\n",
    "            u_A_out = np.divide(u_A.transpose(), u_sum_out)\n",
    "            u_A = np.concatenate([u_A_in, u_A_out]).transpose()\n",
    "            A.append(u_A)\n",
    "            alias_inputs.append([np.where(node == i)[0][0] for i in u_input])\n",
    "        return alias_inputs, A, items, mask, targets\n",
    "\n",
    "    @staticmethod\n",
    "    def data_masks(all_usr_pois, item_tail):\n",
    "        us_lens = [len(upois) for upois in all_usr_pois]\n",
    "        len_max = max(us_lens)\n",
    "        us_pois = [upois + item_tail * (len_max - le) for upois, le in zip(all_usr_pois, us_lens)]\n",
    "        us_msks = [[1] * le + [0] * (len_max - le) for le in us_lens]\n",
    "        return us_pois, us_msks, len_max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([0]), array([1]), array([2]), array([3]), array([4])]\n",
      "[array([0, 1]), array([2, 3]), array([4])]\n",
      "[[1 2 3]\n",
      " [2 3 4]\n",
      " [1 2 4]\n",
      " [2 3 0]\n",
      " [1 0 0]]\n",
      "[array([0]), array([1]), array([2]), array([3]), array([4])]\n",
      "[array([0, 1]), array([2, 3]), array([4])]\n",
      "[[1 2 4]\n",
      " [2 3 4]\n",
      " [1 0 0]\n",
      " [2 3 0]\n",
      " [1 2 3]]\n"
     ]
    }
   ],
   "source": [
    "train_data = ([[1, 2, 3], [2, 3, 4], [1, 2, 4], [2, 3], [1]], \n",
    "              [4, 5, 5, 4, 2])\n",
    "\n",
    "tds = GraphData(train_data, shuffle=False)\n",
    "print(tds.generate_batch(1))\n",
    "print(tds.generate_batch(2))\n",
    "print(tds.inputs)\n",
    "\n",
    "tds = GraphData(train_data, shuffle=True)\n",
    "print(tds.generate_batch(1))\n",
    "print(tds.generate_batch(2))\n",
    "print(tds.inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SessionGraph Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class SessionGraphDataset(Dataset):\n",
    "    \"\"\"\n",
    "    References\n",
    "        1. COTREC session-based recommender model training. https://t.ly/cXTH.\n",
    "    \"\"\"\n",
    "    def __init__(self, root, shuffle=False, n_node=None):\n",
    "        super().__init__(root)\n",
    "        self.n_node = n_node\n",
    "        self.shuffle = shuffle\n",
    "        self.process()\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        return ['data.txt', 'all_seq.txt']\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self) -> str:\n",
    "        pass\n",
    "\n",
    "    def download(self):\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def process(self):\n",
    "        import pickle\n",
    "        data = pickle.load(open(self.raw_paths[0], 'rb'))\n",
    "        all_seq = pickle.load(open(self.raw_paths[1], 'rb'))\n",
    "        self.raw = np.asarray(data[0])\n",
    "        self.targets = np.asarray(data[1])\n",
    "        self.length = len(self.raw)\n",
    "        adj = self.data_masks(all_seq, self.n_node)\n",
    "        self.adjacency = adj.multiply(1.0/adj.sum(axis=0).reshape(1, -1))\n",
    "\n",
    "\n",
    "    def get_overlap(self, sessions):\n",
    "        matrix = np.zeros((len(sessions), len(sessions)))\n",
    "        for i in range(len(sessions)):\n",
    "            seq_a = set(sessions[i])\n",
    "            seq_a.discard(0)\n",
    "            for j in range(i+1, len(sessions)):\n",
    "                seq_b = set(sessions[j])\n",
    "                seq_b.discard(0)\n",
    "                overlap = seq_a.intersection(seq_b)\n",
    "                ab_set = seq_a | seq_b\n",
    "                matrix[i][j] = float(len(overlap))/float(len(ab_set))\n",
    "                matrix[j][i] = matrix[i][j]\n",
    "        # matrix = self.dropout(matrix, 0.2)\n",
    "        matrix = matrix + np.diag([1.0]*len(sessions))\n",
    "        degree = np.sum(np.array(matrix), 1)\n",
    "        degree = np.diag(1.0/degree)\n",
    "        return matrix, degree\n",
    "\n",
    "    def generate_batch(self, batch_size):\n",
    "        if self.shuffle:\n",
    "            shuffled_arg = np.arange(self.length)\n",
    "            np.random.shuffle(shuffled_arg)\n",
    "            self.raw = self.raw[shuffled_arg]\n",
    "            self.targets = self.targets[shuffled_arg]\n",
    "        n_batch = int(self.length / batch_size)\n",
    "        if self.length % batch_size != 0:\n",
    "            n_batch += 1\n",
    "        slices = np.split(np.arange(n_batch * batch_size), n_batch)\n",
    "        slices[-1] = np.arange(self.length-batch_size, self.length)\n",
    "        return slices\n",
    "\n",
    "    def get_slice(self, index):\n",
    "        items, num_node = [], []\n",
    "        inp = self.raw[index]\n",
    "        for session in inp:\n",
    "            num_node.append(len(np.nonzero(session)[0]))\n",
    "        max_n_node = np.max(num_node)\n",
    "        session_len = []\n",
    "        reversed_sess_item = []\n",
    "        mask = []\n",
    "        # item_set = set()\n",
    "        for session in inp:\n",
    "            nonzero_elems = np.nonzero(session)[0]\n",
    "            # item_set.update(set([t-1 for t in session]))\n",
    "            session_len.append([len(nonzero_elems)])\n",
    "            items.append(session + (max_n_node - len(nonzero_elems)) * [0])\n",
    "            mask.append([1]*len(nonzero_elems) + (max_n_node - len(nonzero_elems)) * [0])\n",
    "            reversed_sess_item.append(list(reversed(session)) + (max_n_node - len(nonzero_elems)) * [0])\n",
    "        # item_set = list(item_set)\n",
    "        # index_list = [item_set.index(a) for a in self.targets[index]-1]\n",
    "        diff_mask = np.ones(shape=[100, self.n_node]) * (1/(self.n_node - 1))\n",
    "        for count, value in enumerate(self.targets[index]-1):\n",
    "            diff_mask[count][value] = 1\n",
    "        return self.targets[index]-1, session_len,items, reversed_sess_item, mask, diff_mask\n",
    "    \n",
    "    @staticmethod\n",
    "    def data_masks(all_sessions, n_node):\n",
    "        adj = dict()\n",
    "        for sess in all_sessions:\n",
    "            for i, item in enumerate(sess):\n",
    "                if i == len(sess)-1:\n",
    "                    break\n",
    "                else:\n",
    "                    if sess[i] - 1 not in adj.keys():\n",
    "                        adj[sess[i]-1] = dict()\n",
    "                        adj[sess[i]-1][sess[i]-1] = 1\n",
    "                        adj[sess[i]-1][sess[i+1]-1] = 1\n",
    "                    else:\n",
    "                        if sess[i+1]-1 not in adj[sess[i]-1].keys():\n",
    "                            adj[sess[i] - 1][sess[i + 1] - 1] = 1\n",
    "                        else:\n",
    "                            adj[sess[i]-1][sess[i+1]-1] += 1\n",
    "        row, col, data = [], [], []\n",
    "        for i in adj.keys():\n",
    "            item = adj[i]\n",
    "            for j in item.keys():\n",
    "                row.append(i)\n",
    "                col.append(j)\n",
    "                data.append(adj[i][j])\n",
    "        from scipy.sparse import coo_matrix\n",
    "        coo = coo_matrix((data, (row, col)), shape=(n_node, n_node))\n",
    "        return coo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DigineticaDataset(SessionGraphDataset):\n",
    "    train_url = \"https://github.com/RecoHut-Datasets/diginetica/raw/v2/train.txt\"\n",
    "    test_url = \"https://github.com/RecoHut-Datasets/diginetica/raw/v2/test.txt\"\n",
    "    all_train_seq_url = \"https://github.com/RecoHut-Datasets/diginetica/raw/v2/all_train_seq.txt\"\n",
    "\n",
    "    def __init__(self, root, shuffle=False, n_node=43097, is_train=True):\n",
    "        self.n_node = n_node\n",
    "        self.shuffle = shuffle\n",
    "        self.is_train = is_train\n",
    "        super().__init__(root, shuffle, n_node)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        if self.is_train:\n",
    "            return ['train.txt', 'all_train_seq.txt']\n",
    "        return ['test.txt', 'all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.all_train_seq_url, self.raw_dir)\n",
    "        if self.is_train:\n",
    "            download_url(self.train_url, self.raw_dir)\n",
    "        else:\n",
    "            download_url(self.test_url, self.raw_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/content/diginetica'\n",
    "\n",
    "train_data = DigineticaDataset(root=root, shuffle=True, is_train=True)\n",
    "test_data = DigineticaDataset(root=root, shuffle=False, is_train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TmallDataset(SessionGraphDataset):\n",
    "    train_url = \"https://github.com/RecoHut-Datasets/tmall/raw/v1/train.txt\"\n",
    "    test_url = \"https://github.com/RecoHut-Datasets/tmall/raw/v1/test.txt\"\n",
    "    all_train_seq_url = \"https://github.com/RecoHut-Datasets/tmall/raw/v1/all_train_seq.txt\"\n",
    "\n",
    "    def __init__(self, root, shuffle=False, n_node=40727, is_train=True):\n",
    "        self.n_node = n_node\n",
    "        self.shuffle = shuffle\n",
    "        self.is_train = is_train\n",
    "        super().__init__(root, shuffle, n_node)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        if self.is_train:\n",
    "            return ['train.txt', 'all_train_seq.txt']\n",
    "        return ['test.txt', 'all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.all_train_seq_url, self.raw_dir)\n",
    "        if self.is_train:\n",
    "            download_url(self.train_url, self.raw_dir)\n",
    "        else:\n",
    "            download_url(self.test_url, self.raw_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n",
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    }
   ],
   "source": [
    "root = '/content/tmall'\n",
    "\n",
    "train_data = TmallDataset(root=root, shuffle=True, is_train=True)\n",
    "test_data = TmallDataset(root=root, shuffle=False, is_train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RetailRocketDataset(SessionGraphDataset):\n",
    "    train_url = \"https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/train.txt\"\n",
    "    test_url = \"https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/test.txt\"\n",
    "    all_train_seq_url = \"https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/all_train_seq.txt\"\n",
    "\n",
    "    def __init__(self, root, shuffle=False, n_node=40727, is_train=True):\n",
    "        self.n_node = n_node\n",
    "        self.shuffle = shuffle\n",
    "        self.is_train = is_train\n",
    "        super().__init__(root, shuffle, n_node)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        if self.is_train:\n",
    "            return ['train.txt', 'all_train_seq.txt']\n",
    "        return ['test.txt', 'all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.all_train_seq_url, self.raw_dir)\n",
    "        if self.is_train:\n",
    "            download_url(self.train_url, self.raw_dir)\n",
    "        else:\n",
    "            download_url(self.test_url, self.raw_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/all_train_seq.txt\n",
      "Downloading https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/train.txt\n",
      "Using existing file all_train_seq.txt\n",
      "Downloading https://github.com/RecoHut-Datasets/retail_rocket/raw/v1/test.txt\n"
     ]
    }
   ],
   "source": [
    "root = '/content/retail_rocket'\n",
    "\n",
    "train_data = RetailRocketDataset(root=root, shuffle=True, is_train=True)\n",
    "test_data = RetailRocketDataset(root=root, shuffle=False, is_train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SampleDataset(SessionGraphDataset):\n",
    "    train_url = \"https://github.com/RecoHut-Datasets/sample_session/raw/v2/train.txt\"\n",
    "    test_url = \"https://github.com/RecoHut-Datasets/sample_session/raw/v2/test.txt\"\n",
    "    all_train_seq_url = \"https://github.com/RecoHut-Datasets/sample_session/raw/v2/all_train_seq.txt\"\n",
    "\n",
    "    def __init__(self, root, shuffle=False, n_node=309, is_train=True):\n",
    "        self.n_node = n_node\n",
    "        self.shuffle = shuffle\n",
    "        self.is_train = is_train\n",
    "        super().__init__(root, shuffle, n_node)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self) -> str:\n",
    "        if self.is_train:\n",
    "            return ['train.txt', 'all_train_seq.txt']\n",
    "        return ['test.txt', 'all_train_seq.txt']\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.all_train_seq_url, self.raw_dir)\n",
    "        if self.is_train:\n",
    "            download_url(self.train_url, self.raw_dir)\n",
    "        else:\n",
    "            download_url(self.test_url, self.raw_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/sample_session/raw/v2/all_train_seq.txt\n",
      "Downloading https://github.com/RecoHut-Datasets/sample_session/raw/v2/train.txt\n",
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n",
      "Using existing file all_train_seq.txt\n",
      "Downloading https://github.com/RecoHut-Datasets/sample_session/raw/v2/test.txt\n",
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    }
   ],
   "source": [
    "root = '/content/sample'\n",
    "\n",
    "train_data = SampleDataset(root=root, shuffle=True, is_train=True)\n",
    "test_data = SampleDataset(root=root, shuffle=False, is_train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rating Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class RatingDataset(Dataset):\n",
    "    r\"\"\"Interaction data with rating feedback\n",
    "    \n",
    "    Args:\n",
    "        root: data folder path\n",
    "        min_uc: minimum user count to keep in the data\n",
    "        min_sc: minimum item count to keep in the data\n",
    "        split: data split method - leave_one_out/holdout\n",
    "        min_rating: minimum rating threshold to convert explicit feedback into implicit\n",
    "\n",
    "    References:\n",
    "        1. https://github.com/Yueeeeeeee/RecSys-Extraction-Attack/tree/main/datasets\n",
    "    \"\"\"\n",
    "    def __init__(self, root, min_uc, min_sc, split='leave_one_out', dataset_split_seed=42,\n",
    "                 eval_set_size=None, min_rating=None, iterative_triplet=False):\n",
    "        super().__init__(root)\n",
    "        self.min_rating = min_rating\n",
    "        self.min_uc = min_uc\n",
    "        self.min_sc = min_sc\n",
    "        self.split = split\n",
    "        self.dataset_split_seed = dataset_split_seed\n",
    "        self.eval_set_size = eval_set_size\n",
    "        self.iterative_triplet = iterative_triplet\n",
    "\n",
    "        assert self.min_uc >= 2, 'Need at least 2 ratings per user for validation and test'\n",
    "\n",
    "        self._process()\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        r\"\"\"load raw dataset into pandas dataframe\"\"\"\n",
    "        raise NotImplementedError\n",
    "\n",
    "    def load_processed(self):\n",
    "        return pickle.load(open(self.processed_paths[0], 'rb'))\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return 'dataset.pkl'\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 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 filter_triplets_iteratively(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",
    "            \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.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.split == 'holdout':\n",
    "            print('Splitting')\n",
    "            np.random.seed(self.dataset_split_seed)\n",
    "            eval_set_size = self.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 process(self):\n",
    "        df = self.load_ratings_df()\n",
    "        if self.iterative_triplet:\n",
    "            df = self.filter_triplets_iteratively(df)\n",
    "        else:\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 open(self.processed_paths[0], 'wb') as f:\n",
    "            pickle.dump(dataset, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AmazonGamesDataset(RatingDataset):\n",
    "    url = \"http://snap.stanford.edu/data/amazon/productGraph/categoryFiles/ratings_Video_Games.csv\"\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return 'ratings_Video_Games.csv'\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.url, self.raw_dir)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        df = pd.read_csv(self.raw_paths[0], header=None)\n",
    "        df.columns = ['uid', 'sid', 'rating', 'timestamp']\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7519/7519 [00:02<00:00, 2512.75it/s]\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = AmazonGamesDataset(root='/content/amazon_games', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AmazonBeautyDataset(RatingDataset):\n",
    "    url = \"http://snap.stanford.edu/data/amazon/productGraph/categoryFiles/ratings_Beauty.csv\"\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return 'ratings_Beauty.csv'\n",
    "\n",
    "    def download(self):\n",
    "        download_url(self.url, self.raw_dir)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        df = pd.read_csv(self.raw_paths[0], header=None)\n",
    "        df.columns = ['uid', 'sid', 'rating', 'timestamp']\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading http://snap.stanford.edu/data/amazon/productGraph/categoryFiles/ratings_Video_Games.csv\n",
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7519/7519 [00:02<00:00, 2527.69it/s]\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = AmazonGamesDataset(root='/content/amazon_beauty', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ML1mDataset(RatingDataset):\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",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 6040/6040 [00:02<00:00, 2590.97it/s]\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = ML1mDataset(root='/content/ML1m', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SteamGamesDataset(RatingDataset):\n",
    "    url = \"http://cseweb.ucsd.edu/~wckang/steam_reviews.json.gz\"\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return 'steam_reviews.json'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        extract_gz(path, self.raw_dir)\n",
    "        os.unlink(path)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        data = []\n",
    "        f = open(self.raw_paths[0], 'r', encoding='utf-8')\n",
    "        import ast\n",
    "        for line in f.readlines():\n",
    "            temp = ast.literal_eval(line)\n",
    "            data.append([temp['username'], temp['product_id'], temp['date']])\n",
    "\n",
    "        return pd.DataFrame(data, columns=['uid', 'sid', 'timestamp'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 120145/120145 [01:10<00:00, 1709.62it/s]\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = SteamGamesDataset(root='/content/steam', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class YoochooseDataset(RatingDataset):\n",
    "    url = \"https://s3-eu-west-1.amazonaws.com/yc-rdata/yoochoose-data.7z\"\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return 'yoochoose-clicks.dat'\n",
    "\n",
    "    def download(self):\n",
    "        path = download_url(self.url, self.raw_dir)\n",
    "        # pip install pyunpack patool\n",
    "        import pyunpack\n",
    "        pyunpack.Archive(path).extractall(self.raw_dir)\n",
    "        os.unlink(path)\n",
    "\n",
    "    def load_ratings_df(self):\n",
    "        df = pd.read_csv(self.raw_paths[0], header=None)\n",
    "        df.columns = ['uid', 'timestamp', 'sid', 'category']\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processing...\n",
      "/usr/local/lib/python3.7/dist-packages/ipykernel_launcher.py:28: DtypeWarning: Columns (3) have mixed types.Specify dtype option on import or set low_memory=False.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Filtering triplets\n",
      "Densifying index\n",
      "Splitting\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 449961/449961 [03:55<00:00, 1913.52it/s]\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "ds = YoochooseDataset(root='/content/yoochoose', min_uc=10, min_sc=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Torch Datasets\n",
    "Datasets in PyTorch format."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Session Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TorchSessionDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, fpath, maxlen, is_train=True):\n",
    "            [train, valid, test, itemnum] = self.data_partition(fpath)\n",
    "            print(\"Number of sessions:\",len(train)+len(valid)+len(test))\n",
    "            print(\"Number of items:\", itemnum)\n",
    "\n",
    "            action = 0\n",
    "            for i in train:\n",
    "                action += np.count_nonzero(i)\n",
    "            for i in valid:\n",
    "                action += np.count_nonzero(i)\n",
    "            for i in test:\n",
    "                action += np.count_nonzero(i)\n",
    "\n",
    "            print(\"Number of actions:\", action)\n",
    "            print(\"Average length of sessions:\", action/(len(train)+len(valid)+len(test)))\n",
    "\n",
    "            self.data = train if is_train else test\n",
    "            self.maxlen = maxlen\n",
    "            self.itemnum = itemnum\n",
    "            self.is_train = is_train\n",
    "\n",
    "    def __len__(self):\n",
    "            return len(self.data)\n",
    "\n",
    "    def __train__(self, index):\n",
    "            session = np.asarray(self.data[index], dtype=np.int64)\n",
    "            if len(session) > self.maxlen:\n",
    "                session = session[-self.maxlen:]\n",
    "            else:\n",
    "                session = np.pad(session, (self.maxlen-len(session), 0), 'constant', constant_values=0)\n",
    "            curr_seq = session[:-1]\n",
    "            curr_pos = session[1:]\n",
    "            return curr_seq, curr_pos\n",
    "    \n",
    "    def __test__(self, index):\n",
    "            session = self.data[index]\n",
    "            seq = np.zeros([self.maxlen], dtype=np.int64)\n",
    "            idx = self.maxlen - 1\n",
    "            for i in reversed(session[:-1]): #everything except the last one\n",
    "                seq[idx] = i\n",
    "                idx -= 1\n",
    "                if idx == -1: break\n",
    "            return seq, session[-1]-1 #index of the item in the list of all items\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "            if self.is_train:\n",
    "                return self.__train__(index)\n",
    "            else:\n",
    "                return self.__test__(index)\n",
    "\n",
    "    @staticmethod\n",
    "    def data_partition(fname, percentage=[0.1, 0.2]):\n",
    "        itemnum = 0\n",
    "\n",
    "        sessions = defaultdict(list)\n",
    "        session_train = []\n",
    "        session_valid = []\n",
    "        session_test = []\n",
    "        # assume user/item index starting from 1\n",
    "        session_id = 0\n",
    "        f = open(fname, 'r')\n",
    "        total_length = 0\n",
    "        max_length = 0\n",
    "        for line in f:\n",
    "\n",
    "            items = [int(l) for l in line.rstrip().split(',')]\n",
    "\n",
    "            if len(items) < 5: continue\n",
    "            total_length += len(items)\n",
    "\n",
    "            if max_length< len(items):\n",
    "                max_length = len(items)\n",
    "            \n",
    "            itemnum = max(max(items), itemnum)\n",
    "            sessions[session_id].append(items)\n",
    "            session_id += 1\n",
    "\n",
    "        print(\"Avg length:\", total_length/session_id)\n",
    "        print(\"Maximum length:\", max_length)\n",
    "\n",
    "        valid_perc = percentage[0]\n",
    "        test_perc = percentage[1]\n",
    "\n",
    "        total_sessions = session_id\n",
    "        \n",
    "        shuffle_indices = np.random.permutation(range(total_sessions)) #\n",
    "        \n",
    "        train_index = int(total_sessions*(1 - valid_perc - test_perc))\n",
    "        valid_index = int(total_sessions*(1 - test_perc))\n",
    "\n",
    "        if (train_index == valid_index): valid_index += 1 #break the tie\n",
    "        \n",
    "        train_indices = shuffle_indices[:train_index]\n",
    "        valid_indices = shuffle_indices[train_index:valid_index]\n",
    "        test_indices = shuffle_indices[valid_index:]\n",
    "\n",
    "        for i in train_indices:\n",
    "            session_train.extend(sessions[i])\n",
    "        for i in valid_indices:\n",
    "            session_valid.extend(sessions[i])\n",
    "        for i in test_indices:\n",
    "            session_test.extend(sessions[i])\n",
    "        \n",
    "        return [np.asarray(session_train), np.asarray(session_valid), np.asarray(session_test), itemnum]\n",
    "\n",
    "    @staticmethod\n",
    "    def nextitnet_format(fname, maxlen):\n",
    "            \n",
    "        sessions = []\n",
    "\n",
    "        # assume user/item index starting from 1\n",
    "        f = open(fname, 'r')\n",
    "\n",
    "        for line in f:\n",
    "\n",
    "            items = [int(l) for l in line.rstrip().split(',')]\n",
    "\n",
    "            if len(items) < 5: continue\n",
    "            \n",
    "            seq = np.zeros([maxlen], dtype=np.int32)\n",
    "            \n",
    "            idx = maxlen - 1\n",
    "\n",
    "            for i in reversed(items):\n",
    "                seq[idx] = i\n",
    "                idx -= 1\n",
    "                if idx == -1: break        \n",
    "            \n",
    "            sessions.append(seq)\n",
    "            \n",
    "        print(\"number of session:\", len(sessions))\n",
    "\n",
    "        return sessions\n",
    "\n",
    "    @staticmethod\n",
    "    def gru_format(fname, user_train, user_valid, user_test):\n",
    "        \n",
    "        session_id = 0\n",
    "        train = []\n",
    "        for session in user_train:\n",
    "            for item in session:\n",
    "                train.append([session_id, item, 0])\n",
    "            session_id += 1\n",
    "\n",
    "        valid = []\n",
    "        for session in user_valid:\n",
    "            for item in session:\n",
    "                valid.append([session_id, item, 0])\n",
    "            session_id += 1\n",
    "\n",
    "        test = []\n",
    "        for session in user_test:\n",
    "            for item in session:\n",
    "                test.append([session_id, item, 0])\n",
    "            session_id += 1\n",
    "\n",
    "        train_data = pd.DataFrame(train, columns= ['SessionId', 'ItemId', 'Time'])\n",
    "        valid_data = pd.DataFrame(valid, columns= ['SessionId', 'ItemId', 'Time'])\n",
    "        test_data = pd.DataFrame(test, columns= ['SessionId', 'ItemId', 'Time'])\n",
    "\n",
    "        return train_data, valid_data, test_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class YoochooseDataset(TorchSessionDataset):\n",
    "    url = 'https://github.com/RecoHut-Datasets/yoochoose/raw/v3/yoochoose.csv'\n",
    "\n",
    "    def __init__(self, root, maxlen, is_train=True):\n",
    "        fpath = download_url(url=self.url, folder=root)\n",
    "        super().__init__(fpath, maxlen, is_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using existing file yoochoose.csv\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg length: 10.0\n",
      "Maximum length: 10\n",
      "Number of sessions: 80183\n",
      "Number of items: 12936\n",
      "Number of actions: 406979\n",
      "Average length of sessions: 5.075627003230111\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0, 10309, 10309, 10309],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,   794,  5005,  6891],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0, 10631,  4104,  9852],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,  9469,  9486,  9469],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,   155,  8790,  6931],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,   770, 11239,  6040],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0, 11641, 11610, 12033],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,  8449,  5705, 10331,   170,  8485]]),\n",
       " tensor([[    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0, 10309, 10309, 10309, 10309],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,   794,  5005,  6891,  6501],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0, 10631,  4104,  9852, 10007],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,  9469,  9486,  9469,  9486],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,   155,  8790,  6931,  8821],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,   770, 11239,  6040, 11235],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0, 11641, 11610, 12033, 11638],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,  8449,  5705, 10331,   170,  8485, 10332]])]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = YoochooseDataset(root='/content/yoochoose', maxlen=30)\n",
    "\n",
    "sampler = torch.utils.data.DataLoader(dataset, batch_size=8, num_workers=2, pin_memory=True)\n",
    "samples = next(iter(sampler))\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NowplayingDataset(TorchSessionDataset):\n",
    "    url = 'https://github.com/RecoHut-Datasets/nowplaying/raw/v3/nowplaying.csv'\n",
    "\n",
    "    def __init__(self, root, maxlen, is_train=True):\n",
    "        fpath = download_url(url=self.url, folder=root)\n",
    "        super().__init__(fpath, maxlen, is_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/nowplaying/raw/v3/nowplaying.csv\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg length: 20.0\n",
      "Maximum length: 20\n",
      "Number of sessions: 113918\n",
      "Number of items: 239221\n",
      "Number of actions: 1184815\n",
      "Average length of sessions: 10.400595164943205\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[tensor([[     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0, 114002, 113983,  89621, 113960, 113884, 113926, 114000, 113738,\n",
       "          113930,   3168, 113805, 113800, 113789, 113872, 114018, 113881, 113869,\n",
       "          113776,  21568],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,  13653,  11910,  28131,\n",
       "            4896,  33231],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0, 217911, 218397,  23439,  23684,  40048,  23439,  22123, 218298,\n",
       "           58345, 218399],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,  16292,   3786,  45272,   3574,  28015,  16926,\n",
       "           27992,  33024],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "           76624,  76624, 113070, 113070,  76624,  76624, 113070, 113070,  76624,\n",
       "           76624, 113070],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0, 127696,  39067,  82151,  29706,  29201,  29605,   4791,  29298,\n",
       "          127939,  29456,  29779, 109896,   5945,  73638, 127962,  44011,  29721,\n",
       "          127625, 114913],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,  32544, 155346, 155508, 155074, 155886, 155223, 155360, 155356,\n",
       "          154929, 154914, 155887, 154877, 155115, 155888,  75852, 154969, 155889,\n",
       "          155291, 155890],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,  56331,  56332,\n",
       "           56333,  56334,  56335,  13547,  56336,   7363,  56337,  56338,  56339,\n",
       "            4735,   9554]]),\n",
       " tensor([[     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "          114002, 113983,  89621, 113960, 113884, 113926, 114000, 113738, 113930,\n",
       "            3168, 113805, 113800, 113789, 113872, 114018, 113881, 113869, 113776,\n",
       "           21568, 113925],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,  13653,  11910,  28131,   4896,\n",
       "           33231,   8409],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "          217911, 218397,  23439,  23684,  40048,  23439,  22123, 218298,  58345,\n",
       "          218399, 218330],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,  16292,   3786,  45272,   3574,  28015,  16926,  27992,\n",
       "           33024,  24359],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,  76624,\n",
       "           76624, 113070, 113070,  76624,  76624, 113070, 113070,  76624,  76624,\n",
       "          113070, 113070],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "          127696,  39067,  82151,  29706,  29201,  29605,   4791,  29298, 127939,\n",
       "           29456,  29779, 109896,   5945,  73638, 127962,  44011,  29721, 127625,\n",
       "          114913, 127608],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "           32544, 155346, 155508, 155074, 155886, 155223, 155360, 155356, 154929,\n",
       "          154914, 155887, 154877, 155115, 155888,  75852, 154969, 155889, 155291,\n",
       "          155890, 155891],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,  56331,  56332,  56333,\n",
       "           56334,  56335,  13547,  56336,   7363,  56337,  56338,  56339,   4735,\n",
       "            9554,  56340]])]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = NowplayingDataset(root='/content/nowplaying', maxlen=30)\n",
    "\n",
    "sampler = torch.utils.data.DataLoader(dataset, batch_size=8, num_workers=2, pin_memory=True)\n",
    "samples = next(iter(sampler))\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DigineticaDataset(TorchSessionDataset):\n",
    "    url = 'https://github.com/RecoHut-Datasets/diginetica/raw/v4/diginetica.csv'\n",
    "\n",
    "    def __init__(self, root, maxlen, is_train=True):\n",
    "        fpath = download_url(url=self.url, folder=root)\n",
    "        super().__init__(fpath, maxlen, is_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/diginetica/raw/v4/diginetica.csv\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg length: 8.777109003245833\n",
      "Maximum length: 70\n",
      "Number of sessions: 63466\n",
      "Number of items: 38970\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of actions: 557048\n",
      "Average length of sessions: 8.777109003245833\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[tensor([[    0,     0,     0,     0,     0,     0,     0,     0,  2387,  2245,\n",
       "           9141,  2366,  9142,  9143,  3193,  3193,  1726,  1725,  2366,  1722,\n",
       "           2366,  2366,  9144,  3197,  9145,  1722,  9146,  9147,  9146],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0, 17095, 17101, 17094, 17100, 17096],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0, 13816,\n",
       "          10789,  9204, 11198, 23151,  8289, 30676,  3372, 30678, 14125],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,  4814, 10013, 18788,  9285, 14081],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0, 13064, 36257, 30911, 11052],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0, 28073, 19585,\n",
       "          17214, 13842, 28815, 28815, 13842, 17214, 19585, 16278, 15659],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0, 16291, 17881,  5630, 20969, 20969, 20829, 12938],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,  7298, 25445,  1397, 25445,  7298,  3191,  1735]]),\n",
       " tensor([[    0,     0,     0,     0,     0,     0,     0,  2387,  2245,  9141,\n",
       "           2366,  9142,  9143,  3193,  3193,  1726,  1725,  2366,  1722,  2366,\n",
       "           2366,  9144,  3197,  9145,  1722,  9146,  9147,  9146,  9148],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0, 17095, 17101, 17094, 17100, 17096, 17102],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0, 13816, 10789,\n",
       "           9204, 11198, 23151,  8289, 30676,  3372, 30678, 14125, 14124],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,  4814, 10013, 18788,  9285, 14081,   336],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0, 13064, 36257, 30911, 11052,  7098],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0, 28073, 19585, 17214,\n",
       "          13842, 28815, 28815, 13842, 17214, 19585, 16278, 15659, 28072],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0, 16291, 17881,  5630, 20969, 20969, 20829, 12938, 12938],\n",
       "         [    0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
       "              0,  7298, 25445,  1397, 25445,  7298,  3191,  1735, 25445]])]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = DigineticaDataset(root='/content/diginetica', maxlen=30)\n",
    "\n",
    "sampler = torch.utils.data.DataLoader(dataset, batch_size=8, num_workers=2, pin_memory=True)\n",
    "samples = next(iter(sampler))\n",
    "samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LastfmDataset(TorchSessionDataset):\n",
    "    url = 'https://github.com/RecoHut-Datasets/lastfm/raw/v2/last_fm.csv'\n",
    "\n",
    "    def __init__(self, root, maxlen, is_train=True):\n",
    "        fpath = download_url(url=self.url, folder=root)\n",
    "        super().__init__(fpath, maxlen, is_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://github.com/RecoHut-Datasets/lastfm/raw/v2/last_fm.csv\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Avg length: 17.447849599510228\n",
      "Maximum length: 49\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.7/dist-packages/numpy/core/_asarray.py:83: VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences (which is a list-or-tuple of lists-or-tuples-or ndarrays with different lengths or shapes) is deprecated. If you meant to do this, you must specify 'dtype=object' when creating the ndarray\n",
      "  return array(a, dtype, copy=False, order=order)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of sessions: 196010\n",
      "Number of items: 107391\n",
      "Number of actions: 3419953\n",
      "Average length of sessions: 17.447849599510228\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[tensor([[     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,   3821,     96,   3821,   1600,     96,\n",
       "            3366,   3821,     96,   3366,  18639,   3821,   3280,   3366,     96,\n",
       "            3366,   1600],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,  13864,\n",
       "           35393,  13864,  50765,  13743,  51628,  34165,  44702,  62996,   9504,\n",
       "          106404,  13864],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,   3363,   1875,\n",
       "            2782,   1875],\n",
       "         [   973,   2740,   2712,  17892,   2228,    829,   2740,    128,    744,\n",
       "            1193,   1284,   2755,   1443,   4028,   2712,   7635,    620,   1861,\n",
       "            3978,    790,    916,   1455,    227,  10492,   1257,    633,  29590,\n",
       "            2631,   8419],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,  13576,   9525,   9513,\n",
       "            9857,   9512,   9180,   9335,   1953,   9504,   6466,  19980,  19966,\n",
       "           71404,  90684],\n",
       "         [  8405,   5743,   8106,  35491,    675,   3393,   8239,   3469,   8239,\n",
       "             675,   8405,   8239,   2895,  19781,  16754,   3572,  16754,  16088,\n",
       "             637,   3713,   8171,    675,  15850,   3135,   3713,   8509,   1903,\n",
       "            1900,  13521],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,   8771,   8510,   1409,   3180,   8407,   5464,\n",
       "             479,    707],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,   3366,    987,    988,    987,  47629,    707,   1644,\n",
       "            3775,   1291]]),\n",
       " tensor([[     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,   3821,     96,   3821,   1600,     96,   3366,\n",
       "            3821,     96,   3366,  18639,   3821,   3280,   3366,     96,   3366,\n",
       "            1600,   3366],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,  13864,  35393,\n",
       "           13864,  50765,  13743,  51628,  34165,  44702,  62996,   9504, 106404,\n",
       "           13864,  20385],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,   3363,   1875,   2782,\n",
       "            1875,   9683],\n",
       "         [  2740,   2712,  17892,   2228,    829,   2740,    128,    744,   1193,\n",
       "            1284,   2755,   1443,   4028,   2712,   7635,    620,   1861,   3978,\n",
       "             790,    916,   1455,    227,  10492,   1257,    633,  29590,   2631,\n",
       "            8419,   8555],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,  13576,   9525,   9513,   9857,\n",
       "            9512,   9180,   9335,   1953,   9504,   6466,  19980,  19966,  71404,\n",
       "           90684,  71404],\n",
       "         [  5743,   8106,  35491,    675,   3393,   8239,   3469,   8239,    675,\n",
       "            8405,   8239,   2895,  19781,  16754,   3572,  16754,  16088,    637,\n",
       "            3713,   8171,    675,  15850,   3135,   3713,   8509,   1903,   1900,\n",
       "           13521,   8509],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,   8771,   8510,   1409,   3180,   8407,   5464,    479,\n",
       "             707,    937],\n",
       "         [     0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,      0,      0,      0,      0,      0,      0,      0,      0,\n",
       "               0,   3366,    987,    988,    987,  47629,    707,   1644,   3775,\n",
       "            1291,   3775]])]"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = LastfmDataset(root='/content/lastfm', maxlen=30)\n",
    "\n",
    "sampler = torch.utils.data.DataLoader(dataset, batch_size=8, num_workers=2, pin_memory=True)\n",
    "samples = next(iter(sampler))\n",
    "samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rating Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TorchRatingDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"Wrapper, convert <user, item, rating> Tensor into Pytorch Dataset.\"\"\"\n",
    "\n",
    "    def __init__(self, user_tensor, item_tensor, target_tensor):\n",
    "        \"\"\"Init UserItemRatingDataset Class.\n",
    "        Args:\n",
    "            target_tensor: torch.Tensor, the corresponding rating for <user, item> pair.\n",
    "        \"\"\"\n",
    "        self.user_tensor = user_tensor\n",
    "        self.item_tensor = item_tensor\n",
    "        self.target_tensor = target_tensor\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"Get an item from dataset.\"\"\"\n",
    "        return (\n",
    "            self.user_tensor[index],\n",
    "            self.item_tensor[index],\n",
    "            self.target_tensor[index],\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Get the size of the dataset.\"\"\"\n",
    "        return self.user_tensor.size(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "class TorchPairwiseNegativeDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"Wrapper, convert <user, pos_item, neg_item> Tensor into Pytorch Dataset.\"\"\"\n",
    "\n",
    "    def __init__(self, user_tensor, pos_item_tensor, neg_item_tensor):\n",
    "        \"\"\"Init PairwiseNegativeDataset Class.\n",
    "        Args:\n",
    "            target_tensor: torch.Tensor, the corresponding rating for <user, item> pair.\n",
    "        \"\"\"\n",
    "        self.user_tensor = user_tensor\n",
    "        self.pos_item_tensor = pos_item_tensor\n",
    "        self.neg_item_tensor = neg_item_tensor\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \"\"\"Get an item from the dataset.\"\"\"\n",
    "        return (\n",
    "            self.user_tensor[index],\n",
    "            self.pos_item_tensor[index],\n",
    "            self.neg_item_tensor[index],\n",
    "        )\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\"Get the size of the dataset.\"\"\"\n",
    "        return self.user_tensor.size(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **References:-**\n",
    "- https://github.com/beta-team/beta-recsys/blob/master/beta_rec/data/data_loaders.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Author: Sparsh A.\n",
      "\n",
      "Last updated: 2021-12-31 06:27:08\n",
      "\n",
      "recohut: 0.0.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",
      "sys    : 3.7.12 (default, Sep 10 2021, 00:21:48) \n",
      "[GCC 7.5.0]\n",
      "numpy  : 1.19.5\n",
      "pandas : 1.1.5\n",
      "csv    : 1.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"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
