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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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/25] 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, From 792d83ed781c2b76d4686d2fb4492a79643cd1c6 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:05:14 +0200 Subject: [PATCH 13/25] refactor(kg_emb): drop KGDatasetProtocol, annotate models with SampleKGDataset SampleKGDataset is back in the SampleDataset hierarchy, so the structural protocol introduced for the standalone design no longer has a purpose. As announced on 29 Aug, remove protocols.py and its export, and annotate the `dataset` parameter of KGEBaseModel, TransE, RotatE, DistMult and ComplEx with SampleKGDataset, imported at runtime. There is no import cycle: the dataset modules do not import the models. The postponed-evaluation import is dropped from the five model files so that the annotation is the class itself rather than a string. --- .../kg_emb/datasets/__init__.py | 2 -- .../kg_emb/datasets/protocols.py | 34 ------------------- .../kg_emb/models/complex.py | 10 ++---- .../kg_emb/models/distmult.py | 10 ++---- .../kg_emb/models/kg_base.py | 8 ++--- .../kg_emb/models/rotate.py | 10 ++---- .../kg_emb/models/transe.py | 10 ++---- 7 files changed, 14 insertions(+), 70 deletions(-) delete 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 0dee767bd..02176f5a9 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,12 +1,10 @@ from .base_kg_dataset import BaseKGDataset -from .protocols import KGDatasetProtocol from .sample_kg_dataset import SampleKGDataset from .splitter import split from .umls import UMLSDataset __all__ = [ "BaseKGDataset", - "KGDatasetProtocol", "SampleKGDataset", "UMLSDataset", "split", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py deleted file mode 100644 index e335daa91..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Structural contract between knowledge-graph datasets and embedding models.""" - -from __future__ import annotations - -from collections.abc import Mapping -from typing import Any, Protocol, runtime_checkable - -__all__ = ["KGDatasetProtocol"] - - -@runtime_checkable -class KGDatasetProtocol(Protocol): - """Minimal capability a dataset must expose to parameterise a KGE model. - - Embedding models only need the cardinality of the entity and relation - vocabularies. They never read the samples at construction time. Depending - on this Protocol rather than on a concrete class keeps the model layer - testable with lightweight doubles and free of import-time coupling to the - dataset layer. - - Examples: - >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None - >>> isinstance(_Toy(), KGDatasetProtocol) - True - """ - - entity_num: int - relation_num: int - - @property - def task_spec_param(self) -> Mapping[str, Any] | None: ... diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index ee8837745..143b8dc6f 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,14 +1,10 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any +from typing import Any import torch +from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel -if TYPE_CHECKING: - from ..datasets.protocols import KGDatasetProtocol - class ComplEx(KGEBaseModel): """ComplEx @@ -28,7 +24,7 @@ class ComplEx(KGEBaseModel): def __init__( self, - dataset: KGDatasetProtocol, + dataset: SampleKGDataset, 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 e6f38d4ff..600576206 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,14 +1,10 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any +from typing import Any import torch +from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel -if TYPE_CHECKING: - from ..datasets.protocols import KGDatasetProtocol - class DistMult(KGEBaseModel): """DistMult @@ -27,7 +23,7 @@ class DistMult(KGEBaseModel): """ def __init__( self, - dataset: KGDatasetProtocol, + dataset: SampleKGDataset, 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 eaa535fbf..1a5ccab37 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,11 @@ -from __future__ import annotations - from abc import ABC -from typing import TYPE_CHECKING import numpy as np import torch import torch.nn.functional as F from torch import nn -if TYPE_CHECKING: - from ..datasets.protocols import KGDatasetProtocol +from ..datasets.sample_kg_dataset import SampleKGDataset class KGEBaseModel(ABC, nn.Module): @@ -44,7 +40,7 @@ def device(self): def __init__( self, - dataset: KGDatasetProtocol, + dataset: SampleKGDataset, e_dim: int = 500, r_dim: int = 500, ns: str = "uniform", diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index 1c352e72a..802a08397 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,14 +1,10 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any +from typing import Any import torch +from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel -if TYPE_CHECKING: - from ..datasets.protocols import KGDatasetProtocol - class RotatE(KGEBaseModel): """RotatE @@ -28,7 +24,7 @@ class RotatE(KGEBaseModel): def __init__( self, - dataset: KGDatasetProtocol, + dataset: SampleKGDataset, e_dim: int = 600, r_dim: int = 300, ns='adv', diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 5e6287a32..cda6876ea 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,14 +1,10 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any +from typing import Any import torch +from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel -if TYPE_CHECKING: - from ..datasets.protocols import KGDatasetProtocol - class TransE(KGEBaseModel): """TransE @@ -28,7 +24,7 @@ class TransE(KGEBaseModel): def __init__( self, - dataset: KGDatasetProtocol, + dataset: SampleKGDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", From cb2dc8847e3e402be5c6a8a8c7046cdfa8528dd8 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:11:47 +0200 Subject: [PATCH 14/25] docs(kg_emb): describe the InMemorySampleDataset design accurately The API page, the example and the test name still described the standalone map-style Dataset of the first version, and the module docstring justified KGProcessor by litdata's pickle-based caching, which InMemorySampleDataset does not go through: it transforms every sample once at construction and keeps the result in memory. - docs/api/medcode.rst: SampleKGDataset is an InMemorySampleDataset and get_dataloader accepts it; split() returns plain lists, which is why a fold is wrapped in torch DataLoader directly. - examples/kg_emb_sample_dataset.py: same correction in the docstring. - sample_kg_dataset.py, kg_base.py: KGProcessor pads the ground-truth lists so that they collate into fixed-shape tensors. - tests: rename test_is_a_map_style_dataset, since SampleKGDataset is an IterableDataset subclass. --- docs/api/medcode.rst | 21 +++++++++++----- examples/kg_emb_sample_dataset.py | 16 ++++++++----- .../kg_emb/datasets/sample_kg_dataset.py | 24 +++++++++---------- .../kg_emb/models/kg_base.py | 4 ++-- tests/core/test_kg_emb.py | 2 +- 5 files changed, 40 insertions(+), 27 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index c5b691af9..d79dffd23 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -101,12 +101,21 @@ 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()``. +DistMult and ComplEx on an in-memory list of triples. +:class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset` +is an :class:`~pyhealth.datasets.InMemorySampleDataset`: samples are processed +once, at construction, and kept in memory. +:func:`pyhealth.datasets.get_dataloader` accepts it. Its ``ground_truth_head`` +and ``ground_truth_tail`` fields are padded by +:class:`~pyhealth.processors.kg_processor.KGProcessor` and returned with a +mask. + +:func:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split` returns +plain lists of sample dictionaries that carry the ``train`` and +``hyperparameters`` fields the models read. Wrap a fold in +:class:`torch.utils.data.DataLoader` with +:func:`pyhealth.datasets.collate_fn_dict_with_padding`: ``get_dataloader`` +calls ``set_shuffle()``, which a list does not have. See ``examples/kg_emb_sample_dataset.py`` for a self-contained walk-through. diff --git a/examples/kg_emb_sample_dataset.py b/examples/kg_emb_sample_dataset.py index d2bfe6a34..c534864a5 100644 --- a/examples/kg_emb_sample_dataset.py +++ b/examples/kg_emb_sample_dataset.py @@ -1,12 +1,16 @@ """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`. +This example does not download UMLS. It shows the 2.0 path: + +1. Build a :class:`SampleKGDataset`, an ``InMemorySampleDataset`` whose + ``ground_truth_head`` / ``ground_truth_tail`` fields are padded by + ``KGProcessor`` and returned with a mask. +2. Split into train/val/test folds with :func:`split`. The folds are plain + lists of sample dicts carrying the ``train`` and ``hyperparameters`` + fields the model reads. 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()``. + :func:`collate_fn_dict_with_padding`. ``get_dataloader`` calls + ``set_shuffle()``, which a list does not have. """ import torch 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 9d716fefe..addad928d 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,17 +1,17 @@ """Task-specific sample dataset for knowledge-graph embedding models. -``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 +``SampleKGDataset`` builds on :class:`pyhealth.datasets.InMemorySampleDataset`, +which fits the processors and transforms every sample once, at construction, +and keeps the result in memory. It does not write or read ``litdata`` chunks. + +``ground_truth_head`` and ``ground_truth_tail`` are variable-length lists of +entity ids. The registered ``"kg_entity_list"`` processor +(:class:`~pyhealth.processors.kg_processor.KGProcessor`) pads each field +independently, to its own observed maximum length, and returns a +``{"value": Tensor, "mask": Tensor}`` pair, so that +:func:`~pyhealth.datasets.collate_fn_dict_with_padding` stacks these fields +into batch tensors. 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. """ 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 1a5ccab37..bdcb09c60 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -148,8 +148,8 @@ 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 + fixed length with ``pad_token_id`` (0 by default) so that they collate + into fixed-shape tensors. 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 diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index a350cc7cb..82d65d501 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -113,7 +113,7 @@ 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: + def test_is_an_in_memory_sample_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 From e95c358ae9c90f11b6dd273468a9162268990184 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:14:17 +0200 Subject: [PATCH 15/25] fix(kg_emb): build the models' __main__ demo loaders with get_dataloader Since SampleKGDataset inherits from InMemorySampleDataset, it is an IterableDataset, and torch's DataLoader rejects shuffle=True for it: running any of the four model modules raised "DataLoader with IterableDataset: expected unspecified shuffle option". Use get_dataloader, which calls InMemorySampleDataset.set_shuffle() instead, as the demos did before the 2.0 migration. The local SampleKGDataset import is dropped because the module now imports it at the top. --- .../pretrained_embeddings/kg_emb/models/complex.py | 14 ++------------ .../kg_emb/models/distmult.py | 14 ++------------ .../pretrained_embeddings/kg_emb/models/rotate.py | 14 ++------------ .../pretrained_embeddings/kg_emb/models/transe.py | 14 ++------------ 4 files changed, 8 insertions(+), 48 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 143b8dc6f..8dae884af 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -76,12 +76,7 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - 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, - ) + from pyhealth.datasets import get_dataloader samples: list[dict[str, Any]] = [ { @@ -108,12 +103,7 @@ def calc(self, head, relation, tail, mode='pos'): entity_num=8000, relation_num=8, ) - train_loader = DataLoader( - dataset, - batch_size=2, - shuffle=True, - collate_fn=collate_fn_dict_with_padding, - ) + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) # 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 600576206..c7422e480 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -58,12 +58,7 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - 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, - ) + from pyhealth.datasets import get_dataloader samples: list[dict[str, Any]] = [ { @@ -90,12 +85,7 @@ def calc(self, head, relation, tail, mode='pos'): entity_num=8000, relation_num=8, ) - train_loader = DataLoader( - dataset, - batch_size=2, - shuffle=True, - collate_fn=collate_fn_dict_with_padding, - ) + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) # 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 802a08397..48c254dd3 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -72,12 +72,7 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - 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, - ) + from pyhealth.datasets import get_dataloader samples: list[dict[str, Any]] = [ { @@ -104,12 +99,7 @@ def calc(self, head, relation, tail, mode='pos'): entity_num=8000, relation_num=8, ) - train_loader = DataLoader( - dataset, - batch_size=2, - shuffle=True, - collate_fn=collate_fn_dict_with_padding, - ) + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) # 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 cda6876ea..9ba220ca6 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -61,12 +61,7 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - 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, - ) + from pyhealth.datasets import get_dataloader samples: list[dict[str, Any]] = [ { @@ -93,12 +88,7 @@ def calc(self, head, relation, tail, mode='pos'): entity_num=8000, relation_num=8, ) - train_loader = DataLoader( - dataset, - batch_size=2, - shuffle=True, - collate_fn=collate_fn_dict_with_padding, - ) + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) # model model = TransE( From 478f36857ad0d43caf12f3c1d85d15ea2c2983ec Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 6 Oct 2026 16:24:07 +0200 Subject: [PATCH 16/25] docs(processors): match KGProcessor's doctest output to torch's repr torch pads integer tensor elements to a common width, so the example prints tensor([16, 0, 0]); the docstring showed tensor([16, 0, 0]) and failed under doctest. --- pyhealth/processors/kg_processor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyhealth/processors/kg_processor.py b/pyhealth/processors/kg_processor.py index 817ab8050..ead0dfa03 100644 --- a/pyhealth/processors/kg_processor.py +++ b/pyhealth/processors/kg_processor.py @@ -38,7 +38,7 @@ class KGProcessor(FeatureProcessor): ... ] >>> processor.fit(samples, "ground_truth_tail") >>> processor.process([16]) - {'value': tensor([16, 0, 0]), 'mask': tensor([1, 0, 0])} + {'value': tensor([16, 0, 0]), 'mask': tensor([1, 0, 0])} """ def __init__(self, pad_token_id: int = 0): From 1911c2ca1c0b7c92b05992acabcc1ae71818e3c5 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 7 Oct 2026 21:17:25 +0200 Subject: [PATCH 17/25] refactor(kg_emb): rebuild BaseKGDataset on BaseDataset A knowledge graph is now one `triples` table declared in YAML with `patient_id: null` and `timestamp: null`, so the loader numbers the rows and each triple is one record (the ClinVarDataset pattern). A PatientSplit passed to set_task therefore splits triples. - Entity and relation ids are global: assigned on the full graph, names sorted, ids 0..n-1. The Polars helpers (triples_frame, entity_vocabulary, relation_vocabulary, index_triples) will be reused by the link-prediction task. - Triples with a missing head, relation or tail raise instead of being dropped. - UMLSDataset is the bundled configs/umls.yaml plus a thin subclass. prepare_metadata copies the headerless graph.txt to umls-pyhealth.tsv under a header line; graph.txt is never modified, and the copy is made again when graph.txt is newer or differs in size. URL roots are rejected with a hint: the 1.x bucket serves graph.txt only. - Ids no longer follow the order of first appearance used in 1.x, so 1.x embeddings must be mapped through the id2entity saved with them. - The 1.x API keeps working for one release with DeprecationWarning: set_task(task_fn, ...) returns a SampleKGDataset built from the triples in file order, and triples, stat() and info() are kept. refresh_cache is accepted and ignored. The legacy path rejects split= and the other 2.0 arguments instead of storing them as hyper-parameters. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/base_kg_dataset.py | 539 +++++++++++++----- .../kg_emb/datasets/configs/umls.yaml | 13 + .../kg_emb/datasets/umls.py | 163 ++++-- tests/core/test_kg_emb.py | 193 ++++++- 4 files changed, 717 insertions(+), 191 deletions(-) create mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/configs/umls.yaml 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 6070ad178..0244d8d32 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,167 +1,448 @@ +"""Knowledge graphs on the standard PyHealth dataset backend. + +A knowledge graph is one ``triples`` table, declared in a YAML config with +``patient_id: null`` and ``timestamp: null`` and the attributes ``head``, +``relation`` and ``tail``. The loader then numbers the rows, so each triple +is its own record (the same pattern as ``ClinVarDataset``): a +:class:`~pyhealth.datasets.PatientSplit` passed to ``set_task`` splits +triples. + +Entity and relation ids are global. They are assigned on the full graph, +names sorted, ids ``0..n-1``, so every split shares one id space, as in the +transductive setting of Bordes et al. (2013) and Sun et al. (2019). Knowing +that an entity exists does not reveal any of its edges. +""" + +from __future__ import annotations + import logging -import os -from abc import ABC +import warnings from collections.abc import Callable +from pathlib import Path +from typing import Any + +import polars as pl -from pyhealth.datasets.utils import MODULE_CACHE_PATH, hash_str -from pyhealth.utils import load_pickle, save_pickle +from pyhealth.datasets import BaseDataset +from pyhealth.tasks import BaseTask from .sample_kg_dataset import SampleKGDataset logger = logging.getLogger(__name__) -INFO_MSG = """ -dataset.triples: - Array((, , )) -""" +__all__ = [ + "BaseKGDataset", + "entity_vocabulary", + "index_triples", + "relation_vocabulary", + "triples_frame", +] + +TABLE = "triples" +FIELDS = ("head", "relation", "tail") + + +def triples_frame(global_event_df: pl.LazyFrame) -> pl.LazyFrame: + """Selects the triples of a ``BaseKGDataset`` event frame. + + Args: + global_event_df: The dataset's ``global_event_df``. + + Returns: + A lazy frame with the string columns ``patient_id``, ``head``, + ``relation`` and ``tail``, one row per triple. + + Examples: + >>> import polars as pl + >>> events = pl.LazyFrame({ + ... "patient_id": ["0"], "event_type": ["triples"], + ... "triples/head": ["a"], "triples/relation": ["r"], + ... "triples/tail": ["b"], + ... }) + >>> triples_frame(events).collect().row(0) + ('0', 'a', 'r', 'b') + """ + return global_event_df.filter(pl.col("event_type") == TABLE).select( + "patient_id", *(pl.col(f"{TABLE}/{f}").alias(f) for f in FIELDS) + ) + + +def _count_incomplete(triples: pl.LazyFrame) -> int: + """Counts the triples with a missing head, relation or tail. + + The loader reads every column as a string and turns empty fields into + nulls, so a malformed row shows up here. + """ + return ( + triples.filter(pl.any_horizontal(pl.col(f).is_null() for f in FIELDS)) + .select(pl.len()) + .collect() + .item() + ) + +def entity_vocabulary(triples: pl.LazyFrame) -> pl.LazyFrame: + """Every entity, as a head or a tail, with its id. -class BaseKGDataset(ABC): - """Abstract base Knowledge Graph class - - This abstract class defines a uniform + Names are sorted before numbering, so the ids depend only on the set of + entities, not on the order of the rows. + + Args: + triples: Output of :func:`triples_frame`. + + Returns: + A lazy frame with the columns ``name`` (string) and ``id`` (Int64). + + Examples: + >>> import polars as pl + >>> triples = pl.LazyFrame( + ... {"head": ["c", "a"], "relation": ["r", "r"], "tail": ["a", "b"]} + ... ) + >>> entity_vocabulary(triples).collect().rows() + [('a', 0), ('b', 1), ('c', 2)] + """ + names = pl.concat( + [ + triples.select(pl.col("head").alias("name")), + triples.select(pl.col("tail").alias("name")), + ] + ) + return ( + names.unique() + .sort("name") + .with_row_index("id") + .select("name", pl.col("id").cast(pl.Int64)) + ) + + +def relation_vocabulary(triples: pl.LazyFrame) -> pl.LazyFrame: + """Every relation with its id, numbered like :func:`entity_vocabulary`. + + Examples: + >>> import polars as pl + >>> triples = pl.LazyFrame( + ... {"head": ["a", "b"], "relation": ["s", "r"], "tail": ["b", "a"]} + ... ) + >>> relation_vocabulary(triples).collect().rows() + [('r', 0), ('s', 1)] + """ + return ( + triples.select(pl.col("relation").alias("name")) + .unique() + .sort("name") + .with_row_index("id") + .select("name", pl.col("id").cast(pl.Int64)) + ) - Each specific dataset will be a subclass of this abstract class, which can then - be converted to samples dataset for different tasks by calling `self.set_task()`. + +def index_triples( + triples: pl.LazyFrame, + entities: pl.LazyFrame, + relations: pl.LazyFrame, +) -> pl.LazyFrame: + """Replaces the names of each triple with their ids. Args: - dataset_name: name of the dataset. - root: root directory of the raw data (should contain many csv files). - dev: whether to enable dev mode (only use a small subset of the data). - 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. + triples: Output of :func:`triples_frame`. + entities: Output of :func:`entity_vocabulary` on the same graph. + relations: Output of :func:`relation_vocabulary` on the same graph. + + Returns: + ``triples`` with the Int64 columns ``head_id``, ``relation_id`` and + ``tail_id`` added. + + Examples: + >>> import polars as pl + >>> triples = pl.LazyFrame( + ... {"head": ["b"], "relation": ["r"], "tail": ["a"]} + ... ) + >>> indexed = index_triples( + ... triples, entity_vocabulary(triples), relation_vocabulary(triples) + ... ) + >>> indexed.select("head_id", "relation_id", "tail_id").collect().row(0) + (1, 0, 0) + """ + return ( + triples.join( + entities.rename({"name": "head", "id": "head_id"}), on="head", how="left" + ) + .join( + relations.rename({"name": "relation", "id": "relation_id"}), + on="relation", + how="left", + ) + .join( + entities.rename({"name": "tail", "id": "tail_id"}), on="tail", how="left" + ) + ) + + +class BaseKGDataset(BaseDataset): + """A knowledge graph stored as one ``triples`` table. + + Each record of the table is one triple ``(head, relation, tail)``. Use + ``set_task`` with a task and ``split=PatientSplit(...)`` to get one + streaming :class:`~pyhealth.datasets.SampleDataset` per part. + + Args: + root: Directory (or URL) holding the table file named in the config. + config_path: YAML config with a ``triples`` table whose attributes + include ``head``, ``relation`` and ``tail``. + dataset_name: Name of the dataset. Defaults to the class name. + cache_dir: Cache directory, as in :class:`~pyhealth.datasets.BaseDataset`. + num_workers: Worker processes for loading and ``set_task``. + dev: Keep only the first 1000 triples, as ``BaseDataset`` keeps the + first 1000 patients. + refresh_cache: Deprecated and ignored. The cache key now covers the + config and the source files, so an edited file builds a new cache. + + Raises: + ValueError: If the config has no ``triples`` table with the + ``head``, ``relation`` and ``tail`` attributes. 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 + >>> from pathlib import Path + >>> root = Path(tempfile.mkdtemp()) + >>> _ = (root / "kg.tsv").write_text("head\\trelation\\ttail\\na\\tr\\tb\\n") + >>> _ = (root / "kg.yaml").write_text( + ... 'version: "1.0"\\n' + ... "tables:\\n" + ... " triples:\\n" + ... " file_path: kg.tsv\\n" + ... " patient_id: null\\n" + ... " timestamp: null\\n" + ... " attributes: [head, relation, tail]\\n" + ... ) + >>> ds = BaseKGDataset( # doctest: +SKIP + ... root=str(root), config_path=root / "kg.yaml", cache_dir=root / "cache" + ... ) + >>> ds.entity2id, ds.relation2id # doctest: +SKIP + ({'a': 0, 'b': 1}, {'r': 0}) """ def __init__( self, root: str, + config_path: str | Path, dataset_name: str | None = None, + cache_dir: str | Path | None = None, + num_workers: int = 1, dev: bool = False, - refresh_cache: bool = False - ): - self.dataset_name = ( - self.__class__.__name__ if dataset_name is None else dataset_name + refresh_cache: bool | None = None, + ) -> None: + if refresh_cache is not None: + warnings.warn( + "refresh_cache is deprecated and ignored: the dataset cache is " + "keyed on the config and the source files.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__( + root=root, + tables=[TABLE], + dataset_name=dataset_name, + config_path=str(config_path), + cache_dir=cache_dir, + num_workers=num_workers, + dev=dev, ) - self.root = root - self.dev = dev - self.triples = [] - self.samples = [] - self.task_name = "Null" - self.entity_num = 0 - self.relation_num = 0 - self.entity2id = None - self.relation2id = None - self.refresh_cache = refresh_cache - - # hash filename for cache - args_to_hash = [self.dataset_name, root] + ["dev" if dev else "prod"] - filename = hash_str("+".join([str(arg) for arg in args_to_hash])) - self.filepath = os.path.join(MODULE_CACHE_PATH, filename) - self.raw_graph_process() - - self.len = len(self.triples) - - - def __str__(self): - """Prints some information of the dataset.""" - return f"Base dataset {self.dataset_name}" - - def __len__(self): - return self.len - - def raw_graph_process(self): - """Process the raw graph to triples (a list of triple) - """ - raise NotImplementedError + self._check_config() + self._entity2id: dict[str, int] | None = None + self._relation2id: dict[str, int] | None = None - @staticmethod - def info(): - """Prints the output format.""" - print(INFO_MSG) - - def stat(self): - """Returns some statistics of the base dataset.""" - lines = [] - 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.triples)}") - 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- Number of samples: {len(self.samples)}") - lines.append("") - print("\n".join(lines)) - - - def set_task( - self, - task_fn: Callable, - task_name: str | None = None, - save: bool = True, - **kwargs - ) -> SampleKGDataset: - """Processes the base dataset to generate the task-specific sample dataset. + def _check_config(self) -> None: + table = self.config.tables.get(TABLE) if self.config else None + if table is None: + raise ValueError(f"A knowledge-graph config needs a '{TABLE}' table.") + missing = [f for f in FIELDS if f not in table.attributes] + if missing: + raise ValueError( + f"The '{TABLE}' table must have the attributes {list(FIELDS)}; " + f"missing {missing}." + ) + + def _triples_frame(self) -> pl.LazyFrame: + return triples_frame(self.global_event_df) + + def _load_vocabularies(self) -> None: + triples = self._triples_frame() + incomplete = _count_incomplete(triples) + if incomplete: + # Dropping them would change the graph; that decision is the + # user's, so stop instead. + raise ValueError( + f"{incomplete} triple(s) have a missing head, relation or tail." + ) + entities = entity_vocabulary(triples).collect() + relations = relation_vocabulary(triples).collect() + self._entity2id = dict(zip(entities["name"], entities["id"])) + self._relation2id = dict(zip(relations["name"], relations["id"])) + + @property + def entity2id(self) -> dict[str, int]: + """Entity name to id, on the full graph.""" + if self._entity2id is None: + self._load_vocabularies() + return self._entity2id - This function should be called by the user after the base dataset is - initialized. It will iterate through all patients in the base dataset - and call `task_fn` which should be implemented by the specific task. + @property + def relation2id(self) -> dict[str, int]: + """Relation name to id, on the full graph.""" + if self._relation2id is None: + self._load_vocabularies() + return self._relation2id - Args: - task_fn: a function that takes a single patient and returns a - list of samples (each sample is a dict with patient_id, visit_id, - and other task-specific attributes as key). The samples will be - concatenated to form the sample dataset. - task_name: the name of the task. If None, the name of the task - function will be used. + @property + def id2entity(self) -> dict[int, str]: + """Entity id to name.""" + return {i: name for name, i in self.entity2id.items()} - Returns: - sample_dataset: the task-specific sample (Base) dataset. + @property + def id2relation(self) -> dict[int, str]: + """Relation id to name.""" + return {i: name for name, i in self.relation2id.items()} + @property + def num_entities(self) -> int: + """Number of distinct entities in the graph.""" + return len(self.entity2id) + + @property + def num_relations(self) -> int: + """Number of distinct relations in the graph.""" + return len(self.relation2id) + + @property + def entity_num(self) -> int: + """Alias of ``num_entities``, the name used before PyHealth 2.0.""" + return self.num_entities + + @property + def relation_num(self) -> int: + """Alias of ``num_relations``, the name used before PyHealth 2.0.""" + return self.num_relations + + def _indexed_triples(self) -> list[tuple[int, int, int]]: + """All triples as ``(head_id, relation_id, tail_id)``, in file order.""" + self._load_vocabularies() # raises on incomplete triples + triples = self._triples_frame() + indexed = ( + index_triples(triples, entity_vocabulary(triples), relation_vocabulary(triples)) + # The loader numbers rows in file order but stores them sorted on + # the string id ("0", "1", "10", ...); the numeric sort restores + # the file order. + .sort(pl.col("patient_id").cast(pl.Int64)) + .select("head_id", "relation_id", "tail_id") + .collect() + ) + return list(indexed.iter_rows()) + + @property + def triples(self) -> list[tuple[int, int, int]]: + """Deprecated: all triples as id tuples, as in PyHealth 1.x. + + This loads the whole graph into a Python list. Use a task and + ``set_task`` instead. + """ + warnings.warn( + "BaseKGDataset.triples is deprecated and will be removed in the next " + "release; use set_task with a task instead.", + DeprecationWarning, + stacklevel=2, + ) + return self._indexed_triples() + + def stat(self) -> str: + """Deprecated: prints and returns the size of the graph. + + Use ``stats()``, ``num_entities`` and ``num_relations`` instead. """ + warnings.warn( + "BaseKGDataset.stat() is deprecated and will be removed in the next " + "release; use stats(), num_entities and num_relations.", + DeprecationWarning, + stacklevel=2, + ) + n_triples = self._triples_frame().select(pl.len()).collect().item() + report = "\n".join( + [ + "", + f"Statistics of base dataset (dev={self.dev}):", + f"\t- Dataset: {self.dataset_name}", + f"\t- Number of triples: {n_triples}", + f"\t- Number of entities: {self.num_entities}", + f"\t- Number of relations: {self.num_relations}", + "", + ] + ) + print(report) + return report + + @staticmethod + def info() -> None: + """Deprecated: prints the layout of the triples table.""" + warnings.warn( + "BaseKGDataset.info() is deprecated and will be removed in the next " + "release.", + DeprecationWarning, + stacklevel=2, + ) + print(f"Table '{TABLE}': one row per triple, attributes {list(FIELDS)}.") + + def set_task(self, task: BaseTask | Callable | None = None, *args, **kwargs): + """Builds the sample dataset(s) of ``task``. - if task_name is None: - self.task_name = task_fn.__name__ + With a :class:`~pyhealth.tasks.BaseTask`, this is + :meth:`pyhealth.datasets.BaseDataset.set_task`. A plain function, + such as ``link_prediction_fn``, goes through the deprecated path of + PyHealth 1.x: it receives the list of indexed triples, in file order, + and its samples are wrapped in a :class:`SampleKGDataset`, with the + remaining keyword arguments (e.g. ``negative_sampling``) as task + hyper-parameters. That path takes neither ``split`` nor the other + arguments of the 2.0 ``set_task``. + """ + if task is None and "task_fn" in kwargs: + # PyHealth 1.x named this argument task_fn. + task = kwargs.pop("task_fn") + if task is not None and not isinstance(task, BaseTask) and callable(task): + return self._set_legacy_task(task, *args, **kwargs) + return super().set_task(task, *args, **kwargs) - # check if cache exists or refresh_cache is True - if os.path.exists(self.filepath + '.pkl') and (not self.refresh_cache): - # load from cache - print( - f"Loading {self.dataset_name} base dataset from {self.filepath}.pkl" + def _set_legacy_task( + self, + task_fn: Callable, + task_name: str | None = None, + save: bool | None = None, + **task_spec_param: Any, + ) -> SampleKGDataset: + warnings.warn( + "set_task(task_fn) with a function is deprecated and will be removed " + "in the next release. Pass a BaseTask and split=PatientSplit(...) " + "to get streaming SampleDatasets.", + DeprecationWarning, + stacklevel=3, + ) + # These would otherwise be stored as task hyper-parameters and + # silently have no effect. + unsupported = sorted( + {"split", "num_workers", "input_processors", "output_processors"} + & task_spec_param.keys() + ) + if unsupported: + raise TypeError( + f"set_task with a function does not take {unsupported}; " + "pass a BaseTask to use them." ) - self.samples = load_pickle(self.filepath + ".pkl") - - else: - print(f"Processing {self.dataset_name} base dataset...") - self.samples = task_fn(self.triples) - - # save to cache - print(f"Saving {self.dataset_name} base dataset to {self.filepath}") - if save == True: - save_pickle(self.samples, self.filepath + ".pkl") - - sample_dataset = SampleKGDataset( - samples=self.samples, + if save is not None: + logger.info("set_task(save=...) is ignored: samples are not pickled.") + samples = task_fn(self._indexed_triples()) + return SampleKGDataset( + samples=samples, dataset_name=self.dataset_name, - task_name=self.task_name, + task_name=task_name or task_fn.__name__, dev=self.dev, - entity_num=self.entity_num, - relation_num=self.relation_num, entity2id=self.entity2id, relation2id=self.relation2id, - **kwargs + **task_spec_param, ) - - return sample_dataset diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/configs/umls.yaml b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/configs/umls.yaml new file mode 100644 index 000000000..117799094 --- /dev/null +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/configs/umls.yaml @@ -0,0 +1,13 @@ +version: "1.0" +tables: + triples: + # Written by UMLSDataset.prepare_metadata from the headerless graph.txt. + file_path: "umls-pyhealth.tsv" + # No patient column: the loader numbers the rows, so each triple is its + # own record and PatientSplit splits triples. + patient_id: null + timestamp: null + attributes: + - "head" + - "relation" + - "tail" diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 8180a9b91..802a43671 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -1,25 +1,50 @@ import logging import os +import shutil +from pathlib import Path -import pandas as pd -from tqdm import tqdm +from pyhealth.datasets.base_dataset import is_url from .base_kg_dataset import BaseKGDataset logger = logging.getLogger(__name__) +RAW_FILE = "graph.txt" +PREPARED_FILE = "umls-pyhealth.tsv" +HEADER = b"head\trelation\ttail\n" + + class UMLSDataset(BaseKGDataset): - """Base UMLS knowleddge graph dataset + """UMLS knowledge graph dataset. Dataset is available at https://www.nlm.nih.gov/research/umls/index.html + ``root`` holds ``graph.txt``: one triple per line, ``head``, ``relation`` + and ``tail`` separated by tabs, without a header. On first use, + :meth:`prepare_metadata` writes ``umls-pyhealth.tsv`` next to it, the same + rows under a ``head``/``relation``/``tail`` header, which the bundled + ``configs/umls.yaml`` reads. ``graph.txt`` itself is never modified. The + copy is made again whenever ``graph.txt`` is newer than it or differs from + it in size, so replacing ``graph.txt`` gives a new dataset cache. + + Entity and relation ids are assigned by sorted name (see + :class:`BaseKGDataset`), no longer by order of first appearance as in + PyHealth 1.x: embeddings trained with 1.x must be mapped through the + ``id2entity`` / ``id2relation`` saved with them. + Args: - dataset_name: name of the dataset. - root: root directory of the raw data (should contain many csv files). - dev: whether to enable dev mode (only use a small subset of the data). - 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. + root: Local directory holding ``graph.txt``. URLs are not accepted; + the 1.x bucket ``https://storage.googleapis.com/pyhealth/umls/`` + serves ``graph.txt``, which can be downloaded into ``root``. + dataset_name: Name of the dataset. Defaults to ``"umls"``. + config_path: Config to use instead of the bundled ``umls.yaml``. + **kwargs: Passed to :class:`BaseKGDataset` (``cache_dir``, + ``num_workers``, ``dev``). + + Raises: + ValueError: If ``root`` is a URL. + FileNotFoundError: If neither ``graph.txt`` nor + ``umls-pyhealth.tsv`` is in ``root``. Examples: >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( @@ -28,47 +53,95 @@ class UMLSDataset(BaseKGDataset): ... ) >>> issubclass(UMLSDataset, BaseKGDataset) True + >>> dataset = UMLSDataset(root="/path/to/umls") # doctest: +SKIP + >>> dataset.num_entities # doctest: +SKIP """ - def raw_graph_process(self): - if self.dev == False: - self.graph_path = os.path.join(self.root, "graph.txt") - else: - self.graph_path = os.path.join(self.root, "graph_filtered.txt") - - if os.path.exists(self.graph_path): - logger.debug("umls knowledge graph exists and load umls") - else: - logger.debug("umls does not exist") - - print("Loading UMLS knowledge graph...") - graph_df = pd.read_csv( - self.graph_path, - sep='\t', - names=['e1', 'r', 'e2'] + def __init__( + self, + root: str, + dataset_name: str | None = None, + config_path: str | Path | None = None, + **kwargs, + ) -> None: + if is_url(root): + # The public bucket used by PyHealth 1.x serves graph.txt only, + # and the prepared copy cannot be written next to a URL. + raise ValueError( + f"UMLSDataset needs a local root; download {root.rstrip('/')}/" + f"{RAW_FILE} into a directory and pass that directory." + ) + if config_path is None: + config_path = Path(__file__).parent / "configs" / "umls.yaml" + if self._needs_preparation(root): + self.prepare_metadata(root) + super().__init__( + root=root, + config_path=config_path, + dataset_name=dataset_name or "umls", + **kwargs, ) - print("Processing UMLS knowledge graph...") - entity_list = pd.unique(graph_df[['e1', 'e2']].values.ravel('K')) - relation_list = pd.unique(graph_df['r'].values) - - self.entity2id = {val: i for i, val in enumerate(entity_list)} - self.relation2id = {val: i for i, val in enumerate(relation_list)} - self.entity_num = len(self.entity2id) - self.relation_num = len(self.relation2id) - - print("Building UMLS knowledge graph...") - self.triples = [(self.entity2id[e1], self.relation2id[r], self.entity2id[e2]) - for e1, r, e2 in tqdm(zip(graph_df['e1'], graph_df['r'], graph_df['e2']), total=graph_df.shape[0])] - + @staticmethod + def _needs_preparation(root: str) -> bool: + """Whether ``umls-pyhealth.tsv`` is missing or older than ``graph.txt``. + + The dataset cache is keyed on the prepared copy, so a ``graph.txt`` + replaced after the copy was made must trigger a new copy; otherwise + the old graph would be reused without notice. + """ + source = os.path.join(root, RAW_FILE) + target = os.path.join(root, PREPARED_FILE) + if not os.path.exists(target): + return True + if not os.path.exists(source): + # Only the prepared copy was provided; use it as it is. + return False + src, dst = os.stat(source), os.stat(target) + return ( + dst.st_size != src.st_size + len(HEADER) + or src.st_mtime_ns > dst.st_mtime_ns + ) + @staticmethod + def prepare_metadata(root: str) -> None: + """Writes ``umls-pyhealth.tsv``: ``graph.txt`` under a header line. + + The rows are copied byte for byte; only the header + ``head\\trelation\\ttail`` is added, because the table loader reads the + column names from the first line. + + Args: + root: Directory holding ``graph.txt``. + + Raises: + FileNotFoundError: If ``graph.txt`` is not in ``root``. + + Examples: + >>> import tempfile + >>> from pathlib import Path + >>> root = Path(tempfile.mkdtemp()) + >>> _ = (root / "graph.txt").write_text("C1\\tPAR\\tC2\\n") + >>> UMLSDataset.prepare_metadata(str(root)) + >>> (root / "umls-pyhealth.tsv").read_text() + 'head\\trelation\\ttail\\nC1\\tPAR\\tC2\\n' + """ + source = os.path.join(root, RAW_FILE) + if not os.path.exists(source): + raise FileNotFoundError(f"UMLS graph not found: {source}") + target = os.path.join(root, PREPARED_FILE) + logger.info(f"Writing {target} from {source}") + # Written under a temporary name and renamed when complete: an + # interrupted copy must not leave a truncated file that later runs + # would take as prepared. + partial = target + ".partial" + with open(source, "rb") as src, open(partial, "wb") as dst: + dst.write(HEADER) + shutil.copyfileobj(src, dst) + os.replace(partial, target) if __name__ == "__main__": - dataset = UMLSDataset( - root="https://storage.googleapis.com/pyhealth/umls/", - dev=True, - refresh_cache=True, - ) - dataset.stat() - dataset.info() \ No newline at end of file + dataset = UMLSDataset(root="/path/to/umls", dev=True) + dataset.stats() + print(dataset.num_entities, dataset.num_relations) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 82d65d501..c3b3352e6 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -7,8 +7,11 @@ from __future__ import annotations +import hashlib import tempfile import unittest +import warnings +from pathlib import Path from typing import Any import torch @@ -18,6 +21,7 @@ from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( BaseKGDataset, SampleKGDataset, + UMLSDataset, split, ) from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import link_prediction_fn @@ -36,6 +40,41 @@ def make_samples(n: int = 8) -> list[dict[str, Any]]: ] +# A small graph: entities a..d, relations r1, r2. Names first appear out of +# order (c before a, r2 before r1), so numbering by first appearance would +# differ from the sorted ids a=0, b=1, c=2, d=3 and r1=0, r2=1. +TOY_TRIPLES = [ + ("c", "r2", "a"), + ("a", "r1", "b"), + ("a", "r1", "c"), + ("b", "r2", "c"), + ("d", "r1", "a"), +] + +# Twelve triples: the loader stores rows sorted on the string row number +# ("0", "1", "10", "11", "2", ...), so file order is only recovered by a +# numeric sort once there are more than ten rows. +ORDER_TRIPLES = [(f"e{i}", "r", f"e{i + 1}") for i in range(12)] + +TOY_CONFIG = """version: "1.0" +tables: + triples: + file_path: "kg.tsv" + patient_id: null + timestamp: null + attributes: [head, relation, tail] +""" + + +def write_kg(root: Path, triples: list[tuple[str, str, str]]) -> Path: + """Write ``triples`` as ``kg.tsv`` with its config; return the config path.""" + lines = ["head\trelation\ttail"] + ["\t".join(t) for t in triples] + (root / "kg.tsv").write_text("\n".join(lines) + "\n") + config = root / "kg.yaml" + config.write_text(TOY_CONFIG) + return config + + 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)} @@ -215,27 +254,147 @@ def test_transe_train_step(self) -> None: out["loss"].backward() -class TestSetTask(unittest.TestCase): - """Production path: BaseKGDataset.set_task must return a usable SampleKGDataset.""" +class TestBaseKGDataset(unittest.TestCase): + """BaseKGDataset loads one triples table through the BaseDataset backend.""" - 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)] + @classmethod + def setUpClass(cls) -> None: + cls._tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + cls.root = Path(cls._tmp.name) + config = write_kg(cls.root, TOY_TRIPLES) + cls.dataset = BaseKGDataset( + root=str(cls.root), config_path=config, cache_dir=cls.root / "cache" + ) - 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 - ) + @classmethod + def tearDownClass(cls) -> None: + cls._tmp.cleanup() + + def test_each_triple_is_one_record(self) -> None: + self.assertEqual(len(self.dataset.unique_patient_ids), len(TOY_TRIPLES)) + + def test_vocabularies_are_global_and_sorted(self) -> None: + self.assertEqual(self.dataset.entity2id, {"a": 0, "b": 1, "c": 2, "d": 3}) + self.assertEqual(self.dataset.relation2id, {"r1": 0, "r2": 1}) + self.assertEqual(self.dataset.id2entity[3], "d") + self.assertEqual(self.dataset.id2relation[1], "r2") + self.assertEqual((self.dataset.num_entities, self.dataset.num_relations), (4, 2)) + # Pre-2.0 names, still read by older code. + self.assertEqual((self.dataset.entity_num, self.dataset.relation_num), (4, 2)) + + def test_legacy_function_task_still_works_with_a_warning(self) -> None: + with self.assertWarns(DeprecationWarning): + sample_ds = self.dataset.set_task(link_prediction_fn, negative_sampling=4) self.assertIsInstance(sample_ds, SampleKGDataset) - self.assertEqual(len(sample_ds), 3) + self.assertEqual(len(sample_ds), len(TOY_TRIPLES)) self.assertEqual(sample_ds.task_spec_param, {"negative_sampling": 4}) - self.assertIn("triple", sample_ds[0]) + self.assertEqual(sample_ds.entity_num, 4) + + def test_legacy_task_fn_keyword_is_accepted(self) -> None: + with self.assertWarns(DeprecationWarning): + sample_ds = self.dataset.set_task(task_fn=link_prediction_fn) + self.assertEqual(len(sample_ds), len(TOY_TRIPLES)) + + def test_legacy_path_rejects_2_0_arguments(self) -> None: + from pyhealth.datasets import PatientSplit + + with self.assertRaisesRegex(TypeError, "split"), warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + self.dataset.set_task(link_prediction_fn, split=PatientSplit((0.5, 0.5))) + + def test_legacy_members_warn(self) -> None: + with self.assertWarns(DeprecationWarning): + report = self.dataset.stat() + self.assertIn("Number of triples: 5", report) + with self.assertWarns(DeprecationWarning): + BaseKGDataset.info() + with self.assertWarns(DeprecationWarning): + self.assertEqual(len(self.dataset.triples), len(TOY_TRIPLES)) + + +class TestLegacyTripleOrder(unittest.TestCase): + """The legacy path hands task functions the triples in file order.""" + + def test_more_than_ten_triples_keep_file_order(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_kg(root, ORDER_TRIPLES) + dataset = BaseKGDataset(root=tmp, config_path=config, cache_dir=root / "cache") + e, r = dataset.entity2id, dataset.relation2id + expected = [(e[h], r[rel], e[t]) for h, rel, t in ORDER_TRIPLES] + with self.assertWarns(DeprecationWarning): + self.assertEqual(dataset.triples, expected) + + +class TestBaseKGDatasetValidation(unittest.TestCase): + """Malformed configs and triples are reported, never silently dropped.""" + + def test_config_without_the_triple_attributes_is_rejected(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_kg(root, TOY_TRIPLES) + config.write_text(TOY_CONFIG.replace("[head, relation, tail]", "[head, tail]")) + with self.assertRaisesRegex(ValueError, r"missing \['relation'\]"): + BaseKGDataset(root=tmp, config_path=config, cache_dir=root / "cache") + + def test_a_triple_with_a_missing_field_raises(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_kg(root, TOY_TRIPLES + [("a", "r1", "")]) + dataset = BaseKGDataset(root=tmp, config_path=config, cache_dir=root / "cache") + with self.assertRaisesRegex(ValueError, r"1 triple\(s\) have a missing"): + dataset.entity2id + + def test_refresh_cache_is_deprecated(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_kg(root, TOY_TRIPLES) + with self.assertWarns(DeprecationWarning): + BaseKGDataset( + root=tmp, config_path=config, cache_dir=root / "cache", + refresh_cache=True, + ) + + +class TestUMLSDataset(unittest.TestCase): + """UMLSDataset reads the headerless graph.txt through a prepared copy.""" + + def test_prepared_copy_leaves_graph_txt_untouched(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + raw = root / "graph.txt" + raw.write_text("".join(f"{h}\t{r}\t{t}\n" for h, r, t in TOY_TRIPLES)) + digest = hashlib.sha256(raw.read_bytes()).hexdigest() + + dataset = UMLSDataset(root=tmp, cache_dir=root / "cache") + + self.assertEqual(hashlib.sha256(raw.read_bytes()).hexdigest(), digest) + prepared = (root / "umls-pyhealth.tsv").read_text().splitlines() + self.assertEqual(prepared[0], "head\trelation\ttail") + self.assertEqual(prepared[1:], raw.read_text().splitlines()) + self.assertEqual(dataset.dataset_name, "umls") + self.assertEqual(dataset.entity2id, {"a": 0, "b": 1, "c": 2, "d": 3}) + + def test_a_replaced_graph_txt_is_prepared_again(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + raw = root / "graph.txt" + raw.write_text("A\tPAR\tB\nB\tCHD\tA\n") + first = UMLSDataset(root=tmp, cache_dir=root / "cache") + self.assertEqual(first.num_entities, 2) + + raw.write_text("A\tPAR\tB\nB\tCHD\tA\nC\tRO\tA\n") + second = UMLSDataset(root=tmp, cache_dir=root / "cache") + self.assertEqual((second.num_entities, second.num_relations), (3, 3)) + + def test_url_root_is_rejected_with_a_hint(self) -> None: + with self.assertRaisesRegex(ValueError, "graph.txt"): + UMLSDataset(root="https://storage.googleapis.com/pyhealth/umls/") + + def test_missing_graph_is_reported(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + with self.assertRaises(FileNotFoundError): + UMLSDataset(root=tmp, cache_dir=Path(tmp) / "cache") class TestScoringInvariants(unittest.TestCase): From d349b1c3e3b39218e0ca1ad7a6a3468676536dc7 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 7 Oct 2026 22:13:19 +0200 Subject: [PATCH 18/25] feat(processors): add a kg_triple processor and keep long KG lists whole KGTripleProcessor ("kg_triple") returns a triple of ids as a LongTensor and, from the triples it is fitted on, keeps plain dicts: true_head, true_tail and the word2vec frequencies behind subsampling weights, built as in the reference RotatE implementation of Sun et al. (2019), which computes them from training triples only. With set_task(task, split=PatientSplit(...)), SampleBuilder fits processors on the training part, so these dicts never see validation or test triples. learns_statistics is set, so set_task without split= warns. Entity and relation counts are given, not fitted, because ids are global; size() returns the number of entities. The fitted dicts sit in a holder whose str() is a digest of their content: set_task puts vars() of pre-fitted processors into a JSON cache key, which cannot hold tuple-keyed dicts. KGProcessor no longer truncates lists longer than its fitted length. Fitted on the training split, that length depends on the seed, and a truncated filter set leaves known positives among the ranked candidates. collate_fn_dict_with_padding already pads such a batch with mask 0. Co-Authored-By: Claude Opus 5.5 --- docs/api/processors.rst | 6 +- .../pyhealth.processors.KGProcessor.rst | 16 ++ .../pyhealth.processors.KGTripleProcessor.rst | 23 ++ pyhealth/datasets/utils.py | 9 +- pyhealth/processors/__init__.py | 2 + pyhealth/processors/kg_processor.py | 35 +-- pyhealth/processors/kg_triple_processor.py | 232 ++++++++++++++++++ tests/core/test_kg_emb.py | 159 ++++++++++++ 8 files changed, 461 insertions(+), 21 deletions(-) create mode 100644 docs/api/processors/pyhealth.processors.KGProcessor.rst create mode 100644 docs/api/processors/pyhealth.processors.KGTripleProcessor.rst create mode 100644 pyhealth/processors/kg_triple_processor.py diff --git a/docs/api/processors.rst b/docs/api/processors.rst index ea83659b4..0936b0e16 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -55,6 +55,8 @@ Available Processors - ``MultiHotProcessor``: For multi-hot encoding - ``IgnoreProcessor``: A special feature processor that marks a feature to be ignored. - ``GraphProcessor``: For knowledge graph subgraph extraction (e.g., GraphCare, G-BERT) +- ``KGTripleProcessor``: For knowledge graph triples; keeps training-graph dicts for KG embedding models +- ``KGProcessor``: For variable-length lists of knowledge graph entity ids, padded with a mask **Temporal Multimodal Processors (** :class:`~pyhealth.processors.TemporalFeatureProcessor` **subclasses):** @@ -527,4 +529,6 @@ API Reference processors/pyhealth.processors.MultiHotProcessor processors/pyhealth.processors.StageNetProcessor processors/pyhealth.processors.StageNetTensorProcessor - processors/pyhealth.processors.GraphProcessor \ No newline at end of file + processors/pyhealth.processors.GraphProcessor + processors/pyhealth.processors.KGTripleProcessor + processors/pyhealth.processors.KGProcessor \ No newline at end of file diff --git a/docs/api/processors/pyhealth.processors.KGProcessor.rst b/docs/api/processors/pyhealth.processors.KGProcessor.rst new file mode 100644 index 000000000..7143020dd --- /dev/null +++ b/docs/api/processors/pyhealth.processors.KGProcessor.rst @@ -0,0 +1,16 @@ +pyhealth.processors.KGProcessor +================================= + +Overview +-------- +Processor that pads a variable-length list of knowledge-graph entity ids, +such as the ground-truth sets used for filtered link-prediction ranking, +and returns it with a mask. Registered as ``"kg_entity_list"`` in the +processor registry. A list longer than the fitted length is kept whole. + +API Reference +------------- +.. automodule:: pyhealth.processors.kg_processor + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst new file mode 100644 index 000000000..b673eeace --- /dev/null +++ b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst @@ -0,0 +1,23 @@ +pyhealth.processors.KGTripleProcessor +======================================= + +Overview +-------- +Processor for a knowledge-graph triple ``(head, relation, tail)`` of integer +ids, registered as ``"kg_triple"`` in the processor registry. It returns the +triple as a ``LongTensor`` and, from the triples it is fitted on, keeps plain +dicts used to train KG embedding models: the known heads of each +``(relation, tail)`` pair, the known tails of each ``(head, relation)`` pair, +and the pair frequencies behind subsampling weights. + +Fit it on the training triples only, with +``set_task(task, split=PatientSplit(...))``: the training negatives and +weights then never depend on validation or test triples. ``size()`` returns +the number of entities, which is given (ids are global), not fitted. + +API Reference +------------- +.. automodule:: pyhealth.processors.kg_triple_processor + :members: + :undoc-members: + :show-inheritance: diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index c59a31e07..8dfa62870 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -324,10 +324,11 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: 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. + # from KGProcessor. Stack each sub-key independently. KGProcessor + # pads to its fitted length, so this is usually a plain stack; a + # list longer than that length (e.g. a test filter set after + # fitting on the training split) is kept whole, and the batch is + # then padded here with 0, which gives mask 0 on the padding. # # Restricted to all-tensor dicts so heterogeneous per-sample # dicts (e.g. a "hyperparameters" field of plain Python values) diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index e3c9e82f0..41c2de96d 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -48,6 +48,7 @@ def get_processor(name: str): from .time_image_processor import TimeImageProcessor from .graph_processor import GraphProcessor from .kg_processor import KGProcessor +from .kg_triple_processor import KGTripleProcessor from .audio_processor import AudioProcessor from .ignore_processor import IgnoreProcessor from .temporal_timeseries_processor import TemporalTimeseriesProcessor @@ -82,6 +83,7 @@ def get_processor(name: str): "TimeImageProcessor", "GraphProcessor", "KGProcessor", + "KGTripleProcessor", "AudioProcessor", "TupleTimeTextProcessor", "CehrProcessor", diff --git a/pyhealth/processors/kg_processor.py b/pyhealth/processors/kg_processor.py index ead0dfa03..748f22acb 100644 --- a/pyhealth/processors/kg_processor.py +++ b/pyhealth/processors/kg_processor.py @@ -14,9 +14,13 @@ class KGProcessor(FeatureProcessor): 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`. + KG embedding models). Lists are padded to the maximum list length + observed for that field during ``fit``. A longer list, which can occur + when the processor is fitted on a training split and the list belongs to + a validation or test sample, keeps all its entities: truncating a filter + set would leave known positives among the ranked candidates and bias + filtered metrics. :func:`~pyhealth.datasets.collate_fn_dict_with_padding` + pads such a batch to its longest list, with mask 0 on the padding. Note: The padded ``pad_token_id`` is not a valid entity id on its own: any @@ -39,6 +43,8 @@ class KGProcessor(FeatureProcessor): >>> processor.fit(samples, "ground_truth_tail") >>> processor.process([16]) {'value': tensor([16, 0, 0]), 'mask': tensor([1, 0, 0])} + >>> processor.process([1, 2, 3, 4]) + {'value': tensor([1, 2, 3, 4]), 'mask': tensor([1, 1, 1, 1])} """ def __init__(self, pad_token_id: int = 0): @@ -62,28 +68,25 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: 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. + A list longer than ``max_length`` (e.g. a validation or test filter + set, after ``fit`` on the training split) is kept whole, never + truncated. Args: value: List of entity ids. Returns: Dict with: - - ``"value"``: LongTensor of shape ``(max_length,)``, padded with + - ``"value"``: LongTensor of shape + ``(max(max_length, len(value)),)``, padded with ``pad_token_id``. - - ``"mask"``: LongTensor of shape ``(max_length,)``, 1 for real + - ``"mask"``: LongTensor of the same shape, 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) + n_pad = max(0, self.max_length - len(entities)) + padded = entities + [self.pad_token_id] * n_pad + mask = [1] * len(entities) + [0] * n_pad return { "value": torch.tensor(padded, dtype=torch.long), @@ -91,7 +94,7 @@ def process(self, value: list[int]) -> dict[str, torch.Tensor]: } def size(self) -> int: - """Return the fitted padding length.""" + """Return the fitted padding length, the minimum output length.""" return self.max_length def is_token(self) -> bool: diff --git a/pyhealth/processors/kg_triple_processor.py b/pyhealth/processors/kg_triple_processor.py new file mode 100644 index 000000000..ec2ec7ce6 --- /dev/null +++ b/pyhealth/processors/kg_triple_processor.py @@ -0,0 +1,232 @@ +import hashlib +from collections import defaultdict +from collections.abc import Iterable +from typing import Any + +import torch + +from . import register_processor +from .base_processor import FeatureProcessor + + +class _FittedDicts: + """The dicts a ``KGTripleProcessor`` learns in ``fit``. + + ``BaseDataset.set_task`` puts ``vars()`` of pre-fitted processors into a + JSON cache key. Tuple-keyed dicts cannot be written as JSON, and their + full content would not belong in a key anyway, so they live here and the + key gets ``str()`` of this object: a digest of the content, the same on + every run and different whenever the fitted dicts differ. + """ + + def __init__( + self, + true_head: dict[tuple[int, int], list[int]], + true_tail: dict[tuple[int, int], list[int]], + count: dict[tuple[int, int], int], + ): + self.true_head = true_head + self.true_tail = true_tail + self.count = count + content = repr( + (sorted(true_head.items()), sorted(true_tail.items()), sorted(count.items())) + ) + self.digest = hashlib.sha256(content.encode()).hexdigest() + + def __str__(self) -> str: + return f"fitted-dicts-sha256:{self.digest}" + + +@register_processor("kg_triple") +class KGTripleProcessor(FeatureProcessor): + """Encodes a knowledge-graph triple and keeps training-graph dicts. + + ``process`` turns a triple ``(head, relation, tail)`` of integer ids into + a ``LongTensor`` of shape ``(3,)``. ``fit`` reads the triples it is given + and fills three plain dicts, built as in the reference implementation of + Sun et al. (2019, RotatE), whose training dataset computes them from the + training triples only: + + - ``true_head[(relation, tail)]``: the heads seen with that pair; + - ``true_tail[(head, relation)]``: the tails seen with that pair; + - ``count``: word2vec-style frequencies of ``(head, relation)`` and + ``(tail, -relation - 1)``, starting at ``count_start``, from which + :meth:`subsampling_weight` derives the weight of a training triple. + + The models use ``true_head`` / ``true_tail`` to keep known positives out + of the negatives they draw for training. Evaluation does not use them: + filtered ranking uses the ground-truth sets of the whole graph, carried + by the samples. + + These dicts must describe the training triples only, otherwise the + negatives and weights of training depend on validation and test + triples. ``set_task(task, split=PatientSplit(...))`` fits processors on + the training part only; without ``split``, PyHealth warns, because + ``learns_statistics`` is set. + + Entity and relation ids are global (fixed on the full graph), so + ``num_entities`` and ``num_relations`` are given, not fitted: an entity + absent from the training triples keeps its id. ``size()`` returns + ``num_entities``. + + Args: + num_entities: Number of entities; ids are ``0..num_entities - 1``. + num_relations: Number of relations; ids are + ``0..num_relations - 1``. + count_start: Value of a frequency the first time a pair is seen. + Default is 4, as in Sun et al. (2019). + + Raises: + ValueError: If a count is not a positive integer, or if ``fit`` + meets a malformed triple or an id out of range. + + Examples: + >>> processor = KGTripleProcessor(num_entities=3, num_relations=1) + >>> processor.fit([{"triple": (0, 0, 1)}, {"triple": (0, 0, 2)}], "triple") + >>> processor.process((0, 0, 1)) + tensor([0, 0, 1]) + >>> processor.true_tail[(0, 0)] + [1, 2] + >>> processor.true_head[(0, 2)] + [0] + >>> processor.subsampling_weight(torch.tensor([[0, 0, 1]])) + tensor([0.3333]) + """ + + learns_statistics = True + + def __init__(self, num_entities: int, num_relations: int, count_start: int = 4): + for name, value in ( + ("num_entities", num_entities), + ("num_relations", num_relations), + ("count_start", count_start), + ): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError(f"{name} must be a positive integer, got {value!r}.") + self.num_entities = num_entities + self.num_relations = num_relations + self.count_start = count_start + self._fitted = _FittedDicts({}, {}, {}) + + @property + def true_head(self) -> dict[tuple[int, int], list[int]]: + """``(relation, tail)`` to the sorted heads seen with it in ``fit``.""" + return self._fitted.true_head + + @property + def true_tail(self) -> dict[tuple[int, int], list[int]]: + """``(head, relation)`` to the sorted tails seen with it in ``fit``.""" + return self._fitted.true_tail + + @property + def count(self) -> dict[tuple[int, int], int]: + """Frequencies of ``(head, relation)`` and ``(tail, -relation - 1)``.""" + return self._fitted.count + + def _check(self, value: Any) -> tuple[int, int, int]: + ids = [int(x) for x in value] + if len(ids) != 3: + raise ValueError(f"A triple has 3 ids, got {tuple(value)}.") + head, relation, tail = ids + if not (0 <= head < self.num_entities and 0 <= tail < self.num_entities): + raise ValueError( + f"Triple {tuple(value)} has an entity id outside " + f"[0, {self.num_entities})." + ) + if not 0 <= relation < self.num_relations: + raise ValueError( + f"Triple {tuple(value)} has a relation id outside " + f"[0, {self.num_relations})." + ) + return head, relation, tail + + def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: + """Builds ``true_head``, ``true_tail`` and ``count`` from ``samples``. + + Refitting replaces the dicts; it does not add to them. + + Args: + samples: Samples whose ``field`` holds a triple of ids. With + ``set_task(..., split=...)`` these are the training samples. + field: Name of the triple field. + """ + true_head: dict[tuple[int, int], set[int]] = defaultdict(set) + true_tail: dict[tuple[int, int], set[int]] = defaultdict(set) + count: dict[tuple[int, int], int] = {} + for sample in samples: + head, relation, tail = self._check(sample[field]) + true_head[(relation, tail)].add(head) + true_tail[(head, relation)].add(tail) + # The inverse key (tail, -relation - 1) keeps head-side and + # tail-side frequencies apart in one dict, as in the reference. + for key in ((head, relation), (tail, -relation - 1)): + count[key] = count[key] + 1 if key in count else self.count_start + # Sorted lists make the fitted state independent of sample order. + self._fitted = _FittedDicts( + {k: sorted(v) for k, v in true_head.items()}, + {k: sorted(v) for k, v in true_tail.items()}, + count, + ) + + def process(self, value: Any) -> torch.Tensor: + """Returns the triple as a ``LongTensor`` of shape ``(3,)``. + + Args: + value: A triple ``(head, relation, tail)`` of integer ids. + """ + triple = torch.as_tensor(value, dtype=torch.long) + if triple.shape != (3,): + raise ValueError(f"A triple has 3 ids, got shape {tuple(triple.shape)}.") + return triple + + def subsampling_weight(self, triples: torch.Tensor) -> torch.Tensor: + """Weights of training triples, ``1 / sqrt(count(h, r) + count(t, -r-1))``. + + Args: + triples: ``LongTensor`` of shape ``(batch, 3)``. + + Returns: + Float tensor of shape ``(batch,)``. + + Raises: + KeyError: If ``(head, relation)`` or ``(tail, -relation - 1)`` was + not seen in ``fit``. Only that is checked: a triple absent + from ``fit`` whose two pairs were each seen in other triples + gets a weight. Pass training triples only. + """ + totals = [] + for head, relation, tail in triples.tolist(): + try: + totals.append( + self.count[(head, relation)] + self.count[(tail, -relation - 1)] + ) + except KeyError: + raise KeyError( + f"Triple {(head, relation, tail)} has a pair not seen in " + "fit; subsampling weights exist for training triples only." + ) from None + return torch.sqrt(1.0 / torch.tensor(totals, dtype=torch.float32)) + + def size(self) -> int: + """Number of entities.""" + return self.num_entities + + def is_token(self) -> bool: + """Entity and relation ids are discrete indices.""" + return True + + def schema(self) -> tuple[str, ...]: + return ("value",) + + def dim(self) -> tuple[int, ...]: + return (1,) + + def spatial(self) -> tuple[bool, ...]: + return (False,) + + def __repr__(self) -> str: + return ( + f"KGTripleProcessor(num_entities={self.num_entities}, " + f"num_relations={self.num_relations}, " + f"fitted_pairs={len(self.true_tail)})" + ) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index c3b3352e6..29a0707cb 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -452,3 +452,162 @@ def test_unpad_ground_truth_passes_through_plain_lists_unchanged(self) -> None: 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) + + +def fit_through_builder(samples, split=None): + """Fit a kg_triple processor the way set_task does, via SampleBuilder.""" + from pyhealth.datasets.sample_dataset import SampleBuilder + + builder = SampleBuilder( + input_schema={"triple": ("kg_triple", {"num_entities": 6, "num_relations": 2})}, + output_schema={}, + ) + builder.fit(samples, split=split) + return builder + + +class TestKGTripleProcessor(unittest.TestCase): + """The kg_triple processor encodes triples and keeps training-graph dicts.""" + + def setUp(self) -> None: + from pyhealth.processors import KGTripleProcessor + + self.triples = [(0, 0, 1), (0, 0, 2), (3, 1, 1), (4, 0, 2), (5, 1, 0)] + self.processor = KGTripleProcessor(num_entities=6, num_relations=2) + self.processor.fit([{"triple": t} for t in self.triples], "triple") + + def test_process_returns_a_long_triple(self) -> None: + out = self.processor.process((3, 1, 1)) + self.assertEqual(out.dtype, torch.long) + self.assertEqual(out.tolist(), [3, 1, 1]) + + def test_size_is_the_given_entity_count(self) -> None: + self.assertEqual(self.processor.size(), 6) + self.assertEqual(self.processor.num_relations, 2) + + def test_dicts_match_the_fitted_triples(self) -> None: + self.assertEqual(self.processor.true_tail[(0, 0)], [1, 2]) + self.assertEqual(self.processor.true_head[(0, 2)], [0, 4]) + self.assertEqual(self.processor.true_head[(1, 1)], [3]) + self.assertNotIn((1, 0), self.processor.true_tail) + + def test_counts_and_weights_match_link_prediction_fn(self) -> None: + """Same frequencies and weights as the 1.x task on the same triples.""" + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.link_prediction import ( + count_frequency, + ) + + self.assertEqual(self.processor.count, count_frequency(self.triples)) + legacy = [s["subsampling_weight"].item() for s in link_prediction_fn(self.triples)] + weights = self.processor.subsampling_weight(torch.tensor(self.triples)) + self.assertTrue(torch.allclose(weights, torch.tensor(legacy))) + + def test_weight_of_a_triple_with_an_unseen_pair_raises(self) -> None: + # (1, 1) never occurs as (head, relation) in the fitted triples. + with self.assertRaisesRegex(KeyError, "not seen in fit"): + self.processor.subsampling_weight(torch.tensor([[1, 1, 5]])) + + def test_malformed_and_out_of_range_triples_raise(self) -> None: + from pyhealth.processors import KGTripleProcessor + + processor = KGTripleProcessor(num_entities=3, num_relations=1) + for bad in [(0, 0, 3), (0, 1, 1), (-1, 0, 1), (0, 0)]: + with self.subTest(triple=bad), self.assertRaises(ValueError): + processor.fit([{"triple": bad}], "triple") + with self.assertRaises(ValueError): + KGTripleProcessor(num_entities=0, num_relations=1) + + def test_refit_replaces_the_dicts(self) -> None: + self.processor.fit([{"triple": (5, 1, 0)}], "triple") + self.assertEqual(self.processor.true_tail, {(5, 1): [0]}) + self.assertEqual(self.processor.true_head, {(1, 0): [5]}) + self.assertEqual(self.processor.count, {(5, 1): 4, (0, -2): 4}) + + def test_a_fitted_processor_can_be_passed_back_to_set_task(self) -> None: + """set_task puts vars() of pre-fitted processors in a JSON cache key.""" + import json + + from pyhealth.processors import KGTripleProcessor + + def key(processor): + return json.dumps({"p": vars(processor)}, sort_keys=True, default=str) + + same = KGTripleProcessor(num_entities=6, num_relations=2) + same.fit([{"triple": t} for t in reversed(self.triples)], "triple") + other = KGTripleProcessor(num_entities=6, num_relations=2) + other.fit([{"triple": t} for t in self.triples[:3]], "triple") + + self.assertEqual(key(self.processor), key(same)) + self.assertNotEqual(key(self.processor), key(other)) + + def test_survives_pickling(self) -> None: + import pickle + + clone = pickle.loads(pickle.dumps(self.processor)) + self.assertEqual(clone.true_tail, self.processor.true_tail) + self.assertEqual(clone.count, self.processor.count) + + def test_split_fits_on_the_training_part_only(self) -> None: + """No validation or test triple reaches the dicts or the counts.""" + from pyhealth.datasets import PatientSplit + + samples = [ + {"patient_id": str(i), "triple": t} for i, t in enumerate(self.triples) + ] + split = PatientSplit(ratios=(0.6, 0.4), seed=0) + builder = fit_through_builder(samples, split=split) + processor = builder.input_processors["triple"] + train_idx, held_out_idx = builder.split_indices + + train = {self.triples[i] for i in train_idx} + held_out = {self.triples[i] for i in held_out_idx} + fitted = {(h, r, t) for (h, r), ts in processor.true_tail.items() for t in ts} + self.assertEqual(fitted, train) + self.assertTrue(held_out) + self.assertFalse(fitted & held_out) + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.link_prediction import ( + count_frequency, + ) + + self.assertEqual(processor.count, count_frequency(sorted(train))) + + def test_fitting_without_split_warns(self) -> None: + samples = [{"patient_id": str(i), "triple": t} for i, t in enumerate(self.triples)] + with self.assertLogs("pyhealth.datasets.sample_dataset", level="WARNING") as logs: + fit_through_builder(samples) + self.assertTrue(any("triple" in line for line in logs.output)) + + +class TestKGProcessorKeepsLongLists(unittest.TestCase): + """Filter sets longer than the training maximum are kept whole (no truncation).""" + + def setUp(self) -> None: + from pyhealth.processors import KGProcessor + + self.processor = KGProcessor(pad_token_id=0) + # Fitted on "training" lists of at most 2 entities. + self.processor.fit([{"gt": [1, 2]}, {"gt": [3]}], "gt") + + def test_a_longer_list_keeps_all_entities(self) -> None: + out = self.processor.process([5, 0, 7, 9]) + self.assertEqual(out["value"].tolist(), [5, 0, 7, 9]) + self.assertEqual(out["mask"].tolist(), [1, 1, 1, 1]) + + def test_shorter_lists_are_still_padded_to_the_fitted_length(self) -> None: + out = self.processor.process([4]) + self.assertEqual(out["value"].tolist(), [4, 0]) + self.assertEqual(out["mask"].tolist(), [1, 0]) + + def test_collation_pads_the_batch_with_mask_zero(self) -> None: + batch = [ + {"gt": self.processor.process([5, 0, 7, 9])}, + {"gt": self.processor.process([4])}, + ] + gt = collate_fn_dict_with_padding(batch)["gt"] + self.assertEqual(gt["value"].tolist(), [[5, 0, 7, 9], [4, 0, 0, 0]]) + self.assertEqual(gt["mask"].tolist(), [[1, 1, 1, 1], [1, 0, 0, 0]]) + # The model's unpadding recovers the exact lists, entity 0 included. + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + model = TransE(dataset=make_dataset(n=4), e_dim=8, r_dim=8) + self.assertEqual(model._unpad_ground_truth(gt), [[5, 0, 7, 9], [4]]) From 8607cecaeb411ef0af966430744cc4eb9184c163 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:24:40 +0200 Subject: [PATCH 19/25] feat(kg_emb): add the KGLinkPrediction task KGLinkPrediction(BaseTask) gives one sample per triple: the triple of ids and its ground_truth_head / ground_truth_tail lists. Those lists cover the whole graph, as the filter sets of filtered ranking evaluation (Bordes et al., 2013; Sun et al., 2019); training will draw its negatives from the kg_triple processor fitted on the training part instead. set_task shards records by patient_id across workers, so the graph-wide lists are computed in pre_filter with Polars group_by and join, and __call__ reads one row. pre_filter materialises its frame once: the workers re-run the lazy plan for every batch of records, which would otherwise repeat the aggregations each time. It removes duplicate triples, keeping the first in file order, logs how many, raises on a triple with a missing field, and checks the graph has the entity and relation counts the task was built for. BaseKGDataset gains default_task, and its set_task: - warns at every call when the kg_triple processor would be fitted without split=, i.e. on every triple (not when the processor is supplied, since supplied processors are never refitted); - with split=, logs per held-out part how many triples involve an entity or a relation absent from the training triples. Ids are global, so these triples are kept and scored; their embeddings simply got no training signal. The parts are read sequentially, because litdata's indexed access on a subset is linear in the index. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/base_kg_dataset.py | 83 +++++- .../kg_emb/tasks/__init__.py | 3 + .../kg_emb/tasks/kg_link_prediction.py | 181 +++++++++++++ tests/core/test_kg_emb.py | 248 ++++++++++++++++++ 4 files changed, 512 insertions(+), 3 deletions(-) create mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py 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 0244d8d32..95e70a442 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 @@ -15,6 +15,7 @@ from __future__ import annotations +import inspect import logging import warnings from collections.abc import Callable @@ -177,6 +178,44 @@ def index_triples( ) +def _uses_kg_triple(task: BaseTask) -> bool: + spec = getattr(task, "input_schema", {}).get("triple") + name = spec[0] if isinstance(spec, tuple) else spec + return name == "kg_triple" + + +def _log_unseen_in_training(parts: tuple) -> None: + """Logs, per held-out part, the triples whose entity or relation is not + in any training triple. + + Ids are global, so such triples are kept and scored, but their embeddings + received no training signal for that entity or relation. + """ + fitted = parts[0].input_processors["triple"] + entities = {h for h, _ in fitted.true_tail} | { + t for tails in fitted.true_tail.values() for t in tails + } + relations = {r for _, r in fitted.true_tail} + for k, part in enumerate(parts[1:], start=1): + n_entity = n_relation = 0 + # Sequential iteration: on a subset, litdata's part[i] costs O(i), + # which made this pass quadratic in the size of the part. + for sample in part: + head, relation, tail = sample["triple"].tolist() + n_entity += head not in entities or tail not in entities + n_relation += relation not in relations + level = logging.WARNING if n_entity or n_relation else logging.INFO + logger.log( + level, + "Part %d: %d of %d triples involve an entity absent from the " + "training triples, %d a relation absent from them.", + k, + n_entity, + len(part), + n_relation, + ) + + class BaseKGDataset(BaseDataset): """A knowledge graph stored as one ``triples`` table. @@ -393,8 +432,14 @@ def info() -> None: def set_task(self, task: BaseTask | Callable | None = None, *args, **kwargs): """Builds the sample dataset(s) of ``task``. - With a :class:`~pyhealth.tasks.BaseTask`, this is - :meth:`pyhealth.datasets.BaseDataset.set_task`. A plain function, + With a :class:`~pyhealth.tasks.BaseTask` (by default + :class:`KGLinkPrediction`), this is + :meth:`pyhealth.datasets.BaseDataset.set_task`. For a task whose + ``triple`` field uses the ``kg_triple`` processor, it also warns when + ``split`` is missing, since that processor would then be fitted on + every triple, and with ``split`` it logs, for each held-out part, how + many triples involve an entity or a relation absent from the training + triples. A plain function, such as ``link_prediction_fn``, goes through the deprecated path of PyHealth 1.x: it receives the list of indexed triples, in file order, and its samples are wrapped in a :class:`SampleKGDataset`, with the @@ -407,7 +452,39 @@ def set_task(self, task: BaseTask | Callable | None = None, *args, **kwargs): task = kwargs.pop("task_fn") if task is not None and not isinstance(task, BaseTask) and callable(task): return self._set_legacy_task(task, *args, **kwargs) - return super().set_task(task, *args, **kwargs) + if task is None: + task = self.default_task + + bound = inspect.signature(BaseDataset.set_task).bind( + self, task, *args, **kwargs + ).arguments + split = bound.get("split") + # A supplied processor is used as it is, never refitted. + fits_triples = _uses_kg_triple(task) and "triple" not in ( + bound.get("input_processors") or {} + ) + if fits_triples and split is None: + warnings.warn( + "set_task without split= fits the kg_triple processor on every " + "triple, so training negatives and subsampling weights depend on " + "validation and test triples. Pass split=PatientSplit(...) to fit " + "it on the training triples only.", + UserWarning, + stacklevel=2, + ) + result = super().set_task(task, *args, **kwargs) + if _uses_kg_triple(task) and isinstance(result, tuple): + _log_unseen_in_training(result) + return result + + @property + def default_task(self) -> BaseTask: + """:class:`KGLinkPrediction` on this graph.""" + from ..tasks.kg_link_prediction import KGLinkPrediction + + return KGLinkPrediction( + num_entities=self.num_entities, num_relations=self.num_relations + ) def _set_legacy_task( self, diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py index b86bfcc87..a445a57ac 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py @@ -1 +1,4 @@ +from .kg_link_prediction import KGLinkPrediction from .link_prediction import link_prediction_fn + +__all__ = ["KGLinkPrediction", "link_prediction_fn"] diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py new file mode 100644 index 000000000..bbda112f4 --- /dev/null +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py @@ -0,0 +1,181 @@ +"""Link prediction on a :class:`BaseKGDataset`, one sample per triple.""" + +from __future__ import annotations + +import logging +from typing import Any + +import polars as pl + +from pyhealth.data import Patient +from pyhealth.tasks import BaseTask + +from ..datasets.base_kg_dataset import ( + FIELDS, + TABLE, + _count_incomplete, + entity_vocabulary, + index_triples, + relation_vocabulary, + triples_frame, +) + +logger = logging.getLogger(__name__) + +__all__ = ["KGLinkPrediction"] + + +class KGLinkPrediction(BaseTask): + """Link prediction: predict the tail of ``(h, r, ?)`` and the head of ``(?, r, t)``. + + Each triple of the graph gives one sample:: + + {"patient_id": ..., "triple": (h, r, t), + "ground_truth_head": [h' : (h', r, t) in the graph], + "ground_truth_tail": [t' : (h, r, t') in the graph]} + + The ground-truth lists cover the whole graph. They are the filter sets of + filtered ranking evaluation (Bordes et al., 2013; Sun et al., 2019): + when a test triple is ranked, the other known true entities are removed + from the candidates, whatever part they belong to. Training does not use + them; it draws its negatives from the training triples only, through the + ``kg_triple`` processor fitted on the training part. + + Because ``set_task`` builds samples patient by patient, and each triple + is its own patient, these graph-wide lists are computed beforehand, in + :meth:`pre_filter`, with Polars ``group_by`` and ``join``. + + Duplicate triples are removed in :meth:`pre_filter`, keeping the first in + file order, and their number is logged: a duplicate could otherwise land + in both the training and the test part. + + Args: + num_entities: Number of entities of the graph + (``dataset.num_entities``). + num_relations: Number of relations of the graph + (``dataset.num_relations``). + + Examples: + >>> task = KGLinkPrediction(num_entities=4, num_relations=2) + >>> sorted(task.input_schema) + ['ground_truth_head', 'ground_truth_tail', 'triple'] + >>> task.input_schema["triple"] + ('kg_triple', {'num_entities': 4, 'num_relations': 2}) + >>> train, val, test = dataset.set_task( # doctest: +SKIP + ... task, split=PatientSplit(ratios=(0.8, 0.1, 0.1), seed=0) + ... ) + """ + + task_name: str = "KGLinkPrediction" + + def __init__(self, num_entities: int, num_relations: int): + for name, value in (("num_entities", num_entities), ("num_relations", num_relations)): + if not isinstance(value, int) or isinstance(value, bool) or value < 1: + raise ValueError(f"{name} must be a positive integer, got {value!r}.") + self.num_entities = num_entities + self.num_relations = num_relations + self.input_schema = { + "triple": ( + "kg_triple", + {"num_entities": num_entities, "num_relations": num_relations}, + ), + "ground_truth_head": ("kg_entity_list", {"pad_token_id": 0}), + "ground_truth_tail": ("kg_entity_list", {"pad_token_id": 0}), + } + self.output_schema = {} + super().__init__() + + def pre_filter(self, df: pl.LazyFrame) -> pl.LazyFrame: + """Indexes the triples, removes duplicates and adds the ground truths. + + Args: + df: The dataset's ``global_event_df``. + + Returns: + One event per distinct triple, with the attributes ``head_id``, + ``relation_id``, ``tail_id``, ``ground_truth_head`` and + ``ground_truth_tail``. The frame is computed here, once: + ``set_task`` reads it again for every batch of records, which + would otherwise repeat the graph-wide aggregations each time. + + Raises: + ValueError: If a triple has a missing field, or if the graph + does not have ``num_entities`` entities and + ``num_relations`` relations. + """ + triples = triples_frame(df) + incomplete = _count_incomplete(triples) + if incomplete: + raise ValueError( + f"{incomplete} triple(s) have a missing head, relation or tail." + ) + + entities = entity_vocabulary(triples).collect() + relations = relation_vocabulary(triples).collect() + if (len(entities), len(relations)) != (self.num_entities, self.num_relations): + raise ValueError( + f"The task expects {self.num_entities} entities and " + f"{self.num_relations} relations; the graph has {len(entities)} " + f"and {len(relations)}. Build it from this dataset's " + "num_entities and num_relations." + ) + + indexed = ( + index_triples(triples, entities.lazy(), relations.lazy()) + .with_columns(row=pl.col("patient_id").cast(pl.Int64)) + .sort("row") + .collect() + ) + distinct = indexed.unique(subset=list(FIELDS), keep="first", maintain_order=True) + n_duplicates = len(indexed) - len(distinct) + if n_duplicates: + logger.warning( + "Removed %d duplicate triple(s), keeping the first occurrence of " + "each; %d distinct triples remain.", + n_duplicates, + len(distinct), + ) + + ground_truth_head = distinct.group_by("relation_id", "tail_id").agg( + pl.col("head_id").sort().alias("ground_truth_head") + ) + ground_truth_tail = distinct.group_by("head_id", "relation_id").agg( + pl.col("tail_id").sort().alias("ground_truth_tail") + ) + samples = ( + distinct.join(ground_truth_head, on=["relation_id", "tail_id"], how="left") + .join(ground_truth_tail, on=["head_id", "relation_id"], how="left") + .select( + "patient_id", + pl.lit(TABLE).alias("event_type"), + pl.lit(None, dtype=pl.Datetime("ms")).alias("timestamp"), + *( + pl.col(c).alias(f"{TABLE}/{c}") + for c in ( + "head_id", + "relation_id", + "tail_id", + "ground_truth_head", + "ground_truth_tail", + ) + ), + ) + .sort("patient_id") + ) + return samples.lazy() + + def __call__(self, patient: Patient) -> list[dict[str, Any]]: + """Returns the sample of the record's triple. + + Args: + patient: One record of the frame returned by :meth:`pre_filter`. + """ + return [ + { + "patient_id": patient.patient_id, + "triple": (event.head_id, event.relation_id, event.tail_id), + "ground_truth_head": list(event.ground_truth_head), + "ground_truth_tail": list(event.ground_truth_tail), + } + for event in patient.get_events(event_type=TABLE) + ] diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 29a0707cb..04ff906fa 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -8,6 +8,7 @@ from __future__ import annotations import hashlib +import logging import tempfile import unittest import warnings @@ -611,3 +612,250 @@ def test_collation_pads_the_batch_with_mask_zero(self) -> None: model = TransE(dataset=make_dataset(n=4), e_dim=8, r_dim=8) self.assertEqual(model._unpad_ground_truth(gt), [[5, 0, 7, 9], [4]]) + + +# Twelve rows, one of them a duplicate of the first: 11 distinct triples +# over entities a..f (f appears once, so a split can leave it out of training). +TASK_TRIPLES = [ + ("a", "r1", "b"), + ("a", "r1", "c"), + ("b", "r1", "c"), + ("c", "r2", "a"), + ("d", "r2", "a"), + ("a", "r2", "d"), + ("b", "r2", "d"), + ("e", "r1", "a"), + ("a", "r1", "b"), + ("c", "r1", "e"), + ("d", "r1", "b"), + ("e", "r2", "f"), +] +TASK_SPLIT_RATIOS = (0.6, 0.2, 0.2) + + +def unpad(field: dict[str, torch.Tensor]) -> list[int]: + """The real entities of one processed kg_entity_list field.""" + return field["value"][field["mask"].bool()].tolist() + + +class _ListHandler(logging.Handler): + def __init__(self) -> None: + super().__init__(level=logging.DEBUG) + self.messages: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +class TestKGLinkPrediction(unittest.TestCase): + """KGLinkPrediction through the real set_task, with a triple-level split.""" + + @classmethod + def setUpClass(cls) -> None: + from pyhealth.datasets import PatientSplit + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + cls._tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + cls.root = Path(cls._tmp.name) + config = write_kg(cls.root, TASK_TRIPLES) + cls.dataset = BaseKGDataset( + root=str(cls.root), config_path=config, cache_dir=cls.root / "cache" + ) + cls.task = KGLinkPrediction( + num_entities=cls.dataset.num_entities, + num_relations=cls.dataset.num_relations, + ) + cls.split = PatientSplit(ratios=TASK_SPLIT_RATIOS, seed=3) + handler = _ListHandler() + kg_logger = logging.getLogger("pyhealth.medcode.pretrained_embeddings.kg_emb") + kg_logger.addHandler(handler) + kg_logger.setLevel(logging.INFO) + try: + cls.parts = cls.dataset.set_task(cls.task, split=cls.split, num_workers=1) + finally: + kg_logger.removeHandler(handler) + cls.log = handler.messages + + e, r = cls.dataset.entity2id, cls.dataset.relation2id + cls.distinct = sorted({(e[h], r[rel], e[t]) for h, rel, t in TASK_TRIPLES}) + + @classmethod + def tearDownClass(cls) -> None: + cls._tmp.cleanup() + + def triples_of(self, part) -> list[tuple[int, int, int]]: + return [tuple(part[i]["triple"].tolist()) for i in range(len(part))] + + def test_parts_are_streaming_sample_datasets(self) -> None: + from pyhealth.datasets import SampleDataset + from pyhealth.datasets.sample_dataset import InMemorySampleDataset + + self.assertEqual(len(self.parts), 3) + for part in self.parts: + self.assertIsInstance(part, SampleDataset) + self.assertNotIsInstance(part, InMemorySampleDataset) + + def test_parts_partition_the_distinct_triples(self) -> None: + triples = [t for part in self.parts for t in self.triples_of(part)] + self.assertEqual(sorted(triples), self.distinct) + self.assertEqual(len(triples), 11) + + def test_duplicates_are_removed_and_logged(self) -> None: + self.assertTrue(any("Removed 1 duplicate" in m for m in self.log), self.log) + + def test_the_first_duplicate_in_file_order_is_kept(self) -> None: + # (a, r1, b) is on rows 0 and 8; row 0 must be the record kept, since + # the record id decides which part the triple falls into. + e, r = self.dataset.entity2id, self.dataset.relation2id + duplicate = (e["a"], r["r1"], e["b"]) + kept = [ + part[i]["patient_id"] + for part in self.parts + for i in range(len(part)) + if tuple(part[i]["triple"].tolist()) == duplicate + ] + self.assertEqual(kept, ["0"]) + + def test_ground_truth_covers_the_whole_graph(self) -> None: + for part in self.parts: + for i in range(len(part)): + sample = part[i] + h, r, t = sample["triple"].tolist() + with self.subTest(triple=(h, r, t)): + self.assertEqual( + unpad(sample["ground_truth_head"]), + sorted(x for x, y, z in self.distinct if (y, z) == (r, t)), + ) + self.assertEqual( + unpad(sample["ground_truth_tail"]), + sorted(z for x, y, z in self.distinct if (x, y) == (h, r)), + ) + + def test_triple_processor_is_fitted_on_training_triples_only(self) -> None: + processor = self.parts[0].input_processors["triple"] + fitted = {(h, r, t) for (h, r), ts in processor.true_tail.items() for t in ts} + self.assertEqual(fitted, set(self.triples_of(self.parts[0]))) + held_out = set(self.triples_of(self.parts[1]) + self.triples_of(self.parts[2])) + self.assertTrue(held_out) + self.assertFalse(fitted & held_out) + self.assertEqual(processor.size(), self.dataset.num_entities) + self.assertEqual(processor.num_relations, self.dataset.num_relations) + + def test_unseen_entities_are_logged_per_held_out_part(self) -> None: + processor = self.parts[0].input_processors["triple"] + seen = {h for h, _ in processor.true_tail} | { + t for ts in processor.true_tail.values() for t in ts + } + for k in (1, 2): + expected = sum( + h not in seen or t not in seen for h, _, t in self.triples_of(self.parts[k]) + ) + prefix = f"Part {k}: {expected} of {len(self.parts[k])} triples" + self.assertTrue(any(m.startswith(prefix) for m in self.log), self.log) + + def test_two_workers_give_the_same_samples(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + dataset = BaseKGDataset( + root=str(self.root), + config_path=self.root / "kg.yaml", + cache_dir=Path(tmp), + num_workers=2, + ) + parts = dataset.set_task(self.task, split=self.split, num_workers=2) + + def content(part): + return sorted( + ( + tuple(s["triple"].tolist()), + tuple(unpad(s["ground_truth_head"])), + tuple(unpad(s["ground_truth_tail"])), + ) + for s in (part[i] for i in range(len(part))) + ) + + for mine, theirs in zip(self.parts, parts): + self.assertEqual(content(mine), content(theirs)) + + def test_without_split_warns_at_every_call(self) -> None: + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + dataset = BaseKGDataset( + root=str(self.root), config_path=self.root / "kg.yaml", cache_dir=Path(tmp) + ) + with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): + samples = dataset.set_task(self.task) + self.assertEqual(len(samples), 11) + # Second call: the processed cache exists, the warning still fires. + with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): + dataset.set_task(self.task) + # Without a task, the default KGLinkPrediction is used, and warned about. + with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): + dataset.set_task() + + def test_default_task_matches_the_graph(self) -> None: + task = self.dataset.default_task + self.assertEqual( + (task.num_entities, task.num_relations), + (self.dataset.num_entities, self.dataset.num_relations), + ) + + def test_a_task_for_another_graph_is_rejected(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + task = KGLinkPrediction(num_entities=99, num_relations=2) + with self.assertRaisesRegex(ValueError, "expects 99 entities"): + self.dataset.set_task(task, split=self.split) + + def test_supplied_processors_are_used_as_they_are(self) -> None: + """A supplied kg_triple processor is neither refitted nor warned about.""" + from pyhealth.processors import KGTripleProcessor + + everything = KGTripleProcessor( + num_entities=self.dataset.num_entities, + num_relations=self.dataset.num_relations, + ) + everything.fit([{"triple": t} for t in self.distinct], "triple") + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + parts = self.dataset.set_task( + self.task, split=self.split, input_processors={"triple": everything} + ) + whole = self.dataset.set_task( + self.task, input_processors={"triple": everything} + ) + # Fitted on the 11 triples, not refitted on the 6 training ones. + for result in (parts[0], whole): + self.assertEqual(result.input_processors["triple"].true_tail, everything.true_tail) + + +class TestUnseenInTrainingLog(unittest.TestCase): + """The per-part count of triples whose entity or relation training never saw.""" + + def test_counts_heads_tails_and_relations(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import base_kg_dataset + from pyhealth.processors import KGTripleProcessor + + processor = KGTripleProcessor(num_entities=6, num_relations=3) + processor.fit([{"triple": (0, 0, 1)}, {"triple": (1, 1, 2)}], "triple") + + class _Train(list): + input_processors = {"triple": processor} + + def part(*triples): + return [{"triple": torch.tensor(t)} for t in triples] + + parts = ( + _Train(), + part((0, 0, 2), (4, 0, 1)), # seen; unseen head + part((0, 0, 5), (1, 2, 2), (3, 2, 4)), # unseen tail; unseen relation; both + ) + with self.assertLogs(base_kg_dataset.logger, level="INFO") as logs: + base_kg_dataset._log_unseen_in_training(parts) + self.assertEqual( + [line.split(":", 2)[2] for line in logs.output], + [ + "Part 1: 1 of 2 triples involve an entity absent from the training " + "triples, 0 a relation absent from them.", + "Part 2: 2 of 3 triples involve an entity absent from the training " + "triples, 2 a relation absent from them.", + ], + ) From 349bcbd74c6b19e9408e52b63073ed943a32b3fb Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:28:38 +0200 Subject: [PATCH 20/25] refactor(kg_emb): make the KGE models consume a SampleDataset - The numbers of entities and relations come from the dataset's fitted kg_triple processor, which the model keeps as triple_processor. - Training and evaluation follow model.train() / model.eval(), as the Trainer sets them, instead of the per-sample train flag. A 1.x batch whose flag disagrees with the mode triggers a warning. - negative_sampling is a model argument. It defaults to the value a 1.x SampleKGDataset carries, else to 128, the default of the reference implementation of Sun et al. (2019). - With use_subsampling_weight, the weights come from the processor (training triples) when the samples do not carry them. A 1.x SampleKGDataset is still accepted. Training negatives are still filtered with the samples' ground-truth lists; the next commit changes that. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/models/complex.py | 19 +++- .../kg_emb/models/distmult.py | 17 ++- .../kg_emb/models/kg_base.py | 104 ++++++++++++++---- .../kg_emb/models/rotate.py | 19 +++- .../kg_emb/models/transe.py | 19 +++- tests/core/test_kg_emb.py | 100 +++++++++++++++++ 6 files changed, 233 insertions(+), 45 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 8dae884af..bb9be0888 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -2,6 +2,8 @@ import torch +from pyhealth.datasets import SampleDataset + from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -13,10 +15,11 @@ class ComplEx(KGEBaseModel): Complex embeddings for simple link prediction. In International conference on machine learning (pp. 2071-2080). PMLR Examples: + >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } >>> model = ComplEx(_Toy(), e_dim=4, r_dim=4, ns="uniform") >>> tuple(model.E_emb.shape) (2, 4) @@ -24,16 +27,20 @@ class ComplEx(KGEBaseModel): def __init__( self, - dataset: SampleKGDataset, + dataset: SampleDataset, e_dim: int = 600, r_dim: int = 600, ns: str = "adv", gamma: float = 24.0, use_subsampling_weight: bool = False, use_regularization: str = 'l3', - mode: str = "multiclass" + mode: str = "multiclass", + negative_sampling: int | None = None, ): - super().__init__(dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, use_regularization, mode) + super().__init__( + dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, + use_regularization, mode, negative_sampling=negative_sampling, + ) def regularization(self, sample_batch, mode='pos'): diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index c7422e480..74bcc7141 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -2,6 +2,8 @@ import torch +from pyhealth.datasets import SampleDataset + from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -13,17 +15,18 @@ class DistMult(KGEBaseModel): relations for learning and inference in knowledge bases. ICLR 2015. Examples: + >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } >>> model = DistMult(_Toy(), e_dim=4, r_dim=4, ns="uniform") >>> tuple(model.E_emb.shape) (2, 4) """ def __init__( self, - dataset: SampleKGDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", @@ -31,8 +34,12 @@ def __init__( use_subsampling_weight: bool = False, use_regularization: str = 'l3', mode: str = "multiclass", + negative_sampling: int | None = None, ): - super().__init__(dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, use_regularization, mode) + super().__init__( + dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, + use_regularization, mode, negative_sampling=negative_sampling, + ) def regularization(self, sample_batch, mode='pos'): 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 bdcb09c60..441dc68c6 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,3 +1,4 @@ +import warnings from abc import ABC import numpy as np @@ -5,15 +6,28 @@ import torch.nn.functional as F from torch import nn -from ..datasets.sample_kg_dataset import SampleKGDataset +from pyhealth.datasets import SampleDataset +from pyhealth.processors import KGTripleProcessor + +# Negatives per positive triple when the caller does not choose: the +# default of the reference implementation of Sun et al. (2019). +DEFAULT_NEGATIVE_SAMPLING = 128 class KGEBaseModel(ABC, nn.Module): """Abstract class for Knowledge Graph Embedding models. + The numbers of entities and relations come from the dataset's fitted + ``kg_triple`` processor (``KGLinkPrediction`` with + ``set_task(..., split=PatientSplit(...))``). Training and evaluation + follow ``model.train()`` / ``model.eval()``, as the + :class:`~pyhealth.trainer.Trainer` sets them. + Args: - e_num: the number of entities in the dataset. - r_num: the number of relations in the dataset. + dataset: A sample dataset whose ``triple`` field uses the + ``kg_triple`` processor, typically the training part returned + by ``set_task``. A PyHealth 1.x ``SampleKGDataset`` is still + accepted for one release. e_dim: the hidden embedding size for entity, 500 by default. r_dim: the hidden embedding size for relation, 500 by default. ns: negative sampling technique to use: can be "uniform", "normal" or "adv" (self-adversarial). @@ -21,15 +35,21 @@ class KGEBaseModel(ABC, nn.Module): use_subsampling_weight: whether to use subsampling weight (like in word2vec) or not, False by default. use_regularization: whether to apply regularization or not, False by default. mode: evaluation metric type, one of "binary", "multiclass", or "multilabel", "multiclass" by default + negative_sampling: Number of negative heads and of negative tails drawn + per training triple. Defaults to the ``negative_sampling`` of a + 1.x ``SampleKGDataset`` if it has one, else 128. Examples: + >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } >>> model = KGEBaseModel(_Toy(), e_dim=4, r_dim=4, ns="uniform") - >>> model.e_num, tuple(model.E_emb.shape) - (2, (2, 4)) + >>> model.e_num, model.r_num, tuple(model.E_emb.shape) + (2, 1, (2, 4)) + >>> model.negative_sampling + 128 """ @property @@ -39,19 +59,36 @@ def device(self): def __init__( - self, - dataset: SampleKGDataset, + self, + dataset: SampleDataset, e_dim: int = 500, r_dim: int = 500, ns: str = "uniform", gamma: float | None = None, use_subsampling_weight: bool = False, use_regularization: str | None = None, - mode: str = "multiclass" + mode: str = "multiclass", + negative_sampling: int | None = None, ): super().__init__() - self.e_num = dataset.entity_num - self.r_num = dataset.relation_num + processor = getattr(dataset, "input_processors", {}).get("triple") + if isinstance(processor, KGTripleProcessor): + self.triple_processor: KGTripleProcessor | None = processor + self.e_num = processor.num_entities + self.r_num = processor.num_relations + else: + # A PyHealth 1.x SampleKGDataset carries the counts itself. + self.triple_processor = None + self.e_num = dataset.entity_num + self.r_num = dataset.relation_num + if negative_sampling is None: + legacy = getattr(dataset, "task_spec_param", None) or {} + negative_sampling = legacy.get("negative_sampling", DEFAULT_NEGATIVE_SAMPLING) + if not isinstance(negative_sampling, int) or negative_sampling < 1: + raise ValueError( + f"negative_sampling must be a positive integer, got {negative_sampling!r}." + ) + self.negative_sampling = negative_sampling self.e_dim = e_dim self.r_dim = r_dim self.ns = ns @@ -327,15 +364,36 @@ def calc(self, head, relation, tail, mode='pos'): raise NotImplementedError + def _subsampling_weight(self, data, positive_sample): + if "subsampling_weight" in data: + # PyHealth 1.x samples carry their weight. + return torch.cat([d for d in data["subsampling_weight"]], dim=0) + if self.triple_processor is None: + raise ValueError("use_subsampling_weight needs a kg_triple processor.") + return self.triple_processor.subsampling_weight(positive_sample.cpu()) + def forward(self, **data): - positive_sample = torch.stack([torch.LongTensor(d) for d in data['triple']], dim=0).to(self.device) + triples = data["triple"] + if not isinstance(triples, torch.Tensor): + triples = torch.stack([torch.as_tensor(d) for d in triples], dim=0) + positive_sample = triples.long().to(self.device) + + if "train" in data and bool(data["train"][0]) != self.training: + # PyHealth 1.x samples carry a per-sample train flag, which no + # longer decides the branch. + warnings.warn( + "The samples' 'train' flag is ignored: the model trains or " + "evaluates according to model.train() / model.eval().", + UserWarning, + stacklevel=2, + ) - if data['train'][0]: + if self.training: negative_sample_head, negative_sample_tail = self.train_neg_sample_gen( gt_head=data['ground_truth_head'], gt_tail=data['ground_truth_tail'], - negative_sampling=data['hyperparameters'][0]['negative_sampling'] + negative_sampling=self.negative_sampling, ) negative_sample_head, negative_sample_tail = negative_sample_head.to(self.device), negative_sample_tail.to(self.device) @@ -356,11 +414,13 @@ def forward(self, **data): head, relation, tail = self.data_process((positive_sample), mode="pos") pos_score = F.logsigmoid(self.calc(head=head, relation=relation, tail=tail)).squeeze(dim=1) - subsampling_weight = torch.cat([d for d in data['subsampling_weight']], dim=0).to(self.device) - pos_sample_loss = \ - - (subsampling_weight * pos_score).sum()/subsampling_weight.sum() if self.use_subsampling_weight else (- pos_score.mean()) - neg_sample_loss = \ - - (subsampling_weight * neg_score).sum()/subsampling_weight.sum() if self.use_subsampling_weight else (- neg_score.mean()) + if self.use_subsampling_weight: + subsampling_weight = self._subsampling_weight(data, positive_sample).to(self.device) + pos_sample_loss = - (subsampling_weight * pos_score).sum() / subsampling_weight.sum() + neg_sample_loss = - (subsampling_weight * neg_score).sum() / subsampling_weight.sum() + else: + pos_sample_loss = - pos_score.mean() + neg_sample_loss = - neg_score.mean() loss = (pos_sample_loss + neg_sample_loss) / 2 @@ -373,7 +433,7 @@ def forward(self, **data): else: # valid/test inputs = self.test_neg_sample_filter_bias_gen( - triples=data['triple'], + triples=positive_sample.tolist(), gt_head=data['ground_truth_head'], gt_tail=data['ground_truth_tail'] ) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index 48c254dd3..b057a642a 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -2,6 +2,8 @@ import torch +from pyhealth.datasets import SampleDataset + from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -13,10 +15,11 @@ class RotatE(KGEBaseModel): Rotate: Knowledge graph embedding by relational rotation in complex space. ICLR 2019. Examples: + >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } >>> model = RotatE(_Toy(), e_dim=4, r_dim=2, ns="uniform") >>> tuple(model.E_emb.shape) (2, 4) @@ -24,16 +27,20 @@ class RotatE(KGEBaseModel): def __init__( self, - dataset: SampleKGDataset, + dataset: SampleDataset, e_dim: int = 600, r_dim: int = 300, ns='adv', gamma=24.0, use_subsampling_weight: bool = False, use_regularization: str | None = None, - mode: str = "multiclass" + mode: str = "multiclass", + negative_sampling: int | None = None, ): - super().__init__(dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, use_regularization, mode) + super().__init__( + dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, + use_regularization, mode, negative_sampling=negative_sampling, + ) self.pi = 3.14159265358979323846 def regularization(self, sample_batch, mode='pos'): diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 9ba220ca6..1d875ec8e 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -2,6 +2,8 @@ import torch +from pyhealth.datasets import SampleDataset + from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -13,10 +15,11 @@ class TransE(KGEBaseModel): Translating embeddings for modeling multi-relational data. NIPS 2013. Examples: + >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: - ... entity_num = 2 - ... relation_num = 1 - ... task_spec_param = None + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } >>> model = TransE(_Toy(), e_dim=4, r_dim=4, ns="uniform") >>> tuple(model.E_emb.shape) (2, 4) @@ -24,7 +27,7 @@ class TransE(KGEBaseModel): def __init__( self, - dataset: SampleKGDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", @@ -32,9 +35,13 @@ def __init__( use_subsampling_weight: bool = False, use_regularization: str | None = None, mode: str = "multiclass", - p_norm: int = 1.0 + p_norm: int = 1.0, + negative_sampling: int | None = None, ): - super().__init__(dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, use_regularization, mode) + super().__init__( + dataset, e_dim, r_dim, ns, gamma, use_subsampling_weight, + use_regularization, mode, negative_sampling=negative_sampling, + ) self.p_norm = p_norm diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 04ff906fa..da96e3f17 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -859,3 +859,103 @@ def part(*triples): "triples, 2 a relation absent from them.", ], ) + + +class TestModelsOnSampleDatasets(unittest.TestCase): + """The KGE models consume the SampleDatasets returned by set_task.""" + + @classmethod + def setUpClass(cls) -> None: + from pyhealth.datasets import PatientSplit + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + cls._tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + root = Path(cls._tmp.name) + config = write_kg(root, TASK_TRIPLES) + cls.dataset = BaseKGDataset(root=str(root), config_path=config, cache_dir=root / "cache") + task = KGLinkPrediction( + num_entities=cls.dataset.num_entities, num_relations=cls.dataset.num_relations + ) + cls.train, cls.val, cls.test = cls.dataset.set_task( + task, split=PatientSplit(ratios=TASK_SPLIT_RATIOS, seed=3) + ) + + @classmethod + def tearDownClass(cls) -> None: + cls._tmp.cleanup() + + def batch(self, part, size=4): + from pyhealth.datasets import get_dataloader + + return next(iter(get_dataloader(part, batch_size=size))) + + def test_counts_come_from_the_fitted_processor(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + RotatE, + TransE, + ) + + for cls in (TransE, RotatE, DistMult, ComplEx): + with self.subTest(model=cls.__name__): + model = cls(dataset=self.train, e_dim=8, r_dim=4 if cls is RotatE else 8) + self.assertEqual(model.e_num, self.dataset.num_entities) + self.assertEqual(model.r_num, self.dataset.num_relations) + self.assertIs(model.triple_processor, self.train.input_processors["triple"]) + + def test_negative_sampling_is_a_model_argument(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + self.assertEqual(TransE(dataset=self.train, e_dim=8, r_dim=8).negative_sampling, 128) + model = TransE(dataset=self.train, e_dim=8, r_dim=8, negative_sampling=5) + self.assertEqual(model.negative_sampling, 5) + # A 1.x SampleKGDataset still provides its own value. + self.assertEqual(TransE(dataset=make_dataset(n=4), e_dim=8, r_dim=8).negative_sampling, 4) + with self.assertRaises(ValueError): + TransE(dataset=self.train, e_dim=8, r_dim=8, negative_sampling=0) + + def test_train_mode_returns_a_loss_and_draws_negatives(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + torch.manual_seed(0) + model = TransE(dataset=self.train, e_dim=8, r_dim=8, negative_sampling=5) + model.train() + out = model(**self.batch(self.train)) + self.assertEqual(set(out), {"loss"}) + out["loss"].backward() + self.assertIsNotNone(model.E_emb.grad) + + def test_subsampling_weights_come_from_the_processor(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + model = TransE(dataset=self.train, e_dim=8, r_dim=8, use_subsampling_weight=True) + batch = self.batch(self.train) + weights = model._subsampling_weight(batch, batch["triple"]) + expected = self.train.input_processors["triple"].subsampling_weight(batch["triple"]) + self.assertTrue(torch.equal(weights, expected)) + model.train() + model(**batch)["loss"].backward() + + def test_eval_mode_ranks_every_entity(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + model = TransE(dataset=self.train, e_dim=8, r_dim=8) + model.eval() + batch = self.batch(self.test, size=len(self.test)) + with torch.no_grad(): + out = model(**batch) + n = len(self.test) + self.assertEqual(tuple(out["y_prob"].shape), (2 * n, self.dataset.num_entities)) + self.assertEqual(out["y_true"].tolist(), batch["triple"][:, 0].tolist() + batch["triple"][:, 2].tolist()) + + def test_a_disagreeing_1x_train_flag_warns(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + train, _, _ = split(make_dataset(n=4), [1.0, 0.0, 0.0], seed=0) + batch = collate_fn_dict_with_padding(train[:2]) + model = TransE(dataset=make_dataset(n=4), e_dim=8, r_dim=8) + model.eval() + with self.assertWarnsRegex(UserWarning, "'train' flag is ignored"), torch.no_grad(): + out = model(**batch) + self.assertIn("y_prob", out) # model.eval() decides, not the flag From a141bb422881cccab10f548fb2be737e89b5475b Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:34:31 +0200 Subject: [PATCH 21/25] fix(kg_emb): draw training negatives from training triples only Training negatives were filtered with the samples' ground-truth lists, which cover the whole graph. Every validation and test positive was thus kept out of the training negatives, so training depended on the held-out triples. With a kg_triple processor, the filter now comes from its true_head / true_tail dicts, fitted on the training part only, as in the reference implementation of Sun et al. (2019). The full-graph lists remain what they are for: filtered ranking at evaluation (Bordes et al., 2013; Sun et al., 2019). A 1.x SampleKGDataset has no such processor and keeps its former behaviour until it is removed. The regression tests fail when the full-graph filter is restored: a tail true only in a test triple, and on a real split a head true only via a test triple, must remain drawable as training negatives. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/models/kg_base.py | 30 +++++++- tests/core/test_kg_emb.py | 77 +++++++++++++++++++ 2 files changed, 105 insertions(+), 2 deletions(-) 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 441dc68c6..118e7c711 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -364,6 +364,31 @@ def calc(self, head, relation, tail, mode='pos'): raise NotImplementedError + def _training_filters(self, data, positive_sample): + """The entities kept out of each triple's training negatives. + + With a ``kg_triple`` processor, these are its ``true_head`` / + ``true_tail`` dicts, fitted on the training triples only, as in the + reference implementation of Sun et al. (2019). The samples' + ``ground_truth_*`` lists cover the whole graph and serve filtered + evaluation; using them here would keep every validation and test + positive out of the training negatives, so training would depend on + the held-out triples. A 1.x ``SampleKGDataset`` has no processor and + keeps its former behaviour. + + Returns: + Two lists of entity-id lists, for head and tail negatives. + """ + if self.triple_processor is None: + return data["ground_truth_head"], data["ground_truth_tail"] + true_head = self.triple_processor.true_head + true_tail = self.triple_processor.true_tail + gt_head, gt_tail = [], [] + for head, relation, tail in positive_sample.tolist(): + gt_head.append(true_head.get((relation, tail), [])) + gt_tail.append(true_tail.get((head, relation), [])) + return gt_head, gt_tail + def _subsampling_weight(self, data, positive_sample): if "subsampling_weight" in data: # PyHealth 1.x samples carry their weight. @@ -390,9 +415,10 @@ def forward(self, **data): ) if self.training: + gt_head, gt_tail = self._training_filters(data, positive_sample) negative_sample_head, negative_sample_tail = self.train_neg_sample_gen( - gt_head=data['ground_truth_head'], - gt_tail=data['ground_truth_tail'], + gt_head=gt_head, + gt_tail=gt_tail, negative_sampling=self.negative_sampling, ) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index da96e3f17..255ec0a8e 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -959,3 +959,80 @@ def test_a_disagreeing_1x_train_flag_warns(self) -> None: with self.assertWarnsRegex(UserWarning, "'train' flag is ignored"), torch.no_grad(): out = model(**batch) self.assertIn("y_prob", out) # model.eval() decides, not the flag + + +class TestTrainingNegativesUseTrainingTriplesOnly(unittest.TestCase): + """Regression test for the leak: training negatives were filtered with + ground truths of the whole graph, so held-out positives were never drawn.""" + + def test_a_tail_true_only_in_test_can_be_a_training_negative(self) -> None: + import numpy as np + + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + from pyhealth.processors import KGTripleProcessor + + # Training triple (0, 0, 1); (0, 0, 2) is a test triple. + processor = KGTripleProcessor(num_entities=3, num_relations=1) + processor.fit([{"triple": (0, 0, 1)}], "triple") + + class _Train: + input_processors = {"triple": processor} + + model = TransE(dataset=_Train(), e_dim=4, r_dim=4, negative_sampling=64) + model.train() + # The sample's ground truth covers the whole graph, as KGLinkPrediction gives it. + batch = { + "triple": torch.tensor([[0, 0, 1]]), + "ground_truth_head": [[0]], + "ground_truth_tail": [[1, 2]], + } + drawn = {} + original = model.train_neg_sample_gen + + def record(**kwargs): + drawn["head"], drawn["tail"] = original(**kwargs) + return drawn["head"], drawn["tail"] + + model.train_neg_sample_gen = record + np.random.seed(0) + model(**batch)["loss"].backward() + + tails = set(drawn["tail"].flatten().tolist()) + self.assertNotIn(1, tails) # the training positive stays filtered + self.assertIn(2, tails) # the test positive is a legitimate negative + self.assertNotIn(0, set(drawn["head"].flatten().tolist())) + + +class TestTrainingNegativesOnARealSplit(unittest.TestCase): + """The same property on set_task output: (b, r2, d) is a test triple and + (a, r2, d) a training one, so b must stay drawable as a head negative.""" + + # Same graph and split as TestModelsOnSampleDatasets, without its tests. + setUpClass = classmethod(TestModelsOnSampleDatasets.setUpClass.__func__) + tearDownClass = classmethod(TestModelsOnSampleDatasets.tearDownClass.__func__) + + def test_held_out_heads_are_not_filtered(self) -> None: + import numpy as np + + from pyhealth.datasets import get_dataloader + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + e, r = self.dataset.entity2id, self.dataset.relation2id + target = [e["a"], r["r2"], e["d"]] + batch = next( + b for b in get_dataloader(self.train, batch_size=1) + if b["triple"][0].tolist() == target + ) + test_triples = {tuple(self.test[i]["triple"].tolist()) for i in range(len(self.test))} + self.assertIn((e["b"], r["r2"], e["d"]), test_triples) + self.assertEqual(sorted(unpad({k: v[0] for k, v in batch["ground_truth_head"].items()})), + [e["a"], e["b"]]) + + model = TransE(dataset=self.train, e_dim=4, r_dim=4, negative_sampling=256) + model.train() + gt_head, _ = model._training_filters(batch, batch["triple"]) + self.assertEqual(gt_head, [[e["a"]]]) + np.random.seed(0) + heads, _ = model.train_neg_sample_gen(gt_head, [[]], 256) + self.assertIn(e["b"], heads.flatten().tolist()) + self.assertNotIn(e["a"], heads.flatten().tolist()) From d6dfe338cef36a28ffff1bcd79e2442ed4a3e6e9 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:40:02 +0200 Subject: [PATCH 22/25] chore(kg_emb): deprecate the 1.x entry points for one release link_prediction_fn, kg_emb.datasets.split, SampleKGDataset and building a KGE model from a dataset without a kg_triple processor now raise a DeprecationWarning that names the replacement: BaseKGDataset.set_task with the KGLinkPrediction task and split=PatientSplit(...). Behaviour and import paths are unchanged; they will be removed in the next release. The new import in link_prediction.py goes at the end of the existing, unsorted block so that the PR lint gate, which checks added lines only, does not flag the block's pre-existing violations. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/sample_kg_dataset.py | 14 ++++++++ .../kg_emb/datasets/splitter.py | 12 +++++++ .../kg_emb/models/kg_base.py | 8 +++++ .../kg_emb/tasks/link_prediction.py | 20 +++++++++++ tests/core/test_kg_emb.py | 36 +++++++++++++++++++ 5 files changed, 90 insertions(+) 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 addad928d..68d70962e 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 @@ -18,6 +18,7 @@ from __future__ import annotations +import warnings from collections.abc import Mapping, Sequence from typing import Any @@ -33,6 +34,12 @@ class SampleKGDataset(InMemorySampleDataset): r"""In-memory dataset of knowledge-graph link-prediction samples. + .. deprecated:: + Use :class:`BaseKGDataset` with ``set_task(KGLinkPrediction(...), + split=PatientSplit(...))``, which returns streaming + :class:`~pyhealth.datasets.SampleDataset` parts. This class will be + removed in the next release. + Each sample is a mapping with the following keys: ``triple`` @@ -114,6 +121,13 @@ def __init__( pad_token_id: int = 0, **task_spec_param: Any, ) -> None: + warnings.warn( + "SampleKGDataset is deprecated and will be removed in the next " + "release; use BaseKGDataset.set_task(KGLinkPrediction(...), " + "split=PatientSplit(...)).", + DeprecationWarning, + stacklevel=2, + ) input_schema = { "triple": ("tensor", {"dtype": torch.long}), "ground_truth_head": ("kg_entity_list", {"pad_token_id": pad_token_id}), diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 0ad096461..3f6310da5 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -3,6 +3,7 @@ from __future__ import annotations import math +import warnings from typing import Any import numpy as np @@ -21,6 +22,11 @@ def split( ) -> tuple[Fold, Fold, Fold]: """Split a KG sample dataset into three disjoint folds. + .. deprecated:: + Use ``set_task(KGLinkPrediction(...), split=PatientSplit(...))``, + which splits triples before fitting the processors. This function + will be removed in the next release. + 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 @@ -76,6 +82,12 @@ def split( ... ValueError: ratios must sum to 1.0, got 0.9 """ + warnings.warn( + "kg_emb.datasets.split is deprecated and will be removed in the next " + "release; use set_task(KGLinkPrediction(...), split=PatientSplit(...)).", + DeprecationWarning, + stacklevel=2, + ) 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) 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 118e7c711..cdb3e8174 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -78,6 +78,14 @@ def __init__( self.r_num = processor.num_relations else: # A PyHealth 1.x SampleKGDataset carries the counts itself. + warnings.warn( + "Building a KGE model from a dataset without a kg_triple " + "processor (e.g. a 1.x SampleKGDataset) is deprecated and will " + "be removed in the next release; pass the training part of " + "set_task(KGLinkPrediction(...), split=PatientSplit(...)).", + DeprecationWarning, + stacklevel=3, + ) self.triple_processor = None self.e_num = dataset.entity_num self.r_num = dataset.relation_num diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py index 69bc90477..6d1e40e21 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py @@ -3,6 +3,7 @@ from tqdm import tqdm from typing import Tuple, List from collections import defaultdict +import warnings def link_prediction_fn( triples: List[Tuple] @@ -10,6 +11,11 @@ def link_prediction_fn( """Process a triple list for the link prediction task + .. deprecated:: + Use the :class:`KGLinkPrediction` task with + ``BaseKGDataset.set_task``. This function will be removed in the + next release. + Link prediction is a task to either Tail Prediction: predict tail entity t given a triple query (h, r, ?), or Head Prediction: predict head entity h given a triple query (?, r, t) @@ -19,7 +25,21 @@ def link_prediction_fn( Returns: samples: a list of samples + + Examples: + >>> import warnings + >>> with warnings.catch_warnings(): + ... warnings.simplefilter("ignore", DeprecationWarning) + ... samples = link_prediction_fn([(0, 0, 1), (0, 0, 2)]) + >>> samples[0]["ground_truth_tail"] + [1, 2] """ + warnings.warn( + "link_prediction_fn is deprecated and will be removed in the next " + "release; use the KGLinkPrediction task with BaseKGDataset.set_task.", + DeprecationWarning, + stacklevel=2, + ) samples = [] count = count_frequency(triples) triple_set = set(triples) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 255ec0a8e..ae05a383e 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -1036,3 +1036,39 @@ def test_held_out_heads_are_not_filtered(self) -> None: heads, _ = model.train_neg_sample_gen(gt_head, [[]], 256) self.assertIn(e["b"], heads.flatten().tolist()) self.assertNotIn(e["a"], heads.flatten().tolist()) + + +class TestDeprecatedEntryPoints(unittest.TestCase): + """The 1.x entry points keep their import paths and work, with a warning.""" + + def test_old_import_paths_resolve(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb import datasets, tasks + + for name in ("BaseKGDataset", "SampleKGDataset", "UMLSDataset", "split"): + self.assertTrue(hasattr(datasets, name), name) + self.assertTrue(callable(tasks.link_prediction_fn)) + + def test_link_prediction_fn_warns(self) -> None: + with self.assertWarns(DeprecationWarning): + samples = link_prediction_fn([(0, 0, 1)]) + self.assertEqual(len(samples), 1) + + def test_sample_kg_dataset_warns(self) -> None: + with self.assertWarns(DeprecationWarning): + make_dataset(n=2) + + def test_split_warns(self) -> None: + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + dataset = make_dataset(n=10) + with self.assertWarns(DeprecationWarning): + split(dataset, [0.6, 0.2, 0.2], seed=0) + + def test_a_model_on_a_1x_dataset_warns(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + dataset = make_dataset(n=4) + with self.assertWarnsRegex(DeprecationWarning, "kg_triple"): + TransE(dataset=dataset, e_dim=4, r_dim=4) From a46c3ffafc039cb3244915cee11b041725400189 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 22:52:56 +0200 Subject: [PATCH 23/25] docs(kg_emb): document the 2.0 pipeline and update the examples - docs/api/medcode.rst describes BaseKGDataset / UMLSDataset, the KGLinkPrediction task with split=PatientSplit(...), the kg_triple processor that keeps training negatives free of held-out triples, and the deprecated 1.x entry points, including the change of id order. - examples/kg_emb_sample_dataset.py, added earlier in this PR for the in-memory SampleKGDataset, becomes examples/kg_emb_link_prediction.py: a self-contained run of the 2.0 pipeline with the Trainer on a small synthetic graph whose structure a translation model can learn, so the filtered test metrics end well above those of a random ranking. - kg_emb/examples/train_kge_model.py trains on UMLS with the new API. It no longer loads a 1.x checkpoint, whose rows follow the old id order. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 64 +++-- examples/kg_emb_link_prediction.py | 132 ++++++++++ examples/kg_emb_sample_dataset.py | 62 ----- .../kg_emb/examples/train_kge_model.py | 227 ++++++++---------- 4 files changed, 276 insertions(+), 209 deletions(-) create mode 100644 examples/kg_emb_link_prediction.py delete mode 100644 examples/kg_emb_sample_dataset.py diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index b271cc79d..2b07d0bf5 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -121,23 +121,55 @@ Knowledge graph embeddings -------------------------- ``pyhealth.medcode.pretrained_embeddings.kg_emb`` trains TransE, RotatE, -DistMult and ComplEx on an in-memory list of triples. +DistMult and ComplEx on a knowledge graph with the standard streaming +pipeline: + +1. :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset` + is a :class:`~pyhealth.datasets.BaseDataset` with one ``triples`` table + (``head``, ``relation``, ``tail``) declared in YAML with + ``patient_id: null``, so each triple is one record. + :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.UMLSDataset` + is its bundled UMLS configuration. Entity and relation ids are assigned on + the full graph, by sorted name. +2. ``set_task(KGLinkPrediction(...), split=PatientSplit(...))`` splits the + triples and fits the processors on the training part only, returning one + streaming :class:`~pyhealth.datasets.SampleDataset` per part. Each sample + holds a ``triple`` and its ``ground_truth_head`` / ``ground_truth_tail`` + lists, which cover the whole graph and filter the ranking at evaluation + (filtered setting of Bordes et al., 2013). Duplicate triples are removed + and counted. +3. The ``kg_triple`` processor + (:class:`~pyhealth.processors.KGTripleProcessor`), fitted on the training + triples, gives the models the numbers of entities and relations and the + dicts that keep known training positives out of the training negatives, + as in the reference implementation of Sun et al. (2019). Validation and + test triples therefore never shape training. ``set_task`` warns when + called without ``split``. +4. The models take ``negative_sampling`` as an argument and follow + ``model.train()`` / ``model.eval()``, as set by + :class:`~pyhealth.trainer.Trainer`; use the ``hits@n`` and ``mean_rank`` + metrics. + +See ``examples/kg_emb_link_prediction.py`` for a self-contained walk-through. + +The PyHealth 1.x entry points, ``link_prediction_fn``, :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset` -is an :class:`~pyhealth.datasets.InMemorySampleDataset`: samples are processed -once, at construction, and kept in memory. -:func:`pyhealth.datasets.get_dataloader` accepts it. Its ``ground_truth_head`` -and ``ground_truth_tail`` fields are padded by -:class:`~pyhealth.processors.kg_processor.KGProcessor` and returned with a -mask. - -:func:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split` returns -plain lists of sample dictionaries that carry the ``train`` and -``hyperparameters`` fields the models read. Wrap a fold in -:class:`torch.utils.data.DataLoader` with -:func:`pyhealth.datasets.collate_fn_dict_with_padding`: ``get_dataloader`` -calls ``set_shuffle()``, which a list does not have. - -See ``examples/kg_emb_sample_dataset.py`` for a self-contained walk-through. +and :func:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split`, +still work for one release with a ``DeprecationWarning``. Ids no longer +follow the 1.x order of first appearance: map 1.x embeddings through the +``id2entity`` saved with them. + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset + :members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.UMLSDataset + :members: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction + :members: + :show-inheritance: .. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset :members: diff --git a/examples/kg_emb_link_prediction.py b/examples/kg_emb_link_prediction.py new file mode 100644 index 000000000..0788f026f --- /dev/null +++ b/examples/kg_emb_link_prediction.py @@ -0,0 +1,132 @@ +"""Train and evaluate TransE on a small knowledge graph with the 2.0 pipeline. + +The example is self-contained: it writes a small synthetic graph with a +learnable structure to a temporary directory instead of downloading UMLS. +Its filtered test metrics should end well above those of a random ranking +(mean rank about 30 for 60 entities); they say nothing about real graphs. +Steps: + +1. :class:`BaseKGDataset` reads one ``triples`` table (``head``, + ``relation``, ``tail``) declared in a YAML config; each triple is one + record. :class:`UMLSDataset` does the same with a bundled config. +2. ``set_task(KGLinkPrediction(...), split=PatientSplit(...))`` splits the + triples and fits the processors on the training part only. The + ``kg_triple`` processor keeps the training-graph dicts that filter the + training negatives; the ground-truth lists of each sample cover the + whole graph and filter the ranking at evaluation (filtered setting of + Bordes et al., 2013). +3. The model reads the numbers of entities and relations from that + processor; the :class:`~pyhealth.trainer.Trainer` switches between + training and filtered evaluation through ``model.train()`` / + ``model.eval()``. + +Run from the repository root:: + + python examples/kg_emb_link_prediction.py +""" + +import tempfile +from pathlib import Path + +import numpy as np +import torch + +from pyhealth.datasets import PatientSplit, get_dataloader +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import BaseKGDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE +from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction +from pyhealth.trainer import Trainer + +SEED = 0 +N_ENTITIES = 60 +N_CLUSTERS = 6 +N_RELATIONS = 2 +N_TRIPLES = 600 +SPLIT_RATIOS = (0.8, 0.1, 0.1) +EPOCHS = 30 +BATCH_SIZE = 32 + +CONFIG = """version: "1.0" +tables: + triples: + file_path: "kg.tsv" + patient_id: null + timestamp: null + attributes: [head, relation, tail] +""" + + +def write_synthetic_graph(root: Path) -> Path: + """Writes a learnable synthetic graph and its config; returns the config. + + Entities fall into ``N_CLUSTERS`` ordered clusters, and relation ``rk`` + links a random entity of cluster ``c`` to a random entity of cluster + ``c + k``. A translation model can represent this, so held-out triples + are predictable from the training ones; on uniformly random triples the + metrics would stay at chance level whatever the pipeline does. + """ + rng = np.random.default_rng(SEED) + size = N_ENTITIES // N_CLUSTERS + lines = ["head\trelation\ttail"] + for _ in range(N_TRIPLES): + k = int(rng.integers(1, N_RELATIONS + 1)) + c = int(rng.integers(0, N_CLUSTERS - k)) + head = c * size + int(rng.integers(size)) + tail = (c + k) * size + int(rng.integers(size)) + lines.append(f"e{head}\tr{k}\te{tail}") + (root / "kg.tsv").write_text("\n".join(lines) + "\n") + (root / "kg.yaml").write_text(CONFIG) + return root / "kg.yaml" + + +def main() -> None: + # Training negatives are drawn with NumPy's global generator. + np.random.seed(SEED) + torch.manual_seed(SEED) + + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_synthetic_graph(root) + dataset = BaseKGDataset( + root=str(root), config_path=config, cache_dir=root / "cache" + ) + print(f"{dataset.num_entities} entities, {dataset.num_relations} relations") + + task = KGLinkPrediction( + num_entities=dataset.num_entities, num_relations=dataset.num_relations + ) + train, val, test = dataset.set_task( + task, split=PatientSplit(ratios=SPLIT_RATIOS, seed=SEED) + ) + print(f"train/val/test: {len(train)}/{len(val)}/{len(test)} triples") + + model = TransE(dataset=train, e_dim=32, r_dim=32, negative_sampling=16) + trainer = Trainer( + model=model, + metrics=["hits@n", "mean_rank"], + device="cpu", + output_path=str(root / "output"), + exp_name="kg_emb_link_prediction", + ) + trainer.train( + train_dataloader=get_dataloader(train, batch_size=BATCH_SIZE, shuffle=True), + val_dataloader=get_dataloader(val, batch_size=BATCH_SIZE), + epochs=EPOCHS, + optimizer_params={"lr": 1e-2}, + monitor="mean_reciprocal_rank", + monitor_criterion="max", + ) + scores = trainer.evaluate(get_dataloader(test, batch_size=BATCH_SIZE)) + print("filtered test metrics:", scores) + + # Entity names from ids: the most likely tails of (e0, r1, ?), which + # should come from the second cluster, e10..e19. + model.cpu() + top = model.inference( + head=dataset.entity2id["e0"], relation=dataset.relation2id["r1"], top_k=5 + ) + print("(e0, r1, ?):", [dataset.id2entity[i] for i in top]) + + +if __name__ == "__main__": + main() diff --git a/examples/kg_emb_sample_dataset.py b/examples/kg_emb_sample_dataset.py deleted file mode 100644 index c534864a5..000000000 --- a/examples/kg_emb_sample_dataset.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Train a TransE model on a synthetic knowledge-graph sample dataset. - -This example does not download UMLS. It shows the 2.0 path: - -1. Build a :class:`SampleKGDataset`, an ``InMemorySampleDataset`` whose - ``ground_truth_head`` / ``ground_truth_tail`` fields are padded by - ``KGProcessor`` and returned with a mask. -2. Split into train/val/test folds with :func:`split`. The folds are plain - lists of sample dicts carrying the ``train`` and ``hyperparameters`` - fields the model reads. -3. Wrap the train fold in ``torch.utils.data.DataLoader`` using - :func:`collate_fn_dict_with_padding`. ``get_dataloader`` calls - ``set_shuffle()``, which a list does not have. -""" - -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") 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 fe8f40c9e..6fec4da87 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,149 +1,114 @@ -from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import UMLSDataset, split -from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import link_prediction_fn -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 -import torch -import json -import pickle +"""Train a KG embedding model on UMLS with the PyHealth 2.0 pipeline. +``UMLS_ROOT`` must hold ``graph.txt`` (head, relation and tail separated by +tabs, no header), e.g. downloaded from +https://storage.googleapis.com/pyhealth/umls/graph.txt. Outputs go to +``OUTPUT_DIR``. Ids follow sorted entity names, not the 1.x order, so +checkpoints trained with PyHealth 1.x must not be loaded as they are. """ -This is an example to show you how to train a KG embedding model using our package -""" +import json +import pickle +from pathlib import Path +import numpy as np +import torch -umls_ds = UMLSDataset( - root="/data/pj20/umls/", - # root="https://storage.googleapis.com/pyhealth/umls/", - dev=False, - refresh_cache=True -) - -# check the dataset statistics before setting task -print(umls_ds.stat()) - -# check the relation numbers in the dataset -print("Relations in KG:", umls_ds.relation2id) - -umls_ds = umls_ds.set_task(link_prediction_fn, negative_sampling=64, save=False) - -# save the id2entity, id2relation -with open("/data/pj20/umls_kge/pretrained_model/umls_transe_new/id2entity.json", "w") as f: - json.dump(umls_ds.id2entity, f, indent=6) - -with open("/data/pj20/umls_kge/pretrained_model/umls_transe_new/id2relation.json", "w") as f: - json.dump(umls_ds.id2relation, f, indent=6) - -# check the dataset statistics after setting task -print(umls_ds.stat()) - -# split the dataset and get the dataloaders -train_dataset, val_dataset, test_dataset = split(umls_ds, [0.9, 0.05, 0.05]) -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 -model = TransE( - dataset=umls_ds, - e_dim=512, - r_dim=512, -) - -print('Loaded model: ', model) -state_dict = torch.load("/data/pj20/umls_kge/pretrained_model/umls_transe_new/1_250000_last.ckpt") -model.load_state_dict(state_dict) - -# initialize a trainer and start training -trainer = Trainer( - model=model, - device='cuda:5', - metrics=['hits@n', 'mean_rank'], - output_path='/data/pj20/umls_kge/pretrained_model', - exp_name='umls_transe_new' - ) +from pyhealth.datasets import PatientSplit, get_dataloader +from pyhealth.medcode import InnerMap +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import UMLSDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE +from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction +from pyhealth.trainer import Trainer -trainer.train( - train_dataloader=train_loader, - # val_dataloader=val_loader, - epochs=10, - # steps_per_epoch=100, - # evaluation_steps=1, - optimizer_params={'lr': 1e-3}, - monitor='mean_rank', - monitor_criterion='min' -) - -# save the entity embedding and relation embedding -with open("/data/pj20/umls_kge/pretrained_model/umls_transe_new/entity_embedding.pkl", "wb") as f: - pickle.dump(model.E_emb, f) - -with open("/data/pj20/umls_kge/pretrained_model/umls_transe_new/relation_embedding.pkl", "wb") as f: - pickle.dump(model.R_emb, f) - -# evaluate the trained model -# trainer.evaluate(test_loader) - -# use the trained model to handle head/tail entity prediction, and use the CodeMap to map the code to free text -umls_code_map = InnerMap.load("UMLS", refresh_cache=True) - -head = 'C0000039' -relation = "PAR" - -model.to('cpu') -result_eid = model.inference( - head=umls_ds.entity2id[head], - relation=umls_ds.relation2id[relation], - tail=None, - top_k=3 - ) +UMLS_ROOT = "path/to/umls" +OUTPUT_DIR = Path("output/umls_transe") +SEED = 0 +SPLIT_RATIOS = (0.9, 0.05, 0.05) +NUM_WORKERS = 4 -print(f"Input Head: {head} - {umls_code_map.lookup(head)}") -print(f"Input Relation: {relation}") -print("Tail Prediction:") -for idx, eid in enumerate(result_eid): - tail = umls_ds.id2entity[eid] - print(f"{idx}: {tail} - {umls_code_map.lookup(tail)}") +def _main() -> None: + np.random.seed(SEED) # training negatives use NumPy's global generator + torch.manual_seed(SEED) + OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + umls_ds = UMLSDataset(root=UMLS_ROOT, num_workers=NUM_WORKERS) + umls_ds.stats() + print("Relations in KG:", umls_ds.relation2id) -tail = 'C5162542' -relation = "CHD" + # Ids and names, to read the embeddings later. + with open(OUTPUT_DIR / "id2entity.json", "w") as f: + json.dump(umls_ds.id2entity, f, indent=2) + with open(OUTPUT_DIR / "id2relation.json", "w") as f: + json.dump(umls_ds.id2relation, f, indent=2) -result_eid = model.inference( - head=None, - relation=umls_ds.relation2id[relation], - tail=umls_ds.entity2id[tail], - top_k=3 + # Triple-level split; the processors, including the training-graph dicts + # that filter training negatives, are fitted on the training part only. + task = KGLinkPrediction( + num_entities=umls_ds.num_entities, num_relations=umls_ds.num_relations + ) + train, val, test = umls_ds.set_task( + task, split=PatientSplit(ratios=SPLIT_RATIOS, seed=SEED) ) -print(f"Input Tail: {tail} - {umls_code_map.lookup(tail)}") -print(f"Input Relation: {relation}") - -print("Head Prediction:") -for idx, eid in enumerate(result_eid): - head = umls_ds.id2entity[eid] - print(f"{idx}: {head} - {umls_code_map.lookup(head)}") + model = TransE(dataset=train, e_dim=512, r_dim=512, negative_sampling=64) + trainer = Trainer( + model=model, + metrics=["hits@n", "mean_rank"], + output_path=str(OUTPUT_DIR), + exp_name="umls_transe", + ) + trainer.train( + train_dataloader=get_dataloader(train, batch_size=8, shuffle=True), + val_dataloader=get_dataloader(val, batch_size=8), + epochs=10, + optimizer_params={"lr": 1e-3}, + monitor="mean_rank", + monitor_criterion="min", + ) + print("filtered test metrics:", trainer.evaluate(get_dataloader(test, batch_size=8))) + with open(OUTPUT_DIR / "entity_embedding.pkl", "wb") as f: + pickle.dump(model.E_emb, f) + with open(OUTPUT_DIR / "relation_embedding.pkl", "wb") as f: + pickle.dump(model.R_emb, f) -head = 'C0000039' -tail = 'C0000039' -relation = "SY" + # Head/tail prediction, with the CodeMap mapping codes to free text. + umls_code_map = InnerMap.load("UMLS") + model.to("cpu") -classification_score = model.inference( - head=umls_ds.entity2id[head], - relation=umls_ds.relation2id[relation], - tail=umls_ds.entity2id[tail], + head, relation = "C0000039", "PAR" + result_eid = model.inference( + head=umls_ds.entity2id[head], relation=umls_ds.relation2id[relation], top_k=3 + ) + print(f"Input Head: {head} - {umls_code_map.lookup(head)}") + print(f"Input Relation: {relation}") + print("Tail Prediction:") + for idx, eid in enumerate(result_eid): + tail = umls_ds.id2entity[eid] + print(f"{idx}: {tail} - {umls_code_map.lookup(tail)}") + + tail, relation = "C5162542", "CHD" + result_eid = model.inference( + relation=umls_ds.relation2id[relation], tail=umls_ds.entity2id[tail], top_k=3 ) + print(f"Input Tail: {tail} - {umls_code_map.lookup(tail)}") + print(f"Input Relation: {relation}") + print("Head Prediction:") + for idx, eid in enumerate(result_eid): + head = umls_ds.id2entity[eid] + print(f"{idx}: {head} - {umls_code_map.lookup(head)}") + + head = tail = "C0000039" + relation = "SY" + score = model.inference( + head=umls_ds.entity2id[head], + relation=umls_ds.relation2id[relation], + tail=umls_ds.entity2id[tail], + ) + print(f"Classification Score: {score}") + -print(f"Classification Score: {classification_score}") +if __name__ == "__main__": + _main() From e26916fef208a974309c88c6bd0ebe8b87c699ef Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:06:51 +0200 Subject: [PATCH 24/25] fix(kg_emb): honour 1.x per-sample negative_sampling and NumPy integers Two regressions of the model refactor, for 1.x callers: - 1.x samples may carry hyperparameters["negative_sampling"], which the models read before. The model's value now only defaults to it when neither the model argument nor the dataset chose one (as in the models' __main__ demos); when both disagree, the model's value is used and a warning says so. Before this fix the batch value was ignored and 128 negatives were drawn. - negative_sampling accepts any integral number, e.g. np.int64 from a hyperparameter grid, instead of int only. Tests now check that the chosen number of negatives reaches the sampler. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/models/kg_base.py | 47 +++++++++-- tests/core/test_kg_emb.py | 77 +++++++++++++++++++ 2 files changed, 118 insertions(+), 6 deletions(-) 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 cdb3e8174..332f6ffa8 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,3 +1,4 @@ +import numbers import warnings from abc import ABC @@ -37,7 +38,8 @@ class KGEBaseModel(ABC, nn.Module): mode: evaluation metric type, one of "binary", "multiclass", or "multilabel", "multiclass" by default negative_sampling: Number of negative heads and of negative tails drawn per training triple. Defaults to the ``negative_sampling`` of a - 1.x ``SampleKGDataset`` if it has one, else 128. + 1.x ``SampleKGDataset`` if it has one, else to the value 1.x + samples carry in ``hyperparameters``, else 128. Examples: >>> from pyhealth.processors import KGTripleProcessor @@ -89,14 +91,23 @@ def __init__( self.triple_processor = None self.e_num = dataset.entity_num self.r_num = dataset.relation_num + # Whether the value was chosen (argument or 1.x dataset) rather than + # defaulted: a 1.x batch's own value only overrides a default. + legacy = getattr(dataset, "task_spec_param", None) or {} if negative_sampling is None: - legacy = getattr(dataset, "task_spec_param", None) or {} - negative_sampling = legacy.get("negative_sampling", DEFAULT_NEGATIVE_SAMPLING) - if not isinstance(negative_sampling, int) or negative_sampling < 1: + negative_sampling = legacy.get("negative_sampling") + self._negative_sampling_chosen = negative_sampling is not None + if negative_sampling is None: + negative_sampling = DEFAULT_NEGATIVE_SAMPLING + if ( + not isinstance(negative_sampling, numbers.Integral) + or isinstance(negative_sampling, bool) + or negative_sampling < 1 + ): raise ValueError( f"negative_sampling must be a positive integer, got {negative_sampling!r}." ) - self.negative_sampling = negative_sampling + self.negative_sampling = int(negative_sampling) self.e_dim = e_dim self.r_dim = r_dim self.ns = ns @@ -372,6 +383,30 @@ def calc(self, head, relation, tail, mode='pos'): raise NotImplementedError + def _batch_negative_sampling(self, data) -> int: + """Negatives per triple for this batch. + + PyHealth 1.x samples may carry ``hyperparameters["negative_sampling"]``, + which the models used to read. It still applies when the model's + value was only the default, and is ignored, with a warning, when the + model's value was chosen. + """ + hyperparameters = data.get("hyperparameters") + if not hyperparameters or not isinstance(hyperparameters[0], dict): + return self.negative_sampling + batch_value = hyperparameters[0].get("negative_sampling") + if batch_value is None or batch_value == self.negative_sampling: + return self.negative_sampling + if not self._negative_sampling_chosen: + return int(batch_value) + warnings.warn( + f"The samples' negative_sampling={batch_value} is ignored; the " + f"model's negative_sampling={self.negative_sampling} is used.", + UserWarning, + stacklevel=3, + ) + return self.negative_sampling + def _training_filters(self, data, positive_sample): """The entities kept out of each triple's training negatives. @@ -427,7 +462,7 @@ def forward(self, **data): negative_sample_head, negative_sample_tail = self.train_neg_sample_gen( gt_head=gt_head, gt_tail=gt_tail, - negative_sampling=self.negative_sampling, + negative_sampling=self._batch_negative_sampling(data), ) negative_sample_head, negative_sample_tail = negative_sample_head.to(self.device), negative_sample_tail.to(self.device) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index ae05a383e..c0e3cf51b 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -1072,3 +1072,80 @@ def test_a_model_on_a_1x_dataset_warns(self) -> None: dataset = make_dataset(n=4) with self.assertWarnsRegex(DeprecationWarning, "kg_triple"): TransE(dataset=dataset, e_dim=4, r_dim=4) + + +class TestNegativeSamplingReachesTheSampler(unittest.TestCase): + """The number of negatives drawn is the one the caller chose.""" + + def spy(self, model): + shapes = [] + original = model.train_neg_sample_gen + + def record(**kwargs): + out = original(**kwargs) + shapes.append(tuple(out[0].shape)) + return out + + model.train_neg_sample_gen = record + return shapes + + def legacy_batch(self, per_sample=None, **dataset_kwargs): + samples = make_samples(4) + if per_sample is not None: + for s in samples: + s["hyperparameters"] = {"negative_sampling": per_sample} + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + dataset = SampleKGDataset( + samples=samples, entity_num=10, relation_num=3, **dataset_kwargs + ) + loader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn_dict_with_padding) + return dataset, next(iter(loader)) + + def build(self, dataset, **kwargs): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + model = TransE(dataset=dataset, e_dim=4, r_dim=4, **kwargs) + model.train() + return model + + def test_the_model_argument_sets_the_sample_count(self) -> None: + from pyhealth.processors import KGTripleProcessor + + processor = KGTripleProcessor(num_entities=10, num_relations=3) + processor.fit([{"triple": s["triple"]} for s in make_samples(4)], "triple") + + class _Train: + input_processors = {"triple": processor} + + model = self.build(_Train(), negative_sampling=7) + shapes = self.spy(model) + model(triple=torch.tensor([[0, 0, 4], [1, 1, 5]])) + self.assertEqual(shapes, [(2, 7)]) + + def test_numpy_integers_are_accepted(self) -> None: + import numpy as np + + dataset, _ = self.legacy_batch() + self.assertEqual(self.build(dataset, negative_sampling=np.int64(6)).negative_sampling, 6) + dataset, _ = self.legacy_batch(negative_sampling=np.int32(5)) + self.assertEqual(self.build(dataset).negative_sampling, 5) + with self.assertRaises(ValueError): + self.build(dataset, negative_sampling=True) + + def test_a_1x_per_sample_value_applies_when_none_was_chosen(self) -> None: + dataset, batch = self.legacy_batch(per_sample=3) + model = self.build(dataset) + shapes = self.spy(model) + model(**batch) + self.assertEqual(shapes, [(2, 3)]) + + def test_a_chosen_value_wins_over_the_samples_with_a_warning(self) -> None: + dataset, batch = self.legacy_batch(per_sample=3) + model = self.build(dataset, negative_sampling=5) + shapes = self.spy(model) + with self.assertWarnsRegex(UserWarning, "negative_sampling=3 is ignored"): + model(**batch) + self.assertEqual(shapes, [(2, 5)]) From 50e3a8e86a8acdf03fd2394916df611db0abc5ea Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:33:20 +0200 Subject: [PATCH 25/25] fix(kg_emb): refuse base datasets in the models and tighten the leak tests - A KGE model built from a BaseKGDataset itself, instead of a set_task part, took the 1.x fallback (the dataset has the 1.x count attributes) and filtered training negatives with the batches' full-graph ground truths, i.e. with held-out triples. Base datasets are now refused with a message pointing to the training part. - The leak regression tests go through forward on both sides and check that every training positive of a pair stays filtered while held-out positives are drawable; restoring the full-graph filter on either side, or keeping only the triple's own entity, now fails them. - A legacy set_task(link_prediction_fn) call emits one DeprecationWarning instead of three, two of which pointed inside pyhealth. - Deprecation notes use a warning admonition: Sphinx's deprecated directive needs a version number, which is not known yet. - The docs say that the deprecated 1.x path keeps filtering training negatives with every triple, and the example no longer claims that all top tails of (e0, r1, ?) come from the second cluster. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 9 +- examples/kg_emb_link_prediction.py | 2 +- .../kg_emb/datasets/base_kg_dataset.py | 24 +++--- .../kg_emb/datasets/sample_kg_dataset.py | 3 +- .../kg_emb/datasets/splitter.py | 3 +- .../kg_emb/models/kg_base.py | 12 ++- .../kg_emb/tasks/link_prediction.py | 3 +- tests/core/test_kg_emb.py | 83 ++++++++++++------- 8 files changed, 92 insertions(+), 47 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 2b07d0bf5..cb7579b30 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -155,9 +155,12 @@ See ``examples/kg_emb_link_prediction.py`` for a self-contained walk-through. The PyHealth 1.x entry points, ``link_prediction_fn``, :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset` and :func:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split`, -still work for one release with a ``DeprecationWarning``. Ids no longer -follow the 1.x order of first appearance: map 1.x embeddings through the -``id2entity`` saved with them. +still work for one release with a ``DeprecationWarning``. They keep their +1.x behaviour, including training negatives filtered with every triple of +the graph, held-out ones included; only the 2.0 path above avoids that. A +model must be built from a sample dataset part, not from the +``BaseKGDataset`` itself. Ids no longer follow the 1.x order of first +appearance: map 1.x embeddings through the ``id2entity`` saved with them. .. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset :members: diff --git a/examples/kg_emb_link_prediction.py b/examples/kg_emb_link_prediction.py index 0788f026f..a9283e35e 100644 --- a/examples/kg_emb_link_prediction.py +++ b/examples/kg_emb_link_prediction.py @@ -120,7 +120,7 @@ def main() -> None: print("filtered test metrics:", scores) # Entity names from ids: the most likely tails of (e0, r1, ?), which - # should come from the second cluster, e10..e19. + # should mostly come from the second cluster, e10..e19. model.cpu() top = model.inference( head=dataset.entity2id["e0"], relation=dataset.relation2id["r1"], top_k=5 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 95e70a442..caf596759 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 @@ -513,13 +513,17 @@ def _set_legacy_task( ) if save is not None: logger.info("set_task(save=...) is ignored: samples are not pickled.") - samples = task_fn(self._indexed_triples()) - return SampleKGDataset( - samples=samples, - dataset_name=self.dataset_name, - task_name=task_name or task_fn.__name__, - dev=self.dev, - entity2id=self.entity2id, - relation2id=self.relation2id, - **task_spec_param, - ) + # One deprecation warning per call, the one above: link_prediction_fn + # and SampleKGDataset would each add their own. + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + samples = task_fn(self._indexed_triples()) + return SampleKGDataset( + samples=samples, + dataset_name=self.dataset_name, + task_name=task_name or task_fn.__name__, + dev=self.dev, + entity2id=self.entity2id, + relation2id=self.relation2id, + **task_spec_param, + ) 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 68d70962e..0970b234b 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 @@ -34,7 +34,8 @@ class SampleKGDataset(InMemorySampleDataset): r"""In-memory dataset of knowledge-graph link-prediction samples. - .. deprecated:: + .. warning:: + Deprecated. Use :class:`BaseKGDataset` with ``set_task(KGLinkPrediction(...), split=PatientSplit(...))``, which returns streaming :class:`~pyhealth.datasets.SampleDataset` parts. This class will be diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 3f6310da5..84f7f2806 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -22,7 +22,8 @@ def split( ) -> tuple[Fold, Fold, Fold]: """Split a KG sample dataset into three disjoint folds. - .. deprecated:: + .. warning:: + Deprecated. Use ``set_task(KGLinkPrediction(...), split=PatientSplit(...))``, which splits triples before fitting the processors. This function will be removed in the next release. 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 332f6ffa8..4f6fcd146 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -7,7 +7,7 @@ import torch.nn.functional as F from torch import nn -from pyhealth.datasets import SampleDataset +from pyhealth.datasets import BaseDataset, SampleDataset from pyhealth.processors import KGTripleProcessor # Negatives per positive triple when the caller does not choose: the @@ -73,6 +73,16 @@ def __init__( negative_sampling: int | None = None, ): super().__init__() + if isinstance(dataset, BaseDataset): + # A BaseKGDataset has the 1.x count attributes but no fitted + # processor: accepting it would filter training negatives with + # the batches' full-graph ground truths, i.e. with held-out + # triples. + raise TypeError( + f"{type(dataset).__name__} is a base dataset; pass the training " + "part returned by set_task(KGLinkPrediction(...), " + "split=PatientSplit(...)) instead." + ) processor = getattr(dataset, "input_processors", {}).get("triple") if isinstance(processor, KGTripleProcessor): self.triple_processor: KGTripleProcessor | None = processor diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py index 6d1e40e21..8071fd8b3 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py @@ -11,7 +11,8 @@ def link_prediction_fn( """Process a triple list for the link prediction task - .. deprecated:: + .. warning:: + Deprecated. Use the :class:`KGLinkPrediction` task with ``BaseKGDataset.set_task``. This function will be removed in the next release. diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index c0e3cf51b..b3475cb2a 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -291,6 +291,13 @@ def test_legacy_function_task_still_works_with_a_warning(self) -> None: self.assertEqual(sample_ds.task_spec_param, {"negative_sampling": 4}) self.assertEqual(sample_ds.entity_num, 4) + def test_legacy_task_warns_once_per_call(self) -> None: + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + self.dataset.set_task(link_prediction_fn) + ours = [w for w in caught if "will be removed in the next release" in str(w.message)] + self.assertEqual(len(ours), 1, [str(w.message) for w in ours]) + def test_legacy_task_fn_keyword_is_accepted(self) -> None: with self.assertWarns(DeprecationWarning): sample_ds = self.dataset.set_task(task_fn=link_prediction_fn) @@ -961,51 +968,60 @@ def test_a_disagreeing_1x_train_flag_warns(self) -> None: self.assertIn("y_prob", out) # model.eval() decides, not the flag +def record_negatives(model) -> dict: + """Wraps train_neg_sample_gen to keep what forward draws.""" + drawn = {} + original = model.train_neg_sample_gen + + def record(**kwargs): + drawn["head"], drawn["tail"] = original(**kwargs) + return drawn["head"], drawn["tail"] + + model.train_neg_sample_gen = record + return drawn + + class TestTrainingNegativesUseTrainingTriplesOnly(unittest.TestCase): """Regression test for the leak: training negatives were filtered with ground truths of the whole graph, so held-out positives were never drawn.""" - def test_a_tail_true_only_in_test_can_be_a_training_negative(self) -> None: + def test_held_out_positives_are_drawable_and_training_ones_are_not(self) -> None: import numpy as np from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE from pyhealth.processors import KGTripleProcessor - # Training triple (0, 0, 1); (0, 0, 2) is a test triple. - processor = KGTripleProcessor(num_entities=3, num_relations=1) - processor.fit([{"triple": (0, 0, 1)}], "triple") + # Training: (0, 0, 1), (0, 0, 3), (4, 0, 1). Test: (0, 0, 2), (5, 0, 1). + processor = KGTripleProcessor(num_entities=6, num_relations=1) + processor.fit( + [{"triple": t} for t in [(0, 0, 1), (0, 0, 3), (4, 0, 1)]], "triple" + ) class _Train: input_processors = {"triple": processor} - model = TransE(dataset=_Train(), e_dim=4, r_dim=4, negative_sampling=64) + model = TransE(dataset=_Train(), e_dim=4, r_dim=4, negative_sampling=256) model.train() - # The sample's ground truth covers the whole graph, as KGLinkPrediction gives it. + # Ground truths of the whole graph, as KGLinkPrediction gives them. batch = { "triple": torch.tensor([[0, 0, 1]]), - "ground_truth_head": [[0]], - "ground_truth_tail": [[1, 2]], + "ground_truth_head": [[0, 4, 5]], + "ground_truth_tail": [[1, 2, 3]], } - drawn = {} - original = model.train_neg_sample_gen - - def record(**kwargs): - drawn["head"], drawn["tail"] = original(**kwargs) - return drawn["head"], drawn["tail"] - - model.train_neg_sample_gen = record + drawn = record_negatives(model) np.random.seed(0) model(**batch)["loss"].backward() - tails = set(drawn["tail"].flatten().tolist()) - self.assertNotIn(1, tails) # the training positive stays filtered - self.assertIn(2, tails) # the test positive is a legitimate negative - self.assertNotIn(0, set(drawn["head"].flatten().tolist())) + # Every training positive of the pair stays out (reference + # semantics); the test-only positives 2 and 5 are drawable. + self.assertEqual(set(drawn["tail"].flatten().tolist()), {0, 2, 4, 5}) + self.assertEqual(set(drawn["head"].flatten().tolist()), {1, 2, 3, 5}) class TestTrainingNegativesOnARealSplit(unittest.TestCase): - """The same property on set_task output: (b, r2, d) is a test triple and - (a, r2, d) a training one, so b must stay drawable as a head negative.""" + """The same property on set_task output, through forward: (b, r2, d) is a + test triple and (a, r2, d) a training one, so b must be drawable as a + head negative of (a, r2, d).""" # Same graph and split as TestModelsOnSampleDatasets, without its tests. setUpClass = classmethod(TestModelsOnSampleDatasets.setUpClass.__func__) @@ -1025,17 +1041,26 @@ def test_held_out_heads_are_not_filtered(self) -> None: ) test_triples = {tuple(self.test[i]["triple"].tolist()) for i in range(len(self.test))} self.assertIn((e["b"], r["r2"], e["d"]), test_triples) - self.assertEqual(sorted(unpad({k: v[0] for k, v in batch["ground_truth_head"].items()})), - [e["a"], e["b"]]) + self.assertEqual( + sorted(unpad({k: v[0] for k, v in batch["ground_truth_head"].items()})), + [e["a"], e["b"]], + ) model = TransE(dataset=self.train, e_dim=4, r_dim=4, negative_sampling=256) model.train() - gt_head, _ = model._training_filters(batch, batch["triple"]) - self.assertEqual(gt_head, [[e["a"]]]) + drawn = record_negatives(model) np.random.seed(0) - heads, _ = model.train_neg_sample_gen(gt_head, [[]], 256) - self.assertIn(e["b"], heads.flatten().tolist()) - self.assertNotIn(e["a"], heads.flatten().tolist()) + model(**batch)["loss"].backward() + heads = set(drawn["head"].flatten().tolist()) + self.assertIn(e["b"], heads) + self.assertNotIn(e["a"], heads) + + def test_a_model_cannot_be_built_from_the_base_dataset(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + # It has the 1.x count attributes, but no processor fitted on training triples. + with self.assertRaisesRegex(TypeError, "training part"): + TransE(dataset=self.dataset, e_dim=4, r_dim=4) class TestDeprecatedEntryPoints(unittest.TestCase):