Fix ImportError in kg_emb caused by SampleBaseDataset rename

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:

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
 

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
 
 
@@ -13,7 +13,7 @@ class ComplEx(KGEBaseModel):
 
     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
 
 
@@ -12,7 +12,7 @@ class DistMult(KGEBaseModel):
     """
     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
@@ -32,7 +32,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
 
 
@@ -13,7 +13,7 @@ class RotatE(KGEBaseModel):
 
     def __init__(
         self, 
-        dataset: SampleBaseDataset, 
+        dataset: SampleDataset,
         e_dim: int = 600, 
         r_dim: int = 300, 
         ns='adv', 

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
 
 
@@ -13,7 +13,7 @@ class TransE(KGEBaseModel):
 
     def __init__(
         self, 
-        dataset: SampleBaseDataset, 
+        dataset: SampleDataset, 
         e_dim: int = 300, 
         r_dim: int = 300, 
         ns: str = "adv",