fix(kg_emb): migrate SampleKGDataset off the removed 1.x dataset API

docs/api/medcode.rst

@@ -117,6 +117,94 @@ 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 a ``triple`` and its ``ground_truth_head`` / ``ground_truth_tail``
+   lists, which cover the whole graph and filter the ranking at evaluation
+   (filtered setting of Bordes et al., 2013). Duplicate triples are removed
+   and counted.
+3. The ``kg_triple`` processor
+   (:class:`~pyhealth.processors.KGTripleProcessor`), fitted on the training
+   triples, gives the models the numbers of entities and relations and the
+   dicts that keep known training positives out of the training negatives,
+   as in the reference implementation of Sun et al. (2019). Validation and
+   test triples therefore never shape training. ``set_task`` warns when
+   called without ``split``.
+4. The models take ``negative_sampling`` as an argument and follow
+   ``model.train()`` / ``model.eval()``, as set by
+   :class:`~pyhealth.trainer.Trainer`; use the ``hits@n`` and ``mean_rank``
+   metrics.
+
+See ``examples/kg_emb_link_prediction.py`` for a self-contained walk-through.
+
+The PyHealth 1.x entry points, ``link_prediction_fn``,
+:class:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset`
+and :func:`~pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split`,
+still work for one release with a ``DeprecationWarning``. They keep their
+1.x behaviour, including training negatives filtered with every triple of
+the graph, held-out ones included; only the 2.0 path above avoids that. A
+model must be built from a sample dataset part, not from the
+``BaseKGDataset`` itself. Ids no longer follow the 1.x order of first
+appearance: map 1.x embeddings through the ``id2entity`` saved with them.
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.BaseKGDataset
+    :members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.UMLSDataset
+    :members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.tasks.KGLinkPrediction
+    :members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.SampleKGDataset
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autofunction:: pyhealth.medcode.pretrained_embeddings.kg_emb.datasets.split
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.KGEBaseModel
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.TransE
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.RotatE
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.DistMult
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.medcode.pretrained_embeddings.kg_emb.models.ComplEx
+    :members:
+    :undoc-members:
+    :show-inheritance:
 
 
 

docs/api/processors.rst

@@ -55,6 +55,8 @@ Available Processors
 - ``MultiHotProcessor``: For multi-hot encoding
 - ``IgnoreProcessor``: A special feature processor that marks a feature to be ignored.
 - ``GraphProcessor``: For knowledge graph subgraph extraction (e.g., GraphCare, G-BERT)
+- ``KGTripleProcessor``: For knowledge graph triples; keeps training-graph dicts for KG embedding models
+- ``KGProcessor``: For variable-length lists of knowledge graph entity ids, padded with a mask
 
 **Temporal Multimodal Processors (** :class:`~pyhealth.processors.TemporalFeatureProcessor` **subclasses):**
 
@@ -527,4 +529,6 @@ API Reference
     processors/pyhealth.processors.MultiHotProcessor
     processors/pyhealth.processors.StageNetProcessor
     processors/pyhealth.processors.StageNetTensorProcessor
-    processors/pyhealth.processors.GraphProcessor
+    processors/pyhealth.processors.GraphProcessor
+    processors/pyhealth.processors.KGTripleProcessor
+    processors/pyhealth.processors.KGProcessor

docs/api/processors/pyhealth.processors.KGProcessor.rst

@@ -0,0 +1,16 @@
+pyhealth.processors.KGProcessor
+=================================
+
+Overview
+--------
+Processor that pads a variable-length list of knowledge-graph entity ids,
+such as the ground-truth sets used for filtered link-prediction ranking,
+and returns it with a mask. Registered as ``"kg_entity_list"`` in the
+processor registry. A list longer than the fitted length is kept whole.
+
+API Reference
+-------------
+.. automodule:: pyhealth.processors.kg_processor
+    :members:
+    :undoc-members:
+    :show-inheritance:

docs/api/processors/pyhealth.processors.KGTripleProcessor.rst

@@ -0,0 +1,23 @@
+pyhealth.processors.KGTripleProcessor
+=======================================
+
+Overview
+--------
+Processor for a knowledge-graph triple ``(head, relation, tail)`` of integer
+ids, registered as ``"kg_triple"`` in the processor registry. It returns the
+triple as a ``LongTensor`` and, from the triples it is fitted on, keeps plain
+dicts used to train KG embedding models: the known heads of each
+``(relation, tail)`` pair, the known tails of each ``(head, relation)`` pair,
+and the pair frequencies behind subsampling weights.
+
+Fit it on the training triples only, with
+``set_task(task, split=PatientSplit(...))``: the training negatives and
+weights then never depend on validation or test triples. ``size()`` returns
+the number of entities, which is given (ids are global), not fitted.
+
+API Reference
+-------------
+.. automodule:: pyhealth.processors.kg_triple_processor
+    :members:
+    :undoc-members:
+    :show-inheritance:

examples/kg_emb_link_prediction.py

@@ -0,0 +1,132 @@
+"""Train and evaluate TransE on a small knowledge graph with the 2.0 pipeline.
+
+The example is self-contained: it writes a small synthetic graph with a
+learnable structure to a temporary directory instead of downloading UMLS.
+Its filtered test metrics should end well above those of a random ranking
+(mean rank about 30 for 60 entities); they say nothing about real graphs.
+Steps:
+
+1. :class:`BaseKGDataset` reads one ``triples`` table (``head``,
+   ``relation``, ``tail``) declared in a YAML config; each triple is one
+   record. :class:`UMLSDataset` does the same with a bundled config.
+2. ``set_task(KGLinkPrediction(...), split=PatientSplit(...))`` splits the
+   triples and fits the processors on the training part only. The
+   ``kg_triple`` processor keeps the training-graph dicts that filter the
+   training negatives; the ground-truth lists of each sample cover the
+   whole graph and filter the ranking at evaluation (filtered setting of
+   Bordes et al., 2013).
+3. The model reads the numbers of entities and relations from that
+   processor; the :class:`~pyhealth.trainer.Trainer` switches between
+   training and filtered evaluation through ``model.train()`` /
+   ``model.eval()``.
+
+Run from the repository root::
+
+    python examples/kg_emb_link_prediction.py
+"""
+
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import PatientSplit, get_dataloader
+from pyhealth.medcode.pretrained_embeddings.kg_emb.datasets import BaseKGDataset
+from pyhealth.medcode.pretrained_embeddings.kg_emb.models import TransE
+from pyhealth.medcode.pretrained_embeddings.kg_emb.tasks import KGLinkPrediction
+from pyhealth.trainer import Trainer
+
+SEED = 0
+N_ENTITIES = 60
+N_CLUSTERS = 6
+N_RELATIONS = 2
+N_TRIPLES = 600
+SPLIT_RATIOS = (0.8, 0.1, 0.1)
+EPOCHS = 30
+BATCH_SIZE = 32
+
+CONFIG = """version: "1.0"
+tables:
+  triples:
+    file_path: "kg.tsv"
+    patient_id: null
+    timestamp: null
+    attributes: [head, relation, tail]
+"""
+
+
+def write_synthetic_graph(root: Path) -> Path:
+    """Writes a learnable synthetic graph and its config; returns the config.
+
+    Entities fall into ``N_CLUSTERS`` ordered clusters, and relation ``rk``
+    links a random entity of cluster ``c`` to a random entity of cluster
+    ``c + k``. A translation model can represent this, so held-out triples
+    are predictable from the training ones; on uniformly random triples the
+    metrics would stay at chance level whatever the pipeline does.
+    """
+    rng = np.random.default_rng(SEED)
+    size = N_ENTITIES // N_CLUSTERS
+    lines = ["head\trelation\ttail"]
+    for _ in range(N_TRIPLES):
+        k = int(rng.integers(1, N_RELATIONS + 1))
+        c = int(rng.integers(0, N_CLUSTERS - k))
+        head = c * size + int(rng.integers(size))
+        tail = (c + k) * size + int(rng.integers(size))
+        lines.append(f"e{head}\tr{k}\te{tail}")
+    (root / "kg.tsv").write_text("\n".join(lines) + "\n")
+    (root / "kg.yaml").write_text(CONFIG)
+    return root / "kg.yaml"
+
+
+def main() -> None:
+    # Training negatives are drawn with NumPy's global generator.
+    np.random.seed(SEED)
+    torch.manual_seed(SEED)
+
+    with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+        root = Path(tmp)
+        config = write_synthetic_graph(root)
+        dataset = BaseKGDataset(
+            root=str(root), config_path=config, cache_dir=root / "cache"
+        )
+        print(f"{dataset.num_entities} entities, {dataset.num_relations} relations")
+
+        task = KGLinkPrediction(
+            num_entities=dataset.num_entities, num_relations=dataset.num_relations
+        )
+        train, val, test = dataset.set_task(
+            task, split=PatientSplit(ratios=SPLIT_RATIOS, seed=SEED)
+        )
+        print(f"train/val/test: {len(train)}/{len(val)}/{len(test)} triples")
+
+        model = TransE(dataset=train, e_dim=32, r_dim=32, negative_sampling=16)
+        trainer = Trainer(
+            model=model,
+            metrics=["hits@n", "mean_rank"],
+            device="cpu",
+            output_path=str(root / "output"),
+            exp_name="kg_emb_link_prediction",
+        )
+        trainer.train(
+            train_dataloader=get_dataloader(train, batch_size=BATCH_SIZE, shuffle=True),
+            val_dataloader=get_dataloader(val, batch_size=BATCH_SIZE),
+            epochs=EPOCHS,
+            optimizer_params={"lr": 1e-2},
+            monitor="mean_reciprocal_rank",
+            monitor_criterion="max",
+        )
+        scores = trainer.evaluate(get_dataloader(test, batch_size=BATCH_SIZE))
+        print("filtered test metrics:", scores)
+
+        # Entity names from ids: the most likely tails of (e0, r1, ?), which
+        # should 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()

pyhealth/datasets/utils.py

@@ -260,6 +260,22 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
         A dictionary where each key corresponds to a list of values from the batch.
         Tensor values are padded to the same shape.
         Tuples of (time, values) from temporal processors are collated separately.
+        A nested dict of tensors (e.g. ``{"value": ..., "mask": ...}`` from a
+        processor like ``KGProcessor``) is collated sub-key by sub-key.
+
+    Examples:
+        >>> import torch
+        >>> batch = [
+        ...     {"triple": torch.tensor([0, 0, 1]),
+        ...      "ground_truth_head": {"value": torch.tensor([0, 4]), "mask": torch.tensor([1, 1])}},
+        ...     {"triple": torch.tensor([2, 0, 3]),
+        ...      "ground_truth_head": {"value": torch.tensor([2, 0]), "mask": torch.tensor([1, 0])}},
+        ... ]
+        >>> collated = collate_fn_dict_with_padding(batch)
+        >>> collated["triple"].shape
+        torch.Size([2, 3])
+        >>> collated["ground_truth_head"]["value"].shape
+        torch.Size([2, 2])
     """
     collated = {}
     keys = batch[0].keys()
@@ -304,6 +320,31 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
         elif HAS_PYG and isinstance(values[0], PyGData):
             collated[key] = PyGBatch.from_data_list(values)
 
+        elif isinstance(values[0], dict) and all(
+            isinstance(v, torch.Tensor) for v in values[0].values()
+        ):
+            # Nested all-tensor feature dict, e.g. {"value": Tensor, "mask": Tensor}
+            # from KGProcessor. Stack each sub-key independently. KGProcessor
+            # pads to its fitted length, so this is usually a plain stack; a
+            # list longer than that length (e.g. a test filter set after
+            # fitting on the training split) is kept whole, and the batch is
+            # then padded here with 0, which gives mask 0 on the padding.
+            #
+            # Restricted to all-tensor dicts so heterogeneous per-sample
+            # dicts (e.g. a "hyperparameters" field of plain Python values)
+            # keep falling through to the generic list passthrough below,
+            # preserving their existing list-of-dicts collation shape.
+            sub_collated: dict = {}
+            for sub_key in values[0]:
+                sub_values = [v[sub_key] for v in values]
+                if all(v.shape == sub_values[0].shape for v in sub_values):
+                    sub_collated[sub_key] = torch.stack(sub_values)
+                else:
+                    sub_collated[sub_key] = pad_sequence(
+                        sub_values, batch_first=True, padding_value=0
+                    )
+            collated[key] = sub_collated
+
         elif isinstance(values[0], torch.Tensor):
             # Check if shapes are the same
             shapes = [v.shape for v in values]

pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py

@@ -1,4 +1,11 @@
-from .sample_kg_dataset import SampleKGDataset
 from .base_kg_dataset import BaseKGDataset
+from .sample_kg_dataset import SampleKGDataset
+from .splitter import split
 from .umls import UMLSDataset
-from .splitter import split
+
+__all__ = [
+    "BaseKGDataset",
+    "SampleKGDataset",
+    "UMLSDataset",
+    "split",
+]