fix(kg_emb): port kg_emb to the 2.0 dataset pipeline by AxelNoun · Pull Request #1202 · sunlabuiuc/PyHealth

AxelNoun · 2026-08-25T09:16:51Z

Closes #952. Part of #1201.

Contributor: @AxelNoun

Contribution Type: Bug fix. pyhealth.medcode.pretrained_embeddings.kg_emb fails at import in every 2.x release (ImportError: cannot import name 'SampleBaseDataset'); this PR ports the module to the 2.0 dataset pipeline.

Description:

Following the reviews of #953 and of this PR by @jhnwu3, kg_emb now uses the standard 2.0 pipeline: a BaseDataset, a BaseTask, a registered processor and streaming SampleDataset parts. Only the filtered-ranking sets sit outside it, as an in-memory KnownTriples that set_task attaches to the parts.

BaseKGDataset is a BaseDataset with one triples table declared in YAML with patient_id: null, so each triple is one record and PatientSplit splits triples. Entity and relation ids are assigned on the full graph, by sorted name. UMLSDataset is its bundled UMLS configuration; it reads the headerless graph.txt through a prepared copy and never writes to the source file.
KGLinkPrediction gives one sample per distinct triple. Duplicates are removed, keeping the first in file order, and counted.
The kg_triple processor (KGTripleProcessor), fitted on the training part, holds the dicts that keep training positives out of the training negatives, and the subsampling weights, as in the reference implementation of Sun et al. (2019).
set_task builds the filtered-ranking sets (KnownTriples) once, from the triples of all the parts, and attaches them to every part; the samples carry only the triple. The held-out parts are read by index from the stored samples, so that under torch.distributed the sets are not limited to the share of one rank; the training part is taken from the processor fitted on it, or read the same way when that processor was supplied or is a subclass. set_task warns at every call without split and without a supplied triple processor, and when the supplied kg_triple processor was never fitted. It logs, per held-out part, how many triples involve an entity or a relation absent from the training triples.
The four models read their sizes from the fitted processor, follow model.train() / model.eval() as Trainer sets them, and take negative_sampling as an argument (default 128, as in the reference implementation).

Changes from 1.x, also listed in the migration note of docs/api/medcode.rst:

1.x filtered training negatives with, and computed subsampling weights on, every triple of the graph, validation and test triples included. Both now come from the training triples only. Filtered evaluation still removes the known triples of every part, as in the standard protocol (Bordes et al., 2013).
One sample per distinct triple, where 1.x made one per row of the file.
UMLSDataset(dev=True) keeps 1000 triples of graph.txt, where 1.x loaded the whole of graph_filtered.txt.
Ids follow sorted names, not the 1.x numbering (the order in which pandas.unique returned the head and tail entities).
As agreed with @jhnwu3 on Discord on 27 Oct, the 1.x entry points are removed rather than deprecated: SampleKGDataset, kg_emb.datasets.split, link_prediction_fn, and the 1.x members and arguments of BaseKGDataset. None of them could be imported in a 2.x release.

Measured on UMLS graph_filtered.txt with PatientSplit((0.9, 0.05, 0.05), seed=0) (parts of 79,358, 4,409 and 4,409 triples) and num_workers=1, the cache (dataset and task) takes 50.4 MiB, of which 44.0 MiB of processed samples, against 1,783.9 MiB, of which 1,736.2 MiB of samples, with the filter lists stored in each sample, as in the previous version of this PR. In two interleaved runs of each version on one machine, set_task on this branch, merged with the current master, takes 42.0 to 42.5 s on a fresh cache, against 92.6 to 92.9 s for the previous version, and 2.5 s when it reuses the cache, against 5.4 to 5.5 s, with litdata 0.2.67, which the previous pin (<0.2.68) allowed. With litdata 0.2.76, which master now allows, the figures are 40.4 to 40.5 s against 105.6 to 106.0 s, and 2.5 s against 15.8 to 16.0 s. On a synthetic graph, the four models give bitwise identical scores on the validation and test parts before and after these changes.

Two caveats are documented in docs/api/medcode.rst. Ties with the true entity are ranked in the order np.argsort returns. In graph_filtered.txt, every triple has its exact inverse, so a triple-level split leaves the inverse of most test triples in the training part (3,924 of 4,409 with the split above), the leakage reported for FB15k and WN18 by Toutanova and Chen (2015) and Dettmers et al. (2018).

Files to Review:

pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py: BaseKGDataset, KnownTriples, set_task.
pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py and configs/umls.yaml: UMLS.
pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py: the task.
pyhealth/processors/kg_triple_processor.py: the kg_triple processor.
pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py: training negatives and filtered evaluation. The four model files change their dataset type, take negative_sampling, gain a doctest example, reorder their imports, and lose their 1.x __main__ demos; TransE and RotatE annotate use_regularization as str | None.
tests/core/test_kg_emb.py: reference scores for the four models, a filtered ranking computed by hand, the training and evaluation filters on a real split, and the dataset, task and processor.
examples/kg_emb_link_prediction.py: a self-contained run on a synthetic graph (filtered test MRR 0.237, mean rank 12.5 for 60 entities).
pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py: the UMLS training script, rewritten for the 2.0 pipeline (it needs graph.txt and is not run by the tests).
docs/api/medcode.rst: the pipeline, the caveats and the migration note.

Follow-ups:

Defects of the models that predate this PR and change trained models (the unweighted default L3 penalty of DistMult and ComplEx, head and tail negative scores added before the loss, ns != "adv" in TransE and RotatE): kg_emb models: default L3 penalty has no coefficient, negative scores are summed, and ns other than "adv" breaks TransE and RotatE #1346.
examples/kg_embedding.ipynb still uses the 1.x API. It already failed at its first cell on master and hard-codes private paths; it is left as it is and marked as 1.x in docs/tutorials.rst.
Fix kg_emb broken import: SampleBaseDataset -> SampleDataset, add pan… #1192 (@userjuma) also targeted fix: kg_emb broken import in PyHealth 2.0 #952; it was closed on 28 Oct without being merged.

🤖 Generated with Claude Code

SampleDataset is now a litdata.StreamingDataset that expects schema.pkl. A knowledge-graph task is an in-memory list of triples, so SampleKGDataset subclasses torch.utils.data.Dataset and exposes KGDatasetProtocol for the models. Co-authored-by: Cursor <cursoragent@cursor.com>

KGE models only need entity_num, relation_num and task_spec_param. Annotate that structural contract with KGDatasetProtocol so the model layer no longer imports the removed SampleBaseDataset. Co-authored-by: Cursor <cursoragent@cursor.com>

