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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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/42] 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): From 062e6889ddc4e4143ab625ec952a409b5348e5b5 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:22 +0100 Subject: [PATCH 26/42] fix(kg_emb): replace np.in1d, removed in NumPy 2.4, with np.isin pyproject.toml allows numpy>=2.1,<3 since #1287, and np.in1d no longer exists in NumPy 2.4, so the first training step of any KGE model could raise AttributeError in an environment the declared range allows. assume_unique=True is dropped. The negatives are drawn with replacement, and on NumPy's sort path the flag removes repeated draws of a valid negative. No positive was ever let through; only the sampling scheme depended on the size of the filter set. The RotatE reference implementation passes the same flag. The new tests fail if np.in1d is called and if repeated draws are dropped. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/models/kg_base.py | 17 +++------ tests/core/test_kg_emb.py | 37 +++++++++++++++++++ 2 files changed, 42 insertions(+), 12 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 4f6fcd146..8967448ed 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -268,12 +268,10 @@ def train_neg_sample_gen(self, gt_head, gt_tail, negative_sampling): while negative_sample_size_head < negative_sampling: negative_sample = np.random.randint(self.e_num, size=negative_sampling*2) - mask = np.in1d( - negative_sample, - gt_head[i], - assume_unique=True, - invert=True - ) + # No assume_unique=True (unlike the RotatE reference): the draws are + # made with replacement, and on NumPy's sort path the flag silently + # drops repeated draws of a valid negative. + mask = np.isin(negative_sample, gt_head[i], invert=True) negative_sample = negative_sample[mask] negative_sample_list_head.append(negative_sample) negative_sample_size_head += negative_sample.size @@ -284,12 +282,7 @@ def train_neg_sample_gen(self, gt_head, gt_tail, negative_sampling): while negative_sample_size_tail < negative_sampling: negative_sample = np.random.randint(self.e_num, size=negative_sampling*2) - mask = np.in1d( - negative_sample, - gt_tail[i], - assume_unique=True, - invert=True - ) + mask = np.isin(negative_sample, gt_tail[i], invert=True) negative_sample = negative_sample[mask] negative_sample_list_tail.append(negative_sample) negative_sample_size_tail += negative_sample.size diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index b3475cb2a..6962f842d 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -933,6 +933,20 @@ def test_train_mode_returns_a_loss_and_draws_negatives(self) -> None: out["loss"].backward() self.assertIsNotNone(model.E_emb.grad) + def test_training_does_not_call_np_in1d(self) -> None: + # np.in1d no longer exists in NumPy 2.4, which pyproject.toml allows. + from unittest import mock + + import numpy as np + + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + model = TransE(dataset=self.train, e_dim=8, r_dim=8, negative_sampling=5) + model.train() + removed = AssertionError("np.in1d was removed in NumPy 2.4") + with mock.patch.object(np, "in1d", side_effect=removed, create=True): + model(**self.batch(self.train))["loss"].backward() + def test_subsampling_weights_come_from_the_processor(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE @@ -1017,6 +1031,29 @@ class _Train: self.assertEqual(set(drawn["tail"].flatten().tolist()), {0, 2, 4, 5}) self.assertEqual(set(drawn["head"].flatten().tolist()), {1, 2, 3, 5}) + def test_negatives_are_drawn_with_replacement(self) -> None: + # With assume_unique=True, the repeated draws 5 and 7 would be dropped: + # for these sizes NumPy takes its sort path, where the flag matters. + from types import SimpleNamespace + from unittest import mock + + import numpy as np + + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + from pyhealth.processors import KGTripleProcessor + + processor = KGTripleProcessor(num_entities=2000, num_relations=1) + part = SimpleNamespace(input_processors={"triple": processor}) + model = TransE(dataset=part, e_dim=2, r_dim=2) + known = list(range(100, 2000, 100)) + draws = np.array([5, 5, 7, 7, 100, 200, 8, 9]) + with mock.patch.object(np.random, "randint", return_value=draws): + head, tail = model.train_neg_sample_gen( + gt_head=[known], gt_tail=[known], negative_sampling=4 + ) + self.assertEqual(head.tolist(), [[5, 5, 7, 7]]) + self.assertEqual(tail.tolist(), [[5, 5, 7, 7]]) + class TestTrainingNegativesOnARealSplit(unittest.TestCase): """The same property on set_task output, through forward: (b, r2, d) is a From 08fd5c0edaaad41b29e1ce70a161c1be4ebb87cf Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:23 +0100 Subject: [PATCH 27/42] docs(kg_emb): state what held-out triples do and do not shape With a split, the training filters and the subsampling weights come from the training triples only, but negatives are drawn from every entity of the graph, so an entity that appears only in validation or test triples is still drawn as a training negative. Four statements said that held-out triples never shape training; they now say what holds. The migration note also lists the other changes that make 2.0 metrics not comparable with 1.x runs: deduplicated triples, subsampling weights from the training part, and dev=True keeping 1000 triples of graph.txt where 1.x loaded the whole of graph_filtered.txt. The dev parameter no longer promises "the first 1000 triples", which holds only when the table is read as a single partition. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 16 ++++++++++++---- .../pyhealth.processors.KGTripleProcessor.rst | 9 ++++++--- .../kg_emb/datasets/base_kg_dataset.py | 12 ++++++++---- .../kg_emb/datasets/umls.py | 4 +++- .../kg_emb/tasks/kg_link_prediction.py | 5 +++-- 5 files changed, 32 insertions(+), 14 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 64e263cc5..14081d8dc 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -158,8 +158,12 @@ pipeline: (: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 + as in the reference implementation of Sun et al. (2019). The training + filters and the subsampling weights therefore come from the training + triples only. The entity set is that of the whole graph, however, so an + entity that appears only in validation or test triples is still drawn as a + training negative; such held-out triples are kept and scored, and + ``set_task`` logs how many there are in each part. ``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 @@ -172,8 +176,12 @@ 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``. 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 +1.x behaviour, including training negatives filtered with, and subsampling +weights computed on, every triple of the graph, held-out ones included; only +the 2.0 path above avoids that. The 2.0 path also removes duplicate triples, +where 1.x made one sample per row, so its sample counts, weights and metrics +are not comparable with 1.x runs. ``UMLSDataset(dev=True)`` now keeps 1000 +triples of ``graph.txt``, where 1.x loaded the whole of ``graph_filtered.txt``. 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. diff --git a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst index b673eeace..5a1ec51e2 100644 --- a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst +++ b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst @@ -11,9 +11,12 @@ dicts used to train KG embedding models: the known heads of each 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. +``set_task(task, split=PatientSplit(...))``: the dicts that filter the +training negatives, and the subsampling weights, then come from the training +triples only. The negatives themselves are drawn from every entity of the +graph, including entities that appear only in validation or test triples. +``size()`` returns the number of entities, which is given (ids are global), +not fitted. API Reference ------------- 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 caf596759..87b08c4a9 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 @@ -188,8 +188,11 @@ 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. + Ids are global, so such triples are kept and scored. An entity absent + from the training triples is still drawn as a negative, so the data terms + of the loss only train it to score low; a relation absent from them gets + no gradient from those terms. Regularization, such as the default ``l3`` + of ComplEx and DistMult, still reaches both. """ fitted = parts[0].input_processors["triple"] entities = {h for h, _ in fitted.true_tail} | { @@ -230,8 +233,9 @@ class BaseKGDataset(BaseDataset): 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. + dev: Keep 1000 triples, chosen by the dev filter of ``BaseDataset``: + the first 1000 of the file when the table is read as a single + partition, not necessarily otherwise. refresh_cache: Deprecated and ignored. The cache key now covers the config and the source files, so an edited file builds a new cache. diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 802a43671..14bacb7d9 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -39,7 +39,9 @@ class UMLSDataset(BaseKGDataset): 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``). + ``num_workers``, ``dev``). With ``dev=True`` the dataset still + reads ``graph.txt`` and keeps 1000 of its triples; PyHealth 1.x + read the whole of ``graph_filtered.txt`` instead. Raises: ValueError: If ``root`` is a URL. 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 index bbda112f4..1b055cc0e 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py @@ -38,8 +38,9 @@ class KGLinkPrediction(BaseTask): 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. + them: it filters its negatives with the positives of the training triples + only, held by the ``kg_triple`` processor fitted on the training part. + The negatives themselves are drawn from every entity of the graph. Because ``set_task`` builds samples patient by patient, and each triple is its own patient, these graph-wide lists are computed beforehand, in From f283fa56b1002816cf5bcbce5c52c05193d75e85 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:24 +0100 Subject: [PATCH 28/42] docs(kg_emb): document tie handling and the UMLS inverse-relation leakage Ranks are computed with np.argsort, which follows no stated tie-breaking convention. Every triple of UMLS graph_filtered.txt has its exact inverse under one of four relation pairs, so a triple-level split leaves the inverse of most test triples in training: with PatientSplit((0.9, 0.05, 0.05), seed=0), 3924 of the 4409 test triples have it there, and an inverse look-up rule reaches a filtered tail Hits@1 of at least 0.89 under any tie-breaking. graph.txt has not been checked. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 11 +++++++++++ .../pretrained_embeddings/kg_emb/datasets/umls.py | 12 ++++++++++++ 2 files changed, 23 insertions(+) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 14081d8dc..e15beaeac 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -172,6 +172,17 @@ pipeline: See ``examples/kg_emb_link_prediction.py`` for a self-contained walk-through. +Two caveats apply to the reported metrics. A candidate whose score ties with +the true entity's is ranked in the order returned by ``np.argsort``, not by a +stated convention, so ranks are ill-defined when many scores tie, as with +degenerate embeddings. In ``graph_filtered.txt``, the UMLS file that PyHealth +1.x used in dev mode, every triple has its exact inverse, so a triple-level +split leaves the inverse of most test triples in the training part +(``graph.txt`` has not been checked; see +:class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.UMLSDataset`); +this is the leakage reported for FB15k and WN18 by Toutanova and Chen (2015) +and Dettmers et al. (2018). + 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`, diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 14bacb7d9..3794ea9fa 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -32,6 +32,18 @@ class UMLSDataset(BaseKGDataset): PyHealth 1.x: embeddings trained with 1.x must be mapped through the ``id2entity`` / ``id2relation`` saved with them. + Every triple of ``graph_filtered.txt``, the file of the 1.x bucket that + PyHealth 1.x read with ``dev=True``, has its exact inverse under one of + four pairs of relations, such as ``may_treat`` and ``may_be_treated_by``. + A triple-level split therefore leaves the inverse of most test triples in + the training part: with that file in ``root`` as ``graph.txt`` and + ``PatientSplit(ratios=(0.9, 0.05, 0.05), seed=0)``, 3924 of the 4409 test + triples (0.89) have it there, so a rule that only looks the inverse up + reaches a filtered tail Hits@1 of at least 0.89, whatever the + tie-breaking. ``graph.txt`` itself has not been checked. Report UMLS + metrics next to such a baseline, or keep the two triples of each inverse + pair on the same side of the split. + Args: root: Local directory holding ``graph.txt``. URLs are not accepted; the 1.x bucket ``https://storage.googleapis.com/pyhealth/umls/`` From 8b0950034e501ac78da13d71beeebc9f0f6f1647 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:25 +0100 Subject: [PATCH 29/42] docs(kg_emb): use the default ns in the TransE and RotatE examples With ns != "adv", TransE scores by +distance and RotatE raises AttributeError at its first forward. Both defects predate this PR and are left for a separate issue; the examples and the 1.x train-step test no longer advertise that mode for these two models. Co-Authored-By: Claude Opus 5.5 --- pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py | 2 +- pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py | 2 +- tests/core/test_kg_emb.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index b057a642a..b3eaa5020 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -20,7 +20,7 @@ class RotatE(KGEBaseModel): ... input_processors = { ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) ... } - >>> model = RotatE(_Toy(), e_dim=4, r_dim=2, ns="uniform") + >>> model = RotatE(_Toy(), e_dim=4, r_dim=2) >>> tuple(model.E_emb.shape) (2, 4) """ diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 1d875ec8e..373dbe8c0 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -20,7 +20,7 @@ class TransE(KGEBaseModel): ... input_processors = { ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) ... } - >>> model = TransE(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> model = TransE(_Toy(), e_dim=4, r_dim=4) >>> tuple(model.E_emb.shape) (2, 4) """ diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 6962f842d..0ed41f3ba 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -249,7 +249,7 @@ def test_transe_train_step(self) -> None: 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") + model = TransE(dataset=dataset, e_dim=8, r_dim=8) out = model(**next(iter(loader))) self.assertIn("loss", out) out["loss"].backward() From 074848d2b154fe8ba8994cd8eb3d383ed2dfd42f Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:25 +0100 Subject: [PATCH 30/42] fix(kg_emb): let UMLSDataset jobs prepare the same root at once Jobs started together on a fresh root, such as parallel hyper-parameter runs, wrote the prepared copy under one fixed temporary name, and on Windows one of two concurrent jobs failed with PermissionError. Each call now writes under a unique name and removes it on failure. When Windows refuses os.replace because another job is replacing the same file, the call accepts that job's copy if it is up to date. The class docstring now states the remaining limitation: a graph.txt replaced by an older file of the same size is not detected. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/umls.py | 33 +++++++++++++++---- tests/core/test_kg_emb.py | 31 +++++++++++++++++ 2 files changed, 57 insertions(+), 7 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 3794ea9fa..92616ad06 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -1,6 +1,7 @@ import logging import os import shutil +import uuid from pathlib import Path from pyhealth.datasets.base_dataset import is_url @@ -25,7 +26,10 @@ class UMLSDataset(BaseKGDataset): 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. + it in size, so replacing ``graph.txt`` gives a new dataset cache. A + ``graph.txt`` replaced by an older file of the same size, as ``cp -p`` or + ``rsync -t`` can produce, is not detected: delete ``umls-pyhealth.tsv`` + in that case. Entity and relation ids are assigned by sorted name (see :class:`BaseKGDataset`), no longer by order of first appearance as in @@ -147,12 +151,27 @@ def prepare_metadata(root: str) -> None: 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) + # would take as prepared. The name is unique, so jobs preparing the + # same root at once (parallel hyper-parameter runs) never write into + # the same file; "xb" keeps the default permissions, which mkstemp + # would narrow to the owner. A process killed mid-copy leaves its + # partial file behind, which is never taken as prepared. + partial = f"{target}.{uuid.uuid4().hex}.partial" + try: + with open(source, "rb") as src, open(partial, "xb") as dst: + dst.write(HEADER) + shutil.copyfileobj(src, dst) + try: + os.replace(partial, target) + except PermissionError: + # Windows refuses to replace a file that another process is + # replacing at the same moment. If that process has just + # written an up-to-date copy, there is nothing left to do. + if UMLSDataset._needs_preparation(root): + raise + finally: + if os.path.exists(partial): + os.remove(partial) if __name__ == "__main__": diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 0ed41f3ba..efbcefaaf 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -395,6 +395,37 @@ def test_a_replaced_graph_txt_is_prepared_again(self) -> None: second = UMLSDataset(root=tmp, cache_dir=root / "cache") self.assertEqual((second.num_entities, second.num_relations), (3, 3)) + def test_a_concurrent_up_to_date_copy_is_accepted(self) -> None: + # Windows refuses os.replace while another job replaces the same file. + from unittest import mock + + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import umls + + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + (root / "graph.txt").write_text("A\tPAR\tB\n") + UMLSDataset.prepare_metadata(tmp) # the other job's copy + busy = PermissionError("the target is being replaced") + with mock.patch.object(umls.os, "replace", side_effect=busy): + UMLSDataset.prepare_metadata(tmp) + self.assertEqual( + sorted(f.name for f in root.iterdir()), ["graph.txt", "umls-pyhealth.tsv"] + ) + + def test_a_failed_replacement_without_an_up_to_date_copy_raises(self) -> None: + from unittest import mock + + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import umls + + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + (root / "graph.txt").write_text("A\tPAR\tB\n") + busy = PermissionError("the target is being replaced") + replace = mock.patch.object(umls.os, "replace", side_effect=busy) + with replace, self.assertRaises(PermissionError): + UMLSDataset.prepare_metadata(tmp) + self.assertEqual([f.name for f in root.iterdir()], ["graph.txt"]) + 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/") From 6d3076315370172723aee4e28ba5afd6a1ec81dc Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Tue, 27 Oct 2026 14:33:26 +0100 Subject: [PATCH 31/42] test(kg_emb): reference scores, filtered ranking by hand, all four models - Reference values for the four score functions, chosen so that every sign, the ComplEx conjugate and the TransE norm change the result. - Filtered ranks, MRR and Hits@1 computed by hand through forward and the rank metrics, including a padding 0 next to a real entity 0. - A train and eval step for each of the four models on the 2.0 path, and a training step with ns="uniform" for DistMult. - One shared set_task fixture instead of three identical builds; the kg_emb logger level is restored after it. - Remove test_package_imports, which cannot fail once the module has been collected. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_kg_emb.py | 247 +++++++++++++++++++++++++++++++------- 1 file changed, 206 insertions(+), 41 deletions(-) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index efbcefaaf..28827192b 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -90,12 +90,19 @@ def make_dataset(n: int = 8, **kwargs: Any) -> SampleKGDataset: ) +def processor_dataset(num_entities: int, num_relations: int) -> Any: + """Stands in for a set_task part: the models only read its triple processor.""" + from types import SimpleNamespace + + from pyhealth.processors import KGTripleProcessor + + processor = KGTripleProcessor(num_entities=num_entities, num_relations=num_relations) + return SimpleNamespace(input_processors={"triple": processor}) + + class TestKGEmbImports(unittest.TestCase): """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 - def test_model_classes_are_exported(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( ComplEx, @@ -466,6 +473,61 @@ def test_transe_scores_a_perfect_triple_at_the_margin(self) -> None: ) +class TestScoringReferenceValues(unittest.TestCase): + """Each score function on hand-set embeddings, against its formula worked by hand. + + The values are chosen so that a sign error, a missing conjugate or a wrong + norm gives a different number. Shapes are (batch, candidates, dim). + """ + + def check(self, model, head, relation, tail, expected) -> None: + h, r, t = (torch.tensor([[x]]) for x in (head, relation, tail)) + for mode in ("pos", "head", "tail"): + with self.subTest(model=type(model).__name__, mode=mode): + self.assertAlmostEqual(model.calc(h, r, t, mode=mode).item(), expected, places=5) + + def test_transe(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + # gamma - ||h + r - t||_1 = 24 - (|1 + 0.5 - 2| + |2 + 0.25 - 1|) = 24 - 1.75; + # a negative component, so a norm without absolute values differs. + model = TransE(processor_dataset(2, 1), e_dim=2, r_dim=2) + self.check(model, [1.0, 2.0], [0.5, 0.25], [2.0, 1.0], 22.25) + + def test_distmult(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + + # sum(h * r * t) = 1 * 3 * 0.5 + 2 * (-1) * 4 + model = DistMult(processor_dataset(2, 1), e_dim=2, r_dim=2) + self.check(model, [1.0, 2.0], [3.0, -1.0], [0.5, 4.0], -6.5) + + def test_complex(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ComplEx + + # Re(h * r * conj(t)) with h = 1 + 2i, r = 3 + 5i, t = 7 + 11i: + # Re((-7 + 11i)(7 - 11i)) = 72, every real and imaginary part being + # non-zero. Without the conjugate it would be -170. + model = ComplEx(processor_dataset(2, 1), e_dim=2, r_dim=2) + self.check(model, [1.0, 2.0], [3.0, 5.0], [7.0, 11.0], 72.0) + + def test_rotate(self) -> None: + import cmath + import math + + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import RotatE + + # gamma - |h * exp(i theta) - t| with h = 1 and theta = pi / 2, the phase + # being relation / (r_emb_range / pi): t = i scores 24, t = 1 scores 24 - sqrt 2. + model = RotatE(processor_dataset(2, 1), e_dim=2, r_dim=1) + quarter_turn = [model.r_emb_range.item() / 2] + self.check(model, [1.0, 0.0], quarter_turn, [0.0, 1.0], 24.0) + self.check(model, [1.0, 0.0], quarter_turn, [1.0, 0.0], 24.0 - math.sqrt(2)) + # A general case, with no zero part: h = 1 + 2i, theta = pi / 4, t = 3 - i. + eighth_turn = [model.r_emb_range.item() / 4] + rotated = complex(1, 2) * cmath.exp(1j * math.pi / 4) + self.check(model, [1.0, 2.0], eighth_turn, [3.0, -1.0], 24.0 - abs(rotated - complex(3, -1))) + + 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 @@ -685,41 +747,63 @@ 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.""" +_TASK_SPLIT: dict[str, Any] = {} - @classmethod - def setUpClass(cls) -> None: + +def task_split() -> dict[str, Any]: + """set_task on TASK_TRIPLES with a triple-level split, built once. + + Tests may call set_task on the shared dataset again: another task or other + processors get their own cache entry, so the parts held here are never + rewritten. Tests must not refit or modify the parts' processors. The log + records are captured while it is built, and the kg_emb logger is left as + it was found. + """ + if not _TASK_SPLIT: 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, + tmp = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + root = Path(tmp.name) + config = write_kg(root, TASK_TRIPLES) + dataset = BaseKGDataset(root=str(root), config_path=config, cache_dir=root / "cache") + task = KGLinkPrediction( + num_entities=dataset.num_entities, num_relations=dataset.num_relations ) - cls.split = PatientSplit(ratios=TASK_SPLIT_RATIOS, seed=3) + split = PatientSplit(ratios=TASK_SPLIT_RATIOS, seed=3) handler = _ListHandler() kg_logger = logging.getLogger("pyhealth.medcode.pretrained_embeddings.kg_emb") + level = kg_logger.level kg_logger.addHandler(handler) kg_logger.setLevel(logging.INFO) try: - cls.parts = cls.dataset.set_task(cls.task, split=cls.split, num_workers=1) + parts = dataset.set_task(task, split=split, num_workers=1) finally: kg_logger.removeHandler(handler) - cls.log = handler.messages + kg_logger.setLevel(level) + _TASK_SPLIT.update( + tmp=tmp, root=root, dataset=dataset, task=task, split=split, + parts=parts, log=handler.messages, + ) + return _TASK_SPLIT - e, r = cls.dataset.entity2id, cls.dataset.relation2id - cls.distinct = sorted({(e[h], r[rel], e[t]) for h, rel, t in TASK_TRIPLES}) + +def tearDownModule() -> None: + if _TASK_SPLIT: + _TASK_SPLIT["tmp"].cleanup() + + +class TestKGLinkPrediction(unittest.TestCase): + """KGLinkPrediction through the real set_task, with a triple-level split.""" @classmethod - def tearDownClass(cls) -> None: - cls._tmp.cleanup() + def setUpClass(cls) -> None: + shared = task_split() + cls.root, cls.dataset, cls.task = shared["root"], shared["dataset"], shared["task"] + cls.split, cls.parts, cls.log = shared["split"], shared["parts"], shared["log"] + + e, r = cls.dataset.entity2id, cls.dataset.relation2id + cls.distinct = sorted({(e[h], r[rel], e[t]) for h, rel, t in TASK_TRIPLES}) def triples_of(self, part) -> list[tuple[int, int, int]]: return [tuple(part[i]["triple"].tolist()) for i in range(len(part))] @@ -904,23 +988,9 @@ class TestModelsOnSampleDatasets(unittest.TestCase): @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() + shared = task_split() + cls.dataset = shared["dataset"] + cls.train, cls.val, cls.test = shared["parts"] def batch(self, part, size=4): from pyhealth.datasets import get_dataloader @@ -942,6 +1012,48 @@ def test_counts_come_from_the_fitted_processor(self) -> None: self.assertEqual(model.r_num, self.dataset.num_relations) self.assertIs(model.triple_processor, self.train.input_processors["triple"]) + def test_every_model_trains_and_evaluates(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + RotatE, + TransE, + ) + + train_batch = self.batch(self.train) + test_batch = self.batch(self.test, size=len(self.test)) + for cls in (TransE, RotatE, DistMult, ComplEx): + with self.subTest(model=cls.__name__): + torch.manual_seed(0) + model = cls( + dataset=self.train, e_dim=8, r_dim=4 if cls is RotatE else 8, + negative_sampling=4, + ) + model.train() + loss = model(**train_batch)["loss"] + self.assertTrue(torch.isfinite(loss)) + loss.backward() + self.assertTrue(torch.isfinite(model.E_emb.grad).all()) + + model.eval() + with torch.no_grad(): + y_prob = model(**test_batch)["y_prob"] + n = len(self.test) + self.assertEqual(tuple(y_prob.shape), (2 * n, self.dataset.num_entities)) + self.assertTrue(torch.isfinite(y_prob).all()) + + def test_uniform_negative_weighting_trains(self) -> None: + # ns="uniform" is valid for DistMult, whose score does not depend on ns: + # it reaches the loss branch without self-adversarial weighting. + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + + torch.manual_seed(0) + model = DistMult(dataset=self.train, e_dim=8, r_dim=8, ns="uniform", negative_sampling=4) + model.train() + loss = model(**self.batch(self.train))["loss"] + self.assertTrue(torch.isfinite(loss)) + loss.backward() + def test_negative_sampling_is_a_model_argument(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE @@ -1013,6 +1125,60 @@ def test_a_disagreeing_1x_train_flag_warns(self) -> None: self.assertIn("y_prob", out) # model.eval() decides, not the flag +class TestFilteredRankingByHand(unittest.TestCase): + """forward in eval mode, then the rank metrics, on a TransE scored by hand. + + A filtered candidate must fall below the true entity, and the padding of a + kg_entity_list field must not be taken for entity 0. + """ + + def test_filtered_scores_ranks_and_mrr(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + from pyhealth.metrics import multiclass_metrics_fn + + model = TransE(processor_dataset(4, 1), e_dim=1, r_dim=1, gamma=10.0) + with torch.no_grad(): + model.E_emb.copy_(torch.tensor([[0.0], [1.2], [2.5], [4.0]])) + model.R_emb.copy_(torch.tensor([[1.0]])) + model.eval() + # Both rows rank (0, r, 2); the score is 10 - |h + 1 - t|. + # candidate heads, t = 2.5: 8.5 (true), 9.7, 9.0, 7.5 + # candidate tails, h = 0: 9.0, 9.8, 8.5 (true), 7.0 + # Row 0: heads [0] and tails [2, 1], each followed by a padding 0. + # Row 1: entity 0 is a real member of both lists. + batch = { + "triple": torch.tensor([[0, 0, 2], [0, 0, 2]]), + "ground_truth_head": { + "value": torch.tensor([[0, 0], [0, 2]]), + "mask": torch.tensor([[1, 0], [1, 1]]), + }, + "ground_truth_tail": { + "value": torch.tensor([[2, 1, 0], [2, 1, 0]]), + "mask": torch.tensor([[1, 1, 0], [1, 1, 1]]), + }, + } + with torch.no_grad(): + out = model(**batch) + + # A filtered candidate is scored as the true entity minus 1. + expected = torch.tensor([ + [8.5, 9.7, 9.0, 7.5], # head, row 0: nothing to filter + [8.5, 9.7, 7.5, 7.5], # head, row 1: head 2 filtered + [9.0, 7.5, 8.5, 7.0], # tail, row 0: tail 1 filtered, the padding 0 is not + [7.5, 7.5, 8.5, 7.0], # tail, row 1: tails 1 and 0 filtered + ]) + self.assertTrue(torch.allclose(out["y_prob"], expected, atol=1e-5), out["y_prob"]) + self.assertEqual(out["y_true"].tolist(), [0, 0, 2, 2]) + + scores = multiclass_metrics_fn( + out["y_true"].numpy(), out["y_prob"].numpy(), metrics=["mean_rank", "hits@n"] + ) + # Filtered ranks 3, 2, 2, 1; unfiltered they would be 3, 3, 3, 3. + self.assertAlmostEqual(scores["mean_rank"], 2.0) + self.assertAlmostEqual(scores["mean_reciprocal_rank"], (1 / 3 + 1 / 2 + 1 / 2 + 1) / 4) + self.assertAlmostEqual(scores["HITS@1"], 0.25) + + def record_negatives(model) -> dict: """Wraps train_neg_sample_gen to keep what forward draws.""" drawn = {} @@ -1093,7 +1259,6 @@ class TestTrainingNegativesOnARealSplit(unittest.TestCase): # 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 30a4999cf910deb089ad47fe69b61a77aa5c1d36 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:53 +0100 Subject: [PATCH 32/42] refactor(kg_emb): remove the PyHealth 1.x compatibility layer kg_emb has not been importable in any 2.x release (#952), so no 2.x user depends on its 1.x entry points, and the maintainer chose to remove them rather than deprecate them for a release (#1202). Removed: SampleKGDataset, kg_emb.datasets.split, link_prediction_fn, the 1.x members of BaseKGDataset (set_task with a function or task_fn= and its task_name=, save= and extra keyword arguments, refresh_cache, entity_num, relation_num, triples, stat(), info()), the 1.x branches of KGEBaseModel, the __main__ demos of the four models, and their deprecation warnings and tests. KGEBaseModel now requires a kg_triple processor, and negative_sampling comes only from its argument, 128 by default. Its private helpers _training_filters and _subsampling_weight lose the batch argument that only the 1.x branch read. The tests that used 1.x code are rewritten without it: their models are built from a kg_triple processor rather than a SampleKGDataset, and the comparisons with link_prediction_fn give way to values computed by hand or by a second processor. The migration note in docs/api/medcode.rst lists the names removed here, and len(), removed earlier in this PR, and points to the 2.0 path. The 2.0 path is not changed. examples/kg_embedding.ipynb still uses the 1.x API. It already fails at its first cell on master (SampleBaseDataset no longer exists) and hard-codes private paths, so it is left as it is; docs/tutorials.rst marks it as 1.x and now lists examples/kg_emb_link_prediction.py. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 36 +- docs/tutorials.rst | 5 +- .../kg_emb/datasets/__init__.py | 4 - .../kg_emb/datasets/base_kg_dataset.py | 153 +------ .../kg_emb/datasets/sample_kg_dataset.py | 201 ---------- .../kg_emb/datasets/splitter.py | 113 ------ .../kg_emb/models/complex.py | 51 --- .../kg_emb/models/distmult.py | 51 --- .../kg_emb/models/kg_base.py | 100 +---- .../kg_emb/models/rotate.py | 51 --- .../kg_emb/models/transe.py | 51 --- .../kg_emb/tasks/__init__.py | 3 +- .../kg_emb/tasks/link_prediction.py | 99 ----- tests/core/test_kg_emb.py | 373 ++---------------- 14 files changed, 67 insertions(+), 1224 deletions(-) delete mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py delete mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py delete mode 100644 pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index e15beaeac..1cf20e672 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -183,19 +183,22 @@ split leaves the inverse of most test triples in the training part this is the leakage reported for FB15k and WN18 by Toutanova and Chen (2015) and Dettmers et al. (2018). -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``. They keep their -1.x behaviour, including training negatives filtered with, and subsampling -weights computed on, every triple of the graph, held-out ones included; only -the 2.0 path above avoids that. The 2.0 path also removes duplicate triples, -where 1.x made one sample per row, so its sample counts, weights and metrics -are not comparable with 1.x runs. ``UMLSDataset(dev=True)`` now keeps 1000 -triples of ``graph.txt``, where 1.x loaded the whole of ``graph_filtered.txt``. 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. +The PyHealth 1.x entry points ``link_prediction_fn``, ``SampleKGDataset`` +and ``kg_emb.datasets.split`` have been removed, as have ``set_task`` with a +function or ``task_fn=``, its ``task_name=``, ``save=`` and extra keyword +arguments (``negative_sampling`` is now a model argument, step 4), the +``refresh_cache`` argument, ``len()``, and the ``entity_num``, +``relation_num``, ``triples``, ``stat()`` and ``info()`` members of +``BaseKGDataset`` (use ``num_entities``, ``num_relations`` and +``stats()``); use the 2.0 path above. Its +results are not comparable with 1.x runs: 1.x filtered training negatives +with, and computed subsampling weights on, every triple of the graph, +held-out ones included, and made one sample per row of the file, duplicates +included. ``UMLSDataset(dev=True)`` now keeps 1000 triples of ``graph.txt``, +where 1.x loaded the whole of ``graph_filtered.txt``. 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: @@ -209,13 +212,6 @@ appearance: map 1.x embeddings through the ``id2entity`` saved with them. :members: :show-inheritance: -.. 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: diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 68bcb051d..93c50dee1 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -350,6 +350,8 @@ Advanced Topics - Working with OMOP Common Data Model * - ``medcode.py`` - Medical code vocabulary and mappings + * - ``kg_emb_link_prediction.py`` + - Knowledge graph embeddings: link prediction with filtered ranking * - ``benchmark_ehrshot_xgboost.ipynb`` - EHRShot benchmark with XGBoost (notebook) @@ -383,7 +385,8 @@ Notebooks (Interactive) * - ``drug_recommendation/drug_recommendation_mimic3_micron.ipynb`` - MICRON interactive notebook * - ``kg_embedding.ipynb`` - - Knowledge graph embeddings + - Knowledge graph embeddings (PyHealth 1.x API, removed; see + ``kg_emb_link_prediction.py``) * - ``lm_embedding_huggingface.ipynb`` - Language model embeddings with HuggingFace * - ``lm_embedding_openai.ipynb`` diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 02176f5a9..98b1dfec2 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,11 +1,7 @@ from .base_kg_dataset import BaseKGDataset -from .sample_kg_dataset import SampleKGDataset -from .splitter import split from .umls import UMLSDataset __all__ = [ "BaseKGDataset", - "SampleKGDataset", "UMLSDataset", - "split", ] 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 87b08c4a9..7b1831fcb 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 @@ -18,17 +18,13 @@ import inspect import logging import warnings -from collections.abc import Callable from pathlib import Path -from typing import Any import polars as pl from pyhealth.datasets import BaseDataset from pyhealth.tasks import BaseTask -from .sample_kg_dataset import SampleKGDataset - logger = logging.getLogger(__name__) __all__ = [ @@ -236,8 +232,6 @@ class BaseKGDataset(BaseDataset): dev: Keep 1000 triples, chosen by the dev filter of ``BaseDataset``: the first 1000 of the file when the table is read as a single partition, not necessarily otherwise. - 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 @@ -272,15 +266,7 @@ def __init__( cache_dir: str | Path | None = None, num_workers: int = 1, dev: bool = False, - 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], @@ -356,106 +342,17 @@ 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): + def set_task(self, task: BaseTask | None = None, *args, **kwargs): """Builds the sample dataset(s) of ``task``. - With a :class:`~pyhealth.tasks.BaseTask` (by default - :class:`KGLinkPrediction`), this is - :meth:`pyhealth.datasets.BaseDataset.set_task`. For a task whose + This is :meth:`pyhealth.datasets.BaseDataset.set_task`, with + :class:`KGLinkPrediction` as the default 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 - 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``. + triples. """ - 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) if task is None: task = self.default_task @@ -489,45 +386,3 @@ def default_task(self) -> BaseTask: return KGLinkPrediction( num_entities=self.num_entities, num_relations=self.num_relations ) - - 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." - ) - if save is not None: - logger.info("set_task(save=...) is ignored: samples are not pickled.") - # 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 deleted file mode 100644 index 0970b234b..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py +++ /dev/null @@ -1,201 +0,0 @@ -"""Task-specific sample dataset for knowledge-graph embedding models. - -``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. -""" - -from __future__ import annotations - -import warnings -from collections.abc import Mapping, Sequence -from typing import Any - -import torch - -from pyhealth.datasets.sample_dataset import InMemorySampleDataset - -KGSample = Mapping[str, Any] - -__all__ = ["SampleKGDataset"] - - -class SampleKGDataset(InMemorySampleDataset): - r"""In-memory dataset of knowledge-graph link-prediction samples. - - .. warning:: - 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`` - 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: 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. - 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``). - - 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.entity_num, dataset.relation_num - (5, 2) - >>> dataset.id2entity[2] - 'c' - >>> dataset.task_spec_param - {'negative_sampling': 4} - """ - - def __init__( - 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, - 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}), - "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 {}) - 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) - - def stat(self) -> str: - """Print a human-readable summary and return it.""" - 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)}", - 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})" - ) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py deleted file mode 100644 index 84f7f2806..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Ratio-based splitting of a :class:`SampleKGDataset` into train/val/test folds.""" - -from __future__ import annotations - -import math -import warnings -from typing import Any - -import numpy as np - -from .sample_kg_dataset import SampleKGDataset - -__all__ = ["split"] - -Fold = list[dict[str, Any]] - - -def split( - 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. - - .. warning:: - 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 - samples are flagged so that the model switches to filtered ranking - evaluation. - - Args: - 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: - The train, validation and test folds, each a list of sample - dictionaries. - - 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 - """ - 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) - 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] - ] - 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 diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index bb9be0888..e35543def 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,10 +1,7 @@ -from typing import Any - import torch from pyhealth.datasets import SampleDataset -from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -80,51 +77,3 @@ def calc(self, head, relation, tail, mode='pos'): score = score.sum(dim = 2) return score - - -if __name__ == "__main__": - from pyhealth.datasets import get_dataloader - - 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], - 'ground_truth_tail': [398, 244, 3872, 3053, 1711, 2835, 1348, 2309], - 'subsampling_weight': torch.tensor([0.1857]) - }, - { - 'triple': (4, 2, 6502), - 'ground_truth_head': [4, 69, 1470, 505, 3069], - 'ground_truth_tail': [2517, 907, 4859, 5209, 3680, 273, 6502, 1810, 875, 1794, 1070, 192, 3079, 1420, 5649, 4779, 2348, 4991, 2714, 3202, 120, 1942, 259, 1617, 3203, 292, 1585, 2691, 1512, 2187, 2000, 1935, 5863, 2277, 1635, 4912, 2261, 1367, 2286, 2782, 3750, 6157, 2864, 1506, 4507, 1669, 4044, 1336, 3239, 881, 3264, 2841, 410, 1329, 4029, 1752, 1362, 1216], - 'subsampling_weight': torch.tensor([0.1204]) - }, - ] - - for sample in samples: - sample["train"] = True - sample["hyperparameters"] = {"negative_sampling": 8} - - dataset = SampleKGDataset( - samples=samples, - dataset_name="test", - entity_num=8000, - relation_num=8, - ) - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) - - # model - model = ComplEx( - dataset=dataset, - e_dim=600, - r_dim=600, - ) - - # data batch - data_batch = next(iter(train_loader)) - - # try the model - ret = model(**data_batch) - print(ret) - - # try loss backward - ret["loss"].backward() diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index 74bcc7141..2c82d4352 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,10 +1,7 @@ -from typing import Any - import torch from pyhealth.datasets import SampleDataset -from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -62,51 +59,3 @@ def calc(self, head, relation, tail, mode='pos'): score = score.sum(dim=2) return score - - -if __name__ == "__main__": - from pyhealth.datasets import get_dataloader - - 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], - 'ground_truth_tail': [398, 244, 3872, 3053, 1711, 2835, 1348, 2309], - 'subsampling_weight': torch.tensor([0.1857]) - }, - { - 'triple': (4, 2, 6502), - 'ground_truth_head': [4, 69, 1470, 505, 3069], - 'ground_truth_tail': [2517, 907, 4859, 5209, 3680, 273, 6502, 1810, 875, 1794, 1070, 192, 3079, 1420, 5649, 4779, 2348, 4991, 2714, 3202, 120, 1942, 259, 1617, 3203, 292, 1585, 2691, 1512, 2187, 2000, 1935, 5863, 2277, 1635, 4912, 2261, 1367, 2286, 2782, 3750, 6157, 2864, 1506, 4507, 1669, 4044, 1336, 3239, 881, 3264, 2841, 410, 1329, 4029, 1752, 1362, 1216], - 'subsampling_weight': torch.tensor([0.1204]) - }, - ] - - for sample in samples: - sample["train"] = True - sample["hyperparameters"] = {"negative_sampling": 8} - - dataset = SampleKGDataset( - samples=samples, - dataset_name="test", - entity_num=8000, - relation_num=8, - ) - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) - - # model - model = DistMult( - dataset=dataset, - e_dim=600, - r_dim=600, - ) - - # data batch - data_batch = next(iter(train_loader)) - - # try the model - ret = model(**data_batch) - print(ret) - - # try loss backward - ret["loss"].backward() \ No newline at end of file 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 8967448ed..e9ca78cb5 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,5 +1,4 @@ import numbers -import warnings from abc import ABC import numpy as np @@ -27,8 +26,7 @@ class KGEBaseModel(ABC, nn.Module): Args: 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. + by ``set_task``. 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). @@ -37,9 +35,7 @@ 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 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 to the value 1.x - samples carry in ``hyperparameters``, else 128. + per training triple, 128 by default. Examples: >>> from pyhealth.processors import KGTripleProcessor @@ -74,39 +70,19 @@ def __init__( ): 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. + # A base dataset has no fitted processor; point to the part + # that has one rather than fail on the missing processor. 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 - self.e_num = processor.num_entities - 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 - # 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: - negative_sampling = legacy.get("negative_sampling") - self._negative_sampling_chosen = negative_sampling is not None + if not isinstance(processor, KGTripleProcessor): + raise TypeError("dataset needs a kg_triple processor.") + self.triple_processor = processor + self.e_num = processor.num_entities + self.r_num = processor.num_relations if negative_sampling is None: negative_sampling = DEFAULT_NEGATIVE_SAMPLING if ( @@ -386,47 +362,20 @@ 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): + def _training_filters(self, 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 + These are the ``true_head`` / ``true_tail`` dicts of the ``kg_triple`` + processor, 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. + the held-out triples. 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 = [], [] @@ -435,12 +384,7 @@ def _training_filters(self, data, positive_sample): 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. - 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.") + def _subsampling_weight(self, positive_sample): return self.triple_processor.subsampling_weight(positive_sample.cpu()) def forward(self, **data): @@ -450,22 +394,12 @@ def forward(self, **data): 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 self.training: - gt_head, gt_tail = self._training_filters(data, positive_sample) + gt_head, gt_tail = self._training_filters(positive_sample) negative_sample_head, negative_sample_tail = self.train_neg_sample_gen( gt_head=gt_head, gt_tail=gt_tail, - negative_sampling=self._batch_negative_sampling(data), + negative_sampling=self.negative_sampling, ) negative_sample_head, negative_sample_tail = negative_sample_head.to(self.device), negative_sample_tail.to(self.device) @@ -487,7 +421,7 @@ def forward(self, **data): pos_score = F.logsigmoid(self.calc(head=head, relation=relation, tail=tail)).squeeze(dim=1) if self.use_subsampling_weight: - subsampling_weight = self._subsampling_weight(data, positive_sample).to(self.device) + subsampling_weight = self._subsampling_weight(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: diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index b3eaa5020..618800d98 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,10 +1,7 @@ -from typing import Any - import torch from pyhealth.datasets import SampleDataset -from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -76,51 +73,3 @@ def calc(self, head, relation, tail, mode='pos'): score = self.margin.item() - score.sum(dim=2) return score - - -if __name__ == "__main__": - from pyhealth.datasets import get_dataloader - - 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], - 'ground_truth_tail': [398, 244, 3872, 3053, 1711, 2835, 1348, 2309], - 'subsampling_weight': torch.tensor([0.1857]) - }, - { - 'triple': (4, 2, 6502), - 'ground_truth_head': [4, 69, 1470, 505, 3069], - 'ground_truth_tail': [2517, 907, 4859, 5209, 3680, 273, 6502, 1810, 875, 1794, 1070, 192, 3079, 1420, 5649, 4779, 2348, 4991, 2714, 3202, 120, 1942, 259, 1617, 3203, 292, 1585, 2691, 1512, 2187, 2000, 1935, 5863, 2277, 1635, 4912, 2261, 1367, 2286, 2782, 3750, 6157, 2864, 1506, 4507, 1669, 4044, 1336, 3239, 881, 3264, 2841, 410, 1329, 4029, 1752, 1362, 1216], - 'subsampling_weight': torch.tensor([0.1204]) - }, - ] - - for sample in samples: - sample["train"] = True - sample["hyperparameters"] = {"negative_sampling": 8} - - dataset = SampleKGDataset( - samples=samples, - dataset_name="test", - entity_num=8000, - relation_num=8, - ) - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) - - # model - model = RotatE( - dataset=dataset, - e_dim=600, - r_dim=300, - ) - - # data batch - data_batch = next(iter(train_loader)) - - # try the model - ret = model(**data_batch) - print(ret) - - # try loss backward - ret["loss"].backward() diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 373dbe8c0..546e6d8ae 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,10 +1,7 @@ -from typing import Any - import torch from pyhealth.datasets import SampleDataset -from ..datasets.sample_kg_dataset import SampleKGDataset from .kg_base import KGEBaseModel @@ -65,51 +62,3 @@ def calc(self, head, relation, tail, mode='pos'): score = self.margin.item() - score return score - - -if __name__ == "__main__": - from pyhealth.datasets import get_dataloader - - 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], - 'ground_truth_tail': [398, 244, 3872, 3053, 1711, 2835, 1348, 2309], - 'subsampling_weight': torch.tensor([0.1857]) - }, - { - 'triple': (4, 2, 6502), - 'ground_truth_head': [4, 69, 1470, 505, 3069], - 'ground_truth_tail': [2517, 907, 4859, 5209, 3680, 273, 6502, 1810, 875, 1794, 1070, 192, 3079, 1420, 5649, 4779, 2348, 4991, 2714, 3202, 120, 1942, 259, 1617, 3203, 292, 1585, 2691, 1512, 2187, 2000, 1935, 5863, 2277, 1635, 4912, 2261, 1367, 2286, 2782, 3750, 6157, 2864, 1506, 4507, 1669, 4044, 1336, 3239, 881, 3264, 2841, 410, 1329, 4029, 1752, 1362, 1216], - 'subsampling_weight': torch.tensor([0.1204]) - }, - ] - - for sample in samples: - sample["train"] = True - sample["hyperparameters"] = {"negative_sampling": 8} - - dataset = SampleKGDataset( - samples=samples, - dataset_name="test", - entity_num=8000, - relation_num=8, - ) - train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) - - # model - model = TransE( - dataset=dataset, - e_dim=600, - r_dim=600, - ) - - # data batch - data_batch = next(iter(train_loader)) - - # try the model - ret = model(**data_batch) - print(ret) - - # try loss backward - ret["loss"].backward() \ No newline at end of file diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py index a445a57ac..a4e32cc38 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py @@ -1,4 +1,3 @@ from .kg_link_prediction import KGLinkPrediction -from .link_prediction import link_prediction_fn -__all__ = ["KGLinkPrediction", "link_prediction_fn"] +__all__ = ["KGLinkPrediction"] diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py deleted file mode 100644 index 8071fd8b3..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py +++ /dev/null @@ -1,99 +0,0 @@ -import numpy as np -import torch -from tqdm import tqdm -from typing import Tuple, List -from collections import defaultdict -import warnings - -def link_prediction_fn( - triples: List[Tuple] -): - - """Process a triple list for the link prediction task - - .. warning:: - 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) - - Args: - triples: a list of triples (indexed) from the knowledge graph - - 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) - gt_head, gt_tail = ground_truth_for_query(triple_set) - - for positive_sample in tqdm(triples): - head, relation, tail = positive_sample - subsampling_weight = torch.sqrt(1 / torch.Tensor([count[(head, relation)] + count[(tail, -relation-1)]])) - gt_h = gt_head[(relation, tail)] - gt_t = gt_tail[(head, relation)] - - tmp = { - "triple": positive_sample, - "ground_truth_head": gt_h, - "ground_truth_tail": gt_t, - "subsampling_weight": subsampling_weight - } - - samples.append(tmp) - - return samples - - -def ground_truth_for_query(triple_set): - """ - Search ground truth of either query (h, r, ?) or (?, r, t) in the dataset - """ - gt_head = defaultdict(list) - gt_tail = defaultdict(list) - - for triple in triple_set: - head, relation, tail = triple - gt_head[(relation, tail)].append(head) - gt_tail[(head, relation)].append(tail) - - return gt_head, gt_tail - - -def count_frequency(triples, start=4): - ''' - Get frequency of a partial triple like (head, relation) or (relation, tail) - The frequency will be used for subsampling like word2vec - ''' - count = {} - for head, relation, tail in triples: - if (head, relation) not in count: - count[(head, relation)] = start - else: - count[(head, relation)] += 1 - - if (tail, -relation-1) not in count: - count[(tail, -relation-1)] = start - else: - count[(tail, -relation-1)] += 1 - return count - diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 28827192b..e35abddfe 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -16,30 +16,12 @@ 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, UMLSDataset, - 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) - ] - # 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 @@ -52,11 +34,6 @@ def make_samples(n: int = 8) -> list[dict[str, Any]]: ("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: @@ -76,20 +53,6 @@ def write_kg(root: Path, triples: list[tuple[str, str, str]]) -> Path: 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)} - return SampleKGDataset( - samples=make_samples(n), - dataset_name="synthetic", - task_name="link_prediction", - entity2id=entity2id, - relation2id=relation2id, - negative_sampling=4, - **kwargs, - ) - - def processor_dataset(num_entities: int, num_relations: int) -> Any: """Stands in for a set_task part: the models only read its triple processor.""" from types import SimpleNamespace @@ -116,152 +79,6 @@ def test_model_classes_are_exported(self) -> None: self.assertTrue(issubclass(cls, torch.nn.Module)) -class TestSampleKGDataset(unittest.TestCase): - """Construction and indexing: the failure a 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) - # "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: - 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_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 - isolation check.""" - from pyhealth.datasets.sample_dataset import InMemorySampleDataset - - dataset = make_dataset(n=2) - self.assertIsInstance(dataset, torch.utils.data.Dataset) - self.assertIsInstance(dataset, InMemorySampleDataset) - self.assertTrue(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) - # "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"].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 - - 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): - """Tensor Trick collation: KG fields arrive pre-padded, with masks; one train step runs.""" - - 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"], 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 - - 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) - out = model(**next(iter(loader))) - self.assertIn("loss", out) - out["loss"].backward() - - class TestBaseKGDataset(unittest.TestCase): """BaseKGDataset loads one triples table through the BaseDataset backend.""" @@ -287,58 +104,6 @@ def test_vocabularies_are_global_and_sorted(self) -> None: 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), len(TOY_TRIPLES)) - 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) - 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): @@ -360,16 +125,6 @@ def test_a_triple_with_a_missing_field_raises(self) -> None: 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.""" @@ -449,7 +204,7 @@ class TestScoringInvariants(unittest.TestCase): 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) + model = DistMult(dataset=processor_dataset(5, 2), 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( @@ -462,7 +217,7 @@ def test_distmult_is_symmetric_in_head_and_tail(self) -> None: 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) + model = TransE(dataset=processor_dataset(5, 2), 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 @@ -536,7 +291,7 @@ class TestGroundTruthUnpadding(unittest.TestCase): 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) + model = TransE(dataset=processor_dataset(5, 2), 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. @@ -550,7 +305,7 @@ def test_unpad_ground_truth_keeps_real_entity_zero_and_drops_padding(self) -> No 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) + model = TransE(dataset=processor_dataset(5, 2), e_dim=8, r_dim=8) raw = [[0, 5], [3]] self.assertEqual(model._unpad_ground_truth(raw), raw) @@ -592,16 +347,18 @@ def test_dicts_match_the_fitted_triples(self) -> None: 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, + def test_counts_and_weights_follow_the_reference_formula(self) -> None: + """Worked by hand: a pair's count starts at 4 and grows by one per + further triple, and a triple weighs 1 / sqrt(count[(h, r)] + + count[(t, -r - 1)]), as in the reference implementation of RotatE.""" + self.assertEqual( + self.processor.count, + {(0, 0): 5, (1, -1): 4, (2, -1): 5, (3, 1): 4, + (1, -2): 4, (4, 0): 4, (5, 1): 4, (0, -2): 4}, ) - - 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))) + expected = torch.tensor([9.0, 10.0, 8.0, 9.0, 8.0]).rsqrt() + self.assertTrue(torch.allclose(weights, expected)) def test_weight_of_a_triple_with_an_unseen_pair_raises(self) -> None: # (1, 1) never occurs as (head, relation) in the fitted triples. @@ -666,11 +423,11 @@ def test_split_fits_on_the_training_part_only(self) -> None: 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, - ) + from pyhealth.processors import KGTripleProcessor - self.assertEqual(processor.count, count_frequency(sorted(train))) + on_train_only = KGTripleProcessor(num_entities=6, num_relations=2) + on_train_only.fit([{"triple": t} for t in sorted(train)], "triple") + self.assertEqual(processor.count, on_train_only.count) def test_fitting_without_split_warns(self) -> None: samples = [{"patient_id": str(i), "triple": t} for i, t in enumerate(self.triples)] @@ -710,7 +467,7 @@ def test_collation_pads_the_batch_with_mask_zero(self) -> None: # 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) + model = TransE(dataset=processor_dataset(5, 2), e_dim=8, r_dim=8) self.assertEqual(model._unpad_ground_truth(gt), [[5, 0, 7, 9], [4]]) @@ -1060,8 +817,6 @@ def test_negative_sampling_is_a_model_argument(self) -> None: 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) @@ -1095,7 +850,7 @@ def test_subsampling_weights_come_from_the_processor(self) -> None: 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"]) + weights = model._subsampling_weight(batch["triple"]) expected = self.train.input_processors["triple"].subsampling_weight(batch["triple"]) self.assertTrue(torch.equal(weights, expected)) model.train() @@ -1113,17 +868,6 @@ def test_eval_mode_ranks_every_entity(self) -> None: 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 - class TestFilteredRankingByHand(unittest.TestCase): """forward in eval mode, then the rank metrics, on a TransE scored by hand. @@ -1291,47 +1035,11 @@ def test_held_out_heads_are_not_filtered(self) -> None: 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. + # A base dataset has 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): - """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) - - class TestNegativeSamplingReachesTheSampler(unittest.TestCase): """The number of negatives drawn is the one the caller chose.""" @@ -1347,25 +1055,10 @@ def record(**kwargs): 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 = TransE(dataset=dataset, e_dim=4, r_dim=4, **kwargs) model.train() return model @@ -1373,7 +1066,7 @@ 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") + processor.fit([{"triple": t} for t in [(0, 0, 4), (1, 1, 5), (2, 2, 6)]], "triple") class _Train: input_processors = {"triple": processor} @@ -1386,24 +1079,8 @@ class _Train: def test_numpy_integers_are_accepted(self) -> None: import numpy as np - dataset, _ = self.legacy_batch() + dataset = processor_dataset(10, 3) 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) + self.assertEqual(self.build(dataset, negative_sampling=np.int32(5)).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 e19927353464d19ce598a10502bd21223121b325 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:53 +0100 Subject: [PATCH 33/42] perf(kg_emb): keep the filtered-ranking sets once per task Each sample carried every known head of its (relation, tail) pair and every known tail of its (head, relation) pair, padded by KGProcessor to the largest degree, so the processed cache grew with the number of triples times the largest degree: 1,736 MiB of samples on UMLS graph_filtered.txt. set_task now builds one KnownTriples from the triples of all the parts it returns and attaches it to every part as known_triples; the models read it at evaluation, and the samples carry only the triple. Built from the parts rather than from the raw graph, the sets also match a task whose pre_filter drops triples or numbers a subgraph. The training part is not read again: its triples are those the kg_triple processor was fitted on, unless that processor was supplied. The held-out parts are read once, for both the sets and the log of unseen entities, which now counts against the training part's triples even when the processor was supplied. The sets are attached, and the missing-split warning given, for a KGTripleProcessor named, given as a class or given as an instance in the schema. At evaluation, a model built without known_triples, or given a triple outside them, raises a ValueError that says so. KnownTriples is exported and documented. On graph_filtered.txt with PatientSplit((0.9, 0.05, 0.05), seed=0), the samples take 44.0 MiB instead of 1,736.2 and the whole cache 50.4 MiB instead of 1,783.9. Measured back to back in one session on one machine (num_workers=1; two runs of this version, one of the previous), set_task on a fresh cache takes 71 to 73 s instead of 119 s. A call that reuses the cache takes 25 to 26 s instead of 18 s: both versions read the 8,818 held-out samples once, and litdata spends longer mapping the chunks of the new cache. Training and filtered ranks are unchanged: on a synthetic graph, y_prob of the four models on the validation and test parts is bitwise identical before and after. KGProcessor, the unpadding in KGEBaseModel and the padded branch of the shared collate in pyhealth/datasets/utils.py, which only KGProcessor needed, are removed. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 14 +- docs/api/processors.rst | 4 +- .../pyhealth.processors.KGProcessor.rst | 16 - examples/kg_emb_link_prediction.py | 6 +- pyhealth/datasets/utils.py | 41 -- .../kg_emb/datasets/__init__.py | 3 +- .../kg_emb/datasets/base_kg_dataset.py | 144 +++++-- .../kg_emb/models/kg_base.py | 93 ++-- .../kg_emb/tasks/kg_link_prediction.py | 75 +--- pyhealth/processors/__init__.py | 2 - pyhealth/processors/kg_processor.py | 114 ----- pyhealth/processors/kg_triple_processor.py | 4 +- tests/core/test_kg_emb.py | 399 +++++++++++------- 13 files changed, 445 insertions(+), 470 deletions(-) delete mode 100644 docs/api/processors/pyhealth.processors.KGProcessor.rst delete mode 100644 pyhealth/processors/kg_processor.py diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 1cf20e672..8df1bdcee 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -150,10 +150,13 @@ pipeline: 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. + holds one ``triple``. The known heads and tails of every pair, which + filter the ranking at evaluation (filtered setting of Bordes et al., + 2013), are built once from the triples of all the parts and attached to + every part as ``known_triples`` + (:class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.KnownTriples`), + so the cache grows with the number of triples only. + 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 @@ -204,6 +207,9 @@ the ``id2entity`` saved with them. :members: :show-inheritance: +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.KnownTriples + :members: + .. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.UMLSDataset :members: :show-inheritance: diff --git a/docs/api/processors.rst b/docs/api/processors.rst index 7c01030ce..45361db55 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -56,7 +56,6 @@ Available Processors - ``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):** @@ -565,5 +564,4 @@ API Reference processors/pyhealth.processors.StageNetProcessor processors/pyhealth.processors.StageNetTensorProcessor processors/pyhealth.processors.GraphProcessor - processors/pyhealth.processors.KGTripleProcessor - processors/pyhealth.processors.KGProcessor \ No newline at end of file + processors/pyhealth.processors.KGTripleProcessor \ No newline at end of file diff --git a/docs/api/processors/pyhealth.processors.KGProcessor.rst b/docs/api/processors/pyhealth.processors.KGProcessor.rst deleted file mode 100644 index 7143020dd..000000000 --- a/docs/api/processors/pyhealth.processors.KGProcessor.rst +++ /dev/null @@ -1,16 +0,0 @@ -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/examples/kg_emb_link_prediction.py b/examples/kg_emb_link_prediction.py index a9283e35e..c1588551c 100644 --- a/examples/kg_emb_link_prediction.py +++ b/examples/kg_emb_link_prediction.py @@ -12,9 +12,9 @@ 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). + training negatives. The known triples of all the parts, here the whole + graph, which filter the ranking at evaluation (filtered setting of Bordes + et al., 2013), are built once and attached to every part. 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()`` / diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 58658b84b..fa0acf5f7 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -260,22 +260,6 @@ 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() @@ -320,31 +304,6 @@ 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. 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) - # 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/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 98b1dfec2..b31ea4a53 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,7 +1,8 @@ -from .base_kg_dataset import BaseKGDataset +from .base_kg_dataset import BaseKGDataset, KnownTriples from .umls import UMLSDataset __all__ = [ "BaseKGDataset", + "KnownTriples", "UMLSDataset", ] 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 7b1831fcb..3b5d50a07 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 @@ -18,17 +18,20 @@ import inspect import logging import warnings +from dataclasses import dataclass from pathlib import Path import polars as pl from pyhealth.datasets import BaseDataset +from pyhealth.processors import PROCESSOR_REGISTRY, KGTripleProcessor from pyhealth.tasks import BaseTask logger = logging.getLogger(__name__) __all__ = [ "BaseKGDataset", + "KnownTriples", "entity_vocabulary", "index_triples", "relation_vocabulary", @@ -174,35 +177,105 @@ def index_triples( ) +@dataclass +class KnownTriples: + """The known heads and tails of every pair, over the triples of a task. + + Filtered ranking (Bordes et al., 2013) ranks the true entity of a test + triple against every other entity except those that also complete it, + in any part: ``heads[(relation, tail)]`` and ``tails[(head, relation)]`` + hold them, as sorted lists of ids. :meth:`BaseKGDataset.set_task` builds + them once, from the triples of all the parts it returns, and attaches + them to every part as ``part.known_triples``. Their size therefore grows + with the number of triples; a copy in each sample grew with the number + of triples times the largest degree. One instance is shared by every + part and by every model built from them, so treat it as read-only. + Training never reads them: its filters come from the ``kg_triple`` + processor, fitted on the training part. + + Examples: + >>> import polars as pl + >>> ids = pl.DataFrame( + ... {"head_id": [0, 2, 0], "relation_id": [0, 0, 1], "tail_id": [1, 1, 2]} + ... ) + >>> known = KnownTriples.from_ids(ids) + >>> known.heads[(0, 1)], known.tails[(0, 0)] + ([0, 2], [1]) + """ + + heads: dict[tuple[int, int], list[int]] + tails: dict[tuple[int, int], list[int]] + + @classmethod + def from_ids(cls, ids: pl.DataFrame) -> KnownTriples: + """Builds the sets from triples given as ``head_id``, ``relation_id`` + and ``tail_id`` columns; duplicate triples count once.""" + ids = ids.select("head_id", "relation_id", "tail_id").unique() + heads = ids.group_by("relation_id", "tail_id").agg(pl.col("head_id").sort()) + tails = ids.group_by("head_id", "relation_id").agg(pl.col("tail_id").sort()) + return cls( + heads={(r, t): h for r, t, h in heads.iter_rows()}, + tails={(h, r): t for h, r, t in tails.iter_rows()}, + ) + + +def _triples_of(part) -> list[tuple[int, int, int]]: + # Sequential iteration: on a subset, litdata's part[i] costs O(i), which + # made a pass quadratic in the size of the part. + return [tuple(sample["triple"].tolist()) for sample in part] + + +def _fitted_triples(processor: KGTripleProcessor) -> list[tuple[int, int, int]]: + """The triples a ``kg_triple`` processor was fitted on.""" + return [(h, r, t) for (h, r), tails in processor.true_tail.items() for t in tails] + + +def _known_triples(triples: list[tuple[int, int, int]]) -> KnownTriples: + """The :class:`KnownTriples` of ``triples``. + + ``set_task`` passes the triples of its parts rather than those of the + raw graph, so that a task whose ``pre_filter`` drops triples, or numbers + a subgraph, gets filter sets that match its own samples. + """ + schema = {name: pl.Int64 for name in ("head_id", "relation_id", "tail_id")} + return KnownTriples.from_ids(pl.DataFrame(triples, schema=schema, orient="row")) + + def _uses_kg_triple(task: BaseTask) -> bool: + """Whether the task's ``triple`` field uses a ``KGTripleProcessor``, in any + form the schema accepts: a registered name, a class or an instance, alone + or with keyword arguments.""" spec = getattr(task, "input_schema", {}).get("triple") - name = spec[0] if isinstance(spec, tuple) else spec - return name == "kg_triple" + if isinstance(spec, tuple): + spec = spec[0] + if isinstance(spec, str): + spec = PROCESSOR_REGISTRY.get(spec) + if not inspect.isclass(spec): + spec = type(spec) + return issubclass(spec, KGTripleProcessor) -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. +def _log_unseen_in_training( + train: list[tuple[int, int, int]], held_out: list[list[tuple[int, int, int]]] +) -> None: + """Logs, per held-out part, how many triples involve an entity or a + relation absent from every training triple. Ids are global, so such triples are kept and scored. An entity absent from the training triples is still drawn as a negative, so the data terms of the loss only train it to score low; a relation absent from them gets no gradient from those terms. Regularization, such as the default ``l3`` of ComplEx and DistMult, still reaches both. + + Args: + train: The triples of the training part. + held_out: The triples of each held-out part, in part order. """ - 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 + entities = {h for h, _, _ in train} | {t for _, _, t in train} + relations = {r for _, r, _ in train} + for k, triples in enumerate(held_out, start=1): + n_entity = sum(h not in entities or t not in entities for h, _, t in triples) + n_relation = sum(r not in relations for _, r, _ in triples) level = logging.WARNING if n_entity or n_relation else logging.INFO logger.log( level, @@ -210,7 +283,7 @@ def _log_unseen_in_training(parts: tuple) -> None: "training triples, %d a relation absent from them.", k, n_entity, - len(part), + len(triples), n_relation, ) @@ -349,9 +422,13 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): :class:`KGLinkPrediction` as the default 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. + every triple. When the fitted ``triple`` processor is a + :class:`~pyhealth.processors.KGTripleProcessor`, it logs, with + ``split``, for each held-out part, how many triples involve an entity + or a relation absent from the training triples, and every returned + part carries, as ``known_triples``, the :class:`KnownTriples` of the + triples of all the parts, which the models read to filter the + ranking at evaluation. """ if task is None: task = self.default_task @@ -361,9 +438,8 @@ def set_task(self, task: BaseTask | None = None, *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 {} - ) + supplied = "triple" in (bound.get("input_processors") or {}) + fits_triples = _uses_kg_triple(task) and not supplied if fits_triples and split is None: warnings.warn( "set_task without split= fits the kg_triple processor on every " @@ -374,8 +450,22 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): stacklevel=2, ) result = super().set_task(task, *args, **kwargs) - if _uses_kg_triple(task) and isinstance(result, tuple): - _log_unseen_in_training(result) + parts = result if isinstance(result, tuple) else (result,) + processor = parts[0].input_processors.get("triple") + # Test the fitted processor, as the models do, rather than the schema + # string: a processor class or instance in the schema gives the same + # processor. + if isinstance(processor, KGTripleProcessor): + # A processor fitted here holds exactly the triples of the first + # part, usually the largest, which spares reading it again; a + # supplied one may have been fitted on other triples. + first = _triples_of(parts[0]) if supplied else _fitted_triples(processor) + held_out = [_triples_of(part) for part in parts[1:]] + if held_out: + _log_unseen_in_training(first, held_out) + known = _known_triples(first + [t for part in held_out for t in part]) + for part in parts: + part.known_triples = known return result @property 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 e9ca78cb5..4a423b2bc 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -26,7 +26,9 @@ class KGEBaseModel(ABC, nn.Module): Args: dataset: A sample dataset whose ``triple`` field uses the ``kg_triple`` processor, typically the training part returned - by ``set_task``. + by ``set_task``. Evaluation also reads the task's + ``known_triples`` from it, which ``BaseKGDataset.set_task`` + attaches to every part. 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). @@ -83,6 +85,9 @@ def __init__( self.triple_processor = processor self.e_num = processor.num_entities self.r_num = processor.num_relations + # The task's known triples, for filtered ranking at evaluation; set + # on every part by BaseKGDataset.set_task. Training never reads them. + self.known_triples = getattr(dataset, "known_triples", None) if negative_sampling is None: negative_sampling = DEFAULT_NEGATIVE_SAMPLING if ( @@ -185,36 +190,6 @@ 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 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 - "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) @@ -222,17 +197,10 @@ def train_neg_sample_gen(self, gt_head, gt_tail, 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: For each triple, the entity ids kept out of its head + negatives. + gt_tail: The same, for tail negatives. """ - 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)): @@ -281,15 +249,11 @@ def test_neg_sample_filter_bias_gen(self, triples, gt_head, gt_tail): 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: For each triple, every known head of its + ``(relation, tail)`` pair in the task's ``known_triples``, + the true head included. + gt_tail: The same, for the tails of its ``(head, relation)`` pair. """ - gt_head = self._unpad_ground_truth(gt_head) - gt_tail = self._unpad_ground_truth(gt_tail) - negative_sample_head = [] negative_sample_tail = [] filter_bias_head = [] @@ -367,11 +331,10 @@ def _training_filters(self, positive_sample): These are the ``true_head`` / ``true_tail`` dicts of the ``kg_triple`` processor, 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. + reference implementation of Sun et al. (2019). The task's + ``known_triples`` 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. Returns: Two lists of entity-id lists, for head and tail negatives. @@ -438,10 +401,26 @@ def forward(self, **data): return {"loss": loss} else: # valid/test + known = self.known_triples + if known is None: + raise ValueError( + "Filtered evaluation needs the task's known triples: build " + "the model from a part returned by BaseKGDataset.set_task." + ) + triples = positive_sample.tolist() + for h, r, t in triples: + # h is a known head of (r, t) exactly when the triple is known, + # which also makes tails[(h, r)] defined. + if h not in known.heads.get((r, t), ()): + raise ValueError( + f"Triple {(h, r, t)} is not among the known triples of " + "the model's dataset; filtered ranking covers only the " + "triples of the set_task call that returned it." + ) inputs = self.test_neg_sample_filter_bias_gen( - triples=positive_sample.tolist(), - gt_head=data['ground_truth_head'], - gt_tail=data['ground_truth_tail'] + triples=triples, + gt_head=[known.heads[(r, t)] for _, r, t in triples], + gt_tail=[known.tails[(h, r)] for h, r, _ in triples], ) # inputs, mode = (data['positive_sample'], data['negative_sample'], data['filter_bias']), data['mode'] 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 index 1b055cc0e..983f3f5ec 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py @@ -28,23 +28,19 @@ 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 filters its negatives with the positives of the training triples - only, held by the ``kg_triple`` processor fitted on the training part. - The negatives themselves are drawn from every entity of the graph. - - 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``. + Each distinct triple of the graph gives one sample, + ``{"patient_id": ..., "triple": (h, r, t)}``. + + Filtered ranking evaluation (Bordes et al., 2013; Sun et al., 2019) + removes, when a test triple is ranked, the other entities known to + complete it, whatever part they belong to. Those sets are not carried by + the samples: :meth:`BaseKGDataset.set_task` builds them once, from the + triples of all the parts, as + :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.KnownTriples`, + and attaches them to every part. Training does not use them: it filters + its negatives with the positives of the training triples only, held by + the ``kg_triple`` processor fitted on the training part. The negatives + themselves are drawn from every entity of the graph. Duplicate triples are removed in :meth:`pre_filter`, keeping the first in file order, and their number is logged: a duplicate could otherwise land @@ -59,7 +55,7 @@ class KGLinkPrediction(BaseTask): Examples: >>> task = KGLinkPrediction(num_entities=4, num_relations=2) >>> sorted(task.input_schema) - ['ground_truth_head', 'ground_truth_tail', 'triple'] + ['triple'] >>> task.input_schema["triple"] ('kg_triple', {'num_entities': 4, 'num_relations': 2}) >>> train, val, test = dataset.set_task( # doctest: +SKIP @@ -80,24 +76,21 @@ def __init__(self, num_entities: int, num_relations: int): "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. + """Indexes the triples and removes duplicates. 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: + ``relation_id`` and ``tail_id``. 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. + would otherwise repeat the de-duplication each time. Raises: ValueError: If a triple has a missing field, or if the graph @@ -137,32 +130,12 @@ def pre_filter(self, df: pl.LazyFrame) -> pl.LazyFrame: 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") - ) + samples = distinct.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")), + ).sort("patient_id") return samples.lazy() def __call__(self, patient: Patient) -> list[dict[str, Any]]: @@ -175,8 +148,6 @@ def __call__(self, patient: Patient) -> list[dict[str, Any]]: { "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/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 41c2de96d..c48088db5 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -47,7 +47,6 @@ 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 .kg_triple_processor import KGTripleProcessor from .audio_processor import AudioProcessor from .ignore_processor import IgnoreProcessor @@ -82,7 +81,6 @@ def get_processor(name: str): "TimeseriesProcessor", "TimeImageProcessor", "GraphProcessor", - "KGProcessor", "KGTripleProcessor", "AudioProcessor", "TupleTimeTextProcessor", diff --git a/pyhealth/processors/kg_processor.py b/pyhealth/processors/kg_processor.py deleted file mode 100644 index 748f22acb..000000000 --- a/pyhealth/processors/kg_processor.py +++ /dev/null @@ -1,114 +0,0 @@ -from collections.abc import Iterable -from typing import Any - -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). 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 - 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])} - >>> 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): - 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. - - 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(max_length, len(value)),)``, padded with - ``pad_token_id``. - - ``"mask"``: LongTensor of the same shape, 1 for real - entities and 0 for padding. - """ - entities = list(value) if value is not None else [] - 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), - "mask": torch.tensor(mask, dtype=torch.long), - } - - def size(self) -> int: - """Return the fitted padding length, the minimum output 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/pyhealth/processors/kg_triple_processor.py b/pyhealth/processors/kg_triple_processor.py index ec2ec7ce6..c9d74f5cd 100644 --- a/pyhealth/processors/kg_triple_processor.py +++ b/pyhealth/processors/kg_triple_processor.py @@ -55,8 +55,8 @@ class KGTripleProcessor(FeatureProcessor): 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. + filtered ranking uses the known triples of all the parts, which + ``BaseKGDataset.set_task`` attaches to every part. These dicts must describe the training triples only, otherwise the negatives and weights of training depend on validation and test diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index e35abddfe..e88b98acf 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -17,7 +17,6 @@ import torch -from pyhealth.datasets import collate_fn_dict_with_padding from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( BaseKGDataset, UMLSDataset, @@ -53,14 +52,17 @@ def write_kg(root: Path, triples: list[tuple[str, str, str]]) -> Path: return config -def processor_dataset(num_entities: int, num_relations: int) -> Any: - """Stands in for a set_task part: the models only read its triple processor.""" +def processor_dataset(num_entities: int, num_relations: int, known_triples=None) -> Any: + """Stands in for a set_task part: the models read its triple processor and, + to evaluate, its known triples.""" from types import SimpleNamespace from pyhealth.processors import KGTripleProcessor processor = KGTripleProcessor(num_entities=num_entities, num_relations=num_relations) - return SimpleNamespace(input_processors={"triple": processor}) + return SimpleNamespace( + input_processors={"triple": processor}, known_triples=known_triples + ) class TestKGEmbImports(unittest.TestCase): @@ -283,33 +285,6 @@ def test_rotate(self) -> None: self.check(model, [1.0, 2.0], eighth_turn, [3.0, -1.0], 24.0 - abs(rotated - complex(3, -1))) -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=processor_dataset(5, 2), 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=processor_dataset(5, 2), 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 @@ -436,41 +411,6 @@ def test_fitting_without_split_warns(self) -> None: 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=processor_dataset(5, 2), 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 = [ @@ -490,11 +430,6 @@ def test_collation_pads_the_batch_with_mask_zero(self) -> None: 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) @@ -595,20 +530,22 @@ def test_the_first_duplicate_in_file_order_is_kept(self) -> None: ] self.assertEqual(kept, ["0"]) - def test_ground_truth_covers_the_whole_graph(self) -> None: + def test_samples_carry_only_the_triple(self) -> None: + for part in self.parts: + self.assertEqual(set(part.input_processors), {"triple"}) + self.assertEqual(set(part[0]), {"patient_id", "triple"}) + + def test_every_part_shares_the_known_triples_of_the_whole_graph(self) -> None: + known = self.parts[0].known_triples 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)), - ) + self.assertIs(part.known_triples, known) + heads: dict = {} + tails: dict = {} + for h, r, t in self.distinct: + heads.setdefault((r, t), []).append(h) + tails.setdefault((h, r), []).append(t) + self.assertEqual(known.heads, {k: sorted(v) for k, v in heads.items()}) + self.assertEqual(known.tails, {k: sorted(v) for k, v in tails.items()}) def test_triple_processor_is_fitted_on_training_triples_only(self) -> None: processor = self.parts[0].input_processors["triple"] @@ -643,17 +580,11 @@ def test_two_workers_give_the_same_samples(self) -> None: 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))) - ) + return sorted(tuple(s["triple"].tolist()) for s in part) for mine, theirs in zip(self.parts, parts): self.assertEqual(content(mine), content(theirs)) + self.assertEqual(mine.known_triples, theirs.known_triples) def test_without_split_warns_at_every_call(self) -> None: with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: @@ -663,6 +594,7 @@ def test_without_split_warns_at_every_call(self) -> None: with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): samples = dataset.set_task(self.task) self.assertEqual(len(samples), 11) + self.assertEqual(samples.known_triples, self.parts[0].known_triples) # Second call: the processed cache exists, the warning still fires. with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): dataset.set_task(self.task) @@ -705,30 +637,180 @@ def test_supplied_processors_are_used_as_they_are(self) -> None: for result in (parts[0], whole): self.assertEqual(result.input_processors["triple"].true_tail, everything.true_tail) + def test_known_triples_do_not_depend_on_a_supplied_processor(self) -> None: + """A supplied processor may be fitted on other triples than the + training part; the known triples are still those of the parts.""" + from pyhealth.processors import KGTripleProcessor -class TestUnseenInTrainingLog(unittest.TestCase): - """The per-part count of triples whose entity or relation training never saw.""" + one = KGTripleProcessor( + num_entities=self.dataset.num_entities, + num_relations=self.dataset.num_relations, + ) + one.fit([{"triple": self.distinct[0]}], "triple") + parts = self.dataset.set_task( + self.task, split=self.split, input_processors={"triple": one} + ) + self.assertEqual(parts[0].known_triples, self.parts[0].known_triples) - def test_counts_heads_tails_and_relations(self) -> None: - from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import base_kg_dataset + def test_known_triples_are_those_of_a_task_on_a_subgraph(self) -> None: + """A task that drops triples before numbering gets filter sets in its + own id space, built from its own samples.""" + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + class WithoutA(KGLinkPrediction): + task_name = "KGLinkPredictionWithoutA" + + def pre_filter(self, df): + import polars as pl + + keep = (pl.col("triples/head") != "a") & (pl.col("triples/tail") != "a") + return super().pre_filter(df.filter(keep)) + + # Without a, entities b..f are numbered 0..4, not 1..5 as in the graph. + task = WithoutA(num_entities=5, num_relations=2) + 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) + ) + parts = dataset.set_task(task, split=self.split) + triples = [t for part in parts for t in self.triples_of(part)] + self.assertEqual(len(triples), 5) + heads: dict = {} + for h, r, t in triples: + heads.setdefault((r, t), []).append(h) + self.assertEqual( + parts[0].known_triples.heads, {k: sorted(v) for k, v in heads.items()} + ) + model = TransE(parts[0], e_dim=4, r_dim=4) + model.eval() + for part in parts: + if len(part): + with torch.no_grad(): + out = model(triple=torch.tensor(self.triples_of(part))) + self.assertEqual(tuple(out["y_prob"].shape), (2 * len(part), 5)) + + def test_a_processor_class_in_the_schema_also_gets_known_triples(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction 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 ByClass(KGLinkPrediction): + task_name = "KGLinkPredictionByClass" - class _Train(list): - input_processors = {"triple": processor} + def __init__(self, num_entities, num_relations): + super().__init__(num_entities, num_relations) + self.input_schema = {"triple": (KGTripleProcessor, self.input_schema["triple"][1])} + + task = ByClass(self.dataset.num_entities, self.dataset.num_relations) + 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) + ) + parts = dataset.set_task(task, split=self.split) + for part in parts: + self.assertEqual(part.known_triples, self.parts[0].known_triples) + + def test_the_unseen_log_counts_against_the_training_part(self) -> None: + """With a supplied processor fitted on every triple, the counts are + still those of the training part.""" + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import base_kg_dataset + 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 self.assertLogs(base_kg_dataset.logger, level="INFO") as logs: + self.dataset.set_task( + self.task, split=self.split, input_processors={"triple": everything} + ) + logged = [line.split(":", 2)[2] for line in logs.output if ":Part " in line] + expected = [m for m in self.log if m.startswith("Part ")] + self.assertEqual(len(expected), 2) + self.assertEqual(logged, expected) - def part(*triples): - return [{"triple": torch.tensor(t)} for t in triples] + def test_every_form_of_a_kg_triple_spec_is_recognised(self) -> None: + from types import SimpleNamespace - 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 + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.base_kg_dataset import ( + _uses_kg_triple, ) + from pyhealth.processors import KGTripleProcessor + + class Subclass(KGTripleProcessor): + pass + + kwargs = {"num_entities": 2, "num_relations": 1} + cases = [ + ("kg_triple", True), + (("kg_triple", kwargs), True), + (KGTripleProcessor, True), + ((Subclass, kwargs), True), + (KGTripleProcessor(**kwargs), True), + ("tensor", False), + (("tensor", {}), False), + (None, False), + ] + for spec, expected in cases: + with self.subTest(spec=spec): + task = SimpleNamespace(input_schema={"triple": spec}) + self.assertEqual(_uses_kg_triple(task), expected) + + def test_a_processor_class_in_the_schema_without_split_warns(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + from pyhealth.processors import KGTripleProcessor + + class ByClass(KGLinkPrediction): + task_name = "KGLinkPredictionByClassUnsplit" + + def __init__(self, num_entities, num_relations): + super().__init__(num_entities, num_relations) + self.input_schema = {"triple": (KGTripleProcessor, self.input_schema["triple"][1])} + + task = ByClass(self.dataset.num_entities, self.dataset.num_relations) + 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"): + dataset.set_task(task) + + def test_a_triple_field_with_another_processor_gets_no_known_triples(self) -> None: + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + class AsTensor(KGLinkPrediction): + task_name = "KGLinkPredictionAsTensor" + + def __init__(self, num_entities, num_relations): + super().__init__(num_entities, num_relations) + self.input_schema = {"triple": "tensor"} + + task = AsTensor(self.dataset.num_entities, self.dataset.num_relations) + 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 warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + parts = dataset.set_task(task, split=self.split) + for part in parts: + self.assertFalse(hasattr(part, "known_triples")) + + +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 + + train = [(0, 0, 1), (1, 1, 2)] + held_out = [ + [(0, 0, 2), (4, 0, 1)], # seen; unseen head + [(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) + base_kg_dataset._log_unseen_in_training(train, held_out) self.assertEqual( [line.split(":", 2)[2] for line in logs.output], [ @@ -872,55 +954,83 @@ def test_eval_mode_ranks_every_entity(self) -> None: class TestFilteredRankingByHand(unittest.TestCase): """forward in eval mode, then the rank metrics, on a TransE scored by hand. - A filtered candidate must fall below the true entity, and the padding of a - kg_entity_list field must not be taken for entity 0. + A candidate that is a known true entity of the pair, in any part, must + fall below the true entity; the known triples come from the task-level + KnownTriples, not from the samples. """ - def test_filtered_scores_ranks_and_mrr(self) -> None: + # Embeddings 0, 1.2, 2.5, 4.0, one relation of 1.0, gamma 10, L1: + # the score of (h, r, t) is 10 - |h + 1 - t|. + GRAPH = ((0, 0, 2), (0, 0, 1), (2, 0, 2), (3, 0, 3), (2, 0, 3)) + RANKED = ((0, 0, 2), (2, 0, 2), (3, 0, 3)) + + def model(self, known_triples): from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE - from pyhealth.metrics import multiclass_metrics_fn - model = TransE(processor_dataset(4, 1), e_dim=1, r_dim=1, gamma=10.0) + part = processor_dataset(4, 1, known_triples=known_triples) + model = TransE(part, e_dim=1, r_dim=1, gamma=10.0) with torch.no_grad(): model.E_emb.copy_(torch.tensor([[0.0], [1.2], [2.5], [4.0]])) model.R_emb.copy_(torch.tensor([[1.0]])) model.eval() - # Both rows rank (0, r, 2); the score is 10 - |h + 1 - t|. - # candidate heads, t = 2.5: 8.5 (true), 9.7, 9.0, 7.5 - # candidate tails, h = 0: 9.0, 9.8, 8.5 (true), 7.0 - # Row 0: heads [0] and tails [2, 1], each followed by a padding 0. - # Row 1: entity 0 is a real member of both lists. - batch = { - "triple": torch.tensor([[0, 0, 2], [0, 0, 2]]), - "ground_truth_head": { - "value": torch.tensor([[0, 0], [0, 2]]), - "mask": torch.tensor([[1, 0], [1, 1]]), - }, - "ground_truth_tail": { - "value": torch.tensor([[2, 1, 0], [2, 1, 0]]), - "mask": torch.tensor([[1, 1, 0], [1, 1, 1]]), - }, - } + return model + + def test_filtered_scores_ranks_and_mrr(self) -> None: + import polars as pl + + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import KnownTriples + from pyhealth.metrics import multiclass_metrics_fn + + ids = pl.DataFrame( + list(self.GRAPH), schema=["head_id", "relation_id", "tail_id"], orient="row" + ) + model = self.model(KnownTriples.from_ids(ids)) with torch.no_grad(): - out = model(**batch) + out = model(triple=torch.tensor(self.RANKED)) # A filtered candidate is scored as the true entity minus 1. expected = torch.tensor([ - [8.5, 9.7, 9.0, 7.5], # head, row 0: nothing to filter - [8.5, 9.7, 7.5, 7.5], # head, row 1: head 2 filtered - [9.0, 7.5, 8.5, 7.0], # tail, row 0: tail 1 filtered, the padding 0 is not - [7.5, 7.5, 8.5, 7.0], # tail, row 1: tails 1 and 0 filtered + [8.5, 9.7, 7.5, 7.5], # head of (0, r, 2): head 2 is known, filtered + [8.0, 9.7, 9.0, 7.5], # head of (2, r, 2): head 0 is known, filtered + [7.0, 8.2, 8.0, 9.0], # head of (3, r, 3): head 2 is known, filtered + [9.0, 7.5, 8.5, 7.0], # tail of (0, r, 2): tail 1 is known, filtered + [6.5, 7.7, 9.0, 8.0], # tail of (2, r, 2): tail 3 is known, filtered + [5.0, 6.2, 7.5, 9.0], # tail of (3, r, 3): nothing to filter ]) self.assertTrue(torch.allclose(out["y_prob"], expected, atol=1e-5), out["y_prob"]) - self.assertEqual(out["y_true"].tolist(), [0, 0, 2, 2]) + self.assertEqual(out["y_true"].tolist(), [0, 2, 3, 2, 2, 3]) scores = multiclass_metrics_fn( out["y_true"].numpy(), out["y_prob"].numpy(), metrics=["mean_rank", "hits@n"] ) - # Filtered ranks 3, 2, 2, 1; unfiltered they would be 3, 3, 3, 3. - self.assertAlmostEqual(scores["mean_rank"], 2.0) - self.assertAlmostEqual(scores["mean_reciprocal_rank"], (1 / 3 + 1 / 2 + 1 / 2 + 1) / 4) - self.assertAlmostEqual(scores["HITS@1"], 0.25) + # Filtered ranks 2, 2, 1, 2, 1, 1; unfiltered they would be 3, 2, 2, 3, 2, 1. + self.assertAlmostEqual(scores["mean_rank"], 1.5) + self.assertAlmostEqual(scores["mean_reciprocal_rank"], 0.75) + self.assertAlmostEqual(scores["HITS@1"], 0.5) + + def test_evaluation_without_known_triples_is_refused(self) -> None: + model = self.model(known_triples=None) + with self.assertRaisesRegex(ValueError, "known triples"), torch.no_grad(): + model(triple=torch.tensor(self.RANKED)) + + def test_a_triple_outside_the_known_triples_is_refused(self) -> None: + import polars as pl + + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import KnownTriples + + ids = pl.DataFrame( + list(self.GRAPH), schema=["head_id", "relation_id", "tail_id"], orient="row" + ) + model = self.model(KnownTriples.from_ids(ids)) + # (1, r, 2) has the known pair (r, 2) but not (1, r); (0, r, 0) has + # (0, r) but not (r, 0); (0, r, 3) has both, yet is not a known triple, + # as a triple of another split in the same id space would be. + for triple in ((1, 0, 2), (0, 0, 0), (0, 0, 3)): + with self.subTest(triple=triple), torch.no_grad(): + self.assertRaisesRegex( + ValueError, "is not among the known triples", + model, triple=torch.tensor([triple]), + ) def record_negatives(model) -> dict: @@ -957,12 +1067,7 @@ class _Train: model = TransE(dataset=_Train(), e_dim=4, r_dim=4, negative_sampling=256) model.train() - # Ground truths of the whole graph, as KGLinkPrediction gives them. - batch = { - "triple": torch.tensor([[0, 0, 1]]), - "ground_truth_head": [[0, 4, 5]], - "ground_truth_tail": [[1, 2, 3]], - } + batch = {"triple": torch.tensor([[0, 0, 1]])} drawn = record_negatives(model) np.random.seed(0) model(**batch)["loss"].backward() @@ -1018,10 +1123,8 @@ 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"]], - ) + known_heads = self.train.known_triples.heads[(r["r2"], e["d"])] + self.assertEqual(known_heads, sorted([e["a"], e["b"]])) model = TransE(dataset=self.train, e_dim=4, r_dim=4, negative_sampling=256) model.train() From d4ba5e182daec3621f5baf8d44b96fd47f3025ad Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:54 +0100 Subject: [PATCH 34/42] test(kg_emb): pin the numeric row order of pre_filter Rows are stored sorted on the string row number, so with a duplicate on rows 2 and 10 only the numeric sort in KGLinkPrediction.pre_filter keeps row 2, the first in file order. The record kept decides which part the triple falls into. The existing test puts its duplicate on rows 0 and 8, which sort the same either way, so a string sort went unnoticed. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_kg_emb.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index e88b98acf..8d0b770d9 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -797,6 +797,25 @@ def __init__(self, num_entities, num_relations): for part in parts: self.assertFalse(hasattr(part, "known_triples")) + def test_the_first_duplicate_in_numeric_row_order_is_kept(self) -> None: + # Rows are stored sorted on the string row number ("0", "1", "10", + # "2", ...): with a duplicate on rows 2 and 10, only a numeric sort + # keeps row 2, the first in file order. + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + rows = [(f"e{i}", "r", f"e{i + 1}") for i in range(10)] + rows.append(rows[2]) + with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + root = Path(tmp) + config = write_kg(root, rows) + dataset = BaseKGDataset(root=tmp, config_path=config, cache_dir=root / "cache") + e = dataset.entity2id + duplicate = (e["e2"], 0, e["e3"]) + with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): + samples = dataset.set_task(KGLinkPrediction(num_entities=11, num_relations=1)) + kept = [s["patient_id"] for s in samples if tuple(s["triple"].tolist()) == duplicate] + self.assertEqual(kept, ["2"]) + class TestUnseenInTrainingLog(unittest.TestCase): """The per-part count of triples whose entity or relation training never saw.""" From f2e220ba5aef3d412dd81389424614b8d1c08693 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:55 +0100 Subject: [PATCH 35/42] docs(kg_emb): make the cross-references of the API docs resolve The KGTripleProcessor page used automodule, which documents the class under pyhealth.processors.kg_triple_processor, so the links to pyhealth.processors.KGTripleProcessor did not resolve; with the imported-members option of docs/conf.py, the page also documented the names the module imports. It now uses autoclass, as the other processor pages do, except GraphProcessor's. Three short references to KGLinkPrediction and BaseKGDataset.set_task are written in full, since those classes are documented under other modules. A nitpicky Sphinx 5.2.3 build of the kg_emb and KGTripleProcessor pages reports no unresolved reference. Co-Authored-By: Claude Opus 5.5 --- .../processors/pyhealth.processors.KGTripleProcessor.rst | 2 +- .../kg_emb/datasets/base_kg_dataset.py | 7 +++++-- .../kg_emb/tasks/kg_link_prediction.py | 5 +++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst index 5a1ec51e2..b76373709 100644 --- a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst +++ b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst @@ -20,7 +20,7 @@ not fitted. API Reference ------------- -.. automodule:: pyhealth.processors.kg_triple_processor +.. autoclass:: pyhealth.processors.KGTripleProcessor :members: :undoc-members: :show-inheritance: 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 3b5d50a07..e18065fc8 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 @@ -419,7 +419,8 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): """Builds the sample dataset(s) of ``task``. This is :meth:`pyhealth.datasets.BaseDataset.set_task`, with - :class:`KGLinkPrediction` as the default task. For a task whose + :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction` + as the default 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. When the fitted ``triple`` processor is a @@ -470,7 +471,9 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): @property def default_task(self) -> BaseTask: - """:class:`KGLinkPrediction` on this graph.""" + """The default task: + :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction` + on this graph.""" from ..tasks.kg_link_prediction import KGLinkPrediction return KGLinkPrediction( 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 index 983f3f5ec..0fadf957d 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py @@ -34,8 +34,9 @@ class KGLinkPrediction(BaseTask): Filtered ranking evaluation (Bordes et al., 2013; Sun et al., 2019) removes, when a test triple is ranked, the other entities known to complete it, whatever part they belong to. Those sets are not carried by - the samples: :meth:`BaseKGDataset.set_task` builds them once, from the - triples of all the parts, as + the samples: + :meth:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset.set_task` + builds them once, from the triples of all the parts, as :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.KnownTriples`, and attaches them to every part. Training does not use them: it filters its negatives with the positives of the training triples only, held by From 064041c2b4159f394ac73f15ce7240cdbb6d2265 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:55 +0100 Subject: [PATCH 36/42] fix(kg_emb): read the held-out parts by index to build the filter sets set_task iterated the held-out parts to build KnownTriples. Under torch.distributed, litdata streams each rank only its share of a part, so the sets missed the held-out triples of the other ranks and the filtered ranks came out too high, with no error. On a part whose samples are scattered, iteration also remaps the chunk once per run of consecutive indices, so the reads grew faster than the number of triples. The parts are now read by index from the unsplit samples, at the indices of split.npz from which BaseDataset.set_task builds them; index access is neither sharded nor remapped. Without split, the stored samples are counted from the regions of the dataset, since its len() is the share of the current rank. The first part is still rebuilt from the processor that set_task fitted. A test simulates rank 0 of 2 through litdata's environment detection, on caches built beforehand, with and without split; it fails on the previous code. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/base_kg_dataset.py | 60 +++++++++++++++---- tests/core/test_kg_emb.py | 31 ++++++++++ 2 files changed, 81 insertions(+), 10 deletions(-) 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 e18065fc8..2b00cdc78 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 @@ -21,9 +21,10 @@ from dataclasses import dataclass from pathlib import Path +import numpy as np import polars as pl -from pyhealth.datasets import BaseDataset +from pyhealth.datasets import BaseDataset, SampleDataset from pyhealth.processors import PROCESSOR_REGISTRY, KGTripleProcessor from pyhealth.tasks import BaseTask @@ -219,10 +220,42 @@ def from_ids(cls, ids: pl.DataFrame) -> KnownTriples: ) -def _triples_of(part) -> list[tuple[int, int, int]]: - # Sequential iteration: on a subset, litdata's part[i] costs O(i), which - # made a pass quadratic in the size of the part. - return [tuple(sample["triple"].tolist()) for sample in part] +def _stored_triples( + parts: tuple[SampleDataset, ...], split, start: int = 0 +) -> list[list[tuple[int, int, int]]]: + """The triples of ``parts[start:]``, read by index from the stored samples. + + Iterating a part is not used. Under ``torch.distributed``, litdata + streams to each rank only its share of the part, so filter sets built + from an iteration would miss the triples of the other ranks. On a part + whose samples are scattered, iteration also remaps the chunk once per + run of consecutive indices, which made the reads grow faster than the + number of triples. Index access on the unsplit dataset is neither + sharded nor remapped. The sample indices of each part are those of + ``split.npz``, from which ``BaseDataset.set_task`` builds the parts. + + Args: + parts: The parts returned by ``BaseDataset.set_task``. + split: The ``split`` passed to it, or None. + start: Index of the first part to read. + + Returns: + One list of ``(head, relation, tail)`` per part read, in sample order. + """ + path = Path(parts[0].path) + unsplit = SampleDataset(path=str(path)) + if split is None: + # len() of a streaming dataset is the share of the current rank, so + # count the stored samples instead. + n_samples = sum(end - begin for begin, end in unsplit.region_of_interest) + indices = [np.arange(n_samples)] + else: + with np.load(path / "split.npz") as stored: + indices = [stored[f"arr_{k}"] for k in range(len(stored.files))] + return [ + [tuple(unsplit[int(i)]["triple"].tolist()) for i in part] + for part in indices[start:] + ] def _fitted_triples(processor: KGTripleProcessor) -> list[tuple[int, int, int]]: @@ -429,7 +462,11 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): or a relation absent from the training triples, and every returned part carries, as ``known_triples``, the :class:`KnownTriples` of the triples of all the parts, which the models read to filter the - ranking at evaluation. + ranking at evaluation. The triples of the held-out parts are read by + index from the stored samples rather than streamed, and so are those + of the first part when its processor was supplied, so under + ``torch.distributed`` they are not limited to the share of the + current rank. """ if task is None: task = self.default_task @@ -457,11 +494,14 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): # string: a processor class or instance in the schema gives the same # processor. if isinstance(processor, KGTripleProcessor): - # A processor fitted here holds exactly the triples of the first - # part, usually the largest, which spares reading it again; a + # The first part is usually the largest. A processor fitted here + # holds exactly its triples, which spares reading it again; a # supplied one may have been fitted on other triples. - first = _triples_of(parts[0]) if supplied else _fitted_triples(processor) - held_out = [_triples_of(part) for part in parts[1:]] + if not supplied: + first = _fitted_triples(processor) + held_out = _stored_triples(parts, split, start=1) + else: + first, *held_out = _stored_triples(parts, split) if held_out: _log_unseen_in_training(first, held_out) known = _known_triples(first + [t for part in held_out for t in part]) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 8d0b770d9..c7e1ce75b 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -652,6 +652,37 @@ def test_known_triples_do_not_depend_on_a_supplied_processor(self) -> None: ) self.assertEqual(parts[0].known_triples, self.parts[0].known_triples) + def test_known_triples_cover_every_rank_under_torch_distributed(self) -> None: + """With a process group, litdata streams each rank its share of a + part; the known triples must still hold the triples of all ranks.""" + from unittest import mock + + from litdata.utilities.env import _DistributedEnv + + 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") + # The processed caches are built before the patch: the test covers + # parts read under torch.distributed from caches built beforehand. + self.dataset.set_task(self.task, input_processors={"triple": everything}) + held_out = len(self.triples_of(self.parts[2])) + rank_0_of_2 = _DistributedEnv(world_size=2, global_rank=0, num_nodes=1) + with mock.patch.object(_DistributedEnv, "detect", return_value=rank_0_of_2): + parts = self.dataset.set_task(self.task, split=self.split) + whole = self.dataset.set_task( + self.task, input_processors={"triple": everything} + ) + streamed = sum(1 for _ in parts[2]) + # The patch takes effect: iterating a part gives a share of it. + self.assertLess(streamed, held_out) + self.assertEqual(parts[0].known_triples, self.parts[0].known_triples) + # Without split, every stored sample is read, not the rank's share. + self.assertEqual(whole.known_triples, self.parts[0].known_triples) + def test_known_triples_are_those_of_a_task_on_a_subgraph(self) -> None: """A task that drops triples before numbering gets filter sets in its own id space, built from its own samples.""" From 80b3d56f2b43d2fe5c7f76b07bbe14c1241a8eb3 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:56 +0100 Subject: [PATCH 37/42] fix(kg_emb): read the first part when its processor is a subclass set_task rebuilt the first part from the true_tail of the processor it fitted, which holds exactly the triples given to fit only for KGTripleProcessor itself. With a subclass whose fit records other triples, the filter sets differed from the triples of the parts, and held-out triples were ranked against wrong sets without an error. The first part of a subclass is now read like the held-out parts. A test uses a subclass that records a single triple; it fails on the previous code. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/base_kg_dataset.py | 23 +++++++++---- tests/core/test_kg_emb.py | 33 +++++++++++++++++++ 2 files changed, 50 insertions(+), 6 deletions(-) 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 2b00cdc78..23de312c9 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 @@ -263,6 +263,17 @@ def _fitted_triples(processor: KGTripleProcessor) -> list[tuple[int, int, int]]: return [(h, r, t) for (h, r), tails in processor.true_tail.items() for t in tails] +def _holds_the_first_part(processor: KGTripleProcessor) -> bool: + """Whether a processor fitted by ``set_task`` holds exactly the triples + of the first part, so that reading that part can be skipped. + + ``KGTripleProcessor.fit`` records every triple it is given, and + ``set_task`` gives it the samples of the first part. A subclass may + record other triples, so its first part is read. + """ + return type(processor) is KGTripleProcessor + + def _known_triples(triples: list[tuple[int, int, int]]) -> KnownTriples: """The :class:`KnownTriples` of ``triples``. @@ -464,9 +475,9 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): triples of all the parts, which the models read to filter the ranking at evaluation. The triples of the held-out parts are read by index from the stored samples rather than streamed, and so are those - of the first part when its processor was supplied, so under - ``torch.distributed`` they are not limited to the share of the - current rank. + of the first part when its processor was supplied or is a subclass, + so under ``torch.distributed`` they are not limited to the share of + the current rank. """ if task is None: task = self.default_task @@ -495,9 +506,9 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): # processor. if isinstance(processor, KGTripleProcessor): # The first part is usually the largest. A processor fitted here - # holds exactly its triples, which spares reading it again; a - # supplied one may have been fitted on other triples. - if not supplied: + # by this class holds exactly its triples, which spares reading it + # again; a supplied one may have been fitted on other triples. + if not supplied and _holds_the_first_part(processor): first = _fitted_triples(processor) held_out = _stored_triples(parts, split, start=1) else: diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index c7e1ce75b..4c1757470 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -21,6 +21,7 @@ BaseKGDataset, UMLSDataset, ) +from pyhealth.processors import KGTripleProcessor # 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 @@ -430,6 +431,16 @@ def test_fitting_without_split_warns(self) -> None: TASK_SPLIT_RATIOS = (0.6, 0.2, 0.2) +class FirstTripleOnly(KGTripleProcessor): + """A kg_triple processor that records only the first triple it is given. + + Defined at module level so that set_task can pickle it with the schema. + """ + + def fit(self, samples, field): + super().fit(list(samples)[:1], field) + + class _ListHandler(logging.Handler): def __init__(self) -> None: super().__init__(level=logging.DEBUG) @@ -683,6 +694,28 @@ def test_known_triples_cover_every_rank_under_torch_distributed(self) -> None: # Without split, every stored sample is read, not the rank's share. self.assertEqual(whole.known_triples, self.parts[0].known_triples) + def test_known_triples_do_not_depend_on_a_processor_subclass(self) -> None: + """A KGTripleProcessor subclass may record other triples than those it + is given; the first part is then read rather than rebuilt from it.""" + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + class WithSubclass(KGLinkPrediction): + task_name = "KGLinkPredictionWithSubclass" + + def __init__(self, num_entities, num_relations): + super().__init__(num_entities, num_relations) + kwargs = self.input_schema["triple"][1] + self.input_schema = {"triple": (FirstTripleOnly, kwargs)} + + task = WithSubclass(self.dataset.num_entities, self.dataset.num_relations) + 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) + ) + parts = dataset.set_task(task, split=self.split) + self.assertEqual(len(parts[0].input_processors["triple"].true_tail), 1) + self.assertEqual(parts[0].known_triples, self.parts[0].known_triples) + def test_known_triples_are_those_of_a_task_on_a_subgraph(self) -> None: """A task that drops triples before numbering gets filter sets in its own id space, built from its own samples.""" From 776a5437289283bbd7a8df5dc31a59ee8ed61e9a Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:57 +0100 Subject: [PATCH 38/42] fix(kg_emb): keep the task cache key of a processor instance stable BaseDataset.set_task writes the task's input_schema into the task cache key with json.dumps(..., default=str), which takes str() of a KGTripleProcessor instance given there. str() fell back to __repr__, which reports the number of fitted pairs; set_task fits that instance, so the next call computed another key and rebuilt the cache. The key also lacked count_start and the class, so two instances that differ only in count_start, or a subclass and its base, shared one cache and the processor of the first. __str__ now gives the module and class and every attribute but the fitted dicts; __repr__ is unchanged. Co-Authored-By: Claude Opus 5.5 --- pyhealth/processors/kg_triple_processor.py | 29 ++++++++++++++++ tests/core/test_kg_emb.py | 39 ++++++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/pyhealth/processors/kg_triple_processor.py b/pyhealth/processors/kg_triple_processor.py index c9d74f5cd..7856ca84e 100644 --- a/pyhealth/processors/kg_triple_processor.py +++ b/pyhealth/processors/kg_triple_processor.py @@ -230,3 +230,32 @@ def __repr__(self) -> str: f"num_relations={self.num_relations}, " f"fitted_pairs={len(self.true_tail)})" ) + + def __str__(self) -> str: + """The class and the arguments of the processor, without its fitted + state. + + ``BaseDataset.set_task`` writes the task's ``input_schema`` into the + task cache key with ``json.dumps(..., default=str)``, which takes + ``str()`` of a processor instance given there. ``set_task`` fits + that instance itself, so its fitted state must stay out of the key, + or the next call would miss the cache. Its other attributes, such as + ``count_start`` or those a subclass adds, change what is fitted and + stay in it. + + Examples: + >>> processor = KGTripleProcessor(num_entities=3, num_relations=1) + >>> before = str(processor) + >>> processor.fit([{"triple": (0, 0, 1)}], "triple") + >>> str(processor) == before + True + >>> print(before) + pyhealth.processors.kg_triple_processor.KGTripleProcessor(count_start=4, num_entities=3, num_relations=1) + """ + arguments = ", ".join( + f"{name}={value!r}" + for name, value in sorted(vars(self).items()) + if name != "_fitted" + ) + cls = type(self) + return f"{cls.__module__}.{cls.__qualname__}({arguments})" diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 4c1757470..54a40fada 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -357,6 +357,21 @@ def test_refit_replaces_the_dicts(self) -> None: self.assertEqual(self.processor.true_head, {(1, 0): [5]}) self.assertEqual(self.processor.count, {(5, 1): 4, (0, -2): 4}) + def test_str_holds_the_arguments_and_not_the_fitted_state(self) -> None: + """str() of a processor given in a schema goes into the task cache key.""" + + class WithStep(KGTripleProcessor): + def __init__(self, step, **kwargs): + super().__init__(**kwargs) + self.step = step + + sizes = {"num_entities": 6, "num_relations": 2} + fresh = KGTripleProcessor(**sizes) + self.assertEqual(str(fresh), str(self.processor)) + self.assertNotEqual(str(KGTripleProcessor(**sizes, count_start=1)), str(fresh)) + self.assertNotEqual(str(FirstTripleOnly(**sizes)), str(fresh)) + self.assertNotEqual(str(WithStep(1, **sizes)), str(WithStep(2, **sizes))) + 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 @@ -840,6 +855,30 @@ def __init__(self, num_entities, num_relations): with self.assertWarnsRegex(UserWarning, "split=PatientSplit"): dataset.set_task(task) + def test_a_processor_instance_in_the_schema_keeps_its_task_cache(self) -> None: + """set_task fits an instance given in the schema; the task cache key, + which holds str() of that instance, must not change with that fit.""" + from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction + + class ByInstance(KGLinkPrediction): + task_name = "KGLinkPredictionByInstance" + + def __init__(self, num_entities, num_relations): + super().__init__(num_entities, num_relations) + kwargs = self.input_schema["triple"][1] + self.input_schema = {"triple": KGTripleProcessor(**kwargs)} + + task = ByInstance(self.dataset.num_entities, self.dataset.num_relations) + 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) + ) + first = dataset.set_task(task, split=self.split) + second = dataset.set_task(task, split=self.split) + task_dirs = sorted(p.name for p in Path(tmp).glob("*/tasks/*")) + self.assertEqual(len(task_dirs), 1, task_dirs) + self.assertEqual(second[0].known_triples, first[0].known_triples) + def test_a_triple_field_with_another_processor_gets_no_known_triples(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction From 075d38098a1c927da1118e4cf6d5b9a914f3e821 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:58 +0100 Subject: [PATCH 39/42] fix(kg_emb): give False for a malformed triple spec _uses_kg_triple raised IndexError on an empty tuple spec and AttributeError on a schema that is not a mapping, and gave True for tuples that SampleBuilder rejects, such as (KGTripleProcessor,), so set_task could warn about the missing split before failing on them. It now gives True only for the forms SampleBuilder accepts, and False otherwise; BaseDataset.set_task then raises its own error. Co-Authored-By: Claude Opus 5.5 --- .../kg_emb/datasets/base_kg_dataset.py | 14 +++++++++++--- tests/core/test_kg_emb.py | 9 +++++++++ 2 files changed, 20 insertions(+), 3 deletions(-) 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 23de312c9..63b43b4f8 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 @@ -18,6 +18,7 @@ import inspect import logging import warnings +from collections.abc import Mapping from dataclasses import dataclass from pathlib import Path @@ -287,11 +288,18 @@ def _known_triples(triples: list[tuple[int, int, int]]) -> KnownTriples: def _uses_kg_triple(task: BaseTask) -> bool: """Whether the task's ``triple`` field uses a ``KGTripleProcessor``, in any - form the schema accepts: a registered name, a class or an instance, alone - or with keyword arguments.""" - spec = getattr(task, "input_schema", {}).get("triple") + form the schema accepts: a registered name or a class, alone or with a + mapping of keyword arguments, or an instance. Any other spec, and a schema + that is not a mapping, give False; ``BaseDataset.set_task`` then raises + its own error on them.""" + schema = getattr(task, "input_schema", None) + spec = schema.get("triple") if isinstance(schema, Mapping) else None if isinstance(spec, tuple): + if len(spec) != 2 or not isinstance(spec[1], Mapping): + return False spec = spec[0] + if not (isinstance(spec, str) or inspect.isclass(spec)): + return False if isinstance(spec, str): spec = PROCESSOR_REGISTRY.get(spec) if not inspect.isclass(spec): diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 54a40fada..c6e038a01 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -830,11 +830,20 @@ class Subclass(KGTripleProcessor): ("tensor", False), (("tensor", {}), False), (None, False), + # Malformed: SampleBuilder rejects them, so no warning is due. + ((), False), + ((KGTripleProcessor,), False), + (("kg_triple", kwargs, 1), False), + (("kg_triple", 5), False), + ((KGTripleProcessor(**kwargs), kwargs), False), ] for spec, expected in cases: with self.subTest(spec=spec): task = SimpleNamespace(input_schema={"triple": spec}) self.assertEqual(_uses_kg_triple(task), expected) + for schema in (None, [("triple", "kg_triple")]): + with self.subTest(schema=schema): + self.assertFalse(_uses_kg_triple(SimpleNamespace(input_schema=schema))) def test_a_processor_class_in_the_schema_without_split_warns(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction From 6dc70cbb9069746f082e0ef994a3c9aa0ec747e8 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:58 +0100 Subject: [PATCH 40/42] fix(kg_emb): warn when the supplied kg_triple processor was never fitted set_task warned when it fitted the kg_triple processor without split, but not when the processor supplied in input_processors had never been fitted: training negatives are then not filtered, and subsampling weights raise KeyError. It now warns in that case too. The docstring of set_task and step 3 of docs/api/medcode.rst state both warnings, and that the first one is given at every call without split and without a supplied triple processor. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 3 ++- .../kg_emb/datasets/base_kg_dataset.py | 25 ++++++++++++++++--- tests/core/test_kg_emb.py | 10 ++++++++ 3 files changed, 33 insertions(+), 5 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 8df1bdcee..0d6cc746a 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -167,7 +167,8 @@ pipeline: entity that appears only in validation or test triples is still drawn as a training negative; such held-out triples are kept and scored, and ``set_task`` logs how many there are in each part. ``set_task`` warns when - called without ``split``. + called without ``split`` and without a supplied ``triple`` processor, and + when the supplied ``kg_triple`` processor has not been fitted. 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`` 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 63b43b4f8..da9ed9be4 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 @@ -473,9 +473,11 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): This is :meth:`pyhealth.datasets.BaseDataset.set_task`, with :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction` as the default 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. When the fitted ``triple`` processor is a + ``triple`` field uses the ``kg_triple`` processor, it also warns at + every call without ``split`` and without a supplied ``triple`` + processor, since that processor is then fitted on every triple, and + when the ``kg_triple`` processor passed in ``input_processors`` has + not been fitted. When the fitted ``triple`` processor is a :class:`~pyhealth.processors.KGTripleProcessor`, it logs, with ``split``, for each held-out part, how many triples involve an entity or a relation absent from the training triples, and every returned @@ -495,8 +497,23 @@ def set_task(self, task: BaseTask | None = None, *args, **kwargs): ).arguments split = bound.get("split") # A supplied processor is used as it is, never refitted. - supplied = "triple" in (bound.get("input_processors") or {}) + supplied_processors = bound.get("input_processors") or {} + supplied = "triple" in supplied_processors fits_triples = _uses_kg_triple(task) and not supplied + supplied_triple = supplied_processors.get("triple") + unfitted = ( + isinstance(supplied_triple, KGTripleProcessor) + and not supplied_triple.true_tail + ) + if unfitted: + warnings.warn( + "The supplied kg_triple processor has not been fitted on any " + "triple, so training negatives are not filtered and subsampling " + "weights are undefined. Fit it on the training triples, or let " + "set_task fit it with split=PatientSplit(...).", + UserWarning, + stacklevel=2, + ) if fits_triples and split is None: warnings.warn( "set_task without split= fits the kg_triple processor on every " diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index c6e038a01..222e045dd 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -845,6 +845,16 @@ class Subclass(KGTripleProcessor): with self.subTest(schema=schema): self.assertFalse(_uses_kg_triple(SimpleNamespace(input_schema=schema))) + def test_an_unfitted_supplied_processor_warns(self) -> None: + unfitted = KGTripleProcessor( + num_entities=self.dataset.num_entities, + num_relations=self.dataset.num_relations, + ) + with self.assertWarnsRegex(UserWarning, "has not been fitted on any triple"): + self.dataset.set_task( + self.task, split=self.split, input_processors={"triple": unfitted} + ) + def test_a_processor_class_in_the_schema_without_split_warns(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction from pyhealth.processors import KGTripleProcessor From 88f33f5cf4f92aa8da19848b243fb6d135957676 Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:57:59 +0100 Subject: [PATCH 41/42] test(kg_emb): cover the subsampling weights in the loss and edge cases - A loss that ignored the subsampling weights, or did not divide by their sum, passed the suite. The test now checks, on DistMult, whose two loss terms start far from zero, that weights (2, 0) give the unweighted loss of the first triple alone. - test_a_triple_field_with_another_processor_gets_no_known_triples turned every UserWarning into an error around a set_task that starts a dask cluster, and failed when port 8787 was taken, for instance by another cluster on the machine. It and test_supplied_processors_are_used_as_they_are now ignore the warnings of dask only. - set_task with split is checked to give no warning: a missing-split warning given at every call passed the suite. - The log of unseen triples is checked for its level, WARNING or INFO. - KnownTriples is checked to count once a triple given twice, as when a task's pre_filter keeps a triple in two parts. - An empty held-out part is checked through set_task: its known triples and its log line. - In test_supplied_processors_are_used_as_they_are, the supplied processor is fitted on 3 triples, so that a refit on the 6 training triples, or on all 11 without split, would show. - Two imports and three expected log lines are wrapped as ruff asks. Co-Authored-By: Claude Opus 5.5 --- tests/core/test_kg_emb.py | 134 ++++++++++++++++++++++++++++++++------ 1 file changed, 113 insertions(+), 21 deletions(-) diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 222e045dd..9d1ab1973 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -7,6 +7,7 @@ from __future__ import annotations +import contextlib import hashlib import logging import tempfile @@ -446,6 +447,19 @@ def test_fitting_without_split_warns(self) -> None: TASK_SPLIT_RATIOS = (0.6, 0.2, 0.2) +@contextlib.contextmanager +def user_warnings_as_errors_but_dask(): + """Turns every UserWarning into an error, except those of dask, which + warns on its own, e.g. "Port 8787 is already in use" when another + cluster runs on the machine.""" + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + warnings.filterwarnings( + "ignore", category=UserWarning, module=r"distributed\." + ) + yield + + class FirstTripleOnly(KGTripleProcessor): """A kg_triple processor that records only the first triple it is given. @@ -595,6 +609,23 @@ def test_unseen_entities_are_logged_per_held_out_part(self) -> None: 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_set_task_with_split_gives_no_warning(self) -> None: + with user_warnings_as_errors_but_dask(): + self.dataset.set_task(self.task, split=self.split) + + def test_an_empty_held_out_part_is_handled(self) -> None: + from pyhealth.datasets import PatientSplit + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + base_kg_dataset, + ) + + split = PatientSplit(ratios=(0.7, 0.3, 0.0), seed=3) + with self.assertLogs(base_kg_dataset.logger, level="INFO") as logs: + parts = self.dataset.set_task(self.task, split=split) + self.assertEqual(len(parts[2]), 0) + self.assertEqual(parts[2].known_triples, self.parts[0].known_triples) + self.assertTrue(any("Part 2: 0 of 0 triples" in line for line in logs.output)) + def test_two_workers_give_the_same_samples(self) -> None: with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp: dataset = BaseKGDataset( @@ -644,24 +675,20 @@ def test_a_task_for_another_graph_is_rejected(self) -> None: 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( + some = 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) + some.fit([{"triple": t} for t in self.distinct[:3]], "triple") + with user_warnings_as_errors_but_dask(): 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} + self.task, split=self.split, input_processors={"triple": some} ) - # Fitted on the 11 triples, not refitted on the 6 training ones. + whole = self.dataset.set_task(self.task, input_processors={"triple": some}) + # Fitted on 3 triples: a refit on the 6 training triples, or on all + # 11 without split, would change the dicts. for result in (parts[0], whole): - self.assertEqual(result.input_processors["triple"].true_tail, everything.true_tail) + self.assertEqual(result.input_processors["triple"].true_tail, some.true_tail) def test_known_triples_do_not_depend_on_a_supplied_processor(self) -> None: """A supplied processor may be fitted on other triples than the @@ -792,7 +819,9 @@ def __init__(self, num_entities, num_relations): def test_the_unseen_log_counts_against_the_training_part(self) -> None: """With a supplied processor fitted on every triple, the counts are still those of the training part.""" - from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import base_kg_dataset + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + base_kg_dataset, + ) from pyhealth.processors import KGTripleProcessor everything = KGTripleProcessor( @@ -913,8 +942,7 @@ def __init__(self, num_entities, num_relations): dataset = BaseKGDataset( root=str(self.root), config_path=self.root / "kg.yaml", cache_dir=Path(tmp) ) - with warnings.catch_warnings(): - warnings.simplefilter("error", UserWarning) + with user_warnings_as_errors_but_dask(): parts = dataset.set_task(task, split=self.split) for part in parts: self.assertFalse(hasattr(part, "known_triples")) @@ -943,24 +971,37 @@ 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.medcode.pretrained_embeddings.kg_emb.datasets import ( + base_kg_dataset, + ) train = [(0, 0, 1), (1, 1, 2)] held_out = [ [(0, 0, 2), (4, 0, 1)], # seen; unseen head [(0, 0, 5), (1, 2, 2), (3, 2, 4)], # unseen tail; unseen relation; both + [(1, 1, 0)], # seen ] with self.assertLogs(base_kg_dataset.logger, level="INFO") as logs: base_kg_dataset._log_unseen_in_training(train, held_out) 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.", + ( + "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." + ), + ( + "Part 3: 0 of 1 triples involve an entity absent from the " + "training triples, 0 a relation absent from them." + ), ], ) + levels = [line.split(":", 1)[0] for line in logs.output] + self.assertEqual(levels, ["WARNING", "WARNING", "INFO"]) class TestModelsOnSampleDatasets(unittest.TestCase): @@ -1076,8 +1117,35 @@ def test_subsampling_weights_come_from_the_processor(self) -> None: weights = model._subsampling_weight(batch["triple"]) expected = self.train.input_processors["triple"].subsampling_weight(batch["triple"]) self.assertTrue(torch.equal(weights, expected)) + + # They enter the loss, normalised by their sum: weighted (2, 0), the + # loss of two triples is the unweighted loss of the first alone, + # which draws the same negatives for it. DistMult starts with small + # embeddings, so both terms of its loss are far from zero; with + # TransE's margin the positive term is about 1e-6 and a missing + # normalisation of it would go unseen. + from unittest import mock + + import numpy as np + + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + + torch.manual_seed(0) + model = DistMult( + dataset=self.train, e_dim=8, r_dim=8, ns="uniform", + use_regularization=None, use_subsampling_weight=True, + ) model.train() - model(**batch)["loss"].backward() + pair = batch["triple"][:2] + two_zero = torch.tensor([2.0, 0.0]) + with mock.patch.object(model, "_subsampling_weight", return_value=two_zero): + np.random.seed(0) + weighted = model(triple=pair)["loss"] + model.use_subsampling_weight = False + np.random.seed(0) + alone = model(triple=pair[:1])["loss"] + self.assertTrue(torch.allclose(weighted, alone), (weighted, alone)) + weighted.backward() def test_eval_mode_ranks_every_entity(self) -> None: from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE @@ -1092,6 +1160,30 @@ def test_eval_mode_ranks_every_entity(self) -> None: self.assertEqual(out["y_true"].tolist(), batch["triple"][:, 0].tolist() + batch["triple"][:, 2].tolist()) +class TestKnownTriples(unittest.TestCase): + """The filter sets, built from triples that may repeat across parts.""" + + def test_a_triple_given_twice_counts_once(self) -> None: + import polars as pl + + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + KnownTriples, + base_kg_dataset, + ) + + triples = [(0, 0, 2), (0, 0, 1), (0, 0, 2)] + ids = pl.DataFrame( + triples, schema=["head_id", "relation_id", "tail_id"], orient="row" + ) + expected = KnownTriples(heads={(0, 1): [0], (0, 2): [0]}, tails={(0, 0): [1, 2]}) + self.assertEqual(KnownTriples.from_ids(ids), expected) + # As set_task builds them: the training triples, then those of each + # held-out part, here with a triple in two parts. + train, held_out = [(0, 0, 2)], [[(0, 0, 2), (0, 0, 1)]] + concatenated = train + [t for part in held_out for t in part] + self.assertEqual(base_kg_dataset._known_triples(concatenated), expected) + + class TestFilteredRankingByHand(unittest.TestCase): """forward in eval mode, then the rank metrics, on a TransE scored by hand. From a0ee607a43354bbd96960ce82b19ef408cf8385a Mon Sep 17 00:00:00 2001 From: AxelNoun <150222552+AxelNoun@users.noreply.github.com> Date: Wed, 28 Oct 2026 17:58:00 +0100 Subject: [PATCH 42/42] docs(kg_emb): correct docstrings and the migration note - base_kg_dataset: the module docstring credited Bordes et al. (2013) and Sun et al. (2019) with a "transductive setting", a term neither paper uses; it now cites the protocol they describe, a test triple ranked against every entity. KnownTriples says that subset() and deepcopy copy it. - umls: the 1.x numbering is described as the order in which pandas.unique returned the entities of the head and tail columns, rather than as their order of first appearance. - KGEBaseModel: use_regularization is None by default, "l3" works only in ComplEx and DistMult, which define l3_regularization(), and other values are not supported; the errors of construction and evaluation are listed. - KGTripleProcessor: count_start follows the reference implementation of Sun et al. (2019), not the paper, and the warnings without split are attributed to the code that emits them. - docs/api/medcode.rst: the migration note adds the new signature of BaseKGDataset, whose second position is now config_path, the third position of UMLSDataset, now config_path instead of dev, the removed raw_graph_process() hook and attributes, the read-only entity2id and relation2id, and the same description of the 1.x numbering. - docs/api/processors.rst lists the kg_triple key. Co-Authored-By: Claude Opus 5.5 --- docs/api/medcode.rst | 15 ++++++++++++--- docs/api/processors.rst | 1 + .../kg_emb/datasets/base_kg_dataset.py | 9 ++++++--- .../pretrained_embeddings/kg_emb/datasets/umls.py | 5 +++-- .../kg_emb/models/kg_base.py | 13 ++++++++++++- pyhealth/processors/kg_triple_processor.py | 8 +++++--- 6 files changed, 39 insertions(+), 12 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 0d6cc746a..5c3d25081 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -194,15 +194,24 @@ arguments (``negative_sampling`` is now a model argument, step 4), the ``refresh_cache`` argument, ``len()``, and the ``entity_num``, ``relation_num``, ``triples``, ``stat()`` and ``info()`` members of ``BaseKGDataset`` (use ``num_entities``, ``num_relations`` and -``stats()``); use the 2.0 path above. Its +``stats()``); use the 2.0 path above. ``BaseKGDataset`` now takes +``(root, config_path, dataset_name=None, cache_dir=None, num_workers=1, +dev=False)``, so ``config_path`` holds the second position, formerly +``dataset_name``; its ``raw_graph_process()`` hook and its ``samples``, +``task_name``, ``filepath`` and ``len`` attributes are gone, and +``entity2id`` and ``relation2id`` are read-only properties. +``UMLSDataset`` keeps ``root`` and ``dataset_name`` as its first two +arguments, but its third is now ``config_path``, formerly ``dev``: pass +``dev`` by keyword. Its results are not comparable with 1.x runs: 1.x filtered training negatives with, and computed subsampling weights on, every triple of the graph, held-out ones included, and made one sample per row of the file, duplicates included. ``UMLSDataset(dev=True)`` now keeps 1000 triples of ``graph.txt``, where 1.x loaded the whole of ``graph_filtered.txt``. 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. +longer follow the 1.x numbering, the order in which ``pandas.unique`` +returned the entities of the head and tail columns: map 1.x embeddings +through the ``id2entity`` saved with them. .. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset :members: diff --git a/docs/api/processors.rst b/docs/api/processors.rst index 45361db55..585b7aa91 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -364,6 +364,7 @@ Common string keys for automatic processor selection: - ``"tensor"``: For pre-processed tensors - ``"raw"``: For raw/unprocessed data - ``"graph"``: For knowledge graph subgraphs +- ``"kg_triple"``: For knowledge graph triples, with the training-graph dicts of the KG embedding models Writing Custom FeatureProcessors --------------------------------- 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 da9ed9be4..3d8e8dc18 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 @@ -9,7 +9,8 @@ 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 +link-prediction protocol of Bordes et al. (2013) and Sun et al. (2019), +where a test triple is ranked against every entity of the graph. Knowing that an entity exists does not reveal any of its edges. """ @@ -190,8 +191,10 @@ class KnownTriples: them once, from the triples of all the parts it returns, and attaches them to every part as ``part.known_triples``. Their size therefore grows with the number of triples; a copy in each sample grew with the number - of triples times the largest degree. One instance is shared by every - part and by every model built from them, so treat it as read-only. + of triples times the largest degree. ``set_task`` attaches one instance + to every part it returns, and a model built from a part keeps it, so + treat it as read-only; ``subset()`` and ``copy.deepcopy`` of a part copy + it. Training never reads them: its filters come from the ``kg_triple`` processor, fitted on the training part. diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 92616ad06..9171ae292 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -32,8 +32,9 @@ class UMLSDataset(BaseKGDataset): in that case. 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 + :class:`BaseKGDataset`), no longer as in PyHealth 1.x, which numbered + entities in the order ``pandas.unique`` returned them from the head and + tail columns: embeddings trained with 1.x must be mapped through the ``id2entity`` / ``id2relation`` saved with them. Every triple of ``graph_filtered.txt``, the file of the 1.x bucket that 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 4a423b2bc..acf214358 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -34,11 +34,22 @@ class KGEBaseModel(ABC, nn.Module): ns: negative sampling technique to use: can be "uniform", "normal" or "adv" (self-adversarial). gamma: fixed margin (only need when ns="adv"). 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. + use_regularization: None (the default here and in TransE and RotatE) + adds no penalty; "l3" adds ``l3_regularization()``, which only + ComplEx and DistMult define, and which they use by default. + Other values are not supported: they call ``regularization()``, + which needs a batch that ``forward`` does not pass. 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, 128 by default. + Raises: + TypeError: If ``dataset`` is a base dataset, or has no ``kg_triple`` + processor. + ValueError: If ``negative_sampling`` is not a positive integer. In + evaluation, ``forward`` also raises it when the dataset carried + no ``known_triples``, or for a triple that is not among them. + Examples: >>> from pyhealth.processors import KGTripleProcessor >>> class _Toy: diff --git a/pyhealth/processors/kg_triple_processor.py b/pyhealth/processors/kg_triple_processor.py index 7856ca84e..139d3268e 100644 --- a/pyhealth/processors/kg_triple_processor.py +++ b/pyhealth/processors/kg_triple_processor.py @@ -61,8 +61,9 @@ class KGTripleProcessor(FeatureProcessor): 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. + the training part only. Without ``split``, ``BaseKGDataset.set_task`` + warns at every call, and ``SampleBuilder.fit`` also logs a warning when + it fits the processor, 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 @@ -74,7 +75,8 @@ class KGTripleProcessor(FeatureProcessor): 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). + Default is 4, as in the reference implementation of Sun et + al. (2019). Raises: ValueError: If a count is not a positive integer, or if ``fit``