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") diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 22bda1718..0dee767bd 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,4 +1,13 @@ -from .sample_kg_dataset import SampleKGDataset from .base_kg_dataset import BaseKGDataset +from .protocols import KGDatasetProtocol +from .sample_kg_dataset import SampleKGDataset +from .splitter import split from .umls import UMLSDataset -from .splitter import split \ No newline at end of file + +__all__ = [ + "BaseKGDataset", + "KGDatasetProtocol", + "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 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/protocols.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py new file mode 100644 index 000000000..e335daa91 --- /dev/null +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/protocols.py @@ -0,0 +1,34 @@ +"""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/datasets/sample_kg_dataset.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py index 59d72e888..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 @@ -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: + """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.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})" + ) diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 559ecb7c4..0ad096461 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 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__": 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 8fa2a443a..ee8837745 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, Any + 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", @@ -65,9 +80,14 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + from torch.utils.data import DataLoader - samples = [ + from pyhealth.datasets import collate_fn_dict_with_padding + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + SampleKGDataset, + ) + + 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], @@ -82,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 e7563137c..e6f38d4ff 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, Any + 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", @@ -47,9 +62,14 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + from torch.utils.data import DataLoader - samples = [ + from pyhealth.datasets import collate_fn_dict_with_padding + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + SampleKGDataset, + ) + + 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], @@ -64,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/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..1c352e72a 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, Any + 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) @@ -61,9 +76,14 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + from torch.utils.data import DataLoader - samples = [ + from pyhealth.datasets import collate_fn_dict_with_padding + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + SampleKGDataset, + ) + + 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], @@ -78,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 fbb6e68f6..5e6287a32 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, Any + 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 ): @@ -50,9 +65,14 @@ def calc(self, head, relation, tail, mode='pos'): if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset + from torch.utils.data import DataLoader - samples = [ + from pyhealth.datasets import collate_fn_dict_with_padding + from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + SampleKGDataset, + ) + + 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], @@ -67,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( diff --git a/tests/core/test_kg_emb.py b/tests/core/test_kg_emb.py new file mode 100644 index 000000000..4c664311a --- /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. 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, + 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 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) + 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 + ) + )