SampleKGDataset was never exported from pyhealth.datasets. The examples now import it from kg_emb.datasets and build a torch DataLoader with collate_fn_dict_with_padding, because get_dataloader requires litdata.StreamingDataset.set_shuffle(). Co-authored-by: Cursor <cursoragent@cursor.com>

Validate ratios with ValueError so the check survives python -O, and shuffle with a local Generator so the function no longer mutates global NumPy state. Co-authored-by: Cursor <cursoragent@cursor.com>

pandarallel was never declared in pyproject.toml or pixi.lock. The undeclared import in umls.py and base_kg_dataset.py is what produced the ModuleNotFoundError on the kg_emb import path in issue sunlabuiuc#952. initialize() ran in umls.py with no parallel_apply in kg_emb; mimicextract's parallel_apply calls are unreachable on the empty BaseEHRDataset stub. There is no lockfile entry to regenerate. Co-authored-by: Cursor <cursoragent@cursor.com>

Cover construction, split reproducibility, generic collation of variable-length ground truths, set_task on a synthetic graph, and scoring invariants. Tests instantiate SampleKGDataset so a rename-only fix cannot go green. Co-authored-by: Cursor <cursoragent@cursor.com>

Document the map-style SampleKGDataset path in the MedCode API page and add a synthetic TransE example that uses DataLoader instead of get_dataloader. Co-authored-by: Cursor <cursoragent@cursor.com>

Replace double-hyphen asides in five docstrings with periods or commas so they remain readable in a terminal and under Sphinx. Co-authored-by: Cursor <cursoragent@cursor.com>

SampleKGDataset failed to satisfy its own KGDatasetProtocol under mypy: the Protocol declared task_spec_param as a plain attribute (Mapping[str, Any] | None), which Protocol treats as read-write and therefore invariant, while SampleKGDataset declares it as dict[str, Any] | None. Models only ever read task_spec_param, so declare it as a read-only property instead: read-only Protocol members are covariant, and a concrete dict satisfies it.

Each model's __main__ block builds an untyped list of dict literals and then adds a "train" key with a bool value, which mypy rejects because it infers the dict's value type from the first literal. Annotate samples as list[dict[str, Any]] in all four demo blocks.

AxelNoun · 2026-08-29T18:44:58Z

Update: after discussing direction with @jhnwu3, the target is to keep
SampleKGDataset inside the SampleDataset hierarchy rather than decoupling from it,
with in-memory loading as an interim step for negative sampling. Moving this to draft
while I rework it, and suggesting #1192 merge first as the unblocking fix for #952.

Carrying over regardless of the base class: the split() fixes, the constructor
raising AttributeError on its own default entity2id=None, stat() returning
None, the phantom SampleKGDataset import in the models' __main__ blocks, and the
test suite. Not carrying over: KGDatasetProtocol and the standalone Dataset.

AxelNoun · 2026-08-31T15:14:58Z

Architecture & Implementation Plan Update

Just dropping a quick update here to keep a clear record of the architectural decisions we aligned on via Discord, which will guide the next commits for #1202:

1. Base Architecture (In-Memory)
As discussed, SampleKGDataset will inherit directly from InMemorySampleDataset. This keeps us correctly under the SampleDataset umbrella while cleanly sidestepping the immediate complexities of handling negative sampling in a pure streaming setup.

2. Data Shape & Serialization (The "Tensor Trick")
To handle the KG triples and the variable-length entity lists (ground_truth_head / ground_truth_tail), we are going to avoid slow JSON/Pickle serialization. I dug into the tuple_time_text_processor.py and will adapt their pattern to keep the litdata backend happy:

I will implement a custom KGProcessor.
fit() phase: Calculate the global max_length of the entity lists.
process() phase: Pad the variable-length lists and convert everything directly into pure PyTorch tensors before serialization.

Next Steps:
I will be working on drafting the KGProcessor and the base class integration tonight. I'll push the new commits once the core pipeline is running so we have concrete code to review!

…the Tensor Trick Per the architecture pivot agreed in the PR discussion: SampleKGDataset moves back under InMemorySampleDataset instead of standalone torch.utils.data.Dataset, while keeping its full public surface (entity2id/relation2id, cardinality validation, split(), stat(), dev/task_spec_param) unchanged. - Add KGProcessor ("kg_entity_list"), pre-padding ground_truth_head/tail to each field's own max length and emitting {"value", "mask"} pure tensors ahead of litdata's pickle-based caching, instead of raw variable-length Python lists. "triple" goes through the existing "tensor" processor. - Fix a correctness issue the padding introduces: pad_token_id (0) is not a reserved sentinel and can collide with a real entity id. kg_base.py's train_neg_sample_gen and test_neg_sample_filter_bias_gen now reconstruct the exact unpadded entity list via the mask (_unpad_ground_truth) before doing set-membership filtering, so negative sampling and filtered ranking stay correct whenever entity 0 is legitimate. - Add a nested-dict collation branch to collate_fn_dict_with_padding, restricted to all-tensor dicts, so {"value","mask"} pairs batch via a plain stack (shape is already uniform per field) without disturbing the existing list-of-dicts collation used by heterogeneous per-sample dicts such as "hyperparameters". - Update tests/core/test_kg_emb.py for the new shapes: triple is now a Tensor (not a tuple), ground_truth_* collate to {"value","mask"}, and set_shuffle is now expected (SampleKGDataset is intentionally back under the SampleDataset umbrella). Adds TestGroundTruthUnpadding, a regression test for the padding/entity-0 collision fix above. 24/24 tests in tests/core/test_kg_emb.py pass, plus the sample_kg_dataset.py and splitter.py doctests. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

…mple) The contribution-rules check (tools/check_pr_rules.py) failed on 5f66ff4, something I hadn't run locally before pushing — only pytest, not ruff or the repo's own gate. It enforces ruff-clean added/modified lines plus a '>>>' doctest on every new/modified top-level public class or function. - kg_processor.py: replace typing.Dict/List/Iterable with PEP 585 builtins and collections.abc.Iterable (ruff UP035/UP006, target-version py313). Entirely new file, so every line was in scope. - datasets/utils.py: add a runnable '>>>' example to collate_fn_dict_with_padding's docstring, since the function body was modified (the nested all-tensor-dict branch) and had none. Verified against the same tooling and base/head SHAs CI used (0a75f99..<new commit>): `python tools/check_pr_rules.py --base --head` now reports "All PR contribution rules passed." tests/core/test_kg_emb.py still 24/24, plus the collate_fn_dict_with_padding doctest. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

AxelNoun · 2026-08-31T21:10:26Z

