colab-bst.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "vertical-somewhere",
   "metadata": {
    "id": "vertical-somewhere"
   },
   "source": [
    "# Introduction\n",
    "\n",
    "This Notebooks is a join notebook from both the prepare_data and pytorch-bst in order to be run in google colab."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "engaged-nutrition",
   "metadata": {
    "id": "engaged-nutrition"
   },
   "source": [
    "# Prepare data section"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "q9N2ZJcNGbnK",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "q9N2ZJcNGbnK",
    "outputId": "890c1b04-98f5-41e4-e0de-93e0daef7099"
   },
   "outputs": [],
   "source": [
    "!pip install pytorch_lightning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "indirect-chapter",
   "metadata": {
    "id": "indirect-chapter"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from tqdm import tqdm\n",
    "import torchmetrics\n",
    "import math\n",
    "from urllib.request import urlretrieve\n",
    "from zipfile import ZipFile\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import numpy as np\n",
    "from math import sqrt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cognitive-constraint",
   "metadata": {
    "id": "cognitive-constraint"
   },
   "source": [
    "## Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "latest-medium",
   "metadata": {
    "id": "latest-medium"
   },
   "outputs": [],
   "source": [
    "WINDOW_SIZE = 20"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bibliographic-carbon",
   "metadata": {
    "id": "bibliographic-carbon"
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "medical-budapest",
   "metadata": {
    "id": "medical-budapest"
   },
   "outputs": [],
   "source": [
    "urlretrieve(\"http://files.grouplens.org/datasets/movielens/ml-1m.zip\", \"movielens.zip\")\n",
    "ZipFile(\"movielens.zip\", \"r\").extractall()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "iraqi-rescue",
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "iraqi-rescue",
    "outputId": "69a8d6af-0b9a-4859-a047-b5c56d4727fb"
   },
   "outputs": [],
   "source": [
    "users = pd.read_csv(\n",
    "    \"ml-1m/users.dat\",\n",
    "    sep=\"::\",\n",
    "    names=[\"user_id\", \"sex\", \"age_group\", \"occupation\", \"zip_code\"],\n",
    ")\n",
    "\n",
    "ratings = pd.read_csv(\n",
    "    \"ml-1m/ratings.dat\",\n",
    "    sep=\"::\",\n",
    "    names=[\"user_id\", \"movie_id\", \"rating\", \"unix_timestamp\"],\n",
    ")\n",
    "\n",
    "movies = pd.read_csv(\n",
    "    \"ml-1m/movies.dat\", sep=\"::\", names=[\"movie_id\", \"title\", \"genres\"], encoding=\"ISO-8859-1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "legendary-stomach",
   "metadata": {
    "id": "legendary-stomach"
   },
   "outputs": [],
   "source": [
    "## Movies\n",
    "movies[\"year\"] = movies[\"title\"].apply(lambda x: x[-5:-1])\n",
    "movies.year = pd.Categorical(movies.year)\n",
    "movies[\"year\"] = movies.year.cat.codes\n",
    "## Users\n",
    "users.sex = pd.Categorical(users.sex)\n",
    "users[\"sex\"] = users.sex.cat.codes\n",
    "\n",
    "\n",
    "users.age_group = pd.Categorical(users.age_group)\n",
    "users[\"age_group\"] = users.age_group.cat.codes\n",
    "\n",
    "\n",
    "users.occupation = pd.Categorical(users.occupation)\n",
    "users[\"occupation\"] = users.occupation.cat.codes\n",
    "\n",
    "\n",
    "users.zip_code = pd.Categorical(users.zip_code)\n",
    "users[\"zip_code\"] = users.zip_code.cat.codes\n",
    "\n",
    "#Ratings\n",
    "ratings['unix_timestamp'] = pd.to_datetime(ratings['unix_timestamp'],unit='s')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "voluntary-truth",
   "metadata": {
    "id": "voluntary-truth"
   },
   "outputs": [],
   "source": [
    "# Save primary csv's\n",
    "if not os.path.exists('data'):\n",
    "    os.makedirs('data')\n",
    "    \n",
    "    \n",
    "users.to_csv(\"data/users.csv\",index=False)\n",
    "movies.to_csv(\"data/movies.csv\",index=False)\n",
    "ratings.to_csv(\"data/ratings.csv\",index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "artificial-australia",
   "metadata": {
    "id": "artificial-australia"
   },
   "outputs": [],
   "source": [
    "## Movies\n",
    "movies[\"movie_id\"] = movies[\"movie_id\"].astype(str)\n",
    "## Users\n",
    "users[\"user_id\"] = users[\"user_id\"].astype(str)\n",
    "\n",
    "##Ratings \n",
    "ratings[\"movie_id\"] = ratings[\"movie_id\"].astype(str)\n",
    "ratings[\"user_id\"] = ratings[\"user_id\"].astype(str)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "animated-slovenia",
   "metadata": {
    "id": "animated-slovenia"
   },
   "outputs": [],
   "source": [
    "genres = [\n",
    "    \"Action\",\n",
    "    \"Adventure\",\n",
    "    \"Animation\",\n",
    "    \"Children's\",\n",
    "    \"Comedy\",\n",
    "    \"Crime\",\n",
    "    \"Documentary\",\n",
    "    \"Drama\",\n",
    "    \"Fantasy\",\n",
    "    \"Film-Noir\",\n",
    "    \"Horror\",\n",
    "    \"Musical\",\n",
    "    \"Mystery\",\n",
    "    \"Romance\",\n",
    "    \"Sci-Fi\",\n",
    "    \"Thriller\",\n",
    "    \"War\",\n",
    "    \"Western\",\n",
    "]\n",
    "\n",
    "for genre in genres:\n",
    "    movies[genre] = movies[\"genres\"].apply(\n",
    "        lambda values: int(genre in values.split(\"|\"))\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ranking-calgary",
   "metadata": {
    "id": "ranking-calgary"
   },
   "source": [
    "### Transform the movie ratings data into sequences\n",
    "\n",
    "First, let's sort the the ratings data using the `unix_timestamp`, and then group the\n",
    "`movie_id` values and the `rating` values by `user_id`.\n",
    "\n",
    "The output DataFrame will have a record for each `user_id`, with two ordered lists\n",
    "(sorted by rating datetime): the movies they have rated, and their ratings of these movies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "light-publicity",
   "metadata": {
    "id": "light-publicity"
   },
   "outputs": [],
   "source": [
    "ratings_group = ratings.sort_values(by=[\"unix_timestamp\"]).groupby(\"user_id\")\n",
    "\n",
    "ratings_data = pd.DataFrame(\n",
    "    data={\n",
    "        \"user_id\": list(ratings_group.groups.keys()),\n",
    "        \"movie_ids\": list(ratings_group.movie_id.apply(list)),\n",
    "        \"ratings\": list(ratings_group.rating.apply(list)),\n",
    "        \"timestamps\": list(ratings_group.unix_timestamp.apply(list)),\n",
    "    }\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "amazing-affair",
   "metadata": {
    "id": "amazing-affair"
   },
   "source": [
    "Now, let's split the `movie_ids` list into a set of sequences of a fixed length.\n",
    "We do the same for the `ratings`. Set the `sequence_length` variable to change the length\n",
    "of the input sequence to the model. You can also change the `step_size` to control the\n",
    "number of sequences to generate for each user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "boxed-chick",
   "metadata": {
    "id": "boxed-chick"
   },
   "outputs": [],
   "source": [
    "sequence_length = 8\n",
    "step_size = 1\n",
    "\n",
    "\n",
    "def create_sequences(values, window_size, step_size):\n",
    "    sequences = []\n",
    "    start_index = 0\n",
    "    while True:\n",
    "        end_index = start_index + window_size\n",
    "        seq = values[start_index:end_index]\n",
    "        if len(seq) < window_size:\n",
    "            seq = values[-window_size:]\n",
    "            if len(seq) == window_size:\n",
    "                sequences.append(seq)\n",
    "            break\n",
    "        sequences.append(seq)\n",
    "        start_index += step_size\n",
    "    return sequences\n",
    "\n",
    "\n",
    "ratings_data.movie_ids = ratings_data.movie_ids.apply(\n",
    "    lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    ")\n",
    "\n",
    "ratings_data.ratings = ratings_data.ratings.apply(\n",
    "    lambda ids: create_sequences(ids, sequence_length, step_size)\n",
    ")\n",
    "\n",
    "del ratings_data[\"timestamps\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "starting-catalyst",
   "metadata": {
    "id": "starting-catalyst"
   },
   "source": [
    "After that, we process the output to have each sequence in a separate records in\n",
    "the DataFrame. In addition, we join the user features with the ratings data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "visible-tuner",
   "metadata": {
    "id": "visible-tuner"
   },
   "outputs": [],
   "source": [
    "ratings_data_movies = ratings_data[[\"user_id\", \"movie_ids\"]].explode(\n",
    "    \"movie_ids\", ignore_index=True\n",
    ")\n",
    "ratings_data_rating = ratings_data[[\"ratings\"]].explode(\"ratings\", ignore_index=True)\n",
    "ratings_data_transformed = pd.concat([ratings_data_movies, ratings_data_rating], axis=1)\n",
    "ratings_data_transformed = ratings_data_transformed.join(\n",
    "    users.set_index(\"user_id\"), on=\"user_id\"\n",
    ")\n",
    "ratings_data_transformed.movie_ids = ratings_data_transformed.movie_ids.apply(\n",
    "    lambda x: \",\".join(x)\n",
    ")\n",
    "ratings_data_transformed.ratings = ratings_data_transformed.ratings.apply(\n",
    "    lambda x: \",\".join([str(v) for v in x])\n",
    ")\n",
    "\n",
    "del ratings_data_transformed[\"zip_code\"]\n",
    "\n",
    "ratings_data_transformed.rename(\n",
    "    columns={\"movie_ids\": \"sequence_movie_ids\", \"ratings\": \"sequence_ratings\"},\n",
    "    inplace=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "strategic-relief",
   "metadata": {
    "id": "strategic-relief"
   },
   "source": [
    "With `sequence_length` of 4 and `step_size` of 2, we end up with 498,623 sequences.\n",
    "\n",
    "Finally, we split the data into training and testing splits, with 85% and 15% of\n",
    "the instances, respectively, and store them to CSV files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "subsequent-thanks",
   "metadata": {
    "id": "subsequent-thanks"
   },
   "outputs": [],
   "source": [
    "random_selection = np.random.rand(len(ratings_data_transformed.index)) <= 0.85\n",
    "train_data = ratings_data_transformed[random_selection]\n",
    "test_data = ratings_data_transformed[~random_selection]\n",
    "\n",
    "train_data.to_csv(\"data/train_data.csv\", index=False, sep=\",\")\n",
    "test_data.to_csv(\"data/test_data.csv\", index=False, sep=\",\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "prerequisite-plymouth",
   "metadata": {
    "id": "prerequisite-plymouth"
   },
   "outputs": [],
   "source": [
    "test_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ambient-brother",
   "metadata": {
    "id": "ambient-brother"
   },
   "source": [
    "# BST Implementation and training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adapted-austin",
   "metadata": {
    "id": "adapted-austin"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import pytorch_lightning as pl\n",
    "from tqdm import tqdm\n",
    "import torchmetrics\n",
    "import math\n",
    "from urllib.request import urlretrieve\n",
    "from zipfile import ZipFile\n",
    "import os\n",
    "import torch.nn as nn\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "needed-newark",
   "metadata": {
    "id": "needed-newark"
   },
   "outputs": [],
   "source": [
    "users = pd.read_csv(\n",
    "    \"data/users.csv\",\n",
    "    sep=\",\",\n",
    ")\n",
    "\n",
    "ratings = pd.read_csv(\n",
    "    \"data/ratings.csv\",\n",
    "    sep=\",\",\n",
    ")\n",
    "\n",
    "movies = pd.read_csv(\n",
    "    \"data/movies.csv\", sep=\",\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "revolutionary-responsibility",
   "metadata": {
    "id": "revolutionary-responsibility"
   },
   "source": [
    "## Pytorch dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "confused-console",
   "metadata": {
    "id": "confused-console"
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import torch.utils.data as data\n",
    "from torchvision import transforms\n",
    "import ast\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "\n",
    "class MovieDataset(data.Dataset):\n",
    "    \"\"\"Movie dataset.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self, ratings_file,test=False\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            csv_file (string): Path to the csv file with user,past,future.\n",
    "        \"\"\"\n",
    "        self.ratings_frame = pd.read_csv(\n",
    "            ratings_file,\n",
    "            delimiter=\",\",\n",
    "            # iterator=True,\n",
    "        )\n",
    "        self.test = test\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.ratings_frame)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        data = self.ratings_frame.iloc[idx]\n",
    "        user_id = data.user_id\n",
    "        \n",
    "        movie_history = eval(data.sequence_movie_ids)\n",
    "        movie_history_ratings = eval(data.sequence_ratings)\n",
    "        target_movie_id = movie_history[-1:][0]\n",
    "        target_movie_rating = movie_history_ratings[-1:][0]\n",
    "        \n",
    "        movie_history = torch.LongTensor(movie_history[:-1])\n",
    "        movie_history_ratings = torch.LongTensor(movie_history_ratings[:-1])\n",
    "\n",
    "        \n",
    "        \n",
    "        sex = data.sex\n",
    "        age_group = data.age_group\n",
    "        occupation = data.occupation\n",
    "        \n",
    "        return user_id, movie_history, target_movie_id,  movie_history_ratings, target_movie_rating, sex, age_group, occupation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "difficult-instrument",
   "metadata": {
    "id": "difficult-instrument"
   },
   "outputs": [],
   "source": [
    "genres = [\n",
    "    \"Action\",\n",
    "    \"Adventure\",\n",
    "    \"Animation\",\n",
    "    \"Children's\",\n",
    "    \"Comedy\",\n",
    "    \"Crime\",\n",
    "    \"Documentary\",\n",
    "    \"Drama\",\n",
    "    \"Fantasy\",\n",
    "    \"Film-Noir\",\n",
    "    \"Horror\",\n",
    "    \"Musical\",\n",
    "    \"Mystery\",\n",
    "    \"Romance\",\n",
    "    \"Sci-Fi\",\n",
    "    \"Thriller\",\n",
    "    \"War\",\n",
    "    \"Western\",\n",
    "]\n",
    "\n",
    "for genre in genres:\n",
    "    movies[genre] = movies[\"genres\"].apply(\n",
    "        lambda values: int(genre in values.split(\"|\"))\n",
    "    )\n",
    "    \n",
    "sequence_length = 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pending-negotiation",
   "metadata": {
    "id": "pending-negotiation"
   },
   "outputs": [],
   "source": [
    "class PositionalEmbedding(nn.Module):\n",
    "    \"\"\"\n",
    "    Computes positional embedding following \"Attention is all you need\"\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, max_len, d_model):\n",
    "        super().__init__()\n",
    "\n",
    "        # Compute the positional encodings once in log space.\n",
    "        self.pe = nn.Embedding(max_len, d_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        return self.pe.weight.unsqueeze(0).repeat(batch_size, 1, 1)\n",
    "\n",
    "\n",
    "class BST(pl.LightningModule):\n",
    "    def __init__(\n",
    "        self, args=None,\n",
    "    ):\n",
    "        super().__init__()\n",
    "        super(BST, self).__init__()\n",
    "        \n",
    "        self.save_hyperparameters()\n",
    "        self.args = args\n",
    "        #-------------------\n",
    "        # Embedding layers\n",
    "        ##Users \n",
    "        self.embeddings_user_id = nn.Embedding(\n",
    "            int(users.user_id.max())+1, int(math.sqrt(users.user_id.max()))+1\n",
    "        )\n",
    "        ###Users features embeddings\n",
    "        self.embeddings_user_sex = nn.Embedding(\n",
    "            len(users.sex.unique()), int(math.sqrt(len(users.sex.unique())))\n",
    "        )\n",
    "        self.embeddings_age_group = nn.Embedding(\n",
    "            len(users.age_group.unique()), int(math.sqrt(len(users.age_group.unique())))\n",
    "        )\n",
    "        self.embeddings_user_occupation = nn.Embedding(\n",
    "            len(users.occupation.unique()), int(math.sqrt(len(users.occupation.unique())))\n",
    "        )\n",
    "        self.embeddings_user_zip_code = nn.Embedding(\n",
    "            len(users.zip_code.unique()), int(math.sqrt(len(users.sex.unique())))\n",
    "        )\n",
    "        \n",
    "        ##Movies\n",
    "        self.embeddings_movie_id = nn.Embedding(\n",
    "            int(movies.movie_id.max())+1, int(math.sqrt(movies.movie_id.max()))+1\n",
    "        )\n",
    "        \n",
    "        ###Movies features embeddings\n",
    "        genre_vectors = movies[genres].to_numpy()\n",
    "        self.embeddings_movie_genre = nn.Embedding(\n",
    "            genre_vectors.shape[0], genre_vectors.shape[1]\n",
    "        )\n",
    "        \n",
    "        \n",
    "        \n",
    "        self.embeddings_movie_year = nn.Embedding(\n",
    "            len(movies.year.unique()), int(math.sqrt(len(movies.year.unique())))\n",
    "        )\n",
    "        \n",
    "        self.positional_embedding = PositionalEmbedding(8, 9)\n",
    "        \n",
    "        # Network\n",
    "        self.transfomerlayer = nn.TransformerEncoderLayer(72, 3, dropout=0.2)\n",
    "        self.linear = nn.Sequential(\n",
    "            nn.Linear(\n",
    "                661,\n",
    "                1024,\n",
    "            ),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(1024, 512),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(256, 1),\n",
    "        )\n",
    "        self.criterion = torch.nn.MSELoss()\n",
    "        self.mae = torchmetrics.MeanAbsoluteError()\n",
    "        self.mse = torchmetrics.MeanSquaredError()\n",
    "        \n",
    "\n",
    "\n",
    "    def encode_input(self,inputs):\n",
    "        user_id, movie_history, target_movie_id,  movie_history_ratings, target_movie_rating, sex, age_group, occupation = inputs\n",
    "               \n",
    "        #MOVIES\n",
    "        movie_history = self.embeddings_movie_id(movie_history)\n",
    "        target_movie = self.embeddings_movie_id(target_movie_id)\n",
    "         \n",
    "        target_movie = torch.unsqueeze(target_movie, 1)\n",
    "        transfomer_features = torch.cat((movie_history, target_movie),dim=1)\n",
    "\n",
    "        #USERS\n",
    "        user_id = self.embeddings_user_id(user_id)\n",
    "        \n",
    "        sex = self.embeddings_user_sex(sex)\n",
    "        age_group = self.embeddings_age_group(age_group)\n",
    "        occupation = self.embeddings_user_occupation(occupation)\n",
    "        user_features = torch.cat((user_id, sex, age_group,occupation), 1)\n",
    "        \n",
    "        return transfomer_features, user_features, target_movie_rating.float()\n",
    "    \n",
    "    def forward(self, batch):\n",
    "        transfomer_features, user_features, target_movie_rating = self.encode_input(batch)\n",
    "        positional_embedding = self.positional_embedding(transfomer_features)\n",
    "        transfomer_features = torch.cat((transfomer_features, positional_embedding), dim=2)\n",
    "        transformer_output = self.transfomerlayer(transfomer_features)\n",
    "        transformer_output = torch.flatten(transformer_output,start_dim=1)\n",
    "        \n",
    "        #Concat with other features\n",
    "        features = torch.cat((transformer_output,user_features),dim=1)\n",
    "\n",
    "        output = self.linear(features)\n",
    "        return output, target_movie_rating\n",
    "        \n",
    "    def training_step(self, batch, batch_idx):\n",
    "        out, target_movie_rating = self(batch)\n",
    "        out = out.flatten()\n",
    "        loss = self.criterion(out, target_movie_rating)\n",
    "        \n",
    "        mae = self.mae(out, target_movie_rating)\n",
    "        mse = self.mse(out, target_movie_rating)\n",
    "        rmse =torch.sqrt(mse)\n",
    "        self.log(\n",
    "            \"train/mae\", mae, on_step=True, on_epoch=False, prog_bar=False\n",
    "        )\n",
    "        \n",
    "        self.log(\n",
    "            \"train/rmse\", rmse, on_step=True, on_epoch=False, prog_bar=False\n",
    "        )\n",
    "        \n",
    "        self.log(\"train/step_loss\", loss, on_step=True, on_epoch=False, prog_bar=False)\n",
    "        return loss\n",
    "    \n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        out, target_movie_rating = self(batch)\n",
    "        out = out.flatten()\n",
    "        loss = self.criterion(out, target_movie_rating)\n",
    "        \n",
    "        mae = self.mae(out, target_movie_rating)\n",
    "        mse = self.mse(out, target_movie_rating)\n",
    "        rmse =torch.sqrt(mse)\n",
    "        \n",
    "        return {\"val_loss\": loss, \"mae\": mae.detach(), \"rmse\":rmse.detach()}\n",
    "\n",
    "    def validation_epoch_end(self, outputs):\n",
    "        avg_loss = torch.stack([x[\"val_loss\"] for x in outputs]).mean()\n",
    "        avg_mae = torch.stack([x[\"mae\"] for x in outputs]).mean()\n",
    "        avg_rmse = torch.stack([x[\"rmse\"] for x in outputs]).mean()\n",
    "        \n",
    "        self.log(\"val/loss\", avg_loss, on_step=False, on_epoch=True, prog_bar=False)\n",
    "        self.log(\"val/mae\", avg_mae, on_step=False, on_epoch=True, prog_bar=False)\n",
    "        self.log(\"val/rmse\", avg_rmse, on_step=False, on_epoch=True, prog_bar=False)\n",
    "\n",
    "\n",
    "    def test_epoch_end(self, outputs):\n",
    "        users = torch.cat([x[\"users\"] for x in outputs])\n",
    "        y_hat = torch.cat([x[\"top14\"] for x in outputs])\n",
    "        users = users.tolist()\n",
    "        y_hat = y_hat.tolist()\n",
    "        \n",
    "        data = {\"users\": users, \"top14\": y_hat}\n",
    "        df = pd.DataFrame.from_dict(data)\n",
    "        print(len(df))\n",
    "        df.to_csv(\"lightning_logs/predict.csv\", index=False)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.AdamW(self.parameters(), lr=0.0005)\n",
    "\n",
    "    @staticmethod\n",
    "    def add_model_specific_args(parent_parser):\n",
    "        parser = ArgumentParser(parents=[parent_parser], add_help=False)\n",
    "        parser.add_argument(\"--learning_rate\", type=float, default=0.01)\n",
    "        return parser\n",
    "\n",
    "    ####################\n",
    "    # DATA RELATED HOOKS\n",
    "    ####################\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "        print(\"Loading datasets\")\n",
    "        self.train_dataset = MovieDataset(\"data/train_data.csv\")\n",
    "        self.val_dataset = MovieDataset(\"data/test_data.csv\")\n",
    "        self.test_dataset = MovieDataset(\"data/test_data.csv\")\n",
    "        print(\"Done\")\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.train_dataset,\n",
    "            batch_size=128,\n",
    "            shuffle=False,\n",
    "            num_workers=os.cpu_count(),\n",
    "        )\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.val_dataset,\n",
    "            batch_size=128,\n",
    "            shuffle=False,\n",
    "            num_workers=os.cpu_count(),\n",
    "        )\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return torch.utils.data.DataLoader(\n",
    "            self.test_dataset,\n",
    "            batch_size=128,\n",
    "            shuffle=False,\n",
    "            num_workers=os.cpu_count(),\n",
    "        )\n",
    "        \n",
    "model = BST()\n",
    "trainer = pl.Trainer(gpus=1, max_epochs=50)\n",
    "trainer.fit(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "expanded-split",
   "metadata": {
    "id": "expanded-split"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "colab-bst.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "magnus",
   "language": "python",
   "name": "magnus"
  },
  "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}