GNN_Hands_on_Part_3.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from torch import Tensor\n",
        "print(torch.__version__)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "uh7LaSzfwcj2",
        "outputId": "109d098b-bba1-42fd-f7ac-6c4c3bf90fc4"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "2.5.0+cu121\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AsWUrVh7vZtM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "1d4955a2-be9a-49dc-c75e-2c67740657ad"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in links: https://data.pyg.org/whl/torch-2.5.0+cu121.html\n",
            "Requirement already satisfied: torch-scatter in /usr/local/lib/python3.10/dist-packages (2.1.2+pt25cu121)\n",
            "Looking in links: https://data.pyg.org/whl/torch-2.5.0+cu121.html\n",
            "Requirement already satisfied: torch-sparse in /usr/local/lib/python3.10/dist-packages (0.6.18+pt25cu121)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from torch-sparse) (1.13.1)\n",
            "Requirement already satisfied: numpy<2.3,>=1.22.4 in /usr/local/lib/python3.10/dist-packages (from scipy->torch-sparse) (1.26.4)\n",
            "Looking in links: https://data.pyg.org/whl/nightly/torch-2.5.0+cu121.html\n",
            "Requirement already satisfied: pyg-lib in /usr/local/lib/python3.10/dist-packages (0.4.0.dev20241111+pt25cu121)\n",
            "Collecting git+https://github.com/pyg-team/pytorch_geometric.git\n",
            "  Cloning https://github.com/pyg-team/pytorch_geometric.git to /tmp/pip-req-build-6w07wugn\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/pyg-team/pytorch_geometric.git /tmp/pip-req-build-6w07wugn\n",
            "  Resolved https://github.com/pyg-team/pytorch_geometric.git to commit f5c829344517c823c24abb08ce2fc7cf00ff29f7\n",
            "  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (3.10.10)\n",
            "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (2024.10.0)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (3.1.4)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (1.26.4)\n",
            "Requirement already satisfied: psutil>=5.8.0 in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (5.9.5)\n",
            "Requirement already satisfied: pyparsing in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (3.2.0)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (2.32.3)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.7.0) (4.66.6)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (2.4.3)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (1.5.0)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (6.1.0)\n",
            "Requirement already satisfied: yarl<2.0,>=1.12.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (1.17.1)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.7.0) (4.0.3)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch-geometric==2.7.0) (3.0.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.7.0) (3.4.0)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.7.0) (3.10)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.7.0) (2.2.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.7.0) (2024.8.30)\n",
            "Requirement already satisfied: typing-extensions>=4.1.0 in /usr/local/lib/python3.10/dist-packages (from multidict<7.0,>=4.5->aiohttp->torch-geometric==2.7.0) (4.12.2)\n",
            "Requirement already satisfied: propcache>=0.2.0 in /usr/local/lib/python3.10/dist-packages (from yarl<2.0,>=1.12.0->aiohttp->torch-geometric==2.7.0) (0.2.0)\n"
          ]
        }
      ],
      "source": [
        "# Install required packages.\n",
        "import os\n",
        "os.environ['TORCH'] = torch.__version__\n",
        "\n",
        "!pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
        "!pip install torch-sparse -f https://data.pyg.org/whl/torch-${TORCH}.html\n",
        "!pip install pyg-lib -f https://data.pyg.org/whl/nightly/torch-${TORCH}.html\n",
        "!pip install git+https://github.com/pyg-team/pytorch_geometric.git"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Link Prediction on MovieLens\n",
        "\n",
        "This colab notebook shows how to load a set of `*.csv` files as input and construct a heterogeneous graph from it.\n",
        "We will then use this dataset as input into a [heterogeneous graph model](https://pytorch-geometric.readthedocs.io/en/latest/notes/heterogeneous.html#hgtutorial), and use it for the task of link prediction.\n",
        "A few code cells require user input to let the code run through successfully.\n",
        "If you are stuck on cells that require input, take a look at the fully filled out tutorial [here](https://medium.com/@pytorch_geometric/link-prediction-on-heterogeneous-graphs-with-pyg-6d5c29677c70).\n",
        "\n",
        "We are going to use the [MovieLens dataset](https://grouplens.org/datasets/movielens/) collected by the GroupLens research group.\n",
        "This toy dataset describes ratings and tagging activity from MovieLens.\n",
        "The dataset contains approximately 100k ratings across more than 9k movies from more than 600 users.\n",
        "We are going to use this dataset to generate two node types holding data for movies and users, respectively, and one edge type connecting users and movies, representing the relation of whether a user has rated a specific movie.\n",
        "\n",
        "The link prediction task then tries to predict missing ratings, and can, for example, be used to recommend users new movies."
      ],
      "metadata": {
        "id": "vit8xKCiXAue"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Heterogeneous Graph Creation\n",
        "\n",
        "First, we download the dataset to an arbitrary folder (in this case, the current directory):"
      ],
      "metadata": {
        "id": "N_pshZh8Y3dw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch_geometric.data import download_url, extract_zip\n",
        "\n",
        "url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'\n",
        "extract_zip(download_url(url, '.'), '.')\n",
        "\n",
        "movies_path = './ml-latest-small/movies.csv'\n",
        "ratings_path = './ml-latest-small/ratings.csv'"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ciYr7NWEv5_o",
        "outputId": "caef41a3-18bb-4f40-bbbe-be78592157c5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Using existing file ml-latest-small.zip\n",
            "Extracting ./ml-latest-small.zip\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Before we create the heterogeneous graph, let’s take a look at the data."
      ],
      "metadata": {
        "id": "B89RD_evY7uu"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "\n",
        "print('movies.csv:')\n",
        "print('===========')\n",
        "print(pd.read_csv(movies_path).head())\n",
        "print()\n",
        "print('ratings.csv:')\n",
        "print('============')\n",
        "print(pd.read_csv(ratings_path).head())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D6ixkcCOwCDi",
        "outputId": "67dcf369-04e2-4420-ae99-15bf53c8b573"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "movies.csv:\n",
            "===========\n",
            "   movieId                               title  \\\n",
            "0        1                    Toy Story (1995)   \n",
            "1        2                      Jumanji (1995)   \n",
            "2        3             Grumpier Old Men (1995)   \n",
            "3        4            Waiting to Exhale (1995)   \n",
            "4        5  Father of the Bride Part II (1995)   \n",
            "\n",
            "                                        genres  \n",
            "0  Adventure|Animation|Children|Comedy|Fantasy  \n",
            "1                   Adventure|Children|Fantasy  \n",
            "2                               Comedy|Romance  \n",
            "3                         Comedy|Drama|Romance  \n",
            "4                                       Comedy  \n",
            "\n",
            "ratings.csv:\n",
            "============\n",
            "   userId  movieId  rating  timestamp\n",
            "0       1        1     4.0  964982703\n",
            "1       1        3     4.0  964981247\n",
            "2       1        6     4.0  964982224\n",
            "3       1       47     5.0  964983815\n",
            "4       1       50     5.0  964982931\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "We see that the `movies.csv` file provides two useful columns: `movieId` assigns a unique identifier to each movie, while the `genres` column represent genres of the given movie.\n",
        "We can make use of this column to define a feature representation that can be easily interpreted by machine learning models."
      ],
      "metadata": {
        "id": "LgXipuWTZK39"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the entire movie data frame into memory:\n",
        "movies_df = pd.read_csv(movies_path, index_col='movieId')"
      ],
      "metadata": {
        "id": "dd_BJMNcwJ7k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "movies_df.columns"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9b7M6sIoxFYf",
        "outputId": "49edff50-4717-453e-c1d3-a928a1ef8cd5"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "Index(['title', 'genres'], dtype='object')"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Split genres and convert into indicator variables:\n",
        "genres = movies_df['genres'].str.get_dummies('|')\n",
        "print(genres[[\"Action\", \"Adventure\", \"Drama\", \"Horror\"]].head())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ExOVCJ3mAYuC",
        "outputId": "32f816d7-2fc1-4d7e-d458-478ce448e8d0"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "         Action  Adventure  Drama  Horror\n",
            "movieId                                  \n",
            "1             0          1      0       0\n",
            "2             0          1      0       0\n",
            "3             0          0      0       0\n",
            "4             0          0      1       0\n",
            "5             0          0      0       0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Feature sets should be converted to <font color=\"blue\">Tensors</font>"
      ],
      "metadata": {
        "id": "oVUIaQ21BNEP"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Use genres as movie input features:\n",
        "movie_feat = torch.from_numpy(genres.values).to(torch.float)\n",
        "assert movie_feat.size() == (9742, 20)  # 20 genres in total."
      ],
      "metadata": {
        "id": "uq2hNLggw-5I"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "The `ratings.csv` data connects users (as given by `userId`) and movies (as given by `movieId`).\n",
        "Due to simplicity, we do not make use of the additional `timestamp` and `rating` information.\n",
        "Here, we first read the `*.csv` file from disk, and create a mapping that maps entry IDs to a consecutive value in the range `{ 0, ..., num_rows - 1 }`.\n",
        "This is needed as we want our final data representation to be as compact as possible, *e.g.*, the representation of a movie in the first row should be accessible via `x[0]`.\n",
        "\n",
        "Afterwards, we obtain the final `edge_index` representation of shape `[2, num_ratings]` from `ratings.csv` by merging mapped user and movie indices with the raw indices given by the original data frame."
      ],
      "metadata": {
        "id": "KLaDBVP4plsS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Load the entire ratings data frame into memory:\n",
        "ratings_df = pd.read_csv(ratings_path)\n",
        "ratings_df.head()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 206
        },
        "id": "Mfpz0uEwBcSd",
        "outputId": "404aa263-1bff-4505-859d-fc0ab5e63447"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "   userId  movieId  rating  timestamp\n",
              "0       1        1     4.0  964982703\n",
              "1       1        3     4.0  964981247\n",
              "2       1        6     4.0  964982224\n",
              "3       1       47     5.0  964983815\n",
              "4       1       50     5.0  964982931"
            ],
            "text/html": [
              "\n",
              "  <div id=\"df-3fa1033f-69f7-4a5c-b16a-14fdd221d410\" 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>1</td>\n",
              "      <td>4.0</td>\n",
              "      <td>964982703</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>1</th>\n",
              "      <td>1</td>\n",
              "      <td>3</td>\n",
              "      <td>4.0</td>\n",
              "      <td>964981247</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>2</th>\n",
              "      <td>1</td>\n",
              "      <td>6</td>\n",
              "      <td>4.0</td>\n",
              "      <td>964982224</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>3</th>\n",
              "      <td>1</td>\n",
              "      <td>47</td>\n",
              "      <td>5.0</td>\n",
              "      <td>964983815</td>\n",
              "    </tr>\n",
              "    <tr>\n",
              "      <th>4</th>\n",
              "      <td>1</td>\n",
              "      <td>50</td>\n",
              "      <td>5.0</td>\n",
              "      <td>964982931</td>\n",
              "    </tr>\n",
              "  </tbody>\n",
              "</table>\n",
              "</div>\n",
              "    <div class=\"colab-df-buttons\">\n",
              "\n",
              "  <div class=\"colab-df-container\">\n",
              "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-3fa1033f-69f7-4a5c-b16a-14fdd221d410')\"\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 -960 960 960\">\n",
              "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
              "  </svg>\n",
              "    </button>\n",
              "\n",
              "  <style>\n",
              "    .colab-df-container {\n",
              "      display:flex;\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",
              "    .colab-df-buttons div {\n",
              "      margin-bottom: 4px;\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-3fa1033f-69f7-4a5c-b16a-14fdd221d410 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-3fa1033f-69f7-4a5c-b16a-14fdd221d410');\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",
              "\n",
              "\n",
              "<div id=\"df-3aec1218-e883-4bb3-a611-a672e01ac2d9\">\n",
              "  <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-3aec1218-e883-4bb3-a611-a672e01ac2d9')\"\n",
              "            title=\"Suggest charts\"\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",
              "    <g>\n",
              "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
              "    </g>\n",
              "</svg>\n",
              "  </button>\n",
              "\n",
              "<style>\n",
              "  .colab-df-quickchart {\n",
              "      --bg-color: #E8F0FE;\n",
              "      --fill-color: #1967D2;\n",
              "      --hover-bg-color: #E2EBFA;\n",
              "      --hover-fill-color: #174EA6;\n",
              "      --disabled-fill-color: #AAA;\n",
              "      --disabled-bg-color: #DDD;\n",
              "  }\n",
              "\n",
              "  [theme=dark] .colab-df-quickchart {\n",
              "      --bg-color: #3B4455;\n",
              "      --fill-color: #D2E3FC;\n",
              "      --hover-bg-color: #434B5C;\n",
              "      --hover-fill-color: #FFFFFF;\n",
              "      --disabled-bg-color: #3B4455;\n",
              "      --disabled-fill-color: #666;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart {\n",
              "    background-color: var(--bg-color);\n",
              "    border: none;\n",
              "    border-radius: 50%;\n",
              "    cursor: pointer;\n",
              "    display: none;\n",
              "    fill: var(--fill-color);\n",
              "    height: 32px;\n",
              "    padding: 0;\n",
              "    width: 32px;\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart:hover {\n",
              "    background-color: var(--hover-bg-color);\n",
              "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
              "    fill: var(--button-hover-fill-color);\n",
              "  }\n",
              "\n",
              "  .colab-df-quickchart-complete:disabled,\n",
              "  .colab-df-quickchart-complete:disabled:hover {\n",
              "    background-color: var(--disabled-bg-color);\n",
              "    fill: var(--disabled-fill-color);\n",
              "    box-shadow: none;\n",
              "  }\n",
              "\n",
              "  .colab-df-spinner {\n",
              "    border: 2px solid var(--fill-color);\n",
              "    border-color: transparent;\n",
              "    border-bottom-color: var(--fill-color);\n",
              "    animation:\n",
              "      spin 1s steps(1) infinite;\n",
              "  }\n",
              "\n",
              "  @keyframes spin {\n",
              "    0% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "      border-left-color: var(--fill-color);\n",
              "    }\n",
              "    20% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    30% {\n",
              "      border-color: transparent;\n",
              "      border-left-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    40% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-top-color: var(--fill-color);\n",
              "    }\n",
              "    60% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "    }\n",
              "    80% {\n",
              "      border-color: transparent;\n",
              "      border-right-color: var(--fill-color);\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "    90% {\n",
              "      border-color: transparent;\n",
              "      border-bottom-color: var(--fill-color);\n",
              "    }\n",
              "  }\n",
              "</style>\n",
              "\n",
              "  <script>\n",
              "    async function quickchart(key) {\n",
              "      const quickchartButtonEl =\n",
              "        document.querySelector('#' + key + ' button');\n",
              "      quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
              "      quickchartButtonEl.classList.add('colab-df-spinner');\n",
              "      try {\n",
              "        const charts = await google.colab.kernel.invokeFunction(\n",
              "            'suggestCharts', [key], {});\n",
              "      } catch (error) {\n",
              "        console.error('Error during call to suggestCharts:', error);\n",
              "      }\n",
              "      quickchartButtonEl.classList.remove('colab-df-spinner');\n",
              "      quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
              "    }\n",
              "    (() => {\n",
              "      let quickchartButtonEl =\n",
              "        document.querySelector('#df-3aec1218-e883-4bb3-a611-a672e01ac2d9 button');\n",
              "      quickchartButtonEl.style.display =\n",
              "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
              "    })();\n",
              "  </script>\n",
              "</div>\n",
              "\n",
              "    </div>\n",
              "  </div>\n"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "dataframe",
              "variable_name": "ratings_df"
            }
          },
          "metadata": {},
          "execution_count": 31
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The user_ids are not given sequentially. So we need to convert them into a range of values named as **mappedID** starting from **0** to the total number of **unique** user nodes"
      ],
      "metadata": {
        "id": "VxSGz0fFB1BY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a mapping from unique user indices to range [0, num_user_nodes):\n",
        "unique_user_id = ratings_df['userId'].unique()\n",
        "unique_user_id = pd.DataFrame(data={\n",
        "    'userId': unique_user_id,\n",
        "    'mappedID': pd.RangeIndex(len(unique_user_id)),\n",
        "})\n",
        "\n",
        "print(\"Mapping of user IDs to consecutive values:\")\n",
        "print(\"==========================================\")\n",
        "print(unique_user_id.head())\n",
        "print(f\"\\nThe number of unique Users : {len(unique_user_id)}\")\n",
        "print()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jbpCT8wtByZp",
        "outputId": "2d4edc27-2d27-4a13-e44a-33ae3f22da3f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mapping of user IDs to consecutive values:\n",
            "==========================================\n",
            "   userId  mappedID\n",
            "0       1         0\n",
            "1       2         1\n",
            "2       3         2\n",
            "3       4         3\n",
            "4       5         4\n",
            "\n",
            "The number of unique Users : 610\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Create a mapping from unique movie indices to range [0, num_movie_nodes):\n",
        "unique_movie_id = pd.DataFrame(data={\n",
        "    'movieId': movies_df.index,\n",
        "    'mappedID': pd.RangeIndex(len(movies_df)),\n",
        "})\n",
        "print(\"Mapping of movie IDs to consecutive values:\")\n",
        "print(\"===========================================\")\n",
        "print(unique_movie_id.head())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9wIIJ7fDGAwY",
        "outputId": "f4a78003-dd85-45d0-8f89-cda7f139e076"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mapping of movie IDs to consecutive values:\n",
            "===========================================\n",
            "   movieId  mappedID\n",
            "0        1         0\n",
            "1        2         1\n",
            "2        3         2\n",
            "3        4         3\n",
            "4        5         4\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "First two arguments takes two dfs to join  \n",
        "\n",
        "left_on -> the column from the left dataframe to **join on**  \n",
        "right_on -> the column from the right dataframe  \n",
        "\n",
        "how = 'left' -> perform a left join\n",
        "\n",
        "From these two joins below, we detect the mapping of unique user id to unique movie id for each rating in the **ratings_df**\n",
        "\n",
        "This helps us build the edge_index of **user -> movie**"
      ],
      "metadata": {
        "id": "EVzZH543SIeL"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Perform merge to obtain the edges from users and movies:\n",
        "ratings_user_id = pd.merge(ratings_df['userId'], unique_user_id,\n",
        "                            left_on='userId', right_on='userId', how='left')\n",
        "ratings_user_id = torch.from_numpy(ratings_user_id['mappedID'].values)\n",
        "ratings_user_id.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JpsrKGqTGGdA",
        "outputId": "282e86e6-e55b-4a26-88a8-58237e03eb94"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([100836])"
            ]
          },
          "metadata": {},
          "execution_count": 34
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "ratings_movie_id = pd.merge(ratings_df['movieId'], unique_movie_id,\n",
        "                            left_on='movieId', right_on='movieId', how='left')\n",
        "ratings_movie_id = torch.from_numpy(ratings_movie_id['mappedID'].values)\n",
        "ratings_movie_id.shape"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fC3RJeBRGHSJ",
        "outputId": "129b3fcd-2108-44e6-fd3a-5c1703bbba02"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "torch.Size([100836])"
            ]
          },
          "metadata": {},
          "execution_count": 35
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# With this, we are ready to construct our `edge_index` in COO format\n",
        "# following PyG semantics:\n",
        "edge_index_user_to_movie = torch.stack([ratings_user_id, ratings_movie_id], dim=0)\n",
        "assert edge_index_user_to_movie.size() == (2, 100836)\n",
        "\n",
        "print()\n",
        "print(\"Final edge indices pointing from users to movies:\")\n",
        "print(\"=================================================\")\n",
        "print(edge_index_user_to_movie)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JMGYv83WzSRr",
        "outputId": "69ae3d96-d55b-4c81-e6ed-3ed05e549113"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Final edge indices pointing from users to movies:\n",
            "=================================================\n",
            "tensor([[   0,    0,    0,  ...,  609,  609,  609],\n",
            "        [   0,    2,    5,  ..., 9462, 9463, 9503]])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "With this, we are ready to initialize our `HeteroData` object and pass the necessary information to it.\n",
        "Note that we also pass in a `node_id` vector to each node type in order to reconstruct the original node indices from sampled subgraphs.\n",
        "We also take care of adding reverse edges to the `HeteroData` object.\n",
        "This allows our GNN model to use both directions of the edge for message passing:"
      ],
      "metadata": {
        "id": "9w9fCnjvqmd2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch_geometric.data import HeteroData\n",
        "import torch_geometric.transforms as T\n",
        "\n",
        "data = HeteroData()\n",
        "\n",
        "# Save node indices:\n",
        "data[\"user\"].node_id = torch.arange(len(unique_user_id))\n",
        "data[\"movie\"].node_id = torch.arange(len(movies_df))\n",
        "\n",
        "# Add the node features and edge indices:\n",
        "data[\"movie\"].x = movie_feat\n",
        "data[\"user\", \"rates\", \"movie\"].edge_index = edge_index_user_to_movie\n",
        "\n",
        "# We also need to make sure to add the reverse edges from movies to users\n",
        "# in order to let a GNN be able to pass messages in both directions.\n",
        "# We can leverage the `T.ToUndirected()` transform for this from PyG:\n",
        "data = T.ToUndirected()(data)\n",
        "\n",
        "data.validate(raise_on_error = True)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "X8gZar8WY8Cx",
        "outputId": "33a51688-619d-42d0-85db-c8f591525056"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "metadata": {},
          "execution_count": 43
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(data)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R9R22pEAY9qO",
        "outputId": "289cbb5e-2876-4a38-c9a1-62ad844b2d7d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HeteroData(\n",
            "  user={ node_id=[610] },\n",
            "  movie={\n",
            "    node_id=[9742],\n",
            "    x=[9742, 20],\n",
            "  },\n",
            "  (user, rates, movie)={ edge_index=[2, 100836] },\n",
            "  (movie, rev_rates, user)={ edge_index=[2, 100836] }\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "assert data.node_types == [\"user\", \"movie\"]\n",
        "assert data.edge_types == [(\"user\", \"rates\", \"movie\"),\n",
        "                           (\"movie\", \"rev_rates\", \"user\")]"
      ],
      "metadata": {
        "id": "0d_uwyMLZBra"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "assert data[\"user\"].num_nodes == 610\n",
        "assert data[\"user\"].num_features == 0\n",
        "assert data[\"movie\"].num_nodes == 9742\n",
        "assert data[\"movie\"].num_features == 20\n",
        "assert data[\"user\", \"rates\", \"movie\"].num_edges == 100836\n",
        "assert data[\"movie\", \"rev_rates\", \"user\"].num_edges == 100836"
      ],
      "metadata": {
        "id": "I_63--974srt"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining Edge-level Training Splits\n",
        "\n",
        "Since our data is now ready-to-be-used, we can split the ratings of users into training, validation, and test splits.\n",
        "This is needed in order to ensure that we leak no information about edges used during evaluation into the training phase.\n",
        "For this, we make use of the [`transforms.RandomLinkSplit`](https://pytorch-geometric.readthedocs.io/en/latest/modules/transforms.html#torch_geometric.transforms.RandomLinkSplit) transformation from PyG.\n",
        "This transforms randomly divides the edges in the `(\"user\", \"rates\", \"movie\")` into training, validation and test edges.\n",
        "The `disjoint_train_ratio` parameter further separates edges in the training split into edges used for message passing (`edge_index`) and edges used for supervision (`edge_label_index`).\n",
        "Note that we also need to specify the reverse edge type `(\"movie\", \"rev_rates\", \"user\")`.\n",
        "This allows the `RandomLinkSplit` transform to drop reverse edges accordingly to not leak any information into the training phase."
      ],
      "metadata": {
        "id": "2QGdkLAurBq9"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# For this, we first split the set of edges into\n",
        "# training (80%), validation (10%), and testing edges (10%).\n",
        "# Across the training edges, we use 70% of edges for message passing,\n",
        "# and 30% of edges for supervision.\n",
        "# We further want to generate fixed negative edges for evaluation with a ratio of 2:1.\n",
        "# Negative edges during training will be generated on-the-fly, so we don't want to\n",
        "# add them to the graph right away.\n",
        "# Overall, we can leverage the `RandomLinkSplit()` transform for this from PyG:\n",
        "transform = T.RandomLinkSplit(\n",
        "    num_val=0.1,\n",
        "    num_test=0.1,\n",
        "    disjoint_train_ratio=0.3,\n",
        "    neg_sampling_ratio=2.0,\n",
        "    add_negative_train_samples=False,\n",
        "    edge_types=(\"user\", \"rates\", \"movie\"),\n",
        "    rev_edge_types=(\"movie\", \"rev_rates\", \"user\"),\n",
        ")\n",
        "\n",
        "train_data, val_data, test_data = transform(data)\n",
        "print(\"Training data:\")\n",
        "print(\"==============\")\n",
        "print(train_data)\n",
        "print()\n",
        "print(\"Validation data:\")\n",
        "print(\"================\")\n",
        "print(val_data)\n",
        "\n",
        "# train_edges = (main * (1 - (0.1 + 0.1))) * (1 - 0.3) = 100836 * 0.8 * 0.7 = 56469\n",
        "# train_edge_label_index_size = main * 0.8 * 0.3 = 24201 (no negative samples in the train data)\n",
        "\n",
        "# in the train part, the edge_index consists the edges for message passing (0.7 portion of the train_edges)\n",
        "# the edge_label_index consists the edges for supervision (0.3 portion of the train_edges)\n",
        "# the edge_label_index does not consist of any negative train samples, so no augmentation to the actual number, meaning (main * 0.8 * 0.3) = 24201\n",
        "\n",
        "assert train_data[\"user\", \"rates\", \"movie\"].num_edges == 56469                  # edge_index size\n",
        "assert train_data[\"user\", \"rates\", \"movie\"].edge_label_index.size(1) == 24201   # edge_label_index size\n",
        "assert train_data[\"movie\", \"rev_rates\", \"user\"].num_edges == 56469              # same as edge_index size\n",
        "\n",
        "# No negative edges added: (negative edge means 0)\n",
        "assert train_data[\"user\", \"rates\", \"movie\"].edge_label.min() == 1\n",
        "assert train_data[\"user\", \"rates\", \"movie\"].edge_label.max() == 1\n",
        "\n",
        "assert val_data[\"user\", \"rates\", \"movie\"].num_edges == 80670\n",
        "assert val_data[\"user\", \"rates\", \"movie\"].edge_label_index.size(1) == 30249\n",
        "assert val_data[\"movie\", \"rev_rates\", \"user\"].num_edges == 80670\n",
        "\n",
        "# Negative edges with ratio 2:1:\n",
        "assert val_data[\"user\", \"rates\", \"movie\"].edge_label.long().bincount().tolist() == [20166, 10083]"
      ],
      "metadata": {
        "id": "rwgNwoa26Eja",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "15e28795-f6cf-4076-9a0f-26b1f712ad96"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Training data:\n",
            "==============\n",
            "HeteroData(\n",
            "  user={ node_id=[610] },\n",
            "  movie={\n",
            "    node_id=[9742],\n",
            "    x=[9742, 20],\n",
            "  },\n",
            "  (user, rates, movie)={\n",
            "    edge_index=[2, 56469],\n",
            "    edge_label=[24201],\n",
            "    edge_label_index=[2, 24201],\n",
            "  },\n",
            "  (movie, rev_rates, user)={ edge_index=[2, 56469] }\n",
            ")\n",
            "\n",
            "Validation data:\n",
            "================\n",
            "HeteroData(\n",
            "  user={ node_id=[610] },\n",
            "  movie={\n",
            "    node_id=[9742],\n",
            "    x=[9742, 20],\n",
            "  },\n",
            "  (user, rates, movie)={\n",
            "    edge_index=[2, 80670],\n",
            "    edge_label=[30249],\n",
            "    edge_label_index=[2, 30249],\n",
            "  },\n",
            "  (movie, rev_rates, user)={ edge_index=[2, 80670] }\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Defining Mini-batch Loaders\n",
        "\n",
        "We are now ready to create a mini-batch loader that will generate subgraphs that can be used as input into our GNN.\n",
        "While this step is not strictly necessary for small-scale graphs, it is absolutely necessary to apply GNNs on larger graphs that do not fit onto GPU memory otherwise.\n",
        "Here, we make use of the [`loader.LinkNeighborLoader`](https://pytorch-geometric.readthedocs.io/en/latest/modules/loader.html#torch_geometric.loader.LinkNeighborLoader) which samples multiple hops from both ends of a link and creates a subgraph from it.\n",
        "Here, `edge_label_index` serves as the \"seed links\" to start sampling from."
      ],
      "metadata": {
        "id": "prKLwq6RsYoh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# In the first hop, we sample at most 20 neighbors.\n",
        "# In the second hop, we sample at most 10 neighbors.\n",
        "# In addition, during training, we want to sample negative edges on-the-fly with\n",
        "# a ratio of 2:1.\n",
        "# We can make use of the `loader.LinkNeighborLoader` from PyG:\n",
        "from torch_geometric.loader import LinkNeighborLoader\n",
        "\n",
        "# Define seed edges:\n",
        "edge_label_index = train_data[\"user\", \"rates\", \"movie\"].edge_label_index\n",
        "edge_label = train_data[\"user\", \"rates\", \"movie\"].edge_label\n",
        "\n",
        "train_loader = LinkNeighborLoader(\n",
        "    data=train_data,\n",
        "    num_neighbors=[20, 10],\n",
        "    neg_sampling_ratio=2.0,\n",
        "    edge_label_index=((\"user\", \"rates\", \"movie\"), edge_label_index),\n",
        "    edge_label=edge_label,\n",
        "    batch_size=128,\n",
        "    shuffle=True,\n",
        ")\n",
        "\n",
        "# Inspect a sample:\n",
        "sampled_data = next(iter(train_loader))\n",
        "\n",
        "print(\"Sampled mini-batch:\")\n",
        "print(\"===================\")\n",
        "print(sampled_data)\n",
        "\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label_index.size(1) == 3 * 128\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label.min() == 0\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label.max() == 1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Ogh615ka9I2c",
        "outputId": "60ea5b41-2f19-414f-fa45-6ff9e74461cb"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sampled mini-batch:\n",
            "===================\n",
            "HeteroData(\n",
            "  user={\n",
            "    node_id=[606],\n",
            "    n_id=[606],\n",
            "    num_sampled_nodes=[3],\n",
            "  },\n",
            "  movie={\n",
            "    node_id=[2722],\n",
            "    x=[2722, 20],\n",
            "    n_id=[2722],\n",
            "    num_sampled_nodes=[3],\n",
            "  },\n",
            "  (user, rates, movie)={\n",
            "    edge_index=[2, 17575],\n",
            "    edge_label=[384],\n",
            "    edge_label_index=[2, 384],\n",
            "    e_id=[17575],\n",
            "    num_sampled_edges=[2],\n",
            "    input_id=[128],\n",
            "  },\n",
            "  (movie, rev_rates, user)={\n",
            "    edge_index=[2, 7873],\n",
            "    e_id=[7873],\n",
            "    num_sampled_edges=[2],\n",
            "  }\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Creating a Heterogeneous Link-level GNN\n",
        "\n",
        "We are now ready to create our heterogeneous GNN.\n",
        "The GNN is responsible for learning enriched node representations from the surrounding subgraphs, which can be then used to derive edge-level predictions.\n",
        "For defining our heterogenous GNN, we make use of [`nn.SAGEConv`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.conv.SAGEConv) and the [`nn.to_hetero()`](https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.to_hetero_transformer.to_hetero) function, which transforms a GNN defined on homogeneous graphs to be applied on heterogeneous ones.\n",
        "\n",
        "In addition, we define a final link-level classifier, which simply takes both node embeddings of the link we are trying to predict, and applies a dot-product on them.\n",
        "\n",
        "As users do not have any node-level information, we choose to learn their features jointly via a `torch.nn.Embedding` layer. In order to improve the expressiveness of movie features, we do the same for movie nodes, and simply add their shallow embeddings to the pre-defined genre features."
      ],
      "metadata": {
        "id": "uj7biOtatAmG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from torch_geometric.nn import SAGEConv, to_hetero\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class GNN(torch.nn.Module):\n",
        "    def __init__(self, hidden_channels):\n",
        "        super().__init__()\n",
        "\n",
        "        self.conv1 = SAGEConv(hidden_channels, hidden_channels)\n",
        "        self.conv2 = SAGEConv(hidden_channels, hidden_channels)\n",
        "\n",
        "    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:\n",
        "      x = F.relu(self.conv1(x, edge_index))\n",
        "      x = self.conv2(x, edge_index)\n",
        "      return x\n",
        "\n",
        "# Our final classifier applies the dot-product between source and destination\n",
        "# node embeddings to derive edge-level predictions:\n",
        "class Classifier(torch.nn.Module):\n",
        "    def forward(self, x_user: Tensor, x_movie: Tensor, edge_label_index: Tensor) -> Tensor:\n",
        "        # Convert node embeddings to edge-level representations:\n",
        "        edge_feat_user = x_user[edge_label_index[0]]\n",
        "        edge_feat_movie = x_movie[edge_label_index[1]]\n",
        "\n",
        "        # Apply dot-product to get a prediction per supervision edge:\n",
        "        return (edge_feat_user * edge_feat_movie).sum(dim=-1)\n",
        "\n",
        "\n",
        "class Model(torch.nn.Module):\n",
        "    def __init__(self, hidden_channels):\n",
        "        super().__init__()\n",
        "        # Since the dataset does not come with rich features, we also learn two\n",
        "        # embedding matrices for users and movies:\n",
        "        self.movie_lin = torch.nn.Linear(20, hidden_channels)\n",
        "        self.user_emb = torch.nn.Embedding(data[\"user\"].num_nodes, hidden_channels)\n",
        "        self.movie_emb = torch.nn.Embedding(data[\"movie\"].num_nodes, hidden_channels)\n",
        "\n",
        "        # Instantiate homogeneous GNN:\n",
        "        self.gnn = GNN(hidden_channels)\n",
        "\n",
        "        # Convert GNN model into a heterogeneous variant:\n",
        "        self.gnn = to_hetero(self.gnn, metadata=data.metadata())\n",
        "\n",
        "        self.classifier = Classifier()\n",
        "\n",
        "    def forward(self, data: HeteroData) -> Tensor:\n",
        "        x_dict = {\n",
        "          \"user\": self.user_emb(data[\"user\"].node_id),\n",
        "          \"movie\": self.movie_lin(data[\"movie\"].x) + self.movie_emb(data[\"movie\"].node_id),\n",
        "        }\n",
        "\n",
        "        # `x_dict` holds feature matrices of all node types\n",
        "        # `edge_index_dict` holds all edge indices of all edge types\n",
        "        x_dict = self.gnn(x_dict, data.edge_index_dict)\n",
        "\n",
        "        pred = self.classifier(\n",
        "            x_dict[\"user\"],\n",
        "            x_dict[\"movie\"],\n",
        "            data[\"user\", \"rates\", \"movie\"].edge_label_index,\n",
        "        )\n",
        "\n",
        "        return pred\n",
        "\n",
        "\n",
        "model = Model(hidden_channels=64)\n",
        "\n",
        "print(model)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ebsFf-Pr_4LF",
        "outputId": "7dcbaf9c-75a0-492b-aac1-18af0939ed97"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model(\n",
            "  (movie_lin): Linear(in_features=20, out_features=64, bias=True)\n",
            "  (user_emb): Embedding(610, 64)\n",
            "  (movie_emb): Embedding(9742, 64)\n",
            "  (gnn): GraphModule(\n",
            "    (conv1): ModuleDict(\n",
            "      (user__rates__movie): SAGEConv(64, 64, aggr=mean)\n",
            "      (movie__rev_rates__user): SAGEConv(64, 64, aggr=mean)\n",
            "    )\n",
            "    (conv2): ModuleDict(\n",
            "      (user__rates__movie): SAGEConv(64, 64, aggr=mean)\n",
            "      (movie__rev_rates__user): SAGEConv(64, 64, aggr=mean)\n",
            "    )\n",
            "  )\n",
            "  (classifier): Classifier()\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Training a Heterogeneous Link-level GNN\n",
        "\n",
        "Training our GNN is then similar to training any PyTorch model.\n",
        "We move the model to the desired device, and initialize an optimizer that takes care of adjusting model parameters via stochastic gradient descent.\n",
        "\n",
        "The training loop then iterates over our mini-batches, applies the forward computation of the model, computes the loss from ground-truth labels and obtained predictions (here we make use of binary cross entropy), and adjusts model parameters via back-propagation and stochastic gradient descent."
      ],
      "metadata": {
        "id": "05dfew-WuWHN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import tqdm\n",
        "import torch.nn.functional as F\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(f\"Device: '{device}'\")\n",
        "\n",
        "model = model.to(device)\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
        "\n",
        "for epoch in range(1, 6):\n",
        "    total_loss = total_examples = 0\n",
        "    for sampled_data in tqdm.tqdm(train_loader):\n",
        "        optimizer.zero_grad()\n",
        "\n",
        "        sampled_data.to(device)\n",
        "        pred = model(sampled_data)\n",
        "\n",
        "        ground_truth = sampled_data[\"user\", \"rates\", \"movie\"].edge_label\n",
        "        loss = F.binary_cross_entropy_with_logits(pred, ground_truth)\n",
        "\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        total_loss += float(loss) * pred.numel()\n",
        "        total_examples += pred.numel()\n",
        "    print(f\"Epoch: {epoch:03d}, Loss: {total_loss / total_examples:.4f}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OqLuXEcrAMru",
        "outputId": "3bceb1af-da64-42da-eff4-c41cd0d15fd7"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Device: 'cpu'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:07<00:00, 25.11it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 001, Loss: 0.4339\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:08<00:00, 23.56it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 002, Loss: 0.3505\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:07<00:00, 24.54it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 003, Loss: 0.3299\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:07<00:00, 26.80it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 004, Loss: 0.3135\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:07<00:00, 23.79it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 005, Loss: 0.3018\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Evaluating a Heterogeneous Link-level GNN\n",
        "\n",
        "After training, we evaluate our model on useen data coming from the validation set.\n",
        "For this, we define a new `LinkNeighborLoader` (which now iterates over the edges in the validation set), obtain the predictions on validation edges by running the model, and finally evaluate the performance of the model by computing the AUC score over the set of predictions and their corresponding ground-truth edges (including both positive and negative edges)."
      ],
      "metadata": {
        "id": "Yq-I2xaYueF0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Define the validation seed edges:\n",
        "edge_label_index = val_data[\"user\", \"rates\", \"movie\"].edge_label_index\n",
        "edge_label = val_data[\"user\", \"rates\", \"movie\"].edge_label\n",
        "\n",
        "val_loader = LinkNeighborLoader(\n",
        "    data=val_data,\n",
        "    num_neighbors=[20, 10],\n",
        "    edge_label_index=((\"user\", \"rates\", \"movie\"), edge_label_index),\n",
        "    edge_label=edge_label,\n",
        "    batch_size=3 * 128,\n",
        "    shuffle=False,\n",
        ")\n",
        "\n",
        "sampled_data = next(iter(val_loader))\n",
        "\n",
        "print(\"Sampled mini-batch:\")\n",
        "print(\"===================\")\n",
        "print(sampled_data)\n",
        "\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label_index.size(1) == 3 * 128\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label.min() >= 0\n",
        "assert sampled_data[\"user\", \"rates\", \"movie\"].edge_label.max() <= 1"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZIrRn9YoNllj",
        "outputId": "f326931d-7113-464e-dce4-51fff917b6ab"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sampled mini-batch:\n",
            "===================\n",
            "HeteroData(\n",
            "  user={\n",
            "    node_id=[607],\n",
            "    n_id=[607],\n",
            "    num_sampled_nodes=[3],\n",
            "  },\n",
            "  movie={\n",
            "    node_id=[2648],\n",
            "    x=[2648, 20],\n",
            "    n_id=[2648],\n",
            "    num_sampled_nodes=[3],\n",
            "  },\n",
            "  (user, rates, movie)={\n",
            "    edge_index=[2, 18505],\n",
            "    edge_label=[384],\n",
            "    edge_label_index=[2, 384],\n",
            "    e_id=[18505],\n",
            "    num_sampled_edges=[2],\n",
            "    input_id=[384],\n",
            "  },\n",
            "  (movie, rev_rates, user)={\n",
            "    edge_index=[2, 7463],\n",
            "    e_id=[7463],\n",
            "    num_sampled_edges=[2],\n",
            "  }\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.metrics import roc_auc_score\n",
        "\n",
        "preds = []\n",
        "ground_truths = []\n",
        "for sampled_data in tqdm.tqdm(val_loader):\n",
        "    with torch.no_grad():\n",
        "        sampled_data.to(device)\n",
        "        preds.append(model(sampled_data))\n",
        "        ground_truths.append(sampled_data[\"user\", \"rates\", \"movie\"].edge_label)\n",
        "\n",
        "pred = torch.cat(preds, dim=0).cpu().numpy()\n",
        "ground_truth = torch.cat(ground_truths, dim=0).cpu().numpy()\n",
        "auc = roc_auc_score(ground_truth, pred)\n",
        "print()\n",
        "print(f\"Validation AUC: {auc:.4f}\")"
      ],
      "metadata": {
        "id": "Vi25Z7lFPPjc",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c842e8b0-d633-4bfa-b3b5-a28546328461"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 79/79 [00:03<00:00, 23.83it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Validation AUC: 0.9285\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    }
  ]
}