final_project.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1d1e1eff",
   "metadata": {},
   "source": [
    "## 1. Data Exploration and Pre-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f24176f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: An illegal reflective access operation has occurred\n",
      "WARNING: Illegal reflective access by org.apache.spark.unsafe.Platform (file:/ext3/spark-3.1.2-bin-hadoop3.2/jars/spark-unsafe_2.12-3.1.2.jar) to constructor java.nio.DirectByteBuffer(long,int)\n",
      "WARNING: Please consider reporting this to the maintainers of org.apache.spark.unsafe.Platform\n",
      "WARNING: Use --illegal-access=warn to enable warnings of further illegal reflective access operations\n",
      "WARNING: All illegal access operations will be denied in a future release\n",
      "23/05/16 22:53:18 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n",
      "Using Spark's default log4j profile: org/apache/spark/log4j-defaults.properties\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: pyarrow in /home/vr2229/.local/lib/python3.8/site-packages (12.0.0)\n",
      "Requirement already satisfied: numpy>=1.16.6 in /ext3/pyspark/lib/python3.8/site-packages (from pyarrow) (1.21.2)\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: scikit-learn in /home/vr2229/.local/lib/python3.8/site-packages (1.2.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /home/vr2229/.local/lib/python3.8/site-packages (from scikit-learn) (3.1.0)\n",
      "Requirement already satisfied: scipy>=1.3.2 in /ext3/pyspark/lib/python3.8/site-packages (from scikit-learn) (1.7.1)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /home/vr2229/.local/lib/python3.8/site-packages (from scikit-learn) (1.2.0)\n",
      "Requirement already satisfied: numpy>=1.17.3 in /ext3/pyspark/lib/python3.8/site-packages (from scikit-learn) (1.21.2)\n",
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Requirement already satisfied: implicit in /home/vr2229/.local/lib/python3.8/site-packages (0.6.2)\n",
      "Requirement already satisfied: numpy in /ext3/pyspark/lib/python3.8/site-packages (from implicit) (1.21.2)\n",
      "Requirement already satisfied: tqdm>=4.27 in /home/vr2229/.local/lib/python3.8/site-packages (from implicit) (4.65.0)\n",
      "Requirement already satisfied: scipy>=0.16 in /ext3/pyspark/lib/python3.8/site-packages (from implicit) (1.7.1)\n"
     ]
    }
   ],
   "source": [
    "'''import libraries'''\n",
    "import pyspark\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.feature import StringIndexer\n",
    "from pyspark.ml.linalg import SparseVector\n",
    "from pyspark.ml.linalg import Vectors\n",
    "from pyspark.sql.functions import broadcast\n",
    "from pyspark.sql.types import StructType, StructField, StringType, DoubleType\n",
    "from pyspark.mllib.linalg import Matrices\n",
    "spark = SparkSession.builder.appName('listenBrainz - exploratory data analysis').getOrCreate()\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import pyspark.sql.functions as F\n",
    "from pyspark.sql.functions import col\n",
    "from pyspark.sql.types import StructType, StructField, IntegerType\n",
    "\n",
    "!pip install pyarrow\n",
    "!pip install scikit-learn\n",
    "!pip install implicit\n",
    "\n",
    "from scipy.sparse import coo_matrix\n",
    "import scipy.sparse as sparse\n",
    "import implicit\n",
    "from implicit.als import AlternatingLeastSquares\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_squared_error\n",
    "from scipy.sparse import csr_matrix\n",
    "\n",
    "from pyspark.ml.evaluation import RegressionEvaluator,RankingEvaluator\n",
    "from pyspark.ml.recommendation import ALS\n",
    "from pyspark.ml.tuning import ParamGridBuilder, CrossValidator\n",
    "from pyspark.mllib.evaluation import RankingMetrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c663898",
   "metadata": {},
   "source": [
    "## 2. Import ListenBrainz dataset from given path on the cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6542227b",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''import dataset'''\n",
    "\n",
    "# small dataset\n",
    "interactions_train_small = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/interactions_train_small.parquet')\n",
    "tracks_train_small = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/tracks_train_small.parquet')\n",
    "users_train_small = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/users_train_small.parquet')\n",
    "\n",
    "\n",
    "# complete dataset\n",
    "interactions_train = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/interactions_train.parquet')\n",
    "tracks_train = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/tracks_train.parquet')\n",
    "users_train = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/users_train.parquet')\n",
    "\n",
    "\n",
    "# test dataset\n",
    "interactions_test = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/interactions_test.parquet')\n",
    "tracks_test = pd.read_parquet('/scratch/work/courses/DSGA1004-2021/listenbrainz/tracks_test.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c94114d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 51080140 entries, 0 to 51080139\n",
      "Data columns (total 3 columns):\n",
      " #   Column          Dtype         \n",
      "---  ------          -----         \n",
      " 0   user_id         Int32         \n",
      " 1   recording_msid  string        \n",
      " 2   timestamp       datetime64[ns]\n",
      "dtypes: Int32(1), datetime64[ns](1), string(1)\n",
      "memory usage: 1023.0 MB\n",
      "interactions_train_small:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "Index: 18969111 entries, 2 to 28704238\n",
      "Data columns (total 4 columns):\n",
      " #   Column          Dtype \n",
      "---  ------          ----- \n",
      " 0   recording_msid  string\n",
      " 1   artist_name     string\n",
      " 2   track_name      string\n",
      " 3   recording_mbid  string\n",
      "dtypes: string(4)\n",
      "memory usage: 723.6 MB\n",
      "tracks_train_small:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 6954 entries, 0 to 6953\n",
      "Data columns (total 2 columns):\n",
      " #   Column     Non-Null Count  Dtype \n",
      "---  ------     --------------  ----- \n",
      " 0   user_id    6954 non-null   Int32 \n",
      " 1   user_name  6954 non-null   string\n",
      "dtypes: Int32(1), string(1)\n",
      "memory usage: 88.4 KB\n",
      "users_train_small:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 179466123 entries, 0 to 179466122\n",
      "Data columns (total 3 columns):\n",
      " #   Column          Dtype         \n",
      "---  ------          -----         \n",
      " 0   user_id         Int32         \n",
      " 1   recording_msid  string        \n",
      " 2   timestamp       datetime64[ns]\n",
      "dtypes: Int32(1), datetime64[ns](1), string(1)\n",
      "memory usage: 3.5 GB\n",
      "interactions_train:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 28704239 entries, 0 to 28704238\n",
      "Data columns (total 4 columns):\n",
      " #   Column          Dtype \n",
      "---  ------          ----- \n",
      " 0   recording_msid  string\n",
      " 1   artist_name     string\n",
      " 2   track_name      string\n",
      " 3   recording_mbid  string\n",
      "dtypes: string(4)\n",
      "memory usage: 876.0 MB\n",
      "tracks_train:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 50031760 entries, 0 to 50031759\n",
      "Data columns (total 3 columns):\n",
      " #   Column          Dtype         \n",
      "---  ------          -----         \n",
      " 0   user_id         Int32         \n",
      " 1   recording_msid  string        \n",
      " 2   timestamp       datetime64[ns]\n",
      "dtypes: Int32(1), datetime64[ns](1), string(1)\n",
      "memory usage: 1002.0 MB\n",
      "interactions_test:  None\n",
      "*******\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 11037988 entries, 0 to 11037987\n",
      "Data columns (total 4 columns):\n",
      " #   Column          Dtype \n",
      "---  ------          ----- \n",
      " 0   recording_msid  string\n",
      " 1   artist_name     string\n",
      " 2   track_name      string\n",
      " 3   recording_mbid  string\n",
      "dtypes: string(4)\n",
      "memory usage: 336.9 MB\n",
      "tracks_test:  None\n",
      "*******\n"
     ]
    }
   ],
   "source": [
    "'''print info of pandas dataframes'''\n",
    "\n",
    "print('interactions_train_small: ', interactions_train_small.info())\n",
    "print ('*******')\n",
    "print('tracks_train_small: ', tracks_train_small.info())\n",
    "print ('*******')\n",
    "print('users_train_small: ', users_train_small.info())\n",
    "print ('*******')\n",
    "\n",
    "print('interactions_train: ', interactions_train.info())\n",
    "print ('*******')\n",
    "print('tracks_train: ', tracks_train.info())\n",
    "print ('*******')\n",
    "\n",
    "\n",
    "print('interactions_test: ', interactions_test.info())\n",
    "print ('*******')\n",
    "print('tracks_test: ', tracks_test.info())\n",
    "print ('*******')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "29b6dc1c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "interactions_train_small:  (51080140, 3)\n",
      "tracks_train_small:  (18969111, 4)\n",
      "users_train_small:  (6954, 2)\n",
      "interactions_train:  (179466123, 3)\n",
      "tracks_train:  (28704239, 4)\n",
      "interactions_test:  (50031760, 3)\n",
      "tracks_test:  (11037988, 4)\n"
     ]
    }
   ],
   "source": [
    "'''get datasets shape'''\n",
    "\n",
    "print ('interactions_train_small: ', interactions_train_small.shape)\n",
    "print ('tracks_train_small: ', tracks_train_small.shape)\n",
    "print('users_train_small: ',users_train_small.shape)\n",
    "\n",
    "print ('interactions_train: ', interactions_train.shape)\n",
    "print ('tracks_train: ', tracks_train.shape)\n",
    "\n",
    "print('interactions_test: ',interactions_test.shape)\n",
    "print('tracks_test: ',tracks_test.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24055aa2",
   "metadata": {},
   "source": [
    "## 3. Pre-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5e03d148",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "joined_df:  (51080140, 6)\n",
      "joined_df_train:  (179466123, 6)\n",
      "joined_df_test:  (50031760, 6)\n"
     ]
    }
   ],
   "source": [
    "'''perform the join operation'''\n",
    "\n",
    "# small dataset\n",
    "joined_df_small = pd.merge(interactions_train_small, tracks_train_small, on='recording_msid', how='left')\n",
    "print('joined_df: ',joined_df_small.shape)\n",
    "\n",
    "# complete dataset\n",
    "joined_df_train = pd.merge(interactions_train, tracks_train, on='recording_msid', how='left')\n",
    "print('joined_df_train: ',joined_df_train.shape)\n",
    "\n",
    "# test dataset\n",
    "joined_df_test = pd.merge(interactions_test, tracks_test, on='recording_msid', how='left')\n",
    "print('joined_df_test: ',joined_df_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cb044639",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rec_mbid_present_small shape:  (19226385, 6)\n",
      "rec_mbid_present_train shape:  (75330453, 6)\n",
      "rec_mbid_present_test shape:  (13624362, 6)\n"
     ]
    }
   ],
   "source": [
    "'''create a dataframe where recording_mbid is present'''\n",
    "\n",
    "# small dataset\n",
    "rec_mbid_present_small = joined_df_small[joined_df_small['recording_mbid'].notnull()]\n",
    "print('rec_mbid_present_small shape: ', rec_mbid_present_small.shape)\n",
    "\n",
    "# complete dataset\n",
    "rec_mbid_present_train = joined_df_train[joined_df_train['recording_mbid'].notnull()]\n",
    "print('rec_mbid_present_train shape: ', rec_mbid_present_train.shape)\n",
    "\n",
    "\n",
    "# test dataset\n",
    "rec_mbid_present_test = joined_df_test[joined_df_test['recording_mbid'].notnull()]\n",
    "print('rec_mbid_present_test shape: ', rec_mbid_present_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9ac19f39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rec_mbid_null_small shape:  (31853755, 6)\n",
      "rec_mbid_null_train shape:  (104135670, 6)\n",
      "rec_mbid_null_test shape:  (36407398, 6)\n"
     ]
    }
   ],
   "source": [
    "'''create a dataframe where recording_mbid is null'''\n",
    "\n",
    "# small dataset\n",
    "rec_mbid_null_small = joined_df_small[joined_df_small['recording_mbid'].isnull()]\n",
    "print('rec_mbid_null_small shape: ', rec_mbid_null_small.shape)\n",
    "\n",
    "# complete dataset\n",
    "rec_mbid_null_train = joined_df_train[joined_df_train['recording_mbid'].isnull()]\n",
    "print('rec_mbid_null_train shape: ', rec_mbid_null_train.shape)\n",
    "\n",
    "\n",
    "# test dataset\n",
    "rec_mbid_null_test = joined_df_test[joined_df_test['recording_mbid'].isnull()]\n",
    "print('rec_mbid_null_test shape: ', rec_mbid_null_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b1109560",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "merged_df_small shape:  (19226385, 7)\n",
      "merged_df_train shape:  (75330453, 7)\n",
      "merged_df_test shape:  (13624362, 7)\n"
     ]
    }
   ],
   "source": [
    "# small dataset\n",
    "'''group recording_msid by recording_mbid and get the first recording_msid for each group'''\n",
    "rec_mbid_present_grouped_small = rec_mbid_present_small.groupby('recording_mbid').agg({'recording_msid': 'first'})\n",
    "\n",
    "'''rename the column to first_recording_msid'''\n",
    "rec_mbid_present_grouped_small = rec_mbid_present_grouped_small.rename(columns={'recording_msid': 'first_recording_msid'})\n",
    "\n",
    "'''merge rec_mbid_present_grouped and rec_mbid_present'''\n",
    "merged_df_small = rec_mbid_present_small.merge(rec_mbid_present_grouped_small, on='recording_mbid')\n",
    "print('merged_df_small shape: ', merged_df_small.shape)\n",
    "\n",
    "# train dataset\n",
    "rec_mbid_present_grouped_train = rec_mbid_present_train.groupby('recording_mbid').agg({'recording_msid': 'first'})\n",
    "rec_mbid_present_grouped_train = rec_mbid_present_grouped_train.rename(columns={'recording_msid': 'first_recording_msid'})\n",
    "merged_df_train = rec_mbid_present_train.merge(rec_mbid_present_grouped_train, on='recording_mbid')\n",
    "print('merged_df_train shape: ', merged_df_train.shape)\n",
    "\n",
    "# test dataset\n",
    "rec_mbid_present_grouped_test = rec_mbid_present_test.groupby('recording_mbid').agg({'recording_msid': 'first'})\n",
    "rec_mbid_present_grouped_test = rec_mbid_present_grouped_test.rename(columns={'recording_msid': 'first_recording_msid'})\n",
    "merged_df_test = rec_mbid_present_test.merge(rec_mbid_present_grouped_test, on='recording_mbid')\n",
    "print('merged_df_test shape: ', merged_df_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d61885a6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "merged_df_small shape:  (19226385, 6)\n",
      "merged_df_train shape:  (75330453, 6)\n",
      "merged_df_test shape:  (13624362, 6)\n"
     ]
    }
   ],
   "source": [
    "'''drop recording_msid'''\n",
    "\n",
    "# small dataset\n",
    "merged_df_small  = merged_df_small .drop(['recording_msid'], axis=1)\n",
    "merged_df_small = merged_df_small.rename(columns={'first_recording_msid':'recording_msid'})\n",
    "print('merged_df_small shape: ', merged_df_small.shape)\n",
    "\n",
    "# train dataset\n",
    "merged_df_train  = merged_df_train .drop(['recording_msid'], axis=1)\n",
    "merged_df_train = merged_df_train.rename(columns={'first_recording_msid':'recording_msid'})\n",
    "print('merged_df_train shape: ', merged_df_train.shape)\n",
    "\n",
    "# test dataset\n",
    "merged_df_test  = merged_df_test .drop(['recording_msid'], axis=1)\n",
    "merged_df_test = merged_df_test.rename(columns={'first_recording_msid':'recording_msid'})\n",
    "print('merged_df_test shape: ', merged_df_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3a1eed65",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "concat_df_small:  (51080140, 6)\n",
      "concat_df_train:  (179466123, 6)\n",
      "concat_df_test:  (50031760, 6)\n"
     ]
    }
   ],
   "source": [
    "'''merging dataframes with mbid and without mbid'''\n",
    "\n",
    "# small dataset\n",
    "concat_df_small = pd.concat([merged_df_small, rec_mbid_null_small], axis=0)\n",
    "print('concat_df_small: ',concat_df_small.shape)\n",
    "\n",
    "# train dataset\n",
    "concat_df_train = pd.concat([merged_df_train, rec_mbid_null_train], axis=0)\n",
    "print('concat_df_train: ',concat_df_train.shape)\n",
    "\n",
    "# test dataset\n",
    "concat_df_test = pd.concat([merged_df_test, rec_mbid_null_test], axis=0)\n",
    "print('concat_df_test: ',concat_df_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "45d25a5d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "'''check if there is any null recording_msid'''\n",
    "print(concat_df_small['recording_msid'].isnull().sum())\n",
    "print(concat_df_train['recording_msid'].isnull().sum())\n",
    "print(concat_df_test['recording_msid'].isnull().sum())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bed7f635",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''drop recording_mbid'''\n",
    "concat_df_small = concat_df_small.drop(['recording_mbid'], axis=1)\n",
    "concat_df_train = concat_df_train.drop(['recording_mbid'], axis=1)\n",
    "concat_df_test = concat_df_test.drop(['recording_mbid'], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "87249872",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "listen_counts_small:  (21922976, 3)\n",
      "listen_counts_train:  (60912467, 3)\n",
      "listen_counts_test:  (21037310, 3)\n"
     ]
    }
   ],
   "source": [
    "'''group DataFrame by user_id and recording_msid, and count number of listens'''\n",
    "listen_counts_small = concat_df_small.groupby(['user_id', 'recording_msid']).size().reset_index(name='listens')\n",
    "print('listen_counts_small: ',listen_counts_small.shape)\n",
    "\n",
    "listen_counts_train = concat_df_train.groupby(['user_id', 'recording_msid']).size().reset_index(name='listens')\n",
    "print('listen_counts_train: ',listen_counts_train.shape)\n",
    "\n",
    "listen_counts_test = concat_df_test.groupby(['user_id', 'recording_msid']).size().reset_index(name='listens')\n",
    "print('listen_counts_test: ',listen_counts_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "a9a5fefe",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average listens for listen_counts_small: 2.3299820243383014\n",
      "Median listens for listen_counts_small: 1.0\n",
      "Average listens for listen_counts_train: 2.9462954275025504\n",
      "Median listens for listen_counts_train: 1.0\n",
      "Average listens for listen_counts_test: 2.3782394231962165\n",
      "Median listens for listen_counts_test: 1.0\n"
     ]
    }
   ],
   "source": [
    "'''print average and median listens'''\n",
    "print(f\"Average listens for listen_counts_small: {listen_counts_small['listens'].mean()}\")\n",
    "print(f\"Median listens for listen_counts_small: {listen_counts_small['listens'].median()}\")\n",
    "\n",
    "print(f\"Average listens for listen_counts_train: {listen_counts_train['listens'].mean()}\")\n",
    "print(f\"Median listens for listen_counts_train: {listen_counts_train['listens'].median()}\")\n",
    "\n",
    "print(f\"Average listens for listen_counts_test: {listen_counts_test['listens'].mean()}\")\n",
    "print(f\"Median listens for listen_counts_test: {listen_counts_test['listens'].median()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "df8f9e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''filter rows with listens less than average'''\n",
    "filtered_df_small = listen_counts_small[listen_counts_small['listens'] > listen_counts_small['listens'].median() + 3]\n",
    "filtered_df_train = listen_counts_train[listen_counts_train['listens'] > listen_counts_train['listens'].median() + 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "79524333",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filtered_df_small shape:  (52257, 4)\n",
      "filtered_df_train shape:  (454094, 4)\n"
     ]
    }
   ],
   "source": [
    "# small dataset\n",
    "user_means_small = filtered_df_small.groupby('user_id')['listens'].mean()\n",
    "song_means_small = filtered_df_small.groupby('recording_msid')['listens'].mean()\n",
    "\n",
    "filtered_df_small = filtered_df_small.merge(user_means_small, on='user_id', suffixes=('', '_user_mean'))\n",
    "filtered_df_small = filtered_df_small.merge(song_means_small, on='recording_msid', suffixes=('', '_song_mean'))\n",
    "filtered_df_small['listens_centered'] = filtered_df_small['listens'] - filtered_df_small['listens_user_mean'] - filtered_df_small['listens_song_mean']\n",
    "filtered_df_small = filtered_df_small.drop(['listens_user_mean', 'listens_song_mean'], axis=1)\n",
    "filtered_df_small = filtered_df_small.rename(columns={'listens_centered':'listens_mean'})\n",
    "filtered_df_small = filtered_df_small[filtered_df_small['listens_mean'] >= 0]\n",
    "filtered_df_small['listens_mean'] = filtered_df_small['listens_mean'].round().astype(int)\n",
    "print('filtered_df_small shape: ',filtered_df_small.shape)\n",
    "\n",
    "# complete dataset\n",
    "user_means_train = filtered_df_train.groupby('user_id')['listens'].mean()\n",
    "song_means_train = filtered_df_train.groupby('recording_msid')['listens'].mean()\n",
    "\n",
    "filtered_df_train = filtered_df_train.merge(user_means_small, on='user_id', suffixes=('', '_user_mean'))\n",
    "filtered_df_train = filtered_df_train.merge(song_means_small, on='recording_msid', suffixes=('', '_song_mean'))\n",
    "filtered_df_train['listens_centered'] = filtered_df_train['listens'] - filtered_df_train['listens_user_mean'] - filtered_df_train['listens_song_mean']\n",
    "filtered_df_train = filtered_df_train.drop(['listens_user_mean', 'listens_song_mean'], axis=1)\n",
    "filtered_df_train = filtered_df_train.rename(columns={'listens_centered':'listens_mean'})\n",
    "filtered_df_train = filtered_df_train[filtered_df_train['listens_mean'] >= 0]\n",
    "filtered_df_train['listens_mean'] = filtered_df_train['listens_mean'].round().astype(int)\n",
    "print('filtered_df_train shape: ',filtered_df_train.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "1876fe44",
   "metadata": {},
   "outputs": [],
   "source": [
    "df1 = pd.Series(filtered_df_small['recording_msid'].unique())\n",
    "df2 = pd.Series(filtered_df_train['recording_msid'].unique())\n",
    "df3 = pd.Series(listen_counts_test['recording_msid'].unique())\n",
    "\n",
    "msid_series = pd.unique(pd.concat([df1, df2, df3]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9f143c59",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''assign cat codes to recording_msid'''\n",
    "# Convert union_unique_values to a Pandas Series\n",
    "union_series = pd.Series(msid_series)\n",
    "\n",
    "# Assign categorical codes using pd.Categorical()\n",
    "codes = pd.Categorical(union_series).codes\n",
    "\n",
    "# Convert the codes array to a Pandas DataFrame\n",
    "msid = pd.DataFrame({'recording_msid': union_series, 'codes': codes})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "1f3f96f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_df_small = pd.merge(filtered_df_small, msid, on='recording_msid', how='left')\n",
    "filtered_df_train = pd.merge(filtered_df_train, msid, on='recording_msid', how='left')\n",
    "listen_counts_test = pd.merge(listen_counts_test, msid, on='recording_msid', how='left')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "a1531530",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filtered_df_small shape:  (52257, 5)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>recording_msid</th>\n",
       "      <th>listens</th>\n",
       "      <th>listens_mean</th>\n",
       "      <th>codes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1166</td>\n",
       "      <td>02aa34a9-28bd-4abe-a3c8-b1a4cd5dd1cf</td>\n",
       "      <td>21</td>\n",
       "      <td>1</td>\n",
       "      <td>105873</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>3159</td>\n",
       "      <td>03b0037a-33eb-4b00-b680-a43445df310c</td>\n",
       "      <td>38</td>\n",
       "      <td>16</td>\n",
       "      <td>146558</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3159</td>\n",
       "      <td>03cfae38-14e3-47ff-ae24-3982660a6758</td>\n",
       "      <td>42</td>\n",
       "      <td>21</td>\n",
       "      <td>151508</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>16332</td>\n",
       "      <td>0852c051-2b58-472f-9324-95b5c953e61d</td>\n",
       "      <td>24</td>\n",
       "      <td>4</td>\n",
       "      <td>330240</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1015</td>\n",
       "      <td>09a3a128-b467-4f6c-be8b-72e6c33364f2</td>\n",
       "      <td>152</td>\n",
       "      <td>105</td>\n",
       "      <td>382772</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>22157</td>\n",
       "      <td>09a3a128-b467-4f6c-be8b-72e6c33364f2</td>\n",
       "      <td>36</td>\n",
       "      <td>2</td>\n",
       "      <td>382772</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2881</td>\n",
       "      <td>0a22d1e6-3706-406d-80a9-7e1ea902793b</td>\n",
       "      <td>37</td>\n",
       "      <td>15</td>\n",
       "      <td>402575</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1015</td>\n",
       "      <td>0b6c2d5c-47a9-4972-ab11-804db96e6f63</td>\n",
       "      <td>129</td>\n",
       "      <td>91</td>\n",
       "      <td>453577</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>3816</td>\n",
       "      <td>133712d7-8ce3-422b-9429-831a367fc570</td>\n",
       "      <td>29</td>\n",
       "      <td>5</td>\n",
       "      <td>763906</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>20420</td>\n",
       "      <td>133712d7-8ce3-422b-9429-831a367fc570</td>\n",
       "      <td>20</td>\n",
       "      <td>0</td>\n",
       "      <td>763906</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id                        recording_msid  listens  listens_mean   \n",
       "0     1166  02aa34a9-28bd-4abe-a3c8-b1a4cd5dd1cf       21             1  \\\n",
       "1     3159  03b0037a-33eb-4b00-b680-a43445df310c       38            16   \n",
       "2     3159  03cfae38-14e3-47ff-ae24-3982660a6758       42            21   \n",
       "3    16332  0852c051-2b58-472f-9324-95b5c953e61d       24             4   \n",
       "4     1015  09a3a128-b467-4f6c-be8b-72e6c33364f2      152           105   \n",
       "5    22157  09a3a128-b467-4f6c-be8b-72e6c33364f2       36             2   \n",
       "6     2881  0a22d1e6-3706-406d-80a9-7e1ea902793b       37            15   \n",
       "7     1015  0b6c2d5c-47a9-4972-ab11-804db96e6f63      129            91   \n",
       "8     3816  133712d7-8ce3-422b-9429-831a367fc570       29             5   \n",
       "9    20420  133712d7-8ce3-422b-9429-831a367fc570       20             0   \n",
       "\n",
       "    codes  \n",
       "0  105873  \n",
       "1  146558  \n",
       "2  151508  \n",
       "3  330240  \n",
       "4  382772  \n",
       "5  382772  \n",
       "6  402575  \n",
       "7  453577  \n",
       "8  763906  \n",
       "9  763906  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print ('filtered_df_small shape: ', filtered_df_small.shape)\n",
    "filtered_df_small.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "4e0b6ed0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "filtered_df_train shape:  (454094, 5)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>recording_msid</th>\n",
       "      <th>listens</th>\n",
       "      <th>listens_mean</th>\n",
       "      <th>codes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>002bd589-8913-4b48-a163-162f03990aae</td>\n",
       "      <td>28</td>\n",
       "      <td>15</td>\n",
       "      <td>6863</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>00e1ef0e-e95e-4ba8-af2d-7423c5372c02</td>\n",
       "      <td>57</td>\n",
       "      <td>43</td>\n",
       "      <td>35362</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1166</td>\n",
       "      <td>013d582c-d436-4d5d-a0b7-15fa8b1bfcf1</td>\n",
       "      <td>45</td>\n",
       "      <td>25</td>\n",
       "      <td>49485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1410</td>\n",
       "      <td>013d582c-d436-4d5d-a0b7-15fa8b1bfcf1</td>\n",
       "      <td>65</td>\n",
       "      <td>43</td>\n",
       "      <td>49485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3107</td>\n",
       "      <td>013d582c-d436-4d5d-a0b7-15fa8b1bfcf1</td>\n",
       "      <td>25</td>\n",
       "      <td>4</td>\n",
       "      <td>49485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>4490</td>\n",
       "      <td>013d582c-d436-4d5d-a0b7-15fa8b1bfcf1</td>\n",
       "      <td>65</td>\n",
       "      <td>43</td>\n",
       "      <td>49485</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1</td>\n",
       "      <td>01905d83-c169-42fc-bb93-d4e8431e9b43</td>\n",
       "      <td>21</td>\n",
       "      <td>6</td>\n",
       "      <td>62333</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>408</td>\n",
       "      <td>02083c1f-88cd-40f5-ba60-b364e4ede644</td>\n",
       "      <td>23</td>\n",
       "      <td>7</td>\n",
       "      <td>80735</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>734</td>\n",
       "      <td>02083c1f-88cd-40f5-ba60-b364e4ede644</td>\n",
       "      <td>39</td>\n",
       "      <td>23</td>\n",
       "      <td>80735</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1166</td>\n",
       "      <td>02083c1f-88cd-40f5-ba60-b364e4ede644</td>\n",
       "      <td>41</td>\n",
       "      <td>23</td>\n",
       "      <td>80735</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id                        recording_msid  listens  listens_mean  codes\n",
       "0        1  002bd589-8913-4b48-a163-162f03990aae       28            15   6863\n",
       "1        1  00e1ef0e-e95e-4ba8-af2d-7423c5372c02       57            43  35362\n",
       "2     1166  013d582c-d436-4d5d-a0b7-15fa8b1bfcf1       45            25  49485\n",
       "3     1410  013d582c-d436-4d5d-a0b7-15fa8b1bfcf1       65            43  49485\n",
       "4     3107  013d582c-d436-4d5d-a0b7-15fa8b1bfcf1       25             4  49485\n",
       "5     4490  013d582c-d436-4d5d-a0b7-15fa8b1bfcf1       65            43  49485\n",
       "6        1  01905d83-c169-42fc-bb93-d4e8431e9b43       21             6  62333\n",
       "7      408  02083c1f-88cd-40f5-ba60-b364e4ede644       23             7  80735\n",
       "8      734  02083c1f-88cd-40f5-ba60-b364e4ede644       39            23  80735\n",
       "9     1166  02083c1f-88cd-40f5-ba60-b364e4ede644       41            23  80735"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print ('filtered_df_train shape: ', filtered_df_train.shape)\n",
    "filtered_df_train.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "dc9b1212",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "listen_counts_test shape:  (21037310, 4)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>recording_msid</th>\n",
       "      <th>listens</th>\n",
       "      <th>codes</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0003ce41-37c8-4e05-9ffe-9bff615406db</td>\n",
       "      <td>4</td>\n",
       "      <td>590</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>0005a3ca-84b8-4420-97a1-12dc9cdcf066</td>\n",
       "      <td>1</td>\n",
       "      <td>888</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>000a7629-fdb7-4e5d-bf7d-00f44d373f3f</td>\n",
       "      <td>1</td>\n",
       "      <td>1596</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>00200571-21b0-4d8c-9eb4-bb839dbc8bca</td>\n",
       "      <td>6</td>\n",
       "      <td>4979</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>00212384-c7d2-424f-a0db-5f1ddc1c3028</td>\n",
       "      <td>2</td>\n",
       "      <td>5165</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>1</td>\n",
       "      <td>00348870-5de1-4326-be81-18fdc3dd43eb</td>\n",
       "      <td>45</td>\n",
       "      <td>8192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>1</td>\n",
       "      <td>00360680-df59-453c-9d5a-8c8142a22f2f</td>\n",
       "      <td>1</td>\n",
       "      <td>8425</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>1</td>\n",
       "      <td>0038bbaa-0092-4001-89f2-be7111d38d94</td>\n",
       "      <td>1</td>\n",
       "      <td>8863</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>1</td>\n",
       "      <td>00420ab6-ff80-4a0f-83d8-fa45cd5dd44d</td>\n",
       "      <td>3</td>\n",
       "      <td>10335</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>1</td>\n",
       "      <td>0044788b-8a59-4d7a-9066-16e4ea7a2bb1</td>\n",
       "      <td>1</td>\n",
       "      <td>10723</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id                        recording_msid  listens  codes\n",
       "0        1  0003ce41-37c8-4e05-9ffe-9bff615406db        4    590\n",
       "1        1  0005a3ca-84b8-4420-97a1-12dc9cdcf066        1    888\n",
       "2        1  000a7629-fdb7-4e5d-bf7d-00f44d373f3f        1   1596\n",
       "3        1  00200571-21b0-4d8c-9eb4-bb839dbc8bca        6   4979\n",
       "4        1  00212384-c7d2-424f-a0db-5f1ddc1c3028        2   5165\n",
       "5        1  00348870-5de1-4326-be81-18fdc3dd43eb       45   8192\n",
       "6        1  00360680-df59-453c-9d5a-8c8142a22f2f        1   8425\n",
       "7        1  0038bbaa-0092-4001-89f2-be7111d38d94        1   8863\n",
       "8        1  00420ab6-ff80-4a0f-83d8-fa45cd5dd44d        3  10335\n",
       "9        1  0044788b-8a59-4d7a-9066-16e4ea7a2bb1        1  10723"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print ('listen_counts_test shape: ', listen_counts_test.shape)\n",
    "listen_counts_test.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "53134445",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Small Train set shape: (45985, 5)\n",
      "Samll Validation set shape: (15678, 5)\n",
      "Train set shape: (399602, 5)\n",
      "Validation set shape: (136229, 5)\n"
     ]
    }
   ],
   "source": [
    "'''split filtered_df into train and test sets'''\n",
    "\n",
    "# small datadset\n",
    "train_df_small, val_df_small = train_test_split(filtered_df_small, test_size=0.3, random_state=42)\n",
    "val1_df_small, val2_df_small = train_test_split(val_df_small, test_size=0.4, random_state=42)\n",
    "train_df_small = pd.concat([train_df_small, val1_df_small])\n",
    "\n",
    "'''print shape of train and test sets'''\n",
    "print(f\"Small Train set shape: {train_df_small.shape}\")\n",
    "print(f\"Samll Validation set shape: {val_df_small.shape}\")\n",
    "\n",
    "# complete dataset\n",
    "train_df_train, val_df_train = train_test_split(filtered_df_train, test_size=0.3, random_state=42)\n",
    "val1_df_train, val2_df_train = train_test_split(val_df_train, test_size=0.4, random_state=42)\n",
    "train_df_train = pd.concat([train_df_train, val1_df_train])\n",
    "\n",
    "'''print shape of train and test sets'''\n",
    "print(f\"Train set shape: {train_df_train.shape}\")\n",
    "print(f\"Validation set shape: {val_df_train.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "5e12b455",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "users in small dataset:  2862\n",
      "common users in small train and validation dataframe:  2209\n",
      "users in complete dataset:  4483\n",
      "common users in train and validation dataframe:  4071\n",
      "common users in train and test dataframe:  3790\n"
     ]
    }
   ],
   "source": [
    "'''Get the common user_ids'''\n",
    "\n",
    "# small dataset\n",
    "train_users_small = set(train_df_small['user_id'].unique())\n",
    "val_users_small = set(val_df_small['user_id'].unique())\n",
    "common_users_small = train_users_small.intersection(val_users_small)\n",
    "print('users in small dataset: ', len(train_users_small))\n",
    "print('common users in small train and validation dataframe: ',len(common_users_small))\n",
    "\n",
    "# complete dataset\n",
    "train_users_train = set(train_df_train['user_id'].unique())\n",
    "val_users_train = set(val_df_train['user_id'].unique())\n",
    "common_users_train = train_users_train.intersection(val_users_train)\n",
    "print('users in complete dataset: ', len(train_users_train))\n",
    "print('common users in train and validation dataframe: ',len(common_users_train))\n",
    "\n",
    "# test dataset\n",
    "test_users_train = set(listen_counts_test['user_id'].unique())\n",
    "common_users_test = train_users_train.intersection(test_users_train)\n",
    "print('common users in train and test dataframe: ',len(common_users_test))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d747cac6",
   "metadata": {},
   "source": [
    "## 4. Save processed dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "82e35e1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''save the relevent dataframes'''\n",
    "train_df_small.to_parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/train_df_small.parquet', compression='snappy')\n",
    "val_df_small.to_parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/val_df_small.parquet', compression='snappy')\n",
    "\n",
    "train_df_train.to_parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/train_df_train.parquet', compression='snappy')\n",
    "val_df_train.to_parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/val_df_train.parquet', compression='snappy')\n",
    "\n",
    "listen_counts_test.to_parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/test.parquet', compression='snappy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc7bb0e5",
   "metadata": {},
   "source": [
    "## 5. Alternating Least Square"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88ed71d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''alternating least square'''\n",
    "train_df = spark.read.parquet(\"/home/vr2229/final-project-group-29/final_project/final-project-group-29/train_df_train.parquet\")\n",
    "val_df = spark.read.parquet(\"/home/vr2229/final-project-group-29/final_project/final-project-group-29/test.parquet\")\n",
    "\n",
    "# train_df = spark.read.parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/train_df_small.parquet')\n",
    "# val_df = spark.read.parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/val_df_small.parquet')\n",
    "from pyspark.sql.functions import asc\n",
    "training_df = train_df.sort(asc(\"user_id\"))\n",
    "validation_df = val_df.sort(asc(\"user_id\"))\n",
    "\n",
    "ranks = [20]\n",
    "regParam = [0.1]\n",
    "errors = []\n",
    "min_error = float('inf')\n",
    "training_df = training_df.repartition(20000)\n",
    "print('repartition done')\n",
    "\n",
    "for param in regParam:\n",
    "    for rank in ranks:\n",
    "        \n",
    "        tempALS = ALS(maxIter=20, rank=rank, regParam=param,\n",
    "              userCol='user_id', itemCol='codes', ratingCol='listens_mean', implicitPrefs = True, numUserBlocks = 10000,\n",
    "              coldStartStrategy='drop', seed=1234)\n",
    "        print('als done')\n",
    "        training_df = training_df.coalesce(20000)\n",
    "\n",
    "        model = tempALS.fit(training_df)\n",
    "        print('fitting done')\n",
    "        \n",
    "        predictions = model.transform(validation_df)\n",
    "        evaluator = RegressionEvaluator(metricName='rmse', labelCol='listens', \n",
    "                                predictionCol='prediction')\n",
    "        print('evaluation done')\n",
    "        rmse = evaluator.evaluate(predictions)\n",
    "        print('rmse calculated')\n",
    "        errors.append(rmse)\n",
    "        print('errors appended')\n",
    "        print ('For rank {0} and regParam {1} the RMSE is {2}'.format (rank, param, rmse))\n",
    "        \n",
    "        if rmse < min_error:\n",
    "            min_error = rmse\n",
    "            best_rank = rank\n",
    "            best_regParam = param\n",
    "            best_model = model\n",
    "        \n",
    "print ('The best model was trained with rank {0} and regParam {1}'.format(best_rank, best_regParam))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ea3db4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab436ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "'''save als model parameters for fastSearch'''\n",
    "user_factors = model.userFactors\n",
    "item_factors = model.itemFactors\n",
    "\n",
    "user_factors.write.parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/userFactors.parquet')\n",
    "item_factors.write.parquet('/home/vr2229/final-project-group-29/final_project/final-project-group-29/item_factors.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ee42c3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}