tutorials/train_model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "zNziUzq9nTdU",
    "outputId": "edb40abf-b984-4fec-8033-1ed92fbdb128"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.3.0+cu121\n",
      "Looking in links: https://data.pyg.org/whl/torch-2.3.0+cu121.html\n",
      "Collecting pyg-lib\n",
      "  Downloading https://data.pyg.org/whl/torch-2.3.0%2Bcu121/pyg_lib-0.4.0%2Bpt23cu121-cp310-cp310-linux_x86_64.whl (2.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.5/2.5 MB\u001b[0m \u001b[31m22.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: pyg-lib\n",
      "Successfully installed pyg-lib-0.4.0+pt23cu121\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-0v5d62c1\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/pyg-team/pytorch_geometric.git /tmp/pip-req-build-0v5d62c1\n",
      "  Resolved https://github.com/pyg-team/pytorch_geometric.git to commit fbafbc4fc9181e8759ec1f39d9618992793b5fe1\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.6.0) (3.9.5)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (2023.6.0)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (3.1.4)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (1.25.2)\n",
      "Requirement already satisfied: psutil>=5.8.0 in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (5.9.5)\n",
      "Requirement already satisfied: pyparsing in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (3.1.2)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (2.31.0)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from torch-geometric==2.6.0) (4.66.4)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.6.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.6.0) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.6.0) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.6.0) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.6.0) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch-geometric==2.6.0) (4.0.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch-geometric==2.6.0) (2.1.5)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.6.0) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.6.0) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.6.0) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->torch-geometric==2.6.0) (2024.6.2)\n",
      "Building wheels for collected packages: torch-geometric\n",
      "  Building wheel for torch-geometric (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
      "  Created wheel for torch-geometric: filename=torch_geometric-2.6.0-py3-none-any.whl size=1122975 sha256=942596ab7c5d81703af08d193615963fbfde8806a6840d5f3ccef502768d7d18\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-9vz0psfl/wheels/d3/78/eb/9e26525b948d19533f1688fb6c209cec8a0ba793d39b49ae8f\n",
      "Successfully built torch-geometric\n",
      "Installing collected packages: torch-geometric\n",
      "Successfully installed torch-geometric-2.6.0\n",
      "Collecting pytorch_frame[full]\n",
      "  Downloading pytorch_frame-0.2.2-py3-none-any.whl (140 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m140.8/140.8 kB\u001b[0m \u001b[31m4.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (1.25.2)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (2.0.3)\n",
      "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (2.3.0+cu121)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (4.66.4)\n",
      "Requirement already satisfied: pyarrow in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (14.0.2)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (9.4.0)\n",
      "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (1.2.2)\n",
      "Collecting xgboost<2.0.0,>=1.7.0 (from pytorch_frame[full])\n",
      "  Downloading xgboost-1.7.6-py3-none-manylinux2014_x86_64.whl (200.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m200.3/200.3 MB\u001b[0m \u001b[31m3.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting optuna>=3.0.0 (from pytorch_frame[full])\n",
      "  Downloading optuna-3.6.1-py3-none-any.whl (380 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m380.1/380.1 kB\u001b[0m \u001b[31m50.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: mpmath==1.3.0 in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (1.3.0)\n",
      "Collecting catboost (from pytorch_frame[full])\n",
      "  Downloading catboost-1.2.5-cp310-cp310-manylinux2014_x86_64.whl (98.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.2/98.2 MB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: lightgbm in /usr/local/lib/python3.10/dist-packages (from pytorch_frame[full]) (4.1.0)\n",
      "Collecting datasets (from pytorch_frame[full])\n",
      "  Downloading datasets-2.20.0-py3-none-any.whl (547 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m547.8/547.8 kB\u001b[0m \u001b[31m45.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting alembic>=1.5.0 (from optuna>=3.0.0->pytorch_frame[full])\n",
      "  Downloading alembic-1.13.2-py3-none-any.whl (232 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m233.0/233.0 kB\u001b[0m \u001b[31m30.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting colorlog (from optuna>=3.0.0->pytorch_frame[full])\n",
      "  Downloading colorlog-6.8.2-py3-none-any.whl (11 kB)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from optuna>=3.0.0->pytorch_frame[full]) (24.1)\n",
      "Requirement already satisfied: sqlalchemy>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from optuna>=3.0.0->pytorch_frame[full]) (2.0.31)\n",
      "Requirement already satisfied: PyYAML in /usr/local/lib/python3.10/dist-packages (from optuna>=3.0.0->pytorch_frame[full]) (6.0.1)\n",
      "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from xgboost<2.0.0,>=1.7.0->pytorch_frame[full]) (1.11.4)\n",
      "Requirement already satisfied: graphviz in /usr/local/lib/python3.10/dist-packages (from catboost->pytorch_frame[full]) (0.20.3)\n",
      "Requirement already satisfied: matplotlib in /usr/local/lib/python3.10/dist-packages (from catboost->pytorch_frame[full]) (3.7.1)\n",
      "Requirement already satisfied: plotly in /usr/local/lib/python3.10/dist-packages (from catboost->pytorch_frame[full]) (5.15.0)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.10/dist-packages (from catboost->pytorch_frame[full]) (1.16.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->pytorch_frame[full]) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->pytorch_frame[full]) (2023.4)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas->pytorch_frame[full]) (2024.1)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets->pytorch_frame[full]) (3.15.4)\n",
      "Collecting pyarrow (from pytorch_frame[full])\n",
      "  Downloading pyarrow-16.1.0-cp310-cp310-manylinux_2_28_x86_64.whl (40.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m40.8/40.8 MB\u001b[0m \u001b[31m13.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pyarrow-hotfix in /usr/local/lib/python3.10/dist-packages (from datasets->pytorch_frame[full]) (0.6)\n",
      "Collecting dill<0.3.9,>=0.3.0 (from datasets->pytorch_frame[full])\n",
      "  Downloading dill-0.3.8-py3-none-any.whl (116 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m16.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting requests>=2.32.2 (from datasets->pytorch_frame[full])\n",
      "  Downloading requests-2.32.3-py3-none-any.whl (64 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.9/64.9 kB\u001b[0m \u001b[31m9.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting xxhash (from datasets->pytorch_frame[full])\n",
      "  Downloading xxhash-3.4.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (194 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m19.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting multiprocess (from datasets->pytorch_frame[full])\n",
      "  Downloading multiprocess-0.70.16-py310-none-any.whl (134 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m10.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: fsspec[http]<=2024.5.0,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from datasets->pytorch_frame[full]) (2023.6.0)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets->pytorch_frame[full]) (3.9.5)\n",
      "Requirement already satisfied: huggingface-hub>=0.21.2 in /usr/local/lib/python3.10/dist-packages (from datasets->pytorch_frame[full]) (0.23.4)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->pytorch_frame[full]) (1.4.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->pytorch_frame[full]) (3.5.0)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /usr/local/lib/python3.10/dist-packages (from torch->pytorch_frame[full]) (4.12.2)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->pytorch_frame[full]) (1.12.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->pytorch_frame[full]) (3.3)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->pytorch_frame[full]) (3.1.4)\n",
      "Collecting nvidia-cuda-nvrtc-cu12==12.1.105 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (23.7 MB)\n",
      "Collecting nvidia-cuda-runtime-cu12==12.1.105 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (823 kB)\n",
      "Collecting nvidia-cuda-cupti-cu12==12.1.105 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (14.1 MB)\n",
      "Collecting nvidia-cudnn-cu12==8.9.2.26 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl (731.7 MB)\n",
      "Collecting nvidia-cublas-cu12==12.1.3.1 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl (410.6 MB)\n",
      "Collecting nvidia-cufft-cu12==11.0.2.54 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl (121.6 MB)\n",
      "Collecting nvidia-curand-cu12==10.3.2.106 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl (56.5 MB)\n",
      "Collecting nvidia-cusolver-cu12==11.4.5.107 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl (124.2 MB)\n",
      "Collecting nvidia-cusparse-cu12==12.1.0.106 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl (196.0 MB)\n",
      "Collecting nvidia-nccl-cu12==2.20.5 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_nccl_cu12-2.20.5-py3-none-manylinux2014_x86_64.whl (176.2 MB)\n",
      "Collecting nvidia-nvtx-cu12==12.1.105 (from torch->pytorch_frame[full])\n",
      "  Using cached nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl (99 kB)\n",
      "Requirement already satisfied: triton==2.3.0 in /usr/local/lib/python3.10/dist-packages (from torch->pytorch_frame[full]) (2.3.0)\n",
      "Collecting nvidia-nvjitlink-cu12 (from nvidia-cusolver-cu12==11.4.5.107->torch->pytorch_frame[full])\n",
      "  Downloading nvidia_nvjitlink_cu12-12.5.82-py3-none-manylinux2014_x86_64.whl (21.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.3/21.3 MB\u001b[0m \u001b[31m71.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting Mako (from alembic>=1.5.0->optuna>=3.0.0->pytorch_frame[full])\n",
      "  Downloading Mako-1.3.5-py3-none-any.whl (78 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.6/78.6 kB\u001b[0m \u001b[31m14.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets->pytorch_frame[full]) (4.0.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets->pytorch_frame[full]) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets->pytorch_frame[full]) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets->pytorch_frame[full]) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets->pytorch_frame[full]) (2024.6.2)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy>=1.3.0->optuna>=3.0.0->pytorch_frame[full]) (3.0.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->pytorch_frame[full]) (2.1.5)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->catboost->pytorch_frame[full]) (1.2.1)\n",
      "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.10/dist-packages (from matplotlib->catboost->pytorch_frame[full]) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /usr/local/lib/python3.10/dist-packages (from matplotlib->catboost->pytorch_frame[full]) (4.53.0)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->catboost->pytorch_frame[full]) (1.4.5)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /usr/local/lib/python3.10/dist-packages (from matplotlib->catboost->pytorch_frame[full]) (3.1.2)\n",
      "Requirement already satisfied: tenacity>=6.2.0 in /usr/local/lib/python3.10/dist-packages (from plotly->catboost->pytorch_frame[full]) (8.4.2)\n",
      "Installing collected packages: xxhash, requests, pyarrow, nvidia-nvtx-cu12, nvidia-nvjitlink-cu12, nvidia-nccl-cu12, nvidia-curand-cu12, nvidia-cufft-cu12, nvidia-cuda-runtime-cu12, nvidia-cuda-nvrtc-cu12, nvidia-cuda-cupti-cu12, nvidia-cublas-cu12, Mako, dill, colorlog, xgboost, nvidia-cusparse-cu12, nvidia-cudnn-cu12, multiprocess, alembic, optuna, nvidia-cusolver-cu12, catboost, datasets, pytorch_frame\n",
      "  Attempting uninstall: requests\n",
      "    Found existing installation: requests 2.31.0\n",
      "    Uninstalling requests-2.31.0:\n",
      "      Successfully uninstalled requests-2.31.0\n",
      "  Attempting uninstall: pyarrow\n",
      "    Found existing installation: pyarrow 14.0.2\n",
      "    Uninstalling pyarrow-14.0.2:\n",
      "      Successfully uninstalled pyarrow-14.0.2\n",
      "  Attempting uninstall: xgboost\n",
      "    Found existing installation: xgboost 2.0.3\n",
      "    Uninstalling xgboost-2.0.3:\n",
      "      Successfully uninstalled xgboost-2.0.3\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "cudf-cu12 24.4.1 requires pyarrow<15.0.0a0,>=14.0.1, but you have pyarrow 16.1.0 which is incompatible.\n",
      "google-colab 1.0.0 requires requests==2.31.0, but you have requests 2.32.3 which is incompatible.\n",
      "ibis-framework 8.0.0 requires pyarrow<16,>=2, but you have pyarrow 16.1.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed Mako-1.3.5 alembic-1.13.2 catboost-1.2.5 colorlog-6.8.2 datasets-2.20.0 dill-0.3.8 multiprocess-0.70.16 nvidia-cublas-cu12-12.1.3.1 nvidia-cuda-cupti-cu12-12.1.105 nvidia-cuda-nvrtc-cu12-12.1.105 nvidia-cuda-runtime-cu12-12.1.105 nvidia-cudnn-cu12-8.9.2.26 nvidia-cufft-cu12-11.0.2.54 nvidia-curand-cu12-10.3.2.106 nvidia-cusolver-cu12-11.4.5.107 nvidia-cusparse-cu12-12.1.0.106 nvidia-nccl-cu12-2.20.5 nvidia-nvjitlink-cu12-12.5.82 nvidia-nvtx-cu12-12.1.105 optuna-3.6.1 pyarrow-16.1.0 pytorch_frame-0.2.2 requests-2.32.3 xgboost-1.7.6 xxhash-3.4.1\n",
      "Collecting relbench[full]==1.0.0rc1\n",
      "  Downloading relbench-1.0.0rc1-py3-none-any.whl (57 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.4/57.4 kB\u001b[0m \u001b[31m3.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (2.0.3)\n",
      "Requirement already satisfied: pooch in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (1.8.2)\n",
      "Requirement already satisfied: pyarrow in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (16.1.0)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (1.25.2)\n",
      "Requirement already satisfied: duckdb in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (0.10.3)\n",
      "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (1.2.2)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (4.12.2)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (4.66.4)\n",
      "Requirement already satisfied: torch in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (2.3.0+cu121)\n",
      "Requirement already satisfied: pytorch_frame>=0.2.2 in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (0.2.2)\n",
      "Requirement already satisfied: torch_geometric in /usr/local/lib/python3.10/dist-packages (from relbench[full]==1.0.0rc1) (2.6.0)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from pytorch_frame>=0.2.2->relbench[full]==1.0.0rc1) (9.4.0)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->relbench[full]==1.0.0rc1) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->relbench[full]==1.0.0rc1) (2023.4)\n",
      "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas->relbench[full]==1.0.0rc1) (2024.1)\n",
      "Requirement already satisfied: platformdirs>=2.5.0 in /usr/local/lib/python3.10/dist-packages (from pooch->relbench[full]==1.0.0rc1) (4.2.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from pooch->relbench[full]==1.0.0rc1) (24.1)\n",
      "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.10/dist-packages (from pooch->relbench[full]==1.0.0rc1) (2.32.3)\n",
      "Requirement already satisfied: scipy>=1.3.2 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->relbench[full]==1.0.0rc1) (1.11.4)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->relbench[full]==1.0.0rc1) (1.4.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->relbench[full]==1.0.0rc1) (3.5.0)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (3.15.4)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (1.12.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (3.3)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (3.1.4)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (2023.6.0)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.20.5 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (2.20.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (12.1.105)\n",
      "Requirement already satisfied: triton==2.3.0 in /usr/local/lib/python3.10/dist-packages (from torch->relbench[full]==1.0.0rc1) (2.3.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch->relbench[full]==1.0.0rc1) (12.5.82)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from torch_geometric->relbench[full]==1.0.0rc1) (3.9.5)\n",
      "Requirement already satisfied: psutil>=5.8.0 in /usr/local/lib/python3.10/dist-packages (from torch_geometric->relbench[full]==1.0.0rc1) (5.9.5)\n",
      "Requirement already satisfied: pyparsing in /usr/local/lib/python3.10/dist-packages (from torch_geometric->relbench[full]==1.0.0rc1) (3.1.2)\n",
      "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->relbench[full]==1.0.0rc1) (1.16.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->pooch->relbench[full]==1.0.0rc1) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->pooch->relbench[full]==1.0.0rc1) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->pooch->relbench[full]==1.0.0rc1) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.19.0->pooch->relbench[full]==1.0.0rc1) (2024.6.2)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->torch_geometric->relbench[full]==1.0.0rc1) (4.0.3)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch->relbench[full]==1.0.0rc1) (2.1.5)\n",
      "Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch->relbench[full]==1.0.0rc1) (1.3.0)\n",
      "Installing collected packages: relbench\n",
      "Successfully installed relbench-1.0.0rc1\n"
     ]
    }
   ],
   "source": [
    "# Install required packages.\n",
    "import os\n",
    "import torch\n",
    "\n",
    "!pip install pyg-lib -f https://data.pyg.org/whl/torch-2.3.0+cpu.html # PyG for working with graphs\n",
    "!pip install git+https://github.com/pyg-team/pytorch_geometric.git # more PyG\n",
    "!pip install pytorch_frame[full] #PyTorch Frame for working with tabular data\n",
    "!pip install relbench[full]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-07-26T00:05:04.397024Z",
     "iopub.status.busy": "2024-07-26T00:05:04.396612Z",
     "iopub.status.idle": "2024-07-26T00:05:05.049064Z",
     "shell.execute_reply": "2024-07-26T00:05:05.048407Z",
     "shell.execute_reply.started": "2024-07-26T00:05:04.397003Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.0'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import relbench\n",
    "\n",
    "relbench.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6DWB-Kf6nl2y"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "from torch.nn import BCEWithLogitsLoss, L1Loss\n",
    "from relbench.datasets import get_dataset\n",
    "from relbench.tasks import get_task\n",
    "\n",
    "dataset = get_dataset(\"rel-f1\", download=True)\n",
    "task = get_task(\"rel-f1\", \"driver-position\", download=True)\n",
    "\n",
    "train_table = task.get_table(\"train\")\n",
    "val_table = task.get_table(\"val\")\n",
    "test_table = task.get_table(\"test\")\n",
    "\n",
    "out_channels = 1\n",
    "loss_fn = L1Loss()\n",
    "tune_metric = \"mae\"\n",
    "higher_is_better = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UKFT5H51j_Um"
   },
   "source": [
    "Let's check out the training table just to make sure it looks fine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ABN_fdN3kAB9",
    "outputId": "03d4a31a-124d-45c7-b4dc-9713e5e4b942"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Table(df=\n",
       "           date  driverId  position\n",
       "0    2004-07-05        10     10.75\n",
       "1    2004-07-05        47     12.00\n",
       "2    2004-03-07         7     15.00\n",
       "3    2004-01-07        10      9.00\n",
       "4    2003-09-09        52     13.00\n",
       "...         ...       ...       ...\n",
       "7448 1995-08-22        96     15.75\n",
       "7449 1975-06-08       228      8.00\n",
       "7450 1965-05-31       418     16.00\n",
       "7451 1961-08-20       467     37.00\n",
       "7452 1954-05-29       677     30.00\n",
       "\n",
       "[7453 rows x 3 columns],\n",
       "  fkey_col_to_pkey_table={'driverId': 'drivers'},\n",
       "  pkey_col=None,\n",
       "  time_col=date)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_table"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qQhuHIdHkOxv"
   },
   "source": [
    "Note that to load the data we did not require any deep learning libraries. Now we introduce the PyTorch Frame library, which is useful for encoding individual tables into initial node features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qNzfdwsrkPIo",
    "outputId": "c985185e-b785-405e-bd46-ebf2f48e3ac6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "import torch_geometric\n",
    "import torch_frame\n",
    "\n",
    "# Some book keeping\n",
    "from torch_geometric.seed import seed_everything\n",
    "\n",
    "seed_everything(42)\n",
    "\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)  # check that it's cuda if you want it to run in reasonable time!\n",
    "root_dir = \"./data\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0Y79g5H0kVjX"
   },
   "source": [
    "The first big move is to build a graph out of the database. Here we use our pre-prepared conversion function.\n",
    "\n",
    "The source code can be found at: https://github.com/snap-stanford/relbench/blob/main/relbench/modeling/graph.py\n",
    "\n",
    "Each node in the graph corresonds to a single row in the database. Crucially, PyTorch Frame stores whole tables as objects in a way that is compatibile with PyG minibatch sampling, meaning we can sample subgraphs as in https://arxiv.org/abs/1706.02216, and retrieve the relevant raw features.\n",
    "\n",
    "PyTorch Frame also stores the `stype` (i.e., modality) of each column, and any specialized feature encoders (e.g., text encoders) to be used later. So we need to configure the `stype` for each column, for which we use a function that tries to automatically detect the `stype`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kiV3TGI-kRuy",
    "outputId": "98e88ec3-ab38-4a14-8dd8-24f3ea349893"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading Database object from /root/.cache/relbench/rel-f1/db...\n",
      "Done in 0.04 seconds.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'drivers': {'driverId': <stype.numerical: 'numerical'>,\n",
       "  'driverRef': <stype.text_embedded: 'text_embedded'>,\n",
       "  'code': <stype.text_embedded: 'text_embedded'>,\n",
       "  'forename': <stype.text_embedded: 'text_embedded'>,\n",
       "  'surname': <stype.text_embedded: 'text_embedded'>,\n",
       "  'dob': <stype.timestamp: 'timestamp'>,\n",
       "  'nationality': <stype.text_embedded: 'text_embedded'>},\n",
       " 'races': {'raceId': <stype.numerical: 'numerical'>,\n",
       "  'year': <stype.categorical: 'categorical'>,\n",
       "  'round': <stype.numerical: 'numerical'>,\n",
       "  'circuitId': <stype.numerical: 'numerical'>,\n",
       "  'name': <stype.text_embedded: 'text_embedded'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>,\n",
       "  'time': <stype.timestamp: 'timestamp'>},\n",
       " 'constructor_standings': {'constructorStandingsId': <stype.numerical: 'numerical'>,\n",
       "  'raceId': <stype.numerical: 'numerical'>,\n",
       "  'constructorId': <stype.numerical: 'numerical'>,\n",
       "  'points': <stype.numerical: 'numerical'>,\n",
       "  'position': <stype.numerical: 'numerical'>,\n",
       "  'wins': <stype.numerical: 'numerical'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>},\n",
       " 'constructor_results': {'constructorResultsId': <stype.numerical: 'numerical'>,\n",
       "  'raceId': <stype.numerical: 'numerical'>,\n",
       "  'constructorId': <stype.numerical: 'numerical'>,\n",
       "  'points': <stype.numerical: 'numerical'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>},\n",
       " 'results': {'resultId': <stype.numerical: 'numerical'>,\n",
       "  'raceId': <stype.numerical: 'numerical'>,\n",
       "  'driverId': <stype.numerical: 'numerical'>,\n",
       "  'constructorId': <stype.numerical: 'numerical'>,\n",
       "  'number': <stype.numerical: 'numerical'>,\n",
       "  'grid': <stype.numerical: 'numerical'>,\n",
       "  'position': <stype.numerical: 'numerical'>,\n",
       "  'positionOrder': <stype.numerical: 'numerical'>,\n",
       "  'points': <stype.numerical: 'numerical'>,\n",
       "  'laps': <stype.numerical: 'numerical'>,\n",
       "  'milliseconds': <stype.numerical: 'numerical'>,\n",
       "  'fastestLap': <stype.numerical: 'numerical'>,\n",
       "  'rank': <stype.numerical: 'numerical'>,\n",
       "  'statusId': <stype.numerical: 'numerical'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>},\n",
       " 'qualifying': {'qualifyId': <stype.numerical: 'numerical'>,\n",
       "  'raceId': <stype.numerical: 'numerical'>,\n",
       "  'driverId': <stype.numerical: 'numerical'>,\n",
       "  'constructorId': <stype.numerical: 'numerical'>,\n",
       "  'number': <stype.numerical: 'numerical'>,\n",
       "  'position': <stype.numerical: 'numerical'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>},\n",
       " 'circuits': {'circuitId': <stype.numerical: 'numerical'>,\n",
       "  'circuitRef': <stype.text_embedded: 'text_embedded'>,\n",
       "  'name': <stype.text_embedded: 'text_embedded'>,\n",
       "  'location': <stype.text_embedded: 'text_embedded'>,\n",
       "  'country': <stype.text_embedded: 'text_embedded'>,\n",
       "  'lat': <stype.numerical: 'numerical'>,\n",
       "  'lng': <stype.numerical: 'numerical'>,\n",
       "  'alt': <stype.numerical: 'numerical'>},\n",
       " 'constructors': {'constructorId': <stype.numerical: 'numerical'>,\n",
       "  'constructorRef': <stype.text_embedded: 'text_embedded'>,\n",
       "  'name': <stype.text_embedded: 'text_embedded'>,\n",
       "  'nationality': <stype.text_embedded: 'text_embedded'>},\n",
       " 'standings': {'driverStandingsId': <stype.numerical: 'numerical'>,\n",
       "  'raceId': <stype.numerical: 'numerical'>,\n",
       "  'driverId': <stype.numerical: 'numerical'>,\n",
       "  'points': <stype.numerical: 'numerical'>,\n",
       "  'position': <stype.numerical: 'numerical'>,\n",
       "  'wins': <stype.numerical: 'numerical'>,\n",
       "  'date': <stype.timestamp: 'timestamp'>}}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from relbench.modeling.utils import get_stype_proposal\n",
    "\n",
    "db = dataset.get_db()\n",
    "col_to_stype_dict = get_stype_proposal(db)\n",
    "col_to_stype_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Sm3uYXqXkbZt"
   },
   "source": [
    "If trying a new dataset, you should definitely check through this dict of `stype`s to check that look right, and manually change any mistakes by the auto-detection function.\n",
    "\n",
    "Next we also define our text encoding model, which we use GloVe embeddings for speed and convenience. Feel free to try alternatives here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "QQHYmgIxkX1j",
    "outputId": "857b70dd-e7eb-4b09-a5cd-394fccef758a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting sentence-transformers\n",
      "  Downloading sentence_transformers-3.0.1-py3-none-any.whl (227 kB)\n",
      "\u001b[?25l     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m0.0/227.1 kB\u001b[0m \u001b[31m?\u001b[0m eta \u001b[36m-:--:--\u001b[0m\r",
      "\u001b[2K     \u001b[91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m\u001b[91m╸\u001b[0m \u001b[32m225.3/227.1 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\r",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m227.1/227.1 kB\u001b[0m \u001b[31m5.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: transformers<5.0.0,>=4.34.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (4.41.2)\n",
      "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (4.66.4)\n",
      "Requirement already satisfied: torch>=1.11.0 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (2.3.0+cu121)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (1.25.2)\n",
      "Requirement already satisfied: scikit-learn in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (1.2.2)\n",
      "Requirement already satisfied: scipy in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (1.11.4)\n",
      "Requirement already satisfied: huggingface-hub>=0.15.1 in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (0.23.4)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.10/dist-packages (from sentence-transformers) (9.4.0)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (3.15.4)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (2023.6.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (24.1)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (6.0.1)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (2.32.3)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.15.1->sentence-transformers) (4.12.2)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (1.12.1)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (3.3)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (3.1.4)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.105)\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==8.9.2.26 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (8.9.2.26)\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.1.3.1 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.3.1)\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.0.2.54 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (11.0.2.54)\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.2.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (10.3.2.106)\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.4.5.107 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (11.4.5.107)\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.1.0.106 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.0.106)\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.20.5 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (2.20.5)\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.1.105 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (12.1.105)\n",
      "Requirement already satisfied: triton==2.3.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.11.0->sentence-transformers) (2.3.0)\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12 in /usr/local/lib/python3.10/dist-packages (from nvidia-cusolver-cu12==11.4.5.107->torch>=1.11.0->sentence-transformers) (12.5.82)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers<5.0.0,>=4.34.0->sentence-transformers) (2024.5.15)\n",
      "Requirement already satisfied: tokenizers<0.20,>=0.19 in /usr/local/lib/python3.10/dist-packages (from transformers<5.0.0,>=4.34.0->sentence-transformers) (0.19.1)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /usr/local/lib/python3.10/dist-packages (from transformers<5.0.0,>=4.34.0->sentence-transformers) (0.4.3)\n",
      "Requirement already satisfied: joblib>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->sentence-transformers) (1.4.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /usr/local/lib/python3.10/dist-packages (from scikit-learn->sentence-transformers) (3.5.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.11.0->sentence-transformers) (2.1.5)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.15.1->sentence-transformers) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.15.1->sentence-transformers) (3.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.15.1->sentence-transformers) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->huggingface-hub>=0.15.1->sentence-transformers) (2024.6.2)\n",
      "Requirement already satisfied: mpmath<1.4.0,>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.11.0->sentence-transformers) (1.3.0)\n",
      "Installing collected packages: sentence-transformers\n",
      "Successfully installed sentence-transformers-3.0.1\n"
     ]
    }
   ],
   "source": [
    "!pip install -U sentence-transformers # we need another package for text encoding\n",
    "from typing import List, Optional\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from torch import Tensor\n",
    "\n",
    "\n",
    "class GloveTextEmbedding:\n",
    "    def __init__(self, device: Optional[torch.device\n",
    "                                       ] = None):\n",
    "        self.model = SentenceTransformer(\n",
    "            \"sentence-transformers/average_word_embeddings_glove.6B.300d\",\n",
    "            device=device,\n",
    "        )\n",
    "\n",
    "    def __call__(self, sentences: List[str]) -> Tensor:\n",
    "        return torch.from_numpy(self.model.encode(sentences))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "L-BBpUrakdwY",
    "outputId": "b152bf13-f47d-4728-d58b-fc65f738b03d"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00,  6.11it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00, 177.19it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00, 185.70it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00, 184.06it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00, 157.68it/s]\n",
      "/usr/local/lib/python3.10/dist-packages/torch_frame/data/stats.py:177: UserWarning: Could not infer format, so each element will be parsed individually, falling back to `dateutil`. To ensure parsing is consistent and as-expected, please specify a format.\n",
      "  ser = pd.to_datetime(ser, format=time_format)\n",
      "Embedding raw data in mini-batch: 100%|██████████| 4/4 [00:00<00:00, 144.99it/s]\n",
      "/usr/local/lib/python3.10/dist-packages/torch_frame/data/mapper.py:290: UserWarning: Could not infer format, so each element will be parsed individually, falling back to `dateutil`. To ensure parsing is consistent and as-expected, please specify a format.\n",
      "  ser = pd.to_datetime(ser, format=self.format, errors='coerce')\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 151.76it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 172.90it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 207.29it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 148.61it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 42.85it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 95.30it/s]\n",
      "Embedding raw data in mini-batch: 100%|██████████| 1/1 [00:00<00:00, 93.10it/s]\n"
     ]
    }
   ],
   "source": [
    "from torch_frame.config.text_embedder import TextEmbedderConfig\n",
    "from relbench.modeling.graph import make_pkey_fkey_graph\n",
    "\n",
    "text_embedder_cfg = TextEmbedderConfig(\n",
    "    text_embedder=GloveTextEmbedding(device=device), batch_size=256\n",
    ")\n",
    "\n",
    "data, col_stats_dict = make_pkey_fkey_graph(\n",
    "    db,\n",
    "    col_to_stype_dict=col_to_stype_dict,  # speficied column types\n",
    "    text_embedder_cfg=text_embedder_cfg,  # our chosen text encoder\n",
    "    cache_dir=os.path.join(\n",
    "        root_dir, f\"rel-f1_materialized_cache\"\n",
    "    ),  # store materialized graph for convenience\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mwQejmg0kzOg"
   },
   "source": [
    "We can now check out `data`, our main graph object. `data` is a heterogeneous and temporal graph, with node types given by the table it originates from."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Gt4a8lw1kufy",
    "outputId": "4117959f-6f0d-4c31-9489-49db7d5f3c5d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  drivers={ tf=TensorFrame([857, 6]) },\n",
       "  races={\n",
       "    tf=TensorFrame([820, 5]),\n",
       "    time=[820],\n",
       "  },\n",
       "  constructor_standings={\n",
       "    tf=TensorFrame([10170, 4]),\n",
       "    time=[10170],\n",
       "  },\n",
       "  constructor_results={\n",
       "    tf=TensorFrame([9408, 2]),\n",
       "    time=[9408],\n",
       "  },\n",
       "  results={\n",
       "    tf=TensorFrame([20323, 11]),\n",
       "    time=[20323],\n",
       "  },\n",
       "  qualifying={\n",
       "    tf=TensorFrame([4082, 3]),\n",
       "    time=[4082],\n",
       "  },\n",
       "  circuits={ tf=TensorFrame([77, 7]) },\n",
       "  constructors={ tf=TensorFrame([211, 3]) },\n",
       "  standings={\n",
       "    tf=TensorFrame([28115, 4]),\n",
       "    time=[28115],\n",
       "  },\n",
       "  (races, f2p_circuitId, circuits)={ edge_index=[2, 820] },\n",
       "  (circuits, rev_f2p_circuitId, races)={ edge_index=[2, 820] },\n",
       "  (constructor_standings, f2p_raceId, races)={ edge_index=[2, 10170] },\n",
       "  (races, rev_f2p_raceId, constructor_standings)={ edge_index=[2, 10170] },\n",
       "  (constructor_standings, f2p_constructorId, constructors)={ edge_index=[2, 10170] },\n",
       "  (constructors, rev_f2p_constructorId, constructor_standings)={ edge_index=[2, 10170] },\n",
       "  (constructor_results, f2p_raceId, races)={ edge_index=[2, 9408] },\n",
       "  (races, rev_f2p_raceId, constructor_results)={ edge_index=[2, 9408] },\n",
       "  (constructor_results, f2p_constructorId, constructors)={ edge_index=[2, 9408] },\n",
       "  (constructors, rev_f2p_constructorId, constructor_results)={ edge_index=[2, 9408] },\n",
       "  (results, f2p_raceId, races)={ edge_index=[2, 20323] },\n",
       "  (races, rev_f2p_raceId, results)={ edge_index=[2, 20323] },\n",
       "  (results, f2p_driverId, drivers)={ edge_index=[2, 20323] },\n",
       "  (drivers, rev_f2p_driverId, results)={ edge_index=[2, 20323] },\n",
       "  (results, f2p_constructorId, constructors)={ edge_index=[2, 20323] },\n",
       "  (constructors, rev_f2p_constructorId, results)={ edge_index=[2, 20323] },\n",
       "  (qualifying, f2p_raceId, races)={ edge_index=[2, 4082] },\n",
       "  (races, rev_f2p_raceId, qualifying)={ edge_index=[2, 4082] },\n",
       "  (qualifying, f2p_driverId, drivers)={ edge_index=[2, 4082] },\n",
       "  (drivers, rev_f2p_driverId, qualifying)={ edge_index=[2, 4082] },\n",
       "  (qualifying, f2p_constructorId, constructors)={ edge_index=[2, 4082] },\n",
       "  (constructors, rev_f2p_constructorId, qualifying)={ edge_index=[2, 4082] },\n",
       "  (standings, f2p_raceId, races)={ edge_index=[2, 28115] },\n",
       "  (races, rev_f2p_raceId, standings)={ edge_index=[2, 28115] },\n",
       "  (standings, f2p_driverId, drivers)={ edge_index=[2, 28115] },\n",
       "  (drivers, rev_f2p_driverId, standings)={ edge_index=[2, 28115] }\n",
       ")"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yd6DqCXgk41x"
   },
   "source": [
    "We can also check out the TensorFrame for one table like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "-mMQTQeLk1rl",
    "outputId": "04d698af-a4f4-4a98-8321-b5f8dc6ee10d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorFrame(\n",
       "  num_cols=5,\n",
       "  num_rows=820,\n",
       "  categorical (1): ['year'],\n",
       "  numerical (1): ['round'],\n",
       "  timestamp (2): ['date', 'time'],\n",
       "  embedding (1): ['name'],\n",
       "  has_target=False,\n",
       "  device='cpu',\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[\"races\"].tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1kbysKXMk-3X"
   },
   "source": [
    "This may be a little confusing at first, as in graph ML it is more standard to associate to the graph object `data` a tensor, e.g., `data.x` for which `data.x[idx]` is a 1D array/tensor storing all the features for node with index `idx`.\n",
    "\n",
    "But actually this `data` object behaves similarly. For a given node type, e.g., `races` again, `data['races']` stores two pieces of information\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "cDIcp7L5k6pU",
    "outputId": "be742ecb-02db-43e6-9c12-53fb00e51fec"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['tf', 'time']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(data[\"races\"].keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Z18qPRPllB1H"
   },
   "source": [
    "A `TensorFrame` object, and a timestamp for each node. The `TensorFrame` object acts analogously to the usual tensor of node features, and you can simply use indexing to retrieve the features of a single row (node), or group of nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Im8bhNh5lFG6",
    "outputId": "7ae1e0bd-746a-4164-fc95-124e302c816d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorFrame(\n",
       "  num_cols=5,\n",
       "  num_rows=1,\n",
       "  categorical (1): ['year'],\n",
       "  numerical (1): ['round'],\n",
       "  timestamp (2): ['date', 'time'],\n",
       "  embedding (1): ['name'],\n",
       "  has_target=False,\n",
       "  device='cpu',\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[\"races\"].tf[10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eYZ28pzNlG4s",
    "outputId": "1066a167-2e3b-4ad6-d929-02acf18cad0e"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TensorFrame(\n",
       "  num_cols=5,\n",
       "  num_rows=10,\n",
       "  categorical (1): ['year'],\n",
       "  numerical (1): ['round'],\n",
       "  timestamp (2): ['date', 'time'],\n",
       "  embedding (1): ['name'],\n",
       "  has_target=False,\n",
       "  device='cpu',\n",
       ")"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[\"races\"].tf[10:20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ql15svcelK3A"
   },
   "source": [
    "We can also check the edge indices between two different node types, such as `races` amd `circuits`. Note that the edges are also heterogenous, so we also need to specify which edge type we want to look at. Here we look at `f2p_curcuitId`, which are the directed edges pointing _from_ a race (the `f` stands for `foreign key`), _to_ the circuit at which te race happened (the `p` stands for `primary key`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TynkD36QlInL",
    "outputId": "abc2f80d-5ff4-42b1-f9e3-bd9f004d84ee"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'edge_index': tensor([[  0,   1,   2,  ..., 817, 818, 819],\n",
       "        [  8,   5,  18,  ...,  21,  17,  23]])}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[(\"races\", \"f2p_circuitId\", \"circuits\")]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xx4V5KCelNxl"
   },
   "source": [
    "Now we are ready to instantiate our data loaders. For this we will need to import PyTorch Geometric, our GNN library. Whilst we're at it let's add a seed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "HUHVG-g6lM-b"
   },
   "outputs": [],
   "source": [
    "from relbench.modeling.graph import get_node_train_table_input, make_pkey_fkey_graph\n",
    "from torch_geometric.loader import NeighborLoader\n",
    "\n",
    "loader_dict = {}\n",
    "\n",
    "for split, table in [\n",
    "    (\"train\", train_table),\n",
    "    (\"val\", val_table),\n",
    "    (\"test\", test_table),\n",
    "]:\n",
    "    table_input = get_node_train_table_input(\n",
    "        table=table,\n",
    "        task=task,\n",
    "    )\n",
    "    entity_table = table_input.nodes[0]\n",
    "    loader_dict[split] = NeighborLoader(\n",
    "        data,\n",
    "        num_neighbors=[\n",
    "            128 for i in range(2)\n",
    "        ],  # we sample subgraphs of depth 2, 128 neighbors per node.\n",
    "        time_attr=\"time\",\n",
    "        input_nodes=table_input.nodes,\n",
    "        input_time=table_input.time,\n",
    "        transform=table_input.transform,\n",
    "        batch_size=512,\n",
    "        temporal_strategy=\"uniform\",\n",
    "        shuffle=split == \"train\",\n",
    "        num_workers=0,\n",
    "        persistent_workers=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BQc8BWsGludR"
   },
   "source": [
    "Now we need our model...\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u3m3jEqClQnw"
   },
   "outputs": [],
   "source": [
    "from torch.nn import BCEWithLogitsLoss\n",
    "import copy\n",
    "from typing import Any, Dict, List\n",
    "\n",
    "import torch\n",
    "from torch import Tensor\n",
    "from torch.nn import Embedding, ModuleDict\n",
    "from torch_frame.data.stats import StatType\n",
    "from torch_geometric.data import HeteroData\n",
    "from torch_geometric.nn import MLP\n",
    "from torch_geometric.typing import NodeType\n",
    "\n",
    "from relbench.modeling.nn import HeteroEncoder, HeteroGraphSAGE, HeteroTemporalEncoder\n",
    "\n",
    "\n",
    "class Model(torch.nn.Module):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        data: HeteroData,\n",
    "        col_stats_dict: Dict[str, Dict[str, Dict[StatType, Any]]],\n",
    "        num_layers: int,\n",
    "        channels: int,\n",
    "        out_channels: int,\n",
    "        aggr: str,\n",
    "        norm: str,\n",
    "        # List of node types to add shallow embeddings to input\n",
    "        shallow_list: List[NodeType] = [],\n",
    "        # ID awareness\n",
    "        id_awareness: bool = False,\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = HeteroEncoder(\n",
    "            channels=channels,\n",
    "            node_to_col_names_dict={\n",
    "                node_type: data[node_type].tf.col_names_dict\n",
    "                for node_type in data.node_types\n",
    "            },\n",
    "            node_to_col_stats=col_stats_dict,\n",
    "        )\n",
    "        self.temporal_encoder = HeteroTemporalEncoder(\n",
    "            node_types=[\n",
    "                node_type for node_type in data.node_types if \"time\" in data[node_type]\n",
    "            ],\n",
    "            channels=channels,\n",
    "        )\n",
    "        self.gnn = HeteroGraphSAGE(\n",
    "            node_types=data.node_types,\n",
    "            edge_types=data.edge_types,\n",
    "            channels=channels,\n",
    "            aggr=aggr,\n",
    "            num_layers=num_layers,\n",
    "        )\n",
    "        self.head = MLP(\n",
    "            channels,\n",
    "            out_channels=out_channels,\n",
    "            norm=norm,\n",
    "            num_layers=1,\n",
    "        )\n",
    "        self.embedding_dict = ModuleDict(\n",
    "            {\n",
    "                node: Embedding(data.num_nodes_dict[node], channels)\n",
    "                for node in shallow_list\n",
    "            }\n",
    "        )\n",
    "\n",
    "        self.id_awareness_emb = None\n",
    "        if id_awareness:\n",
    "            self.id_awareness_emb = torch.nn.Embedding(1, channels)\n",
    "        self.reset_parameters()\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.encoder.reset_parameters()\n",
    "        self.temporal_encoder.reset_parameters()\n",
    "        self.gnn.reset_parameters()\n",
    "        self.head.reset_parameters()\n",
    "        for embedding in self.embedding_dict.values():\n",
    "            torch.nn.init.normal_(embedding.weight, std=0.1)\n",
    "        if self.id_awareness_emb is not None:\n",
    "            self.id_awareness_emb.reset_parameters()\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        batch: HeteroData,\n",
    "        entity_table: NodeType,\n",
    "    ) -> Tensor:\n",
    "        seed_time = batch[entity_table].seed_time\n",
    "        x_dict = self.encoder(batch.tf_dict)\n",
    "\n",
    "        rel_time_dict = self.temporal_encoder(\n",
    "            seed_time, batch.time_dict, batch.batch_dict\n",
    "        )\n",
    "\n",
    "        for node_type, rel_time in rel_time_dict.items():\n",
    "            x_dict[node_type] = x_dict[node_type] + rel_time\n",
    "\n",
    "        for node_type, embedding in self.embedding_dict.items():\n",
    "            x_dict[node_type] = x_dict[node_type] + embedding(batch[node_type].n_id)\n",
    "\n",
    "        x_dict = self.gnn(\n",
    "            x_dict,\n",
    "            batch.edge_index_dict,\n",
    "            batch.num_sampled_nodes_dict,\n",
    "            batch.num_sampled_edges_dict,\n",
    "        )\n",
    "\n",
    "        return self.head(x_dict[entity_table][: seed_time.size(0)])\n",
    "\n",
    "    def forward_dst_readout(\n",
    "        self,\n",
    "        batch: HeteroData,\n",
    "        entity_table: NodeType,\n",
    "        dst_table: NodeType,\n",
    "    ) -> Tensor:\n",
    "        if self.id_awareness_emb is None:\n",
    "            raise RuntimeError(\n",
    "                \"id_awareness must be set True to use forward_dst_readout\"\n",
    "            )\n",
    "        seed_time = batch[entity_table].seed_time\n",
    "        x_dict = self.encoder(batch.tf_dict)\n",
    "        # Add ID-awareness to the root node\n",
    "        x_dict[entity_table][: seed_time.size(0)] += self.id_awareness_emb.weight\n",
    "\n",
    "        rel_time_dict = self.temporal_encoder(\n",
    "            seed_time, batch.time_dict, batch.batch_dict\n",
    "        )\n",
    "\n",
    "        for node_type, rel_time in rel_time_dict.items():\n",
    "            x_dict[node_type] = x_dict[node_type] + rel_time\n",
    "\n",
    "        for node_type, embedding in self.embedding_dict.items():\n",
    "            x_dict[node_type] = x_dict[node_type] + embedding(batch[node_type].n_id)\n",
    "\n",
    "        x_dict = self.gnn(\n",
    "            x_dict,\n",
    "            batch.edge_index_dict,\n",
    "        )\n",
    "\n",
    "        return self.head(x_dict[dst_table])\n",
    "\n",
    "\n",
    "model = Model(\n",
    "    data=data,\n",
    "    col_stats_dict=col_stats_dict,\n",
    "    num_layers=2,\n",
    "    channels=128,\n",
    "    out_channels=1,\n",
    "    aggr=\"sum\",\n",
    "    norm=\"batch_norm\",\n",
    ").to(device)\n",
    "\n",
    "\n",
    "# if you try out different RelBench tasks you will need to change these\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.005)\n",
    "epochs = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vl-6So7Llb-p"
   },
   "source": [
    "We also need standard train/test loops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SAHRIr15lVs6"
   },
   "outputs": [],
   "source": [
    "def train() -> float:\n",
    "    model.train()\n",
    "\n",
    "    loss_accum = count_accum = 0\n",
    "    for batch in tqdm(loader_dict[\"train\"]):\n",
    "        batch = batch.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        pred = model(\n",
    "            batch,\n",
    "            task.entity_table,\n",
    "        )\n",
    "        pred = pred.view(-1) if pred.size(1) == 1 else pred\n",
    "\n",
    "        loss = loss_fn(pred.float(), batch[entity_table].y.float())\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        loss_accum += loss.detach().item() * pred.size(0)\n",
    "        count_accum += pred.size(0)\n",
    "\n",
    "    return loss_accum / count_accum\n",
    "\n",
    "\n",
    "@torch.no_grad()\n",
    "def test(loader: NeighborLoader) -> np.ndarray:\n",
    "    model.eval()\n",
    "\n",
    "    pred_list = []\n",
    "    for batch in loader:\n",
    "        batch = batch.to(device)\n",
    "        pred = model(\n",
    "            batch,\n",
    "            task.entity_table,\n",
    "        )\n",
    "        pred = pred.view(-1) if pred.size(1) == 1 else pred\n",
    "        pred_list.append(pred.detach().cpu())\n",
    "    return torch.cat(pred_list, dim=0).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4s-p7dW1ledd"
   },
   "source": [
    "Now we are ready to train!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yF3W68Eqlew_",
    "outputId": "a81a48dc-234a-47f3-8759-8dc9a766661c"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:03<00:00,  4.55it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01, Train loss: 4.5940603298195954, Val metrics: {'r2': 0.26664876365429346, 'mae': 3.192621118001486, 'rmse': 3.970144408166473}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 02, Train loss: 4.568261575973963, Val metrics: {'r2': 0.28401349234585715, 'mae': 3.1377921566297466, 'rmse': 3.9228590932193326}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.33it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 03, Train loss: 4.508331975344442, Val metrics: {'r2': 0.2779053222548966, 'mae': 3.145869382063229, 'rmse': 3.9395567561975118}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 04, Train loss: 4.454095084475588, Val metrics: {'r2': 0.2748767008953672, 'mae': 3.1624791348545886, 'rmse': 3.9478097883334575}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:03<00:00,  3.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 05, Train loss: 4.428226253993716, Val metrics: {'r2': 0.2624319702044985, 'mae': 3.2122242543724435, 'rmse': 3.9815422768300364}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 06, Train loss: 4.3808791889818695, Val metrics: {'r2': 0.23813335297839988, 'mae': 3.2412406909282634, 'rmse': 4.046595277497622}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:03<00:00,  4.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 07, Train loss: 4.338853529823013, Val metrics: {'r2': 0.2456946012502239, 'mae': 3.247610264136621, 'rmse': 4.026464715586905}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 08, Train loss: 4.280953785724709, Val metrics: {'r2': 0.2494808324189115, 'mae': 3.227723037073751, 'rmse': 4.016346595592636}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:03<00:00,  4.67it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 09, Train loss: 4.242680662992767, Val metrics: {'r2': 0.25660482780398064, 'mae': 3.185533819345132, 'rmse': 3.997239384230069}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 15/15 [00:02<00:00,  5.23it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10, Train loss: 4.173513173641142, Val metrics: {'r2': 0.24525833749507364, 'mae': 3.2292002695755078, 'rmse': 4.027628930176982}\n",
      "Best Val metrics: {'r2': 0.2855097827935169, 'mae': 3.132982980591819, 'rmse': 3.918757894088351}\n",
      "Best test metrics: {'r2': -0.026182202542586408, 'mae': 4.38193179076178, 'rmse': 5.278154456184378}\n"
     ]
    }
   ],
   "source": [
    "state_dict = None\n",
    "best_val_metric = -math.inf if higher_is_better else math.inf\n",
    "for epoch in range(1, epochs + 1):\n",
    "    train_loss = train()\n",
    "    val_pred = test(loader_dict[\"val\"])\n",
    "    val_metrics = task.evaluate(val_pred, val_table)\n",
    "    print(f\"Epoch: {epoch:02d}, Train loss: {train_loss}, Val metrics: {val_metrics}\")\n",
    "\n",
    "    if (higher_is_better and val_metrics[tune_metric] > best_val_metric) or (\n",
    "        not higher_is_better and val_metrics[tune_metric] < best_val_metric\n",
    "    ):\n",
    "        best_val_metric = val_metrics[tune_metric]\n",
    "        state_dict = copy.deepcopy(model.state_dict())\n",
    "\n",
    "\n",
    "model.load_state_dict(state_dict)\n",
    "val_pred = test(loader_dict[\"val\"])\n",
    "val_metrics = task.evaluate(val_pred, val_table)\n",
    "print(f\"Best Val metrics: {val_metrics}\")\n",
    "\n",
    "test_pred = test(loader_dict[\"test\"])\n",
    "test_metrics = task.evaluate(test_pred)\n",
    "print(f\"Best test metrics: {test_metrics}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "58i-5Z508liB"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}