{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "Ten notatnik zawiera przykłady implementacji modeli dla grafów heterogenicznych."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch as th\n",
    "import torch.nn as nn\n",
    "import torch_geometric as pyg\n",
    "import torchmetrics as tm\n",
    "from typing import Literal\n",
    "import pytorch_lightning as pl\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics import classification_report\n",
    "\n",
    "from tqdm.auto import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przykład trywialnego grafu heterogenicznego"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  author={ x=[2, 2] },\n",
       "  paper={ x=[1, 3] },\n",
       "  journal={ x=[1, 3] },\n",
       "  redactor={ x=[1, 1] },\n",
       "  (author, writes, paper)={ edge_index=[2, 2] },\n",
       "  (journal, published, paper)={ edge_index=[2, 1] },\n",
       "  (redactor, edited, paper)={ edge_index=[2, 1] }\n",
       ")"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "example_g = pyg.data.HeteroData()\n",
    "example_g['author'].x = th.tensor([\n",
    "    [1., 30.],\n",
    "    [2., 35.]\n",
    "])\n",
    "example_g['paper'].x = th.tensor([\n",
    "    [0.1, 0.2, 0.3]\n",
    "])\n",
    "example_g['journal'].x = th.tensor([\n",
    "    [3., 10., 100.]\n",
    "])\n",
    "example_g['redactor'].x = th.tensor([\n",
    "    [5.]\n",
    "])\n",
    "example_g[('author', 'writes', 'paper')].edge_index = th.tensor([\n",
    "    [0, 1],\n",
    "    [0, 0]\n",
    "])\n",
    "example_g[('journal', 'published', 'paper')].edge_index = th.tensor([\n",
    "    [0],\n",
    "    [0]\n",
    "])\n",
    "example_g[('redactor', 'edited', 'paper')].edge_index = th.tensor([\n",
    "    [0],\n",
    "    [0]\n",
    "])\n",
    "\n",
    "example_g"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[186.6000, 186.6000]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TARGET_DIM = 2\n",
    "\n",
    "node_2_dim = {\n",
    "    \"author\": 2,\n",
    "    \"paper\": 3,\n",
    "    \"journal\": 3,\n",
    "    \"redactor\": 1\n",
    "}\n",
    "\n",
    "class TrivialRelGCN(pyg.nn.MessagePassing):\n",
    "    \n",
    "    def __init__(\n",
    "        self, \n",
    "        node_2_dim: dict,\n",
    "        relations: list,\n",
    "        target_dim: int,\n",
    "        aggr='sum',\n",
    "        **kwargs):\n",
    "        super().__init__(aggr=aggr, **kwargs)\n",
    "        self.node_2_dim = node_2_dim\n",
    "        self.relations = relations\n",
    "        self.target_dim = target_dim\n",
    "        \n",
    "        # Dedykowane macierze wag (=1), dla każdego typu wierzchołka\n",
    "        self.node_weights = {}\n",
    "        for node, dim in node_2_dim.items():\n",
    "            W = nn.Parameter(th.ones(dim, self.target_dim))\n",
    "            self.node_weights[node] = W\n",
    "            \n",
    "        self.node_weights = nn.ParameterDict(self.node_weights)\n",
    "        \n",
    "    def forward(self, x_dict: dict, edge_index_dict: dict):\n",
    "        # Nowe reprezentacje wierzchołków\n",
    "        new_x_dict = {\n",
    "            node: th.zeros(x_dict[node].shape[0], self.target_dim)\n",
    "            for node in x_dict\n",
    "        }\n",
    "        \n",
    "        # Dla każdej relacji\n",
    "        for relname in self.relations:\n",
    "            # Pobierz typy: wysyłającego, relację, odbierającego\n",
    "            u, rel, v = relname.split('__')\n",
    "            \n",
    "            # Pobierz wagi dla typu sąsiada u\n",
    "            W_u = self.node_weights[u]\n",
    "            # Pobierz listę sąsiedztwa dla typu sąsiada i relacji\n",
    "            edges_rel = edge_index_dict[(u, rel, v)]\n",
    "            \n",
    "            h_u = x_dict[u]\n",
    "            \n",
    "            # Wykonaj przekazywanie wiadomości\n",
    "            h_u_new = self.propagate(\n",
    "                x=h_u, \n",
    "                edge_index=edges_rel,\n",
    "                W_rel=W_u)\n",
    "            new_x_dict[v] += h_u_new\n",
    "        \n",
    "        # Dodanie nowych reprezentacji wierzchołków\n",
    "        # odbierających v: W0 * h_v\n",
    "        for node, h_prev in x_dict.items():\n",
    "            W_node = self.node_weights[node]\n",
    "            h_new = th.matmul(h_prev, W_node)\n",
    "            new_x_dict[node] += h_new\n",
    "        return new_x_dict\n",
    "\n",
    "    def message(self, x_j: th.tensor, W_rel: th.tensor):\n",
    "        # Nowa wiadomość od sąsiada: W_rel * h_u\n",
    "        return th.matmul(x_j, W_rel)\n",
    "    \n",
    "    def aggregate(self, inputs):\n",
    "        # Agregacja wiadomości jednego typu\n",
    "        return th.sum(inputs, dim=0)\n",
    "        \n",
    "\n",
    "        \n",
    "trivial_rel_gcn = TrivialRelGCN(\n",
    "    node_2_dim=node_2_dim,\n",
    "    relations=[\n",
    "        'author__writes__paper',\n",
    "        'journal__published__paper',\n",
    "        'redactor__edited__paper'\n",
    "    ],\n",
    "    target_dim=TARGET_DIM,\n",
    "    node_dim=-2\n",
    ")\n",
    "\n",
    "h_new = trivial_rel_gcn(example_g.x_dict, example_g.edge_index_dict)\n",
    "h_new['paper']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Graf cytowań naukowych - DBLP\n",
    "\n",
    "Przygorujemy też znacznie większy graf zwierający cytowania naukowe. Będziemy go używać do trenowania modeli."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://www.dropbox.com/s/yh4grpeks87ugr2/DBLP_processed.zip?dl=1\n",
      "Extracting data/raw/DBLP_processed.zip\n",
      "Processing...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "dblp = pyg.datasets.DBLP(root=\"./data/\",)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  author={\n",
       "    x=[4057, 334],\n",
       "    y=[4057],\n",
       "    train_mask=[4057],\n",
       "    val_mask=[4057],\n",
       "    test_mask=[4057],\n",
       "  },\n",
       "  paper={ x=[14328, 4231] },\n",
       "  term={ x=[7723, 50] },\n",
       "  conference={ num_nodes=20 },\n",
       "  (author, to, paper)={ edge_index=[2, 19645] },\n",
       "  (paper, to, author)={ edge_index=[2, 19645] },\n",
       "  (paper, to, term)={ edge_index=[2, 85810] },\n",
       "  (paper, to, conference)={ edge_index=[2, 14328] },\n",
       "  (term, to, paper)={ edge_index=[2, 85810] },\n",
       "  (conference, to, paper)={ edge_index=[2, 14328] }\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dblp_graph = dblp[0]\n",
    "dblp_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Obiekt typu \"konferencja\" nie posiada atrybutów - dodajemy stałą = 1 dla każdego wierzchołka.\n",
    "dblp_graph['conference'].x = th.ones(dblp_graph['conference'].num_nodes, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Modele predykcyjne\n",
    "\n",
    "W poniższej sekcji sprawdzimy działanie i budowę modeli predykcyjnych dla grafów heterogenicznych."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ogólny klasyfikator\n",
    "\n",
    "Model do klasyfikacji wierzchołków typu \"autor\". Przyjmuje sieci grafową budującą wektory reprezentacji, a następnie (dla wybranego typu wierzchołka) przewiduje etykietę."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NodeClassifier(nn.Module):\n",
    "    \n",
    "    def __init__(\n",
    "        self, \n",
    "        gnn, \n",
    "        num_classes: int=dblp_graph['author'].y.unique().size(0),\n",
    "        lin_sizes: tuple[int, ...] = (256, 128, 64)):\n",
    "        super().__init__()\n",
    "        self.gnn = gnn\n",
    "        \n",
    "        lin_layers = []\n",
    "        for dim_in, dim_out in zip(lin_sizes[:-1], lin_sizes[1:]):\n",
    "            lin_layers.append(nn.Linear(dim_in, dim_out))\n",
    "            lin_layers.append(nn.ReLU())\n",
    "        lin_layers.append(nn.Linear(lin_sizes[-1], num_classes))\n",
    "        self.lin = nn.Sequential(*lin_layers)\n",
    "        \n",
    "    def forward(self, x_dict, edge_index_dict):\n",
    "        embeds = self.gnn(x_dict, edge_index_dict)\n",
    "        author_embed = embeds['author']\n",
    "        lin_out = self.lin(author_embed)\n",
    "        return lin_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Funkcje pomocnicze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def log_metrics(yhat, y, prefix=\"\"):\n",
    "    f1_score = tm.functional.f1_score(yhat, y, num_classes=dblp_graph['author'].y.unique().size(0), task='multiclass')\n",
    "    precision = tm.functional.precision(yhat, y, num_classes=dblp_graph['author'].y.unique().size(0), task='multiclass')\n",
    "    recall = tm.functional.recall(yhat, y, num_classes=dblp_graph['author'].y.unique().size(0), task='multiclass')\n",
    "    auc = tm.functional.auroc(yhat, y, num_classes=dblp_graph['author'].y.unique().size(0), task='multiclass')\n",
    "    \n",
    "    print(f\"{prefix}F1: {f1_score:.2f}  AUC: {auc:.2f}  Precision: {precision:.2f}  Recall: {recall:.2f}\")\n",
    "    return {\n",
    "        f\"{prefix} F1\": f1_score.item(),\n",
    "        f\"{prefix} Precision\": precision.item(),\n",
    "        f\"{prefix} Recall\": recall.item(),\n",
    "        f\"{prefix} AUC\": auc.item()\n",
    "    }\n",
    "\n",
    "\n",
    "def train_model(classifier, data, epochs=10, lr=0.01):\n",
    "    optimizer = th.optim.Adam(classifier.parameters(), lr=lr)\n",
    "    loss_fn = nn.CrossEntropyLoss()\n",
    "    epoch_losses = {\"train\": [], \"validation\": []}\n",
    "    train_metrics = []\n",
    "    val_metrics = []\n",
    "\n",
    "    for epoch in tqdm(range(epochs)):\n",
    "        classifier.train()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        train_mask = data['author'].train_mask\n",
    "        yhat = classifier(data.x_dict, data.edge_index_dict)[train_mask]\n",
    "        y = data['author'].y[train_mask]\n",
    "\n",
    "        loss_val = loss_fn(yhat, y)\n",
    "        loss_val.backward()\n",
    "        optimizer.step()\n",
    "        epoch_losses['train'].append(loss_val.item())\n",
    "\n",
    "        print(f\"Epoch {epoch} Loss: {loss_val.item():.2f}\")\n",
    "        epoch_train_metrics = log_metrics(yhat, y, prefix=\"Train \")\n",
    "        train_metrics.append(epoch_train_metrics)\n",
    "\n",
    "        classifier.eval()\n",
    "        val_mask = data['author'].val_mask\n",
    "        yhat_val = classifier(data.x_dict, data.edge_index_dict)[val_mask]\n",
    "        y_val = data['author'].y[val_mask]\n",
    "        val_loss = loss_fn(yhat_val, y_val).item()\n",
    "        epoch_losses['validation'].append(val_loss)\n",
    "\n",
    "        epoch_val_metrics = log_metrics(yhat_val, y_val, prefix=\"Val \")\n",
    "        val_metrics.append(epoch_val_metrics)\n",
    "\n",
    "        print(\"===\\n\")\n",
    "\n",
    "    return epoch_losses, train_metrics, val_metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_losses(epoch_losses, title='Funkcja kosztu w czasie', train_label='Koszt treningowy', val_label='Koszt walidacyjny', fname: str = \"loss_in_time.tif\"):\n",
    "    plt.figure(figsize=(10, 6))\n",
    "    plt.plot(epoch_losses['train'], label=train_label)\n",
    "    plt.plot(epoch_losses['validation'], label=val_label)\n",
    "    plt.xlabel('Epoka')\n",
    "    plt.ylabel('Koszt')\n",
    "    plt.title(title)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(fname)\n",
    "    plt.show()\n",
    "\n",
    "def plot_metrics(metrics, title='Metryki jakości predykcji', metric_labels=None, fname: str = \"metrics_in_time.tif\"):\n",
    "    if metric_labels is None:\n",
    "        metric_labels = list(metrics[0].keys())\n",
    "    \n",
    "    metrics_df = pd.DataFrame(metrics)\n",
    "    metrics_df.plot()\n",
    "    plt.xlabel('Epoka')\n",
    "    plt.ylabel('Wartość metryki')\n",
    "    plt.title(title)\n",
    "    plt.legend(loc='best')\n",
    "    plt.tight_layout()\n",
    "    plt.savefig(fname, pil_kwargs={\"compression\": \"tiff_lzw\"})\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_model(model, data, mask_key='test_mask', prefix=\"Test \"):\n",
    "    model.eval()\n",
    "    mask = data['author'][mask_key]\n",
    "    yhat = model(data.x_dict, data.edge_index_dict)[mask]\n",
    "    y = data['author'].y[mask]\n",
    "\n",
    "    log_metrics(yhat, y, prefix=prefix)\n",
    "    print(classification_report(y.cpu().numpy(), yhat.argmax(dim=1).cpu().numpy()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uproszczona implementacja\n",
    "\n",
    "Zacznijmy od uproszczonej implementacji modelu - pojedynczej warstwy splotu grafowego dla grafów heterogenicznych."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SimpleHeteroConv(nn.Module):\n",
    "    \"\"\"\n",
    "    Klasa implementująca uproszczony splot dla grafów heterogenicznych.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    base_architecture : pyg.nn.MessagePassing\n",
    "        Podstawowa architektura splotu grafowego.\n",
    "    node_type_dims : dict[str, int]\n",
    "        Słownik zawierający wymiary cech dla poszczególnych typów wierzchołków.\n",
    "    rel_types : list[tuple[str, str, str]]\n",
    "        Lista typów relacji w grafie heterogenicznym.\n",
    "    hidden_dim : int\n",
    "        Wymiar ukryty dla warstwy splotowej.\n",
    "    node_type_aggr : Literal['add', 'mean'], optional\n",
    "        Metoda agregacji dla wierzchołków tego samego typu, domyślnie 'add'.\n",
    "    **model_kwargs\n",
    "        Dodatkowe argumenty przekazywane do warstwy splotowej.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    final_embeddings : dict[str, th.Tensor]\n",
    "        Słownik zawierający końcowe osadzenia dla każdego typu wierzchołka.\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(\n",
    "            self, \n",
    "            base_architecture: pyg.nn.MessagePassing,\n",
    "            node_type_dims: dict[str, int],\n",
    "            rel_types: list[tuple[str, str, str]],\n",
    "            hidden_dim: int,\n",
    "            node_type_aggr: Literal['add', 'mean'] = 'add',\n",
    "            **model_kwargs):\n",
    "\n",
    "        super().__init__()\n",
    "        # Słownik zawierający wymiary cech dla poszczególnych typów wierzchołków\n",
    "        self.node_type_dims = node_type_dims\n",
    "        # Lista relacji występujących w grafie\n",
    "        self.rel_types = rel_types\n",
    "        # Wspólna wymiarowość dla warstwy splotowej.\n",
    "        self.hidden_dim = hidden_dim\n",
    "        # Sposób agregacji sasiedztwa POMIĘDZY relacjami\n",
    "        self.node_type_aggr = node_type_aggr\n",
    "\n",
    "        # Tworzenie słownika modułów splotowych dla każdej relacji z osobna.\n",
    "        self.rel_convs = nn.ModuleDict()\n",
    "\n",
    "        # Iteracja po typach relacji\n",
    "        for rel in rel_types:\n",
    "            # Pobranie wymiarów cech dla wierzchołków \n",
    "            # wysyłającego (u) i odbierającego (v).\n",
    "            src_type, rel_type, trg_type = rel\n",
    "            in_channels_src = node_type_dims[src_type]\n",
    "            in_channels_trg = node_type_dims[trg_type]\n",
    "            \n",
    "            # Inicjalizacja warstwy splotowej dla danej relacji\n",
    "            conv = base_architecture(\n",
    "                (in_channels_src, in_channels_trg), \n",
    "                hidden_dim, \n",
    "                **model_kwargs)\n",
    "            self.rel_convs[\"__\".join([src_type, rel_type, trg_type])] = conv\n",
    "        \n",
    "    def forward(self, x_dict: dict[str, th.Tensor], edge_index_dict: dict[str, th.Tensor]):\n",
    "        \"\"\"\n",
    "        Wykonuje przekazywanie wiadomości.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x_dict : dict[str, th.Tensor]\n",
    "            Słownik zawierający cechy wierzchołków dla każdego typu wierzchołka.\n",
    "        edge_index_dict : dict[str, th.Tensor]\n",
    "            Słownik zawierający indeksy krawędzi dla każdej relacji.\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        final_embeddings : dict[str, th.Tensor]\n",
    "            Słownik zawierający końcowe osadzenia dla każdego typu wierzchołka.\n",
    "        \"\"\"\n",
    "        \n",
    "        # Inicjalizacja słownika na osadzenia wierzchołków\n",
    "        node_type_embeddings = {\n",
    "            node_type: []\n",
    "            for node_type, x in x_dict.items()\n",
    "        }\n",
    "\n",
    "        # Przetwarzanie każdej relacji w grafie\n",
    "        for rel, edge_index in edge_index_dict.items():\n",
    "            src_type, rel_type, trg_type = rel\n",
    "\n",
    "            # Pobranie odpowiedniej warstwy splotu\n",
    "            conv = self.rel_convs[\"__\".join([src_type, rel_type, trg_type])]\n",
    "\n",
    "            # Macierz cech dla wierzchołka wysyłającego i odbierającego\n",
    "            x_src = x_dict[src_type]\n",
    "            x_trg = x_dict[trg_type]\n",
    "\n",
    "            # Wykonanie splotu dla danej relacji\n",
    "            target_embeddings = conv((x_src, x_trg), edge_index)\n",
    "\n",
    "            # Dodanie osadzeń do odpowiedniego typu wierzchołka\n",
    "            node_type_embeddings[trg_type].append(target_embeddings)\n",
    "        \n",
    "        final_embeddings = {}\n",
    "        # Agregacja osadzeń wierzchołków pomiędzy relacjami\n",
    "        # dla tego samego wierzchołka v.\n",
    "        if self.node_type_aggr == 'mean':\n",
    "            for node_type, embeddings in node_type_embeddings.items():\n",
    "                final_embeddings[node_type] = th.stack(embeddings).mean(dim=0)\n",
    "        elif self.node_type_aggr == 'add':\n",
    "            for node_type, embeddings in node_type_embeddings.items():\n",
    "                final_embeddings[node_type] = th.stack(embeddings).sum(dim=0)\n",
    "        \n",
    "        return final_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sprawdzenie uproszczonego modelu na prostych danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "example_g_undirected = pyg.transforms.ToUndirected()(example_g)\n",
    "\n",
    "simple_hetero_gnn = SimpleHeteroConv(\n",
    "    pyg.nn.SAGEConv, \n",
    "    node_type_dims=node_2_dim,\n",
    "    rel_types=example_g_undirected.metadata()[1],\n",
    "    hidden_dim=16\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    out = simple_hetero_gnn(example_g_undirected.x_dict, example_g_undirected.edge_index_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "author torch.Size([2, 16])\n",
      "paper torch.Size([1, 16])\n",
      "journal torch.Size([1, 16])\n",
      "redactor torch.Size([1, 16])\n"
     ]
    }
   ],
   "source": [
    "for k, embed in out.items():\n",
    "    print(k, embed.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------------------------------------+---------------+----------------+----------+\n",
      "| Layer                                          | Input Shape   | Output Shape   |   #Param |\n",
      "|------------------------------------------------+---------------+----------------+----------|\n",
      "| SimpleHeteroConv                               |               |                |      576 |\n",
      "| ├─(rel_convs)ModuleDict                        | --            | --             |      576 |\n",
      "| │    └─(author__writes__paper)SAGEConv         | [2, 2]        | [1, 16]        |       96 |\n",
      "| │    └─(journal__published__paper)SAGEConv     | [2, 1]        | [1, 16]        |      112 |\n",
      "| │    └─(redactor__edited__paper)SAGEConv       | [2, 1]        | [1, 16]        |       80 |\n",
      "| │    └─(paper__rev_writes__author)SAGEConv     | [2, 2]        | [2, 16]        |       96 |\n",
      "| │    └─(paper__rev_published__journal)SAGEConv | [2, 1]        | [1, 16]        |      112 |\n",
      "| │    └─(paper__rev_edited__redactor)SAGEConv   | [2, 1]        | [1, 16]        |       80 |\n",
      "+------------------------------------------------+---------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg.nn.summary(simple_hetero_gnn, example_g_undirected.x_dict, example_g_undirected.edge_index_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------------------------------------------------+---------------+----------------+----------+\n",
      "| Layer                                               | Input Shape   | Output Shape   | #Param   |\n",
      "|-----------------------------------------------------+---------------+----------------+----------|\n",
      "| GraphModule                                         |               |                | 576      |\n",
      "| ├─(convs)ModuleList                                 | --            | --             | 576      |\n",
      "| │    └─(0)ModuleDict                                | --            | --             | 576      |\n",
      "| │    │    └─(author__writes__paper)SAGEConv         | [2, 2]        | [1, 16]        | 96       |\n",
      "| │    │    └─(journal__published__paper)SAGEConv     | [2, 1]        | [1, 16]        | 112      |\n",
      "| │    │    └─(redactor__edited__paper)SAGEConv       | [2, 1]        | [1, 16]        | 80       |\n",
      "| │    │    └─(paper__rev_writes__author)SAGEConv     | [2, 2]        | [2, 16]        | 96       |\n",
      "| │    │    └─(paper__rev_published__journal)SAGEConv | [2, 1]        | [1, 16]        | 112      |\n",
      "| │    │    └─(paper__rev_edited__redactor)SAGEConv   | [2, 1]        | [1, 16]        | 80       |\n",
      "| ├─(dropout)Dropout                                  | --            | --             | --       |\n",
      "| ├─(act)ReLU                                         | --            | --             | --       |\n",
      "| ├─(norms)ModuleList                                 | --            | --             | --       |\n",
      "| │    └─(0)Identity                                  | --            | --             | --       |\n",
      "+-----------------------------------------------------+---------------+----------------+----------+\n"
     ]
    }
   ],
   "source": [
    "auto_model_hetero_g = pyg.nn.to_hetero(pyg.nn.GraphSAGE((-1, -1), hidden_channels=16, num_layers=1), example_g_undirected.metadata())\n",
    "with th.no_grad():\n",
    "    out_auto = auto_model_hetero_g(example_g_undirected.x_dict, example_g_undirected.edge_index_dict)\n",
    "print(pyg.nn.summary(auto_model_hetero_g, example_g_undirected.x_dict, example_g_undirected.edge_index_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sprawdzenie uproszczonego modelu dla DBLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  author={\n",
       "    x=[4057, 334],\n",
       "    y=[4057],\n",
       "    train_mask=[4057],\n",
       "    val_mask=[4057],\n",
       "    test_mask=[4057],\n",
       "  },\n",
       "  paper={ x=[14328, 4231] },\n",
       "  term={ x=[7723, 50] },\n",
       "  conference={\n",
       "    num_nodes=20,\n",
       "    x=[20, 1],\n",
       "  },\n",
       "  (author, to, paper)={ edge_index=[2, 19645] },\n",
       "  (paper, to, author)={ edge_index=[2, 19645] },\n",
       "  (paper, to, term)={ edge_index=[2, 85810] },\n",
       "  (paper, to, conference)={ edge_index=[2, 14328] },\n",
       "  (term, to, paper)={ edge_index=[2, 85810] },\n",
       "  (conference, to, paper)={ edge_index=[2, 14328] }\n",
       ")"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dblp_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "author torch.Size([4057, 128])\n",
      "paper torch.Size([14328, 128])\n",
      "term torch.Size([7723, 128])\n",
      "conference torch.Size([20, 128])\n"
     ]
    }
   ],
   "source": [
    "simple_dblp_hetero_gnn = SimpleHeteroConv(\n",
    "    pyg.nn.SAGEConv,\n",
    "    node_type_dims={node_type: dblp_graph[node_type].x.shape[1] for node_type in dblp_graph.metadata()[0]},\n",
    "    rel_types=dblp_graph.metadata()[1],\n",
    "    hidden_dim=128\n",
    ")\n",
    "\n",
    "with th.no_grad():\n",
    "    out = simple_dblp_hetero_gnn(dblp_graph.x_dict, dblp_graph.edge_index_dict)\n",
    "\n",
    "for k, v in out.items():\n",
    "    print(k, v.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Szkolenie uproszczonego modelu dla DBLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "simple_gnn = pyg.nn.Sequential('x_dict,edge_index_dict', [\n",
    "    (\n",
    "     SimpleHeteroConv(\n",
    "        pyg.nn.SAGEConv,\n",
    "        node_type_dims={node_type: dblp_graph[node_type].x.shape[1] for node_type in dblp_graph.metadata()[0]},\n",
    "        rel_types=dblp_graph.metadata()[1],\n",
    "        node_type_aggr='add',\n",
    "        hidden_dim=512,\n",
    "        aggr='mean'),\n",
    "    'x_dict,edge_index_dict->x_dict'\n",
    "    ),\n",
    "    (\n",
    "     SimpleHeteroConv(\n",
    "        pyg.nn.SAGEConv,\n",
    "        node_type_dims={node_type: 512 for node_type in dblp_graph.metadata()[0]},\n",
    "        rel_types=dblp_graph.metadata()[1],\n",
    "        node_type_aggr='add',\n",
    "        hidden_dim=256,\n",
    "        aggr='mean'),\n",
    "    'x_dict,edge_index_dict->x_dict'\n",
    "    )\n",
    "])\n",
    "\n",
    "pyg.seed_everything(123)\n",
    "simple_classifier = NodeClassifier(simple_gnn, lin_sizes=(256, 128, 64))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c09413f1977f42168fcd4eec3a5e8089",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 Loss: 1.41\n",
      "Train F1: 0.17  AUC: 0.55  Precision: 0.17  Recall: 0.17\n",
      "Val F1: 0.26  AUC: 0.84  Precision: 0.26  Recall: 0.26\n",
      "===\n",
      "\n",
      "Epoch 1 Loss: 2.01\n",
      "Train F1: 0.24  AUC: 0.91  Precision: 0.24  Recall: 0.24\n",
      "Val F1: 0.26  AUC: 0.73  Precision: 0.26  Recall: 0.26\n",
      "===\n",
      "\n",
      "Epoch 2 Loss: 1.85\n",
      "Train F1: 0.29  AUC: 0.78  Precision: 0.29  Recall: 0.29\n",
      "Val F1: 0.57  AUC: 0.75  Precision: 0.57  Recall: 0.57\n",
      "===\n",
      "\n",
      "Epoch 3 Loss: 1.30\n",
      "Train F1: 0.75  AUC: 0.78  Precision: 0.75  Recall: 0.75\n",
      "Val F1: 0.62  AUC: 0.87  Precision: 0.62  Recall: 0.62\n",
      "===\n",
      "\n",
      "Epoch 4 Loss: 0.48\n",
      "Train F1: 0.86  AUC: 0.98  Precision: 0.86  Recall: 0.86\n",
      "Val F1: 0.65  AUC: 0.87  Precision: 0.65  Recall: 0.65\n",
      "===\n",
      "\n",
      "Epoch 5 Loss: 0.29\n",
      "Train F1: 0.91  AUC: 0.98  Precision: 0.91  Recall: 0.91\n",
      "Val F1: 0.70  AUC: 0.92  Precision: 0.70  Recall: 0.70\n",
      "===\n",
      "\n",
      "Epoch 6 Loss: 0.21\n",
      "Train F1: 0.93  AUC: 1.00  Precision: 0.93  Recall: 0.93\n",
      "Val F1: 0.74  AUC: 0.92  Precision: 0.74  Recall: 0.74\n",
      "===\n",
      "\n",
      "Epoch 7 Loss: 0.02\n",
      "Train F1: 0.99  AUC: 1.00  Precision: 0.99  Recall: 0.99\n",
      "Val F1: 0.73  AUC: 0.91  Precision: 0.73  Recall: 0.73\n",
      "===\n",
      "\n",
      "Epoch 8 Loss: 0.02\n",
      "Train F1: 0.99  AUC: 1.00  Precision: 0.99  Recall: 0.99\n",
      "Val F1: 0.71  AUC: 0.92  Precision: 0.71  Recall: 0.71\n",
      "===\n",
      "\n",
      "Epoch 9 Loss: 0.00\n",
      "Train F1: 1.00  AUC: 1.00  Precision: 1.00  Recall: 1.00\n",
      "Val F1: 0.72  AUC: 0.92  Precision: 0.72  Recall: 0.72\n",
      "===\n",
      "\n"
     ]
    }
   ],
   "source": [
    "epoch_losses, train_metrics, val_metrics = train_model(simple_classifier, dblp_graph, epochs=10, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test F1: 0.73  AUC: 0.91  Precision: 0.73  Recall: 0.73\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.91      0.72      0.80       977\n",
      "           1       0.53      0.63      0.57       584\n",
      "           2       0.66      0.88      0.76       892\n",
      "           3       0.89      0.66      0.76       804\n",
      "\n",
      "    accuracy                           0.73      3257\n",
      "   macro avg       0.75      0.72      0.72      3257\n",
      "weighted avg       0.77      0.73      0.74      3257\n",
      "\n"
     ]
    }
   ],
   "source": [
    "evaluate_model(simple_classifier, dblp_graph, mask_key='test_mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 60260 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_losses(epoch_losses, title='Funkcja kosztu w czasie dla prostego modelu', train_label='Koszt treningowy', val_label='Koszt walidacyjny', fname=\"loss_in_time_simple_model.tiff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48272 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 47648 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_metrics(train_metrics, title='Metryki treningowe dla prostego modelu', fname=\"train_metrics_in_time_simple_model.tiff\")\n",
    "plot_metrics(val_metrics, title='Metryki walidacyjne dla prostego modelu', fname=\"val_metrics_in_time_simple_model.tiff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------------------------------+---------------+----------------+------------+\n",
      "| Layer                             | Input Shape   | Output Shape   | #Param     |\n",
      "|-----------------------------------+---------------+----------------+------------|\n",
      "| NodeClassifier                    |               | [4057, 4]      | 15,010,756 |\n",
      "| ├─(gnn)Sequential                 |               |                | 14,969,344 |\n",
      "| │    └─(module_0)SimpleHeteroConv |               |                | 13,394,944 |\n",
      "| │    │    └─(rel_convs)ModuleDict | --            | --             | 13,394,944 |\n",
      "| │    └─(module_1)SimpleHeteroConv |               |                | 1,574,400  |\n",
      "| │    │    └─(rel_convs)ModuleDict | --            | --             | 1,574,400  |\n",
      "| ├─(lin)Sequential                 | [4057, 256]   | [4057, 4]      | 41,412     |\n",
      "| │    └─(0)Linear                  | [4057, 256]   | [4057, 128]    | 32,896     |\n",
      "| │    └─(1)ReLU                    | [4057, 128]   | [4057, 128]    | --         |\n",
      "| │    └─(2)Linear                  | [4057, 128]   | [4057, 64]     | 8,256      |\n",
      "| │    └─(3)ReLU                    | [4057, 64]    | [4057, 64]     | --         |\n",
      "| │    └─(4)Linear                  | [4057, 64]    | [4057, 4]      | 260        |\n",
      "+-----------------------------------+---------------+----------------+------------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg.nn.summary(simple_classifier, dblp_graph.x_dict, dblp_graph.edge_index_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sprawdzenie modelu automatycznego dla DBLP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e8d924fc8dd14f0fb611ba2a005e4643",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/10 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0 Loss: 1.39\n",
      "Train F1: 0.24  AUC: 0.59  Precision: 0.24  Recall: 0.24\n",
      "Val F1: 0.39  AUC: 0.72  Precision: 0.39  Recall: 0.39\n",
      "===\n",
      "\n",
      "Epoch 1 Loss: 1.48\n",
      "Train F1: 0.50  AUC: 0.81  Precision: 0.50  Recall: 0.50\n",
      "Val F1: 0.26  AUC: 0.74  Precision: 0.26  Recall: 0.26\n",
      "===\n",
      "\n",
      "Epoch 2 Loss: 1.38\n",
      "Train F1: 0.24  AUC: 0.76  Precision: 0.24  Recall: 0.24\n",
      "Val F1: 0.40  AUC: 0.85  Precision: 0.40  Recall: 0.40\n",
      "===\n",
      "\n",
      "Epoch 3 Loss: 1.27\n",
      "Train F1: 0.56  AUC: 0.92  Precision: 0.56  Recall: 0.56\n",
      "Val F1: 0.47  AUC: 0.92  Precision: 0.47  Recall: 0.47\n",
      "===\n",
      "\n",
      "Epoch 4 Loss: 1.01\n",
      "Train F1: 0.63  AUC: 0.99  Precision: 0.63  Recall: 0.63\n",
      "Val F1: 0.47  AUC: 0.91  Precision: 0.47  Recall: 0.47\n",
      "===\n",
      "\n",
      "Epoch 5 Loss: 0.73\n",
      "Train F1: 0.59  AUC: 0.99  Precision: 0.59  Recall: 0.59\n",
      "Val F1: 0.60  AUC: 0.90  Precision: 0.60  Recall: 0.60\n",
      "===\n",
      "\n",
      "Epoch 6 Loss: 0.48\n",
      "Train F1: 0.80  AUC: 0.99  Precision: 0.80  Recall: 0.80\n",
      "Val F1: 0.73  AUC: 0.90  Precision: 0.73  Recall: 0.73\n",
      "===\n",
      "\n",
      "Epoch 7 Loss: 0.26\n",
      "Train F1: 0.97  AUC: 0.99  Precision: 0.97  Recall: 0.97\n",
      "Val F1: 0.71  AUC: 0.92  Precision: 0.71  Recall: 0.71\n",
      "===\n",
      "\n",
      "Epoch 8 Loss: 0.14\n",
      "Train F1: 0.98  AUC: 1.00  Precision: 0.98  Recall: 0.98\n",
      "Val F1: 0.75  AUC: 0.94  Precision: 0.75  Recall: 0.75\n",
      "===\n",
      "\n",
      "Epoch 9 Loss: 0.03\n",
      "Train F1: 1.00  AUC: 1.00  Precision: 1.00  Recall: 1.00\n",
      "Val F1: 0.79  AUC: 0.95  Precision: 0.79  Recall: 0.79\n",
      "===\n",
      "\n"
     ]
    }
   ],
   "source": [
    "class SageModel(nn.Module):\n",
    "\n",
    "    def __init__(self, sage_dims: tuple[int, ...]):\n",
    "        super().__init__()\n",
    "        dim_in = (-1, -1)\n",
    "        layers = []\n",
    "        for dim_out in sage_dims:\n",
    "            layers.append(pyg.nn.SAGEConv(dim_in, dim_out))\n",
    "            dim_in = (dim_out, dim_out)\n",
    "        self.layers = nn.ModuleList(layers)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, edge_index).relu()\n",
    "        return x\n",
    "\n",
    "pyg.seed_everything(123)    \n",
    "auto_hetero_model = pyg.nn.to_hetero(SageModel((512, 256)), dblp_graph.metadata())\n",
    "auto_classifier = NodeClassifier(auto_hetero_model, lin_sizes=(256, 128, 64))\n",
    "epoch_losses_auto, train_metrics_auto, val_metrics_auto = train_model(auto_classifier, dblp_graph, epochs=10, lr=0.01)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test F1: 0.81  AUC: 0.95  Precision: 0.81  Recall: 0.81\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "           0       0.89      0.86      0.88       977\n",
      "           1       0.87      0.57      0.69       584\n",
      "           2       0.79      0.85      0.82       892\n",
      "           3       0.72      0.87      0.79       804\n",
      "\n",
      "    accuracy                           0.81      3257\n",
      "   macro avg       0.82      0.79      0.79      3257\n",
      "weighted avg       0.82      0.81      0.80      3257\n",
      "\n"
     ]
    }
   ],
   "source": [
    "evaluate_model(auto_classifier, dblp_graph, mask_key='test_mask')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61076 characters]",
      "text/plain": [
       "<Figure size 1000x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_losses(epoch_losses_auto, title='Funkcja kosztu dla automatycznego modelu', fname=\"loss_in_time_auto_model.tiff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 55856 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 49932 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_metrics(train_metrics_auto, title='Metryki treningowe dla automatycznego modelu', fname=\"train_metrics_in_time_auto_model.tif\")\n",
    "plot_metrics(val_metrics_auto, title='Metryki walidacyjne dla automatycznego modelu', fname=\"val_metrics_in_time_auto_model.tif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---------------------------+---------------+----------------+------------+\n",
      "| Layer                     | Input Shape   | Output Shape   | #Param     |\n",
      "|---------------------------+---------------+----------------+------------|\n",
      "| NodeClassifier            |               | [4057, 4]      | 15,010,756 |\n",
      "| ├─(gnn)GraphModule        |               |                | 14,969,344 |\n",
      "| │    └─(layers)ModuleList | --            | --             | 14,969,344 |\n",
      "| │    │    └─(0)ModuleDict | --            | --             | 13,394,944 |\n",
      "| │    │    └─(1)ModuleDict | --            | --             | 1,574,400  |\n",
      "| ├─(lin)Sequential         | [4057, 256]   | [4057, 4]      | 41,412     |\n",
      "| │    └─(0)Linear          | [4057, 256]   | [4057, 128]    | 32,896     |\n",
      "| │    └─(1)ReLU            | [4057, 128]   | [4057, 128]    | --         |\n",
      "| │    └─(2)Linear          | [4057, 128]   | [4057, 64]     | 8,256      |\n",
      "| │    └─(3)ReLU            | [4057, 64]    | [4057, 64]     | --         |\n",
      "| │    └─(4)Linear          | [4057, 64]    | [4057, 4]      | 260        |\n",
      "+---------------------------+---------------+----------------+------------+\n"
     ]
    }
   ],
   "source": [
    "print(pyg.nn.summary(auto_classifier, dblp_graph.x_dict, dblp_graph.edge_index_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GraphModule(\n",
       "  (layers): ModuleList(\n",
       "    (0): ModuleDict(\n",
       "      (author__to__paper): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "      (paper__to__author): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "      (paper__to__term): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "      (paper__to__conference): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "      (term__to__paper): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "      (conference__to__paper): SAGEConv((-1, -1), 512, aggr=mean)\n",
       "    )\n",
       "    (1): ModuleDict(\n",
       "      (author__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "      (paper__to__author): SAGEConv((512, 512), 256, aggr=mean)\n",
       "      (paper__to__term): SAGEConv((512, 512), 256, aggr=mean)\n",
       "      (paper__to__conference): SAGEConv((512, 512), 256, aggr=mean)\n",
       "      (term__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "      (conference__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "auto_classifier.gnn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0) - SimpleHeteroConv(\n",
       "  (rel_convs): ModuleDict(\n",
       "    (author__to__paper): SAGEConv((334, 4231), 512, aggr=mean)\n",
       "    (paper__to__author): SAGEConv((4231, 334), 512, aggr=mean)\n",
       "    (paper__to__term): SAGEConv((4231, 50), 512, aggr=mean)\n",
       "    (paper__to__conference): SAGEConv((4231, 1), 512, aggr=mean)\n",
       "    (term__to__paper): SAGEConv((50, 4231), 512, aggr=mean)\n",
       "    (conference__to__paper): SAGEConv((1, 4231), 512, aggr=mean)\n",
       "  )\n",
       "): x_dict, edge_index_dict -> x_dict\n",
       "  (1) - SimpleHeteroConv(\n",
       "  (rel_convs): ModuleDict(\n",
       "    (author__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    (paper__to__author): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    (paper__to__term): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    (paper__to__conference): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    (term__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "    (conference__to__paper): SAGEConv((512, 512), 256, aggr=mean)\n",
       "  )\n",
       "): x_dict, edge_index_dict -> x_dict\n",
       ")"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "simple_classifier.gnn"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