CI/CD Green & Ready for Review

Everything is fully green! I've completed the implementation of the architectural plan we discussed:

Reverted SampleKGDataset to correctly inherit from InMemorySampleDataset.
Integrated the KGProcessor (using the "Tensor Trick" with independent max-length padding) to bypass the litdata serialization bottlenecks.
Ensured negative sampling correctness in kg_base.py by safely reconstructing the unpadded entity lists via attention masks prior to np.in1d filtering.
Passed all local repo linting rules (Ruff Python 3.13 typing updates and newly required doctests).

The PR is out of draft and ready for your final review whenever you have time!

…KGDataset SampleKGDataset is back in the SampleDataset hierarchy, so the structural protocol introduced for the standalone design no longer has a purpose. As announced on 29 Aug, remove protocols.py and its export, and annotate the `dataset` parameter of KGEBaseModel, TransE, RotatE, DistMult and ComplEx with SampleKGDataset, imported at runtime. There is no import cycle: the dataset modules do not import the models. The postponed-evaluation import is dropped from the five model files so that the annotation is the class itself rather than a string.

The API page, the example and the test name still described the standalone map-style Dataset of the first version, and the module docstring justified KGProcessor by litdata's pickle-based caching, which InMemorySampleDataset does not go through: it transforms every sample once at construction and keeps the result in memory. - docs/api/medcode.rst: SampleKGDataset is an InMemorySampleDataset and get_dataloader accepts it; split() returns plain lists, which is why a fold is wrapped in torch DataLoader directly. - examples/kg_emb_sample_dataset.py: same correction in the docstring. - sample_kg_dataset.py, kg_base.py: KGProcessor pads the ground-truth lists so that they collate into fixed-shape tensors. - tests: rename test_is_a_map_style_dataset, since SampleKGDataset is an IterableDataset subclass.

Since SampleKGDataset inherits from InMemorySampleDataset, it is an IterableDataset, and torch's DataLoader rejects shuffle=True for it: running any of the four model modules raised "DataLoader with IterableDataset: expected unspecified shuffle option". Use get_dataloader, which calls InMemorySampleDataset.set_shuffle() instead, as the demos did before the 2.0 migration. The local SampleKGDataset import is dropped because the module now imports it at the top.

Brings the branch up to sunlabuiuc/PyHealth master at 5568155. No textual conflicts: docs/api/medcode.rst and pyhealth/processors/__init__.py were merged automatically.

torch pads integer tensor elements to a common width, so the example prints tensor([16, 0, 0]); the docstring showed tensor([16, 0, 0]) and failed under doctest.

AxelNoun · 2026-10-07T11:07:27Z

@jhnwu3 @joshuasteier A status update on this PR, and two questions.

Current state: SampleKGDataset inherits from InMemorySampleDataset, with a registered KGProcessor (kg_entity_list) that pads ground_truth_head and ground_truth_tail and returns {value, mask}; kg_base.py removes the padding through the mask before filtering, so entity id 0 is never treated as padding. KGDatasetProtocol has been removed, as announced on 29 Aug, and KGEBaseModel and the four models are annotated with SampleKGDataset. I also fixed the models' __main__ demos, which had failed since the move to InMemorySampleDataset (torch's DataLoader rejects shuffle=True on an IterableDataset), and rewrote the description, which still described the standalone Dataset design. Current master (e94ca1d) is merged into the branch; tests/core/test_kg_emb.py (24 tests) and tools/check_pr_rules.py pass locally, and CI is green (build and contribution-rules).

A correction to my comments of 31 Aug: InMemorySampleDataset does not go through litdata, so the padding serves to collate these fields into tensors, not serialization. The docstrings now say so.

Relation to your review of #953 (8 Apr: stay in the PyHealth pipeline and move towards graph/): this PR does the first part only. It stays in the SampleDataset hierarchy and does not touch pyhealth/graph; InMemorySampleDataset is an interim base. Is that acceptable as a first step, or should the move to graph/ be part of this PR?

#1192 and this PR both close #952 and overlap on 9 files. There are two options:

A. Merge this PR alone (it also removes the SampleBaseDataset imports). #1192 would then be closed, with credit to @userjuma, if they agree.
B. Merge #1192 first, then this PR. This requires CI to run on the head of #1192 (it has only run on its first commit) and a reconciliation afterwards: #1192 adds pandarallel to pyproject.toml, whereas this PR removes the only pandarallel imports under pyhealth/, and both add tests/core/test_kg_emb.py. The six tests in the #1192 version of that file pass against this branch.

I have no preference between the two.

jhnwu3

Quick thoughts, please don't hesitate to share your thoughts too @AxelNoun

Do you think we can have everything inherit BaseDataset and SampleDataset?

Basically, it'd be nice if we could unify all of our systems into one streaming backend here for time/memory complexity reasons.

Similarly, that way we can assume SampleDataset for each downstream model still that just leverages your kg_processor representations.

AxelNoun · 2026-10-07T15:16:37Z

Quick thoughts, please don't hesitate to share your thoughts too @AxelNoun

Do you think we can have everything inherit BaseDataset and SampleDataset?

Basically, it'd be nice if we could unify all of our systems into one streaming backend here for time/memory complexity reasons.

Similarly, that way we can assume SampleDataset for each downstream model still that just leverages your kg_processor representations.

Thanks @jhnwu3, I agree with the direction. Below is where the branch stands today, a design I have checked on the current master, and two points where I would like your input before implementing it.

Where the branch stands

SampleKGDataset already subclasses SampleDataset, through InMemorySampleDataset. The variable-length ground_truth_head / ground_truth_tail fields go through the registered kg_entity_list processor (KGProcessor), so batches are built by get_dataloader and collate_fn_dict_with_padding like in any other task. The second half of your suggestion is therefore largely in place.
The first half is not. BaseKGDataset is still a standalone ABC: it loads the graph with pandas, runs link_prediction_fn in a single process and pickles the result. InMemorySampleDataset also keeps every transformed sample in memory, and its docstring says it is meant for testing and debugging. So neither the dataset nor the samples use the streaming backend yet.

Proposed design

