gnn/Events2018.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b1e0fc77-53ff-419d-b9b6-d2fe8457b0cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Device (for PyTorch): cpu\n"
     ]
    }
   ],
   "source": [
    "from bidict import bidict\n",
    "from collections import defaultdict\n",
    "from dataclasses import dataclass\n",
    "from datetime import datetime, timedelta\n",
    "import itertools\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from franz.openrdf.connect import ag_connect\n",
    "from franz.openrdf.query.query import QueryLanguage\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch_geometric.data import DataLoader\n",
    "from torch_geometric.io import read_txt_array\n",
    "from torch_geometric.nn.models.re_net import RENet\n",
    "from torch_geometric.datasets.icews import EventDataset\n",
    "from tqdm import tqdm\n",
    "import shutil\n",
    "from pathlib import Path\n",
    "\n",
    "datapath = Path(\".\").joinpath(\"data\")\n",
    "datapath.mkdir(exist_ok=True)\n",
    "\n",
    "@dataclass\n",
    "class _Resource(object):\n",
    "    idx: int\n",
    "    uri: str   \n",
    "    def __hash__(self): return hash(self.uri)    \n",
    "    def __repr__(self): return self.uri\n",
    "        \n",
    "class Entity(_Resource): pass\n",
    "class Relation(_Resource): pass\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print(\"Device (for PyTorch):\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cadfe9c0-e649-44a8-929a-d31a3549c13d",
   "metadata": {},
   "source": [
    "## Construct a Dataframe\n",
    "\n",
    "We will construct a dataframe where each row is represented as `(subject, predicate, object, timestamp)`.\n",
    "\n",
    "The dataframe will be sorted by `timestamp`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e010f9b6-99a6-456b-9ee3-cf929f4899da",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 468558/468558 [01:25<00:00, 5465.97it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 38s, sys: 738 ms, total: 1min 39s\n",
      "Wall time: 2min 6s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>s</th>\n",
       "      <th>p</th>\n",
       "      <th>o</th>\n",
       "      <th>t</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>http://franz.com/examples/Opposition%20Support...</td>\n",
       "      <td>http://franz.com/examples/Property#Accuse</td>\n",
       "      <td>http://franz.com/examples/Citizen%20(Russia)</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>http://franz.com/examples/Government%20(Ukraine)</td>\n",
       "      <td>http://franz.com/examples/Property#Make%20stat...</td>\n",
       "      <td>http://franz.com/examples/Military%20(Ukraine)</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>http://franz.com/examples/Armed%20Rebel%20(Ukr...</td>\n",
       "      <td>http://franz.com/examples/Property#Accuse</td>\n",
       "      <td>http://franz.com/examples/Military%20(Ukraine)</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>http://franz.com/examples/Militia%20(Ukraine)</td>\n",
       "      <td>http://franz.com/examples/Property#Make%20stat...</td>\n",
       "      <td>http://franz.com/examples/Military%20(Ukraine)</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>http://franz.com/examples/Philippine%20Nationa...</td>\n",
       "      <td>http://franz.com/examples/Property#Rally%20opp...</td>\n",
       "      <td>http://franz.com/examples/Military%20Personnel...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>468553</th>\n",
       "      <td>http://franz.com/examples/Opposition%20Support...</td>\n",
       "      <td>http://franz.com/examples/Property#Accuse</td>\n",
       "      <td>http://franz.com/examples/Lawmaker%20(India)</td>\n",
       "      <td>303</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>468554</th>\n",
       "      <td>http://franz.com/examples/Ant%C3%B3nio%20Manue...</td>\n",
       "      <td>http://franz.com/examples/Property#Make%20stat...</td>\n",
       "      <td>http://franz.com/examples/Foreign%20Affairs%20...</td>\n",
       "      <td>303</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>468555</th>\n",
       "      <td>http://franz.com/examples/Staffan%20de%20Mistura</td>\n",
       "      <td>http://franz.com/examples/Property#Make%20stat...</td>\n",
       "      <td>http://franz.com/examples/Foreign%20Affairs%20...</td>\n",
       "      <td>303</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>468556</th>\n",
       "      <td>http://franz.com/examples/Police%20(India)</td>\n",
       "      <td>http://franz.com/examples/Property#Arrest,%20d...</td>\n",
       "      <td>http://franz.com/examples/Saddam%20Hussein</td>\n",
       "      <td>303</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>468557</th>\n",
       "      <td>http://franz.com/examples/Elmar%20Mammadyarov</td>\n",
       "      <td>http://franz.com/examples/Property#Consult</td>\n",
       "      <td>http://franz.com/examples/Ministry%20of%20Educ...</td>\n",
       "      <td>303</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>468558 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                        s  \\\n",
       "0       http://franz.com/examples/Opposition%20Support...   \n",
       "1        http://franz.com/examples/Government%20(Ukraine)   \n",
       "2       http://franz.com/examples/Armed%20Rebel%20(Ukr...   \n",
       "3           http://franz.com/examples/Militia%20(Ukraine)   \n",
       "4       http://franz.com/examples/Philippine%20Nationa...   \n",
       "...                                                   ...   \n",
       "468553  http://franz.com/examples/Opposition%20Support...   \n",
       "468554  http://franz.com/examples/Ant%C3%B3nio%20Manue...   \n",
       "468555   http://franz.com/examples/Staffan%20de%20Mistura   \n",
       "468556         http://franz.com/examples/Police%20(India)   \n",
       "468557      http://franz.com/examples/Elmar%20Mammadyarov   \n",
       "\n",
       "                                                        p  \\\n",
       "0               http://franz.com/examples/Property#Accuse   \n",
       "1       http://franz.com/examples/Property#Make%20stat...   \n",
       "2               http://franz.com/examples/Property#Accuse   \n",
       "3       http://franz.com/examples/Property#Make%20stat...   \n",
       "4       http://franz.com/examples/Property#Rally%20opp...   \n",
       "...                                                   ...   \n",
       "468553          http://franz.com/examples/Property#Accuse   \n",
       "468554  http://franz.com/examples/Property#Make%20stat...   \n",
       "468555  http://franz.com/examples/Property#Make%20stat...   \n",
       "468556  http://franz.com/examples/Property#Arrest,%20d...   \n",
       "468557         http://franz.com/examples/Property#Consult   \n",
       "\n",
       "                                                        o    t  \n",
       "0            http://franz.com/examples/Citizen%20(Russia)    0  \n",
       "1          http://franz.com/examples/Military%20(Ukraine)    0  \n",
       "2          http://franz.com/examples/Military%20(Ukraine)    0  \n",
       "3          http://franz.com/examples/Military%20(Ukraine)    0  \n",
       "4       http://franz.com/examples/Military%20Personnel...    0  \n",
       "...                                                   ...  ...  \n",
       "468553       http://franz.com/examples/Lawmaker%20(India)  303  \n",
       "468554  http://franz.com/examples/Foreign%20Affairs%20...  303  \n",
       "468555  http://franz.com/examples/Foreign%20Affairs%20...  303  \n",
       "468556         http://franz.com/examples/Saddam%20Hussein  303  \n",
       "468557  http://franz.com/examples/Ministry%20of%20Educ...  303  \n",
       "\n",
       "[468558 rows x 4 columns]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "q = \"\"\"SELECT DISTINCT ?s ?s_id ?p ?p_id ?o ?o_id ?t {\n",
    "  ?event a <http://franz.com/Relation> ;\n",
    "           <http://franz.com/date> ?t ;\n",
    "           <http://franz.com/from> ?s ;\n",
    "           <http://franz.com/predicate> ?p ;\n",
    "           <http://franz.com/ID> ?p_id ;\n",
    "           <http://franz.com/to> ?o .\n",
    "  ?s <http://franz.com/ID> ?s_id .\n",
    "  ?o <http://franz.com/ID> ?o_id .\n",
    "}\n",
    "ORDER BY ?t\"\"\"\n",
    "\n",
    "data = defaultdict(list)\n",
    "with ag_connect('Events2018', host=\"localhost\", port=10000, catalog=\"demos\") as conn:\n",
    "    with conn.prepareTupleQuery(QueryLanguage.SPARQL, q).evaluate() as res:\n",
    "        for bs in tqdm(res):\n",
    "            data[\"s\"].append(Entity(idx=bs.getValue(\"s_id\").intValue(), uri=bs.getValue(\"s\").value))\n",
    "            data[\"p\"].append(Relation(idx=bs.getValue(\"p_id\").intValue(), uri=bs.getValue(\"p\").value))\n",
    "            data[\"o\"].append(Entity(idx=bs.getValue(\"o_id\").intValue(), uri=bs.getValue(\"o\").value))\n",
    "            data[\"t\"].append(bs.getValue(\"t\").toPython())\n",
    "df = pd.DataFrame(data=data)\n",
    "start_date = df[\"t\"][0]\n",
    "df.t = df.t.apply(lambda d: (d-start_date).days)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be3b5fcf-74d3-4d74-8f2b-f813f2de15bc",
   "metadata": {},
   "source": [
    "## Construct entity2id and relation2id\n",
    "\n",
    "To reference all entities and relations by ids (integers), we construct dictionaries for both of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7a2cd179-0b4d-4c33-a898-e83146159bd2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "processing entity2id: 937116it [00:04, 211297.96it/s]\n",
      "processing relation2id: 100%|██████████| 468558/468558 [00:02<00:00, 218576.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "23033 Entities, 256 Relations, 468558 Events in total\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "entity2id, relation2id = bidict(), bidict()\n",
    "\n",
    "for ent in tqdm(itertools.chain(df.s, df.o), desc=\"processing entity2id\"):\n",
    "    entity2id[ent] = ent.idx\n",
    "        \n",
    "for rel in tqdm(df.p, desc=\"processing relation2id\"):\n",
    "    relation2id[rel] = rel.idx\n",
    "\n",
    "print(f\"{len(entity2id)} Entities, {len(relation2id)} Relations, {len(df)} Events in total\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b634835-9052-4329-ae70-c71f1fc6576e",
   "metadata": {},
   "source": [
    "## Spliting training, validation and test dataset\n",
    "\n",
    "* train - 80%\n",
    "* test - 10%\n",
    "* validation - 10 Sample Events of Donald Trump only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "e3ac43ba-f455-43a8-bc5d-16983a64b74f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(374846, 46856, 46856)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train, valid, test = np.split(df, [int(.8*len(df)), int(.9*len(df))])\n",
    "len(train), len(valid), len(test)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64d85080-b44d-463e-b0fe-78bd6512e086",
   "metadata": {},
   "source": [
    "## Customising our Events2018 dataset class\n",
    "\n",
    "By using torch_geometric's [Dataset API](https://pytorch-geometric.readthedocs.io/en/latest/modules/data.html#torch_geometric.data.Dataset), we tell torch_geometric:\n",
    "\n",
    "1. How to prepare raw data (see `download` method)\n",
    "2. How to process raw data (see `process` method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a546ed6-96ac-44fc-babe-c4b728c3d298",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Events2018(EventDataset):\n",
    "    splits = [0, len(train), len(train)+len(valid), len(train)+len(valid)+len(test)]  # Train/Val/Test splits.\n",
    "\n",
    "    def __init__(self, root, split='train', transform=None, pre_transform=None, pre_filter=None):\n",
    "        assert split in ['train', 'val', 'test']\n",
    "        super(Events2018, self).__init__(root, transform, pre_transform, pre_filter)\n",
    "        idx = self.processed_file_names.index('{}.pt'.format(split))\n",
    "        self.data, self.slices = torch.load(self.processed_paths[idx])\n",
    "\n",
    "    @property\n",
    "    def num_nodes(self):\n",
    "        return len(entity2id)\n",
    "\n",
    "    @property\n",
    "    def num_rels(self):\n",
    "        return len(relation2id)\n",
    "\n",
    "    @property\n",
    "    def raw_file_names(self):\n",
    "        return ['{}.txt'.format(name) for name in ['train', 'valid', 'test']]\n",
    "\n",
    "    @property\n",
    "    def processed_file_names(self):\n",
    "        return ['train.pt', 'val.pt', 'test.pt']\n",
    "    \n",
    "    def download(self):\n",
    "        d = { \"train.txt\": train, \"valid.txt\": valid, \"test.txt\": test }        \n",
    "        for filename in self.raw_file_names:            \n",
    "            with datapath.joinpath(\"raw\", filename).open(\"w\") as fd:\n",
    "                buffer = []\n",
    "                for row in tqdm(d[filename].itertuples(), desc=f\"writing {filename}\"):\n",
    "                    buffer.append(f\"{entity2id[row.s]}\\t{relation2id[row.p]}\\t{entity2id[row.o]}\\t{row.t}\\n\")\n",
    "                    if len(buffer) == 4096:\n",
    "                        fd.writelines(buffer)\n",
    "                        buffer = []\n",
    "                if len(buffer) != 0:\n",
    "                    fd.writelines(buffer)                    \n",
    "\n",
    "    def process_events(self):\n",
    "        events = []\n",
    "        for path in self.raw_paths:\n",
    "            events += [read_txt_array(path, sep='\\t', end=4, dtype=torch.long)]\n",
    "        return torch.cat(events, dim=0)\n",
    "\n",
    "    def process(self):\n",
    "        s = self.splits\n",
    "        data_list = super(Events2018, self).process()\n",
    "        torch.save(self.collate(data_list[s[0]:s[1]]), self.processed_paths[0])\n",
    "        torch.save(self.collate(data_list[s[1]:s[2]]), self.processed_paths[1])\n",
    "        torch.save(self.collate(data_list[s[2]:s[3]]), self.processed_paths[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff4fcf90-fdd4-4a8b-b899-5cc9e665e698",
   "metadata": {},
   "source": [
    "## Create dataloader for training and test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "549c1245-3267-4f52-bc79-707580b3ca7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Processing...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:The `pre_transform` argument differs from the one used in the pre-processed version of this dataset. If you really want to make use of another pre-processing technique, make sure to delete `data/processed` first.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done!\n",
      "CPU times: user 2min 39s, sys: 728 ms, total: 2min 40s\n",
      "Wall time: 2min 40s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "seq_len = 10 # how many historical events to look back\n",
    "if datapath.joinpath(\"processed\").exists(): shutil.rmtree(datapath.joinpath(\"processed\"))\n",
    "train_dataset = Events2018(datapath, pre_transform=RENet.pre_transform(seq_len))\n",
    "test_dataset = Events2018(datapath, split='test')\n",
    "\n",
    "train_loader = DataLoader(\n",
    "    train_dataset,\n",
    "    batch_size=1024,\n",
    "    shuffle=True,\n",
    "    follow_batch=['h_sub', 'h_obj'],\n",
    "    num_workers=6)\n",
    "\n",
    "test_loader = DataLoader(\n",
    "    test_dataset,\n",
    "    batch_size=1024,\n",
    "    shuffle=False,\n",
    "    follow_batch=['h_sub', 'h_obj'],\n",
    "    num_workers=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "684acc91-51ef-4da4-973d-ed99b9a66ec7",
   "metadata": {},
   "source": [
    "## Define the RENet model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d926c5a0-1e07-49a1-a386-c7d1c3cc1e5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = RENet(\n",
    "    train_dataset.num_nodes,\n",
    "    train_dataset.num_rels,\n",
    "    hidden_channels=200,\n",
    "    seq_len=seq_len,\n",
    "    dropout=0.5,\n",
    ").to(device)\n",
    "\n",
    "optimizer = torch.optim.Adam(\n",
    "    model.parameters(), lr=0.001, weight_decay=0.00001)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61586463-6b86-4b7f-a8f4-873d3c5ce508",
   "metadata": {},
   "source": [
    "## Train the model and save it to disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f5a92883-061c-4e79-a576-f0a91bfcb298",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01, MRR: 0.1946, Hits@1: 0.1228, Hits@3: 0.2158, Hits@10: 0.3353\n",
      "Epoch: 02, MRR: 0.2381, Hits@1: 0.1530, Hits@3: 0.2687, Hits@10: 0.4043\n",
      "Epoch: 03, MRR: 0.2567, Hits@1: 0.1662, Hits@3: 0.2899, Hits@10: 0.4331\n",
      "Epoch: 04, MRR: 0.2667, Hits@1: 0.1736, Hits@3: 0.3021, Hits@10: 0.4485\n",
      "Epoch: 05, MRR: 0.2726, Hits@1: 0.1778, Hits@3: 0.3100, Hits@10: 0.4575\n",
      "Epoch: 06, MRR: 0.2763, Hits@1: 0.1806, Hits@3: 0.3138, Hits@10: 0.4623\n",
      "Epoch: 07, MRR: 0.2788, Hits@1: 0.1831, Hits@3: 0.3164, Hits@10: 0.4647\n",
      "Epoch: 08, MRR: 0.2794, Hits@1: 0.1826, Hits@3: 0.3177, Hits@10: 0.4674\n",
      "Epoch: 09, MRR: 0.2818, Hits@1: 0.1855, Hits@3: 0.3205, Hits@10: 0.4695\n",
      "Epoch: 10, MRR: 0.2820, Hits@1: 0.1853, Hits@3: 0.3208, Hits@10: 0.4708\n",
      "Epoch: 11, MRR: 0.2820, Hits@1: 0.1844, Hits@3: 0.3226, Hits@10: 0.4702\n",
      "Epoch: 12, MRR: 0.2823, Hits@1: 0.1853, Hits@3: 0.3220, Hits@10: 0.4715\n",
      "Epoch: 13, MRR: 0.2827, Hits@1: 0.1857, Hits@3: 0.3224, Hits@10: 0.4717\n",
      "Epoch: 14, MRR: 0.2827, Hits@1: 0.1857, Hits@3: 0.3226, Hits@10: 0.4721\n",
      "Epoch: 15, MRR: 0.2827, Hits@1: 0.1855, Hits@3: 0.3228, Hits@10: 0.4712\n",
      "Epoch: 16, MRR: 0.2830, Hits@1: 0.1860, Hits@3: 0.3226, Hits@10: 0.4718\n",
      "Epoch: 17, MRR: 0.2836, Hits@1: 0.1867, Hits@3: 0.3232, Hits@10: 0.4716\n",
      "Epoch: 18, MRR: 0.2832, Hits@1: 0.1862, Hits@3: 0.3222, Hits@10: 0.4710\n",
      "Epoch: 19, MRR: 0.2821, Hits@1: 0.1854, Hits@3: 0.3218, Hits@10: 0.4700\n",
      "Epoch: 20, MRR: 0.2834, Hits@1: 0.1863, Hits@3: 0.3238, Hits@10: 0.4709\n",
      "CPU times: user 12h 53min 3s, sys: 49min 49s, total: 13h 42min 52s\n",
      "Wall time: 3h 55min 21s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "def train():\n",
    "    model.train()\n",
    "    # Train model via multi-class classification against the corresponding\n",
    "    # object and subject entities.\n",
    "    for data in train_loader:\n",
    "        data = data.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        log_prob_obj, log_prob_sub = model(data)\n",
    "        loss_obj = F.nll_loss(log_prob_obj, data.obj)\n",
    "        loss_sub = F.nll_loss(log_prob_sub, data.sub)\n",
    "        loss = loss_obj + loss_sub\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "\n",
    "def test(loader):\n",
    "    model.eval()\n",
    "    # Compute Mean Reciprocal Rank (MRR) and Hits@1/3/10.\n",
    "    result = torch.tensor([0, 0, 0, 0], dtype=torch.float)\n",
    "    for data in loader:\n",
    "        data = data.to(device)\n",
    "        with torch.no_grad():\n",
    "            log_prob_obj, log_prob_sub = model(data)\n",
    "        result += model.test(log_prob_obj, data.obj) * data.obj.size(0)\n",
    "        result += model.test(log_prob_sub, data.sub) * data.sub.size(0)\n",
    "    result = result / (2 * len(loader.dataset))\n",
    "    return result.tolist()\n",
    "\n",
    "for epoch in range(1, 21):\n",
    "    train()\n",
    "    results = test(test_loader)\n",
    "    print('Epoch: {:02d}, MRR: {:.4f}, Hits@1: {:.4f}, Hits@3: {:.4f}, '\n",
    "          'Hits@10: {:.4f}'.format(epoch, *results))\n",
    "\n",
    "torch.save(model.state_dict(), datapath.joinpath(\"model.pt\")) # save model states"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}