GNN/Link_Prediction_on_MovieLens.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": "08573a01-e064-4c84-8624-440ffb2a2d6f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "1.13.1+cu116\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AsWUrVh7vZtM",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2dbdb078-f500-48f2-9bf2-e5c0b0faa0a5"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://data.pyg.org/whl/torch-1.13.1+cu116.html\n",
            "Requirement already satisfied: torch-scatter in /usr/local/lib/python3.9/dist-packages (2.1.0+pt113cu116)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://data.pyg.org/whl/torch-1.13.1+cu116.html\n",
            "Requirement already satisfied: torch-sparse in /usr/local/lib/python3.9/dist-packages (0.6.16+pt113cu116)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (from torch-sparse) (1.10.1)\n",
            "Requirement already satisfied: numpy<1.27.0,>=1.19.5 in /usr/local/lib/python3.9/dist-packages (from scipy->torch-sparse) (1.22.4)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Looking in links: https://data.pyg.org/whl/torch-1.13.1+cu116.html\n",
            "Collecting pyg-lib\n",
            "  Downloading https://data.pyg.org/whl/torch-1.13.0%2Bcu116/pyg_lib-0.1.0%2Bpt113cu116-cp39-cp39-linux_x86_64.whl (1.9 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m60.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: pyg-lib\n",
            "Successfully installed pyg-lib-0.1.0+pt113cu116\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\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-85ih8pcj\n",
            "  Running command git clone --filter=blob:none --quiet https://github.com/pyg-team/pytorch_geometric.git /tmp/pip-req-build-85ih8pcj\n",
            "  Resolved https://github.com/pyg-team/pytorch_geometric.git to commit 688c7541d5e827ee159839026a35356ed3bc0bce\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: pyparsing in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (3.0.9)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (2.25.1)\n",
            "Requirement already satisfied: psutil>=5.8.0 in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (5.9.4)\n",
            "Requirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (1.10.1)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (4.65.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (1.22.4)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (3.1.2)\n",
            "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.9/dist-packages (from torch-geometric==2.3.0) (1.2.2)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from jinja2->torch-geometric==2.3.0) (2.1.2)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.9/dist-packages (from requests->torch-geometric==2.3.0) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.9/dist-packages (from requests->torch-geometric==2.3.0) (2022.12.7)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.9/dist-packages (from requests->torch-geometric==2.3.0) (4.0.0)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->torch-geometric==2.3.0) (1.26.14)\n",
            "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.9/dist-packages (from scikit-learn->torch-geometric==2.3.0) (1.1.1)\n",
            "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.9/dist-packages (from scikit-learn->torch-geometric==2.3.0) (3.1.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": "5bbe4a86-5d5b-4e42-d876-c7b58d8e030b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "Downloading https://files.grouplens.org/datasets/movielens/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)[[\"movieId\", \"genres\"]].head())\n",
        "print()\n",
        "print('ratings.csv:')\n",
        "print('============')\n",
        "print(pd.read_csv(ratings_path)[[\"userId\", \"movieId\"]].head())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D6ixkcCOwCDi",
        "outputId": "4bc902aa-4707-48ef-c5f0-65a8d0f0faf3"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "movies.csv:\n",
            "===========\n",
            "   movieId                                       genres\n",
            "0        1  Adventure|Animation|Children|Comedy|Fantasy\n",
            "1        2                   Adventure|Children|Fantasy\n",
            "2        3                               Comedy|Romance\n",
            "3        4                         Comedy|Drama|Romance\n",
            "4        5                                       Comedy\n",
            "\n",
            "ratings.csv:\n",
            "============\n",
            "   userId  movieId\n",
            "0       1        1\n",
            "1       1        3\n",
            "2       1        6\n",
            "3       1       47\n",
            "4       1       50\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')\n",
        "\n",
        "# Split genres and convert into indicator variables:\n",
        "genres = movies_df['genres'].str.get_dummies('|')\n",
        "print(genres[[\"Action\", \"Adventure\", \"Drama\", \"Horror\"]].head())\n",
        "\n",
        "# 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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dd_BJMNcwJ7k",
        "outputId": "d9c2d6b4-d968-4226-e2d6-b643683931f2"
      },
      "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": [
        "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",
        "\n",
        "# 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",
        "print(\"Mapping of user IDs to consecutive values:\")\n",
        "print(\"==========================================\")\n",
        "print(unique_user_id.head())\n",
        "print()\n",
        "# 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())\n",
        "\n",
        "# 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_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",
        "\n",
        "# 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": "35e859ca-75cb-407e-9a7c-b0d4b7b76fbb"
      },
      "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",
            "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",
            "\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 = ...  # TODO\n",
        "data[\"user\", \"rates\", \"movie\"].edge_index = ...  # TODO\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",
        "\n",
        "# TODO:\n",
        "raise NotImplementedError\n",
        "\n",
        "print(data)\n",
        "\n",
        "assert data.node_types == [\"user\", \"movie\"]\n",
        "assert data.edge_types == [(\"user\", \"rates\", \"movie\"),\n",
        "                           (\"movie\", \"rev_rates\", \"user\")]\n",
        "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": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I_63--974srt",
        "outputId": "32fe707d-f9ac-4dca-8c7b-52bbb8ddfa07"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "HeteroData(\n",
            "  \u001b[1muser\u001b[0m={ node_id=[610] },\n",
            "  \u001b[1mmovie\u001b[0m={\n",
            "    node_id=[9742],\n",
            "    x=[9742, 20]\n",
            "  },\n",
            "  \u001b[1m(user, rates, movie)\u001b[0m={ edge_index=[2, 100836] }\n",
            ")\n"
          ]
        }
      ]
    },
    {
      "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=...,  # TODO\n",
        "    num_test=...,  # TODO\n",
        "    disjoint_train_ratio=...,  # TODO\n",
        "    neg_sampling_ratio=...,  # TODO\n",
        "    add_negative_train_samples=...,  # TODO\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",
        "assert train_data[\"user\", \"rates\", \"movie\"].num_edges == 56469\n",
        "assert train_data[\"user\", \"rates\", \"movie\"].edge_label_index.size(1) == 24201\n",
        "assert train_data[\"movie\", \"rev_rates\", \"user\"].num_edges == 56469\n",
        "# No negative edges added:\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",
        "# Negative edges with ratio 2:1:\n",
        "assert val_data[\"user\", \"rates\", \"movie\"].edge_label.long().bincount().tolist() == [20166, 10083]"
      ],
      "metadata": {
        "id": "rwgNwoa26Eja"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "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=...,  # TODO\n",
        "    num_neighbors=...,  # TODO\n",
        "    neg_sampling_ratio=...,  # TODO\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": "e4b6d4d3-f59b-4f8a-9b78-a2536ed08516"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sampled mini-batch:\n",
            "===================\n",
            "HeteroData(\n",
            "  \u001b[1muser\u001b[0m={ node_id=[606] },\n",
            "  \u001b[1mmovie\u001b[0m={\n",
            "    node_id=[2769],\n",
            "    x=[2769, 20]\n",
            "  },\n",
            "  \u001b[1m(user, rates, movie)\u001b[0m={\n",
            "    edge_index=[2, 17686],\n",
            "    edge_label=[384],\n",
            "    edge_label_index=[2, 384],\n",
            "    input_id=[128]\n",
            "  },\n",
            "  \u001b[1m(movie, rev_rates, user)\u001b[0m={ edge_index=[2, 7808] }\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",
        "\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",
        "        # Define a 2-layer GNN computation graph.\n",
        "        # Use a *single* `ReLU` non-linearity in-between.\n",
        "        # TODO:\n",
        "        raise NotImplementedError\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": "5756285f-592d-4783-a31d-1bbccf8fd454"
      },
      "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",
        "        # TODO: Move `sampled_data` to the respective `device`\n",
        "        # TODO: Run `forward` pass of the model\n",
        "        # TODO: Apply binary cross entropy via\n",
        "        # `F.binary_cross_entropy_with_logits(pred, ground_truth)`\n",
        "        raise NotImplementedError\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": "65d515e9-00d6-48d9-92f2-2b41207e6d53"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Device: 'cpu'\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:11<00:00, 16.58it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 001, Loss: 0.4392\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:11<00:00, 16.29it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 002, Loss: 0.3498\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:11<00:00, 16.77it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 003, Loss: 0.3290\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:11<00:00, 16.76it/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 004, Loss: 0.3138\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 190/190 [00:13<00:00, 14.46it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Epoch: 005, Loss: 0.2988\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": "174894fc-fc38-4500-a022-f8232c5015b2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Sampled mini-batch:\n",
            "===================\n",
            "HeteroData(\n",
            "  \u001b[1muser\u001b[0m={ node_id=[610] },\n",
            "  \u001b[1mmovie\u001b[0m={\n",
            "    node_id=[2624],\n",
            "    x=[2624, 20]\n",
            "  },\n",
            "  \u001b[1m(user, rates, movie)\u001b[0m={\n",
            "    edge_index=[2, 18595],\n",
            "    edge_label=[384],\n",
            "    edge_label_index=[2, 384],\n",
            "    input_id=[384]\n",
            "  },\n",
            "  \u001b[1m(movie, rev_rates, user)\u001b[0m={ edge_index=[2, 7785] }\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",
        "        # TODO: Collect predictions and ground-truths and write them into\n",
        "        # `preds` and `ground_truths`.\n",
        "        raise NotImplementedError\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": "9cdf67ca-f335-4009-b2df-769697f1c99a"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 79/79 [00:02<00:00, 29.13it/s]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\n",
            "Validation AUC: 0.9299\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    }
  ]
}