fix(kg_emb): port kg_emb to the 2.0 dataset pipeline

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:
 
 
 

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

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:

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``

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()

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
+
+__all__ = [
+    "BaseKGDataset",
+    "KnownTriples",
+    "UMLSDataset",
+]