Using custom cross-fold splits with SurpriseLib

custom_crossfold_surpriselib.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from attr import dataclass\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from collections import namedtuple\n",
    "from typing import List, Optional\n",
    "from surprise import Dataset as SurpriseDataset, Reader, SVD\n",
    "from surprise import accuracy\n",
    "\n",
    "from sklearn.metrics import mean_squared_error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fold 0 RMSE: 0.9495\n",
      "Fold 1 RMSE: 0.9401\n",
      "Fold 2 RMSE: 0.9364\n",
      "Fold 3 RMSE: 0.9294\n",
      "Fold 4 RMSE: 0.9342\n"
     ]
    }
   ],
   "source": [
    "DATA_DIR = \"/Users/khalil/projects/pyrec/pyrec/federated/data/\"\n",
    "Fold = namedtuple(\"Fold\", [\"train\", \"test\"])\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Dataset:\n",
    "    interactions: pd.DataFrame = None\n",
    "    folds: List[Fold] = None\n",
    "\n",
    "    @staticmethod\n",
    "    def from_csv(num_folds: Optional[int] = None):\n",
    "        interactions = pd.read_csv(os.path.join(DATA_DIR, \"movielens-100k.csv\"))\n",
    "        interactions[\"user_id\"] = interactions[\"user_id\"].astype(str)\n",
    "        interactions[\"item_id\"] = interactions[\"item_id\"].astype(str)\n",
    "        if num_folds:\n",
    "            folds = Dataset.split_crossfold(interactions, num_folds)\n",
    "        else:\n",
    "            folds = None\n",
    "        return Dataset(interactions=interactions, folds=folds)\n",
    "\n",
    "    @staticmethod\n",
    "    def split_crossfold(interactions: pd.DataFrame, num_folds: int = 3):\n",
    "        splits = np.array_split(interactions.index, num_folds)\n",
    "        folds = []\n",
    "        for idx, split in enumerate(splits):\n",
    "            test = interactions.loc[split]\n",
    "            train_folds_idxs = [splits[x] for x in range(len(splits)) if x != idx]\n",
    "            train_folds_idxs = np.concatenate(train_folds_idxs)\n",
    "            train = interactions.loc[train_folds_idxs]\n",
    "            assert len(test) + len(train) == len(\n",
    "                interactions\n",
    "            ), f\"{len(test)} + {len(train)} =={len(interactions)}\"\n",
    "            folds.append(Fold(train=train, test=test))\n",
    "        return folds\n",
    "\n",
    "\n",
    "NUM_FOLDS = 5\n",
    "dataset = Dataset.from_csv(num_folds=NUM_FOLDS)\n",
    "assert len(dataset.folds) == NUM_FOLDS\n",
    "\n",
    "for fold_idx, fold in enumerate(dataset.folds):\n",
    "    train = SurpriseDataset.load_from_df(fold.train, reader=Reader())\n",
    "    test = train.construct_testset(\n",
    "        [(x.user_id, x.item_id, x.rating, 0) for x in fold.test.itertuples()]\n",
    "    )\n",
    "    train.build_full_trainset()\n",
    "    algo = SVD()\n",
    "    algo.fit(train.build_full_trainset())\n",
    "    predictions = algo.test(test)\n",
    "    predictions_df = pd.DataFrame(predictions)\n",
    "    rmse = np.sqrt(mean_squared_error(predictions_df[\"r_ui\"], predictions_df[\"est\"]))\n",
    "    print('Fold {} RMSE: {:.4f}'.format(fold_idx, rmse))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}