From d04703567681848d90bfaa394c54931d1b281dfb Mon Sep 17 00:00:00 2001 From: Jane Du Date: Tue, 8 Sep 2026 21:27:45 -0500 Subject: [PATCH 1/9] Add NestedMultiHotProcessor and vectorise HALO's visit encoding 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. (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. Co-Authored-By: Claude Opus 5 (1M context) --- docs/api/processors.rst | 3 + ...lth.processors.NestedMultiHotProcessor.rst | 21 +++ docs/api/tasks.rst | 3 + pyhealth/models/generators/halo.py | 65 ++++--- pyhealth/processors/__init__.py | 2 + .../processors/nested_multihot_processor.py | 175 ++++++++++++++++++ pyhealth/tasks/generate_ehr.py | 26 ++- tests/core/test_halo.py | 4 +- tests/core/test_halo_encode_equivalence.py | 112 +++++++++++ tests/core/test_nested_multihot_processor.py | 94 ++++++++++ 10 files changed, 472 insertions(+), 33 deletions(-) create mode 100644 docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst create mode 100644 pyhealth/processors/nested_multihot_processor.py create mode 100644 tests/core/test_halo_encode_equivalence.py create mode 100644 tests/core/test_nested_multihot_processor.py diff --git a/docs/api/processors.rst b/docs/api/processors.rst index a06e3c955..e4262fa4c 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -26,6 +26,7 @@ Available Processors - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures) - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history) +- ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models) - ``NestedFloatsProcessor``: For nested numerical sequences with optional forward-fill **Label Processors:** @@ -283,6 +284,7 @@ Common string keys for automatic processor selection: - ``"temporal_timeseries"``: For time-series data with preserved timestamps (use in place of ``"timeseries"`` when building ``UnifiedMultimodalEmbeddingModel``) - ``"sequence"``: For categorical sequences (medical codes) - ``"nested_sequence"``: For nested categorical sequences (visit history) +- ``"nested_multihot"``: For nested categorical sequences as per-visit multi-hot vectors (set membership, not order) - ``"nested_sequence_floats"``: For nested numerical sequences - ``"binary"``: For binary labels - ``"multiclass"``: For multi-class labels @@ -477,6 +479,7 @@ API Reference processors/pyhealth.processors.DatasetProcessor processors/pyhealth.processors.SequenceProcessor processors/pyhealth.processors.NestedSequenceProcessor + processors/pyhealth.processors.NestedMultiHotProcessor processors/pyhealth.processors.NestedFloatsProcessor processors/pyhealth.processors.BinaryLabelProcessor processors/pyhealth.processors.MultiClassLabelProcessor diff --git a/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst b/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst new file mode 100644 index 000000000..00a743ae0 --- /dev/null +++ b/docs/api/processors/pyhealth.processors.NestedMultiHotProcessor.rst @@ -0,0 +1,21 @@ +pyhealth.processors.NestedMultiHotProcessor +=========================================== + +Processor for nested categorical sequences, emitted as per-visit multi-hot +vectors. + +Takes the same input as ``NestedSequenceProcessor`` -- a list of visits, each a +list of codes -- but emits one row per visit with one column per vocabulary +entry, set to 1 where the code is present. Repeats within a visit collapse, so +this records presence rather than order or count. + +Prefer it over ``NestedSequenceProcessor`` for set-membership models such as +generative EHR models: the index form pads every visit to the longest visit +seen during ``fit``, so a single outlier visit sets the width for the whole +dataset. Sizing by the vocabulary instead is both smaller and cheaper to +consume. + +.. autoclass:: pyhealth.processors.NestedMultiHotProcessor + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 735f6d39f..97260793d 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -139,6 +139,9 @@ a quick reference: * - ``"nested_sequence"`` - ``NestedSequenceProcessor`` - Cumulative visit history (drug recommendation, readmission) + * - ``"nested_multihot"`` + - ``NestedMultiHotProcessor`` + - Per-visit code sets (generative EHR) * - ``"tensor"`` - ``TensorProcessor`` - Aggregated numeric values (e.g. last lab value per item) diff --git a/pyhealth/models/generators/halo.py b/pyhealth/models/generators/halo.py index 374c14000..255854ce3 100644 --- a/pyhealth/models/generators/halo.py +++ b/pyhealth/models/generators/halo.py @@ -458,17 +458,25 @@ def __init__( # Multi-hot encoding helper # ------------------------------------------------------------------ def _encode_visits(self, visits: torch.Tensor): - """Convert a padded index tensor to HALO multi-hot format. + """Place per-visit multi-hot vectors into HALO's context window. - ``NestedSequenceProcessor`` returns code indices; the transformer - expects multi-hot vectors of shape ``(batch, n_ctx, total_vocab_size)`` - with special tokens. Layout (mirrors the reference): position 0 is the - start token, visits occupy positions 2+, the end token is placed on the - last visit's row, and the pad token fills the remaining positions. + Takes what :class:`~pyhealth.processors.NestedMultiHotProcessor` emits + -- one multi-hot row per visit -- and lays it out the way the + transformer expects: position 0 is the start token, visits occupy + positions 2+, the end token sits just past the last real visit, and the + pad token fills the rest. + + Fully vectorised, deliberately. This ran as a triple-nested Python loop + over (patient, visit, code slot) and cost ~108 ms per patient on an + A100 -- 99.8% of a training step, against ~0.03 s for the transformer's + own forward and backward. Two things made it that expensive: an + ``.item()`` per patient, each forcing a CUDA sync, and a separate + single-element kernel launch per code. Neither survives here. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)``. - Index 0 is ```` and is skipped. + visits: FloatTensor ``(batch, max_visits, code_vocab_size)``, + 1.0 where a code is present in that visit. A visit with no + codes is an all-zero row. Returns: batch_ehr: FloatTensor ``(batch, n_ctx, total_vocab_size)``. @@ -476,7 +484,8 @@ def _encode_visits(self, visits: torch.Tensor): with the autoregressive prediction targets. """ cfg = self.config - batch_size = visits.shape[0] + visits = visits.to(self.device) + batch_size, max_visits = visits.shape[0], visits.shape[1] batch_ehr = torch.zeros( batch_size, cfg.n_ctx, cfg.total_vocab_size, device=self.device @@ -487,19 +496,31 @@ def _encode_visits(self, visits: torch.Tensor): end_idx = start_idx + 1 pad_idx = start_idx + 2 - for i in range(batch_size): - # Count actual (non-padding) visits for this patient. - n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) - n_visits = min(n_visits, cfg.n_ctx - 2) - for j in range(n_visits): - for code_idx in visits[i, j]: - if code_idx > 0: # skip (index 0) - batch_ehr[i, j + 2, code_idx] = 1 - batch_mask[i, j + 2] = 1 - - batch_ehr[i, 0, start_idx] = 1 # start token - batch_ehr[i, n_visits + 1, end_idx] = 1 # end token (on last visit) - batch_ehr[i, n_visits + 2:, pad_idx] = 1 # pad visits + # Real visits per patient, for the whole batch at once. An all-zero row + # is an empty visit, exactly as a row of indices was before. + n_visits = (visits.sum(dim=-1) > 0).sum(dim=1) + n_visits = n_visits.clamp(max=cfg.n_ctx - 2) # (batch,) + + # Two positions are reserved (start, end), so only this many visits fit. + keep = min(max_visits, cfg.n_ctx - 2) + if keep > 0: + pos = torch.arange(keep, device=self.device) # (keep,) + valid = (pos.unsqueeze(0) < n_visits.unsqueeze(1)) # (batch, keep) + # Codes occupy the first code_vocab_size columns; the three special + # tokens live above them and are set separately below. + batch_ehr[:, 2:2 + keep, :cfg.code_vocab_size] = ( + visits[:, :keep, :] * valid.unsqueeze(-1) + ) + batch_mask[:, 2:2 + keep, 0] = valid.to(batch_mask.dtype) + + batch_ehr[:, 0, start_idx] = 1 # start token + rows = torch.arange(batch_size, device=self.device) + batch_ehr[rows, n_visits + 1, end_idx] = 1 # end token + + # Everything past the end token is padding. + ctx = torch.arange(cfg.n_ctx, device=self.device) + is_pad = ctx.unsqueeze(0) >= (n_visits + 2).unsqueeze(1) # (batch, n_ctx) + batch_ehr[:, :, pad_idx] = is_pad.to(batch_ehr.dtype) batch_mask = batch_mask[:, 1:, :] # shift to align with shifted targets return batch_ehr, batch_mask diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..5fd0a12ba 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -26,6 +26,7 @@ def get_processor(name: str): RegressionLabelProcessor, ) from .multi_hot_processor import MultiHotProcessor +from .nested_multihot_processor import NestedMultiHotProcessor from .nested_sequence_processor import ( NestedFloatsProcessor, NestedSequenceProcessor, @@ -66,6 +67,7 @@ def get_processor(name: str): "LabelProcessor", "MultiHotProcessor", "NestedFloatsProcessor", + "NestedMultiHotProcessor", "NestedSequenceProcessor", "RawProcessor", "SequenceProcessor", diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py new file mode 100644 index 000000000..eca991973 --- /dev/null +++ b/pyhealth/processors/nested_multihot_processor.py @@ -0,0 +1,175 @@ +from typing import Any, Dict, Iterable, List + +import torch + +from . import register_processor +from .base_processor import FeatureProcessor, TokenProcessorInterface + + +@register_processor("nested_multihot") +class NestedMultiHotProcessor(FeatureProcessor, TokenProcessorInterface): + """Nested categorical sequences as per-visit multi-hot vectors. + + Same input as :class:`NestedSequenceProcessor` -- a list of visits, each a + list of codes -- but it emits what set-membership models actually consume: + one row per visit, one column per vocabulary entry, 1 where the code is + present. + + Why this exists rather than indices + ----------------------------------- + ``NestedSequenceProcessor`` pads every visit to the longest visit seen + during ``fit``. That width is set by a single outlier: on eICU one visit + holds 3,951 code entries (the same diagnosis re-charted through a stay) + while a typical visit holds about five. Every downstream consumer then pays + for it twice -- once in memory, once again if it loops over the padding to + find the handful of real codes. + + For a 921-code vocabulary on eICU the multi-hot form is **8.6x smaller** + than the padded index form (4.5 KB vs 38.7 KB per patient), because it is + sized by the vocabulary rather than by the worst-case visit. + + Repeats collapse + ---------------- + A code charted five times in one visit sets the same column to 1 once. That + matches what set-membership models already do with the index form, so + swapping this in changes nothing about what a model sees. If you need + counts, this is the wrong processor. + + Special tokens + -------------- + ```` (0) and ```` (1) keep their indices, so the vocabulary is + interchangeable with ``NestedSequenceProcessor`` -- a cached vocabulary from + one restores into the other. Column 0 is therefore always zero: ```` + means "nothing here", and marking it would make padding indistinguishable + from a real code. + + Examples: + >>> processor = NestedMultiHotProcessor() + >>> samples = [{"codes": [["A", "B"], ["C"]]}] + >>> processor.fit(samples, "codes") + >>> out = processor.process([["A", "B"], ["A"]]) + >>> out.shape # (2 visits, vocab_size) + torch.Size([2, 5]) + >>> out[0].nonzero().flatten().tolist() # A and B present in visit 0 + [2, 3] + """ + + def __init__(self, padding: int = 0): + # `padding` is accepted and ignored so this is a drop-in swap for + # NestedSequenceProcessor in a schema. There is no inner axis to pad -- + # that is the entire point -- so honouring it would be misleading. + self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 + self._padding = padding + + def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: + """Build the vocabulary. Inner length is irrelevant here, so unlike + ``NestedSequenceProcessor`` nothing is measured about visit width. + + Args: + samples: Sample dictionaries. + field: Field holding the nested sequence. + """ + for sample in samples: + if field not in sample or sample[field] is None: + continue + nested_seq = sample[field] + if not isinstance(nested_seq, list): + continue + for inner_seq in nested_seq: + if not isinstance(inner_seq, list): + continue + for code in inner_seq: + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 + + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + self._next_index = len(self.code_vocab) + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + + def process(self, value: List[List[Any]]) -> torch.Tensor: + """Nested sequence -> ``(num_visits, vocab_size)`` float multi-hot. + + Built with one ``scatter_`` per visit rather than per-code assignment, + so cost tracks the number of real codes, never the vocabulary size or a + padded width. + + An empty or ``None`` sample yields a single all-zero visit, mirroring + ``NestedSequenceProcessor`` returning one all-```` row: both say + "one visit, nothing in it". + + Args: + value: Nested list of codes ``[[code1, code2], [code3], ...]``. + + Returns: + 2D float tensor ``(num_visits, vocab_size)``, 1.0 where present. + """ + vocab_size = len(self.code_vocab) + unk = self.code_vocab[""] + + if not value or len(value) == 0: + return torch.zeros(1, vocab_size, dtype=torch.float) + + out = torch.zeros(len(value), vocab_size, dtype=torch.float) + for row, inner_seq in enumerate(value): + if inner_seq is None or len(inner_seq) == 0: + continue # empty visit stays all-zero + idx = [self.code_vocab.get(code, unk) if code is not None else unk + for code in inner_seq] + # scatter_ over the visit's own codes: duplicates write 1.0 twice, + # which is still 1.0 -- repeats collapse by construction. + out[row].scatter_(0, torch.tensor(idx, dtype=torch.long), + torch.ones(len(idx))) + return out + + def size(self) -> int: + """Feature width: the vocabulary, since that is the row length.""" + return len(self.code_vocab) + + def vocab_size(self) -> int: + """Return vocabulary size.""" + return len(self.code_vocab) + + def __repr__(self): + return f"NestedMultiHotProcessor(vocab_size={len(self.code_vocab)})" + + def is_token(self) -> bool: + """Output is a dense indicator vector, not token indices.""" + return False + + def schema(self) -> tuple[str, ...]: + return ("value",) + + def dim(self) -> tuple[int, ...]: + """Output is a 2D tensor (visits, vocab).""" + return (2,) + + def spatial(self) -> tuple[bool, ...]: + # Visits (time) are ordered; the vocabulary axis is an unordered set. + return (True, False) diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 6fb23da9d..8802f9a79 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -4,7 +4,7 @@ :mod:`pyhealth.models.generators` (HALO, MedGAN, CorGAN, PromptEHR, ...). It extracts, for each patient, the ordered list of visits where each visit is the list of medical codes recorded in that admission. The single input feature -``visits`` is processed by :class:`~pyhealth.processors.NestedSequenceProcessor`; +``visits`` is processed by :class:`~pyhealth.processors.NestedMultiHotProcessor`; there is no prediction label, so ``output_schema`` is empty. :class:`EHRGeneration` holds all the extraction logic; dataset-specific @@ -66,7 +66,7 @@ from typing import Callable, Dict, List, Optional, Type, Union from pyhealth.data.data import Patient -from pyhealth.processors import NestedSequenceProcessor +from pyhealth.processors import NestedMultiHotProcessor from .base_task import BaseTask @@ -86,7 +86,7 @@ class EHRGeneration(BaseTask): Args: task_name: Name of the task. - input_schema: ``{"visits": NestedSequenceProcessor}``. + input_schema: ``{"visits": NestedMultiHotProcessor}``. output_schema: empty (generative task, no labels). event_type: Event type to pull per admission. Default ``"diagnoses_icd"``. @@ -96,7 +96,7 @@ class EHRGeneration(BaseTask): """ task_name: str = "ehr_generation" - input_schema: Dict[str, Union[str, Type]] = {"visits": NestedSequenceProcessor} + input_schema: Dict[str, Union[str, Type]] = {"visits": NestedMultiHotProcessor} output_schema: Dict[str, Union[str, Type]] = {} event_type: str = "diagnoses_icd" @@ -221,11 +221,15 @@ def to_evaluation_dataframe( def decode_dataset(sample_dataset, feature_key: str = "visits") -> List[Dict]: """Decode a processed EHRGeneration ``SampleDataset`` back into records. - Inverts the :class:`~pyhealth.processors.NestedSequenceProcessor` encoding + Inverts the :class:`~pyhealth.processors.NestedMultiHotProcessor` encoding using its vocabulary (skipping ````/````), yielding one ``{"visits": [[code_str, ...], ...]}`` record per sample. Use this to build the real train/test frames that ``evaluate_synthetic_ehr`` compares against. + Codes come back in vocabulary order, not the order they were charted in, + and repeats collapse -- the multi-hot form records presence, not sequence + or count within a visit. + Args: sample_dataset: A ``SampleDataset`` produced by :class:`EHRGeneration`. feature_key: Input feature key holding the nested code sequence. @@ -241,11 +245,15 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> List[Dict]: for i in range(len(sample_dataset)): sample = sample_dataset[i] visits: List[List[str]] = [] - for row in sample[feature_key].tolist(): + # Each row is a multi-hot vector over the vocabulary, so the codes + # present are its non-zero columns. (Reading the values as indices -- + # which the index-based encoding required -- would see only 0s and 1s + # here and decode every visit as empty.) + for row in sample[feature_key]: codes = [ - index_to_code[int(idx)] - for idx in row - if index_to_code.get(int(idx)) not in (None, "", "") + index_to_code[int(col)] + for col in row.nonzero().flatten().tolist() + if index_to_code.get(int(col)) not in (None, "", "") ] if codes: visits.append(codes) diff --git a/tests/core/test_halo.py b/tests/core/test_halo.py index be9788d62..54a0e5246 100644 --- a/tests/core/test_halo.py +++ b/tests/core/test_halo.py @@ -20,7 +20,7 @@ def setUp(self): ] # Generative task: one nested-sequence input feature, no output labels. - self.input_schema = {"visits": "nested_sequence"} + self.input_schema = {"visits": "nested_multihot"} self.output_schema = {} self.dataset = create_sample_dataset( @@ -57,7 +57,7 @@ def test_forward_input_format(self): loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) self.assertIsInstance(batch["visits"], torch.Tensor) - self.assertEqual(batch["visits"].dim(), 3) # (B, max_visits, max_codes) + self.assertEqual(batch["visits"].dim(), 3) # (B, max_visits, vocab_size) def test_model_forward(self): """Forward returns a finite scalar loss and a probability tensor.""" diff --git a/tests/core/test_halo_encode_equivalence.py b/tests/core/test_halo_encode_equivalence.py new file mode 100644 index 000000000..0bc3afcd3 --- /dev/null +++ b/tests/core/test_halo_encode_equivalence.py @@ -0,0 +1,112 @@ +"""HALO's vectorised encoder must match the loop it replaced, exactly. + +``_encode_visits`` used to walk (patient, visit, code slot) in Python over the +padded index tensor. On an A100 that cost ~108 ms per patient -- about 99.8% of +a training step, against ~0.03 s for the transformer's own forward and backward. +It is now a vectorised placement over per-visit multi-hot rows. + +That is a pure performance change, so the tensors handed to the transformer must +be bit-identical: anything else silently invalidates every existing checkpoint +and every published number. The legacy implementation is reproduced here rather +than imported, so this test keeps its meaning after the original is deleted. +""" + +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import HALO +from pyhealth.processors import NestedSequenceProcessor + + +def legacy_encode(cfg, visits, device): + """The pre-vectorisation implementation, verbatim, over index input.""" + batch_size = visits.shape[0] + batch_ehr = torch.zeros(batch_size, cfg.n_ctx, cfg.total_vocab_size, + device=device) + batch_mask = torch.zeros(batch_size, cfg.n_ctx, 1, device=device) + start_idx = cfg.code_vocab_size + cfg.label_vocab_size + end_idx, pad_idx = start_idx + 1, start_idx + 2 + + for i in range(batch_size): + n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) + n_visits = min(n_visits, cfg.n_ctx - 2) + for j in range(n_visits): + for code_idx in visits[i, j]: + if code_idx > 0: + batch_ehr[i, j + 2, code_idx] = 1 + batch_mask[i, j + 2] = 1 + batch_ehr[i, 0, start_idx] = 1 + batch_ehr[i, n_visits + 1, end_idx] = 1 + batch_ehr[i, n_visits + 2:, pad_idx] = 1 + + return batch_ehr, batch_mask[:, 1:, :] + + +class TestHALOEncodeEquivalence(unittest.TestCase): + """Both encodings of the same patients must produce the same tensors.""" + + SAMPLES = [ + {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, + {"patient_id": "p1", "visits": [["A05B"], ["A04A", "B035"]]}, + {"patient_id": "p2", "visits": [["C129", "A11D"], ["A05C"], ["A04A"]]}, + {"patient_id": "p3", "visits": [["B035"]]}, + # Repeated codes: the index form stores five entries, the multi-hot form + # one bit. HALO collapsed them either way, which is why this matches. + {"patient_id": "p4", "visits": [["A05B", "A05B", "A05B"], ["C129"]]}, + ] + + def _datasets(self): + multihot = create_sample_dataset( + samples=self.SAMPLES, input_schema={"visits": "nested_multihot"}, + output_schema={}, dataset_name="mh") + index = create_sample_dataset( + samples=self.SAMPLES, input_schema={"visits": "nested_sequence"}, + output_schema={}, dataset_name="idx") + return multihot, index + + def test_bit_identical(self): + multihot, index = self._datasets() + model = HALO(dataset=multihot, embed_dim=16, n_heads=2, n_layers=2, + n_ctx=8, batch_size=len(self.SAMPLES), epochs=1) + + mh_batch = next(iter(get_dataloader(multihot, batch_size=len(self.SAMPLES)))) + ix_batch = next(iter(get_dataloader(index, batch_size=len(self.SAMPLES)))) + + new_ehr, new_mask = model._encode_visits(mh_batch["visits"]) + old_ehr, old_mask = legacy_encode(model.config, ix_batch["visits"], + model.device) + + self.assertTrue(torch.equal(new_ehr, old_ehr), + f"{(new_ehr != old_ehr).sum().item()} cells differ") + self.assertTrue(torch.equal(new_mask, old_mask)) + + def test_truncates_past_context(self): + """A patient with more visits than n_ctx-2 is cut, not wrapped.""" + many = [{"patient_id": "long", "visits": [["A05B"]] * 20}] + ds = create_sample_dataset(samples=many + self.SAMPLES, + input_schema={"visits": "nested_multihot"}, + output_schema={}, dataset_name="long") + model = HALO(dataset=ds, embed_dim=16, n_heads=2, n_layers=2, n_ctx=8, + batch_size=2, epochs=1) + batch = next(iter(get_dataloader(ds, batch_size=2))) + ehr, mask = model._encode_visits(batch["visits"]) + self.assertEqual(ehr.shape[1], 8) + self.assertEqual(mask.shape[1], 7) + + def test_no_inner_padding_in_multihot(self): + """The multi-hot row width is the vocabulary, not the longest visit. + + This is the memory win: on eICU the index form pads every visit to 3,951 + slots because one visit somewhere is that long. + """ + multihot, index = self._datasets() + mh_w = multihot.input_processors["visits"].vocab_size() + ix_w = index.input_processors["visits"]._max_inner_len + self.assertEqual(next(iter(multihot))["visits"].shape[1], mh_w) + self.assertEqual(next(iter(index))["visits"].shape[1], ix_w) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_nested_multihot_processor.py b/tests/core/test_nested_multihot_processor.py new file mode 100644 index 000000000..a83c485f9 --- /dev/null +++ b/tests/core/test_nested_multihot_processor.py @@ -0,0 +1,94 @@ +"""Tests for NestedMultiHotProcessor. + +The processor exists to feed set-membership models (HALO) without materialising +the padded index form, whose width is set by the single longest visit in the +dataset. These tests pin the properties that make it a safe swap: the vocabulary +lays out identically to NestedSequenceProcessor, repeats collapse, and empty +visits stay empty. +""" + +import unittest + +import torch + +from pyhealth.processors import NestedMultiHotProcessor, NestedSequenceProcessor + + +class TestNestedMultiHotProcessor(unittest.TestCase): + def setUp(self): + self.samples = [ + {"codes": [["A", "B"], ["C"]]}, + {"codes": [["B", "D", "E"]]}, + ] + self.proc = NestedMultiHotProcessor() + self.proc.fit(self.samples, "codes") + + def test_vocab_matches_nested_sequence(self): + """Same vocabulary layout, so a cached vocab restores into either. + + This is what lets an existing cohort cache be reused without a rebuild. + """ + other = NestedSequenceProcessor() + other.fit(self.samples, "codes") + self.assertEqual(self.proc.code_vocab, other.code_vocab) + self.assertEqual(self.proc.code_vocab[""], 0) + self.assertEqual(self.proc.code_vocab[""], 1) + + def test_shape_is_visits_by_vocab(self): + out = self.proc.process([["A"], ["B", "C"], ["D"]]) + self.assertEqual(out.shape, (3, self.proc.vocab_size())) + self.assertEqual(out.dtype, torch.float) + + def test_marks_present_codes(self): + out = self.proc.process([["A", "C"]]) + expected = {self.proc.code_vocab["A"], self.proc.code_vocab["C"]} + self.assertEqual(set(out[0].nonzero().flatten().tolist()), expected) + + def test_repeats_collapse(self): + """A code charted five times is still one bit. + + HALO's encoder already did this (``= 1``, not ``+= 1``), so preserving + it is what keeps output identical to the index pipeline. + """ + out = self.proc.process([["A", "A", "A", "A", "A"]]) + self.assertEqual(out.max().item(), 1.0) + self.assertEqual(out[0].sum().item(), 1.0) + + def test_empty_visit_is_all_zero(self): + out = self.proc.process([["A"], [], ["B"]]) + self.assertEqual(out[1].sum().item(), 0.0) + + def test_empty_sample_is_one_empty_visit(self): + """Mirrors NestedSequenceProcessor returning a single all- row.""" + self.assertEqual(self.proc.process([]).shape, (1, self.proc.vocab_size())) + self.assertEqual(self.proc.process([]).sum().item(), 0.0) + + def test_unknown_code_maps_to_unk(self): + out = self.proc.process([["NOT_IN_VOCAB"]]) + self.assertEqual(out[0].nonzero().flatten().tolist(), + [self.proc.code_vocab[""]]) + + def test_pad_column_never_set(self): + """Column 0 must stay clear, or padding becomes indistinguishable from + a real code.""" + out = self.proc.process([["A", "B"], [], ["C"]]) + self.assertEqual(out[:, 0].sum().item(), 0.0) + + def test_none_entries_treated_as_unknown(self): + out = self.proc.process([[None]]) + self.assertEqual(out[0].nonzero().flatten().tolist(), + [self.proc.code_vocab[""]]) + + def test_vocab_edit_helpers(self): + p = NestedMultiHotProcessor() + p.fit(self.samples, "codes") + p.add({"Z"}) + self.assertIn("Z", p.tokens()) + p.retain({"A"}) + self.assertEqual(p.tokens(), {"", "", "A"}) + # Width follows the vocabulary, so it must shrink with it. + self.assertEqual(p.process([["A"]]).shape, (1, 3)) + + +if __name__ == "__main__": + unittest.main() From db3112257fd26f013357dac85df6f12532fcf642 Mon Sep 17 00:00:00 2001 From: Jane Du Date: Wed, 9 Sep 2026 16:59:23 -0500 Subject: [PATCH 2/9] Satisfy PR contribution rules: lint, docstring examples, update HALO 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) --- examples/halo_mimic3.py | 14 +++--- .../processors/nested_multihot_processor.py | 9 ++-- pyhealth/tasks/generate_ehr.py | 44 +++++++++++++++---- 3 files changed, 49 insertions(+), 18 deletions(-) diff --git a/examples/halo_mimic3.py b/examples/halo_mimic3.py index 3ea0a71be..5f9c35ef9 100644 --- a/examples/halo_mimic3.py +++ b/examples/halo_mimic3.py @@ -3,7 +3,7 @@ This example demonstrates: 1. Loading MIMIC-III data 2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sequences) -3. Creating a SampleDataset with a NestedSequenceProcessor +3. Creating a SampleDataset with a NestedMultiHotProcessor 4. Training the HALO generator with its custom training loop 5. Generating synthetic patients 6. Evaluating the synthetic data with the generative metrics suite @@ -34,6 +34,7 @@ sample = sample_dataset[0] print("\nSample structure:") print(f" Patient ID: {sample['patient_id']}") + # (num_visits, vocab_size) multi-hot -- there is no padded inner axis. print(f" Visits tensor shape: {tuple(sample['visits'].shape)}") # STEP 3: Split dataset by patient @@ -85,12 +86,15 @@ } def real_subset_to_records(subset): + # NestedMultiHotProcessor encodes each patient as a dense + # (num_visits, vocab_size) multi-hot tensor, so the set of codes in a + # visit is the set of *nonzero column indices* -- not the tensor + # values themselves. for sample in subset: pid = str(sample["patient_id"]) - visits_tensor = sample["visits"] - for t, visit in enumerate(visits_tensor.tolist()): - for idx in visit: - code = index_to_code.get(int(idx)) + for t, visit in enumerate(sample["visits"]): + for idx in visit.nonzero(as_tuple=True)[0].tolist(): + code = index_to_code.get(idx) if code in (None, "", ""): continue yield {"id": pid, "time": t, "visit_codes": code, "labels": 0} diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py index eca991973..3b6c7e10e 100644 --- a/pyhealth/processors/nested_multihot_processor.py +++ b/pyhealth/processors/nested_multihot_processor.py @@ -1,4 +1,5 @@ -from typing import Any, Dict, Iterable, List +from collections.abc import Iterable +from typing import Any import torch @@ -58,11 +59,11 @@ def __init__(self, padding: int = 0): # `padding` is accepted and ignored so this is a drop-in swap for # NestedSequenceProcessor in a schema. There is no inner axis to pad -- # that is the entire point -- so honouring it would be misleading. - self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self.code_vocab: dict[Any, int] = {"": self.PAD, "": self.UNK} self._next_index = 2 self._padding = padding - def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: + def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: """Build the vocabulary. Inner length is irrelevant here, so unlike ``NestedSequenceProcessor`` nothing is measured about visit width. @@ -113,7 +114,7 @@ def tokens(self) -> set[str]: """Return the set of tokens in the processor's vocabulary.""" return set(self.code_vocab.keys()) - def process(self, value: List[List[Any]]) -> torch.Tensor: + def process(self, value: list[list[Any]]) -> torch.Tensor: """Nested sequence -> ``(num_visits, vocab_size)`` float multi-hot. Built with one ``scatter_`` per visit rather than per-code assignment, diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 8802f9a79..764e4de4a 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -63,7 +63,8 @@ """ import logging -from typing import Callable, Dict, List, Optional, Type, Union +from collections.abc import Callable +from typing import ClassVar from pyhealth.data.data import Patient from pyhealth.processors import NestedMultiHotProcessor @@ -93,19 +94,29 @@ class EHRGeneration(BaseTask): code_attr: Event attribute holding the code string. Default ``"icd9_code"``. min_visits: Minimum qualifying visits to keep a patient. Default 2. + + Examples: + >>> from pyhealth.datasets import MIMIC3Dataset + >>> from pyhealth.tasks import EHRGeneration + >>> ds = MIMIC3Dataset(root="...", tables=["diagnoses_icd"], dev=True) + >>> samples = ds.set_task(EHRGeneration()) + >>> samples[0]["visits"].shape # (num_visits, vocab_size) multi-hot + torch.Size([3, 512]) """ task_name: str = "ehr_generation" - input_schema: Dict[str, Union[str, Type]] = {"visits": NestedMultiHotProcessor} - output_schema: Dict[str, Union[str, Type]] = {} + input_schema: ClassVar[dict[str, str | type]] = { + "visits": NestedMultiHotProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" min_visits: int = 2 - def __call__(self, patient: Patient) -> List[Dict]: + def __call__(self, patient: Patient) -> list[dict]: """Extract the per-visit code sequence for a patient.""" - visits: List[List[str]] = [] + visits: list[list[str]] = [] admissions = patient.get_events(event_type="admissions") for admission in admissions: events = patient.get_events( @@ -167,7 +178,7 @@ class EHRGenerationMIMIC4(EHRGeneration): # ---------------------------------------------------------------------------- def to_evaluation_dataframe( records, - label_fn: Optional[Callable[[Dict], int]] = None, + label_fn: Callable[[dict], int] | None = None, subject_col: str = "id", visit_col: str = "time", code_col: str = "visit_codes", @@ -197,6 +208,15 @@ def to_evaluation_dataframe( Returns: ``pandas.DataFrame`` with columns ``[subject_col, visit_col, code_col, label_col]``. + + Examples: + >>> from pyhealth.tasks.generate_ehr import to_evaluation_dataframe + >>> records = [{"visits": [["4019", "25000"], ["4019"]]}] + >>> to_evaluation_dataframe(records) + id time visit_codes labels + 0 0 0 4019 0 + 1 0 0 25000 0 + 2 0 1 4019 0 """ import pandas as pd @@ -218,7 +238,7 @@ def to_evaluation_dataframe( ) -def decode_dataset(sample_dataset, feature_key: str = "visits") -> List[Dict]: +def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: """Decode a processed EHRGeneration ``SampleDataset`` back into records. Inverts the :class:`~pyhealth.processors.NestedMultiHotProcessor` encoding @@ -237,14 +257,20 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> List[Dict]: Returns: List of ``{"visits": [[code_str, ...], ...]}`` records. + + Examples: + >>> from pyhealth.tasks.generate_ehr import decode_dataset + >>> records = decode_dataset(samples) + >>> records[0]["visits"][0] + ['4019', '25000'] """ processor = sample_dataset.input_processors[feature_key] index_to_code = {idx: code for code, idx in processor.code_vocab.items()} - records: List[Dict] = [] + records: list[dict] = [] for i in range(len(sample_dataset)): sample = sample_dataset[i] - visits: List[List[str]] = [] + visits: list[list[str]] = [] # Each row is a multi-hot vector over the vocabulary, so the codes # present are its non-zero columns. (Reading the values as indices -- # which the index-based encoding required -- would see only 0s and 1s From b7e563ebd7250c9826b4f56eb3baba4aab47da0b Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 15:27:51 -0500 Subject: [PATCH 3/9] Drop the encode-equivalence test, keep the truncation coverage 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 --- tests/core/test_halo.py | 23 +++++ tests/core/test_halo_encode_equivalence.py | 112 --------------------- 2 files changed, 23 insertions(+), 112 deletions(-) delete mode 100644 tests/core/test_halo_encode_equivalence.py diff --git a/tests/core/test_halo.py b/tests/core/test_halo.py index 54a0e5246..f4abe96fc 100644 --- a/tests/core/test_halo.py +++ b/tests/core/test_halo.py @@ -59,6 +59,29 @@ def test_forward_input_format(self): self.assertIsInstance(batch["visits"], torch.Tensor) self.assertEqual(batch["visits"].dim(), 3) # (B, max_visits, vocab_size) + def test_encode_truncates_past_context(self): + """A patient with more visits than n_ctx-2 is cut, not wrapped.""" + many = [{"patient_id": "long", "visits": [["A05B"]] * 20}] + dataset = create_sample_dataset( + samples=many + self.samples, + input_schema=self.input_schema, + output_schema=self.output_schema, + dataset_name="test_halo_long", + ) + model = HALO( + dataset=dataset, + embed_dim=16, + n_heads=2, + n_layers=2, + n_ctx=8, + batch_size=2, + epochs=1, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + ehr, mask = model._encode_visits(batch["visits"]) + self.assertEqual(ehr.shape[1], 8) + self.assertEqual(mask.shape[1], 7) + def test_model_forward(self): """Forward returns a finite scalar loss and a probability tensor.""" loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) diff --git a/tests/core/test_halo_encode_equivalence.py b/tests/core/test_halo_encode_equivalence.py deleted file mode 100644 index 0bc3afcd3..000000000 --- a/tests/core/test_halo_encode_equivalence.py +++ /dev/null @@ -1,112 +0,0 @@ -"""HALO's vectorised encoder must match the loop it replaced, exactly. - -``_encode_visits`` used to walk (patient, visit, code slot) in Python over the -padded index tensor. On an A100 that cost ~108 ms per patient -- about 99.8% of -a training step, against ~0.03 s for the transformer's own forward and backward. -It is now a vectorised placement over per-visit multi-hot rows. - -That is a pure performance change, so the tensors handed to the transformer must -be bit-identical: anything else silently invalidates every existing checkpoint -and every published number. The legacy implementation is reproduced here rather -than imported, so this test keeps its meaning after the original is deleted. -""" - -import unittest - -import torch - -from pyhealth.datasets import create_sample_dataset, get_dataloader -from pyhealth.models import HALO -from pyhealth.processors import NestedSequenceProcessor - - -def legacy_encode(cfg, visits, device): - """The pre-vectorisation implementation, verbatim, over index input.""" - batch_size = visits.shape[0] - batch_ehr = torch.zeros(batch_size, cfg.n_ctx, cfg.total_vocab_size, - device=device) - batch_mask = torch.zeros(batch_size, cfg.n_ctx, 1, device=device) - start_idx = cfg.code_vocab_size + cfg.label_vocab_size - end_idx, pad_idx = start_idx + 1, start_idx + 2 - - for i in range(batch_size): - n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) - n_visits = min(n_visits, cfg.n_ctx - 2) - for j in range(n_visits): - for code_idx in visits[i, j]: - if code_idx > 0: - batch_ehr[i, j + 2, code_idx] = 1 - batch_mask[i, j + 2] = 1 - batch_ehr[i, 0, start_idx] = 1 - batch_ehr[i, n_visits + 1, end_idx] = 1 - batch_ehr[i, n_visits + 2:, pad_idx] = 1 - - return batch_ehr, batch_mask[:, 1:, :] - - -class TestHALOEncodeEquivalence(unittest.TestCase): - """Both encodings of the same patients must produce the same tensors.""" - - SAMPLES = [ - {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, - {"patient_id": "p1", "visits": [["A05B"], ["A04A", "B035"]]}, - {"patient_id": "p2", "visits": [["C129", "A11D"], ["A05C"], ["A04A"]]}, - {"patient_id": "p3", "visits": [["B035"]]}, - # Repeated codes: the index form stores five entries, the multi-hot form - # one bit. HALO collapsed them either way, which is why this matches. - {"patient_id": "p4", "visits": [["A05B", "A05B", "A05B"], ["C129"]]}, - ] - - def _datasets(self): - multihot = create_sample_dataset( - samples=self.SAMPLES, input_schema={"visits": "nested_multihot"}, - output_schema={}, dataset_name="mh") - index = create_sample_dataset( - samples=self.SAMPLES, input_schema={"visits": "nested_sequence"}, - output_schema={}, dataset_name="idx") - return multihot, index - - def test_bit_identical(self): - multihot, index = self._datasets() - model = HALO(dataset=multihot, embed_dim=16, n_heads=2, n_layers=2, - n_ctx=8, batch_size=len(self.SAMPLES), epochs=1) - - mh_batch = next(iter(get_dataloader(multihot, batch_size=len(self.SAMPLES)))) - ix_batch = next(iter(get_dataloader(index, batch_size=len(self.SAMPLES)))) - - new_ehr, new_mask = model._encode_visits(mh_batch["visits"]) - old_ehr, old_mask = legacy_encode(model.config, ix_batch["visits"], - model.device) - - self.assertTrue(torch.equal(new_ehr, old_ehr), - f"{(new_ehr != old_ehr).sum().item()} cells differ") - self.assertTrue(torch.equal(new_mask, old_mask)) - - def test_truncates_past_context(self): - """A patient with more visits than n_ctx-2 is cut, not wrapped.""" - many = [{"patient_id": "long", "visits": [["A05B"]] * 20}] - ds = create_sample_dataset(samples=many + self.SAMPLES, - input_schema={"visits": "nested_multihot"}, - output_schema={}, dataset_name="long") - model = HALO(dataset=ds, embed_dim=16, n_heads=2, n_layers=2, n_ctx=8, - batch_size=2, epochs=1) - batch = next(iter(get_dataloader(ds, batch_size=2))) - ehr, mask = model._encode_visits(batch["visits"]) - self.assertEqual(ehr.shape[1], 8) - self.assertEqual(mask.shape[1], 7) - - def test_no_inner_padding_in_multihot(self): - """The multi-hot row width is the vocabulary, not the longest visit. - - This is the memory win: on eICU the index form pads every visit to 3,951 - slots because one visit somewhere is that long. - """ - multihot, index = self._datasets() - mh_w = multihot.input_processors["visits"].vocab_size() - ix_w = index.input_processors["visits"]._max_inner_len - self.assertEqual(next(iter(multihot))["visits"].shape[1], mh_w) - self.assertEqual(next(iter(index))["visits"].shape[1], ix_w) - - -if __name__ == "__main__": - unittest.main() From 57aaa2b9518da9208d8095b4f17031ac1a711838 Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 15:42:24 -0500 Subject: [PATCH 4/9] Share the code vocabulary, and fix GPT2/PromptEHR under the new encoding 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 , so every code in every visit silently became : 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 --- docs/api/processors.rst | 6 + pyhealth/models/generators/gpt2.py | 30 +++-- pyhealth/models/generators/halo.py | 21 +++- pyhealth/models/generators/promptehr.py | 26 +++-- pyhealth/processors/base_processor.py | 82 +++++++++++++ .../deep_nested_sequence_processor.py | 45 +------- .../processors/nested_multihot_processor.py | 66 +++++------ .../processors/nested_sequence_processor.py | 65 ++++------- pyhealth/processors/sequence_processor.py | 35 +----- pyhealth/processors/stagenet_processor.py | 49 ++------ tests/core/test_generator_encodings.py | 108 ++++++++++++++++++ tests/core/test_vocab_processors.py | 64 +++++++++++ 12 files changed, 380 insertions(+), 217 deletions(-) create mode 100644 tests/core/test_generator_encodings.py diff --git a/docs/api/processors.rst b/docs/api/processors.rst index e4262fa4c..91609b1a9 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -27,6 +27,12 @@ Available Processors - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures) - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history) - ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models) + +Processors that map codes to indices share one vocabulary implementation, +``CodeVocabularyMixin`` (``pyhealth.processors.base_processor``): it provides +``add`` / ``remove`` / ``retain`` / ``tokens`` / ``vocab_size`` and the +````=0, ````=1 convention. A new code processor should mix it in +rather than reimplementing them. - ``NestedFloatsProcessor``: For nested numerical sequences with optional forward-fill **Label Processors:** diff --git a/pyhealth/models/generators/gpt2.py b/pyhealth/models/generators/gpt2.py index 99b3bf92a..aec0790e8 100644 --- a/pyhealth/models/generators/gpt2.py +++ b/pyhealth/models/generators/gpt2.py @@ -15,7 +15,7 @@ (``do_sample`` + top-k/top-p) and decodes it back into per-visit code lists, splitting on the ``[VISIT_DELIM]`` token. -The code vocabulary is taken from the dataset's ``NestedSequenceProcessor`` +The code vocabulary is taken from the dataset's ``visits`` processor (which already reserves index 0 for ```` and index 1 for ````); three special tokens (BOS, EOS, VISIT_DELIM) are appended, and ```` (index 0) is reused as the padding token. @@ -39,8 +39,8 @@ class GPT2(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedMultiHotProcessor}`` (or the equivalent + ``NestedSequenceProcessor``) and whose ``output_schema`` is empty. embed_dim: GPT-2 embedding dimension (``n_embd``). Must be divisible by ``n_heads``. Default: 512. n_heads: Number of attention heads. Default: 8. @@ -86,7 +86,17 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "GPT2 expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedSequenceProcessor or NestedMultiHotProcessor." + ) + if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): + # Without this the visit row would be read as raw values. Under a + # multi-hot encoding every value is 1.0, so every code would silently + # become and training would look fine while learning nothing. + raise ValueError( + f"GPT2 needs a 'visits' processor that can invert its own " + f"encoding (a visit_code_ids method); got " + f"{type(dataset.input_processors['visits']).__name__}. Use " + "NestedSequenceProcessor or NestedMultiHotProcessor." ) self.save_dir = save_dir @@ -95,7 +105,7 @@ def __init__( self._lr = lr self.max_len = max_len - # Code vocab from the NestedSequenceProcessor (includes =0, =1). + # Code vocab from the visits processor (includes =0, =1). self.visits_processor = dataset.input_processors["visits"] self.code_vocab_size = self.visits_processor.vocab_size() # Append three special tokens after the code vocab; reuse =0 as PAD. @@ -132,8 +142,8 @@ def _encode_visits(self, visits: torch.Tensor): """Flatten the padded visit-index tensor into causal-LM token streams. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. Index 0 is ```` and is + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. Index 0 is ```` and is skipped. Returns: @@ -147,7 +157,7 @@ def _encode_visits(self, visits: torch.Tensor): n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) seq: List[int] = [self.bos_id] for j in range(n_visits): - codes = [int(c) for c in visits[i, j].tolist() if c > 0] + codes = self.visits_processor.visit_code_ids(visits[i, j]) seq.extend(codes) if j < n_visits - 1: seq.append(self.delim_id) @@ -176,8 +186,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/models/generators/halo.py b/pyhealth/models/generators/halo.py index 255854ce3..a46725109 100644 --- a/pyhealth/models/generators/halo.py +++ b/pyhealth/models/generators/halo.py @@ -365,13 +365,13 @@ class HALO(BaseModel): The model infers its code vocabulary from the fitted ``SampleDataset``: ``code_vocab_size = dataset.input_processors["visits"].vocab_size()`` - (the ``NestedSequenceProcessor`` vocab, which already reserves index 0 for + (the ``NestedMultiHotProcessor`` vocab, which already reserves index 0 for ```` and index 1 for ````). Three special tokens are appended for start-of-sequence, end-of-sequence, and pad-visit. Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` + ``{"visits": NestedMultiHotProcessor}`` and whose ``output_schema`` is empty. embed_dim: Transformer embedding dimension (``n_embd``). Default: 768. n_heads: Number of attention heads. Must divide ``embed_dim``. @@ -420,7 +420,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "HALO expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedMultiHotProcessor." ) self.save_dir = save_dir @@ -428,7 +428,7 @@ def __init__( self._epochs = epochs self._lr = lr - # Code vocab from the NestedSequenceProcessor (includes , ). + # Code vocab from the NestedMultiHotProcessor (includes , ). self.visits_processor = dataset.input_processors["visits"] code_vocab_size = self.visits_processor.vocab_size() label_vocab_size = 0 # unconditional generation -- no output labels @@ -498,6 +498,15 @@ def _encode_visits(self, visits: torch.Tensor): # Real visits per patient, for the whole batch at once. An all-zero row # is an empty visit, exactly as a row of indices was before. + # + # TODO: this counts non-empty rows and then treats the *first* n_visits + # rows as the real ones, so a patient like [codes, empty, codes] would + # silently lose its last visit. The pre-vectorisation loop had the same + # behaviour, and EHRGeneration cannot produce an interior empty visit + # (its __call__ skips admissions with no codes), so nothing hits this + # today -- but NestedMultiHotProcessor does emit all-zero rows for empty + # visits, so a hand-built SampleDataset can. Fix by masking on row + # position rather than row count. n_visits = (visits.sum(dim=-1) > 0).sum(dim=1) n_visits = n_visits.clamp(max=cfg.n_ctx - 2) # (batch,) @@ -532,8 +541,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass. Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: FloatTensor ``(batch, max_visits, code_vocab_size)`` from + the ``NestedMultiHotProcessor``, 1.0 where a code is present. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/models/generators/promptehr.py b/pyhealth/models/generators/promptehr.py index 4622e72c6..e719cd78e 100644 --- a/pyhealth/models/generators/promptehr.py +++ b/pyhealth/models/generators/promptehr.py @@ -38,7 +38,7 @@ each with its own modality prompt token; the PyHealth ``EHRGeneration`` task exposes a single ``visits`` modality, so a single ``[CODE_PROMPT]`` token marks it. The code vocabulary is taken from the dataset's -``NestedSequenceProcessor`` (which already reserves index 0 for ```` and +``visits`` processor (which already reserves index 0 for ```` and index 1 for ````); five special tokens (BOS, EOS, VISIT_DELIM, MASK, CODE_PROMPT) are appended, and ```` (index 0) is reused as the pad token. """ @@ -67,8 +67,8 @@ class PromptEHR(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedSequenceProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedMultiHotProcessor}`` (or the equivalent + ``NestedSequenceProcessor``) and whose ``output_schema`` is empty. embed_dim: BART model dimension (``d_model``). Must be divisible by ``n_heads``. Default: 256. n_heads: Number of attention heads (encoder and decoder). Default: 8. @@ -127,7 +127,17 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "PromptEHR expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor." + "NestedSequenceProcessor or NestedMultiHotProcessor." + ) + if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): + # Without this the visit row would be read as raw values. Under a + # multi-hot encoding every value is 1.0, so every code would silently + # become and training would look fine while learning nothing. + raise ValueError( + f"PromptEHR needs a 'visits' processor that can invert its own " + f"encoding (a visit_code_ids method); got " + f"{type(dataset.input_processors['visits']).__name__}. Use " + "NestedSequenceProcessor or NestedMultiHotProcessor." ) self.save_dir = save_dir @@ -139,7 +149,7 @@ def __init__( self.mean_span_len = mean_span_len self.prompt_length = prompt_length - # Code vocab from the NestedSequenceProcessor (includes =0, =1). + # Code vocab from the visits processor (includes =0, =1). self.visits_processor = dataset.input_processors["visits"] self.code_vocab_size = self.visits_processor.vocab_size() # Append five special tokens after the code vocab; reuse =0 as PAD. @@ -197,7 +207,7 @@ def _serialize(self, visits: torch.Tensor) -> List[List[int]]: n_visits = int((visits[i].sum(dim=-1) > 0).sum().item()) seq: List[int] = [self.code_prompt_id] for j in range(n_visits): - codes = [int(c) for c in visits[i, j].tolist() if c > 0] + codes = self.visits_processor.visit_code_ids(visits[i, j]) seq.extend(codes) if j < n_visits - 1: seq.append(self.delim_id) @@ -313,8 +323,8 @@ def forward(self, visits: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: """Forward pass (denoising seq2seq reconstruction). Args: - visits: LongTensor ``(batch, max_visits, max_codes_per_visit)`` from - the ``NestedSequenceProcessor``. + visits: Processed visit tensor from either nested ``visits`` + processor; the processor's ``visit_code_ids`` inverts a row. **kwargs: Any other batch keys are ignored. Returns: diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 06bbe0c7c..7233fe779 100644 --- a/pyhealth/processors/base_processor.py +++ b/pyhealth/processors/base_processor.py @@ -196,6 +196,88 @@ def vocab_size(self) -> int: pass +class CodeVocabularyMixin(TokenProcessorInterface): + """Concrete ````/```` code vocabulary, shared by code processors. + + Every processor that maps medical codes to indices needs the same four + operations -- build, add, remove, retain -- and each one used to carry its + own copy. ``SequenceProcessor``, ``NestedSequenceProcessor``, + ``DeepNestedSequenceProcessor``, ``NestedMultiHotProcessor`` and + ``StageNetProcessor`` held five verbatim duplicates of the same code. + + Deliberately a mixin and not a base class. Several models dispatch on + ``isinstance(processor, NestedSequenceProcessor)`` to decide whether to + apply ``nn.Embedding``, so making one code processor inherit from another + would silently reroute it. Sharing the implementation without touching the + inheritance chain keeps that dispatch honest. + + Subclasses call :meth:`_init_code_vocab` from ``__init__`` and + :meth:`_observe_code` from their own ``fit`` traversal, which differs by + nesting depth and so is not shared here. + + Examples: + >>> class MyProcessor(FeatureProcessor, CodeVocabularyMixin): + ... def __init__(self): + ... self._init_code_vocab() + ... def fit(self, samples, field): + ... for sample in samples: + ... for code in sample[field]: + ... self._observe_code(code) + ... def process(self, value): + ... return [self.code_vocab.get(c, self.UNK) for c in value] + >>> processor = MyProcessor() + >>> processor.fit([{"codes": ["A", "B"]}], "codes") + >>> processor.vocab_size() + 4 + >>> processor.code_vocab["A"] + 2 + """ + + code_vocab: dict[Any, int] + + def _init_code_vocab(self) -> None: + """Seed the vocabulary with ```` (0) and ```` (1).""" + self.code_vocab = {"": self.PAD, "": self.UNK} + self._next_index = 2 + + def _observe_code(self, code: Any) -> None: + """Add one code to the vocabulary if it is new and not ``None``.""" + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 + + def _reindex(self, keep: set) -> None: + """Rebuild the vocabulary over ``keep``, preserving relative order.""" + order = [ + k + for k, _ in sorted(self.code_vocab.items(), key=lambda kv: kv[1]) + if k in keep + ] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def remove(self, tokens: set[str]) -> None: + """Remove specified vocabularies from the processor.""" + self._reindex(set(self.code_vocab.keys()) - tokens | {"", ""}) + + def retain(self, tokens: set[str]) -> None: + """Retain only the specified vocabularies in the processor.""" + self._reindex(set(self.code_vocab.keys()) & tokens | {"", ""}) + + def add(self, tokens: set[str]) -> None: + """Add specified vocabularies to the processor.""" + for token in tokens: + self._observe_code(token) + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + + def vocab_size(self) -> int: + """Return the size of the processor's vocabulary.""" + return len(self.code_vocab) + + class TemporalFeatureProcessor(FeatureProcessor): """Abstract base class for processors whose features are paired with timestamps. diff --git a/pyhealth/processors/deep_nested_sequence_processor.py b/pyhealth/processors/deep_nested_sequence_processor.py index 24683f54b..2d14f6350 100644 --- a/pyhealth/processors/deep_nested_sequence_processor.py +++ b/pyhealth/processors/deep_nested_sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import FeatureProcessor, TokenProcessorInterface +from .base_processor import CodeVocabularyMixin, FeatureProcessor @register_processor("deep_nested_sequence") -class DeepNestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): +class DeepNestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """ Feature processor for deeply nested categorical sequences with vocabulary. @@ -45,8 +45,7 @@ class DeepNestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): """ def __init__(self): - self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} - self._next_index = 2 + self._init_code_vocab() self._max_middle_len = 1 # Maximum length of middle sequences (e.g. visits) self._max_inner_len = 1 # Maximum length of inner sequences (e.g. codes per visit) @@ -78,39 +77,11 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Build vocabulary for code in inner_seq: - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 + self._observe_code(code) self._max_middle_len = max(1, max_middle_len) self._max_inner_len = max(1, max_inner_len) - def remove(self, tokens: set[str]): - """Remove specified vocabularies from the processor.""" - keep = set(self.code_vocab.keys()) - tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def retain(self, tokens: set[str]): - """Retain only the specified vocabularies in the processor.""" - keep = set(self.code_vocab.keys()) & tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def add(self, tokens: set[str]): - """Add specified vocabularies to the processor.""" - i = len(self.code_vocab) - for token in tokens: - if token not in self.code_vocab: - self.code_vocab[token] = i - i += 1 - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) - def process(self, value: List[List[List[Any]]]) -> torch.Tensor: """Process deep nested sequence into padded 3D tensor. @@ -173,18 +144,10 @@ def process(self, value: List[List[List[Any]]]) -> torch.Tensor: return torch.tensor(encoded_groups, dtype=torch.long) - def vocab_size(self) -> int: - """Return the size of the processor's vocabulary.""" - return len(self.code_vocab) - def size(self) -> int: """Return max inner length (embedding dimension) for unified API.""" return self._max_inner_len - def vocab_size(self) -> int: - """Return vocabulary size.""" - return len(self.code_vocab) - def __repr__(self): return ( f"DeepNestedSequenceProcessor(" diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py index 3b6c7e10e..046ea6e62 100644 --- a/pyhealth/processors/nested_multihot_processor.py +++ b/pyhealth/processors/nested_multihot_processor.py @@ -4,11 +4,11 @@ import torch from . import register_processor -from .base_processor import FeatureProcessor, TokenProcessorInterface +from .base_processor import CodeVocabularyMixin, FeatureProcessor @register_processor("nested_multihot") -class NestedMultiHotProcessor(FeatureProcessor, TokenProcessorInterface): +class NestedMultiHotProcessor(FeatureProcessor, CodeVocabularyMixin): """Nested categorical sequences as per-visit multi-hot vectors. Same input as :class:`NestedSequenceProcessor` -- a list of visits, each a @@ -59,8 +59,7 @@ def __init__(self, padding: int = 0): # `padding` is accepted and ignored so this is a drop-in swap for # NestedSequenceProcessor in a schema. There is no inner axis to pad -- # that is the entire point -- so honouring it would be misleading. - self.code_vocab: dict[Any, int] = {"": self.PAD, "": self.UNK} - self._next_index = 2 + self._init_code_vocab() self._padding = padding def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: @@ -81,38 +80,7 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: if not isinstance(inner_seq, list): continue for code in inner_seq: - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 - - def remove(self, tokens: set[str]): - """Remove specified vocabularies from the processor.""" - keep = set(self.code_vocab.keys()) - tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) - if k in keep] - self.code_vocab = {k: i for i, k in enumerate(order)} - self._next_index = len(self.code_vocab) - - def retain(self, tokens: set[str]): - """Retain only the specified vocabularies in the processor.""" - keep = set(self.code_vocab.keys()) & tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) - if k in keep] - self.code_vocab = {k: i for i, k in enumerate(order)} - self._next_index = len(self.code_vocab) - - def add(self, tokens: set[str]): - """Add specified vocabularies to the processor.""" - i = len(self.code_vocab) - for token in tokens: - if token not in self.code_vocab: - self.code_vocab[token] = i - i += 1 - self._next_index = len(self.code_vocab) - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) + self._observe_code(code) def process(self, value: list[list[Any]]) -> torch.Tensor: """Nested sequence -> ``(num_visits, vocab_size)`` float multi-hot. @@ -149,14 +117,32 @@ def process(self, value: list[list[Any]]) -> torch.Tensor: torch.ones(len(idx))) return out + def visit_code_ids(self, row: torch.Tensor) -> list[int]: + """Code indices present in one processed visit row. + + The inverse of what :meth:`process` writes. Here the row is a multi-hot + vector, so the codes are its *nonzero column indices* -- reading the + values themselves would yield 1.0 (i.e. ````) for every code. + + Codes come back in vocabulary order, not charted order, and repeats are + already collapsed; multi-hot records presence, not sequence or count. + + Args: + row: 1D multi-hot tensor of width ``vocab_size``, one visit. + + Returns: + Code indices present in the visit, ascending. + + Examples: + >>> processor.visit_code_ids(torch.tensor([0., 0., 1., 0., 1.])) + [2, 4] + """ + return row.nonzero(as_tuple=True)[0].tolist() + def size(self) -> int: """Feature width: the vocabulary, since that is the row length.""" return len(self.code_vocab) - def vocab_size(self) -> int: - """Return vocabulary size.""" - return len(self.code_vocab) - def __repr__(self): return f"NestedMultiHotProcessor(vocab_size={len(self.code_vocab)})" diff --git a/pyhealth/processors/nested_sequence_processor.py b/pyhealth/processors/nested_sequence_processor.py index c03c800d0..72690f4d6 100644 --- a/pyhealth/processors/nested_sequence_processor.py +++ b/pyhealth/processors/nested_sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import FeatureProcessor, TokenProcessorInterface +from .base_processor import CodeVocabularyMixin, FeatureProcessor @register_processor("nested_sequence") -class NestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): +class NestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """ Feature processor for nested categorical sequences with vocabulary. @@ -45,8 +45,7 @@ class NestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): """ def __init__(self, padding: int = 0): - self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} - self._next_index = 2 + self._init_code_vocab() self._max_inner_len = 1 # Maximum length of inner sequences self._padding = padding # Additional padding beyond observed max @@ -72,41 +71,13 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Build vocabulary for code in inner_seq: - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 + self._observe_code(code) # Store max inner length: add user-specified padding to observed maximum # This ensures the processor can handle sequences longer than those in training data observed_max = max(1, max_inner_len) self._max_inner_len = observed_max + self._padding - def remove(self, tokens: set[str]): - """Remove specified vocabularies from the processor.""" - keep = set(self.code_vocab.keys()) - tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def retain(self, tokens: set[str]): - """Retain only the specified vocabularies in the processor.""" - keep = set(self.code_vocab.keys()) & tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def add(self, tokens: set[str]): - """Add specified vocabularies to the processor.""" - i = len(self.code_vocab) - for token in tokens: - if token not in self.code_vocab: - self.code_vocab[token] = i - i += 1 - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) - def process(self, value: List[List[Any]]) -> torch.Tensor: """Process nested sequence into padded 2D tensor. @@ -150,18 +121,32 @@ def process(self, value: List[List[Any]]) -> torch.Tensor: return torch.tensor(encoded_sequences, dtype=torch.long) - def vocab_size(self) -> int: - """Return the size of the processor's vocabulary.""" - return len(self.code_vocab) + def visit_code_ids(self, row: torch.Tensor) -> list[int]: + """Code indices present in one processed visit row. + + The inverse of what :meth:`process` writes, for consumers that need a + code *list* rather than the tensor -- sequence generators such as GPT2 + and PromptEHR. Each nested processor implements this for its own + encoding, so a model can accept either without knowing which it has. + + Here the row already holds indices, right-padded with ```` (0). + + Args: + row: 1D tensor, one processed visit. + + Returns: + Code indices in charted order, ```` dropped. + + Examples: + >>> processor.visit_code_ids(torch.tensor([2, 4, 0])) + [2, 4] + """ + return [int(c) for c in row.tolist() if c > 0] def size(self) -> int: """Return max inner length (embedding dimension) for unified API.""" return self._max_inner_len - def vocab_size(self) -> int: - """Return vocabulary size.""" - return len(self.code_vocab) - def __repr__(self): return ( f"NestedSequenceProcessor(" diff --git a/pyhealth/processors/sequence_processor.py b/pyhealth/processors/sequence_processor.py index 47339eefd..ad0a2a0ef 100644 --- a/pyhealth/processors/sequence_processor.py +++ b/pyhealth/processors/sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import FeatureProcessor, TokenProcessorInterface +from .base_processor import CodeVocabularyMixin, FeatureProcessor @register_processor("sequence") -class SequenceProcessor(FeatureProcessor, TokenProcessorInterface): +class SequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """Feature processor for encoding categorical sequences. Encodes medical codes (e.g., diagnoses, procedures) into numerical @@ -29,8 +29,7 @@ class SequenceProcessor(FeatureProcessor, TokenProcessorInterface): """ def __init__(self, code_mapping: Optional[Tuple[str, str]] = None): - self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} - self._next_index = 2 + self._init_code_vocab() self._mapper = None if code_mapping is not None: from pyhealth.medcode import CrossMap @@ -83,34 +82,6 @@ def process(self, value: Any) -> torch.Tensor: indices.append(self.code_vocab[""]) return torch.tensor(indices, dtype=torch.long) - - def remove(self, tokens: set[str]): - """Remove specified vocabularies from the processor.""" - keep = set(self.code_vocab.keys()) - tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - self.code_vocab = { k : i for i, k in enumerate(order) } - - def retain(self, tokens: set[str]): - """Retain only the specified vocabularies in the processor.""" - keep = set(self.code_vocab.keys()) & tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - self.code_vocab = { k : i for i, k in enumerate(order) } - - def add(self, tokens: set[str]): - """Add specified vocabularies to the processor.""" - i = len(self.code_vocab) - for token in tokens: - if token not in self.code_vocab: - self.code_vocab[token] = i - i += 1 - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) - - def vocab_size(self) -> int: - """Return the size of the processor's vocabulary.""" - return len(self.code_vocab) def size(self): return len(self.code_vocab) diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py index 604376ec1..ce302ef36 100644 --- a/pyhealth/processors/stagenet_processor.py +++ b/pyhealth/processors/stagenet_processor.py @@ -3,11 +3,15 @@ import torch from . import register_processor -from .base_processor import FeatureProcessor, ModalityType, TemporalFeatureProcessor, TokenProcessorInterface +from .base_processor import ( + CodeVocabularyMixin, + ModalityType, + TemporalFeatureProcessor, +) @register_processor("stagenet") -class StageNetProcessor(TemporalFeatureProcessor, TokenProcessorInterface): +class StageNetProcessor(TemporalFeatureProcessor, CodeVocabularyMixin): """ Feature processor for StageNet CODE inputs with coupled value/time data. @@ -55,8 +59,7 @@ class StageNetProcessor(TemporalFeatureProcessor, TokenProcessorInterface): """ def __init__(self, padding: int = 0): - self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} - self._next_index = 2 + self._init_code_vocab() self._is_nested = None # Will be determined during fit # Max inner sequence length for nested codes self._max_nested_len = None @@ -101,15 +104,11 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Track max inner length max_inner_len = max(max_inner_len, len(inner_list)) for code in inner_list: - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 + self._observe_code(code) else: # Flat codes for code in value_data: - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 + self._observe_code(code) # Store max nested length: add user-specified padding to observed maximum # This ensures the processor can handle sequences longer than those in training data @@ -117,32 +116,6 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: observed_max = max(1, max_inner_len) self._max_nested_len = observed_max + self._padding - def remove(self, tokens: set[str]): - """Remove specified vocabularies from the processor.""" - keep = set(self.code_vocab.keys()) - tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def retain(self, tokens: set[str]): - """Retain only the specified vocabularies in the processor.""" - keep = set(self.code_vocab.keys()) & tokens | {"", ""} - order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] - - self.code_vocab = { k : i for i, k in enumerate(order) } - - def add(self, tokens: set[str]): - """Add specified vocabularies to the processor.""" - i = len(self.code_vocab) - for token in tokens: - if token not in self.code_vocab: - self.code_vocab[token] = i - i += 1 - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) - def process( self, value: Tuple[Optional[List], List] ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: @@ -221,10 +194,6 @@ def _encode_nested_codes(self, nested_codes: List[List[str]]) -> torch.Tensor: return torch.tensor(encoded_sequences, dtype=torch.long) - def vocab_size(self) -> int: - """Return the size of the processor's vocabulary.""" - return len(self.code_vocab) - def size(self) -> int: """Return vocabulary size.""" return len(self.code_vocab) diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py new file mode 100644 index 000000000..7abc34bb4 --- /dev/null +++ b/tests/core/test_generator_encodings.py @@ -0,0 +1,108 @@ +"""Every generator sharing EHRGeneration must read the task's encoding. + +HALO, GPT2 and PromptEHR all consume ``EHRGeneration``, so changing that task's +``input_schema`` changes what all three receive. The failure mode this guards +is silent: a multi-hot row read as raw values yields 1.0 for every present +code, and 1 is ````, so a model would train happily on nothing but unknown +codes. These tests assert the real codes survive, under either encoding. +""" + +import unittest + +import torch + +from pyhealth.datasets import create_sample_dataset +from pyhealth.models import GPT2, PromptEHR + +SAMPLES = [ + {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, + {"patient_id": "p1", "visits": [["A05B"], ["A04A", "B035"]]}, + {"patient_id": "p2", "visits": [["C129", "A11D"], ["A05C"], ["A04A"]]}, + {"patient_id": "p3", "visits": [["B035"], ["A05B", "C129"]]}, +] + + +def _dataset(schema_key, name): + return create_sample_dataset( + samples=SAMPLES, + input_schema={"visits": schema_key}, + output_schema={}, + dataset_name=name, + ) + + +class TestVisitCodeIds(unittest.TestCase): + """Each nested processor inverts its own encoding to the same code ids.""" + + def test_both_processors_agree(self): + multihot = _dataset("nested_multihot", "vci_mh") + indexed = _dataset("nested_sequence", "vci_ix") + mh_proc = multihot.input_processors["visits"] + ix_proc = indexed.input_processors["visits"] + + # Same samples, same traversal order, so the vocabularies must match -- + # that is what makes the per-visit comparison below meaningful. + self.assertEqual(mh_proc.code_vocab, ix_proc.code_vocab) + + for i in range(len(SAMPLES)): + mh_row = multihot[i]["visits"] + ix_row = indexed[i]["visits"] + for visit in range(mh_row.shape[0]): + # Multi-hot returns vocabulary order, the index form charted + # order, so compare as sets. + self.assertEqual( + set(mh_proc.visit_code_ids(mh_row[visit])), + set(ix_proc.visit_code_ids(ix_row[visit])), + ) + + def test_multihot_ids_are_not_all_unk(self): + """The specific regression: reading values instead of column indices.""" + processor = _dataset("nested_multihot", "vci_unk").input_processors["visits"] + row = processor.process([["A05B", "A05C"]])[0] + ids = processor.visit_code_ids(row) + self.assertEqual(len(ids), 2) + self.assertNotIn(processor.UNK, ids) + + +class TestGeneratorsAcceptEitherEncoding(unittest.TestCase): + """GPT2 and PromptEHR serialise real codes from either processor.""" + + MODELS = [ + (GPT2, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64}), + (PromptEHR, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64, + "prompt_length": 4}), + ] + + def _streams(self, model, visits): + if hasattr(model, "_serialize"): + return model._serialize(visits) + input_ids, _, _ = model._encode_visits(visits) + return [row.tolist() for row in input_ids] + + def test_codes_survive_serialisation(self): + for cls, kwargs in self.MODELS: + with self.subTest(model=cls.__name__): + dataset = _dataset("nested_multihot", f"gen_{cls.__name__}") + model = cls(dataset=dataset, batch_size=2, epochs=1, **kwargs) + visits = torch.stack([dataset[i]["visits"] for i in range(2)]) + streams = self._streams(model, visits) + + code_ids = [ + t + for stream in streams + for t in stream + if t < model.code_vocab_size and t != 0 + ] + self.assertTrue(code_ids, "no code tokens were emitted at all") + # The bug turned every code into ; a real stream carries + # several distinct codes. + self.assertGreater( + len(set(code_ids) - {model.visits_processor.UNK}), + 1, + f"{cls.__name__} emitted only : the visit row was read " + "as raw values instead of via visit_code_ids", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_vocab_processors.py b/tests/core/test_vocab_processors.py index b6d785805..b167f8dcf 100644 --- a/tests/core/test_vocab_processors.py +++ b/tests/core/test_vocab_processors.py @@ -5,6 +5,7 @@ StageNetProcessor, NestedSequenceProcessor, DeepNestedSequenceProcessor, + NestedMultiHotProcessor, ) class TestVocabProcessors(unittest.TestCase): @@ -258,5 +259,68 @@ def test_deep_nested_sequence_processor_add(self): e_idx = processor.code_vocab["E"] self.assertEqual(res[0, 0, 0].item(), e_idx) +class TestSharedCodeVocabulary(unittest.TestCase): + """CodeVocabularyMixin gives every code processor the same vocabulary. + + Before the mixin, five processors carried verbatim copies of these methods + and every copy had the same defect: a removal renumbered the vocabulary but + left ``_next_index`` untouched, so the next ``fit`` allocated past the end + of it. Sharing one implementation fixes all five at once, which is what + these tests pin. + """ + + # Each processor takes a different nesting depth, so the same code list is + # wrapped to match. + PROCESSORS = [ + (SequenceProcessor, lambda codes: codes), + (NestedSequenceProcessor, lambda codes: [codes]), + (DeepNestedSequenceProcessor, lambda codes: [[codes]]), + (NestedMultiHotProcessor, lambda codes: [codes]), + ] + + def test_fit_after_remove_stays_in_range(self): + """The bug: ``fit`` allocates from ``_next_index``, ``remove`` renumbers. + + Removing codes renumbered the vocabulary to 0..n-1 but left + ``_next_index`` at its pre-removal value, so the next ``fit`` handed + out an index past the end of the vocabulary. Anything sizing an + embedding table by ``vocab_size()`` would then index out of bounds. + """ + for cls, wrap in self.PROCESSORS: + with self.subTest(processor=cls.__name__): + processor = cls() + processor.fit([{"codes": wrap(["A", "B", "C"])}], "codes") + processor.remove({"A", "B"}) + processor.fit([{"codes": wrap(["D"])}], "codes") + + indices = list(processor.code_vocab.values()) + self.assertEqual( + max(indices), processor.vocab_size() - 1, + f"{cls.__name__} allocated an index past the vocabulary: " + f"{processor.code_vocab}", + ) + self.assertEqual(len(indices), len(set(indices))) + + def test_indices_stay_contiguous_from_zero(self): + for cls, _ in self.PROCESSORS: + with self.subTest(processor=cls.__name__): + processor = cls() + processor.add({"A", "B", "C"}) + processor.retain({"B"}) + self.assertEqual( + sorted(processor.code_vocab.values()), + list(range(processor.vocab_size())), + ) + + def test_special_tokens_survive_every_edit(self): + for cls, _ in self.PROCESSORS: + with self.subTest(processor=cls.__name__): + processor = cls() + processor.add({"A"}) + processor.retain(set()) + self.assertEqual(processor.code_vocab[""], processor.PAD) + self.assertEqual(processor.code_vocab[""], processor.UNK) + + if __name__ == "__main__": unittest.main() From 83ac1f6b8deb18132aec09b30496ee74b5af10bc Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 15:58:14 -0500 Subject: [PATCH 5/9] Give each generator family its own EHR-generation task 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 --- docs/api/tasks.rst | 2 +- .../api/tasks/pyhealth.tasks.generate_ehr.rst | 44 +++- examples/halo_mimic3.py | 7 +- pyhealth/models/generators/corgan.py | 5 +- pyhealth/models/generators/gpt2.py | 15 +- pyhealth/models/generators/halo.py | 11 +- pyhealth/models/generators/medgan.py | 5 +- pyhealth/models/generators/promptehr.py | 17 +- pyhealth/tasks/__init__.py | 3 + pyhealth/tasks/generate_ehr.py | 222 ++++++++++++++---- tests/core/test_generator_encodings.py | 74 +++++- 11 files changed, 316 insertions(+), 89 deletions(-) diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 97260793d..4f616d7aa 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -141,7 +141,7 @@ a quick reference: - Cumulative visit history (drug recommendation, readmission) * - ``"nested_multihot"`` - ``NestedMultiHotProcessor`` - - Per-visit code sets (generative EHR) + - Per-visit code sets (HALO; see ``VisitMultiHotGeneration``) * - ``"tensor"`` - ``TensorProcessor`` - Aggregated numeric values (e.g. last lab value per item) diff --git a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst index 77c332f33..0ec40d5a4 100644 --- a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst +++ b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst @@ -1,10 +1,31 @@ pyhealth.tasks.generate_ehr =========================================== -Task that turns a longitudinal EHR dataset into per-patient, per-visit code -sequences for training unconditional synthetic-EHR generators (HALO, GPT2, -PromptEHR, MedGAN, CorGAN), plus helpers to flatten generated output into the -long-form dataframe consumed by :mod:`pyhealth.metrics.generative`. +Tasks that turn a longitudinal EHR dataset into training samples for +unconditional synthetic-EHR generators, plus helpers to flatten generated +output into the long-form dataframe consumed by +:mod:`pyhealth.metrics.generative`. + +Extraction is shared; the encoding is not. Each generator family reads its +codes in a different shape, and handing a model the wrong shape fails silently +rather than loudly, so pick the task that matches the model: + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - Task + - Encoding + - Models + * - ``VisitMultiHotGeneration`` + - one multi-hot row per visit + - HALO + * - ``VisitSequenceGeneration`` + - per-visit code indices + - GPT2, PromptEHR + * - ``PatientCodeSetGeneration`` + - one code set per patient + - MedGAN, CorGAN Task Classes ------------ @@ -14,6 +35,21 @@ Task Classes :undoc-members: :show-inheritance: +.. autoclass:: pyhealth.tasks.generate_ehr.VisitMultiHotGeneration + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.tasks.generate_ehr.VisitSequenceGeneration + :members: + :undoc-members: + :show-inheritance: + +.. autoclass:: pyhealth.tasks.generate_ehr.PatientCodeSetGeneration + :members: + :undoc-members: + :show-inheritance: + .. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3 :members: :undoc-members: diff --git a/examples/halo_mimic3.py b/examples/halo_mimic3.py index 5f9c35ef9..ab6133a0a 100644 --- a/examples/halo_mimic3.py +++ b/examples/halo_mimic3.py @@ -2,7 +2,9 @@ This example demonstrates: 1. Loading MIMIC-III data -2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sequences) +2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sets, + multi-hot -- the encoding HALO consumes; GPT2/PromptEHR instead take + 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 @@ -25,7 +27,8 @@ ) # STEP 2: Apply the EHR generation task (unconditional, no labels). - # This task is shared by all generators in pyhealth.models.generators. + # Extraction is shared across the generators; the encoding is not, so this + # is the multi-hot variant that HALO reads directly. sample_dataset = base_dataset.set_task(EHRGenerationMIMIC3()) print(f"Total samples: {len(sample_dataset)}") print(f"Input schema: {sample_dataset.input_schema}") diff --git a/pyhealth/models/generators/corgan.py b/pyhealth/models/generators/corgan.py index 807ed8066..c1636b8b8 100644 --- a/pyhealth/models/generators/corgan.py +++ b/pyhealth/models/generators/corgan.py @@ -7,7 +7,8 @@ ``dataset -> SampleDataset -> model`` pipeline. CorGAN treats each patient as a flat bag-of-codes (no visit structure), so it -expects an input feature named ``visits`` backed by a ``MultiHotProcessor``. +expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, +which the :class:`~pyhealth.tasks.PatientCodeSetGeneration` task provides. Training has two phases (mirroring the reference): * a **convolutional autoencoder** is pre-trained with a sparse-friendly BCE @@ -350,7 +351,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "CorGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor." + "MultiHotProcessor (see PatientCodeSetGeneration)." ) self._batch_size = batch_size diff --git a/pyhealth/models/generators/gpt2.py b/pyhealth/models/generators/gpt2.py index aec0790e8..e2b24b7a7 100644 --- a/pyhealth/models/generators/gpt2.py +++ b/pyhealth/models/generators/gpt2.py @@ -3,8 +3,10 @@ A simple decoder-only baseline that mirrors the standalone reference script ``generate_synthetic_mimic3_gpt2.py`` (``--mode transformer_baseline``) but plugged into the standard PyHealth ``dataset -> set_task -> SampleDataset -> -model`` pipeline. It consumes the same :class:`~pyhealth.tasks.EHRGeneration` -task as :class:`~pyhealth.models.HALO`. +model`` pipeline. It consumes the +:class:`~pyhealth.tasks.VisitSequenceGeneration` task -- the same extraction +:class:`~pyhealth.models.HALO` uses, but emitting code indices rather than +multi-hot rows, since a causal LM reads token ids. Each patient's visits are flattened into a single token stream:: @@ -39,8 +41,9 @@ class GPT2(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedMultiHotProcessor}`` (or the equivalent - ``NestedSequenceProcessor``) and whose ``output_schema`` is empty. + ``{"visits": NestedSequenceProcessor}`` -- use the + :class:`~pyhealth.tasks.VisitSequenceGeneration` task -- and whose + ``output_schema`` is empty. embed_dim: GPT-2 embedding dimension (``n_embd``). Must be divisible by ``n_heads``. Default: 512. n_heads: Number of attention heads. Default: 8. @@ -86,7 +89,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "GPT2 expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor or NestedMultiHotProcessor." + "NestedSequenceProcessor (see VisitSequenceGeneration)." ) if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): # Without this the visit row would be read as raw values. Under a @@ -96,7 +99,7 @@ def __init__( f"GPT2 needs a 'visits' processor that can invert its own " f"encoding (a visit_code_ids method); got " f"{type(dataset.input_processors['visits']).__name__}. Use " - "NestedSequenceProcessor or NestedMultiHotProcessor." + "NestedSequenceProcessor, via the VisitSequenceGeneration task." ) self.save_dir = save_dir diff --git a/pyhealth/models/generators/halo.py b/pyhealth/models/generators/halo.py index a46725109..b3efcd10f 100644 --- a/pyhealth/models/generators/halo.py +++ b/pyhealth/models/generators/halo.py @@ -371,8 +371,9 @@ class HALO(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedMultiHotProcessor}`` and whose ``output_schema`` - is empty. + ``{"visits": NestedMultiHotProcessor}`` -- use the + :class:`~pyhealth.tasks.VisitMultiHotGeneration` task -- and whose + ``output_schema`` is empty. embed_dim: Transformer embedding dimension (``n_embd``). Default: 768. n_heads: Number of attention heads. Must divide ``embed_dim``. Default: 12. @@ -420,7 +421,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "HALO expects an input feature named 'visits' backed by a " - "NestedMultiHotProcessor." + "NestedMultiHotProcessor (see VisitMultiHotGeneration)." ) self.save_dir = save_dir @@ -502,8 +503,8 @@ def _encode_visits(self, visits: torch.Tensor): # TODO: this counts non-empty rows and then treats the *first* n_visits # rows as the real ones, so a patient like [codes, empty, codes] would # silently lose its last visit. The pre-vectorisation loop had the same - # behaviour, and EHRGeneration cannot produce an interior empty visit - # (its __call__ skips admissions with no codes), so nothing hits this + # behaviour, and VisitMultiHotGeneration cannot produce an interior + # empty visit (its __call__ skips codeless admissions), so nothing hits # today -- but NestedMultiHotProcessor does emit all-zero rows for empty # visits, so a hand-built SampleDataset can. Fix by masking on row # position rather than row count. diff --git a/pyhealth/models/generators/medgan.py b/pyhealth/models/generators/medgan.py index 3e0ee06c0..7ce682703 100644 --- a/pyhealth/models/generators/medgan.py +++ b/pyhealth/models/generators/medgan.py @@ -7,7 +7,8 @@ ``dataset -> SampleDataset -> model`` pipeline. MedGAN treats each patient as a flat bag-of-codes (no visit structure), so it -expects an input feature named ``visits`` backed by a ``MultiHotProcessor``. +expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, +which the :class:`~pyhealth.tasks.PatientCodeSetGeneration` task provides. The training procedure has two phases (mirroring the reference): * a **linear autoencoder** is pre-trained with binary cross-entropy @@ -231,7 +232,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "MedGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor." + "MultiHotProcessor (see PatientCodeSetGeneration)." ) # The generator's residual connection (``out + residual`` with diff --git a/pyhealth/models/generators/promptehr.py b/pyhealth/models/generators/promptehr.py index e719cd78e..56b26ed36 100644 --- a/pyhealth/models/generators/promptehr.py +++ b/pyhealth/models/generators/promptehr.py @@ -3,7 +3,7 @@ This is a PyHealth ``BaseModel`` port of PromptEHR (Wang & Sun, EMNLP'22, https://github.com/RyanWangZf/PromptEHR), wrapped so it consumes the standard ``dataset -> set_task -> SampleDataset -> model`` pipeline and shares the same -:class:`~pyhealth.tasks.EHRGeneration` task as +:class:`~pyhealth.tasks.VisitSequenceGeneration` task as :class:`~pyhealth.models.HALO` and :class:`~pyhealth.models.GPT2`. PromptEHR treats sequential EHRs as a *neural database* and learns to fill in @@ -17,7 +17,7 @@ the way :class:`~pyhealth.models.GPT2` wraps ``GPT2LMHeadModel``. * **Prompt learning.** The reference reparameterizes a learnable prompt from patient baseline demographics and prepends it to the encoder/decoder - (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.EHRGeneration` + (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.VisitSequenceGeneration` task is *unconditional* (only ``visits``, no baseline features -- exactly like HALO/GPT2), so the prompt reduces to a learnable continuous **soft prefix** prepended to the encoder. This is the prompt-tuning core without the @@ -35,7 +35,7 @@ [CODE_PROMPT] [VISIT_DELIM] ... [EOS] The reference handles several code types (diagnosis / procedure / drug / lab) -each with its own modality prompt token; the PyHealth ``EHRGeneration`` task +each with its own modality prompt token; the PyHealth ``VisitSequenceGeneration`` task exposes a single ``visits`` modality, so a single ``[CODE_PROMPT]`` token marks it. The code vocabulary is taken from the dataset's ``visits`` processor (which already reserves index 0 for ```` and @@ -63,12 +63,13 @@ class PromptEHR(BaseModel): Trains a BART denoising autoencoder with a learnable soft prompt on patient visit-code streams, then generates synthetic patients by prompt-conditioned encoder-decoder sampling. Generation is **unconditional** (no demographic - conditioning), matching the :class:`~pyhealth.tasks.EHRGeneration` task. + conditioning), matching the :class:`~pyhealth.tasks.VisitSequenceGeneration` task. Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains - ``{"visits": NestedMultiHotProcessor}`` (or the equivalent - ``NestedSequenceProcessor``) and whose ``output_schema`` is empty. + ``{"visits": NestedSequenceProcessor}`` -- use the + :class:`~pyhealth.tasks.VisitSequenceGeneration` task -- and whose + ``output_schema`` is empty. embed_dim: BART model dimension (``d_model``). Must be divisible by ``n_heads``. Default: 256. n_heads: Number of attention heads (encoder and decoder). Default: 8. @@ -127,7 +128,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "PromptEHR expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor or NestedMultiHotProcessor." + "NestedSequenceProcessor (see VisitSequenceGeneration)." ) if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): # Without this the visit row would be read as raw values. Under a @@ -137,7 +138,7 @@ def __init__( f"PromptEHR needs a 'visits' processor that can invert its own " f"encoding (a visit_code_ids method); got " f"{type(dataset.input_processors['visits']).__name__}. Use " - "NestedSequenceProcessor or NestedMultiHotProcessor." + "NestedSequenceProcessor, via the VisitSequenceGeneration task." ) self.save_dir = save_dir diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 612bc2546..54ee000ab 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -55,6 +55,9 @@ EHRGeneration, EHRGenerationMIMIC3, EHRGenerationMIMIC4, + PatientCodeSetGeneration as PatientCodeSetGeneration, + VisitMultiHotGeneration as VisitMultiHotGeneration, + VisitSequenceGeneration as VisitSequenceGeneration, decode_dataset, to_evaluation_dataframe, ) diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 764e4de4a..28840c8fe 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -1,14 +1,28 @@ """EHR sequence-generation tasks for PyHealth generative models. -This is the shared task for every generator in -:mod:`pyhealth.models.generators` (HALO, MedGAN, CorGAN, PromptEHR, ...). It -extracts, for each patient, the ordered list of visits where each visit is the -list of medical codes recorded in that admission. The single input feature -``visits`` is processed by :class:`~pyhealth.processors.NestedMultiHotProcessor`; -there is no prediction label, so ``output_schema`` is empty. - -:class:`EHRGeneration` holds all the extraction logic; dataset-specific -subclasses only declare which event type and code attribute to read. +These back every generator in :mod:`pyhealth.models.generators`. They extract, +for each patient, the ordered list of visits where each visit is the list of +medical codes recorded in that admission. There is no prediction label, so +``output_schema`` is empty. + +Extraction is shared -- :class:`EHRGeneration` holds all of it -- but the +encoding is not, because each generator family reads its codes in a different +shape: + +- :class:`VisitMultiHotGeneration` -- one multi-hot row per visit, for HALO, + whose transformer consumes multi-hot vectors directly. +- :class:`VisitSequenceGeneration` -- per-visit code indices, for the token + models GPT2 and PromptEHR, which flatten visits into streams of code ids. +- :class:`PatientCodeSetGeneration` -- one pooled code set per patient, for the + bag-of-codes models MedGAN and CorGAN, which have no visit axis. + +Match the task to the model. Handing a model the wrong encoding does not raise: +the numbers still have the right shape and dtype, so training runs and produces +a confidently wrong result. That is why these are separate classes rather than +one task with a flag. + +``event_type`` / ``code_attr`` select the dataset's coding columns and can be +passed to any of them, so the encoding and the dataset are independent choices. Evaluating generated data ------------------------- @@ -67,7 +81,11 @@ from typing import ClassVar from pyhealth.data.data import Patient -from pyhealth.processors import NestedMultiHotProcessor +from pyhealth.processors import ( + MultiHotProcessor, + NestedMultiHotProcessor, + NestedSequenceProcessor, +) from .base_task import BaseTask @@ -75,50 +93,78 @@ class EHRGeneration(BaseTask): - """Generic per-visit code-sequence task for unconditional EHR generators. + """Per-visit code extraction for unconditional EHR generators. Builds one sample per qualifying patient: the ordered list of visits, each visit being the list of codes (read from ``code_attr`` on ``event_type`` events) recorded in that admission. Patients with fewer than ``min_visits`` qualifying visits are skipped. - Subclass and override the class attributes for a specific dataset, or set - them on an instance. The defaults read MIMIC-III ICD-9 diagnosis codes. + **This class does not set an** ``input_schema`` **and cannot be used + directly.** Extraction is shared, but each generator family wants the codes + in a different shape, and handing a model the wrong one fails silently + rather than loudly. Pick the subclass that matches your model: + + ============================== ===================== ================== + Task Encoding Models + ============================== ===================== ================== + :class:`VisitMultiHotGeneration` per-visit multi-hot HALO + :class:`VisitSequenceGeneration` per-visit indices GPT2, PromptEHR + :class:`PatientCodeSetGeneration` one set per patient MedGAN, CorGAN + ============================== ===================== ================== Args: - task_name: Name of the task. - input_schema: ``{"visits": NestedMultiHotProcessor}``. - output_schema: empty (generative task, no labels). - event_type: Event type to pull per admission. Default - ``"diagnoses_icd"``. - code_attr: Event attribute holding the code string. Default - ``"icd9_code"``. - min_visits: Minimum qualifying visits to keep a patient. Default 2. + code_mapping: Optional vocabulary mapping, see :class:`BaseTask`. + event_type: Event type to pull per admission. Defaults to the class + attribute (``"diagnoses_icd"``). + code_attr: Event attribute holding the code string. Defaults to the + class attribute (``"icd9_code"``). + min_visits: Minimum qualifying visits to keep a patient. Defaults to + the class attribute (2). Examples: - >>> from pyhealth.datasets import MIMIC3Dataset - >>> from pyhealth.tasks import EHRGeneration - >>> ds = MIMIC3Dataset(root="...", tables=["diagnoses_icd"], dev=True) - >>> samples = ds.set_task(EHRGeneration()) - >>> samples[0]["visits"].shape # (num_visits, vocab_size) multi-hot - torch.Size([3, 512]) + >>> from pyhealth.tasks import VisitSequenceGeneration + >>> task = VisitSequenceGeneration(code_attr="icd_code") + >>> task.code_attr + 'icd_code' """ task_name: str = "ehr_generation" - input_schema: ClassVar[dict[str, str | type]] = { - "visits": NestedMultiHotProcessor - } output_schema: ClassVar[dict[str, str | type]] = {} event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" min_visits: int = 2 - def __call__(self, patient: Patient) -> list[dict]: - """Extract the per-visit code sequence for a patient.""" + def __init__( + self, + code_mapping=None, + event_type: str | None = None, + code_attr: str | None = None, + min_visits: int | None = None, + ) -> None: + if not hasattr(type(self), "input_schema"): + raise TypeError( + f"{type(self).__name__} does not declare an encoding. " + "EHRGeneration only holds the shared extraction logic -- use " + "VisitMultiHotGeneration (HALO), VisitSequenceGeneration " + "(GPT2, PromptEHR) or PatientCodeSetGeneration (MedGAN, " + "CorGAN), whichever matches your model." + ) + super().__init__(code_mapping=code_mapping) + # Per-instance overrides, so a dataset preset is a constructor argument + # rather than yet another subclass in the encoding x dataset grid. + if event_type is not None: + self.event_type = event_type + if code_attr is not None: + self.code_attr = code_attr + if min_visits is not None: + self.min_visits = min_visits + + def _visits(self, patient: Patient) -> list[list[str]]: + """Ordered per-admission code lists, empty admissions dropped.""" visits: list[list[str]] = [] - admissions = patient.get_events(event_type="admissions") - for admission in admissions: + for admission in patient.get_events(event_type="admissions"): events = patient.get_events( event_type=self.event_type, filters=[("hadm_id", "==", admission.hadm_id)], @@ -130,15 +176,92 @@ def __call__(self, patient: Patient) -> list[dict]: ] if codes: visits.append(codes) + return visits + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = self._visits(patient) if len(visits) < self.min_visits: return [] - return [{"patient_id": patient.patient_id, "visits": visits}] -class EHRGenerationMIMIC3(EHRGeneration): - """EHR generation task for MIMIC-III (ICD-9 diagnosis codes). +class VisitMultiHotGeneration(EHRGeneration): + """Per-visit code sets as multi-hot rows. For HALO. + + HALO's transformer consumes a multi-hot vector per context position, so + this hands it exactly that and no repacking happens on the way in. + + Examples: + >>> from pyhealth.tasks import VisitMultiHotGeneration + >>> samples = dataset.set_task(VisitMultiHotGeneration()) + >>> samples[0]["visits"].shape # (num_visits, vocab_size) + torch.Size([3, 512]) + """ + + task_name: str = "ehr_generation_visit_multihot" + input_schema: ClassVar[dict[str, str | type]] = { + "visits": NestedMultiHotProcessor + } + + +class VisitSequenceGeneration(EHRGeneration): + """Per-visit code indices, right-padded. For GPT2 and PromptEHR. + + Both are token-sequence models: they flatten each visit into a stream of + code ids. Indices are what they need, so this avoids encoding to multi-hot + and decoding straight back. + + Examples: + >>> from pyhealth.tasks import VisitSequenceGeneration + >>> samples = dataset.set_task(VisitSequenceGeneration()) + >>> samples[0]["visits"].shape # (num_visits, max_codes_per_visit) + torch.Size([3, 12]) + """ + + task_name: str = "ehr_generation_visit_sequence" + input_schema: ClassVar[dict[str, str | type]] = { + "visits": NestedSequenceProcessor + } + + +class PatientCodeSetGeneration(EHRGeneration): + """One code set per patient, visit structure discarded. For MedGAN/CorGAN. + + Bag-of-codes generators emit a single aggregate vector per patient, so the + visit axis is collapsed here rather than inside the model. ``min_visits`` + still applies -- it filters on the patient's real visit count before the + codes are pooled. + + Note: + Because the visit axis is gone, the next-visit utility metric in + :mod:`pyhealth.metrics.generative` is not meaningful for these models; + see this module's header. + + Examples: + >>> from pyhealth.tasks import PatientCodeSetGeneration + >>> samples = dataset.set_task(PatientCodeSetGeneration()) + >>> samples[0]["visits"].shape # (vocab_size,) + torch.Size([512]) + """ + + task_name: str = "ehr_generation_patient_codeset" + input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + + def __call__(self, patient: Patient) -> list[dict]: + """Pool every visit's codes into one per-patient set.""" + visits = self._visits(patient) + if len(visits) < self.min_visits: + return [] + codes = sorted({code for visit in visits for code in visit}) + return [{"patient_id": patient.patient_id, "visits": codes}] + + +class EHRGenerationMIMIC3(VisitMultiHotGeneration): + """EHR generation task for MIMIC-III (ICD-9 diagnosis codes), for HALO. + + A :class:`VisitMultiHotGeneration` preset. For GPT2/PromptEHR on MIMIC-III + use ``VisitSequenceGeneration()``, whose defaults are already MIMIC-III's. Examples: >>> from pyhealth.datasets import MIMIC3Dataset @@ -155,8 +278,11 @@ class EHRGenerationMIMIC3(EHRGeneration): code_attr: str = "icd9_code" -class EHRGenerationMIMIC4(EHRGeneration): - """EHR generation task for MIMIC-IV (ICD diagnosis codes). +class EHRGenerationMIMIC4(VisitMultiHotGeneration): + """EHR generation task for MIMIC-IV (ICD diagnosis codes), for HALO. + + A :class:`VisitMultiHotGeneration` preset. For another encoding on MIMIC-IV + pass the same columns, e.g. ``VisitSequenceGeneration(code_attr="icd_code")``. Examples: >>> from pyhealth.datasets import MIMIC4Dataset @@ -265,21 +391,25 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: ['4019', '25000'] """ processor = sample_dataset.input_processors[feature_key] + if not hasattr(processor, "visit_code_ids"): + raise ValueError( + "decode_dataset needs a per-visit processor that can invert its own " + f"encoding (a visit_code_ids method); got {type(processor).__name__}. " + "PatientCodeSetGeneration has no visit axis to decode." + ) index_to_code = {idx: code for code, idx in processor.code_vocab.items()} records: list[dict] = [] for i in range(len(sample_dataset)): sample = sample_dataset[i] visits: list[list[str]] = [] - # Each row is a multi-hot vector over the vocabulary, so the codes - # present are its non-zero columns. (Reading the values as indices -- - # which the index-based encoding required -- would see only 0s and 1s - # here and decode every visit as empty.) + # The processor knows how to read its own rows -- multi-hot columns or + # padded indices -- so this works for either per-visit task. for row in sample[feature_key]: codes = [ - index_to_code[int(col)] - for col in row.nonzero().flatten().tolist() - if index_to_code.get(int(col)) not in (None, "", "") + index_to_code[idx] + for idx in processor.visit_code_ids(row) + if index_to_code.get(idx) not in (None, "", "") ] if codes: visits.append(codes) diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py index 7abc34bb4..07702695d 100644 --- a/tests/core/test_generator_encodings.py +++ b/tests/core/test_generator_encodings.py @@ -1,10 +1,11 @@ -"""Every generator sharing EHRGeneration must read the task's encoding. - -HALO, GPT2 and PromptEHR all consume ``EHRGeneration``, so changing that task's -``input_schema`` changes what all three receive. The failure mode this guards -is silent: a multi-hot row read as raw values yields 1.0 for every present -code, and 1 is ````, so a model would train happily on nothing but unknown -codes. These tests assert the real codes survive, under either encoding. +"""Each generator family gets the encoding it actually consumes. + +Extraction is shared across the EHR-generation tasks, the encoding is not: +HALO reads multi-hot rows, GPT2/PromptEHR read code indices, MedGAN/CorGAN read +one pooled set per patient. Pairing a model with the wrong task does not raise +-- the shapes and dtypes stay valid -- so these tests pin the pairing itself, +and check that real codes survive the round trip rather than collapsing to +. """ import unittest @@ -13,6 +14,17 @@ from pyhealth.datasets import create_sample_dataset from pyhealth.models import GPT2, PromptEHR +from pyhealth.processors import ( + MultiHotProcessor, + NestedMultiHotProcessor, + NestedSequenceProcessor, +) +from pyhealth.tasks import ( + EHRGenerationMIMIC3, + PatientCodeSetGeneration, + VisitMultiHotGeneration, + VisitSequenceGeneration, +) SAMPLES = [ {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, @@ -31,6 +43,44 @@ def _dataset(schema_key, name): ) +class TestTaskEncodings(unittest.TestCase): + """Each task declares the processor its models consume.""" + + def test_each_family_gets_its_own_encoding(self): + self.assertIs( + VisitMultiHotGeneration.input_schema["visits"], NestedMultiHotProcessor + ) + self.assertIs( + VisitSequenceGeneration.input_schema["visits"], NestedSequenceProcessor + ) + self.assertIs( + PatientCodeSetGeneration.input_schema["visits"], MultiHotProcessor + ) + + def test_base_task_refuses_to_be_used_directly(self): + """EHRGeneration is extraction only, and says so instead of failing late.""" + from pyhealth.tasks import EHRGeneration + + self.assertFalse(hasattr(EHRGeneration, "input_schema")) + with self.assertRaises(TypeError) as ctx: + EHRGeneration() + self.assertIn("VisitMultiHotGeneration", str(ctx.exception)) + + def test_dataset_presets_stay_multihot(self): + """The MIMIC presets were HALO tasks and must remain so.""" + self.assertIs( + EHRGenerationMIMIC3.input_schema["visits"], NestedMultiHotProcessor + ) + self.assertTrue(issubclass(EHRGenerationMIMIC3, VisitMultiHotGeneration)) + + def test_columns_are_settable_per_instance(self): + """Encoding and dataset are independent choices, not a class grid.""" + task = VisitSequenceGeneration(code_attr="icd_code", min_visits=3) + self.assertEqual(task.code_attr, "icd_code") + self.assertEqual(task.min_visits, 3) + self.assertEqual(VisitSequenceGeneration.code_attr, "icd9_code") + + class TestVisitCodeIds(unittest.TestCase): """Each nested processor inverts its own encoding to the same code ids.""" @@ -56,7 +106,7 @@ def test_both_processors_agree(self): ) def test_multihot_ids_are_not_all_unk(self): - """The specific regression: reading values instead of column indices.""" + """Reading a multi-hot row's values instead of its column indices.""" processor = _dataset("nested_multihot", "vci_unk").input_processors["visits"] row = processor.process([["A05B", "A05C"]])[0] ids = processor.visit_code_ids(row) @@ -64,8 +114,8 @@ def test_multihot_ids_are_not_all_unk(self): self.assertNotIn(processor.UNK, ids) -class TestGeneratorsAcceptEitherEncoding(unittest.TestCase): - """GPT2 and PromptEHR serialise real codes from either processor.""" +class TestTokenGeneratorsOnIndices(unittest.TestCase): + """GPT2 and PromptEHR serialise real codes from their own encoding.""" MODELS = [ (GPT2, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64}), @@ -82,7 +132,7 @@ def _streams(self, model, visits): def test_codes_survive_serialisation(self): for cls, kwargs in self.MODELS: with self.subTest(model=cls.__name__): - dataset = _dataset("nested_multihot", f"gen_{cls.__name__}") + dataset = _dataset("nested_sequence", f"gen_{cls.__name__}") model = cls(dataset=dataset, batch_size=2, epochs=1, **kwargs) visits = torch.stack([dataset[i]["visits"] for i in range(2)]) streams = self._streams(model, visits) @@ -94,8 +144,6 @@ def test_codes_survive_serialisation(self): if t < model.code_vocab_size and t != 0 ] self.assertTrue(code_ids, "no code tokens were emitted at all") - # The bug turned every code into ; a real stream carries - # several distinct codes. self.assertGreater( len(set(code_ids) - {model.visits_processor.UNK}), 1, From 79c4c3fdf6a06848d43e4b2717552105f7b8956d Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 16:00:02 -0500 Subject: [PATCH 6/9] Revert the shared code-vocabulary mixin 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 --- docs/api/processors.rst | 6 -- pyhealth/models/generators/gpt2.py | 4 +- pyhealth/processors/base_processor.py | 82 ------------------- .../deep_nested_sequence_processor.py | 45 +++++++++- .../processors/nested_multihot_processor.py | 44 +++++++++- .../processors/nested_sequence_processor.py | 52 ++++++++++-- pyhealth/processors/sequence_processor.py | 35 +++++++- pyhealth/processors/stagenet_processor.py | 49 +++++++++-- tests/core/test_vocab_processors.py | 64 --------------- 9 files changed, 200 insertions(+), 181 deletions(-) diff --git a/docs/api/processors.rst b/docs/api/processors.rst index 91609b1a9..e4262fa4c 100644 --- a/docs/api/processors.rst +++ b/docs/api/processors.rst @@ -27,12 +27,6 @@ Available Processors - ``SequenceProcessor``: For categorical sequences (e.g., medical codes like diagnoses, procedures) - ``NestedSequenceProcessor``: For nested categorical sequences (e.g., drug recommendation with visit history) - ``NestedMultiHotProcessor``: For nested categorical sequences as per-visit multi-hot vectors (e.g., generative EHR models) - -Processors that map codes to indices share one vocabulary implementation, -``CodeVocabularyMixin`` (``pyhealth.processors.base_processor``): it provides -``add`` / ``remove`` / ``retain`` / ``tokens`` / ``vocab_size`` and the -````=0, ````=1 convention. A new code processor should mix it in -rather than reimplementing them. - ``NestedFloatsProcessor``: For nested numerical sequences with optional forward-fill **Label Processors:** diff --git a/pyhealth/models/generators/gpt2.py b/pyhealth/models/generators/gpt2.py index e2b24b7a7..0014e3ca7 100644 --- a/pyhealth/models/generators/gpt2.py +++ b/pyhealth/models/generators/gpt2.py @@ -146,8 +146,8 @@ def _encode_visits(self, visits: torch.Tensor): Args: visits: Processed visit tensor from either nested ``visits`` - processor; the processor's ``visit_code_ids`` inverts a row. Index 0 is ```` and is - skipped. + processor; the processor's ``visit_code_ids`` inverts a row. + Index 0 is ```` and is skipped. Returns: input_ids: LongTensor ``(batch, L)`` token streams, right-padded. diff --git a/pyhealth/processors/base_processor.py b/pyhealth/processors/base_processor.py index 7233fe779..06bbe0c7c 100644 --- a/pyhealth/processors/base_processor.py +++ b/pyhealth/processors/base_processor.py @@ -196,88 +196,6 @@ def vocab_size(self) -> int: pass -class CodeVocabularyMixin(TokenProcessorInterface): - """Concrete ````/```` code vocabulary, shared by code processors. - - Every processor that maps medical codes to indices needs the same four - operations -- build, add, remove, retain -- and each one used to carry its - own copy. ``SequenceProcessor``, ``NestedSequenceProcessor``, - ``DeepNestedSequenceProcessor``, ``NestedMultiHotProcessor`` and - ``StageNetProcessor`` held five verbatim duplicates of the same code. - - Deliberately a mixin and not a base class. Several models dispatch on - ``isinstance(processor, NestedSequenceProcessor)`` to decide whether to - apply ``nn.Embedding``, so making one code processor inherit from another - would silently reroute it. Sharing the implementation without touching the - inheritance chain keeps that dispatch honest. - - Subclasses call :meth:`_init_code_vocab` from ``__init__`` and - :meth:`_observe_code` from their own ``fit`` traversal, which differs by - nesting depth and so is not shared here. - - Examples: - >>> class MyProcessor(FeatureProcessor, CodeVocabularyMixin): - ... def __init__(self): - ... self._init_code_vocab() - ... def fit(self, samples, field): - ... for sample in samples: - ... for code in sample[field]: - ... self._observe_code(code) - ... def process(self, value): - ... return [self.code_vocab.get(c, self.UNK) for c in value] - >>> processor = MyProcessor() - >>> processor.fit([{"codes": ["A", "B"]}], "codes") - >>> processor.vocab_size() - 4 - >>> processor.code_vocab["A"] - 2 - """ - - code_vocab: dict[Any, int] - - def _init_code_vocab(self) -> None: - """Seed the vocabulary with ```` (0) and ```` (1).""" - self.code_vocab = {"": self.PAD, "": self.UNK} - self._next_index = 2 - - def _observe_code(self, code: Any) -> None: - """Add one code to the vocabulary if it is new and not ``None``.""" - if code is not None and code not in self.code_vocab: - self.code_vocab[code] = self._next_index - self._next_index += 1 - - def _reindex(self, keep: set) -> None: - """Rebuild the vocabulary over ``keep``, preserving relative order.""" - order = [ - k - for k, _ in sorted(self.code_vocab.items(), key=lambda kv: kv[1]) - if k in keep - ] - self.code_vocab = {k: i for i, k in enumerate(order)} - self._next_index = len(self.code_vocab) - - def remove(self, tokens: set[str]) -> None: - """Remove specified vocabularies from the processor.""" - self._reindex(set(self.code_vocab.keys()) - tokens | {"", ""}) - - def retain(self, tokens: set[str]) -> None: - """Retain only the specified vocabularies in the processor.""" - self._reindex(set(self.code_vocab.keys()) & tokens | {"", ""}) - - def add(self, tokens: set[str]) -> None: - """Add specified vocabularies to the processor.""" - for token in tokens: - self._observe_code(token) - - def tokens(self) -> set[str]: - """Return the set of tokens in the processor's vocabulary.""" - return set(self.code_vocab.keys()) - - def vocab_size(self) -> int: - """Return the size of the processor's vocabulary.""" - return len(self.code_vocab) - - class TemporalFeatureProcessor(FeatureProcessor): """Abstract base class for processors whose features are paired with timestamps. diff --git a/pyhealth/processors/deep_nested_sequence_processor.py b/pyhealth/processors/deep_nested_sequence_processor.py index 2d14f6350..24683f54b 100644 --- a/pyhealth/processors/deep_nested_sequence_processor.py +++ b/pyhealth/processors/deep_nested_sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import CodeVocabularyMixin, FeatureProcessor +from .base_processor import FeatureProcessor, TokenProcessorInterface @register_processor("deep_nested_sequence") -class DeepNestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): +class DeepNestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): """ Feature processor for deeply nested categorical sequences with vocabulary. @@ -45,7 +45,8 @@ class DeepNestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """ def __init__(self): - self._init_code_vocab() + self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 self._max_middle_len = 1 # Maximum length of middle sequences (e.g. visits) self._max_inner_len = 1 # Maximum length of inner sequences (e.g. codes per visit) @@ -77,11 +78,39 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Build vocabulary for code in inner_seq: - self._observe_code(code) + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 self._max_middle_len = max(1, max_middle_len) self._max_inner_len = max(1, max_inner_len) + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + def process(self, value: List[List[List[Any]]]) -> torch.Tensor: """Process deep nested sequence into padded 3D tensor. @@ -144,10 +173,18 @@ def process(self, value: List[List[List[Any]]]) -> torch.Tensor: return torch.tensor(encoded_groups, dtype=torch.long) + def vocab_size(self) -> int: + """Return the size of the processor's vocabulary.""" + return len(self.code_vocab) + def size(self) -> int: """Return max inner length (embedding dimension) for unified API.""" return self._max_inner_len + def vocab_size(self) -> int: + """Return vocabulary size.""" + return len(self.code_vocab) + def __repr__(self): return ( f"DeepNestedSequenceProcessor(" diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py index 046ea6e62..735049cec 100644 --- a/pyhealth/processors/nested_multihot_processor.py +++ b/pyhealth/processors/nested_multihot_processor.py @@ -4,11 +4,11 @@ import torch from . import register_processor -from .base_processor import CodeVocabularyMixin, FeatureProcessor +from .base_processor import FeatureProcessor, TokenProcessorInterface @register_processor("nested_multihot") -class NestedMultiHotProcessor(FeatureProcessor, CodeVocabularyMixin): +class NestedMultiHotProcessor(FeatureProcessor, TokenProcessorInterface): """Nested categorical sequences as per-visit multi-hot vectors. Same input as :class:`NestedSequenceProcessor` -- a list of visits, each a @@ -59,7 +59,8 @@ def __init__(self, padding: int = 0): # `padding` is accepted and ignored so this is a drop-in swap for # NestedSequenceProcessor in a schema. There is no inner axis to pad -- # that is the entire point -- so honouring it would be misleading. - self._init_code_vocab() + self.code_vocab: dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 self._padding = padding def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: @@ -80,7 +81,38 @@ def fit(self, samples: Iterable[dict[str, Any]], field: str) -> None: if not isinstance(inner_seq, list): continue for code in inner_seq: - self._observe_code(code) + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 + + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) + if k in keep] + self.code_vocab = {k: i for i, k in enumerate(order)} + self._next_index = len(self.code_vocab) + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + self._next_index = len(self.code_vocab) + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) def process(self, value: list[list[Any]]) -> torch.Tensor: """Nested sequence -> ``(num_visits, vocab_size)`` float multi-hot. @@ -143,6 +175,10 @@ def size(self) -> int: """Feature width: the vocabulary, since that is the row length.""" return len(self.code_vocab) + def vocab_size(self) -> int: + """Return vocabulary size.""" + return len(self.code_vocab) + def __repr__(self): return f"NestedMultiHotProcessor(vocab_size={len(self.code_vocab)})" diff --git a/pyhealth/processors/nested_sequence_processor.py b/pyhealth/processors/nested_sequence_processor.py index 72690f4d6..060b5e3f8 100644 --- a/pyhealth/processors/nested_sequence_processor.py +++ b/pyhealth/processors/nested_sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import CodeVocabularyMixin, FeatureProcessor +from .base_processor import FeatureProcessor, TokenProcessorInterface @register_processor("nested_sequence") -class NestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): +class NestedSequenceProcessor(FeatureProcessor, TokenProcessorInterface): """ Feature processor for nested categorical sequences with vocabulary. @@ -45,7 +45,8 @@ class NestedSequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """ def __init__(self, padding: int = 0): - self._init_code_vocab() + self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 self._max_inner_len = 1 # Maximum length of inner sequences self._padding = padding # Additional padding beyond observed max @@ -71,13 +72,41 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Build vocabulary for code in inner_seq: - self._observe_code(code) + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 # Store max inner length: add user-specified padding to observed maximum # This ensures the processor can handle sequences longer than those in training data observed_max = max(1, max_inner_len) self._max_inner_len = observed_max + self._padding + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + def process(self, value: List[List[Any]]) -> torch.Tensor: """Process nested sequence into padded 2D tensor. @@ -121,13 +150,18 @@ def process(self, value: List[List[Any]]) -> torch.Tensor: return torch.tensor(encoded_sequences, dtype=torch.long) + def vocab_size(self) -> int: + """Return the size of the processor's vocabulary.""" + return len(self.code_vocab) + def visit_code_ids(self, row: torch.Tensor) -> list[int]: """Code indices present in one processed visit row. The inverse of what :meth:`process` writes, for consumers that need a - code *list* rather than the tensor -- sequence generators such as GPT2 - and PromptEHR. Each nested processor implements this for its own - encoding, so a model can accept either without knowing which it has. + code *list* rather than the tensor -- the sequence generators GPT2 and + PromptEHR, and :func:`pyhealth.tasks.decode_dataset`. Each nested + processor implements this for its own encoding, so a consumer never has + to guess how to read a row. Here the row already holds indices, right-padded with ```` (0). @@ -147,6 +181,10 @@ def size(self) -> int: """Return max inner length (embedding dimension) for unified API.""" return self._max_inner_len + def vocab_size(self) -> int: + """Return vocabulary size.""" + return len(self.code_vocab) + def __repr__(self): return ( f"NestedSequenceProcessor(" diff --git a/pyhealth/processors/sequence_processor.py b/pyhealth/processors/sequence_processor.py index ad0a2a0ef..47339eefd 100644 --- a/pyhealth/processors/sequence_processor.py +++ b/pyhealth/processors/sequence_processor.py @@ -3,11 +3,11 @@ import torch from . import register_processor -from .base_processor import CodeVocabularyMixin, FeatureProcessor +from .base_processor import FeatureProcessor, TokenProcessorInterface @register_processor("sequence") -class SequenceProcessor(FeatureProcessor, CodeVocabularyMixin): +class SequenceProcessor(FeatureProcessor, TokenProcessorInterface): """Feature processor for encoding categorical sequences. Encodes medical codes (e.g., diagnoses, procedures) into numerical @@ -29,7 +29,8 @@ class SequenceProcessor(FeatureProcessor, CodeVocabularyMixin): """ def __init__(self, code_mapping: Optional[Tuple[str, str]] = None): - self._init_code_vocab() + self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 self._mapper = None if code_mapping is not None: from pyhealth.medcode import CrossMap @@ -82,6 +83,34 @@ def process(self, value: Any) -> torch.Tensor: indices.append(self.code_vocab[""]) return torch.tensor(indices, dtype=torch.long) + + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + self.code_vocab = { k : i for i, k in enumerate(order) } + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + self.code_vocab = { k : i for i, k in enumerate(order) } + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + + def vocab_size(self) -> int: + """Return the size of the processor's vocabulary.""" + return len(self.code_vocab) def size(self): return len(self.code_vocab) diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py index ce302ef36..604376ec1 100644 --- a/pyhealth/processors/stagenet_processor.py +++ b/pyhealth/processors/stagenet_processor.py @@ -3,15 +3,11 @@ import torch from . import register_processor -from .base_processor import ( - CodeVocabularyMixin, - ModalityType, - TemporalFeatureProcessor, -) +from .base_processor import FeatureProcessor, ModalityType, TemporalFeatureProcessor, TokenProcessorInterface @register_processor("stagenet") -class StageNetProcessor(TemporalFeatureProcessor, CodeVocabularyMixin): +class StageNetProcessor(TemporalFeatureProcessor, TokenProcessorInterface): """ Feature processor for StageNet CODE inputs with coupled value/time data. @@ -59,7 +55,8 @@ class StageNetProcessor(TemporalFeatureProcessor, CodeVocabularyMixin): """ def __init__(self, padding: int = 0): - self._init_code_vocab() + self.code_vocab: Dict[Any, int] = {"": self.PAD, "": self.UNK} + self._next_index = 2 self._is_nested = None # Will be determined during fit # Max inner sequence length for nested codes self._max_nested_len = None @@ -104,11 +101,15 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Track max inner length max_inner_len = max(max_inner_len, len(inner_list)) for code in inner_list: - self._observe_code(code) + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 else: # Flat codes for code in value_data: - self._observe_code(code) + if code is not None and code not in self.code_vocab: + self.code_vocab[code] = self._next_index + self._next_index += 1 # Store max nested length: add user-specified padding to observed maximum # This ensures the processor can handle sequences longer than those in training data @@ -116,6 +117,32 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: observed_max = max(1, max_inner_len) self._max_nested_len = observed_max + self._padding + def remove(self, tokens: set[str]): + """Remove specified vocabularies from the processor.""" + keep = set(self.code_vocab.keys()) - tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def retain(self, tokens: set[str]): + """Retain only the specified vocabularies in the processor.""" + keep = set(self.code_vocab.keys()) & tokens | {"", ""} + order = [k for k, v in sorted(self.code_vocab.items(), key=lambda x: x[1]) if k in keep] + + self.code_vocab = { k : i for i, k in enumerate(order) } + + def add(self, tokens: set[str]): + """Add specified vocabularies to the processor.""" + i = len(self.code_vocab) + for token in tokens: + if token not in self.code_vocab: + self.code_vocab[token] = i + i += 1 + + def tokens(self) -> set[str]: + """Return the set of tokens in the processor's vocabulary.""" + return set(self.code_vocab.keys()) + def process( self, value: Tuple[Optional[List], List] ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: @@ -194,6 +221,10 @@ def _encode_nested_codes(self, nested_codes: List[List[str]]) -> torch.Tensor: return torch.tensor(encoded_sequences, dtype=torch.long) + def vocab_size(self) -> int: + """Return the size of the processor's vocabulary.""" + return len(self.code_vocab) + def size(self) -> int: """Return vocabulary size.""" return len(self.code_vocab) diff --git a/tests/core/test_vocab_processors.py b/tests/core/test_vocab_processors.py index b167f8dcf..b6d785805 100644 --- a/tests/core/test_vocab_processors.py +++ b/tests/core/test_vocab_processors.py @@ -5,7 +5,6 @@ StageNetProcessor, NestedSequenceProcessor, DeepNestedSequenceProcessor, - NestedMultiHotProcessor, ) class TestVocabProcessors(unittest.TestCase): @@ -259,68 +258,5 @@ def test_deep_nested_sequence_processor_add(self): e_idx = processor.code_vocab["E"] self.assertEqual(res[0, 0, 0].item(), e_idx) -class TestSharedCodeVocabulary(unittest.TestCase): - """CodeVocabularyMixin gives every code processor the same vocabulary. - - Before the mixin, five processors carried verbatim copies of these methods - and every copy had the same defect: a removal renumbered the vocabulary but - left ``_next_index`` untouched, so the next ``fit`` allocated past the end - of it. Sharing one implementation fixes all five at once, which is what - these tests pin. - """ - - # Each processor takes a different nesting depth, so the same code list is - # wrapped to match. - PROCESSORS = [ - (SequenceProcessor, lambda codes: codes), - (NestedSequenceProcessor, lambda codes: [codes]), - (DeepNestedSequenceProcessor, lambda codes: [[codes]]), - (NestedMultiHotProcessor, lambda codes: [codes]), - ] - - def test_fit_after_remove_stays_in_range(self): - """The bug: ``fit`` allocates from ``_next_index``, ``remove`` renumbers. - - Removing codes renumbered the vocabulary to 0..n-1 but left - ``_next_index`` at its pre-removal value, so the next ``fit`` handed - out an index past the end of the vocabulary. Anything sizing an - embedding table by ``vocab_size()`` would then index out of bounds. - """ - for cls, wrap in self.PROCESSORS: - with self.subTest(processor=cls.__name__): - processor = cls() - processor.fit([{"codes": wrap(["A", "B", "C"])}], "codes") - processor.remove({"A", "B"}) - processor.fit([{"codes": wrap(["D"])}], "codes") - - indices = list(processor.code_vocab.values()) - self.assertEqual( - max(indices), processor.vocab_size() - 1, - f"{cls.__name__} allocated an index past the vocabulary: " - f"{processor.code_vocab}", - ) - self.assertEqual(len(indices), len(set(indices))) - - def test_indices_stay_contiguous_from_zero(self): - for cls, _ in self.PROCESSORS: - with self.subTest(processor=cls.__name__): - processor = cls() - processor.add({"A", "B", "C"}) - processor.retain({"B"}) - self.assertEqual( - sorted(processor.code_vocab.values()), - list(range(processor.vocab_size())), - ) - - def test_special_tokens_survive_every_edit(self): - for cls, _ in self.PROCESSORS: - with self.subTest(processor=cls.__name__): - processor = cls() - processor.add({"A"}) - processor.retain(set()) - self.assertEqual(processor.code_vocab[""], processor.PAD) - self.assertEqual(processor.code_vocab[""], processor.UNK) - - if __name__ == "__main__": unittest.main() From d947b630d868f42b4262eb6fdff791eee233d489 Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 16:08:13 -0500 Subject: [PATCH 7/9] Drop the unused multi-hot inverse; use the task helpers in the example 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 --- examples/halo_mimic3.py | 45 ++++++------------- .../processors/nested_multihot_processor.py | 22 --------- pyhealth/tasks/generate_ehr.py | 29 +++++++----- tests/core/test_generator_encodings.py | 26 ++++++----- 4 files changed, 48 insertions(+), 74 deletions(-) diff --git a/examples/halo_mimic3.py b/examples/halo_mimic3.py index ab6133a0a..e5188271d 100644 --- a/examples/halo_mimic3.py +++ b/examples/halo_mimic3.py @@ -11,12 +11,14 @@ 6. Evaluating the synthetic data with the generative metrics suite """ -import pandas as pd - from pyhealth.datasets import MIMIC3Dataset, split_by_patient from pyhealth.metrics.generative import evaluate_synthetic_ehr from pyhealth.models import HALO -from pyhealth.tasks import EHRGenerationMIMIC3 +from pyhealth.tasks import ( + EHRGenerationMIMIC3, + decode_dataset, + to_evaluation_dataframe, +) if __name__ == "__main__": # STEP 1: Load MIMIC-III base dataset @@ -84,35 +86,16 @@ # train_df, test_df and syn_df below all share this exact schema. `labels` # is a placeholder here: privacy metrics ignore it and the utility metric # overwrites it with the next-visit prediction target. - index_to_code = { - v: k for k, v in sample_dataset.input_processors["visits"].code_vocab.items() - } - - def real_subset_to_records(subset): - # NestedMultiHotProcessor encodes each patient as a dense - # (num_visits, vocab_size) multi-hot tensor, so the set of codes in a - # visit is the set of *nonzero column indices* -- not the tensor - # values themselves. - for sample in subset: - pid = str(sample["patient_id"]) - for t, visit in enumerate(sample["visits"]): - for idx in visit.nonzero(as_tuple=True)[0].tolist(): - code = index_to_code.get(idx) - if code in (None, "", ""): - continue - yield {"id": pid, "time": t, "visit_codes": code, "labels": 0} - - def synthetic_to_records(patients): - for p in patients: - pid = str(p["patient_id"]) - for t, visit in enumerate(p["visits"]): - for code in visit: - yield {"id": pid, "time": t, "visit_codes": code, "labels": 0} - + # Both conversions already live in the task module, so use them rather + # than re-deriving the encoding here: + # decode_dataset processed (or split) SampleDataset -> code records + # to_evaluation_dataframe records (real or synthetic) -> the long form + # Patients are renumbered 0, 1, 2, ... because synthetic patients do not + # correspond to real ones; the metrics only need a grouping key. schema = {"visit_codes": str, "labels": int, "time": int, "id": str} - train_df = pd.DataFrame(real_subset_to_records(train_dataset)).astype(schema) - test_df = pd.DataFrame(real_subset_to_records(test_dataset)).astype(schema) - syn_df = pd.DataFrame(synthetic_to_records(synthetic)).astype(schema) + train_df = to_evaluation_dataframe(decode_dataset(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(decode_dataset(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) print( f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " f"synthetic: {len(syn_df)}" diff --git a/pyhealth/processors/nested_multihot_processor.py b/pyhealth/processors/nested_multihot_processor.py index 735049cec..3b6c7e10e 100644 --- a/pyhealth/processors/nested_multihot_processor.py +++ b/pyhealth/processors/nested_multihot_processor.py @@ -149,28 +149,6 @@ def process(self, value: list[list[Any]]) -> torch.Tensor: torch.ones(len(idx))) return out - def visit_code_ids(self, row: torch.Tensor) -> list[int]: - """Code indices present in one processed visit row. - - The inverse of what :meth:`process` writes. Here the row is a multi-hot - vector, so the codes are its *nonzero column indices* -- reading the - values themselves would yield 1.0 (i.e. ````) for every code. - - Codes come back in vocabulary order, not charted order, and repeats are - already collapsed; multi-hot records presence, not sequence or count. - - Args: - row: 1D multi-hot tensor of width ``vocab_size``, one visit. - - Returns: - Code indices present in the visit, ascending. - - Examples: - >>> processor.visit_code_ids(torch.tensor([0., 0., 1., 0., 1.])) - [2, 4] - """ - return row.nonzero(as_tuple=True)[0].tolist() - def size(self) -> int: """Feature width: the vocabulary, since that is the row length.""" return len(self.code_vocab) diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 28840c8fe..738ca3f1c 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -384,18 +384,25 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: Returns: List of ``{"visits": [[code_str, ...], ...]}`` records. + Raises: + TypeError: If ``feature_key`` is not backed by a + :class:`~pyhealth.processors.NestedMultiHotProcessor`. + Examples: >>> from pyhealth.tasks.generate_ehr import decode_dataset >>> records = decode_dataset(samples) >>> records[0]["visits"][0] ['4019', '25000'] """ - processor = sample_dataset.input_processors[feature_key] - if not hasattr(processor, "visit_code_ids"): - raise ValueError( - "decode_dataset needs a per-visit processor that can invert its own " - f"encoding (a visit_code_ids method); got {type(processor).__name__}. " - "PatientCodeSetGeneration has no visit axis to decode." + # split_by_patient hands back a torch Subset, which carries no processors + # of its own -- decoding a split is the common case, so resolve through it. + source = getattr(sample_dataset, "dataset", sample_dataset) + processor = source.input_processors[feature_key] + if not isinstance(processor, NestedMultiHotProcessor): + raise TypeError( + f"decode_dataset inverts the multi-hot encoding, but '{feature_key}' " + f"is a {type(processor).__name__}. Use VisitMultiHotGeneration, or " + "read the codes off the index tensor directly." ) index_to_code = {idx: code for code, idx in processor.code_vocab.items()} @@ -403,13 +410,13 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: for i in range(len(sample_dataset)): sample = sample_dataset[i] visits: list[list[str]] = [] - # The processor knows how to read its own rows -- multi-hot columns or - # padded indices -- so this works for either per-visit task. + # Each row is a multi-hot vector over the vocabulary, so the codes + # present are its nonzero columns -- the values are all 1.0. for row in sample[feature_key]: codes = [ - index_to_code[idx] - for idx in processor.visit_code_ids(row) - if index_to_code.get(idx) not in (None, "", "") + index_to_code[int(col)] + for col in row.nonzero(as_tuple=True)[0].tolist() + if index_to_code.get(int(col)) not in (None, "", "") ] if codes: visits.append(codes) diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py index 07702695d..b9ee77df0 100644 --- a/tests/core/test_generator_encodings.py +++ b/tests/core/test_generator_encodings.py @@ -82,35 +82,41 @@ def test_columns_are_settable_per_instance(self): class TestVisitCodeIds(unittest.TestCase): - """Each nested processor inverts its own encoding to the same code ids.""" + """NestedSequenceProcessor inverts its own rows for the token generators.""" - def test_both_processors_agree(self): + def test_matches_the_multihot_columns(self): + """Both encodings of the same visit name the same codes. + + NestedMultiHotProcessor has no visit_code_ids -- nothing consumes one -- + so its codes are read here the way decode_dataset reads them, as the + row's nonzero columns. + """ multihot = _dataset("nested_multihot", "vci_mh") indexed = _dataset("nested_sequence", "vci_ix") - mh_proc = multihot.input_processors["visits"] ix_proc = indexed.input_processors["visits"] # Same samples, same traversal order, so the vocabularies must match -- # that is what makes the per-visit comparison below meaningful. - self.assertEqual(mh_proc.code_vocab, ix_proc.code_vocab) + self.assertEqual(multihot.input_processors["visits"].code_vocab, + ix_proc.code_vocab) for i in range(len(SAMPLES)): mh_row = multihot[i]["visits"] ix_row = indexed[i]["visits"] for visit in range(mh_row.shape[0]): - # Multi-hot returns vocabulary order, the index form charted - # order, so compare as sets. + # Multi-hot columns come out in vocabulary order, the index form + # in charted order, so compare as sets. self.assertEqual( - set(mh_proc.visit_code_ids(mh_row[visit])), + set(mh_row[visit].nonzero(as_tuple=True)[0].tolist()), set(ix_proc.visit_code_ids(ix_row[visit])), ) - def test_multihot_ids_are_not_all_unk(self): - """Reading a multi-hot row's values instead of its column indices.""" - processor = _dataset("nested_multihot", "vci_unk").input_processors["visits"] + def test_padding_is_dropped_not_read_as_a_code(self): + processor = _dataset("nested_sequence", "vci_pad").input_processors["visits"] row = processor.process([["A05B", "A05C"]])[0] ids = processor.visit_code_ids(row) self.assertEqual(len(ids), 2) + self.assertNotIn(processor.PAD, ids) self.assertNotIn(processor.UNK, ids) From bd683b216b8d37ffd7dfd7d605abea7ad5fe2c85 Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 16:10:38 -0500 Subject: [PATCH 8/9] Cover the task extraction and the code-set pooling 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 --- tests/core/test_generator_encodings.py | 69 ++++++++++++++++++++++++-- 1 file changed, 64 insertions(+), 5 deletions(-) diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py index b9ee77df0..c4d109fa0 100644 --- a/tests/core/test_generator_encodings.py +++ b/tests/core/test_generator_encodings.py @@ -9,10 +9,9 @@ """ import unittest +from typing import ClassVar -import torch - -from pyhealth.datasets import create_sample_dataset +from pyhealth.datasets import create_sample_dataset, get_dataloader from pyhealth.models import GPT2, PromptEHR from pyhealth.processors import ( MultiHotProcessor, @@ -43,6 +42,63 @@ def _dataset(schema_key, name): ) +class _Event: + def __init__(self, hadm_id, code=None): + self.hadm_id = hadm_id + self.icd9_code = code + + +class _Patient: + """Minimal stand-in for pyhealth.data.Patient: admissions plus coded events.""" + + def __init__(self, patient_id, visits): + self.patient_id = patient_id + self._admissions = [_Event(f"h{i}") for i in range(len(visits))] + self._codes = [ + _Event(f"h{i}", code) + for i, codes in enumerate(visits) + for code in codes + ] + + def get_events(self, event_type, filters=None): + if event_type == "admissions": + return self._admissions + hadm = filters[0][2] + return [e for e in self._codes if e.hadm_id == hadm] + + +class TestExtraction(unittest.TestCase): + """The shared __call__, and the pooling PatientCodeSetGeneration adds.""" + + VISITS: ClassVar[list] = [["A05B", "A05C"], ["A11D"], ["A05B"]] + + def test_per_visit_tasks_keep_visit_structure(self): + patient = _Patient("p0", self.VISITS) + samples = VisitMultiHotGeneration()(patient) + self.assertEqual(len(samples), 1) + self.assertEqual(samples[0]["visits"], self.VISITS) + # Same extraction regardless of encoding -- only input_schema differs. + self.assertEqual(VisitSequenceGeneration()(patient)[0]["visits"], + self.VISITS) + + def test_codeset_task_pools_and_dedupes(self): + samples = PatientCodeSetGeneration()(_Patient("p0", self.VISITS)) + self.assertEqual(len(samples), 1) + # One flat set: A05B appears in two visits and survives once. + self.assertEqual(samples[0]["visits"], ["A05B", "A05C", "A11D"]) + + def test_min_visits_counts_real_visits_before_pooling(self): + """Pooling must not let a 1-visit patient past a min_visits=2 filter.""" + one_visit = _Patient("p1", [["A05B", "A05C", "A11D"]]) + self.assertEqual(PatientCodeSetGeneration()(one_visit), []) + self.assertEqual(VisitMultiHotGeneration()(one_visit), []) + + def test_codeless_admissions_are_dropped(self): + patient = _Patient("p2", [["A05B"], [], ["A11D"]]) + self.assertEqual(VisitMultiHotGeneration()(patient)[0]["visits"], + [["A05B"], ["A11D"]]) + + class TestTaskEncodings(unittest.TestCase): """Each task declares the processor its models consume.""" @@ -123,7 +179,7 @@ def test_padding_is_dropped_not_read_as_a_code(self): class TestTokenGeneratorsOnIndices(unittest.TestCase): """GPT2 and PromptEHR serialise real codes from their own encoding.""" - MODELS = [ + MODELS: ClassVar[list] = [ (GPT2, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64}), (PromptEHR, {"embed_dim": 16, "n_heads": 2, "n_layers": 2, "max_len": 64, "prompt_length": 4}), @@ -140,7 +196,10 @@ def test_codes_survive_serialisation(self): with self.subTest(model=cls.__name__): dataset = _dataset("nested_sequence", f"gen_{cls.__name__}") model = cls(dataset=dataset, batch_size=2, epochs=1, **kwargs) - visits = torch.stack([dataset[i]["visits"] for i in range(2)]) + # Patients have different visit counts, so let the dataloader + # pad the visit dimension rather than stacking raw samples. + batch = next(iter(get_dataloader(dataset, batch_size=2))) + visits = batch["visits"] streams = self._streams(model, visits) code_ids = [ From db9e4c448333ac45c8597fc9e34a8ac15567182a Mon Sep 17 00:00:00 2001 From: Jane Du Date: Sat, 12 Sep 2026 16:40:28 -0500 Subject: [PATCH 9/9] Review: flatten the generation tasks, name the dataset, add examples 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 --- docs/api/tasks.rst | 2 +- .../api/tasks/pyhealth.tasks.generate_ehr.rst | 31 +- examples/gpt2_mimic3.py | 115 ++++++ examples/halo_mimic3.py | 6 +- examples/medgan_mimic3.py | 112 ++++++ pyhealth/models/generators/corgan.py | 4 +- pyhealth/models/generators/gpt2.py | 8 +- pyhealth/models/generators/halo.py | 6 +- pyhealth/models/generators/medgan.py | 4 +- pyhealth/models/generators/promptehr.py | 14 +- pyhealth/tasks/__init__.py | 8 +- pyhealth/tasks/generate_ehr.py | 333 ++++++++++-------- tests/core/test_generator_encodings.py | 91 ++--- 13 files changed, 509 insertions(+), 225 deletions(-) create mode 100644 examples/gpt2_mimic3.py create mode 100644 examples/medgan_mimic3.py diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst index 4f616d7aa..b250cd271 100644 --- a/docs/api/tasks.rst +++ b/docs/api/tasks.rst @@ -141,7 +141,7 @@ a quick reference: - Cumulative visit history (drug recommendation, readmission) * - ``"nested_multihot"`` - ``NestedMultiHotProcessor`` - - Per-visit code sets (HALO; see ``VisitMultiHotGeneration``) + - Per-visit code sets (HALO; see ``EHRGenerationMIMIC3``) * - ``"tensor"`` - ``TensorProcessor`` - Aggregated numeric values (e.g. last lab value per item) diff --git a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst index 0ec40d5a4..796201c67 100644 --- a/docs/api/tasks/pyhealth.tasks.generate_ehr.rst +++ b/docs/api/tasks/pyhealth.tasks.generate_ehr.rst @@ -6,56 +6,61 @@ unconditional synthetic-EHR generators, plus helpers to flatten generated output into the long-form dataframe consumed by :mod:`pyhealth.metrics.generative`. -Extraction is shared; the encoding is not. Each generator family reads its -codes in a different shape, and handing a model the wrong shape fails silently -rather than loudly, so pick the task that matches the model: +The classes are flat and independent -- one per (model family, dataset). The +extraction is MIMIC-shaped, assuming an ``admissions`` event type and a +``hadm_id`` linking codes to an admission, so the dataset is named in the class +and a task for eICU/OMOP/MEDS belongs alongside these rather than below them. + +Match the task to the model: each generator family reads its codes in a +different shape, and handing a model the wrong shape fails silently rather than +loudly. .. list-table:: :header-rows: 1 - :widths: 30 35 35 + :widths: 38 32 30 * - Task - Encoding - Models - * - ``VisitMultiHotGeneration`` + * - ``EHRGenerationMIMIC3`` / ``MIMIC4`` - one multi-hot row per visit - HALO - * - ``VisitSequenceGeneration`` + * - ``EHRSequenceGenerationMIMIC3`` / ``MIMIC4`` - per-visit code indices - GPT2, PromptEHR - * - ``PatientCodeSetGeneration`` + * - ``EHRCodeSetGenerationMIMIC3`` / ``MIMIC4`` - one code set per patient - MedGAN, CorGAN Task Classes ------------ -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGeneration +.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.VisitMultiHotGeneration +.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC4 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.VisitSequenceGeneration +.. autoclass:: pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC3 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.PatientCodeSetGeneration +.. autoclass:: pyhealth.tasks.generate_ehr.EHRSequenceGenerationMIMIC4 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC3 +.. autoclass:: pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC3 :members: :undoc-members: :show-inheritance: -.. autoclass:: pyhealth.tasks.generate_ehr.EHRGenerationMIMIC4 +.. autoclass:: pyhealth.tasks.generate_ehr.EHRCodeSetGenerationMIMIC4 :members: :undoc-members: :show-inheritance: diff --git a/examples/gpt2_mimic3.py b/examples/gpt2_mimic3.py new file mode 100644 index 000000000..d70382030 --- /dev/null +++ b/examples/gpt2_mimic3.py @@ -0,0 +1,115 @@ +"""Example: train the GPT-2 baseline on MIMIC-III and generate patients. + +GPT2 and PromptEHR are token-sequence models: each patient becomes one flat +stream of code ids, ``[BOS] codes_v1 [DELIM] codes_v2 ... [EOS]``. They read +per-visit code *indices*, so this uses EHRSequenceGenerationMIMIC3 -- not the +multi-hot EHRGenerationMIMIC3 that HALO takes (see halo_mimic3.py). Pairing a +model with the wrong task does not raise; it trains on nonsense. + +Swap GPT2 for PromptEHR below and the rest of the script is unchanged. + +This example demonstrates: +1. Loading MIMIC-III data +2. Applying the EHRSequenceGenerationMIMIC3 task (per-visit ICD-9 code indices) +3. Training GPT2 with its custom training loop +4. Generating synthetic patients +5. Evaluating the synthetic data with the generative metrics suite +""" + +from pyhealth.datasets import MIMIC3Dataset, split_by_patient +from pyhealth.metrics.generative import evaluate_synthetic_ehr +from pyhealth.models import GPT2 +from pyhealth.tasks import EHRSequenceGenerationMIMIC3, to_evaluation_dataframe + +if __name__ == "__main__": + # STEP 1: Load MIMIC-III. dev=True keeps this to a small subset -- start + # here, and only drop it once the whole script runs end to end. + base_dataset = MIMIC3Dataset( + root="/srv/local/data/MIMIC-III/mimic-iii-clinical-database-1.4", + tables=["diagnoses_icd"], + dev=True, + ) + + # STEP 2: Apply the sequence-encoded generation task (no labels). + sample_dataset = base_dataset.set_task(EHRSequenceGenerationMIMIC3()) + print(f"Total samples: {len(sample_dataset)}") + + sample = sample_dataset[0] + # (num_visits, max_codes_per_visit) of vocabulary indices, right-padded + # with (0) -- NOT a multi-hot vector. + print(f"Visits tensor shape: {tuple(sample['visits'].shape)}") + + # STEP 3: Split by patient so no patient appears in two splits. + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + # STEP 4: Initialize GPT2 (small config for the dev subset). + model = GPT2( + dataset=sample_dataset, + embed_dim=128, + n_heads=4, + n_layers=4, + max_len=256, + batch_size=16, + epochs=5, + lr=1e-4, + save_dir="./gpt2_save", + ) + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params} parameters") + + # STEP 5: Train (saves the best checkpoint to save_dir). + model.train_model(train_dataset, val_dataset=val_dataset) + + # STEP 6: Generate one synthetic patient per real training patient. + synthetic = model.generate(num_samples=len(train_dataset)) + print("\nGenerated synthetic patients (first 3):") + for patient in synthetic[:3]: + print(f" {patient['patient_id']}: {len(patient['visits'])} visits") + print(f" {patient['visits']}") + + # STEP 7: Evaluate. The metrics want one row per (patient, visit, code); + # to_evaluation_dataframe produces exactly that from either real records or + # a generator's output. Real visits are read straight off the index tensor + # through the processor, which knows how to invert its own encoding. + processor = sample_dataset.input_processors["visits"] + index_to_code = {idx: code for code, idx in processor.code_vocab.items()} + + def to_records(subset): + for item in subset: + visits = [] + for row in item["visits"]: + codes = [ + index_to_code[idx] + for idx in processor.visit_code_ids(row) + if index_to_code.get(idx) not in (None, "", "") + ] + if codes: + visits.append(codes) + yield {"visits": visits} + + schema = {"visit_codes": str, "labels": int, "time": int, "id": str} + train_df = to_evaluation_dataframe(to_records(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(to_records(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) + print( + f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " + f"synthetic: {len(syn_df)}" + ) + + # Small settings for the dev subset; raise them on the full cohort. + results = evaluate_synthetic_ehr( + train_ehr=train_df, + test_ehr=test_df, + syn_ehr=syn_df, + sample_size=min(30, len(train_dataset), len(test_dataset)), + mode="lstm", + metrics="all", + lstm_params={"embed_dim": 16, "hidden_dim": 16, "batch_size": 16, "epochs": 3}, + n_bootstraps=5, + n_runs=3, + ) + print("\nGenerative metrics (mean +/- std):") + for name, (mean, std) in results.items(): + print(f" {name:30s} {mean:.4f} +/- {std:.4f}") diff --git a/examples/halo_mimic3.py b/examples/halo_mimic3.py index e5188271d..1a2e11b80 100644 --- a/examples/halo_mimic3.py +++ b/examples/halo_mimic3.py @@ -2,9 +2,9 @@ This example demonstrates: 1. Loading MIMIC-III data -2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sets, - multi-hot -- the encoding HALO consumes; GPT2/PromptEHR instead take - VisitSequenceGeneration, and MedGAN/CorGAN PatientCodeSetGeneration) +2. Applying the EHRGenerationMIMIC3 task (per-visit ICD-9 code sets, multi-hot + -- the encoding HALO consumes). See gpt2_mimic3.py for the sequence form + GPT2/PromptEHR need, and medgan_mimic3.py for the bag-of-codes form. 3. Creating a SampleDataset with a NestedMultiHotProcessor 4. Training the HALO generator with its custom training loop 5. Generating synthetic patients diff --git a/examples/medgan_mimic3.py b/examples/medgan_mimic3.py new file mode 100644 index 000000000..4ff14a795 --- /dev/null +++ b/examples/medgan_mimic3.py @@ -0,0 +1,112 @@ +"""Example: train MedGAN on MIMIC-III and generate synthetic patients. + +MedGAN and CorGAN are bag-of-codes generators: a patient is one multi-hot +vector over the code vocabulary, with no visit axis at all. They read +EHRCodeSetGenerationMIMIC3, which pools every admission's codes into a single +set per patient -- not the per-visit tasks HALO (halo_mimic3.py) or +GPT2/PromptEHR (gpt2_mimic3.py) take. + +Swap MedGAN for CorGAN below and the rest of the script is unchanged. + +This example demonstrates: +1. Loading MIMIC-III data +2. Applying the EHRCodeSetGenerationMIMIC3 task (one ICD-9 code set per patient) +3. Training MedGAN (autoencoder pre-training, then adversarial training) +4. Generating synthetic patients +5. Evaluating with the privacy metrics + +Note on metrics: the utility metric in pyhealth.metrics.generative scores +next-visit prediction, which is meaningless without a visit axis. This example +therefore requests ``metrics="privacy"``. Asking for ``"all"`` here would +produce a utility number that looks real and is not. +""" + +from pyhealth.datasets import MIMIC3Dataset, split_by_patient +from pyhealth.metrics.generative import evaluate_synthetic_ehr +from pyhealth.models import MedGAN +from pyhealth.tasks import EHRCodeSetGenerationMIMIC3, to_evaluation_dataframe + +if __name__ == "__main__": + # STEP 1: Load MIMIC-III. dev=True keeps this small -- start here. + base_dataset = MIMIC3Dataset( + root="/srv/local/data/MIMIC-III/mimic-iii-clinical-database-1.4", + tables=["diagnoses_icd"], + dev=True, + ) + + # STEP 2: Pool each patient's admissions into one code set. min_visits + # still counts real admissions, so single-admission patients are dropped + # before the visit axis is collapsed. + sample_dataset = base_dataset.set_task(EHRCodeSetGenerationMIMIC3()) + print(f"Total samples: {len(sample_dataset)}") + + sample = sample_dataset[0] + # (vocab_size,) -- one row per patient, not per visit. + print(f"Visits tensor shape: {tuple(sample['visits'].shape)}") + + # STEP 3: Split by patient. + train_dataset, val_dataset, test_dataset = split_by_patient( + sample_dataset, [0.8, 0.1, 0.1] + ) + + # STEP 4: Initialize MedGAN (small config for the dev subset). + model = MedGAN( + dataset=sample_dataset, + latent_dim=32, + hidden_dim=32, + discriminator_hidden_dim=64, + batch_size=32, + ae_epochs=10, + gan_epochs=20, + save_dir="./medgan_save", + ) + num_params = sum(p.numel() for p in model.parameters()) + print(f"\nModel initialized with {num_params} parameters") + + # STEP 5: Train (autoencoder first, then the GAN). + model.train_model(train_dataset, val_dataset=val_dataset) + + # STEP 6: Generate one synthetic patient per real training patient. + synthetic = model.generate(num_samples=len(train_dataset)) + print("\nGenerated synthetic patients (first 3):") + for patient in synthetic[:3]: + # visits is a single-element list: the aggregate bag of codes. + print(f" {patient['patient_id']}: {len(patient['visits'][0])} codes") + + # STEP 7: Evaluate. MultiHotProcessor keeps its vocabulary in label_vocab + # (the nested processors call theirs code_vocab), and each row is a + # multi-hot vector, so the codes present are its nonzero columns. + processor = sample_dataset.input_processors["visits"] + index_to_code = {idx: code for code, idx in processor.label_vocab.items()} + + def to_records(subset): + for item in subset: + codes = [ + index_to_code[int(col)] + for col in item["visits"].nonzero(as_tuple=True)[0].tolist() + if index_to_code.get(int(col)) is not None + ] + # One "visit" per patient, matching the generator's output shape. + yield {"visits": [codes]} + + schema = {"visit_codes": str, "labels": int, "time": int, "id": str} + train_df = to_evaluation_dataframe(to_records(train_dataset)).astype(schema) + test_df = to_evaluation_dataframe(to_records(test_dataset)).astype(schema) + syn_df = to_evaluation_dataframe(synthetic).astype(schema) + print( + f"\nEval rows -- train: {len(train_df)}, test: {len(test_df)}, " + f"synthetic: {len(syn_df)}" + ) + + # privacy only -- see the note at the top of this file. + results = evaluate_synthetic_ehr( + train_ehr=train_df, + test_ehr=test_df, + syn_ehr=syn_df, + sample_size=min(30, len(train_dataset), len(test_dataset)), + metrics="privacy", + n_bootstraps=5, + ) + print("\nPrivacy metrics (mean +/- std):") + for name, (mean, std) in results.items(): + print(f" {name:30s} {mean:.4f} +/- {std:.4f}") diff --git a/pyhealth/models/generators/corgan.py b/pyhealth/models/generators/corgan.py index c1636b8b8..0c1c64610 100644 --- a/pyhealth/models/generators/corgan.py +++ b/pyhealth/models/generators/corgan.py @@ -8,7 +8,7 @@ CorGAN treats each patient as a flat bag-of-codes (no visit structure), so it expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, -which the :class:`~pyhealth.tasks.PatientCodeSetGeneration` task provides. +which the :class:`~pyhealth.tasks.EHRCodeSetGenerationMIMIC3` task provides. Training has two phases (mirroring the reference): * a **convolutional autoencoder** is pre-trained with a sparse-friendly BCE @@ -351,7 +351,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "CorGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor (see PatientCodeSetGeneration)." + "MultiHotProcessor (see EHRCodeSetGenerationMIMIC3)." ) self._batch_size = batch_size diff --git a/pyhealth/models/generators/gpt2.py b/pyhealth/models/generators/gpt2.py index 0014e3ca7..ac1d541ac 100644 --- a/pyhealth/models/generators/gpt2.py +++ b/pyhealth/models/generators/gpt2.py @@ -4,7 +4,7 @@ ``generate_synthetic_mimic3_gpt2.py`` (``--mode transformer_baseline``) but plugged into the standard PyHealth ``dataset -> set_task -> SampleDataset -> model`` pipeline. It consumes the -:class:`~pyhealth.tasks.VisitSequenceGeneration` task -- the same extraction +:class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- the same extraction :class:`~pyhealth.models.HALO` uses, but emitting code indices rather than multi-hot rows, since a causal LM reads token ids. @@ -42,7 +42,7 @@ class GPT2(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains ``{"visits": NestedSequenceProcessor}`` -- use the - :class:`~pyhealth.tasks.VisitSequenceGeneration` task -- and whose + :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- and whose ``output_schema`` is empty. embed_dim: GPT-2 embedding dimension (``n_embd``). Must be divisible by ``n_heads``. Default: 512. @@ -89,7 +89,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "GPT2 expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor (see VisitSequenceGeneration)." + "NestedSequenceProcessor (see EHRSequenceGenerationMIMIC3)." ) if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): # Without this the visit row would be read as raw values. Under a @@ -99,7 +99,7 @@ def __init__( f"GPT2 needs a 'visits' processor that can invert its own " f"encoding (a visit_code_ids method); got " f"{type(dataset.input_processors['visits']).__name__}. Use " - "NestedSequenceProcessor, via the VisitSequenceGeneration task." + "NestedSequenceProcessor, via the EHRSequenceGenerationMIMIC3 task." ) self.save_dir = save_dir diff --git a/pyhealth/models/generators/halo.py b/pyhealth/models/generators/halo.py index b3efcd10f..2a51b59d2 100644 --- a/pyhealth/models/generators/halo.py +++ b/pyhealth/models/generators/halo.py @@ -372,7 +372,7 @@ class HALO(BaseModel): Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains ``{"visits": NestedMultiHotProcessor}`` -- use the - :class:`~pyhealth.tasks.VisitMultiHotGeneration` task -- and whose + :class:`~pyhealth.tasks.EHRGenerationMIMIC3` task -- and whose ``output_schema`` is empty. embed_dim: Transformer embedding dimension (``n_embd``). Default: 768. n_heads: Number of attention heads. Must divide ``embed_dim``. @@ -421,7 +421,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "HALO expects an input feature named 'visits' backed by a " - "NestedMultiHotProcessor (see VisitMultiHotGeneration)." + "NestedMultiHotProcessor (see EHRGenerationMIMIC3)." ) self.save_dir = save_dir @@ -503,7 +503,7 @@ def _encode_visits(self, visits: torch.Tensor): # TODO: this counts non-empty rows and then treats the *first* n_visits # rows as the real ones, so a patient like [codes, empty, codes] would # silently lose its last visit. The pre-vectorisation loop had the same - # behaviour, and VisitMultiHotGeneration cannot produce an interior + # behaviour, and EHRGenerationMIMIC3 cannot produce an interior # empty visit (its __call__ skips codeless admissions), so nothing hits # today -- but NestedMultiHotProcessor does emit all-zero rows for empty # visits, so a hand-built SampleDataset can. Fix by masking on row diff --git a/pyhealth/models/generators/medgan.py b/pyhealth/models/generators/medgan.py index 7ce682703..6faea9214 100644 --- a/pyhealth/models/generators/medgan.py +++ b/pyhealth/models/generators/medgan.py @@ -8,7 +8,7 @@ MedGAN treats each patient as a flat bag-of-codes (no visit structure), so it expects an input feature named ``visits`` backed by a ``MultiHotProcessor``, -which the :class:`~pyhealth.tasks.PatientCodeSetGeneration` task provides. +which the :class:`~pyhealth.tasks.EHRCodeSetGenerationMIMIC3` task provides. The training procedure has two phases (mirroring the reference): * a **linear autoencoder** is pre-trained with binary cross-entropy @@ -232,7 +232,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "MedGAN expects an input feature named 'visits' backed by a " - "MultiHotProcessor (see PatientCodeSetGeneration)." + "MultiHotProcessor (see EHRCodeSetGenerationMIMIC3)." ) # The generator's residual connection (``out + residual`` with diff --git a/pyhealth/models/generators/promptehr.py b/pyhealth/models/generators/promptehr.py index 56b26ed36..6c50690d4 100644 --- a/pyhealth/models/generators/promptehr.py +++ b/pyhealth/models/generators/promptehr.py @@ -3,7 +3,7 @@ This is a PyHealth ``BaseModel`` port of PromptEHR (Wang & Sun, EMNLP'22, https://github.com/RyanWangZf/PromptEHR), wrapped so it consumes the standard ``dataset -> set_task -> SampleDataset -> model`` pipeline and shares the same -:class:`~pyhealth.tasks.VisitSequenceGeneration` task as +:class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task as :class:`~pyhealth.models.HALO` and :class:`~pyhealth.models.GPT2`. PromptEHR treats sequential EHRs as a *neural database* and learns to fill in @@ -17,7 +17,7 @@ the way :class:`~pyhealth.models.GPT2` wraps ``GPT2LMHeadModel``. * **Prompt learning.** The reference reparameterizes a learnable prompt from patient baseline demographics and prepends it to the encoder/decoder - (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.VisitSequenceGeneration` + (``ConditionalPrompt``). PyHealth's :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task is *unconditional* (only ``visits``, no baseline features -- exactly like HALO/GPT2), so the prompt reduces to a learnable continuous **soft prefix** prepended to the encoder. This is the prompt-tuning core without the @@ -35,7 +35,7 @@ [CODE_PROMPT] [VISIT_DELIM] ... [EOS] The reference handles several code types (diagnosis / procedure / drug / lab) -each with its own modality prompt token; the PyHealth ``VisitSequenceGeneration`` task +each with its own modality prompt token; the PyHealth ``EHRSequenceGenerationMIMIC3`` task exposes a single ``visits`` modality, so a single ``[CODE_PROMPT]`` token marks it. The code vocabulary is taken from the dataset's ``visits`` processor (which already reserves index 0 for ```` and @@ -63,12 +63,12 @@ class PromptEHR(BaseModel): Trains a BART denoising autoencoder with a learnable soft prompt on patient visit-code streams, then generates synthetic patients by prompt-conditioned encoder-decoder sampling. Generation is **unconditional** (no demographic - conditioning), matching the :class:`~pyhealth.tasks.VisitSequenceGeneration` task. + conditioning), matching the :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task. Args: dataset: A fitted ``SampleDataset`` whose ``input_schema`` contains ``{"visits": NestedSequenceProcessor}`` -- use the - :class:`~pyhealth.tasks.VisitSequenceGeneration` task -- and whose + :class:`~pyhealth.tasks.EHRSequenceGenerationMIMIC3` task -- and whose ``output_schema`` is empty. embed_dim: BART model dimension (``d_model``). Must be divisible by ``n_heads``. Default: 256. @@ -128,7 +128,7 @@ def __init__( if "visits" not in dataset.input_processors: raise ValueError( "PromptEHR expects an input feature named 'visits' backed by a " - "NestedSequenceProcessor (see VisitSequenceGeneration)." + "NestedSequenceProcessor (see EHRSequenceGenerationMIMIC3)." ) if not hasattr(dataset.input_processors["visits"], "visit_code_ids"): # Without this the visit row would be read as raw values. Under a @@ -138,7 +138,7 @@ def __init__( f"PromptEHR needs a 'visits' processor that can invert its own " f"encoding (a visit_code_ids method); got " f"{type(dataset.input_processors['visits']).__name__}. Use " - "NestedSequenceProcessor, via the VisitSequenceGeneration task." + "NestedSequenceProcessor, via the EHRSequenceGenerationMIMIC3 task." ) self.save_dir = save_dir diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 54ee000ab..2e7bd6bab 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -52,12 +52,12 @@ MortalityPredictionStageNetMIMIC4, ) from .generate_ehr import ( - EHRGeneration, + EHRCodeSetGenerationMIMIC3 as EHRCodeSetGenerationMIMIC3, + EHRCodeSetGenerationMIMIC4 as EHRCodeSetGenerationMIMIC4, EHRGenerationMIMIC3, EHRGenerationMIMIC4, - PatientCodeSetGeneration as PatientCodeSetGeneration, - VisitMultiHotGeneration as VisitMultiHotGeneration, - VisitSequenceGeneration as VisitSequenceGeneration, + EHRSequenceGenerationMIMIC3 as EHRSequenceGenerationMIMIC3, + EHRSequenceGenerationMIMIC4 as EHRSequenceGenerationMIMIC4, decode_dataset, to_evaluation_dataframe, ) diff --git a/pyhealth/tasks/generate_ehr.py b/pyhealth/tasks/generate_ehr.py index 738ca3f1c..3f7fcd99d 100644 --- a/pyhealth/tasks/generate_ehr.py +++ b/pyhealth/tasks/generate_ehr.py @@ -1,28 +1,30 @@ """EHR sequence-generation tasks for PyHealth generative models. -These back every generator in :mod:`pyhealth.models.generators`. They extract, -for each patient, the ordered list of visits where each visit is the list of -medical codes recorded in that admission. There is no prediction label, so -``output_schema`` is empty. - -Extraction is shared -- :class:`EHRGeneration` holds all of it -- but the -encoding is not, because each generator family reads its codes in a different -shape: - -- :class:`VisitMultiHotGeneration` -- one multi-hot row per visit, for HALO, - whose transformer consumes multi-hot vectors directly. -- :class:`VisitSequenceGeneration` -- per-visit code indices, for the token - models GPT2 and PromptEHR, which flatten visits into streams of code ids. -- :class:`PatientCodeSetGeneration` -- one pooled code set per patient, for the - bag-of-codes models MedGAN and CorGAN, which have no visit axis. +These back every generator in :mod:`pyhealth.models.generators`. Each one +extracts, for each patient, the codes recorded across their admissions. There +is no prediction label, so ``output_schema`` is empty. + +The classes are flat and independent -- one per (model family, dataset) -- with +the dataset named in the class, because the extraction is MIMIC-shaped: it +assumes an ``admissions`` event type and a ``hadm_id`` linking codes to an +admission. eICU, OMOP and MEDS do not look like that, so a task for those +datasets belongs alongside these rather than inheriting from them. + +============================= ========================== ================== +Task Encoding Models +============================= ========================== ================== +``EHRGenerationMIMIC3/4`` per-visit multi-hot rows HALO +``EHRSequenceGenerationMIMIC3/4`` per-visit code indices GPT2, PromptEHR +``EHRCodeSetGenerationMIMIC3/4`` one code set per patient MedGAN, CorGAN +============================= ========================== ================== Match the task to the model. Handing a model the wrong encoding does not raise: the numbers still have the right shape and dtype, so training runs and produces a confidently wrong result. That is why these are separate classes rather than one task with a flag. -``event_type`` / ``code_attr`` select the dataset's coding columns and can be -passed to any of them, so the encoding and the dataset are independent choices. +``event_type`` / ``code_attr`` are class attributes naming the dataset's coding +columns; override them on an instance to read a different table. Evaluating generated data ------------------------- @@ -92,211 +94,253 @@ logger = logging.getLogger(__name__) -class EHRGeneration(BaseTask): - """Per-visit code extraction for unconditional EHR generators. +def _mimic_visits( + patient: Patient, event_type: str, code_attr: str +) -> list[list[str]]: + """Ordered per-admission code lists for a MIMIC-style patient. - Builds one sample per qualifying patient: the ordered list of visits, each - visit being the list of codes (read from ``code_attr`` on ``event_type`` - events) recorded in that admission. Patients with fewer than ``min_visits`` - qualifying visits are skipped. + Deliberately MIMIC-specific and deliberately not a base-class method: it + assumes an ``admissions`` event type and a ``hadm_id`` linking codes to an + admission. eICU, OMOP and MEDS do not look like this, so a task for those + writes its own extraction rather than inheriting one that would silently + return nothing. - **This class does not set an** ``input_schema`` **and cannot be used - directly.** Extraction is shared, but each generator family wants the codes - in a different shape, and handing a model the wrong one fails silently - rather than loudly. Pick the subclass that matches your model: + Args: + patient: Patient to read. + event_type: Event type carrying the codes (e.g. ``"diagnoses_icd"``). + code_attr: Attribute on those events holding the code string. + + Returns: + One list of codes per admission, admissions with no codes dropped. + """ + visits: list[list[str]] = [] + for admission in patient.get_events(event_type="admissions"): + events = patient.get_events( + event_type=event_type, + filters=[("hadm_id", "==", admission.hadm_id)], + ) + codes = [ + getattr(event, code_attr) + for event in events + if getattr(event, code_attr, None) + ] + if codes: + visits.append(codes) + return visits - ============================== ===================== ================== - Task Encoding Models - ============================== ===================== ================== - :class:`VisitMultiHotGeneration` per-visit multi-hot HALO - :class:`VisitSequenceGeneration` per-visit indices GPT2, PromptEHR - :class:`PatientCodeSetGeneration` one set per patient MedGAN, CorGAN - ============================== ===================== ================== - Args: - code_mapping: Optional vocabulary mapping, see :class:`BaseTask`. - event_type: Event type to pull per admission. Defaults to the class - attribute (``"diagnoses_icd"``). - code_attr: Event attribute holding the code string. Defaults to the - class attribute (``"icd9_code"``). - min_visits: Minimum qualifying visits to keep a patient. Defaults to - the class attribute (2). +class EHRGenerationMIMIC3(BaseTask): + """Per-visit ICD-9 code sets from MIMIC-III as multi-hot rows. For HALO. + + HALO's transformer consumes a multi-hot vector per context position, so + this hands it exactly that and nothing is repacked on the way in. + + Patients with fewer than ``min_visits`` coded admissions are skipped. Examples: - >>> from pyhealth.tasks import VisitSequenceGeneration - >>> task = VisitSequenceGeneration(code_attr="icd_code") - >>> task.code_attr - 'icd_code' + >>> from pyhealth.datasets import MIMIC3Dataset + >>> from pyhealth.tasks import EHRGenerationMIMIC3 + >>> dataset = MIMIC3Dataset( + ... root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"] + ... ) + >>> samples = dataset.set_task(EHRGenerationMIMIC3()) + >>> samples[0]["visits"].shape # (num_visits, vocab_size) + torch.Size([3, 512]) """ - task_name: str = "ehr_generation" + task_name: str = "ehr_generation_mimic3" + input_schema: ClassVar[dict[str, str | type]] = { + "visits": NestedMultiHotProcessor + } output_schema: ClassVar[dict[str, str | type]] = {} event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" min_visits: int = 2 - def __init__( - self, - code_mapping=None, - event_type: str | None = None, - code_attr: str | None = None, - min_visits: int | None = None, - ) -> None: - if not hasattr(type(self), "input_schema"): - raise TypeError( - f"{type(self).__name__} does not declare an encoding. " - "EHRGeneration only holds the shared extraction logic -- use " - "VisitMultiHotGeneration (HALO), VisitSequenceGeneration " - "(GPT2, PromptEHR) or PatientCodeSetGeneration (MedGAN, " - "CorGAN), whichever matches your model." - ) - super().__init__(code_mapping=code_mapping) - # Per-instance overrides, so a dataset preset is a constructor argument - # rather than yet another subclass in the encoding x dataset grid. - if event_type is not None: - self.event_type = event_type - if code_attr is not None: - self.code_attr = code_attr - if min_visits is not None: - self.min_visits = min_visits - - def _visits(self, patient: Patient) -> list[list[str]]: - """Ordered per-admission code lists, empty admissions dropped.""" - visits: list[list[str]] = [] - for admission in patient.get_events(event_type="admissions"): - events = patient.get_events( - event_type=self.event_type, - filters=[("hadm_id", "==", admission.hadm_id)], - ) - codes = [ - getattr(event, self.code_attr) - for event in events - if getattr(event, self.code_attr, None) - ] - if codes: - visits.append(codes) - return visits - def __call__(self, patient: Patient) -> list[dict]: """Extract the per-visit code sequence for a patient.""" - visits = self._visits(patient) + visits = _mimic_visits(patient, self.event_type, self.code_attr) if len(visits) < self.min_visits: return [] return [{"patient_id": patient.patient_id, "visits": visits}] -class VisitMultiHotGeneration(EHRGeneration): - """Per-visit code sets as multi-hot rows. For HALO. +class EHRGenerationMIMIC4(BaseTask): + """Per-visit ICD code sets from MIMIC-IV as multi-hot rows. For HALO. - HALO's transformer consumes a multi-hot vector per context position, so - this hands it exactly that and no repacking happens on the way in. + MIMIC-IV's diagnosis codes live on ``icd_code`` rather than MIMIC-III's + ``icd9_code``; otherwise identical to :class:`EHRGenerationMIMIC3`. Examples: - >>> from pyhealth.tasks import VisitMultiHotGeneration - >>> samples = dataset.set_task(VisitMultiHotGeneration()) + >>> from pyhealth.datasets import MIMIC4Dataset + >>> from pyhealth.tasks import EHRGenerationMIMIC4 + >>> dataset = MIMIC4Dataset( + ... ehr_root="/path/to/mimiciv/2.2/", + ... ehr_tables=["patients", "admissions", "diagnoses_icd"], + ... ) + >>> samples = dataset.set_task(EHRGenerationMIMIC4()) >>> samples[0]["visits"].shape # (num_visits, vocab_size) torch.Size([3, 512]) """ - task_name: str = "ehr_generation_visit_multihot" + task_name: str = "ehr_generation_mimic4" input_schema: ClassVar[dict[str, str | type]] = { "visits": NestedMultiHotProcessor } + output_schema: ClassVar[dict[str, str | type]] = {} + + event_type: str = "diagnoses_icd" + code_attr: str = "icd_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + return [{"patient_id": patient.patient_id, "visits": visits}] -class VisitSequenceGeneration(EHRGeneration): - """Per-visit code indices, right-padded. For GPT2 and PromptEHR. +class EHRSequenceGenerationMIMIC3(BaseTask): + """Per-visit ICD-9 code indices from MIMIC-III. For GPT2 and PromptEHR. Both are token-sequence models: they flatten each visit into a stream of - code ids. Indices are what they need, so this avoids encoding to multi-hot - and decoding straight back. + code ids, so indices are what they want. Handing them the multi-hot form + means encoding a code set and decoding it straight back. Examples: - >>> from pyhealth.tasks import VisitSequenceGeneration - >>> samples = dataset.set_task(VisitSequenceGeneration()) + >>> from pyhealth.datasets import MIMIC3Dataset + >>> from pyhealth.tasks import EHRSequenceGenerationMIMIC3 + >>> dataset = MIMIC3Dataset( + ... root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"] + ... ) + >>> samples = dataset.set_task(EHRSequenceGenerationMIMIC3()) >>> samples[0]["visits"].shape # (num_visits, max_codes_per_visit) torch.Size([3, 12]) """ - task_name: str = "ehr_generation_visit_sequence" + task_name: str = "ehr_sequence_generation_mimic3" input_schema: ClassVar[dict[str, str | type]] = { "visits": NestedSequenceProcessor } + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" + code_attr: str = "icd9_code" + min_visits: int = 2 -class PatientCodeSetGeneration(EHRGeneration): - """One code set per patient, visit structure discarded. For MedGAN/CorGAN. + def __call__(self, patient: Patient) -> list[dict]: + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + return [{"patient_id": patient.patient_id, "visits": visits}] - Bag-of-codes generators emit a single aggregate vector per patient, so the - visit axis is collapsed here rather than inside the model. ``min_visits`` - still applies -- it filters on the patient's real visit count before the - codes are pooled. - Note: - Because the visit axis is gone, the next-visit utility metric in - :mod:`pyhealth.metrics.generative` is not meaningful for these models; - see this module's header. +class EHRSequenceGenerationMIMIC4(BaseTask): + """Per-visit ICD code indices from MIMIC-IV. For GPT2 and PromptEHR. Examples: - >>> from pyhealth.tasks import PatientCodeSetGeneration - >>> samples = dataset.set_task(PatientCodeSetGeneration()) - >>> samples[0]["visits"].shape # (vocab_size,) - torch.Size([512]) + >>> from pyhealth.datasets import MIMIC4Dataset + >>> from pyhealth.tasks import EHRSequenceGenerationMIMIC4 + >>> dataset = MIMIC4Dataset( + ... ehr_root="/path/to/mimiciv/2.2/", + ... ehr_tables=["patients", "admissions", "diagnoses_icd"], + ... ) + >>> samples = dataset.set_task(EHRSequenceGenerationMIMIC4()) + >>> samples[0]["visits"].shape # (num_visits, max_codes_per_visit) + torch.Size([3, 12]) """ - task_name: str = "ehr_generation_patient_codeset" - input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + task_name: str = "ehr_sequence_generation_mimic4" + input_schema: ClassVar[dict[str, str | type]] = { + "visits": NestedSequenceProcessor + } + output_schema: ClassVar[dict[str, str | type]] = {} + + event_type: str = "diagnoses_icd" + code_attr: str = "icd_code" + min_visits: int = 2 def __call__(self, patient: Patient) -> list[dict]: - """Pool every visit's codes into one per-patient set.""" - visits = self._visits(patient) + """Extract the per-visit code sequence for a patient.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) if len(visits) < self.min_visits: return [] - codes = sorted({code for visit in visits for code in visit}) - return [{"patient_id": patient.patient_id, "visits": codes}] + return [{"patient_id": patient.patient_id, "visits": visits}] -class EHRGenerationMIMIC3(VisitMultiHotGeneration): - """EHR generation task for MIMIC-III (ICD-9 diagnosis codes), for HALO. +class EHRCodeSetGenerationMIMIC3(BaseTask): + """One pooled ICD-9 code set per MIMIC-III patient. For MedGAN and CorGAN. - A :class:`VisitMultiHotGeneration` preset. For GPT2/PromptEHR on MIMIC-III - use ``VisitSequenceGeneration()``, whose defaults are already MIMIC-III's. + Bag-of-codes generators emit a single aggregate vector per patient, so the + visit axis is collapsed here rather than inside the model. ``min_visits`` + still counts real admissions, before the codes are pooled. + + Note: + With the visit axis gone, the next-visit utility metric in + :mod:`pyhealth.metrics.generative` is not meaningful for these models; + see this module's header. Examples: >>> from pyhealth.datasets import MIMIC3Dataset - >>> from pyhealth.tasks import EHRGenerationMIMIC3 + >>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC3 >>> dataset = MIMIC3Dataset( - ... root="/path/to/mimic-iii/1.4", - ... tables=["diagnoses_icd"], + ... root="/path/to/mimic-iii/1.4", tables=["diagnoses_icd"] ... ) - >>> samples = dataset.set_task(EHRGenerationMIMIC3()) + >>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC3()) + >>> samples[0]["visits"].shape # (vocab_size,) + torch.Size([512]) """ - task_name: str = "ehr_generation_mimic3" + task_name: str = "ehr_codeset_generation_mimic3" + input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" code_attr: str = "icd9_code" + min_visits: int = 2 + def __call__(self, patient: Patient) -> list[dict]: + """Pool every visit's codes into one per-patient set.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + codes = sorted({code for visit in visits for code in visit}) + return [{"patient_id": patient.patient_id, "visits": codes}] -class EHRGenerationMIMIC4(VisitMultiHotGeneration): - """EHR generation task for MIMIC-IV (ICD diagnosis codes), for HALO. - A :class:`VisitMultiHotGeneration` preset. For another encoding on MIMIC-IV - pass the same columns, e.g. ``VisitSequenceGeneration(code_attr="icd_code")``. +class EHRCodeSetGenerationMIMIC4(BaseTask): + """One pooled ICD code set per MIMIC-IV patient. For MedGAN and CorGAN. Examples: >>> from pyhealth.datasets import MIMIC4Dataset - >>> from pyhealth.tasks import EHRGenerationMIMIC4 + >>> from pyhealth.tasks import EHRCodeSetGenerationMIMIC4 >>> dataset = MIMIC4Dataset( ... ehr_root="/path/to/mimiciv/2.2/", ... ehr_tables=["patients", "admissions", "diagnoses_icd"], ... ) - >>> samples = dataset.set_task(EHRGenerationMIMIC4()) + >>> samples = dataset.set_task(EHRCodeSetGenerationMIMIC4()) + >>> samples[0]["visits"].shape # (vocab_size,) + torch.Size([512]) """ - task_name: str = "ehr_generation_mimic4" + task_name: str = "ehr_codeset_generation_mimic4" + input_schema: ClassVar[dict[str, str | type]] = {"visits": MultiHotProcessor} + output_schema: ClassVar[dict[str, str | type]] = {} + event_type: str = "diagnoses_icd" code_attr: str = "icd_code" + min_visits: int = 2 + + def __call__(self, patient: Patient) -> list[dict]: + """Pool every visit's codes into one per-patient set.""" + visits = _mimic_visits(patient, self.event_type, self.code_attr) + if len(visits) < self.min_visits: + return [] + codes = sorted({code for visit in visits for code in visit}) + return [{"patient_id": patient.patient_id, "visits": codes}] # ---------------------------------------------------------------------------- @@ -322,7 +366,7 @@ def to_evaluation_dataframe( Args: records: Iterable of ``{"visits": [[code, ...], ...]}`` dicts. Both the - :class:`EHRGeneration` task output and a generator's ``generate()`` + generation tasks' output and a generator's ``generate()`` output have this shape. label_fn: Optional callable mapping a record to a binary patient label (0/1) used by the utility metrics. Defaults to all-zeros. @@ -365,7 +409,7 @@ def to_evaluation_dataframe( def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: - """Decode a processed EHRGeneration ``SampleDataset`` back into records. + """Decode a processed multi-hot ``SampleDataset`` back into code records. Inverts the :class:`~pyhealth.processors.NestedMultiHotProcessor` encoding using its vocabulary (skipping ````/````), yielding one @@ -377,7 +421,8 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: or count within a visit. Args: - sample_dataset: A ``SampleDataset`` produced by :class:`EHRGeneration`. + sample_dataset: A ``SampleDataset`` (or a split of one) produced by + :class:`EHRGenerationMIMIC3` / :class:`EHRGenerationMIMIC4`. feature_key: Input feature key holding the nested code sequence. Default ``"visits"``. @@ -401,7 +446,7 @@ def decode_dataset(sample_dataset, feature_key: str = "visits") -> list[dict]: if not isinstance(processor, NestedMultiHotProcessor): raise TypeError( f"decode_dataset inverts the multi-hot encoding, but '{feature_key}' " - f"is a {type(processor).__name__}. Use VisitMultiHotGeneration, or " + f"is a {type(processor).__name__}. Use one of the multi-hot tasks\n (EHRGenerationMIMIC3/4), or " "read the codes off the index tensor directly." ) index_to_code = {idx: code for code, idx in processor.code_vocab.items()} diff --git a/tests/core/test_generator_encodings.py b/tests/core/test_generator_encodings.py index c4d109fa0..fee993523 100644 --- a/tests/core/test_generator_encodings.py +++ b/tests/core/test_generator_encodings.py @@ -19,11 +19,14 @@ NestedSequenceProcessor, ) from pyhealth.tasks import ( + EHRCodeSetGenerationMIMIC3, + EHRCodeSetGenerationMIMIC4, EHRGenerationMIMIC3, - PatientCodeSetGeneration, - VisitMultiHotGeneration, - VisitSequenceGeneration, + EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, + EHRSequenceGenerationMIMIC4, ) +from pyhealth.tasks.base_task import BaseTask SAMPLES = [ {"patient_id": "p0", "visits": [["A05B", "A05C"], ["A11D"], ["C129"]]}, @@ -68,21 +71,21 @@ def get_events(self, event_type, filters=None): class TestExtraction(unittest.TestCase): - """The shared __call__, and the pooling PatientCodeSetGeneration adds.""" + """Extraction, and the pooling the code-set tasks add.""" VISITS: ClassVar[list] = [["A05B", "A05C"], ["A11D"], ["A05B"]] def test_per_visit_tasks_keep_visit_structure(self): patient = _Patient("p0", self.VISITS) - samples = VisitMultiHotGeneration()(patient) + samples = EHRGenerationMIMIC3()(patient) self.assertEqual(len(samples), 1) self.assertEqual(samples[0]["visits"], self.VISITS) # Same extraction regardless of encoding -- only input_schema differs. - self.assertEqual(VisitSequenceGeneration()(patient)[0]["visits"], + self.assertEqual(EHRSequenceGenerationMIMIC3()(patient)[0]["visits"], self.VISITS) def test_codeset_task_pools_and_dedupes(self): - samples = PatientCodeSetGeneration()(_Patient("p0", self.VISITS)) + samples = EHRCodeSetGenerationMIMIC3()(_Patient("p0", self.VISITS)) self.assertEqual(len(samples), 1) # One flat set: A05B appears in two visits and survives once. self.assertEqual(samples[0]["visits"], ["A05B", "A05C", "A11D"]) @@ -90,12 +93,12 @@ def test_codeset_task_pools_and_dedupes(self): def test_min_visits_counts_real_visits_before_pooling(self): """Pooling must not let a 1-visit patient past a min_visits=2 filter.""" one_visit = _Patient("p1", [["A05B", "A05C", "A11D"]]) - self.assertEqual(PatientCodeSetGeneration()(one_visit), []) - self.assertEqual(VisitMultiHotGeneration()(one_visit), []) + self.assertEqual(EHRCodeSetGenerationMIMIC3()(one_visit), []) + self.assertEqual(EHRGenerationMIMIC3()(one_visit), []) def test_codeless_admissions_are_dropped(self): patient = _Patient("p2", [["A05B"], [], ["A11D"]]) - self.assertEqual(VisitMultiHotGeneration()(patient)[0]["visits"], + self.assertEqual(EHRGenerationMIMIC3()(patient)[0]["visits"], [["A05B"], ["A11D"]]) @@ -103,38 +106,42 @@ class TestTaskEncodings(unittest.TestCase): """Each task declares the processor its models consume.""" def test_each_family_gets_its_own_encoding(self): - self.assertIs( - VisitMultiHotGeneration.input_schema["visits"], NestedMultiHotProcessor - ) - self.assertIs( - VisitSequenceGeneration.input_schema["visits"], NestedSequenceProcessor - ) - self.assertIs( - PatientCodeSetGeneration.input_schema["visits"], MultiHotProcessor - ) - - def test_base_task_refuses_to_be_used_directly(self): - """EHRGeneration is extraction only, and says so instead of failing late.""" - from pyhealth.tasks import EHRGeneration - - self.assertFalse(hasattr(EHRGeneration, "input_schema")) - with self.assertRaises(TypeError) as ctx: - EHRGeneration() - self.assertIn("VisitMultiHotGeneration", str(ctx.exception)) - - def test_dataset_presets_stay_multihot(self): - """The MIMIC presets were HALO tasks and must remain so.""" - self.assertIs( - EHRGenerationMIMIC3.input_schema["visits"], NestedMultiHotProcessor - ) - self.assertTrue(issubclass(EHRGenerationMIMIC3, VisitMultiHotGeneration)) - - def test_columns_are_settable_per_instance(self): - """Encoding and dataset are independent choices, not a class grid.""" - task = VisitSequenceGeneration(code_attr="icd_code", min_visits=3) - self.assertEqual(task.code_attr, "icd_code") - self.assertEqual(task.min_visits, 3) - self.assertEqual(VisitSequenceGeneration.code_attr, "icd9_code") + for task in (EHRGenerationMIMIC3, EHRGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], NestedMultiHotProcessor) + for task in (EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], NestedSequenceProcessor) + for task in (EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4): + self.assertIs(task.input_schema["visits"], MultiHotProcessor) + + def test_tasks_are_flat(self): + """No task inherits from another: the MIMIC extraction is not a base. + + A parent class would invite subclassing it for eICU/OMOP, where there + is no ``admissions`` event type and no ``hadm_id`` -- which would + return no samples rather than fail. + """ + tasks = [ + EHRGenerationMIMIC3, EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4, + EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4, + ] + for task in tasks: + with self.subTest(task=task.__name__): + self.assertEqual(task.__bases__, (BaseTask,)) + self.assertIn("MIMIC", task.__name__) + + def test_task_names_are_unique(self): + tasks = [ + EHRGenerationMIMIC3, EHRGenerationMIMIC4, + EHRSequenceGenerationMIMIC3, EHRSequenceGenerationMIMIC4, + EHRCodeSetGenerationMIMIC3, EHRCodeSetGenerationMIMIC4, + ] + names = [t.task_name for t in tasks] + self.assertEqual(len(names), len(set(names))) + + def test_mimic4_reads_the_mimic4_code_column(self): + self.assertEqual(EHRGenerationMIMIC3.code_attr, "icd9_code") + self.assertEqual(EHRGenerationMIMIC4.code_attr, "icd_code") class TestVisitCodeIds(unittest.TestCase):