BaseKGDataset(BaseDataset) with a single triples table, declared in YAML with patient_id: null and timestamp: null, so that each triple is one record. ClinVarDataset and MedicalTranscriptionsDataset already do this. UMLSDataset then becomes a config file plus a thin subclass.
A KGLinkPrediction(BaseTask). link_prediction_fn needs statistics over the whole graph: the filter sets gt_head[(r, t)] and gt_tail[(h, r)], and the frequencies behind subsampling_weight. _task_transform, however, splits the work by patient_id across workers. Since BaseTask.pre_filter runs on the global LazyFrame before that split, these statistics can be computed there with Polars group_by / join. __call__ then reads a single row.
set_task(..., split=PatientSplit(...)) replaces the custom splitter. With one record per triple, a patient split is a triple-level split, which is the usual transductive protocol. Processors are then fitted on the training part only (Fit processors on the training split in set_task (split=PatientSplit) #1273).
The models stop depending on attributes specific to the KG dataset and on the train / hyperparameters keys that the current splitter adds to each sample. e_num / r_num come from the fitted processors, train vs. eval from self.training, and negative_sampling becomes a model argument.

I checked points 1 to 3 on a toy graph against the current master. set_task with two workers returns a streaming SampleDataset, the filter sets match those produced by link_prediction_fn, and PatientSplit returns three disjoint parts. One detail: triple needs ("tensor", {"dtype": torch.long}), because the default tensor processor casts to float.

Two points for your input

Filtered negatives during training. gt_head / gt_tail are currently computed on all triples before the split, and train_neg_sample_gen uses them to remove false negatives during training. Validation and test triples therefore affect training. The standard protocol (Bordes et al., 2013; Sun et al., 2019) uses all known triples only for filtered evaluation, and only training triples for training. I suggest keeping the full sets for evaluation and building the training filter from the training part. Do you agree with this change of behavior, or should it go in a separate PR?
Memory used by the filter sets. KGProcessor pads each list to the largest observed length, so storage for these fields grows as N·d_max, where d_max is the largest set size. On UMLS, hub entities make d_max large. Streaming moves this cost to disk but does not remove it. An alternative is to keep a single CSR index of the graph, O(N), alongside the dataset and let the model look up the filter sets there. This is a bigger change, so I would only do it if you think the cost matters.

Scope. This touches the dataset, the task, the four models and the tests. I suggest doing it in this PR, in separate commits, and keeping link_prediction_fn and the old import paths for one release with a deprecation warning. Does that work for you, or would you prefer to merge the current state and do the migration in a follow-up PR?

A knowledge graph is now one `triples` table declared in YAML with `patient_id: null` and `timestamp: null`, so the loader numbers the rows and each triple is one record (the ClinVarDataset pattern). A PatientSplit passed to set_task therefore splits triples. - Entity and relation ids are global: assigned on the full graph, names sorted, ids 0..n-1. The Polars helpers (triples_frame, entity_vocabulary, relation_vocabulary, index_triples) will be reused by the link-prediction task. - Triples with a missing head, relation or tail raise instead of being dropped. - UMLSDataset is the bundled configs/umls.yaml plus a thin subclass. prepare_metadata copies the headerless graph.txt to umls-pyhealth.tsv under a header line; graph.txt is never modified, and the copy is made again when graph.txt is newer or differs in size. URL roots are rejected with a hint: the 1.x bucket serves graph.txt only. - Ids no longer follow the order of first appearance used in 1.x, so 1.x embeddings must be mapped through the id2entity saved with them. - The 1.x API keeps working for one release with DeprecationWarning: set_task(task_fn, ...) returns a SampleKGDataset built from the triples in file order, and triples, stat() and info() are kept. refresh_cache is accepted and ignored. The legacy path rejects split= and the other 2.0 arguments instead of storing them as hyper-parameters. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

KGTripleProcessor ("kg_triple") returns a triple of ids as a LongTensor and, from the triples it is fitted on, keeps plain dicts: true_head, true_tail and the word2vec frequencies behind subsampling weights, built as in the reference RotatE implementation of Sun et al. (2019), which computes them from training triples only. With set_task(task, split=PatientSplit(...)), SampleBuilder fits processors on the training part, so these dicts never see validation or test triples. learns_statistics is set, so set_task without split= warns. Entity and relation counts are given, not fitted, because ids are global; size() returns the number of entities. The fitted dicts sit in a holder whose str() is a digest of their content: set_task puts vars() of pre-fitted processors into a JSON cache key, which cannot hold tuple-keyed dicts. KGProcessor no longer truncates lists longer than its fitted length. Fitted on the training split, that length depends on the seed, and a truncated filter set leaves known positives among the ranked candidates. collate_fn_dict_with_padding already pads such a batch with mask 0. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- docs/api/medcode.rst describes BaseKGDataset / UMLSDataset, the KGLinkPrediction task with split=PatientSplit(...), the kg_triple processor that keeps training negatives free of held-out triples, and the deprecated 1.x entry points, including the change of id order. - examples/kg_emb_sample_dataset.py, added earlier in this PR for the in-memory SampleKGDataset, becomes examples/kg_emb_link_prediction.py: a self-contained run of the 2.0 pipeline with the Trainer on a small synthetic graph whose structure a translation model can learn, so the filtered test metrics end well above those of a random ranking. - kg_emb/examples/train_kge_model.py trains on UMLS with the new API. It no longer loads a 1.x checkpoint, whose rows follow the old id order. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Two regressions of the model refactor, for 1.x callers: - 1.x samples may carry hyperparameters["negative_sampling"], which the models read before. The model's value now only defaults to it when neither the model argument nor the dataset chose one (as in the models' __main__ demos); when both disagree, the model's value is used and a warning says so. Before this fix the batch value was ignored and 128 negatives were drawn. - negative_sampling accepts any integral number, e.g. np.int64 from a hyperparameter grid, instead of int only. Tests now check that the chosen number of negatives reaches the sampler. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…tests - A KGE model built from a BaseKGDataset itself, instead of a set_task part, took the 1.x fallback (the dataset has the 1.x count attributes) and filtered training negatives with the batches' full-graph ground truths, i.e. with held-out triples. Base datasets are now refused with a message pointing to the training part. - The leak regression tests go through forward on both sides and check that every training positive of a pair stays filtered while held-out positives are drawable; restoring the full-graph filter on either side, or keeping only the triple's own entity, now fails them. - A legacy set_task(link_prediction_fn) call emits one DeprecationWarning instead of three, two of which pointed inside pyhealth. - Deprecation notes use a warning admonition: Sphinx's deprecated directive needs a version number, which is not known yet. - The docs say that the deprecated 1.x path keeps filtering training negatives with every triple, and the example no longer claims that all top tails of (e0, r1, ?) come from the second cluster. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

jhnwu3

Will look more later.

jhnwu3 · 2026-10-11T22:25:26Z

+__all__ = ["SampleKGDataset"]
+
+
+class SampleKGDataset(InMemorySampleDataset):


Do we need this?

