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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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/12] 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], From 5f66ff48da31a4634b6b17c2c2abab5da9ac76a9 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:44:18 +0200 Subject: [PATCH 11/12] feat(kg_emb): rebase SampleKGDataset onto InMemorySampleDataset with the Tensor Trick Per the architecture pivot agreed in the PR discussion: SampleKGDataset moves back under InMemorySampleDataset instead of standalone torch.utils.data.Dataset, while keeping its full public surface (entity2id/relation2id, cardinality validation, split(), stat(), dev/task_spec_param) unchanged. - Add KGProcessor ("kg_entity_list"), pre-padding ground_truth_head/tail to each field's own max length and emitting {"value", "mask"} pure tensors ahead of litdata's pickle-based caching, instead of raw variable-length Python lists. "triple" goes through the existing "tensor" processor. - Fix a correctness issue the padding introduces: pad_token_id (0) is not a reserved sentinel and can collide with a real entity id. kg_base.py's train_neg_sample_gen and test_neg_sample_filter_bias_gen now reconstruct the exact unpadded entity list via the mask (_unpad_ground_truth) before doing set-membership filtering, so negative sampling and filtered ranking stay correct whenever entity 0 is legitimate. - Add a nested-dict collation branch to collate_fn_dict_with_padding, restricted to all-tensor dicts, so {"value","mask"} pairs batch via a plain stack (shape is already uniform per field) without disturbing the existing list-of-dicts collation used by heterogeneous per-sample dicts such as "hyperparameters". - Update tests/core/test_kg_emb.py for the new shapes: triple is now a Tensor (not a tuple), ground_truth_* collate to {"value","mask"}, and set_shuffle is now expected (SampleKGDataset is intentionally back under the SampleDataset umbrella). Adds TestGroundTruthUnpadding, a regression test for the padding/entity-0 collision fix above. 24/24 tests in tests/core/test_kg_emb.py pass, plus the sample_kg_dataset.py and splitter.py doctests. Co-Authored-By: Claude Sonnet 5 --- pyhealth/datasets/utils.py | 24 ++++ .../kg_emb/datasets/sample_kg_dataset.py | 58 +++++---- .../kg_emb/models/kg_base.py | 56 ++++++++- pyhealth/processors/__init__.py | 2 + pyhealth/processors/kg_processor.py | 110 ++++++++++++++++++ tests/core/test_kg_emb.py | 73 ++++++++++-- 6 files changed, 288 insertions(+), 35 deletions(-) create mode 100644 pyhealth/processors/kg_processor.py diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 24c87a1d5..e960f2dd7 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -304,6 +304,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] 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 1dea1bff9..9d716fefe 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,12 +1,19 @@ """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. +``SampleKGDataset`` builds on :class:`pyhealth.datasets.InMemorySampleDataset` +so that KG triples and the variable-length ``ground_truth_head`` / +``ground_truth_tail`` filter sets are converted to pure PyTorch tensors ahead +of ``litdata``'s pickle-based caching (the "Tensor Trick"), instead of being +serialized as raw Python lists on every access. + +``ground_truth_head`` and ``ground_truth_tail`` are padded independently, each +to its own field's observed maximum length, by the registered +``"kg_entity_list"`` processor +(:class:`~pyhealth.processors.kg_processor.KGProcessor`), which returns a +``{"value": Tensor, "mask": Tensor}`` pair. Because the padding value is not +a valid entity id on its own, any code that filters on these fields (see +:class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.models.kg_base.KGEBaseModel`) +must use the mask to recover the true, unpadded entity list first. """ from __future__ import annotations @@ -14,14 +21,16 @@ from collections.abc import Mapping, Sequence from typing import Any -from torch.utils.data import Dataset +import torch + +from pyhealth.datasets.sample_dataset import InMemorySampleDataset KGSample = Mapping[str, Any] __all__ = ["SampleKGDataset"] -class SampleKGDataset(Dataset): +class SampleKGDataset(InMemorySampleDataset): r"""In-memory dataset of knowledge-graph link-prediction samples. Each sample is a mapping with the following keys: @@ -54,6 +63,8 @@ class SampleKGDataset(Dataset): index. relation2id: Mapping from surface relation identifier to integer index. + pad_token_id: Entity id used to pad ``ground_truth_head`` / + ``ground_truth_tail`` to their fitted per-field max length. **task_spec_param: Task hyper-parameters forwarded to the model at training time (e.g. ``negative_sampling=128``). @@ -82,8 +93,6 @@ class SampleKGDataset(Dataset): ... ) >>> len(dataset) 10 - >>> dataset[0]["triple"] - (0, 0, 1) >>> dataset.entity_num, dataset.relation_num (5, 2) >>> dataset.id2entity[2] @@ -102,11 +111,22 @@ def __init__( relation_num: int = 0, entity2id: Mapping[Any, int] | None = None, relation2id: Mapping[Any, int] | None = None, + pad_token_id: int = 0, **task_spec_param: Any, ) -> None: - self.samples: list[KGSample] = list(samples) - self.dataset_name = dataset_name - self.task_name = task_name + input_schema = { + "triple": ("tensor", {"dtype": torch.long}), + "ground_truth_head": ("kg_entity_list", {"pad_token_id": pad_token_id}), + "ground_truth_tail": ("kg_entity_list", {"pad_token_id": pad_token_id}), + } + super().__init__( + samples=list(samples), + input_schema=input_schema, + output_schema={}, + dataset_name=dataset_name, + task_name=task_name, + ) + self.dev = dev self.entity2id: dict[Any, int] = dict(entity2id or {}) @@ -139,13 +159,7 @@ def _validate_cardinalities(self) -> None: @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] + return len(self) def stat(self) -> str: """Print a human-readable summary and return it.""" @@ -154,7 +168,7 @@ def stat(self) -> str: 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 triples: {len(self)}", 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}", 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 6de31761c..eaa535fbf 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -147,12 +147,54 @@ def data_process(self, sample_batch, mode): return head, relation, tail + @staticmethod + def _unpad_ground_truth(ground_truth): + """Recover the exact, unpadded per-sample entity-id lists. + + ``KGProcessor`` pads ``ground_truth_head``/``ground_truth_tail`` to a + fixed length with ``pad_token_id`` (0 by default) so they can be + tensorized ahead of serialization. That padding value is not + necessarily an invalid entity id, so it must be stripped via the + accompanying mask before doing any set-membership filtering here; + otherwise a real entity 0 would be spuriously treated as always + "known true" (or, symmetrically, padding would be treated as a real + entity to exclude/replace). + + Args: + ground_truth: Either a ``{"value": Tensor(B, L), "mask": Tensor(B, L)}`` + pair (collated ``KGProcessor`` output), or already a list of + raw per-sample entity-id lists (e.g. when a caller bypasses + the processor and supplies unpadded lists directly). + + Returns: + List of length ``B``, each entry the unpadded list of entity ids + for that sample. + """ + if isinstance(ground_truth, dict): + value, mask = ground_truth["value"], ground_truth["mask"] + return [ + value[i][mask[i].bool()].tolist() for i in range(value.size(0)) + ] + return ground_truth + def train_neg_sample_gen(self, gt_head, gt_tail, negative_sampling): """ - (only run in train batch) + (only run in train batch) This function creates negative triples for training (sampling size: negative_sampling) with ground truth masked. + + Args: + gt_head: Either a list of raw (unpadded) entity-id lists, or a + ``{"value": Tensor(B, L), "mask": Tensor(B, L)}`` pair produced + by ``KGProcessor``. The padded ``value`` is not usable on its + own for membership filtering, since ``pad_token_id`` may + collide with a real entity id (e.g. 0); the ``mask`` recovers + the exact unpadded list first. + gt_tail: Same shape as ``gt_head``, for tail entities. """ + gt_head = self._unpad_ground_truth(gt_head) + gt_tail = self._unpad_ground_truth(gt_tail) + negative_sample_head = [] negative_sample_tail = [] for i in range(len(gt_head)): @@ -203,9 +245,19 @@ def train_neg_sample_gen(self, gt_head, gt_tail, negative_sampling): def test_neg_sample_filter_bias_gen(self, triples, gt_head, gt_tail): """ - (only run in val/test batch) + (only run in val/test batch) This function creates negative triples for validation/testing with ground truth masked. + + Args: + triples: Batch of ``(head, relation, tail)`` triples. + gt_head: Either a list of raw (unpadded) entity-id lists, or a + ``{"value": Tensor(B, L), "mask": Tensor(B, L)}`` pair produced + by ``KGProcessor``. See ``_unpad_ground_truth`` for why the + mask matters. + gt_tail: Same shape as ``gt_head``, for tail entities. """ + gt_head = self._unpad_ground_truth(gt_head) + gt_tail = self._unpad_ground_truth(gt_tail) negative_sample_head = [] negative_sample_tail = [] diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..8e52b1223 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -46,6 +46,7 @@ def get_processor(name: str): from .timeseries_processor import TimeseriesProcessor from .time_image_processor import TimeImageProcessor from .graph_processor import GraphProcessor +from .kg_processor import KGProcessor from .audio_processor import AudioProcessor from .ignore_processor import IgnoreProcessor from .temporal_timeseries_processor import TemporalTimeseriesProcessor @@ -78,6 +79,7 @@ def get_processor(name: str): "TimeseriesProcessor", "TimeImageProcessor", "GraphProcessor", + "KGProcessor", "AudioProcessor", "TupleTimeTextProcessor", "CehrProcessor", diff --git a/pyhealth/processors/kg_processor.py b/pyhealth/processors/kg_processor.py new file mode 100644 index 000000000..0b1bfc9a6 --- /dev/null +++ b/pyhealth/processors/kg_processor.py @@ -0,0 +1,110 @@ +from typing import Any, Dict, Iterable, List + +import torch + +from . import register_processor +from .base_processor import FeatureProcessor + + +@register_processor("kg_entity_list") +class KGProcessor(FeatureProcessor): + """Pads a variable-length list of knowledge-graph entity ids to a fixed length. + + Intended for fields holding lists of entity ids, such as the known-true + ``ground_truth_head`` / ``ground_truth_tail`` filter sets used for + filtered link-prediction ranking evaluation (e.g. TransE, RotatE style + KG embedding models). The target length is the maximum list length + observed for that field during ``fit``, following the same per-field + convention as :class:`~pyhealth.processors.nested_sequence_processor.NestedSequenceProcessor`. + + Note: + The padded ``pad_token_id`` is not a valid entity id on its own: any + code consuming ``ground_truth_head``/``ground_truth_tail`` must use + the accompanying ``mask`` to recover the true (unpadded) entity list + before doing membership filtering, since ``pad_token_id`` may collide + with a real entity id (e.g. 0). See + :meth:`~pyhealth.medcode.pretrained_embeddings.kg_emb.models.kg_base.KGEBaseModel.train_neg_sample_gen`. + + Args: + pad_token_id: Entity id used to pad lists shorter than ``max_length``. + Default is 0. + + Example: + >>> processor = KGProcessor(pad_token_id=0) + >>> samples = [ + ... {"ground_truth_tail": [6, 7, 8]}, + ... {"ground_truth_tail": [16]}, + ... ] + >>> processor.fit(samples, "ground_truth_tail") + >>> processor.process([16]) + {'value': tensor([16, 0, 0]), 'mask': tensor([1, 0, 0])} + """ + + def __init__(self, pad_token_id: int = 0): + self.pad_token_id = pad_token_id + self.max_length = 1 + + def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: + """Determine the maximum list length observed for ``field``. + + Args: + samples: Iterable of sample dictionaries. + field: Name of the field holding a list of entity ids. + """ + max_len = 0 + for sample in samples: + value = sample.get(field) + if value is not None: + max_len = max(max_len, len(value)) + self.max_length = max(1, max_len) + + def process(self, value: List[int]) -> Dict[str, torch.Tensor]: + """Pad a list of entity ids to ``max_length`` and build its attention mask. + + Lists longer than ``max_length`` (e.g. seen only at inference time, + after ``fit`` was called on a different split) are truncated. + + Args: + value: List of entity ids. + + Returns: + Dict with: + - ``"value"``: LongTensor of shape ``(max_length,)``, padded with + ``pad_token_id``. + - ``"mask"``: LongTensor of shape ``(max_length,)``, 1 for real + entities and 0 for padding. + """ + entities = list(value) if value is not None else [] + seq_len = len(entities) + + if seq_len >= self.max_length: + padded = entities[: self.max_length] + mask = [1] * self.max_length + else: + padded = entities + [self.pad_token_id] * (self.max_length - seq_len) + mask = [1] * seq_len + [0] * (self.max_length - seq_len) + + return { + "value": torch.tensor(padded, dtype=torch.long), + "mask": torch.tensor(mask, dtype=torch.long), + } + + def size(self) -> int: + """Return the fitted padding length.""" + return self.max_length + + def is_token(self) -> bool: + """Entity ids are discrete token indices.""" + return True + + def schema(self) -> tuple: + return ("value", "mask") + + def dim(self) -> tuple: + return (1, 1) + + def spatial(self) -> tuple: + return (False, False) + + def __repr__(self) -> str: + return f"KGProcessor(max_length={self.max_length}, pad_token_id={self.pad_token_id})" diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 4c664311a..a350cc7cb 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -80,7 +80,9 @@ def test_construction_and_length(self) -> None: def test_getitem_returns_the_sample(self) -> None: dataset = make_dataset(n=3) - self.assertEqual(dataset[0]["triple"], (0, 0, 4)) + # "triple" is now a pure LongTensor (the Tensor Trick), not the raw + # tuple, so this is a torch.equal check rather than a tuple ==. + self.assertTrue(torch.equal(dataset[0]["triple"], torch.tensor([0, 0, 4]))) self.assertIn("ground_truth_head", dataset[1]) def test_inverse_vocabularies(self) -> None: @@ -112,9 +114,16 @@ def test_stat_returns_a_report(self) -> None: self.assertIn("Number of triples: 2", report) def test_is_a_map_style_dataset(self) -> None: + """SampleKGDataset is deliberately back under the InMemorySampleDataset + umbrella (see PR discussion), so `set_shuffle` is now expected to be + present rather than absent — this supersedes the old standalone-Dataset + isolation check.""" + from pyhealth.datasets.sample_dataset import InMemorySampleDataset + dataset = make_dataset(n=2) self.assertIsInstance(dataset, torch.utils.data.Dataset) - self.assertFalse(hasattr(dataset, "set_shuffle")) + self.assertIsInstance(dataset, InMemorySampleDataset) + self.assertTrue(hasattr(dataset, "set_shuffle")) class TestSplit(unittest.TestCase): @@ -126,14 +135,19 @@ def test_partition_sizes(self) -> None: 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] + # "triple" is now a Tensor, which is neither hashable-by-value nor + # comparable the way a tuple is; compare/hash via .tolist() instead. + triples = [tuple(s["triple"].tolist()) 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]) + self.assertEqual( + [s["triple"].tolist() for s in first], + [s["triple"].tolist() for s in second], + ) def test_global_numpy_state_is_untouched(self) -> None: import numpy as np @@ -162,20 +176,30 @@ def test_ratio_sum_error_message(self) -> None: class TestCollateAndForward(unittest.TestCase): - """Generic padding collation leaves KG lists intact; one train step runs.""" + """Tensor Trick collation: KG fields arrive pre-padded, with masks; one train step runs.""" - def test_variable_length_ground_truth_stays_a_python_list(self) -> None: + def test_ground_truth_collates_to_padded_tensor_with_mask(self) -> None: + """Supersedes the old "stays a python list" expectation: since the + Tensor Trick (KGProcessor), triple/ground_truth_* are pre-padded + pure tensors by the time they leave SampleKGDataset, not raw Python + lists collated dynamically per batch.""" 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)) + + self.assertIsInstance(batch["triple"], torch.Tensor) + self.assertEqual(tuple(batch["triple"].shape), (2, 3)) + + for field in ("ground_truth_head", "ground_truth_tail"): + gt = batch[field] + self.assertIsInstance(gt, dict) + self.assertEqual(gt["value"].shape, gt["mask"].shape) + self.assertEqual(gt["value"].shape[0], 2) # batch size + # Every sample has at least one real (unmasked) entity. + self.assertTrue(gt["mask"].bool().any(dim=1).all()) def test_transe_train_step(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE @@ -242,3 +266,30 @@ def test_transe_scores_a_perfect_triple_at_the_margin(self) -> None: model.calc(head, relation, tail), torch.tensor(24.0), atol=1e-6 ) ) + + +class TestGroundTruthUnpadding(unittest.TestCase): + """Regression test for the padding-sentinel collision the Tensor Trick + introduces: pad_token_id (0) is not a reserved value, so a real entity id + of 0 must survive unpadding while an actual padding slot does not.""" + + def test_unpad_ground_truth_keeps_real_entity_zero_and_drops_padding(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) + + # Row 0: entities {0, 5} are real (mask=1), trailing slot is padding. + # Row 1: entity {3} is real, two trailing slots are padding. + value = torch.tensor([[0, 5, 0], [3, 0, 0]]) + mask = torch.tensor([[1, 1, 0], [1, 0, 0]]) + + unpadded = model._unpad_ground_truth({"value": value, "mask": mask}) + + self.assertEqual(unpadded, [[0, 5], [3]]) + + def test_unpad_ground_truth_passes_through_plain_lists_unchanged(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) + raw = [[0, 5], [3]] + self.assertEqual(model._unpad_ground_truth(raw), raw) From 343ca7f8ad0b85645cbf9e4bcc9dd4abc179ffbf Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Mon, 31 Aug 2026 22:55:49 +0200 Subject: [PATCH 12/12] fix(kg_emb): satisfy contribution-rules CI gate (ruff + docstring example) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The contribution-rules check (tools/check_pr_rules.py) failed on 5f66ff4, something I hadn't run locally before pushing — only pytest, not ruff or the repo's own gate. It enforces ruff-clean added/modified lines plus a '>>>' doctest on every new/modified top-level public class or function. - kg_processor.py: replace typing.Dict/List/Iterable with PEP 585 builtins and collections.abc.Iterable (ruff UP035/UP006, target-version py313). Entirely new file, so every line was in scope. - datasets/utils.py: add a runnable '>>>' example to collate_fn_dict_with_padding's docstring, since the function body was modified (the nested all-tensor-dict branch) and had none. Verified against the same tooling and base/head SHAs CI used (0a75f99..): `python tools/check_pr_rules.py --base --head` now reports "All PR contribution rules passed." tests/core/test_kg_emb.py still 24/24, plus the collate_fn_dict_with_padding doctest. Co-Authored-By: Claude Sonnet 5 --- pyhealth/datasets/utils.py | 16 ++++++++++++++++ pyhealth/processors/kg_processor.py | 7 ++++--- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index e960f2dd7..c59a31e07 100644 --- a/pyhealth/datasets/utils.py +++ b/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() diff --git a/pyhealth/processors/kg_processor.py b/pyhealth/processors/kg_processor.py index 0b1bfc9a6..817ab8050 100644 --- a/pyhealth/processors/kg_processor.py +++ b/pyhealth/processors/kg_processor.py @@ -1,4 +1,5 @@ -from typing import Any, Dict, Iterable, List +from collections.abc import Iterable +from typing import Any import torch @@ -44,7 +45,7 @@ def __init__(self, pad_token_id: int = 0): self.pad_token_id = pad_token_id self.max_length = 1 - def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: + def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: """Determine the maximum list length observed for ``field``. Args: @@ -58,7 +59,7 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: max_len = max(max_len, len(value)) self.max_length = max(1, max_len) - def process(self, value: List[int]) -> Dict[str, torch.Tensor]: + def process(self, value: list[int]) -> dict[str, torch.Tensor]: """Pad a list of entity ids to ``max_length`` and build its attention mask. Lists longer than ``max_length`` (e.g. seen only at inference time,