{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wprowadzenie\n",
    "\n",
    "Ten notatnik przedstawia przykład wykorzystania sieci grafowych jako systemów rekomendacyjnych.\n",
    "\n",
    "W przykładach przedstawionych w niniejszej książce będziemy pracować ze zbiorem \"latest-small\", czyli najmniejszym dostępnym podzbiorem MovieLens. Wybór ten podyktowany jest względami wydajnościowymi – mniejszy zbiór danych pozwala na szybsze trenowanie modeli i eksperymentowanie z różnymi architekturami sieci, co jest szczególnie istotne w kontekście edukacyjnym. Pomimo mniejszego rozmiaru, zbiór ten zachowuje wszystkie kluczowe cechy większych wersji, umożliwiając efektywne zademonstrowanie technik rekomendacji opartych na grafowych sieciach neuronowych.\n",
    "MovieLens jest powszechnie stosowany z kilku powodów:\n",
    "1.\tStanowi stabilny punkt odniesienia (benchmark) dla badań nad systemami rekomendacyjnymi.\n",
    "2.\tZawiera bogaty zestaw danych o ocenach, użytkownikach i filmach, co umożliwia testowanie różnych podejść do rekomendacji.\n",
    "3.\tPozwala na budowanie i porównywanie różnych typów systemów rekomendacyjnych: opartych na wiedzy, zawartości i filtracji kolaboratywnej.\n",
    "4.\tUmożliwia modelowanie problemu rekomendacji jako heterogenicznego grafu, gdzie węzły reprezentują użytkowników i filmy, a krawędzie - oceny.\n",
    "Graf MovieLens jest grafem heterogenicznym o dwóch typach obiektów (film i użytkownik), oraz jednym typie relacji (użytkownik – ocenia – film). Tego rodzaju grafy noszą nazwę dwudzielnych. Dodatkowo posiada:\n",
    "1.\t9742 wierzchołki odpowiadające filmom, z których każdy jest opisany wektorem 404 atrybutów – są to wyniki kodowania tytułu jako wektora numerycznego (ang. text embedding) za pomocą modelu językowego typu sentence transformer, all-MiniLM-L6-v2. \n",
    "2.\t610 wierzchołków odpowiadających użytkownikom. Nie posiadają oni żadnych wektorów cech, dlatego dodano dla nich stałą = 1.\n",
    "3.\t100836 krawędzi odpowiadających ocenom wystawionym dla filmów, w skali 0.5-5. PyG w procesie wstępnego przetwarzania danych zaokrągla wartości w dół do pełnych liczb, dlatego przywrócono oryginalną wartość 0.5, dla najgorszych ocen.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import bibliotek"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch_geometric as pyg\n",
    "import torch as th\n",
    "import pytorch_lightning as pl\n",
    "import torchmetrics as tm\n",
    "import optuna as opt\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie danych"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAIN_REL = ('user', 'rates', 'movie')\n",
    "REV_REL = ('movie', 'rev_rates', 'user')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-warning\">\n",
    "<b>Uwaga</b> Zbiór danych MovieLens jest wstępnie przetwarzany wewnątrz biblioteki\n",
    "za pomocą tokenizerów z biblioteki `HuggingFace`. Przetwarzanie może zająć więcej czasu.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading https://files.grouplens.org/datasets/movielens/ml-latest-small.zip\n",
      "Extracting data/MovieLens/raw/ml-latest-small.zip\n",
      "Processing...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0753327bd8eb4f2aac2ea38a9c71d478",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "modules.json:   0%|          | 0.00/349 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7852f2854695453c9a301996c9b362b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config_sentence_transformers.json:   0%|          | 0.00/116 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4c526803c7844ab7833cbe5f20c6f255",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "414548c1d7c246959f05cb288254b9a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentence_bert_config.json:   0%|          | 0.00/53.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a071f0989a754891a6d8dcfbb530b386",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/612 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cd4626adf61f4aafa13c2f779c5d8e42",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/90.9M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c85e35e54b2a455b94d4954c2f1ebd95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/350 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9b619c47e064c86b4682257e15df79f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fdef03d99f794395be912bd7986caa27",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bb16f8280bd84dfc8a651441776ffd30",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/112 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "498f283b31534f4eb243a09411f90359",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/190 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9bf13043a724c10a7cd1ce50a8b285b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/305 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Done!\n"
     ]
    }
   ],
   "source": [
    "data = pyg.datasets.MovieLens(\"./data/MovieLens\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph = data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "HeteroData(\n",
       "  movie={ x=[9742, 404] },\n",
       "  user={ num_nodes=610 },\n",
       "  (user, rates, movie)={\n",
       "    edge_index=[2, 100836],\n",
       "    edge_label=[100836],\n",
       "    time=[100836],\n",
       "  }\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph['user'].x = th.ones((graph['user'].num_nodes, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Zamiana grafu na dwukierunkowy.\n",
    "graph = pyg.transforms.to_undirected.ToUndirected()(graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "graph[MAIN_REL].edge_label = graph[MAIN_REL].edge_label.to(th.float)\n",
    "graph[MAIN_REL].edge_label[graph[MAIN_REL].edge_label == 0.0] = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Podział na zbiór treningowy i testowy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter = pyg.transforms.RandomLinkSplit(\n",
    "        num_val=0.1,\n",
    "        num_test=0.1,\n",
    "        neg_sampling_ratio=0.0,\n",
    "        add_negative_train_samples=False,\n",
    "        is_undirected=True,\n",
    "        key='edge_label',\n",
    "        edge_types=[('user', 'rates', 'movie')],\n",
    "        rev_edge_types=[('movie', 'rev_rates', 'user')],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, val, test = splitter(graph)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = pyg.loader.LinkNeighborLoader(\n",
    "    train,\n",
    "    neg_sampling_ratio=0,\n",
    "    batch_size=2048,\n",
    "    num_neighbors={MAIN_REL: [-1, -1], REV_REL: [-1, -1]},\n",
    "    edge_label_index=(MAIN_REL, train[MAIN_REL].edge_label_index),\n",
    "    edge_label=train[MAIN_REL].edge_label,\n",
    "    shuffle=True,\n",
    ")\n",
    "\n",
    "val_loader = pyg.loader.DataLoader(\n",
    "    [val],\n",
    "    batch_size=1)\n",
    "\n",
    "test_loader = pyg.loader.DataLoader(\n",
    "    [test], batch_size=1)\n",
    "\n",
    "\n",
    "\n",
    "train_b = next(iter(train_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Przygotowanie modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GraphNN(th.nn.Module):\n",
    "    \"\"\"Heterogeniczna sieć grafowa, budująca osadzenia wierzchołków.\"\"\"\n",
    "    def __init__(self, num_layers: int, hidden_channels: int, **kwargs):\n",
    "        super().__init__()\n",
    "        layers = []\n",
    "        dim_in = (-1, -1)\n",
    "        for _ in range(num_layers):\n",
    "            conv = pyg.nn.SAGEConv(dim_in, hidden_channels, **kwargs)\n",
    "            layers.append(conv)\n",
    "            dim_in = (hidden_channels, hidden_channels)\n",
    "        self.convs = th.nn.ModuleList(layers)\n",
    "\n",
    "    def forward(self, x, edge_index):\n",
    "        for i, conv in enumerate(self.convs):\n",
    "            x = conv(x, edge_index)\n",
    "            if i != len(self.convs) - 1:\n",
    "                x = x.relu()\n",
    "        return x\n",
    "    \n",
    "\n",
    "class Recommender(th.nn.Module):\n",
    "    \"\"\"Hybrydowy system rekomendacyjny - jego głównym składnikiem jest heterogeniczna sieć\n",
    "    grafowa, budująca osadzenia wierzchołków, oraz kilka gęsto połączonych warstw\n",
    "    sieci neuronowej, które przewidują ocenę użytkownika dla filmu.\"\"\"\n",
    "    def __init__(\n",
    "            self, \n",
    "            gnn_layers: int, \n",
    "            gnn_dim: int, \n",
    "            linkpred_dim: int,\n",
    "            metadata: tuple,\n",
    "            **kwargs):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Sieć grafowa\n",
    "        encoder = GraphNN(gnn_layers, gnn_dim, **kwargs)\n",
    "        self.encoder = pyg.nn.to_hetero(encoder, metadata)\n",
    "        # Warstwy przewidujące ocenę\n",
    "        self.decoder = self._build_linkpred(gnn_dim, linkpred_dim)\n",
    "\n",
    "    def _build_linkpred(self, gnn_dim, linkpred_dim):\n",
    "        linkpred_layers = []\n",
    "        dim_in = gnn_dim * 2\n",
    "        for link_dim in linkpred_dim:\n",
    "            linkpred_layers.append(th.nn.Linear(dim_in, link_dim))\n",
    "            linkpred_layers.append(th.nn.ReLU())\n",
    "            dim_in = link_dim\n",
    "        linkpred_layers.append(th.nn.Linear(dim_in, 1))\n",
    "        return th.nn.Sequential(*linkpred_layers)\n",
    "\n",
    "    def forward(self, data: pyg.data.HeteroData, supervised: bool = True):\n",
    "        # Krok 1: zbudowanie numerycznej reprezentacji wierzchołków\n",
    "        enc = self.encoder(data.x_dict, data.edge_index_dict)\n",
    "        # Działanie w trybie treningu lub rekomendacji\n",
    "        if supervised:\n",
    "            index = data[MAIN_REL].edge_label_index\n",
    "        else:\n",
    "            index = data[MAIN_REL].edge_index\n",
    "\n",
    "        # Pobranie dwóch wektorów reprezentacji:\n",
    "        # użytkownika oraz filmu\n",
    "        src_enc = enc['user'][index[0, :]]\n",
    "        dst_enc = enc['movie'][index[1, :]]\n",
    "        z = th.cat([src_enc, dst_enc], dim=-1)\n",
    "        # Przetworzenie przez sieć przewidującą ocenę\n",
    "        logits = self.decoder(z).clamp(min=0.0, max=5.0)\n",
    "        return logits.squeeze()\n",
    "\n",
    "class RecommenderLightning(pl.LightningModule):\n",
    "    def __init__(self, recommender, lr=1e-3):\n",
    "        super().__init__()\n",
    "        self.recommender = recommender\n",
    "        self.lr = lr\n",
    "        self.criterion = th.nn.MSELoss()\n",
    "\n",
    "    def forward(self, data, supervised=True):\n",
    "        return self.recommender(data, supervised)\n",
    "\n",
    "    def _step(self, batch, prefix):\n",
    "        logits = self(batch, supervised=True)\n",
    "        \n",
    "        # RMSE\n",
    "        loss = self.criterion(logits, batch[MAIN_REL].edge_label).sqrt()\n",
    "        self.log(f'{prefix}_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=batch[MAIN_REL].edge_label.size(0))\n",
    "\n",
    "        # Kalkulacja dodatkowych metryk\n",
    "        mse = tm.functional.mean_squared_error(logits, batch[MAIN_REL].edge_label)\n",
    "        mape = tm.functional.mean_absolute_percentage_error(logits, batch[MAIN_REL].edge_label)\n",
    "        rmse = tm.functional.mean_squared_error(logits, batch[MAIN_REL].edge_label).sqrt()\n",
    "\n",
    "        self.log(f'{prefix}_mse', mse, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=batch[MAIN_REL].edge_label.size(0))\n",
    "        self.log(f'{prefix}_mape', mape, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=batch[MAIN_REL].edge_label.size(0))\n",
    "        self.log(f'{prefix}_rmse', rmse, on_step=True, on_epoch=True, prog_bar=True, logger=True, batch_size=batch[MAIN_REL].edge_label.size(0))\n",
    "        return loss\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        return self._step(batch, 'train')\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        return self._step(batch, 'val')\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        return self._step(batch, 'test')\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return th.optim.Adam(self.parameters(), lr=self.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "gnn = GraphNN(num_layers=2, hidden_channels=64, normalize=True, aggr=['sum', 'mean'])\n",
    "gnn_hetero = pyg.nn.to_hetero(gnn, graph.metadata(), aggr='sum')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "with th.no_grad():\n",
    "    out = gnn_hetero(train.x_dict, train.edge_index_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10083])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    recommender = Recommender(gnn_layers=2, gnn_dim=64, linkpred_dim=[64, 64], metadata=graph.metadata(), normalize=True)\n",
    "    out = recommender(val, supervised=True)\n",
    "\n",
    "out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/pytorch_lightning/core/module.py:449: You are trying to `self.log()` but the `self.trainer` reference is not registered on the model yet. This is most likely because the model hasn't been passed to the `Trainer`\n"
     ]
    }
   ],
   "source": [
    "with th.no_grad():\n",
    "    model = RecommenderLightning(recommender)\n",
    "    out = model.training_step(train_b, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optymalizacja hiperparametrów"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def params2model(best_params: dict, graph: pyg.data.HeteroData):\n",
    "    recommender = Recommender(\n",
    "        gnn_layers=best_params['gnn_layers'], \n",
    "        gnn_dim=best_params['gnn_dim'], \n",
    "        linkpred_dim=[best_params[f'linkpred_dim_{i}'] for i in range(best_params['linkpred_num_layers'])], \n",
    "        metadata=graph.metadata(), \n",
    "        normalize=best_params['normalize'], \n",
    "        aggr=best_params['aggr'] if isinstance(best_params['aggr'], list) else [best_params['aggr']]\n",
    "    )\n",
    "    return recommender\n",
    "\n",
    "def objective(trial):\n",
    "    # Suggest hyperparameters\n",
    "    gnn_layers = trial.suggest_int('gnn_layers', 2, 4)\n",
    "    gnn_dim = trial.suggest_categorical('gnn_dim', [16, 32, 64])\n",
    "    linkpred_num_layers = trial.suggest_int('linkpred_num_layers', 1, 3)\n",
    "    linkpred_dim = [trial.suggest_categorical(f'linkpred_dim_{i}', [16, 32, 64]) for i in range(linkpred_num_layers)]\n",
    "    aggr = trial.suggest_categorical('aggr', [('sum'), ('mean'), ('sum', 'mean')])\n",
    "    normalize = trial.suggest_categorical('normalize', [True, False])\n",
    "    \n",
    "    pyg.seed_everything(123)\n",
    "    # Create the recommender model\n",
    "    recommender = Recommender(\n",
    "        gnn_layers=gnn_layers,\n",
    "        gnn_dim=gnn_dim,\n",
    "        linkpred_dim=linkpred_dim,\n",
    "        metadata=graph.metadata(),\n",
    "        normalize=normalize,\n",
    "        aggr=aggr\n",
    "    )\n",
    "    recommender_lit = RecommenderLightning(recommender, lr=0.01)\n",
    "\n",
    "    # Early stopping callback\n",
    "    early_stopping = pl.callbacks.EarlyStopping(\n",
    "        monitor='val_rmse_epoch',\n",
    "        patience=2,\n",
    "        mode='min'\n",
    "    )\n",
    "\n",
    "    # Trainer\n",
    "    trainer = pl.Trainer(max_epochs=4, callbacks=[early_stopping], deterministic=True, enable_progress_bar=False, enable_checkpointing=False)\n",
    "\n",
    "    # Train the model\n",
    "    trainer.fit(recommender_lit, train_loader, val_loader)\n",
    "\n",
    "    # Validate the model\n",
    "    val_scores = trainer.validate(recommender_lit, val_loader)\n",
    "    val_rmse = val_scores[0]['val_rmse_epoch']\n",
    "\n",
    "    return val_rmse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimize = False\n",
    "\n",
    "if optimize:\n",
    "    # Create a study and optimize the objective function\n",
    "    study = opt.create_study(direction='minimize')\n",
    "    study.optimize(objective, n_trials=50)\n",
    "    best_params = study.best_params\n",
    "    json.dump(best_params, open('best_params_recommender.json', 'w'))\n",
    "else:\n",
    "    best_params = json.load(open('best_params_recommender.json', 'r'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'gnn_layers': 2,\n",
       " 'gnn_dim': 32,\n",
       " 'linkpred_num_layers': 1,\n",
       " 'linkpred_dim_0': 32,\n",
       " 'aggr': ['sum'],\n",
       " 'normalize': True}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "best_params"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trenowanie modelu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyg.seed_everything(123)\n",
    "recommender = params2model(best_params, graph)\n",
    "recommender_lit = RecommenderLightning(recommender, lr=0.01)\n",
    "early_stopping = pl.callbacks.EarlyStopping(\n",
    "    monitor='val_rmse_epoch',\n",
    "    patience=5,\n",
    "    mode='min'\n",
    ")\n",
    "\n",
    "checkpoint_callback = pl.callbacks.ModelCheckpoint(\n",
    "    monitor='val_rmse_epoch',\n",
    "    mode='min',\n",
    "    save_top_k=1,\n",
    "    filename='best-checkpoint'\n",
    ")\n",
    "\n",
    "trainer = pl.Trainer(max_epochs=50, callbacks=[early_stopping, checkpoint_callback], deterministic=True)\n",
    "\n",
    "val_scores_before = trainer.validate(recommender_lit, val_loader)\n",
    "\n",
    "print(\"Val scores before\")\n",
    "print(val_scores_before)\n",
    "\n",
    "trainer.fit(recommender_lit, train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = RecommenderLightning.load_from_checkpoint(checkpoint_callback.best_model_path, recommender=recommender)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b4929f946f048f2af03b432d314af22",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "     Validate metric           DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "     val_loss_epoch         0.8918048739433289\n",
      "     val_mape_epoch         0.31205517053604126\n",
      "      val_mse_epoch         0.7953158617019653\n",
      "     val_rmse_epoch         0.8918048739433289\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    }
   ],
   "source": [
    "val_scores = trainer.validate(best_model, val_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "/home/filip/projects/grafowe_sieci_neuronowe/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/data_connector.py:433: The 'test_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=11` in the `DataLoader` to improve performance.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "68851a1b0384499881e95d7ae6121eb9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "     test_loss_epoch        0.9039353132247925\n",
      "     test_mape_epoch        0.32364144921302795\n",
      "     test_mse_epoch         0.8170990347862244\n",
      "     test_rmse_epoch        0.9039353132247925\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    }
   ],
   "source": [
    "test_scores = trainer.test(best_model, test_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Rekomendacja elementów"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "def recommend_items(\n",
    "    user_idx: int, \n",
    "    recommender: Recommender, \n",
    "    graph: pyg.data.HeteroData,\n",
    "    top_k: int = 5):\n",
    "    # Indeksy ocenionych dotąd filmów\n",
    "    rated_movies_idx = graph[MAIN_REL].edge_index[\n",
    "        1, \n",
    "        graph[MAIN_REL].edge_index[0] == user_idx]\n",
    "    # Indeksy wszystkich filmów\n",
    "    all_movies_idx = th.arange(graph['movie'].num_nodes)\n",
    "    # Indeksy nieocenionych dotąd filmów\n",
    "    unrated_movies_idx = all_movies_idx[~th.isin(all_movies_idx, rated_movies_idx)]\n",
    "\n",
    "    # Zbudowanie grafu heterogenicznego dla użytkownika i nieocenionych filmów\n",
    "    grading_data = pyg.data.HeteroData()\n",
    "    grading_data['user'].x = graph['user'].x\n",
    "    grading_data['movie'].x = graph['movie'].x\n",
    "    grading_data[MAIN_REL].edge_index = th.tensor([\n",
    "                [user_idx] * unrated_movies_idx.size(0),\n",
    "                unrated_movies_idx\n",
    "            ])\n",
    "    grading_data[REV_REL].edge_index = grading_data[MAIN_REL].edge_index.flip(0)\n",
    "    recommender.eval()\n",
    "    # Wykonanie predykcji\n",
    "    with th.no_grad():\n",
    "        predictions = recommender(grading_data, supervised=False)\n",
    "    top_k_idx = th.argsort(predictions, descending=True)[:top_k]\n",
    "    return top_k_idx, predictions[top_k_idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([7713, 6595, 9340, 4843, 2147]),\n",
       " tensor([4.1309, 4.1240, 4.1206, 4.0989, 4.0955]))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id, score = recommend_items(0, best_model, graph, top_k=5)\n",
    "id, score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}
