diff --git a/docs/api/medcode.rst b/docs/api/medcode.rst index 9038eaccd..5c3d25081 100644 --- a/docs/api/medcode.rst +++ b/docs/api/medcode.rst @@ -133,6 +133,125 @@ Medication codes: :undoc-members: :show-inheritance: +Knowledge graph embeddings +-------------------------- + +``pyhealth.medcode.pretrained_embeddings.kg_emb`` trains TransE, RotatE, +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 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 + dicts that keep known training positives out of the training negatives, + 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`` 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`` + metrics. + +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``, ``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. ``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 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: + :show-inheritance: + +.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.KnownTriples + :members: + +.. 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.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/docs/api/processors.rst b/docs/api/processors.rst index f4f2996fb..303923f9c 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -57,6 +57,7 @@ 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 **Temporal Multimodal Processors (** :class:`~pyhealth.processors.TemporalFeatureProcessor` **subclasses):** @@ -365,6 +366,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 --------------------------------- @@ -568,4 +570,5 @@ API Reference processors/pyhealth.processors.StageNetProcessor processors/pyhealth.processors.StageNetTensorProcessor processors/pyhealth.processors.GraphProcessor + processors/pyhealth.processors.KGTripleProcessor processors/pyhealth.processors.utils \ No newline at end of file diff --git a/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst new file mode 100644 index 000000000..b76373709 --- /dev/null +++ b/docs/api/processors/pyhealth.processors.KGTripleProcessor.rst @@ -0,0 +1,26 @@ +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 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 +------------- +.. autoclass:: pyhealth.processors.KGTripleProcessor + :members: + :undoc-members: + :show-inheritance: 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/examples/kg_emb_link_prediction.py b/examples/kg_emb_link_prediction.py new file mode 100644 index 000000000..c1588551c --- /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 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()`` / + ``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 mostly 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/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py index 22bda1718..b31ea4a53 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py @@ -1,4 +1,8 @@ -from .sample_kg_dataset import SampleKGDataset -from .base_kg_dataset import BaseKGDataset +from .base_kg_dataset import BaseKGDataset, KnownTriples from .umls import UMLSDataset -from .splitter import split \ No newline at end of file + +__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 c368fedc6..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 @@ -1,161 +1,560 @@ +"""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 +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. +""" + +from __future__ import annotations + +import inspect import logging -import os -from abc import ABC +import warnings +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path -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 +import numpy as np +import polars as pl +from pyhealth.datasets import BaseDataset, SampleDataset +from pyhealth.processors import PROCESSOR_REGISTRY, KGTripleProcessor +from pyhealth.tasks import BaseTask logger = logging.getLogger(__name__) -INFO_MSG = """ -dataset.triples: - Array((, , )) -""" +__all__ = [ + "BaseKGDataset", + "KnownTriples", + "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. + + 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)) + ) + + +def index_triples( + triples: pl.LazyFrame, + entities: pl.LazyFrame, + relations: pl.LazyFrame, +) -> pl.LazyFrame: + """Replaces the names of each triple with their ids. + + Args: + 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" + ) + ) + + +@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. ``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. + + 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 _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]]: + """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 _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``. + + ``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")) -class BaseKGDataset(ABC): - """Abstract base Knowledge Graph class - - This abstract class defines a uniform +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 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): + spec = type(spec) + return issubclass(spec, KGTripleProcessor) - 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 _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: - 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. - + train: The triples of the training part. + held_out: The triples of each held-out part, in part order. + """ + 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, + "Part %d: %d of %d triples involve an entity absent from the " + "training triples, %d a relation absent from them.", + k, + n_entity, + len(triples), + n_relation, + ) + + +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 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. + + Raises: + ValueError: If the config has no ``triples`` table with the + ``head``, ``relation`` and ``tail`` attributes. + + Examples: + >>> import tempfile + >>> 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, - dataset_name: Optional[str] = None, + 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 + ) -> None: + 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 - - @staticmethod - def info(): - """Prints the output format.""" - print(INFO_MSG) - - 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.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)) - return - - - def set_task( - self, - task_fn: Callable, - task_name: Optional[str] = None, - save: bool = True, - **kwargs - ) -> SampleKGDataset: - """Processes the base dataset to generate the task-specific sample dataset. - - 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. - - 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. - - Returns: - sample_dataset: the task-specific sample (Base) dataset. + self._check_config() + self._entity2id: dict[str, int] | None = None + self._relation2id: dict[str, int] | None = None - """ + 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}." + ) - if task_name is None: - self.task_name = task_fn.__name__ + def _triples_frame(self) -> pl.LazyFrame: + return triples_frame(self.global_event_df) - # 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 _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." ) - 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, - dataset_name=self.dataset_name, - task_name=self.task_name, - dev=self.dev, - entity_num=self.entity_num, - relation_num=self.relation_num, - entity2id=self.entity2id, - relation2id=self.relation2id, - **kwargs + 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 + + @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 + + @property + def id2entity(self) -> dict[int, str]: + """Entity id to name.""" + return {i: name for name, i in self.entity2id.items()} + + @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) + + 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:`~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 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 + 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. 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 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 + + 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. + 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 " + "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) + 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): + # The first part is usually the largest. A processor fitted here + # 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: + 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]) + for part in parts: + part.known_triples = known + return result - return sample_dataset + @property + def default_task(self) -> BaseTask: + """The default task: + :class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction` + on this graph.""" + from ..tasks.kg_link_prediction import KGLinkPrediction + + return KGLinkPrediction( + num_entities=self.num_entities, num_relations=self.num_relations + ) 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/sample_kg_dataset.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py deleted file mode 100644 index 59d72e888..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/sample_kg_dataset.py +++ /dev/null @@ -1,81 +0,0 @@ -from pyhealth.datasets import SampleBaseDataset - - -class SampleKGDataset(SampleBaseDataset): - """Sample KG dataset class. - - This class inherits from `SampleBaseDataset` and is specifically designed - for KG datasets. - - 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. - """ - 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.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 - } - """ - 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 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 71681b16e..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ /dev/null @@ -1,54 +0,0 @@ -from itertools import chain -from typing import Optional, Tuple, Union, List - -import numpy as np -import torch - -from pyhealth.datasets import SampleBaseDataset - - -def split( - dataset: SampleBaseDataset, - 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 - ratios: a list/tuple of ratios for train / val / test - seed: random seed for shuffling the dataset - - Returns: - train_dataset, val_dataset, test_dataset: three subsets of the dataset of - type `torch.utils.data.Subset`. - - Note: - The original dataset can be accessed by `train_dataset.dataset`, - `val_dataset.dataset`, and `test_dataset.dataset`. - The seed shuffles with a private generator; NumPy's global random - state is not changed. - - Examples: - >>> train, val, test = split(dataset, [0.8, 0.1, 0.1], seed=0) # doctest: +SKIP - """ - - assert sum(ratios) == 1.0, "ratios must sum to 1.0" - index = np.arange(len(dataset)) - # A private generator: same order as seeding the global one, without - # changing the caller's global NumPy state. - rng = np.random.RandomState(seed) if seed is not None else np.random - rng.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])) - ] - 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 diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py index 22d3cbb5e..9171ae292 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py @@ -1,69 +1,181 @@ 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 +import shutil +import uuid +from pathlib import Path + +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. 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 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 + 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: - 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``). 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. + FileNotFoundError: If neither ``graph.txt`` nor + ``umls-pyhealth.tsv`` is in ``root``. + + Examples: + >>> from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + ... BaseKGDataset, + ... UMLSDataset, + ... ) + >>> issubclass(UMLSDataset, BaseKGDataset) + True + >>> dataset = UMLSDataset(root="/path/to/umls") # doctest: +SKIP + >>> dataset.num_entities # doctest: +SKIP """ - 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: - 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, + ) + + @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 ) - print("Processing UMLS knowledge graph...") - entity_list = pd.unique(graph_df[['e1', 'e2']].values.ravel('K')) - relation_list = pd.unique(graph_df['r'].values) + @staticmethod + def prepare_metadata(root: str) -> None: + """Writes ``umls-pyhealth.tsv``: ``graph.txt`` under a header line. - 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) + 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. - 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])] + Args: + root: Directory holding ``graph.txt``. + Raises: + FileNotFoundError: If ``graph.txt`` is not in ``root``. - return + 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. 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__": - 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/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..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,142 +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 pyhealth.datasets import get_dataloader -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 - -""" -This is an example to show you how to train a KG embedding model using our package +"""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. """ +import json +import pickle +from pathlib import Path -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) +import numpy as np +import torch -with open("/data/pj20/umls_kge/pretrained_model/umls_transe_new/id2relation.json", "w") as f: - json.dump(umls_ds.id2relation, f, indent=6) +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 -# check the dataset statistics after setting task -print(umls_ds.stat()) +UMLS_ROOT = "path/to/umls" +OUTPUT_DIR = Path("output/umls_transe") +SEED = 0 +SPLIT_RATIOS = (0.9, 0.05, 0.05) +NUM_WORKERS = 4 -# 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) +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) -# initialize a KGE model -model = TransE( - dataset=umls_ds, - e_dim=512, - r_dim=512, -) + umls_ds = UMLSDataset(root=UMLS_ROOT, num_workers=NUM_WORKERS) + umls_ds.stats() + print("Relations in KG:", umls_ds.relation2id) -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) + # 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) -# 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' + # 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 ) - -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 + train, val, test = umls_ds.set_task( + task, split=PatientSplit(ratios=SPLIT_RATIOS, seed=SEED) ) -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 = 'C5162542' -relation = "CHD" - -result_eid = model.inference( - head=None, - relation=umls_ds.relation2id[relation], - tail=umls_ds.entity2id[tail], - top_k=3 + 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))) -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)}") - + 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() diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py index 8fa2a443a..e35543def 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py @@ -1,28 +1,43 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset import torch +from pyhealth.datasets import SampleDataset + +from .kg_base import KGEBaseModel + 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: + >>> from pyhealth.processors import KGTripleProcessor + >>> class _Toy: + ... 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) """ def __init__( self, - dataset: SampleBaseDataset, + 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'): @@ -62,47 +77,3 @@ def calc(self, head, relation, tail, mode='pos'): score = score.sum(dim = 2) return score - - -if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset - - samples = [ - { - '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]) - }, - ] - - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader - - 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 e7563137c..2c82d4352 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py @@ -1,18 +1,29 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset import torch +from pyhealth.datasets import SampleDataset + +from .kg_base import KGEBaseModel + 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: + >>> from pyhealth.processors import KGTripleProcessor + >>> class _Toy: + ... 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: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 300, r_dim: int = 300, ns: str = "adv", @@ -20,8 +31,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'): @@ -44,47 +59,3 @@ def calc(self, head, relation, tail, mode='pos'): score = score.sum(dim=2) return score - - -if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset - - samples = [ - { - '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]) - }, - ] - - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader - - 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 2de13afe2..acf214358 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py @@ -1,27 +1,66 @@ +import numbers from abc import ABC -from pyhealth.datasets import SampleBaseDataset -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 + +from pyhealth.datasets import BaseDataset, 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. + """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``. 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). 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: + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } + >>> model = KGEBaseModel(_Toy(), e_dim=4, r_dim=4, ns="uniform") + >>> model.e_num, model.r_num, tuple(model.E_emb.shape) + (2, 1, (2, 4)) + >>> model.negative_sampling + 128 """ @property @@ -31,19 +70,46 @@ def device(self): def __init__( - self, - dataset: SampleBaseDataset, + self, + dataset: SampleDataset, 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, - mode: str = "multiclass" + use_regularization: str | None = None, + mode: str = "multiclass", + negative_sampling: int | None = None, ): - super(KGEBaseModel, self).__init__() - self.e_num = dataset.entity_num - self.r_num = dataset.relation_num + super().__init__() + if isinstance(dataset, BaseDataset): + # 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 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 + # 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 ( + 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 = int(negative_sampling) self.e_dim = e_dim self.r_dim = r_dim self.ns = ns @@ -137,9 +203,14 @@ def data_process(self, sample_batch, mode): 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: For each triple, the entity ids kept out of its head + negatives. + gt_tail: The same, for tail negatives. """ negative_sample_head = [] negative_sample_tail = [] @@ -152,12 +223,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 @@ -168,12 +237,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 @@ -191,10 +255,16 @@ 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: 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. + """ negative_sample_head = [] negative_sample_tail = [] filter_bias_head = [] @@ -267,15 +337,43 @@ def calc(self, head, relation, tail, mode='pos'): raise NotImplementedError + def _training_filters(self, positive_sample): + """The entities kept out of each triple's training negatives. + + 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 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. + """ + 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, positive_sample): + 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 data['train'][0]: + if self.training: + gt_head, gt_tail = self._training_filters(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'], - negative_sampling=data['hyperparameters'][0]['negative_sampling'] + gt_head=gt_head, + gt_tail=gt_tail, + negative_sampling=self.negative_sampling, ) negative_sample_head, negative_sample_tail = negative_sample_head.to(self.device), negative_sample_tail.to(self.device) @@ -296,11 +394,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(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 @@ -312,10 +412,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=data['triple'], - 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'] @@ -392,7 +508,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 +523,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..618800d98 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py @@ -1,28 +1,43 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset import torch +from pyhealth.datasets import SampleDataset + +from .kg_base import KGEBaseModel + 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: + >>> from pyhealth.processors import KGTripleProcessor + >>> class _Toy: + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } + >>> model = RotatE(_Toy(), e_dim=4, r_dim=2) + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, e_dim: int = 600, r_dim: int = 300, ns='adv', gamma=24.0, use_subsampling_weight: bool = False, - use_regularization: str = None, - mode: str = "multiclass" + use_regularization: str | None = None, + 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'): @@ -58,47 +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 SampleKGDataset - - samples = [ - { - '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]) - }, - ] - - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader - - 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 fbb6e68f6..546e6d8ae 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py @@ -1,29 +1,44 @@ -from.kg_base import KGEBaseModel -from pyhealth.datasets import SampleBaseDataset import torch +from pyhealth.datasets import SampleDataset + +from .kg_base import KGEBaseModel + 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: + >>> from pyhealth.processors import KGTripleProcessor + >>> class _Toy: + ... input_processors = { + ... "triple": KGTripleProcessor(num_entities=2, num_relations=1) + ... } + >>> model = TransE(_Toy(), e_dim=4, r_dim=4) + >>> tuple(model.E_emb.shape) + (2, 4) """ def __init__( self, - dataset: SampleBaseDataset, + dataset: SampleDataset, 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 + 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 @@ -47,47 +62,3 @@ def calc(self, head, relation, tail, mode='pos'): score = self.margin.item() - score return score - - -if __name__ == "__main__": - from pyhealth.datasets import SampleKGDataset - - samples = [ - { - '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]) - }, - ] - - # dataset - dataset = SampleKGDataset(samples=samples, dataset_name="test") - - # data loader - from pyhealth.datasets import get_dataloader - - 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 b86bfcc87..a4e32cc38 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py @@ -1 +1,3 @@ -from .link_prediction import link_prediction_fn +from .kg_link_prediction import KGLinkPrediction + +__all__ = ["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 new file mode 100644 index 000000000..0fadf957d --- /dev/null +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py @@ -0,0 +1,154 @@ +"""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 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:`~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 + 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 + 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) + ['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}, + ), + } + self.output_schema = {} + super().__init__() + + def pre_filter(self, df: pl.LazyFrame) -> pl.LazyFrame: + """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`` and ``tail_id``. The frame is computed here, once: + ``set_task`` reads it again for every batch of records, which + would otherwise repeat the de-duplication 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), + ) + + 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]]: + """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), + } + for event in patient.get_events(event_type=TABLE) + ] 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 69bc90477..000000000 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/link_prediction.py +++ /dev/null @@ -1,78 +0,0 @@ -import numpy as np -import torch -from tqdm import tqdm -from typing import Tuple, List -from collections import defaultdict - -def link_prediction_fn( - triples: List[Tuple] -): - - """Process a triple list for the link prediction task - - 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 - """ - 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/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index b0738b05a..3ec9d311b 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -52,6 +52,7 @@ def get_processor(name: str): ) from .time_image_processor import TimeImageProcessor from .graph_processor import GraphProcessor +from .kg_triple_processor import KGTripleProcessor from .audio_processor import AudioProcessor from .ignore_processor import IgnoreProcessor from .temporal_timeseries_processor import TemporalTimeseriesProcessor @@ -88,6 +89,7 @@ def get_processor(name: str): "TimeseriesRegressionLabelProcessor", "TimeImageProcessor", "GraphProcessor", + "KGTripleProcessor", "AudioProcessor", "TupleTimeTextProcessor", "CehrProcessor", diff --git a/pyhealth/processors/kg_triple_processor.py b/pyhealth/processors/kg_triple_processor.py new file mode 100644 index 000000000..139d3268e --- /dev/null +++ b/pyhealth/processors/kg_triple_processor.py @@ -0,0 +1,263 @@ +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 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 + triples. ``set_task(task, split=PatientSplit(...))`` fits processors on + 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 + 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 the reference implementation of 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)})" + ) + + 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 new file mode 100644 index 000000000..9d1ab1973 --- /dev/null +++ b/tests/core/test_kg_emb.py @@ -0,0 +1,1422 @@ +"""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 contextlib +import hashlib +import logging +import tempfile +import unittest +import warnings +from pathlib import Path +from typing import Any + +import torch + +from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import ( + 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 +# 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"), +] + +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 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}, known_triples=known_triples + ) + + +class TestKGEmbImports(unittest.TestCase): + """The module must import cleanly. This was the original symptom of issue #952.""" + + 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 TestBaseKGDataset(unittest.TestCase): + """BaseKGDataset loads one triples table through the BaseDataset backend.""" + + @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" + ) + + @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)) + + +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 + + +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_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/") + + 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): + """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=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( + 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=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 + self.assertTrue( + torch.allclose( + model.calc(head, relation, tail), torch.tensor(24.0), atol=1e-6 + ) + ) + + +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))) + + +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_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}, + ) + weights = self.processor.subsampling_weight(torch.tensor(self.triples)) + 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. + 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_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 + + 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.processors import KGTripleProcessor + + 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)] + 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)) + + +# 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) + + +@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. + + 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) + self.messages: list[str] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.messages.append(record.getMessage()) + + +_TASK_SPLIT: dict[str, Any] = {} + + +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 + + 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 + ) + 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: + parts = dataset.set_task(task, split=split, num_workers=1) + finally: + kg_logger.removeHandler(handler) + 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 + + +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 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))] + + 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_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: + 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"] + 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_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( + 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()) 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: + 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) + 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) + # 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.""" + some = KGTripleProcessor( + num_entities=self.dataset.num_entities, + num_relations=self.dataset.num_relations, + ) + 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": some} + ) + 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, 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 + training part; the known triples are still those of the parts.""" + from pyhealth.processors import KGTripleProcessor + + 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_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_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 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 + + class ByClass(KGLinkPrediction): + task_name = "KGLinkPredictionByClass" + + 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 test_every_form_of_a_kg_triple_spec_is_recognised(self) -> None: + from types import SimpleNamespace + + 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), + # 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_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 + + 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_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 + + 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 user_warnings_as_errors_but_dask(): + parts = dataset.set_task(task, split=self.split) + 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.""" + + 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 + [(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 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): + """The KGE models consume the SampleDatasets returned by set_task.""" + + @classmethod + def setUpClass(cls) -> None: + 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 + + 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_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 + + 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) + 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_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 + + 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["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() + 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 + + 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()) + + +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. + + 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. + """ + + # 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 + + 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() + 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(triple=torch.tensor(self.RANKED)) + + # A filtered candidate is scored as the true entity minus 1. + expected = torch.tensor([ + [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, 2, 3, 2, 2, 3]) + + scores = multiclass_metrics_fn( + out["y_true"].numpy(), out["y_prob"].numpy(), metrics=["mean_rank", "hits@n"] + ) + # 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: + """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_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: (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=256) + model.train() + batch = {"triple": torch.tensor([[0, 0, 1]])} + drawn = record_negatives(model) + np.random.seed(0) + model(**batch)["loss"].backward() + + # 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}) + + 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 + 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__) + + 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) + 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() + drawn = record_negatives(model) + np.random.seed(0) + 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 + + # 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 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 build(self, dataset, **kwargs): + from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE + + 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": t} for t in [(0, 0, 4), (1, 1, 5), (2, 2, 6)]], "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 = processor_dataset(10, 3) + self.assertEqual(self.build(dataset, negative_sampling=np.int64(6)).negative_sampling, 6) + self.assertEqual(self.build(dataset, negative_sampling=np.int32(5)).negative_sampling, 5) + with self.assertRaises(ValueError): + self.build(dataset, negative_sampling=True)