From 98e496d6db1008b96c0e9a000c5dd3570423f97c Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 24 Aug 2026 22:57:32 +0200 Subject: [PATCH 01/10] fix(kg_emb): decouple SampleKGDataset from the 2.0 streaming pipeline SampleDataset is now a litdata.StreamingDataset that expects schema.pkl. A knowledge-graph task is an in-memory list of triples, so SampleKGDataset subclasses torch.utils.data.Dataset and exposes KGDatasetProtocol for the models. Co-authored-by: Cursor --- .../kg_emb/datasets/__init__.py | 12 +- .../kg_emb/datasets/protocols.py | 32 +++ .../kg_emb/datasets/sample_kg_dataset.py | 231 ++++++++++++------ 3 files changed, 204 insertions(+), 71 deletions(-) create mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 22bda1718..15df00d9f 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,4 +1,14 @@ +# ruff: noqa: I001 +# BaseKGDataset imports SampleKGDataset from this package, so the sample +# dataset must be bound before the base class is loaded. +from .protocols import KGDatasetProtocol from .sample_kg_dataset import SampleKGDataset from .base_kg_dataset import BaseKGDataset from .umls import UMLSDataset -from .splitter import split \ No newline at end of file + +__all__ = [ + "BaseKGDataset", + "KGDatasetProtocol", + "SampleKGDataset", + "UMLSDataset", +] diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py new file mode 100644 index 000000000..6b093bcc6 --- /dev/null +++ b/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 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..778cc88a8 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,81 +1,172 @@ -from pyhealth.datasets import SampleBaseDataset +"""Task-specific sample dataset for knowledge-graph embedding models. +This module deliberately does **not** build on +:class:`pyhealth.datasets.SampleDataset`. Since PyHealth 2.0, +``SampleDataset`` is a ``litdata.StreamingDataset`` whose contract is "a +directory containing ``schema.pkl`` plus optimized chunks". A knowledge +graph task produces an in-memory list of triple-level records and has no +feature schema, no processors and no patient/visit index: the two +abstractions are unrelated. +""" -class SampleKGDataset(SampleBaseDataset): - """Sample KG dataset class. +from __future__ import annotations - This class inherits from `SampleBaseDataset` and is specifically designed - for KG datasets. +from collections.abc import Mapping, Sequence +from typing import Any + +from torch.utils.data import Dataset + +KGSample = Mapping[str, Any] + +__all__ = ["SampleKGDataset"] + + +class SampleKGDataset(Dataset): + r"""In-memory dataset of knowledge-graph link-prediction samples. + + Each sample is a mapping with the following keys: + + ``triple`` + A positive triple :math:`(h, r, t)` given as integer indices, + e.g. ``(0, 0, 2835)``. + ``ground_truth_head`` + All entities :math:`h'` such that :math:`(h', r, t)` is observed + in the graph. Used to filter false negatives when scoring the + query :math:`(?, r, t)`. + ``ground_truth_tail`` + All entities :math:`t'` such that :math:`(h, r, t')` is observed + in the graph. + ``subsampling_weight`` + The word2vec-style subsampling weight of the triple, a scalar + tensor. Args: - samples: a list of samples - A sample is a dict containing following data: - { - 'triple': a positive triple e.g., (0, 0, 2835) - 'ground_truth_head': a list of ground truth of the head entity in the dataset given - query (e.g., (?, 0, 2835)) with current relation r and tail entity t. - e.g., [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029] - 'ground_truth_tail': a list of ground truth of the tail entity in the dataset given - query (e.g., (0, 0, ?)) with current head entity h and relation r. - e.g., [398, 244, 3872, 3053, 1711, 2835, 1348, 2309] - 'subsampling_weight': the subsampling weight (a scalar) of this triple, which may be applied for loss calculation - } - dataset_name: the name of the dataset. Default is None. - task_name: the name of the task. Default is None. + samples: The task samples, typically produced by + ``link_prediction_fn``. + dataset_name: Human-readable name of the source dataset. + task_name: Human-readable name of the task. + dev: Whether the samples come from a development subset. + entity_num: Number of entities. Inferred from ``entity2id`` when + omitted. + relation_num: Number of relations. Inferred from ``relation2id`` + when omitted. + entity2id: Mapping from surface entity identifier to integer + index. + relation2id: Mapping from surface relation identifier to integer + index. + **task_spec_param: Task hyper-parameters forwarded to the model + at training time (e.g. ``negative_sampling=128``). + + Raises: + ValueError: If the declared cardinalities contradict the provided + vocabularies. + + Examples: + >>> import torch + >>> samples = [ + ... { + ... "triple": (i, i % 2, (i + 1) % 5), + ... "ground_truth_head": [i, (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, + ... ) + >>> len(dataset) + 10 + >>> dataset[0]["triple"] + (0, 0, 1) + >>> dataset.entity_num, dataset.relation_num + (5, 2) + >>> dataset.id2entity[2] + 'c' + >>> dataset.task_spec_param + {'negative_sampling': 4} """ + def __init__( - self, - samples, - dataset_name="", - task_name="", - dev=False, - entity_num=0, - relation_num=0, - entity2id=None, - relation2id=None, - **kwargs - ): - - super().__init__(samples, dataset_name, task_name) + self, + samples: Sequence[KGSample], + dataset_name: str = "", + task_name: str = "", + dev: bool = False, + entity_num: int = 0, + relation_num: int = 0, + entity2id: Mapping[Any, int] | None = None, + relation2id: Mapping[Any, int] | None = None, + **task_spec_param: Any, + ) -> None: + self.samples: list[KGSample] = list(samples) + self.dataset_name = dataset_name + self.task_name = task_name self.dev = dev - self.entity_num = entity_num - self.relation_num = relation_num - self.sample_size = len(samples) - self.task_spec_param = None - self.entity2id = entity2id - self.id2entity = {v: k for k, v in entity2id.items()} - self.relation2id = relation2id - self.id2relation = {v: k for k, v in relation2id.items()} - if kwargs != None: - self.task_spec_param = kwargs - - def __getitem__(self, index): - """ - A sample is a dict containing following data: - { - 'triple': a positive triple e.g., (0, 0, 2835) - 'ground_truth_head': a list of ground truth of the head entity in the dataset given - query (e.g., (?, 0, 2835)) with current relation r and tail entity t. - e.g., [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029] - 'ground_truth_tail': a list of ground truth of the tail entity in the dataset given - query (e.g., (0, 0, ?)) with current head entity h and relation r. - e.g., [398, 244, 3872, 3053, 1711, 2835, 1348, 2309] - 'subsampling_weight': the subsampling weight (a scalar) of this triple, which may be applied for loss calculation + + self.entity2id: dict[Any, int] = dict(entity2id or {}) + self.relation2id: dict[Any, int] = dict(relation2id or {}) + self.id2entity: dict[int, Any] = {v: k for k, v in self.entity2id.items()} + self.id2relation: dict[int, Any] = { + v: k for k, v in self.relation2id.items() } - """ + + self.entity_num = entity_num or len(self.entity2id) + self.relation_num = relation_num or len(self.relation2id) + self._validate_cardinalities() + + # ``None`` rather than ``{}`` preserves the historical sentinel + # used by the models. + self.task_spec_param: dict[str, Any] | None = task_spec_param or None + + def _validate_cardinalities(self) -> None: + if self.entity2id and self.entity_num != len(self.entity2id): + raise ValueError( + f"entity_num={self.entity_num} contradicts len(entity2id)=" + f"{len(self.entity2id)}" + ) + if self.relation2id and self.relation_num != len(self.relation2id): + raise ValueError( + f"relation_num={self.relation_num} contradicts " + f"len(relation2id)={len(self.relation2id)}" + ) + + @property + def sample_size(self) -> int: + """Number of samples. Kept as a property for backward compatibility.""" + return len(self.samples) + + def __len__(self) -> int: + return len(self.samples) + + def __getitem__(self, index: int) -> KGSample: return self.samples[index] - def stat(self): - """Returns some statistics of the base dataset.""" - lines = list() - lines.append("") - lines.append(f"Statistics of base dataset (dev={self.dev}):") - lines.append(f"\t- Dataset: {self.dataset_name}") - lines.append(f"\t- Number of triples: {len(self.samples)}") - lines.append(f"\t- Number of entities: {self.entity_num}") - lines.append(f"\t- Number of relations: {self.relation_num}") - lines.append(f"\t- Task name: {self.task_name}") - lines.append(f"\t- Task-specific hyperparameters: {self.task_spec_param}") - lines.append("") - print("\n".join(lines)) - return + def stat(self) -> str: + """Return -- and print -- a human-readable summary of the dataset.""" + lines = [ + "", + f"Statistics of sample KG dataset (dev={self.dev}):", + f"\t- Dataset: {self.dataset_name}", + f"\t- Task name: {self.task_name}", + f"\t- Number of triples: {len(self.samples)}", + f"\t- Number of entities: {self.entity_num}", + f"\t- Number of relations: {self.relation_num}", + f"\t- Task-specific hyperparameters: {self.task_spec_param}", + "", + ] + report = "\n".join(lines) + print(report) + return report + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(dataset_name={self.dataset_name!r}, " + f"task_name={self.task_name!r}, size={len(self)}, " + f"entity_num={self.entity_num}, relation_num={self.relation_num})" + ) From 6b5af443364b8070e17cb7af438638f17456d148 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 24 Aug 2026 22:59:58 +0200 Subject: [PATCH 02/10] fix(kg_emb): correct dataset type hints across models KGE models only need entity_num, relation_num and task_spec_param. Annotate that structural contract with KGDatasetProtocol so the model layer no longer imports the removed SampleBaseDataset. Co-authored-by: Cursor --- .../kg_emb/models/complex.py | 25 ++++++++++++--- .../kg_emb/models/distmult.py | 25 ++++++++++++--- .../kg_emb/models/kg_base.py | 32 ++++++++++++------- .../kg_emb/models/rotate.py | 27 ++++++++++++---- .../kg_emb/models/transe.py | 25 ++++++++++++--- 5 files changed, 102 insertions(+), 32 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 8fa2a443a..28b0d78aa 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,19 +1,34 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch +from .kg_base import KGEBaseModel + +if TYPE_CHECKING: + from ..datasets.protocols import KGDatasetProtocol + class ComplEx(KGEBaseModel): - """ ComplEx + """ComplEx - Paper: Trouillon, T., Welbl, J., Riedel, S., Gaussier, É. and Bouchard, G., 2016, June. + 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: + >>> class _Toy: + ... entity_num = 2 + ... relation_num = 1 + ... task_spec_param = None + >>> model = ComplEx(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: KGDatasetProtocol, 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..dd2973328 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,18 +1,33 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch +from .kg_base import KGEBaseModel + +if TYPE_CHECKING: + from ..datasets.protocols import KGDatasetProtocol + class DistMult(KGEBaseModel): - """ DistMult + """DistMult - Paper: Yang, B., Yih, W.T., He, X., Gao, J. and Deng, L. Embedding entities and + 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: + >>> class _Toy: + ... entity_num = 2 + ... relation_num = 1 + ... task_spec_param = None + >>> model = DistMult(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: KGDatasetProtocol, 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..6de31761c 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,15 +1,19 @@ +from __future__ import annotations + from abc import ABC -from pyhealth.datasets import SampleBaseDataset +from typing import TYPE_CHECKING -import torch -import time import numpy as np -import torch.nn as nn +import torch import torch.nn.functional as F +from torch import nn + +if TYPE_CHECKING: + from ..datasets.protocols import KGDatasetProtocol class KGEBaseModel(ABC, nn.Module): - """ Abstract class for Knowledge Graph Embedding models. + """Abstract class for Knowledge Graph Embedding models. Args: e_num: the number of entities in the dataset. @@ -22,6 +26,14 @@ 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: + >>> class _Toy: + ... entity_num = 2 + ... relation_num = 1 + ... task_spec_param = None + >>> model = KGEBaseModel(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> model.e_num, tuple(model.E_emb.shape) + (2, (2, 4)) """ @property @@ -32,16 +44,16 @@ def device(self): def __init__( self, - dataset: SampleBaseDataset, + dataset: KGDatasetProtocol, e_dim: int = 500, r_dim: int = 500, ns: str = "uniform", - gamma: float = None, + gamma: float | None = None, use_subsampling_weight: bool = False, - use_regularization: str = None, + use_regularization: str | None = None, mode: str = "multiclass" ): - super(KGEBaseModel, self).__init__() + super().__init__() self.e_num = dataset.entity_num self.r_num = dataset.relation_num self.e_dim = e_dim @@ -392,7 +404,6 @@ def from_pretrained(self, path): state_dict = torch.load(path, map_location=self.device, weights_only=True) self.update_embedding_size(state_dict) self.load_state_dict(state_dict) - return def update_embedding_size(self, state_dict): e_emb_key = 'E_emb' @@ -408,7 +419,6 @@ def update_embedding_size(self, state_dict): self.E_emb = nn.Parameter(torch.zeros(self.e_num, self.e_dim)) self.R_emb = nn.Parameter(torch.zeros(self.r_num, self.r_dim)) - return diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index df7143a6e..cf7a4d004 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,25 +1,40 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch +from .kg_base import KGEBaseModel + +if TYPE_CHECKING: + from ..datasets.protocols import KGDatasetProtocol + class RotatE(KGEBaseModel): - """ RotatE + """RotatE - Paper: Sun, Z., Deng, Z.H., Nie, J.Y. and Tang, J., 2019. + 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: + >>> class _Toy: + ... entity_num = 2 + ... relation_num = 1 + ... task_spec_param = None + >>> model = RotatE(_Toy(), e_dim=4, r_dim=2, ns="uniform") + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: KGDatasetProtocol, e_dim: int = 600, r_dim: int = 300, ns='adv', gamma=24.0, use_subsampling_weight: bool = False, - use_regularization: str = None, + use_regularization: str | None = None, mode: str = "multiclass" ): super().__init__(dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, use_regularization, mode) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index fbb6e68f6..21125fb3b 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,25 +1,40 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from __future__ import annotations + +from typing import TYPE_CHECKING + import torch +from .kg_base import KGEBaseModel + +if TYPE_CHECKING: + from ..datasets.protocols import KGDatasetProtocol + class TransE(KGEBaseModel): - """ TransE + """TransE Paper: Bordes, A., Usunier, N., Garcia-Duran, A., Weston, J. and Yakhnenko, Translating embeddings for modeling multi-relational data. NIPS 2013. + Examples: + >>> class _Toy: + ... entity_num = 2 + ... relation_num = 1 + ... task_spec_param = None + >>> model = TransE(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: KGDatasetProtocol, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", gamma: float = 24.0, use_subsampling_weight: bool = False, - use_regularization: str = None, + use_regularization: str | None = None, mode: str = "multiclass", p_norm: int = 1.0 ): From 587d9d63872bc6fc14334843e6adbcdede6dabc4 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:02:03 +0200 Subject: [PATCH 03/10] fix(kg_emb): repair the module examples' dataset import SampleKGDataset was never exported from pyhealth.datasets. The examples now import it from kg_emb.datasets and build a torch DataLoader with collate_fn_dict_with_padding, because get_dataloader requires litdata.StreamingDataset.set_shuffle(). Co-authored-by: Cursor --- .../kg_emb/examples/train_kge_model.py | 15 +++++++--- .../kg_emb/models/complex.py | 28 ++++++++++++++----- .../kg_emb/models/distmult.py | 28 ++++++++++++++----- .../kg_emb/models/rotate.py | 28 ++++++++++++++----- .../kg_emb/models/transe.py | 28 ++++++++++++++----- 5 files changed, 95 insertions(+), 32 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py index 53eadb7f1..fe8f40c9e 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py @@ -1,6 +1,7 @@ from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import UMLSDataset, split from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import link_prediction_fn -from pyhealth.datasets import get_dataloader +from torch.utils.data import DataLoader +from pyhealth.datasets import collate_fn_dict_with_padding from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE, RotatE, ComplEx, DistMult from pyhealth.trainer import Trainer from pyhealth.medcode import InnerMap @@ -41,9 +42,15 @@ # split the dataset and get the dataloaders train_dataset, val_dataset, test_dataset = split(umls_ds, [0.9, 0.05, 0.05]) -train_loader = get_dataloader(train_dataset, batch_size=8, shuffle=True) -# val_loader = get_dataloader(val_dataset, batch_size=2, shuffle=False) -# test_loader = get_dataloader(test_dataset, batch_size=2, shuffle=False) +train_loader = DataLoader( + train_dataset, batch_size=8, shuffle=True, collate_fn=collate_fn_dict_with_padding +) +# val_loader = DataLoader( +# val_dataset, batch_size=2, shuffle=False, collate_fn=collate_fn_dict_with_padding +# ) +# test_loader = DataLoader( +# test_dataset, batch_size=2, shuffle=False, collate_fn=collate_fn_dict_with_padding +# ) # initialize a KGE model diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 28b0d78aa..7dce35567 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -80,7 +80,12 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + 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, + ) samples = [ { @@ -97,13 +102,22 @@ def calc(self, head, relation, tail, mode='pos'): }, ] - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader + for sample in samples: + sample["train"] = True + sample["hyperparameters"] = {"negative_sampling": 8} - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + dataset = SampleKGDataset( + samples=samples, + dataset_name="test", + entity_num=8000, + relation_num=8, + ) + train_loader = DataLoader( + dataset, + batch_size=2, + shuffle=True, + collate_fn=collate_fn_dict_with_padding, + ) # model model = ComplEx( diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index dd2973328..71d47ca76 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -62,7 +62,12 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + 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, + ) samples = [ { @@ -79,13 +84,22 @@ def calc(self, head, relation, tail, mode='pos'): }, ] - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader + for sample in samples: + sample["train"] = True + sample["hyperparameters"] = {"negative_sampling": 8} - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + dataset = SampleKGDataset( + samples=samples, + dataset_name="test", + entity_num=8000, + relation_num=8, + ) + train_loader = DataLoader( + dataset, + batch_size=2, + shuffle=True, + collate_fn=collate_fn_dict_with_padding, + ) # model model = DistMult( diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index cf7a4d004..7ce74fb33 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -76,7 +76,12 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + 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, + ) samples = [ { @@ -93,13 +98,22 @@ def calc(self, head, relation, tail, mode='pos'): }, ] - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader + for sample in samples: + sample["train"] = True + sample["hyperparameters"] = {"negative_sampling": 8} - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + dataset = SampleKGDataset( + samples=samples, + dataset_name="test", + entity_num=8000, + relation_num=8, + ) + train_loader = DataLoader( + dataset, + batch_size=2, + shuffle=True, + collate_fn=collate_fn_dict_with_padding, + ) # model model = RotatE( diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 21125fb3b..d714248e0 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -65,7 +65,12 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + 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, + ) samples = [ { @@ -82,13 +87,22 @@ def calc(self, head, relation, tail, mode='pos'): }, ] - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader + for sample in samples: + sample["train"] = True + sample["hyperparameters"] = {"negative_sampling": 8} - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + dataset = SampleKGDataset( + samples=samples, + dataset_name="test", + entity_num=8000, + relation_num=8, + ) + train_loader = DataLoader( + dataset, + batch_size=2, + shuffle=True, + collate_fn=collate_fn_dict_with_padding, + ) # model model = TransE( From 83af7af4b5f44daa66367a3766a0898b7e037c4c Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:03:09 +0200 Subject: [PATCH 04/10] refactor(kg_emb): make split() side-effect free and reproducible Validate ratios with ValueError so the check survives python -O, and shuffle with a local Generator so the function no longer mutates global NumPy state. Co-authored-by: Cursor --- .../kg_emb/datasets/__init__.py | 2 + .../kg_emb/datasets/splitter.py | 122 +++++++++++++----- 2 files changed, 89 insertions(+), 35 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 15df00d9f..954f71490 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -5,10 +5,12 @@ from .sample_kg_dataset import SampleKGDataset from .base_kg_dataset import BaseKGDataset from .umls import UMLSDataset +from .splitter import split __all__ = [ "BaseKGDataset", "KGDatasetProtocol", "SampleKGDataset", "UMLSDataset", + "split", ] diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 559ecb7c4..a583f5f5d 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -1,48 +1,100 @@ -from itertools import chain -from typing import Optional, Tuple, Union, List +"""Ratio-based splitting of a :class:`SampleKGDataset` into train/val/test folds.""" + +from __future__ import annotations + +import math +from typing import Any import numpy as np -import torch -from pyhealth.datasets import SampleBaseDataset +from .sample_kg_dataset import SampleKGDataset + +__all__ = ["split"] + +Fold = list[dict[str, Any]] def split( - dataset: SampleBaseDataset, - ratios: Union[Tuple[float, float, float], List[float]], - seed: Optional[int] = None, -): - """Splits the dataset by its outermost indexed items + dataset: SampleKGDataset, + ratios: list[float] | tuple[float, float, float], + seed: int | None = None, +) -> tuple[Fold, Fold, Fold]: + """Split a KG sample dataset into three disjoint folds. + + The split is uniform over triples: each sample is assigned to exactly one + fold, so the three folds partition the dataset. Training samples carry the + task hyper-parameters needed by the negative sampler; validation and test + samples are flagged so that the model switches to filtered ranking + evaluation. Args: - dataset: a `SampleBaseDataset` object - ratios: a list/tuple of ratios for train / val / test - seed: random seed for shuffling the dataset + dataset: The dataset to split. + ratios: Three non-negative floats summing to 1, in train/val/test + order. + seed: Seed of a local random generator. The global NumPy state is + left untouched, which keeps the function reproducible without + side effects. Returns: - train_dataset, val_dataset, test_dataset: three subsets of the dataset of - type `torch.utils.data.Subset`. + The train, validation and test folds, each a list of sample + dictionaries. - Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, and `test_dataset.dataset`. + Raises: + ValueError: If ``ratios`` is malformed. Validation of user input is + raised rather than asserted, because ``assert`` statements are + stripped under ``python -O`` and this check must survive + optimised runs. The tolerance comparison guards the rare + triplets -- about 0.9% of two-decimal ratios -- for which + floating-point summation does not land exactly on 1. + + Examples: + >>> import torch + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + ... SampleKGDataset, + ... ) + >>> samples = [ + ... { + ... "triple": (i, i % 2, (i + 1) % 5), + ... "ground_truth_head": [i, (i + 1) % 5], + ... "ground_truth_tail": [(i + 1) % 5], + ... "subsampling_weight": torch.tensor([0.25]), + ... } + ... for i in range(10) + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=5, relation_num=2, negative_sampling=4 + ... ) + >>> train, val, test = split(dataset, [0.6, 0.2, 0.2], seed=0) + >>> len(train), len(val), len(test) + (6, 2, 2) + >>> train[0]["train"], train[0]["hyperparameters"] + (True, {'negative_sampling': 4}) + >>> val[0]["train"] + False + >>> split(dataset, [0.5, 0.2, 0.2], seed=0) + Traceback (most recent call last): + ... + ValueError: ratios must sum to 1.0, got 0.9 """ - - if seed is not None: - np.random.seed(seed) - assert sum(ratios) == 1.0, "ratios must sum to 1.0" - index = np.arange(len(dataset)) - np.random.shuffle(index) - train_index = index[: int(len(dataset) * ratios[0])] - val_index = index[ - int(len(dataset) * ratios[0]) : int(len(dataset) * (ratios[0] + ratios[1])) + if len(ratios) != 3 or any(r < 0 for r in ratios): + raise ValueError(f"ratios must be three non-negative floats, got {ratios!r}") + total = sum(ratios) + if not math.isclose(total, 1.0, rel_tol=0.0, abs_tol=1e-9): + raise ValueError(f"ratios must sum to 1.0, got {total}") + + rng = np.random.default_rng(seed) + n = len(dataset) + index = rng.permutation(n) + + n_train = int(n * ratios[0]) + n_val = int(n * (ratios[0] + ratios[1])) + slices = (index[:n_train], index[n_train:n_val], index[n_val:]) + + hyperparameters = dataset.task_spec_param + train = [ + {**dataset[int(i)], "train": True, "hyperparameters": hyperparameters} + for i in slices[0] ] - test_index = index[int(len(dataset) * (ratios[0] + ratios[1])) :] - train_dataset = torch.utils.data.Subset(dataset, train_index) - train_dataset = [{**train_dataset[i], **{'train': True, 'hyperparameters': dataset.task_spec_param}} for i in range(len(train_dataset))] - - val_dataset = torch.utils.data.Subset(dataset, val_index) - val_dataset = [{**val_dataset[i], 'train': False} for i in range(len(val_dataset))] - test_dataset = torch.utils.data.Subset(dataset, test_index) - test_dataset = [{**test_dataset[i], 'train': False} for i in range(len(test_dataset))] - return train_dataset, val_dataset, test_dataset \ No newline at end of file + val = [{**dataset[int(i)], "train": False} for i in slices[1]] + test = [{**dataset[int(i)], "train": False} for i in slices[2]] + return train, val, test From ca7c070195245a2dc311b6edc6b408cdc6ec7301 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 25 Aug 2026 10:46:24 +0200 Subject: [PATCH 05/10] chore(kg_emb): remove dead imports of the undeclared pandarallel package pandarallel was never declared in pyproject.toml or pixi.lock. The undeclared import in umls.py and base_kg_dataset.py is what produced the ModuleNotFoundError on the kg_emb import path in issue #952. initialize() ran in umls.py with no parallel_apply in kg_emb; mimicextract's parallel_apply calls are unreachable on the empty BaseEHRDataset stub. There is no lockfile entry to regenerate. Co-authored-by: Cursor --- .../kg_emb/datasets/__init__.py | 7 ++--- .../kg_emb/datasets/base_kg_dataset.py | 26 ++++++++++++------- .../kg_emb/datasets/umls.py | 19 +++++++++----- 3 files changed, 30 insertions(+), 22 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 954f71490..0dee767bd 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,11 +1,8 @@ -# ruff: noqa: I001 -# BaseKGDataset imports SampleKGDataset from this package, so the sample -# dataset must be bound before the base class is loaded. +from .base_kg_dataset import BaseKGDataset from .protocols import KGDatasetProtocol from .sample_kg_dataset import SampleKGDataset -from .base_kg_dataset import BaseKGDataset -from .umls import UMLSDataset from .splitter import split +from .umls import UMLSDataset __all__ = [ "BaseKGDataset", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py index c368fedc6..6070ad178 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py +++ b/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: diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 22d3cbb5e..8180a9b91 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -1,10 +1,10 @@ import logging import os -from tqdm import tqdm -import numpy as np + import pandas as pd -from pandarallel import pandarallel -from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import BaseKGDataset +from tqdm import tqdm + +from .base_kg_dataset import BaseKGDataset logger = logging.getLogger(__name__) @@ -20,11 +20,17 @@ class UMLSDataset(BaseKGDataset): 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: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + ... BaseKGDataset, + ... UMLSDataset, + ... ) + >>> issubclass(UMLSDataset, BaseKGDataset) + True """ def raw_graph_process(self): - pandarallel.initialize(progress_bar=False) if self.dev == False: self.graph_path = os.path.join(self.root, "graph.txt") else: @@ -56,7 +62,6 @@ def raw_graph_process(self): for e1, r, e2 in tqdm(zip(graph_df['e1'], graph_df['r'], graph_df['e2']), total=graph_df.shape[0])] - return if __name__ == "__main__": From 4019acb1f041d3bf4cbe260367cbf319b0f3bcd8 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 24 Aug 2026 23:07:27 +0200 Subject: [PATCH 06/10] test(kg_emb): add behavioural regression tests for issue #952 Cover construction, split reproducibility, generic collation of variable-length ground truths, set_task on a synthetic graph, and scoring invariants. Tests instantiate SampleKGDataset so a rename-only fix cannot go green. Co-authored-by: Cursor --- tests/core/test_kg_emb.py | 244 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 244 insertions(+) create mode 100644 tests/core/test_kg_emb.py diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py new file mode 100644 index 000000000..028a72476 --- /dev/null +++ b/tests/core/test_kg_emb.py @@ -0,0 +1,244 @@ +"""Regression tests for ``pyhealth.medcode.pretrained_embeddings.kg_emb``. + +The suite is behavioural: it exercises construction, indexing and splitting +rather than asserting on type annotations, which are metadata and not a +contract. +""" + +from __future__ import annotations + +import tempfile +import unittest +from typing import Any + +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 ( + BaseKGDataset, + SampleKGDataset, + split, +) +from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import link_prediction_fn + + +def make_samples(n: int = 8) -> list[dict[str, Any]]: + """Build ``n`` synthetic link-prediction samples over a 10-entity graph.""" + return [ + { + "triple": (i % 10, i % 3, (i + 4) % 10), + "ground_truth_head": [i % 10, (i + 1) % 10], + "ground_truth_tail": [(i + 4) % 10], + "subsampling_weight": torch.tensor([0.25]), + } + for i in range(n) + ] + + +def make_dataset(n: int = 8, **kwargs: Any) -> SampleKGDataset: + entity2id = {f"e{i}": i for i in range(10)} + relation2id = {f"r{i}": i for i in range(3)} + return SampleKGDataset( + samples=make_samples(n), + dataset_name="synthetic", + task_name="link_prediction", + entity2id=entity2id, + relation2id=relation2id, + negative_sampling=4, + **kwargs, + ) + + +class TestKGEmbImports(unittest.TestCase): + """The module must import cleanly -- the original symptom of issue #952.""" + + def test_package_imports(self) -> None: + import pyhealth.medcode.pretrained_embeddings # noqa: F401 + + def test_model_classes_are_exported(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + KGEBaseModel, + RotatE, + TransE, + ) + + for cls in (KGEBaseModel, TransE, RotatE, DistMult, ComplEx): + self.assertTrue(issubclass(cls, torch.nn.Module)) + + +class TestSampleKGDataset(unittest.TestCase): + """Construction and indexing -- the failure the rename alone does not fix.""" + + def test_construction_and_length(self) -> None: + dataset = make_dataset(n=8) + self.assertEqual(len(dataset), 8) + self.assertEqual(dataset.entity_num, 10) + self.assertEqual(dataset.relation_num, 3) + + def test_getitem_returns_the_sample(self) -> None: + dataset = make_dataset(n=3) + self.assertEqual(dataset[0]["triple"], (0, 0, 4)) + self.assertIn("ground_truth_head", dataset[1]) + + def test_inverse_vocabularies(self) -> None: + dataset = make_dataset(n=2) + self.assertEqual(dataset.id2entity[0], "e0") + self.assertEqual(dataset.id2relation[2], "r2") + + def test_task_specific_hyperparameters_are_captured(self) -> None: + dataset = make_dataset(n=2) + self.assertEqual(dataset.task_spec_param, {"negative_sampling": 4}) + + def test_missing_vocabularies_do_not_crash(self) -> None: + dataset = SampleKGDataset( + samples=make_samples(2), entity_num=10, relation_num=3 + ) + self.assertEqual(dataset.id2entity, {}) + self.assertIsNone(dataset.task_spec_param) + + def test_contradictory_cardinalities_are_rejected(self) -> None: + with self.assertRaises(ValueError): + SampleKGDataset( + samples=make_samples(1), + entity_num=99, + entity2id={f"e{i}": i for i in range(10)}, + ) + + def test_stat_returns_a_report(self) -> None: + report = make_dataset(n=2).stat() + self.assertIn("Number of triples: 2", report) + + def test_is_a_map_style_dataset(self) -> None: + dataset = make_dataset(n=2) + self.assertIsInstance(dataset, torch.utils.data.Dataset) + self.assertFalse(hasattr(dataset, "set_shuffle")) + + +class TestSplit(unittest.TestCase): + """The splitter must partition the dataset and stay reproducible.""" + + def test_partition_sizes(self) -> None: + train, val, test = split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=0) + self.assertEqual((len(train), len(val), len(test)), (6, 2, 2)) + + def test_folds_are_disjoint_and_exhaustive(self) -> None: + train, val, test = split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=0) + triples = [s["triple"] for s in train + val + test] + self.assertEqual(len(triples), 10) + self.assertEqual(len(set(triples)), 10) + + def test_is_reproducible_under_a_fixed_seed(self) -> None: + first = split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=7)[0] + second = split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=7)[0] + self.assertEqual([s["triple"] for s in first], [s["triple"] for s in second]) + + def test_global_numpy_state_is_untouched(self) -> None: + import numpy as np + + np.random.seed(1234) + before = np.random.rand() + np.random.seed(1234) + split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=99) + self.assertEqual(before, np.random.rand()) + + def test_training_fold_carries_hyperparameters(self) -> None: + train, val, _ = split(make_dataset(n=10), [0.6, 0.2, 0.2], seed=0) + self.assertTrue(train[0]["train"]) + self.assertEqual(train[0]["hyperparameters"], {"negative_sampling": 4}) + self.assertFalse(val[0]["train"]) + + def test_malformed_ratios_are_rejected(self) -> None: + dataset = make_dataset(n=10) + for bad in ([0.5, 0.2, 0.2], [0.5, 0.5], [1.2, -0.2, 0.0]): + with self.subTest(ratios=bad), self.assertRaises(ValueError): + split(dataset, bad, seed=0) + + def test_ratio_sum_error_message(self) -> None: + with self.assertRaisesRegex(ValueError, "ratios must sum to 1.0, got 0.9"): + split(make_dataset(n=10), [0.5, 0.2, 0.2], seed=0) + + +class TestCollateAndForward(unittest.TestCase): + """Generic padding collation leaves KG lists intact; one train step runs.""" + + def test_variable_length_ground_truth_stays_a_python_list(self) -> None: + dataset = make_dataset(n=4) + train, _, _ = split(dataset, [1.0, 0.0, 0.0], seed=0) + loader = DataLoader( + train, batch_size=2, shuffle=False, collate_fn=collate_fn_dict_with_padding + ) + batch = next(iter(loader)) + self.assertIsInstance(batch["triple"], list) + self.assertIsInstance(batch["ground_truth_head"], list) + self.assertIsInstance(batch["ground_truth_head"][0], list) + lengths = [len(h) for h in batch["ground_truth_head"]] + self.assertTrue(all(length >= 1 for length in lengths)) + + def test_transe_train_step(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + dataset = make_dataset(n=4) + train, _, _ = split(dataset, [1.0, 0.0, 0.0], seed=0) + loader = DataLoader( + train, batch_size=2, shuffle=False, collate_fn=collate_fn_dict_with_padding + ) + model = TransE(dataset=dataset, e_dim=8, r_dim=8, ns="uniform") + out = model(**next(iter(loader))) + self.assertIn("loss", out) + out["loss"].backward() + + +class TestSetTask(unittest.TestCase): + """Production path: BaseKGDataset.set_task must return a usable SampleKGDataset.""" + + def test_set_task_on_a_synthetic_graph(self) -> None: + class _ToyKG(BaseKGDataset): + def raw_graph_process(self): + self.entity2id = {"a": 0, "b": 1, "c": 2} + self.relation2id = {"r": 0} + self.entity_num = 3 + self.relation_num = 1 + self.triples = [(0, 0, 1), (1, 0, 2), (2, 0, 0)] + + with tempfile.TemporaryDirectory() as root: + base = _ToyKG(root=root, dataset_name="toy", refresh_cache=True) + sample_ds = base.set_task( + link_prediction_fn, negative_sampling=4, save=False + ) + self.assertIsInstance(sample_ds, SampleKGDataset) + self.assertEqual(len(sample_ds), 3) + self.assertEqual(sample_ds.task_spec_param, {"negative_sampling": 4}) + self.assertIn("triple", sample_ds[0]) + + +class TestScoringInvariants(unittest.TestCase): + """Mathematical properties the scoring functions must satisfy by construction.""" + + def test_distmult_is_symmetric_in_head_and_tail(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + + model = DistMult(dataset=make_dataset(n=4), e_dim=8, r_dim=8, gamma=12.0) + head, relation, tail = (torch.randn(2, 1, 8) for _ in range(3)) + self.assertTrue( + torch.allclose( + model.calc(head, relation, tail), + model.calc(tail, relation, head), + atol=1e-6, + ) + ) + + def test_transe_scores_a_perfect_triple_at_the_margin(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + model = TransE(dataset=make_dataset(n=4), e_dim=8, r_dim=8, gamma=24.0) + head = torch.zeros(1, 1, 8) + relation = torch.ones(1, 1, 8) + tail = torch.ones(1, 1, 8) # h + r - t == 0 + self.assertTrue( + torch.allclose( + model.calc(head, relation, tail), torch.tensor(24.0), atol=1e-6 + ) + ) From 58033889249d331f35209718e2658d49715af5d2 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:13:10 +0200 Subject: [PATCH 07/10] docs(kg_emb): satisfy contribution-rules Document the map-style SampleKGDataset path in the MedCode API page and add a synthetic TransE example that uses DataLoader instead of get_dataloader. Co-authored-by: Cursor --- docs/api/medcode.rst | 44 +++++++++++++++++++++++ examples/kg_emb_sample_dataset.py | 58 +++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+) create mode 100644 examples/kg_emb_sample_dataset.py diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index bfc4edcea..c5b691af9 100644 --- a/docs/api/medcode.rst +++ b/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: diff --git a/examples/kg_emb_sample_dataset.py b/examples/kg_emb_sample_dataset.py new file mode 100644 index 000000000..d2bfe6a34 --- /dev/null +++ b/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") From 1e7d1b03c6dcfab5b9757fbab36b92d109caf9a4 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 25 Aug 2026 11:15:19 +0200 Subject: [PATCH 08/10] style(kg_emb): tidy docstring punctuation Replace double-hyphen asides in five docstrings with periods or commas so they remain readable in a terminal and under Sphinx. Co-authored-by: Cursor --- .../pretrained_embeddings/kg_emb/datasets/protocols.py | 2 +- .../kg_emb/datasets/sample_kg_dataset.py | 2 +- .../medcode/pretrained_embeddings/kg_emb/datasets/splitter.py | 2 +- tests/core/test_kg_emb.py | 4 ++-- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py index 6b093bcc6..bf17d099a 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py @@ -13,7 +13,7 @@ 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 + 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. 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 778cc88a8..1dea1bff9 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 @@ -148,7 +148,7 @@ def __getitem__(self, index: int) -> KGSample: return self.samples[index] def stat(self) -> str: - """Return -- and print -- a human-readable summary of the dataset.""" + """Print a human-readable summary and return it.""" lines = [ "", f"Statistics of sample KG dataset (dev={self.dev}):", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index a583f5f5d..0ad096461 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -44,7 +44,7 @@ def split( raised rather than asserted, because ``assert`` statements are stripped under ``python -O`` and this check must survive optimised runs. The tolerance comparison guards the rare - triplets -- about 0.9% of two-decimal ratios -- for which + triplets, about 0.9% of two-decimal ratios, for which floating-point summation does not land exactly on 1. Examples: diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 028a72476..4c664311a 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -51,7 +51,7 @@ def make_dataset(n: int = 8, **kwargs: Any) -> SampleKGDataset: class TestKGEmbImports(unittest.TestCase): - """The module must import cleanly -- the original symptom of issue #952.""" + """The module must import cleanly. This was the original symptom of issue #952.""" def test_package_imports(self) -> None: import pyhealth.medcode.pretrained_embeddings # noqa: F401 @@ -70,7 +70,7 @@ def test_model_classes_are_exported(self) -> None: class TestSampleKGDataset(unittest.TestCase): - """Construction and indexing -- the failure the rename alone does not fix.""" + """Construction and indexing: the failure a rename alone does not fix.""" def test_construction_and_length(self) -> None: dataset = make_dataset(n=8) From 53c2a4ee6422bd64d6c1169a816a157bfe3ce12b Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:47:30 +0200 Subject: [PATCH 09/10] fix(kg_emb): make task_spec_param covariant in KGDatasetProtocol SampleKGDataset failed to satisfy its own KGDatasetProtocol under mypy: the Protocol declared task_spec_param as a plain attribute (Mapping[str, Any] | None), which Protocol treats as read-write and therefore invariant, while SampleKGDataset declares it as dict[str, Any] | None. Models only ever read task_spec_param, so declare it as a read-only property instead: read-only Protocol members are covariant, and a concrete dict satisfies it. --- .../pretrained_embeddings/kg_emb/datasets/protocols.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py index bf17d099a..e335daa91 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py @@ -29,4 +29,6 @@ class KGDatasetProtocol(Protocol): entity_num: int relation_num: int - task_spec_param: Mapping[str, Any] | None + + @property + def task_spec_param(self) -> Mapping[str, Any] | None: ... From 319af2d34f7e35683228728efbaba3f46071bfa9 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Sat, 29 Aug 2026 16:47:38 +0200 Subject: [PATCH 10/10] fix(kg_emb): annotate __main__ demo samples for mypy Each model's __main__ block builds an untyped list of dict literals and then adds a "train" key with a bool value, which mypy rejects because it infers the dict's value type from the first literal. Annotate samples as list[dict[str, Any]] in all four demo blocks. --- .../medcode/pretrained_embeddings/kg_emb/models/complex.py | 4 ++-- .../medcode/pretrained_embeddings/kg_emb/models/distmult.py | 4 ++-- .../medcode/pretrained_embeddings/kg_emb/models/rotate.py | 4 ++-- .../medcode/pretrained_embeddings/kg_emb/models/transe.py | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 7dce35567..ee8837745 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -87,7 +87,7 @@ def calc(self, head, relation, tail, mode='pos'): SampleKGDataset, ) - samples = [ + samples: list[dict[str, Any]] = [ { 'triple': (0, 0, 2835), 'ground_truth_head': [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029], diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index 71d47ca76..e6f38d4ff 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -69,7 +69,7 @@ def calc(self, head, relation, tail, mode='pos'): SampleKGDataset, ) - samples = [ + samples: list[dict[str, Any]] = [ { 'triple': (0, 0, 2835), 'ground_truth_head': [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029], diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index 7ce74fb33..1c352e72a 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -83,7 +83,7 @@ def calc(self, head, relation, tail, mode='pos'): SampleKGDataset, ) - samples = [ + samples: list[dict[str, Any]] = [ { 'triple': (0, 0, 2835), 'ground_truth_head': [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029], diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index d714248e0..5e6287a32 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -72,7 +72,7 @@ def calc(self, head, relation, tail, mode='pos'): SampleKGDataset, ) - samples = [ + samples: list[dict[str, Any]] = [ { 'triple': (0, 0, 2835), 'ground_truth_head': [1027, 1293, 5264, 1564, 7416, 6434, 2610, 4094, 2717, 5007, 5277, 5949, 0, 6870, 6029],