notebooks/0_custom_KG_dataset.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Using BESS-KGE with your Own Data\n",
                "\n",
                "<em>Copyright (c) 2023 Graphcore Ltd. All rights reserved.</em>\n",
                "\n",
                "BESS-KGE (`besskge`) is a PyTorch library for knowledge graph embedding (KGE) models on IPUs implementing the distribution framework [BESS](https://arxiv.org/abs/2211.12281), with embedding tables stored in the IPU SRAM.\n",
                "\n",
                "In this notebook we will show how to use the `besskge.dataset.KGDataset` class to easily preprocess a custom knowledge graph dataset for use with BESS-KGE.\n",
                "\n",
                "As an example, we will download and build the [OGBL-BioKG](https://ogb.stanford.edu/docs/linkprop/#ogbl-biokg) biomedical knowledge graph. While BESS-KGE provides a built-in dataloader for this dataset (see [besskge.dataset.KGDataset.build_ogbl_biokg](https://graphcore-research.github.io/bess-kge/generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_ogbl_biokg)), in this notebook — for didactic purposes — we will show how to import the dataset from scratch."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Dependencies"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We recommend that you install `besskge` directly from the GitHub sources:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Found existing installation: besskge 0.1\n",
                        "Uninstalling besskge-0.1:\n",
                        "  Successfully uninstalled besskge-0.1\n"
                    ]
                }
            ],
            "source": [
                "import sys\n",
                "!{sys.executable} -m pip uninstall -y besskge\n",
                "!pip install -q git+https://github.com/graphcore-research/bess-kge.git\n",
                "!pip install -q jupyter ipywidgets"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Next, import the necessary dependencies:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os\n",
                "from pathlib import Path\n",
                "\n",
                "import ogb\n",
                "import pandas as pd\n",
                "import torch\n",
                "from tqdm.notebook import tqdm\n",
                "\n",
                "from besskge.dataset import KGDataset\n",
                "\n",
                "dataset_directory = os.getenv(\"DATASET_DIR\", \"../datasets/\") + \"/biokg/\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## The dataset\n",
                "\n",
                "We download the OGBL-BioKG knowledge graph using the `ogb` package (see the OGB [description of the data loader](https://ogb.stanford.edu/docs/linkprop/#data-loader) for details on how to use it). It shouldn't take more than a couple of minutes for the dataset to download."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Downloading http://snap.stanford.edu/ogb/data/linkproppred/biokg.zip\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "Downloaded 0.90 GB: 100%|██████████| 920/920 [01:18<00:00, 11.67it/s]\n"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Extracting ../datasets//biokg_test/biokg.zip\n",
                        "Loading necessary files...\n",
                        "This might take a while.\n",
                        "Processing graphs...\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "100%|██████████| 1/1 [00:00<00:00, 4860.14it/s]"
                    ]
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Saving...\n"
                    ]
                },
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "\n"
                    ]
                }
            ],
            "source": [
                "dataset = ogb.linkproppred.LinkPropPredDataset(\n",
                "    name=\"ogbl-biokg\", root=dataset_directory\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Since the objective here is to show how to build the BESS-KGE `KGDataset` class from scratch, we will not use any of the preprocessing utilities provided by `ogb`. Instead we will use the raw source files directly. To start from the most generic case, we will actually undo some of the preprocessing already performed on the data, namely the mapping from entity labels to entity indices for the different entity types."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "4c74b206b8ad44138310c6e59c715ce6",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "split:   0%|          | 0/3 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "# train triples: 4,762,678\n",
                        "# validation triples: 162,886\n",
                        "# test triples: 162,870\n"
                    ]
                }
            ],
            "source": [
                "label_dict = {}\n",
                "for l in [\"disease\", \"drug\", \"function\", \"protein\", \"sideeffect\"]:\n",
                "    # labels for the entities of type l\n",
                "    # we append the entity type to the label to prevent label collisions across different types\n",
                "    # notice that some entities, like proteins, have labels that are still numerical\n",
                "    label_dict[l] = (\n",
                "        pd.read_csv(\n",
                "            Path(dataset_directory).joinpath(\n",
                "                f\"ogbl_biokg/mapping/{l}_entidx2name.csv.gz\"\n",
                "            )\n",
                "        ).set_index(\"ent idx\")[\"ent name\"]\n",
                "        + f\" ({l})\"\n",
                "    ).values\n",
                "\n",
                "# labels for relation types\n",
                "rel_dict = (\n",
                "    pd.read_csv(\n",
                "        Path(dataset_directory).joinpath(f\"ogbl_biokg/mapping/relidx2relname.csv.gz\")\n",
                "    )\n",
                "    .set_index(\"rel idx\")[\"rel name\"]\n",
                "    .values\n",
                ")\n",
                "\n",
                "# collect triples in train, valid, test DataFrames\n",
                "# replacing the entity IDs with their original labels.\n",
                "df_dict = {}\n",
                "for split in tqdm({\"test\", \"train\", \"valid\"}, desc=\"split\"):\n",
                "    triples = []\n",
                "    data = torch.load(\n",
                "        Path(dataset_directory).joinpath(f\"ogbl_biokg/split/random/{split}.pt\")\n",
                "    )\n",
                "    for h, h_type, t, t_type, r in zip(\n",
                "        data[\"head\"],\n",
                "        data[\"head_type\"],\n",
                "        data[\"tail\"],\n",
                "        data[\"tail_type\"],\n",
                "        data[\"relation\"],\n",
                "    ):\n",
                "        h_label = label_dict[h_type][h]\n",
                "        t_label = label_dict[t_type][t]\n",
                "        r_label = rel_dict[r]\n",
                "        triples.append((h_label, r_label, t_label))\n",
                "    df_dict[split] = pd.DataFrame(\n",
                "        triples, columns=[\"head_label\", \"relation_label\", \"tail_label\"]\n",
                "    )\n",
                "\n",
                "print(f\"# train triples: {df_dict['train'].shape[0]:,}\")\n",
                "print(f\"# validation triples: {df_dict['valid'].shape[0]:,}\")\n",
                "print(f\"# test triples: {df_dict['test'].shape[0]:,}\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [
                {
                    "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>head_label</th>\n",
                            "      <th>relation_label</th>\n",
                            "      <th>tail_label</th>\n",
                            "    </tr>\n",
                            "  </thead>\n",
                            "  <tbody>\n",
                            "    <tr>\n",
                            "      <th>0</th>\n",
                            "      <td>C0038586 (disease)</td>\n",
                            "      <td>disease-protein</td>\n",
                            "      <td>1653 (protein)</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>1</th>\n",
                            "      <td>C0751849 (disease)</td>\n",
                            "      <td>disease-protein</td>\n",
                            "      <td>718 (protein)</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>2</th>\n",
                            "      <td>C1320474 (disease)</td>\n",
                            "      <td>disease-protein</td>\n",
                            "      <td>8622 (protein)</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>3</th>\n",
                            "      <td>C0270844 (disease)</td>\n",
                            "      <td>disease-protein</td>\n",
                            "      <td>3569 (protein)</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>4</th>\n",
                            "      <td>C4279912 (disease)</td>\n",
                            "      <td>disease-protein</td>\n",
                            "      <td>8856 (protein)</td>\n",
                            "    </tr>\n",
                            "  </tbody>\n",
                            "</table>\n",
                            "</div>"
                        ],
                        "text/plain": [
                            "           head_label   relation_label      tail_label\n",
                            "0  C0038586 (disease)  disease-protein  1653 (protein)\n",
                            "1  C0751849 (disease)  disease-protein   718 (protein)\n",
                            "2  C1320474 (disease)  disease-protein  8622 (protein)\n",
                            "3  C0270844 (disease)  disease-protein  3569 (protein)\n",
                            "4  C4279912 (disease)  disease-protein  8856 (protein)"
                        ]
                    },
                    "execution_count": 3,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "df_dict[\"train\"].head()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We put ourselves in a very generic starting point, where the **edges of the knowledge graph are represented as a list of (head, relation, tail) triples**, with *unique labels* for entities and relations.\n",
                "\n",
                "In our knowledge graph, moreover, entities are of different types (disease, drug, function, protein and side-effect). This is not always the case, but BESS-KGE can leverage this additional information, for instance by constructing negative samples corrupting entities only with entities of the same type (see [besskge.negative_sampler.TypeBasedNegativeSampler](https://graphcore-research.github.io/bess-kge/API_ref/negative_sampler.html#besskge.negative_sampler.TypeBasedShardedNegativeSampler)). We store this data by creating a mapping from entity labels to entity types."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {},
            "outputs": [
                {
                    "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>ent_label</th>\n",
                            "      <th>ent_type</th>\n",
                            "    </tr>\n",
                            "  </thead>\n",
                            "  <tbody>\n",
                            "    <tr>\n",
                            "      <th>0</th>\n",
                            "      <td>C0000737 (disease)</td>\n",
                            "      <td>disease</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>1</th>\n",
                            "      <td>C0000744 (disease)</td>\n",
                            "      <td>disease</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>2</th>\n",
                            "      <td>C0000768 (disease)</td>\n",
                            "      <td>disease</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>3</th>\n",
                            "      <td>C0000771 (disease)</td>\n",
                            "      <td>disease</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>4</th>\n",
                            "      <td>C0000772 (disease)</td>\n",
                            "      <td>disease</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>...</th>\n",
                            "      <td>...</td>\n",
                            "      <td>...</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>93768</th>\n",
                            "      <td>C3665444 (sideeffect)</td>\n",
                            "      <td>sideeffect</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>93769</th>\n",
                            "      <td>C3665596 (sideeffect)</td>\n",
                            "      <td>sideeffect</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>93770</th>\n",
                            "      <td>C3665609 (sideeffect)</td>\n",
                            "      <td>sideeffect</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>93771</th>\n",
                            "      <td>C3665624 (sideeffect)</td>\n",
                            "      <td>sideeffect</td>\n",
                            "    </tr>\n",
                            "    <tr>\n",
                            "      <th>93772</th>\n",
                            "      <td>C3665888 (sideeffect)</td>\n",
                            "      <td>sideeffect</td>\n",
                            "    </tr>\n",
                            "  </tbody>\n",
                            "</table>\n",
                            "<p>93773 rows × 2 columns</p>\n",
                            "</div>"
                        ],
                        "text/plain": [
                            "                   ent_label    ent_type\n",
                            "0         C0000737 (disease)     disease\n",
                            "1         C0000744 (disease)     disease\n",
                            "2         C0000768 (disease)     disease\n",
                            "3         C0000771 (disease)     disease\n",
                            "4         C0000772 (disease)     disease\n",
                            "...                      ...         ...\n",
                            "93768  C3665444 (sideeffect)  sideeffect\n",
                            "93769  C3665596 (sideeffect)  sideeffect\n",
                            "93770  C3665609 (sideeffect)  sideeffect\n",
                            "93771  C3665624 (sideeffect)  sideeffect\n",
                            "93772  C3665888 (sideeffect)  sideeffect\n",
                            "\n",
                            "[93773 rows x 2 columns]"
                        ]
                    },
                    "execution_count": 4,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "entity_types = {l: t for t in label_dict.keys() for l in label_dict[t]}\n",
                "entity_types = pd.DataFrame(\n",
                "    {\"ent_label\": entity_types.keys(), \"ent_type\": entity_types.values()}\n",
                ")\n",
                "entity_types"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "ent_type\n",
                            "disease       10687\n",
                            "drug          10533\n",
                            "function      45085\n",
                            "protein       17499\n",
                            "sideeffect     9969\n",
                            "Name: ent_type, dtype: int64"
                        ]
                    },
                    "execution_count": 5,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# entity type counts\n",
                "entity_types.groupby(\"ent_type\")[\"ent_type\"].count()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## The KGDataset class in BESS-KGE\n",
                "\n",
                "When using BESS-KGE, the knowledge graph data is stored in an instance of the [besskge.dataset.KGDataset](https://graphcore-research.github.io/bess-kge/API_ref/dataset.html#besskge.dataset.KGDataset) class. This class has built-in methods to download and build some commonly-used knowledge graph datasets, or it can be instantiated manually with custom data by specifying all the required attributes.\n",
                "\n",
                "The `besskge.dataset.KGDataset.from_dataframe` method is perfect to **build a custom dataset starting from labelled triples with minimum effort**. It simply requires a pandas DataFrame containing all labelled triples (or a dictionary of DataFrames, one for each of the dataset splits). If entities are of different types, like for OGBL-BioKG, this can be communicated to `KGDataset` by providing it with a mapping of entity labels to entity types, in the form of a dictionary or a pandas Series, indexed over the entity labels."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "ent_label\n",
                            "C0000737 (disease)          disease\n",
                            "C0000744 (disease)          disease\n",
                            "C0000768 (disease)          disease\n",
                            "C0000771 (disease)          disease\n",
                            "C0000772 (disease)          disease\n",
                            "                            ...    \n",
                            "C3665444 (sideeffect)    sideeffect\n",
                            "C3665596 (sideeffect)    sideeffect\n",
                            "C3665609 (sideeffect)    sideeffect\n",
                            "C3665624 (sideeffect)    sideeffect\n",
                            "C3665888 (sideeffect)    sideeffect\n",
                            "Name: ent_type, Length: 93773, dtype: object"
                        ]
                    },
                    "execution_count": 6,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "entity_types_series = entity_types.set_index(\"ent_label\")[\"ent_type\"]\n",
                "entity_types_series"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "To build `KGDataset` we only need to call `KGDataset.from_dataframe`, specifying the names of the columns of the dataframes which contain the entity and relation labels."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [],
            "source": [
                "biokg = KGDataset.from_dataframe(\n",
                "    df_dict,\n",
                "    head_column=\"head_label\",\n",
                "    relation_column=\"relation_label\",\n",
                "    tail_column=\"tail_label\",\n",
                "    entity_types=entity_types_series,\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "That was easy, wasn't it? Let us have a closer look at the attributes of the `KGDataset` class we just created."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Number of entities: 93,773\n",
                        "\n",
                        "Number of relation types: 51\n",
                        "\n",
                        "Number of triples: \n",
                        " training: 4,762,678 \n",
                        " validation: 162,886\n",
                        " test: 162,870\n"
                    ]
                }
            ],
            "source": [
                "print(f\"Number of entities: {biokg.n_entity:,}\\n\")\n",
                "print(f\"Number of relation types: {biokg.n_relation_type}\\n\")\n",
                "print(\n",
                "    f\"Number of triples: \\n training: {biokg.triples['train'].shape[0]:,} \\n validation: {biokg.triples['valid'].shape[0]:,}\"\n",
                "    f\"\\n test: {biokg.triples['test'].shape[0]:,}\"\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Note that only the entities that appear as head or tail in at least one of the dataset's triples are counted in `KGDataset.n_entity` (and the same for relations in `KGDataset.n_relation_type`).\n",
                "\n",
                "If we take a look at the `triples` attribute of `KGDataset`, we see that it is still structured as a dictionary, with the same keys that we used in `df_dict` to identify the different dataset splits."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "dict_keys(['test', 'valid', 'train'])"
                        ]
                    },
                    "execution_count": 9,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "biokg.triples.keys()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "array([[ 9768,     0, 82518],\n",
                            "       [ 5375,     0, 74685],\n",
                            "       [ 9433,     0, 77223],\n",
                            "       ...,\n",
                            "       [67129,    50, 74418],\n",
                            "       [81566,    50, 77737],\n",
                            "       [67073,    50, 82459]], dtype=int32)"
                        ]
                    },
                    "execution_count": 10,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "biokg.triples[\"train\"]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Each row of these NumPy arrays corresponds to a (h,r,t) triple in the dataset, but now the **labels for entities and relations have been replaced by numerical IDs**. We can use `KGDataset.entity_dict` and `KGDataset.relation_dict` to recover the mapping from IDs to labels for entities and relation types respectively."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "['C0031117 (disease)',\n",
                            " 'C4225263 (disease)',\n",
                            " 'C0393814 (disease)',\n",
                            " 'C0342883 (disease)',\n",
                            " 'C4017556 (disease)',\n",
                            " 'C1862103 (disease)',\n",
                            " 'C0751004 (disease)',\n",
                            " 'C4310790 (disease)',\n",
                            " 'C0027339 (disease)',\n",
                            " 'C3150412 (disease)']"
                        ]
                    },
                    "execution_count": 11,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "biokg.entity_dict[:10]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "meaning that the entity with ID 0 is `biokg.entity_dict[0]`, the entity with ID 1 is `biokg.entity_dict[1]`, etc. (and similarly for `biokg.relation_dict`).\n",
                "\n",
                "Let's do a quick sanity check."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "('C0751495 (disease)', 'disease-protein', '7249 (protein)')"
                        ]
                    },
                    "execution_count": 12,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "part = \"test\"\n",
                "triple_number = 1234\n",
                "\n",
                "head_label = biokg.entity_dict[biokg.triples[part][triple_number, 0]]\n",
                "relation_label = biokg.relation_dict[biokg.triples[part][triple_number, 1]]\n",
                "tail_label = biokg.entity_dict[biokg.triples[part][triple_number, 2]]\n",
                "\n",
                "head_label, relation_label, tail_label"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "In the original dataframes, this should coincide with:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "head_label        C0751495 (disease)\n",
                            "relation_label       disease-protein\n",
                            "tail_label            7249 (protein)\n",
                            "Name: 1234, dtype: object"
                        ]
                    },
                    "execution_count": 13,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "df_dict[part].iloc[triple_number][[\"head_label\", \"relation_label\", \"tail_label\"]]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "It checks out!\n",
                "\n",
                "It is important to note that, when entities have different types, the numerical entity IDs need to be assigned so that **entities of the same type have contiguous IDs**! This is done automatically when using `KGDataset.from_dataframe`, but it needs to be kept in mind if you are instantiating the `KGDataset` class manually.\n",
                "\n",
                "Since entity IDs are now clustered by type, we only need to know the ID ranges corresponding to the different types, which are stored in `KGDataset.type_offsets`:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "{'disease': 0,\n",
                            " 'drug': 10687,\n",
                            " 'function': 21220,\n",
                            " 'protein': 66305,\n",
                            " 'sideeffect': 83804}"
                        ]
                    },
                    "execution_count": 14,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "biokg.type_offsets"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "This means that entities with ID from 0 to 10686 are of type 'disease', from 10687 to 21219 are of type 'drug' and so on.\n",
                "\n",
                "The type IDs (assigned following the order of the keys in `KGDataset.type_offsets`) for heads and tails of all triples in the dataset can be immediately recovered using `KGDataset.ht_types`:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "dict_keys(['test', 'valid', 'train'])"
                        ]
                    },
                    "execution_count": 15,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "biokg.ht_types.keys()"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "(4762678, 2) (162886, 2) (162870, 2)\n"
                    ]
                }
            ],
            "source": [
                "print(\n",
                "    biokg.ht_types[\"train\"].shape,\n",
                "    biokg.ht_types[\"valid\"].shape,\n",
                "    biokg.ht_types[\"test\"].shape,\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Each row of these arrays stores the ID of the type of the head entity and tail entity of the corresponding triple, for example:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 17,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "array([0, 3])"
                        ]
                    },
                    "execution_count": 17,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "part = \"valid\"\n",
                "triple_number = 0\n",
                "\n",
                "biokg.ht_types[part][triple_number]"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "This means that the first validation triple has a head entity of type 'disease' (the 0-th key of `biokg.type_offsets`) and a tail entity of type 'protein' (the key of `biokg.type_offsets` with index 3). Indeed, we can check:"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "('disease', 'protein')"
                        ]
                    },
                    "execution_count": 18,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "type_head = entity_types_series[df_dict[part].iloc[triple_number][\"head_label\"]]\n",
                "type_tail = entity_types_series[df_dict[part].iloc[triple_number][\"tail_label\"]]\n",
                "\n",
                "type_head, type_tail"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Random dataset split\n",
                "\n",
                "What if no custom train/validation/test split is provided? `KGDataset.from_dataframe` can perform a random split, with the desired ratios between the three parts. This happens whenever it is provided with a single pandas DataFrame, instead of a dictionary of DataFrames as before."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 19,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Total number of triples: 5,088,434\n"
                    ]
                }
            ],
            "source": [
                "# Merge all triples in a single DataFrame\n",
                "df_all_triples = pd.concat(\n",
                "    [df_dict[\"train\"], df_dict[\"valid\"], df_dict[\"test\"]], axis=0\n",
                ")\n",
                "\n",
                "print(f\"Total number of triples: {df_all_triples.shape[0]:,}\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Number of entities: 93,773\n",
                        "\n",
                        "Number of relation types: 51\n",
                        "\n",
                        "Number of triples: \n",
                        " training: 4,070,747 \n",
                        " validation: 508,843\n",
                        " test: 508,844\n"
                    ]
                }
            ],
            "source": [
                "# 80/10/10 train/valid/test split\n",
                "split_ratios = (0.8, 0.1, 0.1)\n",
                "\n",
                "biokg_random = KGDataset.from_dataframe(\n",
                "    df_all_triples,\n",
                "    head_column=\"head_label\",\n",
                "    relation_column=\"relation_label\",\n",
                "    tail_column=\"tail_label\",\n",
                "    entity_types=entity_types_series,\n",
                "    split=split_ratios,\n",
                ")\n",
                "\n",
                "print(f\"Number of entities: {biokg_random.n_entity:,}\\n\")\n",
                "print(f\"Number of relation types: {biokg_random.n_relation_type}\\n\")\n",
                "print(\n",
                "    f\"Number of triples: \\n training: {biokg_random.triples['train'].shape[0]:,} \\n validation: {biokg_random.triples['valid'].shape[0]:,}\"\n",
                "    f\"\\n test: {biokg_random.triples['test'].shape[0]:,}\"\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Saving and loading\n",
                "\n",
                "If you plan to run multiple experiments on the dataset, it is convenient to save the final `KGDataset` class to file using the `KGDataset.save` method. You can then load it, whenever you need to use it, with `KGDataset.load`. This will ensure that all experiments are run on the same train/validation/test split, and maintaining the same mapping of labels to IDs for entities are relations."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 27,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "KGDataset saved to biokg.pkl\n"
                    ]
                }
            ],
            "source": [
                "save_file = \"biokg.pkl\"\n",
                "\n",
                "biokg.save(save_file)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 28,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Loaded KGDataset at biokg.pkl\n",
                        "Number of entities: 93,773\n",
                        "\n",
                        "Number of relation types: 51\n",
                        "\n",
                        "Number of triples: \n",
                        " training: 4,762,678 \n",
                        " validation: 162,886\n",
                        " test: 162,870\n"
                    ]
                }
            ],
            "source": [
                "reloaded_biokg = KGDataset.load(save_file)\n",
                "\n",
                "print(f\"Number of entities: {reloaded_biokg.n_entity:,}\\n\")\n",
                "print(f\"Number of relation types: {reloaded_biokg.n_relation_type}\\n\")\n",
                "print(\n",
                "    f\"Number of triples: \\n training: {reloaded_biokg.triples['train'].shape[0]:,} \\n validation: {reloaded_biokg.triples['valid'].shape[0]:,}\"\n",
                "    f\"\\n test: {reloaded_biokg.triples['test'].shape[0]:,}\"\n",
                ")"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Conclusions and next steps\n",
                "\n",
                "The `KGDataset` class has a few additional attributes, which can be defined when instantiating the class manually. For instance, for each triple it allows you to specify a set of negative heads and tails that should be used to corrupt that triple. This is useful when you want to find the best completion for a (h,r,?)/(?,r,t) query only among a specific set of candidate nodes, or you have already identified good negative samples that you want to use during training.\n",
                "\n",
                "For more information on the `KGDataset` class, have a look at the [BESS-KGE documentation](https://graphcore-research.github.io/bess-kge/API_ref/dataset.html#besskge.dataset.KGDataset).\n",
                "\n",
                "Once you have built your dataset as a `KGDataset` class, you are ready to use BESS to train your preferred KGE model and perform inference with it! To learn how, we suggest starting from the introductory [KGE Training and Inference on OGBL-BioKG](1_biokg_training_inference.ipynb) notebook."
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": ".venv_3.2",
            "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.10"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}