Add NestedMultiHotProcessor and vectorise HALO's visit encoding - #1233
Merged
Merged
Conversation
EHR generation feeds HALO a nested list of per-visit codes. The existing NestedSequenceProcessor emits code indices padded to the longest visit seen during fit, so one outlier sets the width for the whole dataset: on eICU a single visit holds 3,951 code entries (the same diagnosis re-charted through a stay) while a typical visit holds about five. HALO then unpacked that back into multi-hot vectors with a triple-nested Python loop. Add NestedMultiHotProcessor (registered as "nested_multihot"), which emits one multi-hot row per visit, sized by the vocabulary rather than by the worst-case visit -- 8.6x smaller per patient on a 921-code eICU vocabulary (4.5 KB vs 38.7 KB). Repeats within a visit collapse to a single 1, which is what set-membership models already did with the index form. Switch EHRGeneration to it and rewrite HALO._encode_visits as a vectorised placement into the context window. The old loop cost an .item() per patient -- a CUDA sync each -- and a single-element kernel launch per code, which measured at 84.5% of a training step at batch 128 on an A100 (13.887 s encode vs 2.553 s forward+backward), and ~99.8% in steady state once CUDA warmup is excluded. decode_dataset is updated to invert the new encoding: reading multi-hot rows as indices would see only 0s and 1s and decode every visit as empty. <pad> (0) and <unk> (1) keep their indices, so the vocabulary is interchangeable between the two processors. Tests: test_nested_multihot_processor.py covers the processor directly, and test_halo_encode_equivalence.py asserts HALO sees identical tensors from either processor, so the switch changes no results. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…example - Modernise annotations in nested_multihot_processor.py and generate_ehr.py (PEP 585/604, ClassVar for the schema dicts) to clear ruff UP006/UP007/ UP035/RUF012 on the lines this PR touches. - Add '>>>' usage examples to EHRGeneration and decode_dataset. - Update examples/halo_mimic3.py for the multi-hot encoding: the per-visit code set is now the nonzero column indices of a (num_visits, vocab_size) tensor, not the tensor values. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jhnwu3
requested changes
Sep 11, 2026
jhnwu3
left a comment
jhnwu3
left a comment
Collaborator
There was a problem hiding this comment.
Tldr; We may not need test_halo_encode_equivalence.py here. It seems redundant here.
| from pyhealth.processors import NestedSequenceProcessor | ||
|
|
||
|
|
||
| def legacy_encode(cfg, visits, device): |
Collaborator
There was a problem hiding this comment.
Do we need to have a legacy check here? I think it's ok to just use our patched/fixed generative modeling setup here.
test_bit_identical was a one-time migration check: it pinned the vectorised _encode_visits against a frozen copy of the loop it replaced. With the original deleted, it only re-asserts a dead implementation. test_no_inner_padding_in_multihot duplicated test_nested_multihot_processor.test_shape_is_visits_by_vocab. test_truncates_past_context was not redundant -- nothing else covered _encode_visits cutting a patient at n_ctx-2 -- so it moves to test_halo.py as test_encode_truncates_past_context. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
EHRGeneration is shared by HALO, GPT2 and PromptEHR, so switching its
input_schema to NestedMultiHotProcessor changed what all three receive. Only
HALO was updated. GPT2._encode_visits and PromptEHR._serialize both did
codes = [int(c) for c in visits[i, j].tolist() if c > 0]
which reads the row's *values* as code ids. Under a multi-hot row every value
is 1.0, and 1 is <unk>, so every code in every visit silently became <unk>:
no crash, loss still falls, generated patients are noise. Their tests missed
it by building datasets with 'nested_sequence' spelled out instead of going
through the task.
Each nested processor now inverts its own encoding via visit_code_ids(), and
both models call that, so either encoding works. Their constructors reject a
'visits' processor that cannot.
Separately, five processors (Sequence, NestedSequence, DeepNestedSequence,
StageNet, NestedMultiHot) carried verbatim copies of remove/retain/add/tokens/
vocab_size. Those move to CodeVocabularyMixin in base_processor.py. A mixin,
not a base class: models dispatch on isinstance(p, NestedSequenceProcessor) to
decide whether to apply nn.Embedding, so making one code processor inherit
from another would reroute it.
The copies all shared a bug -- remove() renumbered the vocabulary but left
_next_index stale, so a later fit() allocated an index past vocab_size().
The shared version fixes it for all five.
Also refreshes the halo/gpt2/promptehr docstrings that still described the
index encoding, and leaves a TODO on _encode_visits about interior empty
visits (unreachable from EHRGeneration today).
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
GPT2 and PromptEHR are token models: they flatten each visit into a stream of
code ids. Feeding them multi-hot meant encoding a code set to a vocabulary-wide
indicator vector and immediately decoding it back to the indices they wanted.
HALO is the opposite -- its transformer consumes multi-hot directly. MedGAN and
CorGAN have no visit axis at all and had no task whatsoever; callers hand-built
{'visits': 'multi_hot'} datasets.
EHRGeneration keeps the extraction and drops its input_schema, becoming the
shared base. Three subclasses declare the encodings:
VisitMultiHotGeneration per-visit multi-hot rows HALO
VisitSequenceGeneration per-visit code indices GPT2, PromptEHR
PatientCodeSetGeneration one pooled set per patient MedGAN, CorGAN
EHRGenerationMIMIC3/4 keep their behaviour as VisitMultiHotGeneration presets.
event_type/code_attr/min_visits are now constructor arguments, so the encoding
and the dataset stay independent instead of becoming a class grid.
Instantiating EHRGeneration itself now raises and names the three subclasses,
rather than failing later with an AttributeError inside set_task.
decode_dataset goes through the processor's visit_code_ids, so it handles
either per-visit encoding and refuses the bag-of-codes one explicitly.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Each processor keeps its own remove/retain/add/tokens/vocab_size again, and CodeVocabularyMixin is gone from base_processor.py. The duplication is real but local: a processor stays readable and editable on its own, without a base class to consult, which is the structure this package already had. visit_code_ids stays on the two nested processors -- it is not shared vocabulary machinery but each processor's own inverse, which GPT2, PromptEHR and decode_dataset need in order to read a visit row without assuming an encoding. This also restores the pre-existing _next_index defect in all five copies: remove() renumbers the vocabulary but leaves _next_index stale, so a later fit() allocates an index past vocab_size(). Untouched here rather than fixed five times over, since it predates this PR. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
NestedMultiHotProcessor.visit_code_ids had no production caller: after the task split GPT2 and PromptEHR only ever see NestedSequenceProcessor, and decode_dataset was its only other consumer. decode_dataset reads nonzero columns directly again and rejects a non-multi-hot processor by name. visit_code_ids stays on NestedSequenceProcessor, where GPT2 and PromptEHR use it. decode_dataset now resolves through a torch Subset, so a split from split_by_patient can be decoded -- which is why halo_mimic3.py had its own copy of the decoding. That copy is gone; the example calls decode_dataset and to_evaluation_dataframe, which is what they exist for. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
PatientCodeSetGeneration is the only task with a custom __call__ and had no test that ran it. A minimal stand-in patient exercises all three: the shared per-visit extraction, the pooling and dedupe, that min_visits still counts real visits before the visit axis is collapsed, and that codeless admissions drop out. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
Author
|
Noticed the GPT2 which used to use NestedSequenceProcessor is now taking in the NestedMultiHotProcessor and converting it back into the sequence. Changed it so it just takes NestedSequenceProcessor as before. Also bundled in new tasks for MedGAN/CorGAN which are bag-of-codes, since they don't have a task yet. Result:
note now the task |
jhnwu3
requested changes
Sep 12, 2026
|
|
||
| def _visits(self, patient: Patient) -> list[list[str]]: | ||
| """Ordered per-admission code lists, empty admissions dropped.""" | ||
| visits: list[list[str]] = [] |
Collaborator
There was a problem hiding this comment.
My only concern here is that this type of parent class only works on MIMIC and not other datasets like eICU, so should anyone try to extend this parent class. It would fail.
I think it would be good to simplify and simply just have a flat list of different task classes for each model/dataset. We could maybe rename all of this to the MIMIC-series of generative tasks here and get rid of the inheritance as it's unnecessary.
| VisitSequenceGeneration, and MedGAN/CorGAN PatientCodeSetGeneration) | ||
| 3. Creating a SampleDataset with a NestedMultiHotProcessor | ||
| 4. Training the HALO generator with its custom training loop | ||
| 5. Generating synthetic patients |
Collaborator
There was a problem hiding this comment.
If you could vibe generate more examples using the revamped tasks, that would be awesome!
Addresses jhnwu3's review. The parent class only worked on MIMIC: it assumed an 'admissions' event type and a hadm_id linking codes to an admission. Subclassing it for eICU or OMOP would not have raised -- it would have returned zero samples, which is worse. EHRGeneration, VisitMultiHotGeneration, VisitSequenceGeneration and PatientCodeSetGeneration are replaced by six flat classes, each subclassing BaseTask directly and each naming its dataset: EHRGenerationMIMIC3/4 per-visit multi-hot HALO EHRSequenceGenerationMIMIC3/4 per-visit indices GPT2, PromptEHR EHRCodeSetGenerationMIMIC3/4 one set per patient MedGAN, CorGAN The shared extraction survives as a module-level _mimic_visits() helper rather than six copies. No task inherits from another, and the helper's name and docstring both say it is MIMIC-shaped, so a task for another dataset writes its own instead of reaching for it. A test asserts every task's only base is BaseTask. Removing EHRGeneration also retires the breaking change this PR introduced earlier -- there is no longer a base class to call by mistake. Adds examples/gpt2_mimic3.py and examples/medgan_mimic3.py alongside halo_mimic3.py, one per encoding, each cross-referencing the others. The MedGAN example asks for metrics='privacy' rather than 'all': the utility metric scores next-visit prediction, which is meaningless once the visit axis is pooled away, and would otherwise return a number that looks real. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Collaborator
Author
|
Flattened, got rid of EHRGeneration. now everything is EHR{generation,codeset}{mimic3,mimic4} and I'll add eICU later. |
janezdu
added a commit
to janezdu/PyHealth
that referenced
this pull request
Sep 13, 2026
Brings in the merged NestedMultiHotProcessor / HALO encoding work, which
replaced EHRGeneration with six flat per-(model, dataset) task classes.
Conflict resolutions:
- generate_ehr.py: took upstream's six flat tasks. EHRGenerationEICU is kept
and now subclasses BaseTask directly, declaring
input_schema = {'visits': NestedMultiHotProcessor} explicitly -- the class it
used to inherit that from no longer exists. Its __call__ already did its own
eICU extraction (patientunitstayid, not hadm_id), so no logic changed. Its
annotations were modernised because upstream's file dropped typing.Dict/List.
- tasks/__init__.py: upstream's exports plus EHRGenerationEICU.
- halo.py, test_halo.py: took upstream; the adapter/IRM work here was outside
the conflicted region and is unchanged.
- nested_multihot_processor.py: took upstream (differences were typing-only).
- .gitignore: kept both sides.
In-progress examples/fedpyhealth work is deliberately not part of this commit.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
EHR generation feeds HALO a nested list of per-visit codes. The existing NestedSequenceProcessor emits code indices padded to the longest visit seen during `it, so one outlier sets the width for the whole dataset: on eICU a single visit holds 3,951 code entries (the same diagnosis re-charted through a stay) while a typical visit holds about five. HALO then unpacked that back into multi-hot vectors with a triple-nested Python loop.
Add NestedMultiHotProcessor (registered as
nested_multihot), which emits one multi-hot row per visit, sized by the vocabulary rather than by the worst-case visit -- 8.6x smaller per patient on a 921-code eICU vocabulary (4.5 KB vs 38.7 KB). Repeats within a visit collapse to a single 1, which is what set-membership models already did with the index form.Switch EHRGeneration to it and rewrite
HALO._encode_visitsas a vectorised placement into the context window. The old loop cost an.item()per patient -- a CUDA sync each -- and a single-element kernel launch per code, which measured at 84.5% of a training step at batch 128 on an A100 (13.887 s encode vs 2.553 s forward+backward), and ~99.8% in steady state once CUDA warmup is excluded.decode_dataset is updated to invert the new encoding: reading multi-hot rows as indices would see only 0s and 1s and decode every visit as empty.
(0) and (1) keep their indices, so the vocabulary is interchangeable between the two processors.
Tests: test_nested_multihot_processor.py covers the processor directly, and test_halo_encode_equivalence.py asserts HALO sees identical tensors from either processor, so the switch changes no results.
tl;dr NestedMultiHotProcessor + vectorised HALO._encode_visits. 8.6× less memory per patient, and encoding drops from 84.5% of a training step to negligible. Equivalence test proves results don't change. 10 files, +472/−33.