Conflit unique dans kg_emb/datasets/splitter.py. Depuis le point de divergence, master n'y a apporté qu'une chose : ne plus réensemencer l'état global de NumPy, via np.random.RandomState(seed) if seed is not None else np.random. Cette branche le fait déjà et plus complètement, puisque np.random.default_rng(seed) reste un générateur privé même avec seed=None, alors que le repli np.random de master retombe sur l'état global. Elle y ajoute la validation par ValueError, qui survit à python -O contrairement à l'assert de master, l'avertissement de dépréciation et les doctests. Le côté de la branche est donc conservé intégralement. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

github-actions · 2026-10-26T21:55:56Z

Test Coverage

Total: 75.15% (20205 / 26887 lines)

Modified file	Coverage	Lines
`pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/__init__.py`	100.00%	3 / 3
`pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/base_kg_dataset.py`	98.69%	151 / 153
`pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/umls.py`	92.45%	49 / 53
`pyhealth/medcode/pretrained_embeddings/kg_emb/examples/train_kge_model.py`	not imported by tests	-
`pyhealth/medcode/pretrained_embeddings/kg_emb/models/complex.py`	79.31%	23 / 29
`pyhealth/medcode/pretrained_embeddings/kg_emb/models/distmult.py`	84.21%	16 / 19
`pyhealth/medcode/pretrained_embeddings/kg_emb/models/kg_base.py`	77.69%	188 / 242
`pyhealth/medcode/pretrained_embeddings/kg_emb/models/rotate.py`	90.00%	27 / 30
`pyhealth/medcode/pretrained_embeddings/kg_emb/models/transe.py`	84.21%	16 / 19
`pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/__init__.py`	100.00%	2 / 2
`pyhealth/medcode/pretrained_embeddings/kg_emb/tasks/kg_link_prediction.py`	94.74%	36 / 38
`pyhealth/processors/__init__.py`	94.29%	33 / 35
`pyhealth/processors/kg_triple_processor.py`	94.19%	81 / 86

Workflow run · c66ead9

pyproject.toml allows numpy>=2.1,<3 since sunlabuiuc#1287, and np.in1d no longer exists in NumPy 2.4, so the first training step of any KGE model could raise AttributeError in an environment the declared range allows. assume_unique=True is dropped. The negatives are drawn with replacement, and on NumPy's sort path the flag removes repeated draws of a valid negative. No positive was ever let through; only the sampling scheme depended on the size of the filter set. The RotatE reference implementation passes the same flag. The new tests fail if np.in1d is called and if repeated draws are dropped. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

With a split, the training filters and the subsampling weights come from the training triples only, but negatives are drawn from every entity of the graph, so an entity that appears only in validation or test triples is still drawn as a training negative. Four statements said that held-out triples never shape training; they now say what holds. The migration note also lists the other changes that make 2.0 metrics not comparable with 1.x runs: deduplicated triples, subsampling weights from the training part, and dev=True keeping 1000 triples of graph.txt where 1.x loaded the whole of graph_filtered.txt. The dev parameter no longer promises "the first 1000 triples", which holds only when the table is read as a single partition. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…kage Ranks are computed with np.argsort, which follows no stated tie-breaking convention. Every triple of UMLS graph_filtered.txt has its exact inverse under one of four relation pairs, so a triple-level split leaves the inverse of most test triples in training: with PatientSplit((0.9, 0.05, 0.05), seed=0), 3924 of the 4409 test triples have it there, and an inverse look-up rule reaches a filtered tail Hits@1 of at least 0.89 under any tie-breaking. graph.txt has not been checked. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

With ns != "adv", TransE scores by +distance and RotatE raises AttributeError at its first forward. Both defects predate this PR and are left for a separate issue; the examples and the 1.x train-step test no longer advertise that mode for these two models. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Jobs started together on a fresh root, such as parallel hyper-parameter runs, wrote the prepared copy under one fixed temporary name, and on Windows one of two concurrent jobs failed with PermissionError. Each call now writes under a unique name and removes it on failure. When Windows refuses os.replace because another job is replacing the same file, the call accepts that job's copy if it is up to date. The class docstring now states the remaining limitation: a graph.txt replaced by an older file of the same size is not detected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…dels - Reference values for the four score functions, chosen so that every sign, the ComplEx conjugate and the TransE norm change the result. - Filtered ranks, MRR and Hits@1 computed by hand through forward and the rank metrics, including a padding 0 next to a real entity 0. - A train and eval step for each of the four models on the 2.0 path, and a training step with ns="uniform" for DistMult. - One shared set_task fixture instead of three identical builds; the kg_emb logger level is restored after it. - Remove test_package_imports, which cannot fail once the module has been collected. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

AxelNoun · 2026-10-27T13:47:40Z

Do we need this?

@jhnwu3 No. I propose to remove SampleKGDataset, together with the rest of the 1.x compatibility layer, in this PR. Sorry for the slow answer. The reasons, what I have fixed since your review, and three questions for you are below.

First, a correction: the PR description is out of date. It still describes the design of 7 Oct, with SampleKGDataset on InMemorySampleDataset. Since then the branch has followed the plan in my comment of 7 Oct: BaseKGDataset on BaseDataset, a KGLinkPrediction task, a kg_triple processor, and set_task(KGLinkPrediction(), split=PatientSplit(...)) returning streaming SampleDataset parts. At HEAD, SampleKGDataset is only a deprecated shim. I will rewrite the title and the description once the scope is settled.

Why the 1.x layer can go

