diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index da0ec10..2fb654f 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -1,27 +1,36 @@
# How to contribute to the BESS-KGE project
-You can contribute to the development of the BESS-KGE project, even if you don't have access to IPUs (you can use the [IPUModel](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.2.0/reference.html#poptorch.Options.useIpuModel) to emulate most functionalities of the physical hardware).
+
+
+You can contribute to the development of the BESS-KGE project, even if you don't have access to IPUs (you can use the [IPUModel](https://docs.graphcore.ai/projects/poptorch-user-guide/en/3.2.0/reference.html#poptorch.Options.useIpuModel) to emulate most functionalities of the physical hardware).
## VS Code server on Paperspace
-Setting up a VS Code server on [Paperspace](https://www.paperspace.com/graphcore) will allow you to tunnel into a machine with IPUs from the VS Code web editor or the desktop app. This requires minimum effort and is an excellent solution for developing and testing code directly on IPU hardware.
+Setting up a VS Code server on [Paperspace](https://www.paperspace.com/graphcore) will allow you to tunnel into a machine with IPUs from the VS Code web editor or the desktop app. This requires minimum effort and is an excellent solution for developing and testing code directly on IPU hardware. Here's how to do it.
-You can launch a 6-hours session on a Paperspace machine with access to 4 IPUs **for free** by clicking on this button: 
+1. Fork the [BESS-KGE repository](https://github.com/graphcore-research/bess-kge).
-Start the machine (this will also clone the repo for you) and open up a terminal from the left pane.
+2. You can launch a 6-hours session on a Paperspace machine with access to 4 IPUs **for free** by using a link of the form:
+ ```
+ https://console.paperspace.com/github/{USERID}/{REPONAME}?container=graphcore%2Fpytorch-paperspace%3A3.3.0-ubuntu-20.04-20230703&machine=Free-IPU-POD4
+ ```
-
+ where `{USERID}/{REPOPNAME}` is the github address of the forked repository (e.g. `graphcore-research/bess-kge` for the original repo).
-In the terminal, run the command
-```shell
-bash .gradient/launch_vscode_server.sh {tunnel-name}
-```
+3. Start the machine (this will also clone the repo for you) and open up a terminal from the left pane.
+
+ 
-where `tunnel-name` is an optional argument that you can use to define the name of the remote tunnel (if not set, it will default to `ipu-paperspace`).
+4. In the terminal, run the command
+ ```shell
+ bash .gradient/launch_vscode_server.sh {tunnel-name}
+ ```
-The script will download and install all dependencies and start the tunnel. You will be asked to authorize the tunnel through GitHub, before being provided with the tunnel link. Please refer to [this notebook](https://ipu.dev/fmo4AZ) for additional details on these steps and to connect the VS Code desktop app to the remote tunnel.
+ where `tunnel-name` is an optional argument that you can use to define the name of the remote tunnel (if not set, it will default to `ipu-paperspace`). The script will download and install all dependencies and start the tunnel.
-Once VS Code is connected to the Paperspace machine, run `./dev build` to build all custom ops. You are now ready to create a new git branch and start developing!
+5. When asked, authorize the tunnel through GitHub (with an account having writing privileges to the forked repository). You will be then provided with the tunnel link. Please refer to [this notebook](https://ipu.dev/fmo4AZ) for additional details on these steps and to connect the VS Code desktop app to the remote tunnel.
+
+6. Once VS Code is connected to the Paperspace machine, run `./dev build` to build all custom ops. You are now ready to start developing!
When closing a session and stopping the Paperspace machine, remember to unregister the tunnel in VS Code as explained in the "Common Issues" paragraph of the [notebook](https://ipu.dev/fmo4AZ). To resume your work, just access the clone of the BESS-KGE repo in the "Projects" section of your Paperspace profile, start a new machine and repeat the operations above. All code changes to the local repo, as well as VS Code settings and extensions installed, will persist across sessions.
@@ -43,10 +52,14 @@ pip install $POPLAR_SDK_ENABLED/../poptorch-*.whl
pip install -r requirements-dev.txt
```
-Finally, build all custom ops by running `./dev build`
+Finally, clone your fork of the BESS-KGE repository and build all custom ops by running `./dev build`
+
+## Development tips
+
+The `./dev` command can be used to run several utility scripts during development. Check `./dev --help` for a list of dev options.
-## Tips
+Before submitting a PR to the upstream repo, use `./dev ci` to run all CI checks locally. In particular, be mindful of our formatting requirements: you can check for formatting errors by running `./dev format` and `./dev lint` (both commands are automatically run inside `./dev ci`).
-Run `./dev --help` for a list of dev options. In particular, use `./dev ci` to run all CI checks locally. Run individual tests with pattern matching filtering `./dev tests -k FILTER`.
+Add unit tests to the `tests` folder. You can run individual unit tests with pattern matching filtering `./dev tests -k FILTER`.
Add `.cpp` custom ops to `besskge/custom_ops`. Also, update the [Makefile](Makefile) when adding custom ops.
diff --git a/NOTICE.md b/NOTICE.md
index cd44ef3..18e28d2 100644
--- a/NOTICE.md
+++ b/NOTICE.md
@@ -2,9 +2,7 @@ Copyright (c) 2023 Graphcore Ltd. Licensed under the MIT License.
The included code is released under an MIT license, (see [LICENSE](LICENSE)).
-The ogbl-biokg and ogbl-wikikg2 datasets are licensed under CC-0.
-
-The [YAGO3 dataset](https://yago-knowledge.org/downloads/yago-3) by the [YAGO team](https://yago-knowledge.org/contributors) of the [Max-Planck Institute for Informatics](https://www.mpi-inf.mpg.de/home/) and [Telcom Paris](https://www.telecom-paris.fr/) is licensed under [CC BY 3.0](https://creativecommons.org/licenses/by/3.0/).
+## Dependencies
Our dependencies are (see [requirements.txt](requirements.txt)):
@@ -18,8 +16,21 @@ Our dependencies are (see [requirements.txt](requirements.txt)):
We also use additional Python dependencies for development/testing/documentation (see [requirements-dev.txt](requirements-dev.txt)).
+## Dataset disclaimer
+
+This repository provides dataloaders for third party datasets. The use of these datasets is at own risk and Graphcore offers no warranties of any kind. It is the user's responsibility to comply with all license requirements for datasets downloaded with dataloaders in this repository.
+
+The tutorial notebooks make use of the following datasets:
+
+* [ogbl-biokg](https://ogb.stanford.edu/docs/linkprop/#ogbl-biokg), licensed under CC-0;
+
+* [ogbl-wikikg2](https://ogb.stanford.edu/docs/linkprop/#ogbl-wikikg2), licensed under CC-0;
+
+* [YAGO3 dataset](https://yago-knowledge.org/downloads/yago-3) by the [YAGO team](https://yago-knowledge.org/contributors) of the [Max-Planck Institute for Informatics](https://www.mpi-inf.mpg.de/home/) and [Telcom Paris](https://www.telecom-paris.fr/), licensed under [CC BY 3.0](https://creativecommons.org/licenses/by/3.0/).
+
+## Derived work
-**This directory includes derived work from the following:**
+This directory includes derived work from the following:
---
diff --git a/README.md b/README.md
index 909b229..20cd1ff 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,6 @@
# BESS-KGE
+
[**Installation guide**](#usage)
| [**Tutorials**](#paperspace-notebook-tutorials)
@@ -77,6 +78,17 @@ Additional variations of the distribution scheme are detailed in the [BESS-KGE d
All APIs are documented in the [BESS-KGE API documentation](https://graphcore-research.github.io/bess-kge/API_reference.html).
+### Datasets
+
+BESS-KGE provides built-in dataloaders for the following datasets. Notice that the use of these datasets is at own risk and Graphcore offers no warranties of any kind. It is the user's responsibility to comply with all license requirements for datasets downloaded with dataloaders in this repository.
+
+| Dataset | Builder method | Entities | Entity types | Relation types | Triples | License |
+| --- | --- | --- | --- | --- | --- | --- |
+| [ogbl-biokg](https://ogb.stanford.edu/docs/linkprop/#ogbl-biokg) | [KGDataset.build_ogbl_biokg](https://graphcore-research.github.io/bess-kge/generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_ogbl_biokg) | 93,773 | 5 | 51 | 5,088,434 | CC-0 |
+| [ogbl-wikikg2](https://ogb.stanford.edu/docs/linkprop/#ogbl-wikikg2) | [KGDataset.build_ogbl_wikikg2](https://graphcore-research.github.io/bess-kge/generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_ogbl_wikikg2) | 2,500,604 | 1 | 535 | 16,968,094 | CC-0 |
+| [YAGO3-10](https://yago-knowledge.org/downloads/yago-3) | [KGDataset.build_yago310](https://graphcore-research.github.io/bess-kge/generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_yago310) | 123,182 | 1 | 37 | 1,089,040 | CC BY 3.0 |
+| [OpenBioLink2020](https://github.com/openbiolink/openbiolink#benchmark-dataset) | [KGDataset.build_openbiolink](https://graphcore-research.github.io/bess-kge/generated/besskge.dataset.KGDataset.html#besskge.dataset.KGDataset.build_openbiolink) | 184,635 | 7 | 28 | 4,563,405 | [link](https://github.com/openbiolink/openbiolink#Source-databases-and-their-licenses) |
+
### Known limitations
* BESS-KGE supports distribution for up to 16 IPUs.
@@ -178,10 +190,11 @@ For a walkthrough of the `besskge` library functionalities, see our Jupyter note
2. [Link prediction on the YAGO3-10 dataset](notebooks/2_yago_topk_prediction.ipynb) [](https://console.paperspace.com/github/graphcore-research/bess-kge?container=graphcore%2Fpytorch-paperspace%3A3.3.0-ubuntu-20.04-20230703&machine=Free-IPU-POD4&file=%2Fnotebooks%2F2_yago_topk_prediction.ipynb)
3. [FP16 weights and compute on the OGBL-WikiKG2 dataset](notebooks/3_wikikg2_fp16.ipynb) [](https://console.paperspace.com/github/graphcore-research/bess-kge?container=graphcore%2Fpytorch-paperspace%3A3.3.0-ubuntu-20.04-20230703&machine=Free-IPU-POD4&file=%2Fnotebooks%2F3_wikikg2_fp16.ipynb)
+For pointers on how to run BESS-KGE on a custom Knowledge Graph dataset, see the notebook [Using BESS-KGE with your own data](notebooks/0_custom_KG_dataset.ipynb) [](https://console.paperspace.com/github/graphcore-research/bess-kge?container=graphcore%2Fpytorch-paperspace%3A3.3.0-ubuntu-20.04-20230703&machine=Free-IPU-POD4&file=%2Fnotebooks%2F0_custom_KG_dataset.ipynb)
## Contributing
-You can contribute to the BESS-KGE project. See [How to contribute to the BESS-KGE project](CONTRIBUTING.md)
+You can contribute to the BESS-KGE project: PRs are welcome! For details, see [How to contribute to the BESS-KGE project](CONTRIBUTING.md).
## References
BESS: Balanced Entity Sampling and Sharing for Large-Scale Knowledge Graph Completion ([arXiv](https://arxiv.org/abs/2211.12281))
@@ -190,6 +203,6 @@ BESS: Balanced Entity Sampling and Sharing for Large-Scale Knowledge Graph Compl
Copyright (c) 2023 Graphcore Ltd. Licensed under the MIT License.
-The included code is released under the MIT license, (see [details of the license](LICENSE)).
+The included code is released under the MIT license (see [details of the license](LICENSE)).
See [notices](NOTICE.md) for dependencies, credits, derived work and further details.
\ No newline at end of file
diff --git a/besskge/dataset.py b/besskge/dataset.py
index 85e760a..17a12d1 100644
--- a/besskge/dataset.py
+++ b/besskge/dataset.py
@@ -8,9 +8,10 @@
import dataclasses
import pickle
import tarfile
+import zipfile
from io import BytesIO
from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
import numpy as np
import ogb.linkproppred
@@ -75,7 +76,153 @@ def ht_types(self) -> Optional[Dict[str, NDArray[np.int32]]]:
return None
@classmethod
- def build_biokg(cls, root: Path) -> "KGDataset":
+ def from_triples(
+ cls,
+ data: NDArray[np.int32],
+ split: Tuple[float, float, float] = (0.7, 0.15, 0.15),
+ seed: int = 1234,
+ entity_dict: Optional[List[str]] = None,
+ relation_dict: Optional[List[str]] = None,
+ type_offsets: Optional[Dict[str, int]] = None,
+ ) -> "KGDataset":
+ """
+ Build a dataset from an array of triples, where IDs for entities
+ and relations have already been assigned. Note that, if entities have
+ types, entities of the same type need to have contiguous IDs.
+ Triples are randomly split in train/validation/test sets.
+ If a pre-defined train/validation/test split is wanted, the KGDataset
+ class should be instantiated manually.
+
+ :param data:
+ Numpy array of triples [head_id, relation_id, tail_id]. Shape
+ (num_triples, 3).
+ :param split:
+ Tuple to set the train/validation/test split.
+ :param seed:
+ Random seed for the train/validation/test split.
+ :param entity_dict:
+ Optional entity labels by ID.
+ :param relation_dict:
+ Optional relation labels by ID.
+ :param type_offsets:
+ Offset of entity types
+
+ :return: Instance of the KGDataset class.
+ """
+ num_triples = data.shape[0]
+ num_train = int(num_triples * split[0])
+ num_valid = int(num_triples * split[1])
+
+ rng = np.random.default_rng(seed=seed)
+ rng.shuffle(data, axis=0)
+
+ triples = dict()
+ triples["train"], triples["valid"], triples["test"] = np.split(
+ data, (num_train, num_train + num_valid), axis=0
+ )
+
+ return cls(
+ n_entity=data[:, [0, 2]].max() + 1,
+ n_relation_type=data[:, 1].max() + 1,
+ entity_dict=entity_dict,
+ relation_dict=relation_dict,
+ type_offsets=type_offsets,
+ triples=triples,
+ )
+
+ @classmethod
+ def from_dataframe(
+ cls,
+ df: Union[pd.DataFrame, Dict[str, pd.DataFrame]],
+ head_column: Union[int, str],
+ relation_column: Union[int, str],
+ tail_column: Union[int, str],
+ entity_types: Optional[Union["pd.Series[str]", Dict[str, str]]] = None,
+ split: Tuple[float, float, float] = (0.7, 0.15, 0.15),
+ seed: int = 1234,
+ ) -> "KGDataset":
+ """
+ Build a KGDataset from a pandas DataFrame of labeled (h,r,t) triples.
+ IDs for entities and relations are automatically assigned based on labels
+ in such a way that entities of the same type have contiguous IDs.
+
+ :param df:
+ Pandas DataFrame of all triples in the knowledge graph dataset,
+ or dictionary of DataFrames of triples for each part of the dataset split
+ :param head_column:
+ Name of the DataFrame column storing head entities
+ :param relation_column:
+ Name of the DataFrame column storing relations
+ :param tail_column:
+ Name of the DataFrame column storing tail entities
+ :param entity_types:
+ If entities have types, dictionary or pandas Series of mappings
+ entity label -> entity type.
+ :param split:
+ Tuple to set the train/validation/test split.
+ Only used if no pre-defined dataset split is specified,
+ i.e. if `df` is not a dictionary.
+ :param seed:
+ Random seed for the train/validation/test split.
+ Only used if no pre-defined dataset split is specified,
+ i.e. if `df` is not a dictionary.
+
+ :return: Instance of the KGDataset class.
+ """
+
+ df_dict = {"all": df} if isinstance(df, pd.DataFrame) else df
+ unique_ent = pd.concat(
+ [
+ pd.concat([dfp[head_column], dfp[tail_column]])
+ for dfp in df_dict.values()
+ ]
+ ).unique()
+ ent2id = pd.Series(np.arange(len(unique_ent)), index=unique_ent, name="ent_id")
+ unique_rel = pd.concat(
+ [dfp[relation_column] for dfp in df_dict.values()]
+ ).unique()
+ rel2id = pd.Series(np.arange(len(unique_rel)), index=unique_rel, name="rel_id")
+
+ if entity_types is not None:
+ ent2type = pd.Series(entity_types, name="ent_type")
+ ent2id_type = pd.merge(
+ ent2id, ent2type, how="left", left_index=True, right_index=True
+ ).sort_values("ent_type")
+ ent2id.index = ent2id_type.index
+ type_off = (
+ ent2id_type.groupby("ent_type")["ent_type"].count().cumsum().shift(1)
+ )
+ type_off.iloc[0] = 0
+ type_offsets = type_off.astype("int64").to_dict()
+ else:
+ type_offsets = None
+
+ entity_dict = ent2id.index.tolist()
+ relation_dict = rel2id.index.tolist()
+
+ triples = {}
+ for part, dfp in df_dict.items():
+ heads = dfp[head_column].map(ent2id).values.astype(np.int32)
+ tails = dfp[tail_column].map(ent2id).values.astype(np.int32)
+ rels = dfp[relation_column].map(rel2id).values.astype(np.int32)
+ triples[part] = np.stack([heads, rels, tails], axis=1)
+
+ if isinstance(df, pd.DataFrame):
+ return KGDataset.from_triples(
+ triples["all"], split, seed, entity_dict, relation_dict, type_offsets
+ )
+ else:
+ return cls(
+ n_entity=len(entity_dict),
+ n_relation_type=len(relation_dict),
+ entity_dict=entity_dict,
+ relation_dict=relation_dict,
+ type_offsets=type_offsets,
+ triples=triples,
+ )
+
+ @classmethod
+ def build_ogbl_biokg(cls, root: Path) -> "KGDataset":
"""
Build the ogbl-biokg dataset :cite:p:`OGB`
@@ -138,7 +285,7 @@ def build_biokg(cls, root: Path) -> "KGDataset":
)
@classmethod
- def build_wikikg2(cls, root: Path) -> "KGDataset":
+ def build_ogbl_wikikg2(cls, root: Path) -> "KGDataset":
"""
Build the ogbl-wikikg2 dataset :cite:p:`OGB`
@@ -215,130 +362,82 @@ def build_yago310(cls, root: Path) -> "KGDataset":
with tarfile.open(fileobj=BytesIO(res.content)) as tarf:
tarf.extractall(path=root)
- train = np.loadtxt(root.joinpath("train.txt"), delimiter="\t", dtype=str)
- valid = np.loadtxt(root.joinpath("valid.txt"), delimiter="\t", dtype=str)
- test = np.loadtxt(root.joinpath("test.txt"), delimiter="\t", dtype=str)
-
- entity_dict, entity_id = np.unique(
- np.concatenate(
- [
- train[:, 0],
- train[:, 2],
- valid[:, 0],
- valid[:, 2],
- test[:, 0],
- test[:, 2],
- ]
- ),
- return_inverse=True,
- )
- entity_split_limits = np.cumsum(
- [
- train.shape[0],
- train.shape[0],
- valid.shape[0],
- valid.shape[0],
- test.shape[0],
- ]
+ train_triples = pd.read_csv(
+ root.joinpath("train.txt"), delimiter="\t", dtype=str, header=None
)
- (
- train_head_id,
- train_tail_id,
- validation_head_id,
- validation_tail_id,
- test_head_id,
- test_tail_id,
- ) = np.split(entity_id, entity_split_limits)
-
- rel_dict, rel_id = np.unique(
- np.concatenate([train[:, 1], valid[:, 1], test[:, 1]]),
- return_inverse=True,
+ valid_triples = pd.read_csv(
+ root.joinpath("valid.txt"), delimiter="\t", dtype=str, header=None
)
- relation_split_limits = np.cumsum([train.shape[0], valid.shape[0]])
- train_rel_id, validation_rel_id, test_rel_id = np.split(
- rel_id, relation_split_limits
+ test_triples = pd.read_csv(
+ root.joinpath("test.txt"), delimiter="\t", dtype=str, header=None
)
- triples = {
- "train": np.concatenate(
- [train_head_id[:, None], train_rel_id[:, None], train_tail_id[:, None]],
- axis=1,
- ),
- "validation": np.concatenate(
- [
- validation_head_id[:, None],
- validation_rel_id[:, None],
- validation_tail_id[:, None],
- ],
- axis=1,
- ),
- "test": np.concatenate(
- [test_head_id[:, None], test_rel_id[:, None], test_tail_id[:, None]],
- axis=1,
- ),
- }
-
- return cls(
- n_entity=len(entity_dict),
- n_relation_type=len(rel_dict),
- entity_dict=entity_dict.tolist(),
- relation_dict=rel_dict.tolist(),
- type_offsets=None,
- triples=triples,
- neg_heads=None,
- neg_tails=None,
+ return cls.from_dataframe(
+ {"train": train_triples, "valid": valid_triples, "test": test_triples},
+ head_column=0,
+ relation_column=1,
+ tail_column=2,
)
@classmethod
- def from_triples(
- cls,
- data: NDArray[np.int32],
- split: Tuple[float, float, float] = (0.7, 0.15, 0.15),
- seed: int = 1234,
- entity_dict: Optional[List[str]] = None,
- relation_dict: Optional[List[str]] = None,
- type_offsets: Optional[Dict[str, int]] = None,
- ) -> "KGDataset":
+ def build_openbiolink(cls, root: Path) -> "KGDataset":
"""
- Build a dataset from an array of triples. Note that if a pre-defined
- train/validation/test split is wanted the KGDataset class should be instantiated
- manually.
+ Build the high-quality version of the OpenBioLink2020
+ dataset :cite:p:`openbiolink`
- :param data:
- Numpy array of triples [head_id, relation_id, tail_id]. Shape
- (num_triples, 3).
- :param split:
- Tuple to set the train/validation/test split.
- :param seed:
- Random seed for the train/validation/test split.
- :param entity_dict:
- Optional entity labels by ID.
- :param relation_dict:
- Optional relation labels by ID.
- :param type_offsets:
- Offset of entity types
+ .. seealso:: https://github.com/openbiolink/openbiolink#benchmark-dataset
- :return: Instance of the KGDataset class.
- """
- num_triples = data.shape[0]
- num_train = int(num_triples * split[0])
- num_valid = int(num_triples * split[1])
+ :param root:
+ Local path to the dataset. If the dataset is not present in this
+ location, then it is downloaded and stored here.
- rng = np.random.default_rng(seed=seed)
- rng.shuffle(data, axis=0)
+ :return: The HQ OpenBioLink2020 KGDataset.
+ """
- triples = dict()
- triples["train"], triples["valid"], triples["test"] = np.split(
- data, (num_train, num_train + num_valid), axis=0
+ if not (
+ root.joinpath("HQ_DIR/train_test_data/train_sample.csv").is_file()
+ and root.joinpath("HQ_DIR/train_test_data/val_sample.csv").is_file()
+ and root.joinpath("HQ_DIR/train_test_data/test_sample.csv").is_file()
+ and root.joinpath("HQ_DIR/train_test_data/train_val_nodes.csv").is_file()
+ ):
+ print("Downloading dataset...")
+ res = requests.get(url="https://zenodo.org/record/3834052/files/HQ_DIR.zip")
+ with zipfile.ZipFile(BytesIO(res.content)) as zip_f:
+ zip_f.extractall(path=root)
+
+ column_names = ["h_label", "r_label", "t_label", "quality", "TP/TN", "source"]
+ train_triples = pd.read_csv(
+ root.joinpath("HQ_DIR/train_test_data/train_sample.csv"),
+ header=None,
+ names=column_names,
+ sep="\t",
+ )
+ valid_triples = pd.read_csv(
+ root.joinpath("HQ_DIR/train_test_data/val_sample.csv"),
+ header=None,
+ names=column_names,
+ sep="\t",
+ )
+ test_triples = pd.read_csv(
+ root.joinpath("HQ_DIR/train_test_data/test_sample.csv"),
+ header=None,
+ names=column_names,
+ sep="\t",
)
- return cls(
- n_entity=data[:, [0, 2]].max() + 1,
- n_relation_type=data[:, 1].max() + 1,
- entity_dict=entity_dict,
- relation_dict=relation_dict,
- type_offsets=type_offsets,
- triples=triples,
+ entity_types = pd.read_csv(
+ root.joinpath("HQ_DIR/train_test_data/train_val_nodes.csv"),
+ header=None,
+ names=["ent_label", "ent_type"],
+ sep="\t",
+ ).set_index("ent_label")["ent_type"]
+
+ return cls.from_dataframe(
+ {"train": train_triples, "valid": valid_triples, "test": test_triples},
+ head_column="h_label",
+ relation_column="r_label",
+ tail_column="t_label",
+ entity_types=entity_types,
)
def save(self, out_file: Path) -> None:
diff --git a/docs/source/KGbib.bib b/docs/source/KGbib.bib
index fa9e948..d1ad271 100644
--- a/docs/source/KGbib.bib
+++ b/docs/source/KGbib.bib
@@ -192,3 +192,13 @@ @inproceedings{TranS
pages ={1202--1208},
year ={2022}
}
+
+@article{openbiolink,
+ author = {Breit Anna and Ott Simon and Agibetov Asan and Samwald Matthias},
+ title = {{OpenBioLink: a benchmarking framework for large-scale biomedical link prediction}},
+ journal = {Bioinformatics},
+ volume = {36},
+ number = {13},
+ pages = {4097-4098},
+ year = {2020},
+}
diff --git a/docs/source/images/Terminal1.png b/docs/source/images/Terminal1.png
deleted file mode 100644
index 98ca224..0000000
Binary files a/docs/source/images/Terminal1.png and /dev/null differ
diff --git a/docs/source/images/Terminal3.png b/docs/source/images/Terminal3.png
new file mode 100644
index 0000000..fc684dc
Binary files /dev/null and b/docs/source/images/Terminal3.png differ
diff --git a/notebooks/0_custom_KG_dataset.ipynb b/notebooks/0_custom_KG_dataset.ipynb
new file mode 100644
index 0000000..f788047
--- /dev/null
+++ b/notebooks/0_custom_KG_dataset.ipynb
@@ -0,0 +1,988 @@
+{
+ "cells": [
+ {
+ "attachments": {},
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Using BESS-KGE with your Own Data\n",
+ "\n",
+ "Copyright (c) 2023 Graphcore Ltd. All rights reserved.\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 pre-process 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": [
+ "## Environment setup\n",
+ "\n",
+ "While this notebook doesn't contain any code that needs to be run on IPUs or other accelerating hardware, the best way to run it is on Paperspace Gradient's cloud IPUs because everything is already set up for you.\n",
+ "\n",
+ " [](https://console.paperspace.com/github/graphcore-research/bess-kge?container=graphcore%2Fpytorch-paperspace%3A3.3.0-ubuntu-20.04-20230703&machine=Free-IPU-POD4&file=%2Fnotebooks%2F0_custom_KG_dataset.ipynb)"
+ ]
+ },
+ {
+ "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"
+ ]
+ },
+ {
+ "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",
+ "\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 pre-processing 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": 3,
+ "metadata": {},
+ "outputs": [
+ {
+ "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 {\"test\", \"train\", \"valid\"}:\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": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " head_label | \n",
+ " relation_label | \n",
+ " tail_label | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " C0038586 (disease) | \n",
+ " disease-protein | \n",
+ " 1653 (protein) | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " C0751849 (disease) | \n",
+ " disease-protein | \n",
+ " 718 (protein) | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " C1320474 (disease) | \n",
+ " disease-protein | \n",
+ " 8622 (protein) | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " C0270844 (disease) | \n",
+ " disease-protein | \n",
+ " 3569 (protein) | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " C4279912 (disease) | \n",
+ " disease-protein | \n",
+ " 8856 (protein) | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " ent_label | \n",
+ " ent_type | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " C0000737 (disease) | \n",
+ " disease | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " C0000744 (disease) | \n",
+ " disease | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " C0000768 (disease) | \n",
+ " disease | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " C0000771 (disease) | \n",
+ " disease | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " C0000772 (disease) | \n",
+ " disease | \n",
+ "
\n",
+ " \n",
+ " | ... | \n",
+ " ... | \n",
+ " ... | \n",
+ "
\n",
+ " \n",
+ " | 93768 | \n",
+ " C3665444 (sideeffect) | \n",
+ " sideeffect | \n",
+ "
\n",
+ " \n",
+ " | 93769 | \n",
+ " C3665596 (sideeffect) | \n",
+ " sideeffect | \n",
+ "
\n",
+ " \n",
+ " | 93770 | \n",
+ " C3665609 (sideeffect) | \n",
+ " sideeffect | \n",
+ "
\n",
+ " \n",
+ " | 93771 | \n",
+ " C3665624 (sideeffect) | \n",
+ " sideeffect | \n",
+ "
\n",
+ " \n",
+ " | 93772 | \n",
+ " C3665888 (sideeffect) | \n",
+ " sideeffect | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "
93773 rows × 2 columns
\n",
+ "