fix(kg_emb): migrate SampleKGDataset off the removed 1.x dataset API

docs/api/medcode.rst

@@ -97,6 +97,50 @@ Medication codes:
     :undoc-members:
     :show-inheritance:
 
+Knowledge graph embeddings
+--------------------------
+
+``pyhealth.medcode.pretrained_embeddings.kg_emb`` trains TransE, RotatE,
+DistMult and ComplEx on an in-memory list of triples. Since PyHealth 2.0
+the sample dataset is a map-style :class:`torch.utils.data.Dataset`. Build
+the loader with :class:`torch.utils.data.DataLoader` and
+:func:`pyhealth.datasets.collate_fn_dict_with_padding` --
+:func:`pyhealth.datasets.get_dataloader` is streaming-only and calls
+``set_shuffle()``.
+
+See ``examples/kg_emb_sample_dataset.py`` for a self-contained walk-through.
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autofunction:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split
+
+.. 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/kg_emb_sample_dataset.py

@@ -0,0 +1,58 @@
+"""Train a TransE model on a synthetic knowledge-graph sample dataset.
+
+This example does not download UMLS. It shows the 2.0-safe path:
+
+1. Build an in-memory :class:`SampleKGDataset` (not ``SampleDataset``).
+2. Split into train/val/test folds with :func:`split`.
+3. Wrap the train fold in ``torch.utils.data.DataLoader`` using
+   :func:`collate_fn_dict_with_padding`. Do not call ``get_dataloader``:
+   that helper requires ``litdata.StreamingDataset.set_shuffle()``.
+"""
+
+import torch
+from torch.utils.data import DataLoader
+
+from pyhealth.datasets import collate_fn_dict_with_padding
+from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import (
+    SampleKGDataset,
+    split,
+)
+from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE
+
+samples = [
+    {
+        "triple": (i % 5, i % 2, (i + 1) % 5),
+        "ground_truth_head": [i % 5, (i + 1) % 5],
+        "ground_truth_tail": [(i + 1) % 5],
+        "subsampling_weight": torch.tensor([0.25]),
+    }
+    for i in range(10)
+]
+
+dataset = SampleKGDataset(
+    samples=samples,
+    dataset_name="toy",
+    task_name="link_prediction",
+    entity2id={"a": 0, "b": 1, "c": 2, "d": 3, "e": 4},
+    relation2id={"treats": 0, "causes": 1},
+    negative_sampling=4,
+)
+print(dataset)
+print(dataset.stat())
+
+train, val, test = split(dataset, [0.6, 0.2, 0.2], seed=0)
+print("fold sizes", len(train), len(val), len(test))
+
+train_loader = DataLoader(
+    train,
+    batch_size=2,
+    shuffle=True,
+    collate_fn=collate_fn_dict_with_padding,
+)
+
+model = TransE(dataset=dataset, e_dim=8, r_dim=8, ns="uniform")
+batch = next(iter(train_loader))
+out = model(**batch)
+print("loss", float(out["loss"]))
+out["loss"].backward()
+print("backward ok")

pyhealth/datasets/utils.py

@@ -260,6 +260,22 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
         A dictionary where each key corresponds to a list of values from the batch.
         Tensor values are padded to the same shape.
         Tuples of (time, values) from temporal processors are collated separately.
