From 5029998a01150e55a1a00f50641b36a7f0d99c5b Mon Sep 17 00:00:00 2001 From: B1llsmith Date: Thu, 20 Aug 2026 01:56:24 +0300 Subject: [PATCH 1/4] Fix kg_emb broken import: SampleBaseDataset -> SampleDataset, add pandarallel dependency --- .../kg_emb/datasets/sample_kg_dataset.py | 6 +-- .../kg_emb/datasets/splitter.py | 6 +-- .../kg_emb/models/complex.py | 4 +- .../kg_emb/models/distmult.py | 4 +- .../kg_emb/models/kg_base.py | 4 +- .../kg_emb/models/rotate.py | 4 +- .../kg_emb/models/transe.py | 4 +- pyproject.toml | 3 +- tests/core/test_kg_emb.py | 48 +++++++++++++++++++ 9 files changed, 66 insertions(+), 17 deletions(-) create mode 100644 tests/core/test_kg_emb.py 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..478eba748 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,10 +1,10 @@ -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset -class SampleKGDataset(SampleBaseDataset): +class SampleKGDataset(SampleDataset): """Sample KG dataset class. - This class inherits from `SampleBaseDataset` and is specifically designed + This class inherits from `SampleDataset` and is specifically designed for KG datasets. Args: diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 559ecb7c4..58ecff1d5 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -4,18 +4,18 @@ import numpy as np import torch -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset def split( - dataset: SampleBaseDataset, + dataset: SampleDataset, ratios: Union[Tuple[float, float, float], List[float]], seed: Optional[int] = None, ): """Splits the dataset by its outermost indexed items Args: - dataset: a `SampleBaseDataset` object + dataset: a `SampleDataset` object ratios: a list/tuple of ratios for train / val / test seed: random seed for shuffling the dataset diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 8fa2a443a..7b2f32658 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -13,7 +13,7 @@ class ComplEx(KGEBaseModel): def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, 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..bce297139 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -12,7 +12,7 @@ class DistMult(KGEBaseModel): """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, 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..1be30c23d 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,5 @@ from abc import ABC -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch import time @@ -32,7 +32,7 @@ def device(self): def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, 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 df7143a6e..20536f2f5 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -13,7 +13,7 @@ class RotatE(KGEBaseModel): def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, 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 fbb6e68f6..6cf9a4f89 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset +from pyhealth.datasets import SampleDataset import torch @@ -13,7 +13,7 @@ class TransE(KGEBaseModel): def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", diff --git a/pyproject.toml b/pyproject.toml index b4626e649..d69f40284 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,8 @@ dependencies = [ "narwhals~=2.13.0", "more-itertools~=10.8.0", "einops>=0.8.0", - "linear-attention-transformer>=0.19.1", + "linear-attention-transformer>=0.19.1", + "pandarallel>=1.6.5", ] license = "BSD-3-Clause" license-files = ["LICENSE.md"] diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py new file mode 100644 index 000000000..891d1ef9e --- /dev/null +++ b/tests/core/test_kg_emb.py @@ -0,0 +1,48 @@ +import inspect +import unittest + +from pyhealth.datasets import SampleDataset + + +class TestKgEmbImports(unittest.TestCase): + def test_import_pretrained_embeddings(self): + try: + import pyhealth.medcode.pretrained_embeddings # noqa: F401 + except ImportError as e: + self.fail( + f"Importing pyhealth.medcode.pretrained_embeddings failed: {e}" + ) + + def test_model_classes_importable(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + KGEBaseModel, + RotatE, + TransE, + ) + + for cls in (KGEBaseModel, TransE, RotatE, DistMult, ComplEx): + self.assertTrue( + isinstance(cls, type), + msg=f"{cls} was not importable as a class", + ) + + def test_transe_uses_sample_dataset(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + sig = inspect.signature(TransE.__init__) + dataset_annotation = sig.parameters["dataset"].annotation + self.assertEqual( + dataset_annotation, + SampleDataset, + msg=( + "TransE.__init__'s `dataset` parameter is not annotated as " + "SampleDataset — regression check for the " + "SampleBaseDataset -> SampleDataset rename" + ), + ) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From 19b59297b101e1a5ecbccf8e9a8dc03fc8d7172f Mon Sep 17 00:00:00 2001 From: B1llsmith Date: Wed, 26 Aug 2026 02:15:08 +0300 Subject: [PATCH 2/4] Add usage examples, docs, and example script per PR contribution rules - Add >>> usage examples to docstrings of SampleKGDataset, split(), KGEBaseModel, TransE, RotatE, DistMult, ComplEx - Fix a separate, pre-existing bug found while verifying the examples: SampleKGDataset called SampleDataset.__init__ with the old in-memory list API, but SampleDataset is now litdata-backed and expects a path to a directory built by SampleBuilder.save(). Set base attributes directly instead and added __len__ (previously inherited, now needed since we no longer call the incompatible super().__init__()) - Document the kg_emb classes in docs/api/medcode.rst - Add a runnable KG embedding example to examples/medcode.py --- docs/api/medcode.rst | 37 +++++++++++++++++++ examples/medcode.py | 31 ++++++++++++++++ .../kg_emb/datasets/sample_kg_dataset.py | 37 ++++++++++++++++++- .../kg_emb/datasets/splitter.py | 15 ++++++++ .../kg_emb/models/complex.py | 20 ++++++++++ .../kg_emb/models/distmult.py | 20 ++++++++++ .../kg_emb/models/kg_base.py | 20 ++++++++++ .../kg_emb/models/rotate.py | 20 ++++++++++ .../kg_emb/models/transe.py | 20 ++++++++++ 9 files changed, 218 insertions(+), 2 deletions(-) diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index bfc4edcea..3f1443ff0 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -97,6 +97,43 @@ Medication codes: :undoc-members: :show-inheritance: + +Pretrained Knowledge Graph Embeddings: +--------------------------------------- +Utilities for training knowledge graph embedding (KGE) models over medical +code knowledge graphs, so entity/relation embeddings can be reused as +pretrained inputs for downstream tasks. + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset + :members: + :undoc-members: + :show-inheritance: + +.. 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/medcode.py b/examples/medcode.py index 5fcb27625..a7d22f0ef 100644 --- a/examples/medcode.py +++ b/examples/medcode.py @@ -11,3 +11,34 @@ atc = InnerMap.load("ATC") print("Looking up for ATC code G04CA02") print(atc.lookup("G04CA02")) + + +# Knowledge Graph Embedding (KGE) training over medical code knowledge graphs. +# See pyhealth.medcode.pretrained_embeddings.kg_emb for TransE, RotatE, +# DistMult, and ComplEx. +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + +entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} +relation2id = {"treats": 0} +samples = [ + { + "triple": (0, 0, 1), + "ground_truth_head": [0, 2], + "ground_truth_tail": [1], + "subsampling_weight": 1.0, + }, +] +kg_dataset = SampleKGDataset( + samples=samples, + dataset_name="toy_kg", + entity_num=len(entity2id), + relation_num=len(relation2id), + entity2id=entity2id, + relation2id=relation2id, +) +print("KG dataset stats:") +kg_dataset.stat() + +kge_model = TransE(kg_dataset, e_dim=32, r_dim=32) +print("Entity embedding shape:", kge_model.E_emb.shape) \ No newline at end of file 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 478eba748..338ce0562 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 @@ -22,6 +22,28 @@ class SampleKGDataset(SampleDataset): } dataset_name: the name of the dataset. Default is None. task_name: the name of the task. Default is None. + + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, + ... dataset_name="toy_kg", + ... entity_num=len(entity2id), + ... relation_num=len(relation2id), + ... entity2id=entity2id, + ... relation2id=relation2id, + ... ) + >>> dataset.stat() """ def __init__( self, @@ -36,7 +58,14 @@ def __init__( **kwargs ): - super().__init__(samples, dataset_name, task_name) + # SampleDataset's __init__ now expects a path to a directory built + # by SampleBuilder.save() (it's litdata-backed), not an in-memory + # list of samples. SampleKGDataset predates that change and still + # works with samples held directly in memory, so we set the base + # attributes ourselves instead of delegating to super().__init__(). + self.samples = samples + self.dataset_name = dataset_name + self.task_name = task_name self.dev = dev self.entity_num = entity_num self.relation_num = relation_num @@ -49,6 +78,10 @@ def __init__( if kwargs != None: self.task_spec_param = kwargs + def __len__(self): + """Returns the number of samples in the dataset.""" + return self.sample_size + def __getitem__(self, index): """ A sample is a dict containing following data: @@ -64,7 +97,7 @@ def __getitem__(self, index): } """ return self.samples[index] - + def stat(self): """Returns some statistics of the base dataset.""" lines = list() diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 58ecff1d5..a676f6599 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -26,6 +26,21 @@ def split( Note: The original dataset can be accessed by `train_dataset.dataset`, `val_dataset.dataset`, and `test_dataset.dataset`. + + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.splitter import split + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... {"triple": (0, 0, 1), "ground_truth_head": [0], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + ... {"triple": (2, 0, 1), "ground_truth_head": [2], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> train_dataset, val_dataset, test_dataset = split(dataset, ratios=[0.5, 0.5, 0.0], seed=42) """ if seed is not None: diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 7b2f32658..135ee42ba 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -9,6 +9,26 @@ class ComplEx(KGEBaseModel): 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: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ComplEx + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = ComplEx(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py index bce297139..2cfa2a19f 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -9,6 +9,26 @@ class DistMult(KGEBaseModel): 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: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import DistMult + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = DistMult(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( self, 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 1be30c23d..51c582ac8 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -22,6 +22,26 @@ 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: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import KGEBaseModel + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = KGEBaseModel(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ @property diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py index 20536f2f5..b2a02d202 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -9,6 +9,26 @@ class RotatE(KGEBaseModel): 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: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import RotatE + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = RotatE(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py index 6cf9a4f89..09ad34db3 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -9,6 +9,26 @@ class TransE(KGEBaseModel): Paper: Bordes, A., Usunier, N., Garcia-Duran, A., Weston, J. and Yakhnenko, Translating embeddings for modeling multi-relational data. NIPS 2013. + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + >>> entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + >>> relation2id = {"treats": 0} + >>> samples = [ + ... { + ... "triple": (0, 0, 1), + ... "ground_truth_head": [0, 2], + ... "ground_truth_tail": [1], + ... "subsampling_weight": 1.0, + ... }, + ... ] + >>> dataset = SampleKGDataset( + ... samples=samples, entity_num=len(entity2id), relation_num=len(relation2id), + ... entity2id=entity2id, relation2id=relation2id, + ... ) + >>> model = TransE(dataset, e_dim=32, r_dim=32) + >>> model.E_emb.shape + torch.Size([3, 32]) """ def __init__( From b892693401621b972924d2c0395f07001414d4a0 Mon Sep 17 00:00:00 2001 From: B1llsmith Date: Sat, 29 Aug 2026 16:02:58 +0300 Subject: [PATCH 3/4] Fix dataset type hints (SampleDataset -> SampleKGDataset) and add real instantiation tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per AxelNoun's review on this PR: - KGEBaseModel and its 4 subclasses (TransE, RotatE, DistMult, ComplEx), plus split(), typed their dataset parameter as SampleDataset, but they actually require SampleKGDataset-specific attributes (entity_num, relation_num). Fixed the type hints across all 6 locations. - tests/core/test_kg_emb.py previously only checked imports and a type annotation, never actually instantiating SampleKGDataset or any model — exactly the gap that let the SampleDataset.__init__ incompatibility (fixed in the previous commit) go undetected. Added tests that build a real SampleKGDataset, run split() on it, and construct all 5 model classes from it, checking embedding shapes. splitter.py uses a relative import (from .sample_kg_dataset import SampleKGDataset) rather than importing through the package __init__, since kg_emb/datasets/__init__.py imports split from splitter.py — importing back through the package would be circular. --- .../kg_emb/datasets/splitter.py | 6 +- .../kg_emb/models/complex.py | 4 +- .../kg_emb/models/distmult.py | 4 +- .../kg_emb/models/kg_base.py | 4 +- .../kg_emb/models/rotate.py | 4 +- .../kg_emb/models/transe.py | 4 +- tests/core/test_kg_emb.py | 88 +++++++++++++++++-- 7 files changed, 95 insertions(+), 19 deletions(-) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index a676f6599..6a9f17199 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -4,18 +4,18 @@ import numpy as np import torch -from pyhealth.datasets import SampleDataset +from .sample_kg_dataset import SampleKGDataset def split( - dataset: SampleDataset, + dataset: SampleKGDataset, ratios: Union[Tuple[float, float, float], List[float]], seed: Optional[int] = None, ): """Splits the dataset by its outermost indexed items Args: - dataset: a `SampleDataset` object + dataset: a `SampleKGDataset` object ratios: a list/tuple of ratios for train / val / test seed: random seed for shuffling the dataset diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 135ee42ba..ba5abb570 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset import torch @@ -33,7 +33,7 @@ class ComplEx(KGEBaseModel): def __init__( self, - dataset: SampleDataset, + 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 2cfa2a19f..7ce071bd6 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset import torch @@ -32,7 +32,7 @@ class DistMult(KGEBaseModel): """ def __init__( self, - dataset: SampleDataset, + 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 51c582ac8..94b9c4555 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,5 @@ from abc import ABC -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset import torch import time @@ -52,7 +52,7 @@ def device(self): def __init__( self, - dataset: SampleDataset, + 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 b2a02d202..1f99de1a6 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset import torch @@ -33,7 +33,7 @@ class RotatE(KGEBaseModel): def __init__( self, - dataset: SampleDataset, + 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 09ad34db3..dc37d34eb 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,5 +1,5 @@ from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset import torch @@ -33,7 +33,7 @@ class TransE(KGEBaseModel): def __init__( self, - dataset: SampleDataset, + dataset: SampleKGDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index 891d1ef9e..f5fb4e39c 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -1,7 +1,30 @@ import inspect import unittest -from pyhealth.datasets import SampleDataset +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import SampleKGDataset + + +def _toy_kg_dataset(): + """A tiny, fixed knowledge graph used across these tests: three + entities, one relation, one triple (aspirin treats headache).""" + entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + relation2id = {"treats": 0} + samples = [ + { + "triple": (0, 0, 1), + "ground_truth_head": [0, 2], + "ground_truth_tail": [1], + "subsampling_weight": 1.0, + }, + ] + return SampleKGDataset( + samples=samples, + dataset_name="toy_kg", + entity_num=len(entity2id), + relation_num=len(relation2id), + entity2id=entity2id, + relation2id=relation2id, + ) class TestKgEmbImports(unittest.TestCase): @@ -28,21 +51,74 @@ def test_model_classes_importable(self): msg=f"{cls} was not importable as a class", ) - def test_transe_uses_sample_dataset(self): + def test_transe_uses_sample_kg_dataset(self): from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE sig = inspect.signature(TransE.__init__) dataset_annotation = sig.parameters["dataset"].annotation self.assertEqual( dataset_annotation, - SampleDataset, + SampleKGDataset, msg=( - "TransE.__init__'s `dataset` parameter is not annotated as " - "SampleDataset — regression check for the " - "SampleBaseDataset -> SampleDataset rename" + "TransE.__init__'s `dataset` parameter should be typed as " + "SampleKGDataset, not the base SampleDataset — " + "KGEBaseModel reads dataset.entity_num / dataset.relation_num, " + "which only exist on SampleKGDataset" ), ) +class TestKgEmbInstantiation(unittest.TestCase): + """These tests actually build a SampleKGDataset and construct each + model from it, so a broken __init__ (e.g. an incompatible call to + SampleDataset.__init__) fails loudly here instead of only surfacing + downstream during training.""" + + def test_sample_kg_dataset_constructs_and_reports_correct_length(self): + dataset = _toy_kg_dataset() + self.assertEqual(len(dataset), 1) + self.assertEqual(dataset.entity_num, 3) + self.assertEqual(dataset.relation_num, 1) + + def test_split_works_on_a_real_sample_kg_dataset(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import split + + entity2id = {"aspirin": 0, "headache": 1, "ibuprofen": 2} + relation2id = {"treats": 0} + samples = [ + {"triple": (0, 0, 1), "ground_truth_head": [0], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + {"triple": (2, 0, 1), "ground_truth_head": [2], "ground_truth_tail": [1], "subsampling_weight": 1.0}, + ] + dataset = SampleKGDataset( + samples=samples, entity_num=3, relation_num=1, + entity2id=entity2id, relation2id=relation2id, + ) + train, val, test = split(dataset, ratios=[0.5, 0.5, 0.0], seed=42) + self.assertEqual(len(train) + len(val) + len(test), 2) + + def test_each_kge_model_constructs_from_a_real_dataset(self): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import ( + ComplEx, + DistMult, + KGEBaseModel, + RotatE, + TransE, + ) + + dataset = _toy_kg_dataset() + for cls in (KGEBaseModel, TransE, RotatE, DistMult, ComplEx): + model = cls(dataset, e_dim=32, r_dim=32) + self.assertEqual( + tuple(model.E_emb.shape), + (3, 32), + msg=f"{cls.__name__}'s entity embedding has the wrong shape", + ) + self.assertEqual( + tuple(model.R_emb.shape), + (1, 32), + msg=f"{cls.__name__}'s relation embedding has the wrong shape", + ) + + if __name__ == "__main__": unittest.main() \ No newline at end of file From bc8e24493b4a4576e0c3d465c3e3af117f7fce17 Mon Sep 17 00:00:00 2001 From: B1llsmith Date: Wed, 2 Sep 2026 01:21:04 +0300 Subject: [PATCH 4/4] Fix indentation nit in pyproject.toml and add missing trailing newlines Per AxelNoun's review: linear-attention-transformer picked up an extra 2-space indent when pandarallel was added nearby, and examples/medcode.py and tests/core/test_kg_emb.py were both missing a trailing newline. --- examples/medcode.py | 2 +- pyproject.toml | 2 +- tests/core/test_kg_emb.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/medcode.py b/examples/medcode.py index a7d22f0ef..ddb45ec05 100644 --- a/examples/medcode.py +++ b/examples/medcode.py @@ -41,4 +41,4 @@ kg_dataset.stat() kge_model = TransE(kg_dataset, e_dim=32, r_dim=32) -print("Entity embedding shape:", kge_model.E_emb.shape) \ No newline at end of file +print("Entity embedding shape:", kge_model.E_emb.shape) diff --git a/pyproject.toml b/pyproject.toml index d69f40284..9e2fb9095 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,7 +45,7 @@ dependencies = [ "narwhals~=2.13.0", "more-itertools~=10.8.0", "einops>=0.8.0", - "linear-attention-transformer>=0.19.1", + "linear-attention-transformer>=0.19.1", "pandarallel>=1.6.5", ] license = "BSD-3-Clause" diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py index f5fb4e39c..8f2266cae 100644 --- a/tests/core/test_kg_emb.py +++ b/tests/core/test_kg_emb.py @@ -121,4 +121,4 @@ def test_each_kge_model_constructs_from_a_real_dataset(self): if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main()