Nothing on the 2.0 path uses it. Removing it (sample_kg_dataset.py, splitter.py, tasks/link_prediction.py, the 1.x members of BaseKGDataset, the 1.x branches of KGEBaseModel and the models' __main__ demos) deletes 805 lines of production code and adds 7. On a synthetic graph, the training and evaluation losses of the four models and the Trainer test metrics are identical before and after the removal, and none of the 1.x deprecation warnings is raised on the 2.0 path.
No 2.x user can depend on it. kg_emb does not import in any 2.x release: v2.0.0 to v2.1.0, and current master, all import SampleBaseDataset, which no module defines any more. Only users coming straight from 1.1.6 could reach the shim.
It would not preserve their results. Entity and relation ids are now assigned by sorted name rather than by order of first appearance, so a 1.x checkpoint loads without error but with shifted ids. The UMLSDataset docstring says so, but nothing checks it. The shim also keeps the 1.x training-negative filter, which is built from all triples, validation and test included.
The removal date it announces is not one you have agreed to. I found no written deprecation policy. For the 1.x classes that 2.0 dropped, the precedent is mixed: BaseEHRDataset and four others remain as stubs that only warn ("will be removed in a future version"), while SampleBaseDataset was removed outright. This PR promises removal "in the next release" in eight warnings.

If you would rather give 1.x users a pointer, a module-level __getattr__ that raises ImportError with a migration message would follow the stub precedent in about ten lines, while failing at import instead of returning an object that does nothing.

Your criterion of 7 Oct: one streaming backend

It is only partly met, and I would rather say so. The samples stream, but the filter sets used for filtered ranking do not. KGLinkPrediction.pre_filter materialises the graph and gives each triple its own copy of its filter lists, and KGProcessor pads every sample's lists to the largest degree in the graph. On UMLS graph_filtered.txt (2.9 MB, 88,176 triples, largest degree 591), the processed samples take 1,736 MiB, the padded lists are about 8 times longer than the real ones, and the 105 MiB frame returned by pre_filter is pickled once per set_task worker. The training branch does not read these per-sample lists.

I propose to keep the evaluation filter sets once, at graph level, as the RotatE reference implementation does with all_true_triples, so that samples carry only the triple. This would also remove the change this PR makes to the shared collate_fn_dict_with_padding, which at present also changes how any field whose values are dicts of tensors is batched, temporal_timeseries features included: master returns a list of dicts, this branch a dict of padded, stacked tensors. It changes what the models consume, hence question 2.

Fixed since your review

np.in1d, which NumPy 2.4 removed while pyproject.toml allows numpy>=2.1,<3 since Loosen numpy/pandas/pydantic/scikit-learn/narwhals pins so Colab installs don't downgrade #1287, is replaced by np.isin in the negative sampler (062e688). I also dropped assume_unique=True, which the RotatE reference implementation passes too: the negatives are drawn with replacement, and on NumPy's sort path the flag removed repeated draws of a valid negative. No positive was ever let through.
Four docstrings said that held-out triples never shape training. That holds for the training filters and the subsampling weights, not for the negatives, which are drawn from every entity of the graph (08fd5c0). medcode.rst now also lists what makes 2.0 metrics not comparable with 1.x: deduplicated triples, subsampling weights from the training part, and dev=True.
medcode.rst and UMLSDataset document how ties are ranked, and the UMLS inverse-relation leakage described below (f283fa5).
The TransE and RotatE examples no longer use ns="uniform" (8b09500), and UMLSDataset jobs started together on the same root no longer fail on Windows (074848d).
New tests check the four score functions against values worked out by hand, the filtered ranks and MRR returned through forward (including a padding 0 next to a real entity 0), and a training and evaluation step for each model on the 2.0 path (6d30763). Each of them fails under the corresponding mutation of the code.

Outside this PR

Three defects sit in code this PR does not change (unchanged since 2023). Fixing them changes results, so I suggest a separate issue:

DistMult and ComplEx default to an L3 penalty with no coefficient, summed over the full embedding tables (kg_base.py:499-500). On a synthetic graph of 200 entities, the loss stays at ln 2, the embeddings shrink towards zero, and filtered MRR stays at its untrained value: 0.024 for DistMult, against 0.031 untrained and 0.319 without the penalty (one seed). The reference implementation weights this term, with a default of 0.
Head and tail negative scores are added before the loss (kg_base.py:478), so the loss depends on s(h', r, t) + s(h, r, t'), which is the score of no triple.
TransE with ns != "adv" scores by +distance, and RotatE with ns != "adv" raises AttributeError.

One evaluation caveat for UMLS, now in the docs. In graph_filtered.txt every triple has its exact inverse under one of four pairs of relations. With a triple-level split (PatientSplit((0.9, 0.05, 0.05), seed=0)), 3924 of the 4409 test triples have their inverse in the training part, so a rule that only looks the inverse up reaches a filtered tail Hits@1 of at least 0.89, whatever the tie-breaking. This is the inverse-relation leakage reported for FB15k and WN18 (Toutanova and Chen, 2015; Dettmers et al., 2018). I have not measured graph.txt. Whether to keep inverse pairs on the same side of the split is a choice for users, so I did not change the default.

Questions

Remove the 1.x layer here, or keep an ImportError stub?
Graph-level filter sets in this PR, or in a follow-up?
Should I open a separate issue for the three model defects?

The branch merges cleanly with current master, and CI passes on 6d30763. The status: conflict label is stale: could you remove it?

kg_emb has not been importable in any 2.x release (sunlabuiuc#952), so no 2.x user depends on its 1.x entry points, and the maintainer chose to remove them rather than deprecate them for a release (sunlabuiuc#1202). Removed: SampleKGDataset, kg_emb.datasets.split, link_prediction_fn, the 1.x members of BaseKGDataset (set_task with a function or task_fn= and its task_name=, save= and extra keyword arguments, refresh_cache, entity_num, relation_num, triples, stat(), info()), the 1.x branches of KGEBaseModel, the __main__ demos of the four models, and their deprecation warnings and tests. KGEBaseModel now requires a kg_triple processor, and negative_sampling comes only from its argument, 128 by default. Its private helpers _training_filters and _subsampling_weight lose the batch argument that only the 1.x branch read. The tests that used 1.x code are rewritten without it: their models are built from a kg_triple processor rather than a SampleKGDataset, and the comparisons with link_prediction_fn give way to values computed by hand or by a second processor. The migration note in docs/api/medcode.rst lists the names removed here, and len(), removed earlier in this PR, and points to the 2.0 path. The 2.0 path is not changed. examples/kg_embedding.ipynb still uses the 1.x API. It already fails at its first cell on master (SampleBaseDataset no longer exists) and hard-codes private paths, so it is left as it is; docs/tutorials.rst marks it as 1.x and now lists examples/kg_emb_link_prediction.py. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Each sample carried every known head of its (relation, tail) pair and every known tail of its (head, relation) pair, padded by KGProcessor to the largest degree, so the processed cache grew with the number of triples times the largest degree: 1,736 MiB of samples on UMLS graph_filtered.txt. set_task now builds one KnownTriples from the triples of all the parts it returns and attaches it to every part as known_triples; the models read it at evaluation, and the samples carry only the triple. Built from the parts rather than from the raw graph, the sets also match a task whose pre_filter drops triples or numbers a subgraph. The training part is not read again: its triples are those the kg_triple processor was fitted on, unless that processor was supplied. The held-out parts are read once, for both the sets and the log of unseen entities, which now counts against the training part's triples even when the processor was supplied. The sets are attached, and the missing-split warning given, for a KGTripleProcessor named, given as a class or given as an instance in the schema. At evaluation, a model built without known_triples, or given a triple outside them, raises a ValueError that says so. KnownTriples is exported and documented. On graph_filtered.txt with PatientSplit((0.9, 0.05, 0.05), seed=0), the samples take 44.0 MiB instead of 1,736.2 and the whole cache 50.4 MiB instead of 1,783.9. Measured back to back in one session on one machine (num_workers=1; two runs of this version, one of the previous), set_task on a fresh cache takes 71 to 73 s instead of 119 s. A call that reuses the cache takes 25 to 26 s instead of 18 s: both versions read the 8,818 held-out samples once, and litdata spends longer mapping the chunks of the new cache. Training and filtered ranks are unchanged: on a synthetic graph, y_prob of the four models on the validation and test parts is bitwise identical before and after. KGProcessor, the unpadding in KGEBaseModel and the padded branch of the shared collate in pyhealth/datasets/utils.py, which only KGProcessor needed, are removed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Rows are stored sorted on the string row number, so with a duplicate on rows 2 and 10 only the numeric sort in KGLinkPrediction.pre_filter keeps row 2, the first in file order. The record kept decides which part the triple falls into. The existing test puts its duplicate on rows 0 and 8, which sort the same either way, so a string sort went unnoticed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The KGTripleProcessor page used automodule, which documents the class under pyhealth.processors.kg_triple_processor, so the links to pyhealth.processors.KGTripleProcessor did not resolve; with the imported-members option of docs/conf.py, the page also documented the names the module imports. It now uses autoclass, as the other processor pages do, except GraphProcessor's. Three short references to KGLinkPrediction and BaseKGDataset.set_task are written in full, since those classes are documented under other modules. A nitpicky Sphinx 5.2.3 build of the kg_emb and KGTripleProcessor pages reports no unresolved reference. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

set_task iterated the held-out parts to build KnownTriples. Under torch.distributed, litdata streams each rank only its share of a part, so the sets missed the held-out triples of the other ranks and the filtered ranks came out too high, with no error. On a part whose samples are scattered, iteration also remaps the chunk once per run of consecutive indices, so the reads grew faster than the number of triples. The parts are now read by index from the unsplit samples, at the indices of split.npz from which BaseDataset.set_task builds them; index access is neither sharded nor remapped. Without split, the stored samples are counted from the regions of the dataset, since its len() is the share of the current rank. The first part is still rebuilt from the processor that set_task fitted. A test simulates rank 0 of 2 through litdata's environment detection, on caches built beforehand, with and without split; it fails on the previous code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

set_task rebuilt the first part from the true_tail of the processor it fitted, which holds exactly the triples given to fit only for KGTripleProcessor itself. With a subclass whose fit records other triples, the filter sets differed from the triples of the parts, and held-out triples were ranked against wrong sets without an error. The first part of a subclass is now read like the held-out parts. A test uses a subclass that records a single triple; it fails on the previous code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

BaseDataset.set_task writes the task's input_schema into the task cache key with json.dumps(..., default=str), which takes str() of a KGTripleProcessor instance given there. str() fell back to __repr__, which reports the number of fitted pairs; set_task fits that instance, so the next call computed another key and rebuilt the cache. The key also lacked count_start and the class, so two instances that differ only in count_start, or a subclass and its base, shared one cache and the processor of the first. __str__ now gives the module and class and every attribute but the fitted dicts; __repr__ is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

_uses_kg_triple raised IndexError on an empty tuple spec and AttributeError on a schema that is not a mapping, and gave True for tuples that SampleBuilder rejects, such as (KGTripleProcessor,), so set_task could warn about the missing split before failing on them. It now gives True only for the forms SampleBuilder accepts, and False otherwise; BaseDataset.set_task then raises its own error. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

set_task warned when it fitted the kg_triple processor without split, but not when the processor supplied in input_processors had never been fitted: training negatives are then not filtered, and subsampling weights raise KeyError. It now warns in that case too. The docstring of set_task and step 3 of docs/api/medcode.rst state both warnings, and that the first one is given at every call without split and without a supplied triple processor. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- A loss that ignored the subsampling weights, or did not divide by their sum, passed the suite. The test now checks, on DistMult, whose two loss terms start far from zero, that weights (2, 0) give the unweighted loss of the first triple alone. - test_a_triple_field_with_another_processor_gets_no_known_triples turned every UserWarning into an error around a set_task that starts a dask cluster, and failed when port 8787 was taken, for instance by another cluster on the machine. It and test_supplied_processors_are_used_as_they_are now ignore the warnings of dask only. - set_task with split is checked to give no warning: a missing-split warning given at every call passed the suite. - The log of unseen triples is checked for its level, WARNING or INFO. - KnownTriples is checked to count once a triple given twice, as when a task's pre_filter keeps a triple in two parts. - An empty held-out part is checked through set_task: its known triples and its log line. - In test_supplied_processors_are_used_as_they_are, the supplied processor is fitted on 3 triples, so that a refit on the 6 training triples, or on all 11 without split, would show. - Two imports and three expected log lines are wrapped as ruff asks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

- base_kg_dataset: the module docstring credited Bordes et al. (2013) and Sun et al. (2019) with a "transductive setting", a term neither paper uses; it now cites the protocol they describe, a test triple ranked against every entity. KnownTriples says that subset() and deepcopy copy it. - umls: the 1.x numbering is described as the order in which pandas.unique returned the entities of the head and tail columns, rather than as their order of first appearance. - KGEBaseModel: use_regularization is None by default, "l3" works only in ComplEx and DistMult, which define l3_regularization(), and other values are not supported; the errors of construction and evaluation are listed. - KGTripleProcessor: count_start follows the reference implementation of Sun et al. (2019), not the paper, and the warnings without split are attributed to the code that emits them. - docs/api/medcode.rst: the migration note adds the new signature of BaseKGDataset, whose second position is now config_path, the third position of UMLSDataset, now config_path instead of dev, the removed raw_graph_process() hook and attributes, the read-only entity2id and relation2id, and the same description of the 1.x numbering. - docs/api/processors.rst lists the kg_triple key. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

master gained 10 commits since 0bfc345, up to 4be413d. The only conflict is in the toctree of docs/api/processors.rst, where this branch adds pyhealth.processors.KGTripleProcessor and master adds pyhealth.processors.utils after the same line; both are kept. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

AxelNoun · 2026-10-28T21:47:25Z

@jhnwu3 Following your answer on Discord, the three points are done:

The 1.x layer is removed: SampleKGDataset, kg_emb.datasets.split, link_prediction_fn, the 1.x members and arguments of BaseKGDataset and the 1.x branches of the models. The migration note in docs/api/medcode.rst lists every removed name, with its 2.0 replacement where there is one.
The filter sets are now kept once per task, in this PR. set_task builds one KnownTriples from the triples of all the parts and attaches it to every part, so the samples carry only the triple. On UMLS graph_filtered.txt with PatientSplit((0.9, 0.05, 0.05), seed=0), the cache goes from 1,783.9 MiB to 50.4 MiB. With litdata 0.2.67, set_task takes 42.0 to 42.5 s instead of 92.6 to 92.9 s on a fresh cache, and 2.5 s instead of 5.4 to 5.5 s when it reuses the cache (two interleaved runs of each version; details, and the figures with litdata 0.2.76, in the description). On a synthetic graph, the four models give bitwise identical scores before and after.
The defects of the models that predate this PR are in kg_emb models: default L3 penalty has no coefficient, negative scores are summed, and ns other than "adv" breaks TransE and RotatE #1346.

A further review of this state before pushing found more points, fixed in the seven commits just before the merge: under torch.distributed, the filter sets missed the held-out triples of the other ranks; with a KGTripleProcessor subclass, the training part was rebuilt from what its fit recorded; a processor instance in the schema changed the task cache key after its first fit, and the key ignored count_start and the class; a malformed triple spec made kg_emb raise an IndexError or give a spurious warning; a supplied processor that was never fitted gave no warning; and some tests and docstrings were incomplete.

The last commit merges the current master, whose new commits conflicted with this branch in the toctree of docs/api/processors.rst; with it, the status: conflict label no longer applies, but I do not have the rights to remove it. I have also rewritten the PR title and description to match the current state.

AxelNoun and others added 8 commits August 25, 2026 10:46

refactor(kg_emb): make split() side-effect free and reproducible

83af7af

Validate ratios with ValueError so the check survives python -O, and shuffle with a local Generator so the function no longer mutates global NumPy state. Co-authored-by: Cursor <cursoragent@cursor.com>

docs(kg_emb): satisfy contribution-rules

5803388

Document the map-style SampleKGDataset path in the MedCode API page and add a synthetic TransE example that uses DataLoader instead of get_dataloader. Co-authored-by: Cursor <cursoragent@cursor.com>

style(kg_emb): tidy docstring punctuation

1e7d1b0

Replace double-hyphen asides in five docstrings with periods or commas so they remain readable in a terminal and under Sphinx. Co-authored-by: Cursor <cursoragent@cursor.com>

This was referenced Aug 25, 2026

PyHealth 2.0: several modules still target the 1.x dataset API #1201

Open

Fix kg_emb broken import: SampleBaseDataset -> SampleDataset, add pan… #1192

Closed

AxelNoun added 2 commits August 29, 2026 16:47

AxelNoun marked this pull request as draft August 29, 2026 16:46

AxelNoun and others added 2 commits August 31, 2026 22:46

AxelNoun marked this pull request as ready for review August 31, 2026 21:10

AxelNoun added 6 commits October 6, 2026 16:05

Merge upstream/master into fix/kg-emb-2.0-migration

e2f5f82

Brings the branch up to sunlabuiuc/PyHealth master at 5568155. No textual conflicts: docs/api/medcode.rst and pyhealth/processors/__init__.py were merged automatically.

docs(processors): match KGProcessor's doctest output to torch's repr

478f368

torch pads integer tensor elements to a common width, so the example prints tensor([16, 0, 0]); the docstring showed tensor([16, 0, 0]) and failed under doctest.

Merge upstream/master into fix/kg-emb-2.0-migration

95aeca5

jhnwu3 requested changes Oct 7, 2026

View reviewed changes

AxelNoun and others added 2 commits October 7, 2026 21:17

AxelNoun and others added 3 commits October 9, 2026 22:52

userjuma approved these changes Oct 10, 2026

View reviewed changes

jhnwu3 requested changes Oct 11, 2026

View reviewed changes

Logiquo added the status: conflict There is a merge conflict label Oct 26, 2026

AxelNoun and others added 6 commits October 27, 2026 14:33

AxelNoun and others added 12 commits October 28, 2026 17:57

AxelNoun mentioned this pull request Oct 28, 2026

kg_emb models: default L3 penalty has no coefficient, negative scores are summed, and ns other than "adv" breaks TransE and RotatE #1346

Open

AxelNoun changed the title ~~fix(kg_emb): migrate SampleKGDataset off the removed 1.x dataset API~~ fix(kg_emb): port kg_emb to the 2.0 dataset pipeline Oct 28, 2026

Navigation Menu

Repository navigation

fix(kg_emb): port kg_emb to the 2.0 dataset pipeline - #1202

fix(kg_emb): port kg_emb to the 2.0 dataset pipeline#1202
AxelNoun wants to merge 46 commits into
sunlabuiuc:masterfrom
AxelNoun:fix/kg-emb-2.0-migration

AxelNoun commented Aug 25, 2026 •

edited

Loading

Uh oh!

AxelNoun commented Aug 29, 2026

Uh oh!

AxelNoun commented Aug 31, 2026 •

edited

Loading

Uh oh!

AxelNoun commented Aug 31, 2026

Uh oh!

AxelNoun commented Oct 7, 2026

Uh oh!

jhnwu3 left a comment •

edited

Loading

Uh oh!

AxelNoun commented Oct 7, 2026

Uh oh!

jhnwu3 left a comment

Uh oh!

jhnwu3 Oct 11, 2026

Uh oh!

github-actions Bot commented Oct 26, 2026 •

edited

Loading

Uh oh!

AxelNoun commented Oct 27, 2026

Uh oh!

AxelNoun commented Oct 28, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		__all__ = ["SampleKGDataset"]


		class SampleKGDataset(InMemorySampleDataset):

Conversation

AxelNoun commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AxelNoun commented Aug 29, 2026

Uh oh!

AxelNoun commented Aug 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Architecture & Implementation Plan Update

Uh oh!

AxelNoun commented Aug 31, 2026

CI/CD Green & Ready for Review

Uh oh!

AxelNoun commented Oct 7, 2026

Uh oh!

jhnwu3 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AxelNoun commented Oct 7, 2026

Uh oh!

jhnwu3 left a comment

Choose a reason for hiding this comment

Uh oh!

jhnwu3 Oct 11, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Test Coverage

Uh oh!

AxelNoun commented Oct 27, 2026

Why the 1.x layer can go

Your criterion of 7 Oct: one streaming backend

Fixed since your review

Outside this PR

Questions

Uh oh!

AxelNoun commented Oct 28, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

AxelNoun commented Aug 25, 2026 •

edited

Loading

AxelNoun commented Aug 31, 2026 •

edited

Loading

jhnwu3 left a comment •

edited

Loading

github-actions Bot commented Oct 26, 2026 •

edited

Loading