+        A nested dict of tensors (e.g. ``{"value": ..., "mask": ...}`` from a
+        processor like ``KGProcessor``) is collated sub-key by sub-key.
+
+    Examples:
+        >>> import torch
+        >>> batch = [
+        ...     {"triple": torch.tensor([0, 0, 1]),
+        ...      "ground_truth_head": {"value": torch.tensor([0, 4]), "mask": torch.tensor([1, 1])}},
+        ...     {"triple": torch.tensor([2, 0, 3]),
+        ...      "ground_truth_head": {"value": torch.tensor([2, 0]), "mask": torch.tensor([1, 0])}},
+        ... ]
+        >>> collated = collate_fn_dict_with_padding(batch)
+        >>> collated["triple"].shape
+        torch.Size([2, 3])
+        >>> collated["ground_truth_head"]["value"].shape
+        torch.Size([2, 2])
     """
     collated = {}
     keys = batch[0].keys()
@@ -304,6 +320,30 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
         elif HAS_PYG and isinstance(values[0], PyGData):
             collated[key] = PyGBatch.from_data_list(values)
 
+        elif isinstance(values[0], dict) and all(
+            isinstance(v, torch.Tensor) for v in values[0].values()
+        ):
+            # Nested all-tensor feature dict, e.g. {"value": Tensor, "mask": Tensor}
+            # from KGProcessor. Stack each sub-key independently; sub-values
+            # share shape across samples when the processor pads to a fixed
+            # field-wide length (KGProcessor.fit), so this is typically a
+            # plain stack rather than dynamic padding.
+            #
+            # Restricted to all-tensor dicts so heterogeneous per-sample
+            # dicts (e.g. a "hyperparameters" field of plain Python values)
+            # keep falling through to the generic list passthrough below,
+            # preserving their existing list-of-dicts collation shape.
+            sub_collated: dict = {}
+            for sub_key in values[0]:
+                sub_values = [v[sub_key] for v in values]
+                if all(v.shape == sub_values[0].shape for v in sub_values):
+                    sub_collated[sub_key] = torch.stack(sub_values)
+                else:
+                    sub_collated[sub_key] = pad_sequence(
+                        sub_values, batch_first=True, padding_value=0
+                    )
+            collated[key] = sub_collated
+
         elif isinstance(values[0], torch.Tensor):
             # Check if shapes are the same
             shapes = [v.shape for v in values]

pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py

@@ -1,4 +1,13 @@
-from .sample_kg_dataset import SampleKGDataset
 from .base_kg_dataset import BaseKGDataset
+from .protocols import KGDatasetProtocol
+from .sample_kg_dataset import SampleKGDataset
+from .splitter import split
 from .umls import UMLSDataset
-from .splitter import split
+
+__all__ = [
+    "BaseKGDataset",
+    "KGDatasetProtocol",
+    "SampleKGDataset",
+    "UMLSDataset",
+    "split",
+]

pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py

@@ -1,15 +1,12 @@
 import logging
 import os
 from abc import ABC
+from collections.abc import Callable
 
-from tqdm import tqdm
-import pandas as pd
-from pandarallel import pandarallel
-from typing import Callable, Optional
 from pyhealth.datasets.utils import MODULE_CACHE_PATH, hash_str
-from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset
 from pyhealth.utils import load_pickle, save_pickle
 
+from .sample_kg_dataset import SampleKGDataset
 
 logger = logging.getLogger(__name__)
 
@@ -34,13 +31,23 @@ class BaseKGDataset(ABC):
             Default is False.
         refresh_cache: whether to refresh the cache; if true, the dataset will
             be processed from scratch and the cache will be updated. Default is False.
-    
+
+    Examples:
+        >>> import tempfile
+        >>> class _ToyKG(BaseKGDataset):
+        ...     def raw_graph_process(self):
+        ...         self.triples = [(0, 0, 1)]
+        ...         self.entity_num = 2
+        ...         self.relation_num = 1
+        >>> ds = _ToyKG(root=tempfile.mkdtemp(), dataset_name="toy")
+        >>> len(ds)
+        1
     """
 
     def __init__(
         self,
         root: str,
-        dataset_name: Optional[str] = None,
+        dataset_name: str | None = None,
         dev: bool = False,
         refresh_cache: bool = False
     ):
@@ -86,7 +93,7 @@ def info():
 
     def stat(self):
         """Returns some statistics of the base dataset."""
-        lines = list()
+        lines = []
         lines.append("")
         lines.append(f"Statistics of base dataset (dev={self.dev}):")
         lines.append(f"\t- Dataset: {self.dataset_name}")
@@ -97,13 +104,12 @@ def stat(self):
         lines.append(f"\t- Number of samples: {len(self.samples)}")
         lines.append("")
         print("\n".join(lines))
-        return 
 
 
     def set_task(
         self,
         task_fn: Callable,
-        task_name: Optional[str] = None,
+        task_name: str | None = None,
         save: bool = True,
         **kwargs
     ) -> SampleKGDataset:

pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py

@@ -0,0 +1,34 @@
+"""Structural contract between knowledge-graph datasets and embedding models."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any, Protocol, runtime_checkable
+
+__all__ = ["KGDatasetProtocol"]
+
+
+@runtime_checkable
+class KGDatasetProtocol(Protocol):
+    """Minimal capability a dataset must expose to parameterise a KGE model.
+
+    Embedding models only need the cardinality of the entity and relation
+    vocabularies. They never read the samples at construction time. Depending
+    on this Protocol rather than on a concrete class keeps the model layer
+    testable with lightweight doubles and free of import-time coupling to the
+    dataset layer.
+
+    Examples:
+        >>> class _Toy:
+        ...     entity_num = 2
+        ...     relation_num = 1
+        ...     task_spec_param = None
+        >>> isinstance(_Toy(), KGDatasetProtocol)
+        True
+    """
+
+    entity_num: int
+    relation_num: int
+
+    @property
+    def task_spec_param(self) -> Mapping[str, Any] | None: ...