diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index bfc4edcea..3f1443ff0 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -97,6 +97,43 @@ Medication codes: :undoc-members: :show-inheritance: + +Pretrained Knowledge Graph Embeddings: +--------------------------------------- +Utilities for training knowledge graph embedding (KGE) models over medical +code knowledge graphs, so entity/relation embeddings can be reused as +pretrained inputs for downstream tasks. + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.KGEBaseModel + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.TransE + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.RotatE + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.DistMult + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.ComplEx + :members: + :undoc-members: + :show-inheritance: + diff --git a/examples/medcode.py b/examples/medcode.py index 5fcb27625..a7d22f0ef 100644 --- a/examples/medcode.py +++ b/examples/medcode.py @@ -11,3 +11,34 @@ atc = InnerMap.load("ATC") print("Looking up for ATC code G04CA02") print(atc.lookup("G04CA02")) + + +# Knowledge Graph Embedding (KGE) training over medical code knowledge graphs. +# See pyhealth.medcode.pretrained_embeddings.kg_emb for TransE, RotatE, +# DistMult, and ComplEx. +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + +entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} +relation2id = {"treats": 0} +samples = [ + { + "triple": (0, 0, 1), + "ground_truth_head": [0, 2], + "ground_truth_tail": [1], + "subsampling_weight": 1.0, + }, +] +kg_dataset = SampleKGDataset( + samples=samples, + dataset_name="toy_kg", + entity_num=len(entity2id), + relation_num=len(relation2id), + entity2id=entity2id, + relation2id=relation2id, +) +print("KG dataset stats:") +kg_dataset.stat() + +kge_model = TransE(kg_dataset, e_dim=32, r_dim=32) +print("Entity embedding shape:", kge_model.E_emb.shape) \ No newline at end of file diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py index 59d72e888..338ce0562 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py @@ -1,10 +1,10 @@ -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset -class SampleKGDataset(SampleBaseDataset): +class SampleKGDataset(SampleDataset): """Sample KG dataset class. - This class inherits from `SampleBaseDataset` and is specifically designed + This class inherits from `SampleDataset` and is specifically designed for KG datasets. Args: @@ -22,6 +22,28 @@ class SampleKGDataset(SampleBaseDataset): } dataset_name: the name of the dataset. Default is None. task_name: the name of the task. Default is None. + + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, + ... dataset_name="toy_kg", + ... entity_num=len(entity2id), + ... relation_num=len(relation2id), + ... entity2id=entity2id, + ... relation2id=relation2id, + ... ) + >>> dataset.stat() """ def __init__( self, @@ -36,7 +58,14 @@ def __init__( **kwargs ): - super().__init__(samples, dataset_name, task_name) + # SampleDataset's __init__ now expects a path to a directory built + # by SampleBuilder.save() (it's litdata-backed), not an in-memory + # list of samples. SampleKGDataset predates that change and still + # works with samples held directly in memory, so we set the base + # attributes ourselves instead of delegating to super().__init__(). + self.samples = samples + self.dataset_name = dataset_name + self.task_name = task_name self.dev = dev self.entity_num = entity_num self.relation_num = relation_num @@ -49,6 +78,10 @@ def __init__( if kwargs != None: self.task_spec_param = kwargs + def __len__(self): + """Returns the number of samples in the dataset.""" + return self.sample_size + def __getitem__(self, index): """ A sample is a dict containing following data: @@ -64,7 +97,7 @@ def __getitem__(self, index): } """ return self.samples[index] - + def stat(self): """Returns some statistics of the base dataset.""" lines = list() diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 559ecb7c4..a676f6599 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -4,18 +4,18 @@ import numpy as np import torch -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset def split( - dataset: SampleBaseDataset, + dataset: SampleDataset, ratios: Union[Tuple[float, float, float], List[float]], seed: Optional[int] = None, ): """Splits the dataset by its outermost indexed items Args: - dataset: a `SampleBaseDataset` object + dataset: a `SampleDataset` object ratios: a list/tuple of ratios for train / val / test seed: random seed for shuffling the dataset @@ -26,6 +26,21 @@ def split( Note: The original dataset can be accessed by `train_dataset.dataset`, `val_dataset.dataset`, and `test_dataset.dataset`. + + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.splitter import split + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... {"triple": (0, 0, 1), "ground_truth_head": [0], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + ... {"triple": (2, 0, 1), "ground_truth_head": [2], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> train_dataset, val_dataset, test_dataset = split(dataset, ratios=[0.5, 0.5, 0.0], seed=42) """ if seed is not None: diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 8fa2a443a..135ee42ba 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -9,11 +9,31 @@ class ComplEx(KGEBaseModel): Paper: Trouillon, T., Welbl, J., Riedel, S., Gaussier, É. and Bouchard, G., 2016, June. Complex embeddings for simple link prediction. In International conference on machine learning (pp. 2071-2080). PMLR + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ComplEx + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = ComplEx(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 600, r_dim: int = 600, ns: str = "adv", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index e7563137c..2cfa2a19f 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -9,10 +9,30 @@ class DistMult(KGEBaseModel): Paper: Yang, B., Yih, W.T., He, X., Gao, J. and Deng, L. Embedding entities and relations for learning and inference in knowledge bases. ICLR 2015. + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = DistMult(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py index 2de13afe2..51c582ac8 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,5 +1,5 @@ from abc import ABC -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch import time @@ -22,6 +22,26 @@ class KGEBaseModel(ABC, nn.Module): use_regularization: whether to apply regularization or not, False by default. mode: evaluation metric type, one of "binary", "multiclass", or "multilabel", "multiclass" by default + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import KGEBaseModel + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = KGEBaseModel(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ @property @@ -32,7 +52,7 @@ def device(self): def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 500, r_dim: int = 500, ns: str = "uniform", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index df7143a6e..b2a02d202 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -9,11 +9,31 @@ class RotatE(KGEBaseModel): Paper: Sun, Z., Deng, Z.H., Nie, J.Y. and Tang, J., 2019. Rotate: Knowledge graph embedding by relational rotation in complex space. ICLR 2019. + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import RotatE + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = RotatE(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 600, r_dim: int = 300, ns='adv', diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index fbb6e68f6..09ad34db3 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -9,11 +9,31 @@ class TransE(KGEBaseModel): Paper: Bordes, A., Usunier, N., Garcia-Duran, A., Weston, J. and Yakhnenko, Translating embeddings for modeling multi-relational data. NIPS 2013. + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = TransE(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", diff --git a/pyproject.toml b/pyproject.toml index b4626e649..d69f40284 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,8 @@ dependencies = [ "narwhals~=2.13.0", "more-itertools~=10.8.0", "einops>=0.8.0", - "linear-attention-transformer>=0.19.1", + "linear-attention-transformer>=0.19.1", + "pandarallel>=1.6.5", ] license = "BSD-3-Clause" license-files = ["LICENSE.md"] diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py new file mode 100644 index 000000000..891d1ef9e --- /dev/null +++ b/tests/core/test_kg_emb.py @@ -0,0 +1,48 @@ +import inspect +import unittest + +from pyhealth.datasets import SampleDataset + + +class TestKgEmbImports(unittest.TestCase): + def test_import_pretrained_embeddings(self): + try: + import pyhealth.medcode.pretrained_embeddings # noqa: F401 + except ImportError as e: + self.fail( + f"Importing pyhealth.medcode.pretrained_embeddings failed: {e}" + ) + + def test_model_classes_importable(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + KGEBaseModel, + RotatE, + TransE, + ) + + for cls in (KGEBaseModel, TransE, RotatE, DistMult, ComplEx): + self.assertTrue( + isinstance(cls, type), + msg=f"{cls} was not importable as a class", + ) + + def test_transe_uses_sample_dataset(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + sig = inspect.signature(TransE.__init__) + dataset_annotation = sig.parameters["dataset"].annotation + self.assertEqual( + dataset_annotation, + SampleDataset, + msg=( + "TransE.__init__'s `dataset` parameter is not annotated as " + "SampleDataset — regression check for the " + "SampleBaseDataset -> SampleDataset rename" + ), + ) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file