Transformer_For_Sequential_Recommendation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyPGQo5lyjwVZ17k8skCphtk",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "premium"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/januverma/transformers-for-sequential-recommendation/blob/main/Transformer_For_Sequential_Recommendation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Sequential Recommendation\n",
        "\n",
        "Existing works on recommendation systems can be divided into two paradigms:\n",
        "- **General recommendation** mine static pair-wise correlations between entities from the historical interactions, and are modeled by the collaborative filtering models. \n",
        "- **Sequential recommendation** aims to predict users’ next actions based on the sequential interactions in the past. Most sequential recommendation models involve sequential data mining techniques e.g. Markov chains, RNNs, self-attention etc.\n",
        "\n",
        "While general recommendation methods consider user-item relationships as static and ignore the dynamics & the evolution of users’ preferences, the sequential recommendation models take into consideration the short-term user behaviour and are oblivious to the global user preferences which have stabilized over time.\n",
        "\n",
        "Recently, Transformer architecture has been shown to have superior performance for sequential data modelling. The transformer architecture lends itself to efficient parallelization and is effective at modeling long-range sequences.In this notebook, we will build a sequential recommendation model using transformer architecture. "
      ],
      "metadata": {
        "id": "MMnCGaLCtz-l"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import math"
      ],
      "metadata": {
        "id": "7hzibcB4cA_k"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from torch.nn.utils.rnn import pad_sequence"
      ],
      "metadata": {
        "id": "Woz3uqjD8aHO"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will use the [MovieLens]( https://files.grouplens.org/datasets/movielens) dataset which is a popular benchmark for recommendation systems research. We will use the 1-Million review dataset. \n",
        "\n"
      ],
      "metadata": {
        "id": "eg01x687vnKw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data_path = 'ml-1m'"
      ],
      "metadata": {
        "id": "Kqe0PpBXsL_F"
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "! wget https://files.grouplens.org/datasets/movielens/ml-1m.zip"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h8AkjVqxX2XS",
        "outputId": "9aa08643-3081-4d39-af20-ba3753b9124e"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2023-01-10 13:03:27--  https://files.grouplens.org/datasets/movielens/ml-1m.zip\n",
            "Resolving files.grouplens.org (files.grouplens.org)... 128.101.65.152\n",
            "Connecting to files.grouplens.org (files.grouplens.org)|128.101.65.152|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 5917549 (5.6M) [application/zip]\n",
            "Saving to: ‘ml-1m.zip.1’\n",
            "\n",
            "ml-1m.zip.1         100%[===================>]   5.64M  3.37MB/s    in 1.7s    \n",
            "\n",
            "2023-01-10 13:03:30 (3.37 MB/s) - ‘ml-1m.zip.1’ saved [5917549/5917549]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "! unzip ml-1m.zip"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HN6TbhdsYGrd",
        "outputId": "c75ea39c-a1c3-4114-e8d9-db80c99c6d05"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Archive:  ml-1m.zip\n",
            "   creating: ml-1m/\n",
            "  inflating: ml-1m/movies.dat        \n",
            "  inflating: ml-1m/ratings.dat       \n",
            "  inflating: ml-1m/README            \n",
            "  inflating: ml-1m/users.dat         \n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Loading data into `pandas` "
      ],
      "metadata": {
        "id": "tq7WFNZPv928"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ratings_path = data_path + '/ratings.dat'\n",
        "ratings = pd.read_csv(ratings_path, header=None, sep='::', engine='python', names=['userId', 'movieId', 'rating', 'timestamp'])\n",
        "ratings.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "3F35k9hNuBeK",
        "outputId": "1ab360c7-75d5-4a00-c039-497a519d58d0"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "   userId  movieId  rating  timestamp\n",
              "0       1     1193       5  978300760\n",
              "1       1      661       3  978302109\n",
              "2       1      914       3  978301968\n",
              "3       1     3408       4  978300275\n",
              "4       1     2355       5  978824291"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-5ff51162-818f-40d3-8d7e-a41dc0a11427\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <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>userId</th>\n",
              "      <th>movieId</th>\n",
              "      <th>rating</th>\n",
              "      <th>timestamp</th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>0</th>\n",
              "      <td>1</td>\n",
              "      <td>1193</td>\n",
              "      <td>5</td>\n",
              "      <td>978300760</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>661</td>\n",
              "      <td>3</td>\n",
              "      <td>978302109</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>1</td>\n",
              "      <td>914</td>\n",
              "      <td>3</td>\n",
              "      <td>978301968</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>1</td>\n",
              "      <td>3408</td>\n",
              "      <td>4</td>\n",
              "      <td>978300275</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>1</td>\n",
              "      <td>2355</td>\n",
              "      <td>5</td>\n",
              "      <td>978824291</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-5ff51162-818f-40d3-8d7e-a41dc0a11427')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-5ff51162-818f-40d3-8d7e-a41dc0a11427 button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-5ff51162-818f-40d3-8d7e-a41dc0a11427');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ratings.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kvUFA4QHwK07",
        "outputId": "3bdc816b-05ec-491e-8828-10d1d71076f4"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(1000209, 4)"
            ]
          },
          "metadata": {},
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The data has over 1M ratings by more than 6k users for 3.7k movies.\n",
        "\n",
        "## User Interaction Sequences\n",
        "Following the standard protocol for research in sequential recommendation, we treat the presence of a rating as implict feedback i.e. the user interacted with the movie and use the `timestamps` to determine the sequential structure of the interactions. "
      ],
      "metadata": {
        "id": "Ba5Jn9xUwYCa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "ratings = ratings.sort_values(by=['timestamp'])"
      ],
      "metadata": {
        "id": "9YldcWxTAeh0"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "user_seq_df = ratings.groupby('userId').agg({'movieId': lambda x:list(x), 'timestamp': lambda x:list(x)})\n",
        "user_seq_df.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 238
        },
        "id": "nLb2YZWu5hTq",
        "outputId": "e7afe2bb-e489-4cb4-8e66-e8615e2c247f"
      },
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "                                                  movieId  \\\n",
              "userId                                                      \n",
              "1       [3186, 1721, 1270, 1022, 2340, 1836, 3408, 120...   \n",
              "2       [1198, 1217, 1210, 2717, 1293, 2943, 1225, 119...   \n",
              "3       [593, 2858, 3534, 1968, 1961, 1431, 1266, 3671...   \n",
              "4       [1210, 1097, 480, 3468, 3527, 1196, 260, 1198,...   \n",
              "5       [2717, 919, 908, 356, 1250, 2188, 2858, 1127, ...   \n",
              "\n",
              "                                                timestamp  \n",
              "userId                                                     \n",
              "1       [978300019, 978300055, 978300055, 978300055, 9...  \n",
              "2       [978298124, 978298151, 978298151, 978298196, 9...  \n",
              "3       [978297018, 978297039, 978297068, 978297068, 9...  \n",
              "4       [978293924, 978293964, 978294008, 978294008, 9...  \n",
              "5       [978241072, 978241072, 978241072, 978241112, 9...  "
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-5a1297ec-f5a6-4ce6-92ac-d2703adf8688\">\n",
              "    <div class=\"colab-df-container\">\n",
              "      <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>movieId</th>\n",
              "      <th>timestamp</th>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>userId</th>\n",
              "      <th></th>\n",
              "      <th></th>\n",
              "    </tr>\n",
              "  </thead>\n",
              "  <tbody>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>[3186, 1721, 1270, 1022, 2340, 1836, 3408, 120...</td>\n",
              "      <td>[978300019, 978300055, 978300055, 978300055, 9...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>[1198, 1217, 1210, 2717, 1293, 2943, 1225, 119...</td>\n",
              "      <td>[978298124, 978298151, 978298151, 978298196, 9...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>[593, 2858, 3534, 1968, 1961, 1431, 1266, 3671...</td>\n",
              "      <td>[978297018, 978297039, 978297068, 978297068, 9...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>[1210, 1097, 480, 3468, 3527, 1196, 260, 1198,...</td>\n",
              "      <td>[978293924, 978293964, 978294008, 978294008, 9...</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>5</th>\n",
              "      <td>[2717, 919, 908, 356, 1250, 2188, 2858, 1127, ...</td>\n",
              "      <td>[978241072, 978241072, 978241072, 978241112, 9...</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "      <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-5a1297ec-f5a6-4ce6-92ac-d2703adf8688')\"\n",
              "              title=\"Convert this dataframe to an interactive table.\"\n",
              "              style=\"display:none;\">\n",
              "        \n",
              "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
              "       width=\"24px\">\n",
              "    <path d=\"M0 0h24v24H0V0z\" fill=\"none\"/>\n",
              "    <path d=\"M18.56 5.44l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94zm-11 1L8.5 8.5l.94-2.06 2.06-.94-2.06-.94L8.5 2.5l-.94 2.06-2.06.94zm10 10l.94 2.06.94-2.06 2.06-.94-2.06-.94-.94-2.06-.94 2.06-2.06.94z\"/><path d=\"M17.41 7.96l-1.37-1.37c-.4-.4-.92-.59-1.43-.59-.52 0-1.04.2-1.43.59L10.3 9.45l-7.72 7.72c-.78.78-.78 2.05 0 2.83L4 21.41c.39.39.9.59 1.41.59.51 0 1.02-.2 1.41-.59l7.78-7.78 2.81-2.81c.8-.78.8-2.07 0-2.86zM5.41 20L4 18.59l7.72-7.72 1.47 1.35L5.41 20z\"/>\n",
              "  </svg>\n",
              "      </button>\n",
              "      \n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\n",
              "      flex-wrap:wrap;\n",
              "      gap: 12px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert {\n",
              "      background-color: #E8F0FE;\n",
              "      border: none;\n",
              "      border-radius: 50%;\n",
              "      cursor: pointer;\n",
              "      display: none;\n",
              "      fill: #1967D2;\n",
              "      height: 32px;\n",
              "      padding: 0 0 0 0;\n",
              "      width: 32px;\n",
              "    }\n",
              "\n",
              "    .colab-df-convert:hover {\n",
              "      background-color: #E2EBFA;\n",
              "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "      fill: #174EA6;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert {\n",
              "      background-color: #3B4455;\n",
              "      fill: #D2E3FC;\n",
              "    }\n",
              "\n",
              "    [theme=dark] .colab-df-convert:hover {\n",
              "      background-color: #434B5C;\n",
              "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
              "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
              "      fill: #FFFFFF;\n",
              "    }\n",
              "  </style>\n",
              "\n",
              "      <script>\n",
              "        const buttonEl =\n",
              "          document.querySelector('#df-5a1297ec-f5a6-4ce6-92ac-d2703adf8688 button.colab-df-convert');\n",
              "        buttonEl.style.display =\n",
              "          google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "\n",
              "        async function convertToInteractive(key) {\n",
              "          const element = document.querySelector('#df-5a1297ec-f5a6-4ce6-92ac-d2703adf8688');\n",
              "          const dataTable =\n",
              "            await google.colab.kernel.invokeFunction('convertToInteractive',\n",
              "                                                     [key], {});\n",
              "          if (!dataTable) return;\n",
              "\n",
              "          const docLinkHtml = 'Like what you see? Visit the ' +\n",
              "            '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
              "            + ' to learn more about interactive tables.';\n",
              "          element.innerHTML = '';\n",
              "          dataTable['output_type'] = 'display_data';\n",
              "          await google.colab.output.renderOutput(dataTable, element);\n",
              "          const docLink = document.createElement('div');\n",
              "          docLink.innerHTML = docLinkHtml;\n",
              "          element.appendChild(docLink);\n",
              "        }\n",
              "      </script>\n",
              "    </div>\n",
              "  </div>\n",
              "  "
            ]
          },
          "metadata": {},
          "execution_count": 7
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We have converted users' interactions with movies into sequences (time-ordered) of movies. "
      ],
      "metadata": {
        "id": "lzGkmFkbFwXa"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Pre-processing Interaction Sequences"
      ],
      "metadata": {
        "id": "NiTyWLA5Cijh"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "To keep things simple, we will work with shorter sequences (`len <= 50`). "
      ],
      "metadata": {
        "id": "0nhxa37GwP5C"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "MAX_LENGTH = 50\n",
        "BATCH_SIZE = 8"
      ],
      "metadata": {
        "id": "PFE8t_G0gALH"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "For partitioning into train-val-test sets, we split the historical sequence for each user into three parts:\n",
        "- the most recent action for testing\n",
        "- the second most recent action for validation\n",
        "- all remaining actions for training. \n",
        "\n",
        "Note that during testing, the input sequences contain training actions\n",
        "and the validation action."
      ],
      "metadata": {
        "id": "R4wfm93XF8bz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_seqs = user_seq_df['movieId'].apply(lambda x:x[:-2]).values\n",
        "train_seqs.shape, len(train_seqs[1])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "raYWmxy3kAx2",
        "outputId": "c623fabf-e4d0-4cb7-869e-2e439c3fdd5a"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((6040,), 127)"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def create_valid_seqs(seq):\n",
        "  if len(seq) > MAX_LENGTH:\n",
        "    seq_len = len(seq) - MAX_LENGTH - 2\n",
        "    return seq[seq_len: -1]\n",
        "  else:\n",
        "    return seq[:-1]\n",
        "\n",
        "len(create_valid_seqs(list(range(58))))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "amP-PbHrQ7Wy",
        "outputId": "b9843d4d-edbe-4595-9b6b-540e2022ec6b"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "51"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "valid_seqs = user_seq_df['movieId'].apply(lambda x:create_valid_seqs(x)).values\n",
        "valid_seqs.shape, len(valid_seqs[0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nInAL1x6nD8Q",
        "outputId": "7e79bf7e-4f2b-45cb-f690-1196be5fbd62"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((6040,), 51)"
            ]
          },
          "metadata": {},
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def create_test_seqs(seq):\n",
        "  if len(seq) > MAX_LENGTH:\n",
        "    seq_len = len(seq) - MAX_LENGTH - 1\n",
        "    return seq[seq_len:]\n",
        "  else:\n",
        "    return seq\n",
        "\n",
        "len(create_test_seqs(list(range(58))))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7B7VdYOxedn2",
        "outputId": "563f3e76-89a6-4f5f-de57-b9b46e0100d7"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "51"
            ]
          },
          "metadata": {},
          "execution_count": 12
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Dataloaders for train and valid data\n",
        "\n",
        "We leverage PyTorch `dataloaders` to generate batches of sequences for training and evaluation. "
      ],
      "metadata": {
        "id": "qhETuWtMCuY6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class TrainSequenceData(Dataset):\n",
        "    def __init__(self, user_seqs, maxlen):\n",
        "        '''\n",
        "        Input is a list of lists where each list contains movie ids for a specific user. \n",
        "        '''\n",
        "        self.seqs = []\n",
        "        self.labels = []\n",
        "        self.maxlen = maxlen\n",
        "        for x in user_seqs:\n",
        "          seqs, labels = self.build_examples(x, maxlen)\n",
        "          self.seqs.extend(seqs)\n",
        "          self.labels.extend(labels)\n",
        "\n",
        "    def build_examples(self, x, max_len):\n",
        "        seqs = []\n",
        "        labels = []\n",
        "        for i in range(0, len(x) - 1, max_len):\n",
        "          seq_len = min(max_len, len(x) - 1 - i)\n",
        "          src = x[i:i+seq_len]\n",
        "          tgt = x[i+1:i+1+seq_len]\n",
        "          seqs.append(src)\n",
        "          labels.append(tgt)\n",
        "        return seqs, labels\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.seqs)\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        seq = self.seqs[index]\n",
        "        label = self.labels[index]\n",
        "        return (seq, label)\n"
      ],
      "metadata": {
        "id": "_qNyLnFdZxB3"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def generate_square_subsequent_mask(sz):\n",
        "    \"\"\"\n",
        "    Generates an upper-triangular matrix of ones, with zeros on diag.\n",
        "    Shape max_length * max_length\n",
        "    \"\"\"\n",
        "    return torch.triu(torch.ones(sz, sz, dtype=torch.bool), diagonal=1)\n",
        "\n",
        "def generate_key_padding_mask(max_len, seq_lengths):\n",
        "    \"\"\"\n",
        "    Generates a tensor of shape batch_size * max_length\n",
        "    \"\"\"\n",
        "    padding_mask = torch.zeros(seq_lengths.shape[0], max_len)\n",
        "    for i,x in enumerate(seq_lengths):\n",
        "        padding_mask[i][x:] = 1\n",
        "    return padding_mask"
      ],
      "metadata": {
        "id": "8YndaQPVB-Zl"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def collate_fn(batch, max_len=MAX_LENGTH):\n",
        "    seqs = []\n",
        "    labels = []\n",
        "    seq_lens = []\n",
        "\n",
        "    for (_seq, _label) in batch:\n",
        "        label_tensor = torch.tensor(_label, dtype=torch.int64)\n",
        "        labels.append(label_tensor)\n",
        "        seq_tensor = torch.tensor(_seq, dtype=torch.int64)\n",
        "        seqs.append(seq_tensor)\n",
        "        seq_lens.append(seq_tensor.size(0))\n",
        "\n",
        "    seq_lens = torch.tensor(seq_lens, dtype=torch.int64)\n",
        "    padded_seqs = pad_sequence(seqs, batch_first=True, padding_value=0)\n",
        "    padded_labels = pad_sequence(labels, batch_first=True, padding_value=0).reshape(-1)\n",
        "    \n",
        "    # max_len = seq_lens.max().item()\n",
        "    attn_mask = generate_square_subsequent_mask(max_len)\n",
        "    key_padding_mask = generate_key_padding_mask(max_len, seq_lens)\n",
        "    return padded_seqs, attn_mask, key_padding_mask, padded_labels"
      ],
      "metadata": {
        "id": "OhvqkR9gBvue"
      },
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_ds = TrainSequenceData(train_seqs, maxlen=MAX_LENGTH)\n",
        "train_dataloader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)\n",
        "len(train_dataloader)"
      ],
      "metadata": {
        "id": "NJOCN6VX_3j3",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "83b074b1-c954-481e-f401-b645ecd10f33"
      },
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "2819"
            ]
          },
          "metadata": {},
          "execution_count": 16
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class ValidTestSequenceData(Dataset):\n",
        "    def __init__(self, user_seqs, maxlen):\n",
        "        '''\n",
        "        Input is a list of lists where each list contains movie ids for a specific user. \n",
        "        '''\n",
        "        self.seqs = []\n",
        "        self.labels = []\n",
        "        self.maxlen = maxlen\n",
        "        for x in user_seqs:\n",
        "          if len(x) > 2:\n",
        "            seq, label = self.build_examples(x, maxlen)\n",
        "            self.seqs.append(seq)\n",
        "            self.labels.append(label)\n",
        "\n",
        "    def build_examples(self, x, max_len):\n",
        "        seq = x[:-1]\n",
        "        label = x[1:]\n",
        "        return seq, label\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.seqs)\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        seq = self.seqs[index]\n",
        "        label = self.labels[index]\n",
        "        return (seq, label)\n"
      ],
      "metadata": {
        "id": "Y2kihMHMSDWT"
      },
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "valid_ds = ValidTestSequenceData(valid_seqs, maxlen=MAX_LENGTH)\n",
        "valid_dataloader = DataLoader(valid_ds, batch_size=BATCH_SIZE, shuffle=False, collate_fn=collate_fn)\n",
        "len(valid_dataloader)"
      ],
      "metadata": {
        "id": "vb_M8UIRfRsT",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "ac2c1ccc-4ae9-49b9-dabb-0bf0f3544933"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "750"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model Architecture\n",
        "\n",
        "Here we will implement a transformer-based sequential recommendation model from scratch using PyTorch. For more detailed description of the model components and their implementation from scratch, refer to notebooks in this [GitHub repository](https://github.com/januverma/transformers-stuff) "
      ],
      "metadata": {
        "id": "dhYdq2EmtyDB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class MultiHeadSelfAttention(nn.Module):\n",
        "    '''\n",
        "    Implements MHSA using the PyTorch MultiheadAttention Layer.\n",
        "    '''\n",
        "    def __init__(self, hidden_dim, num_heads, dropout):\n",
        "        '''\n",
        "        Arguments:\n",
        "            hidden_dim: Dimension of the output of the self-attention.\n",
        "            num_heads: Number of heads for the multi-head attention. \n",
        "            dropout: Dropout probability for the self-attention. If `0.0` then no dropout will be used.\n",
        "            \n",
        "        Returns:\n",
        "            A tensor of shape `num_tokens x hidden_size` containing output of the MHSA for each token.\n",
        "        '''\n",
        "        super().__init__()\n",
        "        if hidden_dim % num_heads != 0:\n",
        "            print('The hidden size {} is not a multiple of the number of heads {}'.format(hidden_dim, num_heads))\n",
        "        self.attention_layer = nn.MultiheadAttention(hidden_dim, num_heads, dropout=dropout, batch_first=True)\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        '''\n",
        "        Arguments:\n",
        "            x: Tensor containing input token embeddings.\n",
        "            key_padding_mask: Mask indicating which elements within the input sequence to be considered as padding and ignored for the computation of self-attention scores.  \n",
        "            attention_mask: Mask indicating which relative positions are allowed to attend.  \n",
        "        '''\n",
        "        return self.attention_layer(query=x, key=x, value=x, key_padding_mask=key_padding_mask, attn_mask=attention_mask)\n"
      ],
      "metadata": {
        "id": "LFB7XrpLcCNm"
      },
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class FeedForward(nn.Module):\n",
        "    '''\n",
        "    Implements the feed-forward component of the transfomer model.\n",
        "    '''\n",
        "    def __init__(self, input_dim, hidden_dim, dropout=0.0):\n",
        "        '''\n",
        "        Arguments:\n",
        "            input_dim: Dimension of the token embedding, output of the MHSA layer.\n",
        "            hidden_dim: Hidden size of the Transformer that this feed-forward layer is part of.\n",
        "            dropout: Dropout probability to use for the projected activations. If `0.0` then no dropout will be used.\n",
        "        Returns:\n",
        "            A tensor of shape `num_tokens x hidden_dim` containing projections for each token.\n",
        "        '''\n",
        "        super().__init__()\n",
        "        self.activation = nn.ReLU()\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.layer_1 = nn.Linear(input_dim, hidden_dim)\n",
        "        self.layer_2 = nn.Linear(hidden_dim, input_dim)\n",
        "    def forward(self, x):\n",
        "        x = self.layer_1(x)\n",
        "        x = self.activation(x)\n",
        "        x = self.dropout(x)\n",
        "        x = self.layer_2(x)\n",
        "        return x"
      ],
      "metadata": {
        "id": "9dNbs2tgtzeu"
      },
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerLayerNorm(nn.Module):\n",
        "    '''\n",
        "    Implements LayerNorm for self-attention and feed-forward networks.\n",
        "\n",
        "    Arguments:\n",
        "        input_dim: Input dimension.\n",
        "    \n",
        "    Returns:\n",
        "        A normalized tensor of the same dimension as the input. \n",
        "    '''\n",
        "    def __init__(self, input_dim):\n",
        "        super().__init__()\n",
        "        self.layer_norm = nn.LayerNorm(input_dim)\n",
        "    def forward(self, x):\n",
        "        x = x.to(self.layer_norm.weight.dtype)\n",
        "        return self.layer_norm(x) "
      ],
      "metadata": {
        "id": "BwdXnM0Lt230"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerLayer(nn.Module):\n",
        "    '''\n",
        "    A transformer layer which is a sequential model consisting of self-attention, layer norm, residual connection, feed-forward projection, layer norm, residual connection. \n",
        "    \n",
        "    Arguments:\n",
        "        hidden_dim: Hidden dimension transformer layers.  \n",
        "        num_heads: Number of attention heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    Returns:\n",
        "        A tensor containing attention scores for each token. \n",
        "        attn_weights: A tensor of shape `num_tokens x num_tokens` containing the attention weights. \n",
        "    '''\n",
        "    def __init__(self, hidden_dim, num_heads, attn_dropout=0.0, ffn_dropout=0.0):\n",
        "        super().__init__()\n",
        "        self.attn_layer = MultiHeadSelfAttention(hidden_dim, num_heads, dropout=attn_dropout)\n",
        "        self.ffn_layer = FeedForward(hidden_dim, hidden_dim, dropout=ffn_dropout)\n",
        "        self.layer_norm = TransformerLayerNorm(hidden_dim)\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        attn_out, attn_weights = self.attn_layer(x, key_padding_mask, attention_mask)\n",
        "        x = self.layer_norm(x + attn_out)\n",
        "        ffn_out = self.ffn_layer(x)\n",
        "        x = self.layer_norm(x + ffn_out)\n",
        "        return x, attn_weights"
      ],
      "metadata": {
        "id": "8hk_0L35t6lu"
      },
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class TransformerEncoder(nn.Module):\n",
        "    '''\n",
        "    Transformer Encoder which is composed for a stack of TransformerLayers. \n",
        "\n",
        "    Arguments:\n",
        "        num_layers: Number of Transformer layers in the encoder. \n",
        "        hidden_dim: Hidden dimension of the transformer layers.  \n",
        "        num_heads: Number of heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    '''\n",
        "    def __init__(self, num_layers, hidden_dim, num_heads, attn_dropout=0.0, ffn_dropout=0.0):\n",
        "        super().__init__()\n",
        "        self.layers = nn.ModuleList([TransformerLayer(hidden_dim, num_heads, attn_dropout, ffn_dropout) for _ in range(num_layers)])\n",
        "        self.attn_weights = []\n",
        "    def forward(self, x, key_padding_mask=None, attention_mask=None):\n",
        "        for layer in self.layers:\n",
        "            x, weights = layer(x, key_padding_mask, attention_mask)\n",
        "            self.attn_weights.append(weights)\n",
        "        return x\n",
        "    def get_attention_weights(self):\n",
        "        if len(self.attn_weights) != 0:\n",
        "            return self.attn_weights\n",
        "        else:\n",
        "            print(\"The model hasn't been training yet\")"
      ],
      "metadata": {
        "id": "tVXYBsaruESF"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class PositionalEncoding(nn.Module):\n",
        "    '''\n",
        "    Implements the sinusoidal positional encoding for the input tokens. \n",
        "\n",
        "    Arguments:\n",
        "        embed_dim: Dimension of the positional encoding, should be the same as input token embedding. \n",
        "        dropout: Dropout probability to be used for positional encoding. \n",
        "        max_len: Maximum length of the input token sequences. \n",
        "    Returns:\n",
        "      A tensor containing positional embeddings for each token.\n",
        "    '''\n",
        "    def __init__(self, embed_dim, dropout=0.1, max_len=5000):\n",
        "        super().__init__()\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        position = torch.arange(max_len).unsqueeze(1)\n",
        "        div_term = torch.exp(torch.arange(0, embed_dim, 2) * (-math.log(10000.0)/embed_dim))\n",
        "        pe = torch.zeros(max_len, 1, embed_dim)\n",
        "        pe[:, 0, 0::2] = torch.sin(position * div_term)\n",
        "        pe[:, 0, 1::2] = torch.cos(position * div_term)\n",
        "        self.register_buffer('pe', pe)\n",
        "    def forward(self, x):\n",
        "        x = x + self.pe[:x.size(0)]\n",
        "        return self.dropout(x)"
      ],
      "metadata": {
        "id": "IcEw32f8uHCD"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class MySASRec(nn.Module):\n",
        "    '''\n",
        "    SASRec-like recommendation model.\n",
        "\n",
        "    Arguments:\n",
        "        vocab_size: Size of vocabulary.\n",
        "        embed_dim: Dimension of the input token embedding. \n",
        "        num_layers: Number of Transformer layers in the encoder. \n",
        "        hidden_dim: Hidden dimension of the transformer layers.  \n",
        "        ffn_hidden_dim: Hidden dimension of the Feed-forward layers. \n",
        "        num_heads: Number of heads. \n",
        "        attn_dropout: Dropout for MHSA layers. \n",
        "        ffn_dropout: Dropout for feed-forward layers.\n",
        "    '''\n",
        "    def __init__(self, vocab_size, embed_dim, num_layers, num_heads, attn_dropout, ffn_dropout):\n",
        "        super(MySASRec, self).__init__()\n",
        "        self.embedding_layer = nn.Embedding(vocab_size, embed_dim)\n",
        "        self.pos_encoder = PositionalEncoding(embed_dim)\n",
        "        self.transformer_encoder = TransformerEncoder(num_layers, embed_dim, num_heads, attn_dropout, ffn_dropout)\n",
        "        self.decoder = nn.Linear(embed_dim, vocab_size)\n",
        "        self.embed_layer_norm = nn.LayerNorm(embed_dim)\n",
        "        self.init_weights()\n",
        "\n",
        "    def init_weights(self):\n",
        "        init_range = 0.5\n",
        "        self.decoder.weight.data.uniform_(-init_range, init_range)\n",
        "        self.decoder.bias.data.zero_()\n",
        "        \n",
        "    \n",
        "    def forward(self, seqs, attn_mask=None, key_padding_mask=None):\n",
        "        embedded_seq = self.embedding_layer(seqs)\n",
        "        embedded_seq = self.pos_encoder(embedded_seq)\n",
        "        embedded_seq = self.embed_layer_norm(embedded_seq)\n",
        "        out = self.transformer_encoder(x=embedded_seq, key_padding_mask=key_padding_mask, attention_mask=attn_mask)\n",
        "        results = self.decoder(out)\n",
        "        return results\n",
        "      \n",
        "    def get_attn_weights(self):\n",
        "      return self.transformer_encoder.get_attention_weights()"
      ],
      "metadata": {
        "id": "UX-v1YUfuK2a"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training Module"
      ],
      "metadata": {
        "id": "zK0E3ZZhubm_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "num_users = max(ratings.userId)\n",
        "num_movies = max(ratings.movieId) + 1\n",
        "\n",
        "num_users, num_movies"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8u3snpflDU1Z",
        "outputId": "95dfd7d9-9edc-4397-be91-08e0b36aa407"
      },
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(6040, 3953)"
            ]
          },
          "metadata": {},
          "execution_count": 26
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from torch.nn import CrossEntropyLoss\n",
        "from torch.optim import SGD, lr_scheduler, Adam\n",
        "from torch.nn.utils import clip_grad_norm_"
      ],
      "metadata": {
        "id": "vM0y3ksy0a6B"
      },
      "execution_count": 27,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train(dataloader, model, criterion, optimizer, device):\n",
        "    \"\"\"\n",
        "    Trains an epoch.\n",
        "    \"\"\"\n",
        "    losses = []\n",
        "    running_losses = []\n",
        "    model.train()\n",
        "    for i, batch in enumerate(dataloader):\n",
        "        padded_seqs, attn_mask, key_padding_mask, out = batch\n",
        "        padded_seqs = padded_seqs.to(device)\n",
        "        attn_mask = attn_mask.to(device)\n",
        "        key_padding_mask = key_padding_mask.to(device)\n",
        "        out = out.to(device)\n",
        "\n",
        "        logits = model(padded_seqs, attn_mask, key_padding_mask)\n",
        "        J = criterion(logits.view(-1, num_movies), out)\n",
        "        losses.append(J.item())\n",
        "        running_losses.append(J.item())\n",
        "        optimizer.zero_grad()\n",
        "        J.backward()\n",
        "        clip_grad_norm_(model.parameters(), 0.5)\n",
        "        optimizer.step()\n",
        "\n",
        "        if i%1000 == 0:\n",
        "            print('|{:5d}/{:5d} batches done | loss {:8.3f}'.format(i, len(dataloader), torch.tensor(running_losses).mean()))\n",
        "            running_losses = []\n",
        "    \n",
        "    epoch_loss = torch.tensor(losses).mean()\n",
        "    return epoch_loss"
      ],
      "metadata": {
        "id": "LplQqC7BuU4y"
      },
      "execution_count": 28,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def evaluate(dataloader, model, criterion, device):\n",
        "    \"\"\"\n",
        "    Trains an epoch.\n",
        "    \"\"\"\n",
        "    losses = []\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "      for i, batch in enumerate(dataloader):\n",
        "          padded_seqs, attn_mask, key_padding_mask, out = batch\n",
        "          padded_seqs = padded_seqs.to(device)\n",
        "          attn_mask = attn_mask.to(device)\n",
        "          key_padding_mask = key_padding_mask.to(device)\n",
        "          out = out.to(device)\n",
        "\n",
        "          logits = model(padded_seqs, attn_mask, key_padding_mask)\n",
        "          J = criterion(logits.view(-1, num_movies), out)\n",
        "          losses.append(J.item())\n",
        "\n",
        "    epoch_loss = torch.tensor(losses).mean()\n",
        "    return epoch_loss"
      ],
      "metadata": {
        "id": "5mwaKUePmSA0"
      },
      "execution_count": 29,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training"
      ],
      "metadata": {
        "id": "gpGOpDpeDnf-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ppUtWf1ZibfM",
        "outputId": "3276c438-5b05-4536-9c28-7c9dcffe915b"
      },
      "execution_count": 30,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda', index=0)"
            ]
          },
          "metadata": {},
          "execution_count": 30
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "torch.cuda.empty_cache()"
      ],
      "metadata": {
        "id": "9b4cmaZJigfd"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import time"
      ],
      "metadata": {
        "id": "Ppu7-RYH80cw"
      },
      "execution_count": 32,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "EPOCHS = 10\n",
        "LR = 0.001\n",
        "rec_model = MySASRec(vocab_size=num_movies, embed_dim=50, num_layers=2, num_heads=2, attn_dropout=0.2, ffn_dropout=0.2).to(device)\n",
        "\n",
        "criterion = CrossEntropyLoss(ignore_index=0)\n",
        "optimizer = Adam(rec_model.parameters(), lr=LR, betas=(0.9, 0.98), eps=1e-08)\n",
        "scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 1.0, gamma=0.95)"
      ],
      "metadata": {
        "id": "43se8g829W_-"
      },
      "execution_count": 33,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "for epoch in range(EPOCHS):\n",
        "        epoch_start_time = time.time()\n",
        "        train_loss = train(train_dataloader, rec_model, criterion, optimizer, device)\n",
        "        val_loss = evaluate(valid_dataloader, rec_model, criterion, device)\n",
        "\n",
        "        print('-' * 59)\n",
        "        print('| end of epoch {:3d} | time: {:5.2f}s | val loss {:8.3f}'.\\\n",
        "            format(epoch, time.time() - epoch_start_time, val_loss))\n",
        "        print('-' * 59)\n",
        "        scheduler.step()\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ry6-5BBHGqoC",
        "outputId": "4a42b7fd-c027-405f-89a7-0911ac9197c9"
      },
      "execution_count": 40,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "|    0/ 2819 batches done | loss    5.680\n",
            "| 1000/ 2819 batches done | loss    6.050\n",
            "| 2000/ 2819 batches done | loss    6.049\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   0 | time: 25.29s | val loss    6.049\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.663\n",
            "| 1000/ 2819 batches done | loss    6.021\n",
            "| 2000/ 2819 batches done | loss    6.025\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   1 | time: 25.63s | val loss    6.024\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.590\n",
            "| 1000/ 2819 batches done | loss    6.000\n",
            "| 2000/ 2819 batches done | loss    6.004\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   2 | time: 25.10s | val loss    6.002\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.613\n",
            "| 1000/ 2819 batches done | loss    5.980\n",
            "| 2000/ 2819 batches done | loss    5.985\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   3 | time: 25.62s | val loss    5.983\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.517\n",
            "| 1000/ 2819 batches done | loss    5.964\n",
            "| 2000/ 2819 batches done | loss    5.970\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   4 | time: 25.57s | val loss    5.961\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.581\n",
            "| 1000/ 2819 batches done | loss    5.950\n",
            "| 2000/ 2819 batches done | loss    5.954\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   5 | time: 25.58s | val loss    5.951\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.604\n",
            "| 1000/ 2819 batches done | loss    5.936\n",
            "| 2000/ 2819 batches done | loss    5.941\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   6 | time: 25.85s | val loss    5.937\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.574\n",
            "| 1000/ 2819 batches done | loss    5.923\n",
            "| 2000/ 2819 batches done | loss    5.929\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   7 | time: 26.30s | val loss    5.924\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.530\n",
            "| 1000/ 2819 batches done | loss    5.915\n",
            "| 2000/ 2819 batches done | loss    5.919\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   8 | time: 25.87s | val loss    5.910\n",
            "-----------------------------------------------------------\n",
            "|    0/ 2819 batches done | loss    5.546\n",
            "| 1000/ 2819 batches done | loss    5.903\n",
            "| 2000/ 2819 batches done | loss    5.911\n",
            "-----------------------------------------------------------\n",
            "| end of epoch   9 | time: 25.83s | val loss    5.900\n",
            "-----------------------------------------------------------\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This concludes the training process. We trained a simplified model (2 transformer layers with 2 attention heads) model for 50 epochs. In practice, much deeper versions (e.g. BERT has 12 transformer layers each with 12 heads) are trained for a longer periods of time to achieve desired results.  "
      ],
      "metadata": {
        "id": "URZpmIfZtyCX"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Evaluation on the test data"
      ],
      "metadata": {
        "id": "QqvLmzHsEBgc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "test_seqs = user_seq_df['movieId'].apply(lambda x:create_test_seqs(x)).values\n",
        "test_seqs.shape, len(test_seqs[0])"
      ],
      "metadata": {
        "id": "Vso1bQAThnnw",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "abd45681-eec7-4fab-f018-d09ab6c46c28"
      },
      "execution_count": 35,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "((6040,), 51)"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_ds = ValidTestSequenceData(test_seqs, maxlen=MAX_LENGTH)\n",
        "test_dataloader = DataLoader(test_ds, batch_size=1, shuffle=False, collate_fn=collate_fn)\n",
        "len(test_dataloader)"
      ],
      "metadata": {
        "id": "sEDpUzYJIDls",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "45fb2622-8ae1-4e94-fa00-f8b3b7417640"
      },
      "execution_count": 36,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "6040"
            ]
          },
          "metadata": {},
          "execution_count": 36
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Our metric of choice is HitRate@k which is defined as "
      ],
      "metadata": {
        "id": "rYygS7iLKEp0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def calcualte_hit_rate(rec_model, test_loader, k=10):\n",
        "  total = 0\n",
        "  cnt = 0\n",
        "  for batch in test_dataloader:\n",
        "    padded_seqs, attn_mask, key_padding_mask, out = batch\n",
        "    logits = rec_model(padded_seqs.to(device))\n",
        "    logits = logits.squeeze()\n",
        "    pred = logits[-1].argsort(dim=0, descending=True)\n",
        "    if out[-1].item() in pred[:k].cpu().numpy():\n",
        "      cnt += 1\n",
        "      total += 1\n",
        "    else:\n",
        "      total += 1\n",
        "  return {'total instances': total, 'hits@k': cnt}\n"
      ],
      "metadata": {
        "id": "8uf7hv1lkJyX"
      },
      "execution_count": 37,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "calcualte_hit_rate(rec_model, test_dataloader, k=20)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I3diEYePn7ss",
        "outputId": "c3fee653-bea8-44af-89f5-77d1ccd2ad67"
      },
      "execution_count": 41,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "{'total instances': 6040, 'hits@k': 1264}"
            ]
          },
          "metadata": {},
          "execution_count": 41
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print('hit rate @ 20 after 50 epochs of training is : {}'.format(1264/6040))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "owXBVPjM6xEL",
        "outputId": "df42dc26-8cd0-46e3-a915-7e8b00fcefaf"
      },
      "execution_count": 42,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "hit rate @ 20 after 50 epochs of training is : 0.20927152317880796\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "ihSrj0RVe0I5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}