Fix kg_emb broken import: SampleBaseDataset -> SampleDataset, add pan…

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:
+
 
 
 

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)

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()

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:

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", 

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", 

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",

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',