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/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,32 @@
+"""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
+    task_spec_param: Mapping[str, Any] | None