From a92c25a556df94de363196c843a45fad64c08c53 Mon Sep 17 00:00:00 2001 From: William Pang Date: Mon, 10 Aug 2026 20:27:57 -0700 Subject: [PATCH 01/27] Initial Push --- .../unified_embedding_e2e_mimic4.py | 500 ++++++++ pyhealth/models/bottleneck_transformer.py | 520 ++++++++ pyhealth/models/ehrmamba.py | 101 +- pyhealth/models/embedding/base.py | 52 + pyhealth/models/embedding/unified.py | 595 +++++++++ pyhealth/models/embedding/vanilla.py | 366 ++++++ pyhealth/models/embedding/vision.py | 384 ++++++ pyhealth/models/jamba_ehr.py | 93 +- pyhealth/models/rnn.py | 88 +- pyhealth/models/transformer.py | 140 ++- pyhealth/processors/time_image_processor.py | 105 +- .../processors/tuple_time_text_processor.py | 56 +- .../will/condor/labs_only/labs_only_rnn.sub | 45 + .../condor/labs_only/run_labs_only_rnn.sh | 168 +++ pyhealth/tasks/multimodal_mimic4.py | 1064 +++++++++++++++++ pyhealth/trainer.py | 157 ++- 16 files changed, 4292 insertions(+), 142 deletions(-) create mode 100644 examples/mortality_prediction/unified_embedding_e2e_mimic4.py create mode 100644 pyhealth/models/bottleneck_transformer.py create mode 100644 pyhealth/models/embedding/base.py create mode 100644 pyhealth/models/embedding/unified.py create mode 100644 pyhealth/models/embedding/vanilla.py create mode 100644 pyhealth/models/embedding/vision.py create mode 100644 pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub create mode 100644 pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh create mode 100644 pyhealth/tasks/multimodal_mimic4.py diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py new file mode 100644 index 000000000..d7182955a --- /dev/null +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -0,0 +1,500 @@ +"""End-to-end protocol runner for Unified Embedding on MIMIC-IV. + +Trains and evaluates a unified-embedding model (RNN / Transformer / +BottleneckTransformer / EHRMamba / JambaEHR) on a MIMIC-IV mortality task, +then writes per-sample predictions to CSV. + +Tasks +----- +--task labs (default) + LabsMIMIC4: 10-dim lab vectors only. + +--task notes_labs (recommended for multimodal) + NotesLabsMIMIC4: notes + 10-dim lab vectors. + +--task notes_labs_cxr + NotesLabsCXRMIMIC4: notes + labs + chest-xray. + +Example +------- + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /path/to/mimiciv/2.2 \\ + --task labs \\ + --model transformer \\ + --heads 4 --num-layers 2 \\ + --dev --device cpu \\ + --epochs 10 --batch-size 32 --lr 1e-3 \\ + --output-dir ./output/unified_e2e + + # EHRMamba on full dataset (no --dev): + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\ + --task notes_labs --model ehrmamba \\ + --embedding-dim 128 --num-layers 2 --seed 42 + + # JambaEHR: + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\ + --task notes_labs --model jambaehr \\ + --embedding-dim 128 --jamba-transformer-layers 2 --jamba-mamba-layers 6 +""" + +from __future__ import annotations + +import argparse +import csv +from pathlib import Path +from typing import Any, Dict, Optional, Tuple + +import numpy as np + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import RNN, Transformer, UnifiedMultimodalEmbeddingModel +from pyhealth.models.bottleneck_transformer import BottleneckTransformer +from pyhealth.models.ehrmamba import EHRMamba +from pyhealth.models.jamba_ehr import JambaEHR +from pyhealth.tasks.multimodal_mimic4 import ( + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, +) +from pyhealth.trainer import Trainer +from pyhealth.utils import set_seed + + +class WandbLogger: + + def __init__( + self, + enabled: bool, + project: str, + entity: Optional[str], + run_name: str, + tags: list[str], + config: Dict[str, Any], + ) -> None: + self.enabled = enabled + self._run = None + if self.enabled: + import wandb + + self._run = wandb.init( + project=project, + entity=entity, + name=run_name, + tags=tags, + config=config, + ) + + def log(self, data: Dict[str, Any], step: Optional[int] = None) -> None: + if self.enabled: + self._run.log(data, step=step) + + def finish(self) -> None: + if self.enabled: + self._run.finish() + + +def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: + ehr_tables = ["labevents"] + note_tables = None + cxr_kwargs = {} + + if args.task == "notes_labs": + if not args.note_root: + raise ValueError("--task notes_labs requires --note-root.") + note_tables = ["discharge", "radiology"] + + if args.task == "notes_labs_cxr": + if not args.note_root: + raise ValueError("--task notes_labs_cxr requires --note-root.") + if not args.cxr_root: + raise ValueError("--task notes_labs_cxr requires --cxr-root.") + note_tables = ["discharge", "radiology"] + cxr_kwargs = dict( + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + cxr_tables=["metadata", "negbio", "chexpert", "split"], + ) + + return MIMIC4Dataset( + ehr_root=args.ehr_root, + ehr_tables=ehr_tables, + note_root=args.note_root if note_tables else None, + note_tables=note_tables, + cache_dir=args.cache_dir, + dev=args.dev if args.dev else False, + num_workers=args.num_workers, + **cxr_kwargs, + ) + + +def _build_task(args: argparse.Namespace): + if args.task == "notes_labs": + return NotesLabsMIMIC4( + window_hours=args.observation_window_hours, + ) + if args.task == "notes_labs_cxr": + return NotesLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + if args.task == "labs": + return LabsMIMIC4(window_hours=args.observation_window_hours) + raise ValueError(f"Unknown task: {args.task}") + + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) + if len(train_ds) == 0 or len(test_ds) == 0: + train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) + return train_ds, val_ds, test_ds + + +def _build_model(args: argparse.Namespace, sample_dataset: Any): + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + freeze_text_encoder=args.freeze_encoder, + ) + + if args.model == "rnn": + return RNN( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + unified_embedding=unified, + rnn_type=args.rnn_type, + num_layers=args.rnn_layers, + dropout=args.dropout, + bidirectional=args.bidirectional, + ) + if args.model == "transformer": + return Transformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + heads=args.heads, + num_layers=args.num_layers, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "bottleneck_transformer": + return BottleneckTransformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + bottlenecks_n=args.bottlenecks_n, + fusion_startidx=args.fusion_startidx, + num_layers=args.num_layers, + heads=args.heads, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "ehrmamba": + return EHRMamba( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_layers=args.num_layers, + state_size=args.mamba_state_size, + conv_kernel=args.mamba_conv_kernel, + dropout=args.dropout, + unified_embedding=unified, + ) + if args.model == "jambaehr": + return JambaEHR( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_transformer_layers=args.jamba_transformer_layers, + num_mamba_layers=args.jamba_mamba_layers, + heads=args.heads, + dropout=args.dropout, + state_size=args.mamba_state_size, + conv_kernel=args.mamba_conv_kernel, + unified_embedding=unified, + ) + raise ValueError(f"Unknown model: {args.model}") + + +def _write_predictions( + output_csv: Path, + patient_ids: list[str], + y_true: np.ndarray, + y_prob: np.ndarray, +) -> None: + output_csv.parent.mkdir(parents=True, exist_ok=True) + + y_true_flat = y_true.reshape(-1).tolist() + y_prob_flat = y_prob.reshape(-1).tolist() + + with output_csv.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter( + f, + fieldnames=["patient_id", "y_true", "y_prob", "y_pred_threshold_0_5"], + ) + writer.writeheader() + for idx, prob in enumerate(y_prob_flat): + writer.writerow( + { + "patient_id": patient_ids[idx], + "y_true": int(y_true_flat[idx]), + "y_prob": float(prob), + "y_pred_threshold_0_5": int(float(prob) >= 0.5), + } + ) + + +def run(args: argparse.Namespace) -> Path: + set_seed(args.seed) + + base_dataset = _build_base_dataset(args) + task = _build_task(args) + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + + if len(sample_dataset) == 0: + raise RuntimeError( + "Task produced zero samples. Check roots/tables or adjust settings." + ) + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + model = _build_model(args, sample_dataset) + + train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True) + val_loader = ( + get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False) + if len(val_ds) > 0 + else None + ) + test_loader = ( + get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + if len(test_ds) > 0 + else None + ) + + exp_name = f"{args.model}_seed{args.seed}" + output_dir = Path(args.output_dir) + + wandb_logger = WandbLogger( + enabled=args.wandb, + project=args.wandb_project, + entity=args.wandb_entity, + run_name=args.wandb_run_name or exp_name, + tags=args.wandb_tags.split(",") if args.wandb_tags else [args.task, args.model], + config=vars(args), + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc", "f1", "accuracy"], + device=args.device, + enable_logging=True, + output_path=str(output_dir), + exp_name=exp_name, + ) + + # BottleneckTransformer is more fragile on full MIMIC-IV with no warmup. + # Use safer defaults unless explicitly overridden from CLI. + effective_lr = args.lr + effective_max_grad_norm = args.max_grad_norm + optimizer_params = {} + + if args.model == "bottleneck_transformer": + if effective_lr is None: + effective_lr = 1e-4 + if effective_max_grad_norm is None: + effective_max_grad_norm = 0.5 + optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6 + else: + if effective_lr is None: + effective_lr = 1e-4 + if effective_max_grad_norm is None: + effective_max_grad_norm = 1.0 + if args.adam_eps is not None: + optimizer_params["eps"] = args.adam_eps + + optimizer_params["lr"] = effective_lr + + if args.epochs > 0 and len(train_ds) > 0: + metrics_history = trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params=optimizer_params, + weight_decay=args.weight_decay, + max_grad_norm=effective_max_grad_norm, + monitor="pr_auc", + load_best_model_at_last=True, + patience=args.patience, + use_amp=args.use_amp, + amp_dtype=args.amp_dtype, + ) + for epoch_record in metrics_history: + wandb_logger.log(epoch_record, step=epoch_record["epoch"]) + + if wandb_logger.enabled and test_loader is not None: + test_scores = trainer.evaluate(test_loader) + wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + + output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" + _write_predictions(output_csv, patient_ids, y_true, y_prob) + + wandb_logger.finish() + + return output_csv + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run E2E unified embedding on MIMIC-IV with any of six sequence heads." + ) + parser.add_argument("--ehr-root", type=str, required=True) + parser.add_argument("--note-root", type=str, default=None) + parser.add_argument("--cxr-root", type=str, default=None) + parser.add_argument("--cxr-variant", type=str, default="sunlab", choices=["default", "sunlab"]) + parser.add_argument("--cache-dir", type=str, default=None) + parser.add_argument("--output-dir", type=str, default="./output/unified_e2e") + + parser.add_argument( + "--task", + type=str, + choices=["labs", "notes_labs", "notes_labs_cxr"], + default="labs", + help=( + "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " + "Recommended for multimodal. " + "notes_labs_cxr: notes_labs plus in-window chest X-rays; requires " + "--note-root and --cxr-root." + ), + ) + parser.add_argument( + "--model", + type=str, + choices=["rnn", "transformer", "bottleneck_transformer", + "ehrmamba", "jambaehr"], + default="rnn", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=64) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument( + "--lr", + type=float, + default=None, + help="Learning rate. Default is 1e-4 for all models.", + ) + parser.add_argument( + "--adam-eps", + type=float, + default=None, + help=( + "Adam epsilon. Default is model-specific: 1e-8 for non-BT models, " + "1e-6 for bottleneck_transformer." + ), + ) + parser.add_argument("--weight-decay", type=float, default=0.0) + parser.add_argument("--device", type=str, default=None) + parser.add_argument( + "--use-amp", + action="store_true", + help="Enable automatic mixed precision training to reduce GPU memory usage.", + ) + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + help="AMP dtype when --use-amp is set. bf16 is more stable (default).", + ) + parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--patience", type=int, default=None) + parser.add_argument( + "--dev", + nargs="?", + type=int, + const=1000, + default=0, + help=( + "Dev mode: limit dataset to N patients for fast iteration. " + "--dev (no value) defaults to 1000 patients. " + "--dev 5000 limits to 5000. Omit for full dataset." + ), + ) + parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument( + "--freeze-encoder", + action="store_true", + default=False, + help=( + "Freeze pretrained BERT text encoder weights and train only the " + "downstream backbone (RNN/Transformer head + projection layer). " + ), + ) + parser.add_argument("--rnn-type", type=str, default="GRU") + parser.add_argument("--rnn-layers", type=int, default=1) + parser.add_argument("--bidirectional", action="store_true") + + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--num-layers", type=int, default=2) + + parser.add_argument("--bottlenecks-n", type=int, default=4) + parser.add_argument("--fusion-startidx", type=int, default=1) + + parser.add_argument( + "--max-grad-norm", + type=float, + default=None, + help=( + "Gradient clipping max norm. Default is model-specific: None for " + "non-BT models, 0.5 for bottleneck_transformer." + ), + ) + + parser.add_argument( + "--wandb", + action="store_true", + default=False, + help="Log training/eval metrics to Weights & Biases.", + ) + parser.add_argument("--wandb-project", type=str, default="pyhealth-mortality") + parser.add_argument("--wandb-entity", type=str, default=None) + parser.add_argument( + "--wandb-run-name", + type=str, + default=None, + help="Defaults to '{model}_seed{seed}' if unset.", + ) + parser.add_argument( + "--wandb-tags", + type=str, + default=None, + help="Comma-separated wandb tags, e.g. 'labs,rnn'. Defaults to '{task},{model}' if unset.", + ) + + parser.add_argument("--mamba-state-size", type=int, default=16, + help="SSM state size for EHRMamba and JambaEHR blocks.") + parser.add_argument("--mamba-conv-kernel", type=int, default=4, + help="Causal conv kernel size for EHRMamba and JambaEHR blocks.") + parser.add_argument("--jamba-transformer-layers", type=int, default=2, + help="Number of Transformer (attention) layers in JambaEHR.") + parser.add_argument("--jamba-mamba-layers", type=int, default=6, + help="Number of Mamba (SSM) layers in JambaEHR.") + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + output_csv_path = run(args) + print(f"Saved predictions to: {output_csv_path}") \ No newline at end of file diff --git a/pyhealth/models/bottleneck_transformer.py b/pyhealth/models/bottleneck_transformer.py new file mode 100644 index 000000000..9d20f2549 --- /dev/null +++ b/pyhealth/models/bottleneck_transformer.py @@ -0,0 +1,520 @@ +from typing import Any, Dict, List, Optional, Tuple, Union, cast + +import torch +import torch.nn as nn + +from pyhealth.datasets import SampleDataset +from pyhealth.models import BaseModel +from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel + + +class MultimodalBottleneckTransformerEncoder(nn.Module): + """ + Generalized Bottleneck Transformer Encoder for N modalities. + Based on "Attention Bottlenecks for Multimodal Fusion" (Nagrani et al., NeurIPS 2021). + """ + + def __init__( + self, + n_modality: int, + bottlenecks_n: int, + fusion_startidx: int, + n_layers: int, + n_head: int, + d_model: int, + d_ff: int, + dropout: float = 0.1, + ): + super(MultimodalBottleneckTransformerEncoder, self).__init__() + + self.n_modality = n_modality + self.fusion_startidx = fusion_startidx + self.n_layers = n_layers + self.n_fusion_layers = n_layers - fusion_startidx + self.n_prefusion = fusion_startidx + self.d_model = d_model + self.n_bottlenecks = bottlenecks_n + + # Shared Bottleneck Tokens — small init to avoid early gradient explosion + self.bottlenecks = nn.Parameter(torch.randn(1, bottlenecks_n, d_model) * 0.02) + + # Prefusion Stacks: independent layers per modality + self.prefusion_stacks = nn.ModuleList([ + nn.ModuleList([ + nn.TransformerEncoderLayer( + d_model=d_model, + nhead=n_head, + dim_feedforward=d_ff, + dropout=dropout, + batch_first=True + ) for _ in range(n_modality) + ]) for _ in range(self.n_prefusion) + ]) + + # Fusion Stacks: processes [bottleneck_tokens || modality_tokens] + self.fusion_stacks = nn.ModuleList([ + nn.ModuleList([ + nn.TransformerEncoderLayer( + d_model=d_model, + nhead=n_head, + dim_feedforward=d_ff, + dropout=dropout, + batch_first=True + ) for _ in range(n_modality) + ]) for _ in range(self.n_fusion_layers) + ]) + + def forward_prefusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]: + for enc_layers in self.prefusion_stacks: + enc_outputs = [] + for modal_idx, enc_layer in enumerate(enc_layers): + # Apply mask to padding tokens (src_key_padding_mask requires True for ignoring) + # True in mask = invalid/padding + enc_out = enc_layer(enc_inputs[modal_idx], src_key_padding_mask=~masks[modal_idx] if masks[modal_idx] is not None else None) + enc_outputs.append(enc_out) + enc_inputs = enc_outputs + return enc_inputs + + def forward_fusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor], bottleneck_tokens: torch.Tensor, valid_modalities: List[torch.Tensor]) -> List[torch.Tensor]: + # valid_modalities: [B] list of boolean/float tensors indicating if modality is present + batch_size = enc_inputs[0].size(0) + + for modality_encoders in self.fusion_stacks: + enc_outputs = [] + bottleneck_tokens_modality_sum = torch.zeros_like(bottleneck_tokens) + sum_of_modalities = torch.zeros(batch_size, 1, 1, device=bottleneck_tokens.device) + + for idx, enc_layer in enumerate(modality_encoders): + # Concatenate bottleneck tokens with modality tokens + # bottleneck_tokens: [B, num_bottlenecks, d_model] + # enc_inputs[idx]: [B, seq_len, d_model] + fused_input = torch.cat([bottleneck_tokens, enc_inputs[idx]], dim=1) + + # Padding mask for bottleneck tokens is always False (i.e. valid) + # [B, num_bottlenecks] of False + b_mask = torch.zeros(batch_size, self.n_bottlenecks, dtype=torch.bool, device=fused_input.device) + + # Modality padding mask + m_mask = ~masks[idx] if masks[idx] is not None else torch.zeros(batch_size, enc_inputs[idx].size(1), dtype=torch.bool, device=fused_input.device) + + combined_mask = torch.cat([b_mask, m_mask], dim=1) + + # Pass through the layer + enc_out = enc_layer(fused_input, src_key_padding_mask=combined_mask) + + # The output consists of processed bottleneck tokens and modality tokens + # [B, num_bottlenecks, d_model] and [B, seq_len, d_model] + bottleneck_hidden_tokens = enc_out[:, :self.n_bottlenecks, :] + modality_hidden_tokens = enc_out[:, self.n_bottlenecks:, :] + enc_outputs.append(modality_hidden_tokens) + + # Average updated bottlenecks from valid modalities + modality_is_valid = valid_modalities[idx].view(batch_size, 1, 1) + bottleneck_tokens_modality_sum += bottleneck_hidden_tokens * modality_is_valid + sum_of_modalities += modality_is_valid + + # Prevent division by zero if all modalities are missing + # If sum_of_modalities is 0, just pass zeros (or keep previous bottleneck_tokens) + # sum_of_modalities = torch.clamp(sum_of_modalities, min=1.0) + avg_divisor = sum_of_modalities.clone() + avg_divisor[avg_divisor == 0] = 1.0 + + bottleneck_tokens = bottleneck_tokens_modality_sum / avg_divisor + enc_inputs = enc_outputs + + return enc_inputs + + def forward(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]: + batch_size = enc_inputs[0].size(0) + + # Determine if a modality is valid for each instance in the batch + # A modality is valid if it has at least one True in its mask + valid_modalities = [] + for mask, inp in zip(masks, enc_inputs): + if mask is not None: + # [B] - True if there's any valid token (1/True) + valid = mask.any(dim=1).float() + else: + valid = torch.ones(batch_size, device=inp.device) + valid_modalities.append(valid) + + bottleneck_tokens = self.bottlenecks.expand(batch_size, -1, -1) + + enc_inputs = self.forward_prefusion(enc_inputs, masks) + enc_inputs = self.forward_fusion(enc_inputs, masks, bottleneck_tokens, valid_modalities) + + return enc_inputs + + +class BottleneckTransformer(BaseModel): + """Bottleneck Transformer model for PyHealth datasets. + + Per-field mode: each feature stream is embedded with :class:`EmbeddingModel`, + prefixed with a learnable per-modality ``[CLS]`` token, processed by + independent prefusion layers, then fused via shared bottleneck tokens. + The per-modality ``[CLS]`` embeddings are averaged and fed to the + classification head. + + Unified mode (``unified_embedding`` supplied): all temporal fields are + jointly embedded and time-sorted by + :class:`~pyhealth.models.embedding.unified.UnifiedMultimodalEmbeddingModel` + into a single interleaved sequence. A single ``[CLS]`` token is prepended + and the encoder runs with ``n_modality=1``, so the bottleneck tokens attend + over the full cross-modal timeline. + + Args: + dataset (SampleDataset): dataset providing processed inputs. + embedding_dim (int): shared embedding dimension. + bottlenecks_n (int): number of shared bottleneck tokens. + fusion_startidx (int): layer index at which bottleneck fusion starts. + Must satisfy ``0 <= fusion_startidx <= num_layers``. + num_layers (int): total transformer layers (prefusion + fusion). + heads (int): number of attention heads per transformer block. + dropout (float): dropout rate inside transformer blocks. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, switches to unified mode. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> samples = [ + ... { + ... "patient_id": "patient-0", + ... "visit_id": "visit-0", + ... "conditions": ["A", "B", "C"], + ... "procedures": ["X", "Y"], + ... "label": 1, + ... }, + ... { + ... "patient_id": "patient-1", + ... "visit_id": "visit-0", + ... "conditions": ["D"], + ... "procedures": ["Z", "Y"], + ... "label": 0, + ... }, + ... ] + >>> input_schema = {"conditions": "sequence", "procedures": "sequence"} + >>> output_schema = {"label": "binary"} + >>> dataset = create_sample_dataset( + ... samples, + ... input_schema, + ... output_schema, + ... dataset_name="demo", + ... ) + >>> model = BottleneckTransformer(dataset=dataset, num_layers=3, fusion_startidx=1, bottlenecks_n=4) + >>> loader = get_dataloader(dataset, batch_size=2, shuffle=True) + >>> batch = next(iter(loader)) + >>> output = model(**batch) + >>> sorted(output.keys()) + ['logit', 'loss', 'y_prob', 'y_true'] + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + bottlenecks_n: int = 4, + fusion_startidx: int = 1, + num_layers: int = 3, + heads: int = 4, + dropout: float = 0.5, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, + ): + super().__init__(dataset=dataset) + self.embedding_dim = embedding_dim + self.bottlenecks_n = bottlenecks_n + self.fusion_startidx = fusion_startidx + self.num_layers = num_layers + self.heads = heads + self.dropout = dropout + self._use_unified = unified_embedding is not None + + assert ( + len(self.label_keys) == 1 + ), "Only one label key is supported if BottleneckTransformer is initialized" + self.label_key = self.label_keys[0] + self.mode = self.dataset.output_schema[self.label_key] + + assert 0 <= fusion_startidx <= num_layers, ( + f"fusion_startidx must be in [0, num_layers], got {fusion_startidx}" + ) + + output_size = self.get_output_size() + + if self._use_unified: + self.embedding_model = unified_embedding + # Single CLS token for the unified interleaved sequence + self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02) + self.encoder = MultimodalBottleneckTransformerEncoder( + n_modality=1, + bottlenecks_n=bottlenecks_n, + fusion_startidx=fusion_startidx, + n_layers=num_layers, + n_head=heads, + d_model=embedding_dim, + d_ff=embedding_dim * 4, + dropout=dropout, + ) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.n_modality = len(self.feature_keys) + # Per-modality CLS tokens + self.cls_token_per_modality = nn.ParameterList([ + nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02) + for _ in range(self.n_modality) + ]) + self.encoder = MultimodalBottleneckTransformerEncoder( + n_modality=self.n_modality, + bottlenecks_n=bottlenecks_n, + fusion_startidx=fusion_startidx, + n_layers=num_layers, + n_head=heads, + d_model=embedding_dim, + d_ff=embedding_dim * 4, + dropout=dropout, + ) + + # fc input is embedding_dim in both modes (CLS token, not concat) + self.fc = nn.Linear(embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Extract value/time/mask tensors for UnifiedMultimodalEmbeddingModel.""" + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified( + self, + **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + ) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single time-sorted + sequence, prepends a CLS token, encodes with the bottleneck encoder + (n_modality=1), and classifies from the CLS output. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + event_mask = out["mask"].bool() # (B, S) + + # Prepend CLS token + batch_size = sequence.size(0) + cls = self.cls_token.expand(batch_size, -1, -1) + sequence = torch.cat([cls, sequence], dim=1) + cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=sequence.device) + event_mask = torch.cat([cls_mask, event_mask], dim=1) + + enc_outputs = self.encoder([sequence], [event_mask]) + patient_emb = enc_outputs[0][:, 0, :] # CLS token output + + logits = self.fc(patient_emb) + y_prob = self.prepare_y_prob(logits) + results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob} + if self.label_key in kwargs: + y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + return results + + @staticmethod + def _pool_embedding(x: torch.Tensor) -> torch.Tensor: + if x.dim() == 4: + x = x.sum(dim=2) + if x.dim() == 2: + x = x.unsqueeze(1) + return x + + @staticmethod + def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor: + mask = torch.any(torch.abs(x) > 0, dim=-1) + if mask.dim() == 1: + mask = mask.unsqueeze(1) + invalid_rows = ~mask.any(dim=1) + if invalid_rows.any(): + mask[invalid_rows, 0] = True + return mask.bool() + + def forward( + self, + **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]], + ) -> Dict[str, torch.Tensor]: + """Forward propagation. + + In unified mode dispatches to :meth:`_forward_unified`. Otherwise runs + per-field embedding + bottleneck fusion. + + Args: + **kwargs: keyword arguments for the model. + + Returns: + A dictionary with the following keys: + loss: a scalar tensor representing the final loss. + y_prob: a tensor of predicted probabilities. + y_true: a tensor representing the true labels. + logit: the raw logits before activation. + """ + if self._use_unified: + return self._forward_unified(**kwargs) + + enc_inputs = [] + masks = [] + + for idx, feature_key in enumerate(self.feature_keys): + feature = kwargs[feature_key] + + if isinstance(feature, torch.Tensor): + feature = (feature,) + + schema = self.dataset.input_processors[feature_key].schema() + + value = feature[schema.index("value")] if "value" in schema else None + mask = feature[schema.index("mask")] if "mask" in schema else None + + if len(feature) == len(schema) + 1 and mask is None: + mask = feature[-1] + + if value is None: + raise ValueError( + f"Feature '{feature_key}' must contain 'value' " + f"in the schema." + ) + else: + value = value.to(self.device) + + if mask is not None: + mask = mask.to(self.device) + value = self.embedding_model({feature_key: value}, masks={feature_key: mask})[feature_key] + else: + value = self.embedding_model({feature_key: value})[feature_key] + + value = self._pool_embedding(value) + + if mask is not None: + mask = mask.bool() + if mask.dim() == value.dim(): + mask = mask.any(dim=-1) + else: + mask = self._mask_from_embeddings(value) + + # Prepend Modality CLS token + batch_size = value.size(0) + cls_token = self.cls_token_per_modality[idx].expand(batch_size, -1, -1) + value = torch.cat([cls_token, value], dim=1) + + # Update mask for CLS token (always valid) + cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=value.device) + mask = torch.cat([cls_mask, mask], dim=1) + + enc_inputs.append(value) + masks.append(mask) + + # Pass through Bottleneck Transformer Encoder + enc_outputs = self.encoder(enc_inputs, masks) + + # Extract CLS tokens + cls_tokens = [out[:, 0, :].unsqueeze(1) for out in enc_outputs] + cls_tokens = torch.cat(cls_tokens, dim=1) # [B, n_modality, embedding_dim] + + # Average CLS tokens across valid modalities + b_size = cls_tokens.size(0) + valid_modalities = [] + for mask in masks: + # We check if there's any valid token aside from the CLS token (index 0) + if mask.size(1) > 1: + valid = mask[:, 1:].any(dim=1).float() + else: + valid = mask[:, 0].float() # fallback + valid_modalities.append(valid.view(b_size, 1, 1)) + + valid_modality_tensor = torch.cat(valid_modalities, dim=1) # [B, n_modality, 1] + + # Apply valid mask + masked_cls = cls_tokens * valid_modality_tensor + sum_valid = valid_modality_tensor.sum(dim=1) # [B, 1] + + # Avoid division by zero + sum_valid[sum_valid == 0] = 1.0 + patient_emb = masked_cls.sum(dim=1) / sum_valid # [B, embedding_dim] + + logits = self.fc(patient_emb) + y_prob = self.prepare_y_prob(logits) + + results = { + "logit": logits, + "y_prob": y_prob, + } + + if self.label_key in kwargs: + y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) + loss = self.get_loss_function()(logits, y_true) + results["loss"] = loss + results["y_true"] = y_true + + return results + +if __name__ == "__main__": + from pyhealth.datasets import create_sample_dataset, get_dataloader + + samples = [ + { + "patient_id": "patient-0", + "visit_id": "visit-0", + "conditions": ["A", "B", "C"], + "procedures": ["X", "Y"], + "label": 1, + }, + { + "patient_id": "patient-1", + "visit_id": "visit-0", + "conditions": ["D"], + "procedures": ["Z", "Y"], + "label": 0, + }, + ] + + input_schema = { + "conditions": "sequence", + "procedures": "sequence", + } + output_schema = {"label": "binary"} + + dataset = create_sample_dataset( + samples=samples, + input_schema=input_schema, + output_schema=output_schema, + dataset_name="test", + ) + + train_loader = get_dataloader(dataset, batch_size=2, shuffle=True) + + model = BottleneckTransformer( + dataset=dataset, + embedding_dim=64, + bottlenecks_n=2, + fusion_startidx=1, + num_layers=3, + heads=2 + ) + + data_batch = next(iter(train_loader)) + + result = model(**data_batch) + print(result) + + result["loss"].backward() + print("Test completed successfully.") \ No newline at end of file diff --git a/pyhealth/models/ehrmamba.py b/pyhealth/models/ehrmamba.py index e24c5595f..afdbc8f35 100644 --- a/pyhealth/models/ehrmamba.py +++ b/pyhealth/models/ehrmamba.py @@ -6,6 +6,7 @@ from pyhealth.datasets import SampleDataset from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.models.utils import get_last_visit from pyhealth.processors import ( MultiHotProcessor, @@ -111,6 +112,11 @@ class EHRMamba(BaseModel): Electronic Health Records (arxiv 2405.14567). Uses Mamba (SSM) for linear complexity in sequence length; supports long EHR sequences. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + stack of :class:`MambaBlock` layers rather than one stack per field. + Args: dataset: SampleDataset for token/embedding setup. embedding_dim: Embedding and hidden dimension. Default 128. @@ -118,6 +124,8 @@ class EHRMamba(BaseModel): state_size: SSM state size per channel. Default 16. conv_kernel: Causal conv kernel size in block. Default 4. dropout: Dropout before classification head. Default 0.1. + unified_embedding: Optional pre-built UnifiedMultimodalEmbeddingModel. + When provided, enables unified multi-modal mode. """ def __init__( @@ -128,6 +136,7 @@ def __init__( state_size: int = 16, conv_kernel: int = 4, dropout: float = 0.1, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super().__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -135,19 +144,18 @@ def __init__( self.state_size = state_size self.conv_kernel = conv_kernel self.dropout_rate = dropout + self._use_unified = unified_embedding is not None assert len(self.label_keys) == 1, "EHRMamba supports single label key only" self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - self.feature_processors = { - k: self.dataset.input_processors[k] for k in self.feature_keys - } + output_size = self.get_output_size() + self.dropout = nn.Dropout(dropout) - self.blocks = nn.ModuleDict() - for feature_key in self.feature_keys: - self.blocks[feature_key] = nn.ModuleList( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_blocks = nn.ModuleList( [ MambaBlock( d_model=embedding_dim, @@ -157,10 +165,76 @@ def __init__( for _ in range(num_layers) ] ) - - output_size = self.get_output_size() - self.dropout = nn.Dropout(dropout) - self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.feature_processors = { + k: self.dataset.input_processors[k] for k in self.feature_keys + } + self.blocks = nn.ModuleDict() + for feature_key in self.feature_keys: + self.blocks[feature_key] = nn.ModuleList( + [ + MambaBlock( + d_model=embedding_dim, + state_size=state_size, + conv_kernel=conv_kernel, + ) + for _ in range(num_layers) + ] + ) + self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single + temporally-sorted event sequence, then encodes it with one shared + MambaBlock stack and pools to the last valid event. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + x = out["sequence"] # (B, S_total, E) + mask = out["mask"].bool() # (B, S_total) + + for blk in self._unified_blocks: + x = blk(x) + + last_h = get_last_visit(x, mask) + logits = self.fc(self.dropout(last_h)) + y_prob = self.prepare_y_prob(logits) + results: Dict[str, torch.Tensor] = { + "loss": torch.tensor(0.0), # placeholder, overwritten below + "y_prob": y_prob, + "logit": logits, + } + if self.label_key in kwargs: + y_true = kwargs[self.label_key].to(self.device) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = last_h + return results @staticmethod def _split_temporal(feature: Any) -> Tuple[Optional[torch.Tensor], Any]: @@ -211,6 +285,9 @@ def _pool_embedding(x: torch.Tensor) -> torch.Tensor: return x def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] embedding_inputs: Dict[str, torch.Tensor] = {} masks: Dict[str, torch.Tensor] = {} @@ -261,4 +338,4 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: batch = next(iter(loader)) out = model(**batch) print("keys:", sorted(out.keys())) - out["loss"].backward() + out["loss"].backward() \ No newline at end of file diff --git a/pyhealth/models/embedding/base.py b/pyhealth/models/embedding/base.py new file mode 100644 index 000000000..5ff1ab2a6 --- /dev/null +++ b/pyhealth/models/embedding/base.py @@ -0,0 +1,52 @@ +from abc import ABC, abstractmethod + + +class BaseEmbeddingModel(ABC): + """Abstract base class for all embedding models in PyHealth. + + All embedding models share a common contract: + + - They expose an ``embedding_dim`` property indicating the output vector dimension. + - Their ``forward`` method accepts processor output tensors and returns + vector embeddings. + + Concrete subclasses: + + - :class:`EmbeddingModel` – generic encoder for codes, sequences, timeseries + - :class:`VisionEmbeddingModel` – patch-based encoder for medical images (Josh) + - :class:`TextEmbeddingModel` – BERT-based encoder for clinical text (Rian) + - :class:`UnifiedMultimodalEmbeddingModel` – temporally-aligned multi-modal encoder + """ + + @property + @abstractmethod + def embedding_dim(self) -> int: + """Output embedding dimension shared across all modalities.""" + ... + + @abstractmethod + def forward(self, *args, **kwargs): + """Transform processor outputs into embeddings. + + Subclass return types + --------------------- + EmbeddingModel + ``Dict[str, Tensor]`` mapping each field name to its embedded + tensor. When ``output_mask=True`` is passed, returns a + ``(Dict[str, Tensor], Dict[str, Tensor])`` tuple of + (embeddings, masks). + + VisionEmbeddingModel + ``Tensor`` of shape ``[batch, embedding_dim]``. + + TextEmbeddingModel + ``(Tensor, BoolTensor)`` of shapes ``([B, T, E], [B, T])`` + when ``return_mask=True`` (default), or a plain + ``Tensor [B, T, E]`` when ``return_mask=False``. + + UnifiedMultimodalEmbeddingModel + ``Dict[str, Tensor]`` with keys ``"sequence"`` ``[B, S, E]``, + ``"mask"`` ``[B, S]``, ``"time"`` ``[B, S]``, and + ``"type_ids"`` ``[B, S]``. + """ + ... diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py new file mode 100644 index 000000000..30f83a76d --- /dev/null +++ b/pyhealth/models/embedding/unified.py @@ -0,0 +1,595 @@ +"""UnifiedMultimodalEmbeddingModel, temporally aligned multimodal embedding. + +Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor`` +subclasses ), embeds each event with a modality-specific encoder, then +interleaves all events on a shared timeline by sorting on timestamp and adding +sinusoidal time embeddings + learned modality-type embeddings. + +Output shape: ``(B, S_total, E')``, a single sequence of events usable by +any downstream sequence model (Transformer, Mamba, RNN, …). + +IMAGE encoding delegates to :class:`PatchEmbedding` from +:mod:`pyhealth.models.embedding.vision` (Josh's model), pooling patch tokens +to a single per-image vector via global mean pooling. + +TEXT encoding uses a pretrained BERT tokenizer model directly, extracting the +[CLS] token per note, the same BERT-based approach as +:class:`TextEmbeddingModel` (Rian's model). + +Unimodal model reuse via ``field_embeddings``:: + + vision_model = VisionEmbeddingModel(dataset, embedding_dim=128) + text_model = TextEmbeddingModel(embedding_dim=128) + unified = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + field_embeddings={ + "chest_xray": vision_model, # reuses trained backbone + "notes": text_model, # reuses BERT + projection + }, + ) + +Quickstart:: + + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.datasets.collate import collate_temporal + model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128) + # inside forward: + # inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...} + out = model(inputs) + # out["sequence"]: (B, S_total, 128) + # out["mask"]: (B, S_total) , 1 = real event, 0 = padding + # out["time"]: (B, S_total) , hours from first event +""" + +from __future__ import annotations + +import math +import warnings +from typing import Any, Optional + +import torch +import torch.nn as nn + +from ...processors.base_processor import ModalityType, TemporalFeatureProcessor +from .base import BaseEmbeddingModel +from .vision import PatchEmbedding + + +# ── Helpers ─────────────────────────────────────────────────────────────────── + + +class SinusoidalTimeEmbedding(nn.Module): + """Continuous sinusoidal embedding for scalar time values (in hours). + + Identical in spirit to the positional encoding in "Attention is All You + Need" but operating on real-valued timestamps rather than integer positions. + + Args: + dim: Output embedding dimension (must be even). + max_hours: Maximum expected time value in hours. Values are normalised + to ``[0, 2π]`` before the sin/cos projection. Default 720 (30 days). + + Shape: + Input: ``(*, )`` float tensor of times in hours + Output: ``(*, dim)`` + """ + + def __init__(self, dim: int, max_hours: float = 720.0): + super().__init__() + assert dim % 2 == 0, f"dim must be even, got {dim}" + self.dim = dim + self.max_hours = max_hours + half = dim // 2 + freqs = torch.exp( + -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1) + ) + self.register_buffer("freqs", freqs) # (dim//2,) + + def forward(self, t: torch.Tensor) -> torch.Tensor: + """:param t: ``(...,)`` float, times in hours.""" + t_norm = t / self.max_hours * 2 * math.pi # (...,) + args = t_norm.unsqueeze(-1) * self.freqs # (..., dim//2) + return torch.cat([args.sin(), args.cos()], dim=-1) # (..., dim) + + +class _MeanPool(nn.Module): + """Pool a sequence of patch embeddings to a single vector via global mean.""" + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # (B, num_patches, E) -> (B, E) + return x.mean(dim=1) + + +# ── Main model ─────────────────────────────────────────────────────────────── + + +class UnifiedMultimodalEmbeddingModel(nn.Module, BaseEmbeddingModel): + """Embed heterogeneous temporal features into a single aligned sequence. + + **All** input processors must be ``TemporalFeatureProcessor`` subclasses. + Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``) + are rejected with a clear error, use :class:`EmbeddingModel` for those fields. + + Modality routing: + + - **CODE**: ``nn.Embedding`` lookup. + - **TEXT**: Pretrained BERT (same approach as :class:`TextEmbeddingModel`), + CLS token extracted per note. + - **IMAGE**: :class:`PatchEmbedding` (from :class:`VisionEmbeddingModel`) + followed by global mean pooling to produce one vector per image event. + - **NUMERIC / SIGNAL**: ``nn.Linear`` projection. + + Unimodal model reuse: + + Pass pre-built :class:`EmbeddingModel`, :class:`VisionEmbeddingModel`, or + :class:`TextEmbeddingModel` instances via ``field_embeddings`` to reuse + their trained encoder weights instead of building new ones from scratch. + The core encoder module is extracted from each pre-built model: + + - ``EmbeddingModel`` → ``embedding_layers[field_name]`` (``nn.Embedding`` / + ``nn.Linear``) + - ``VisionEmbeddingModel`` → ``embedding_layers[field_name]`` backbone + + global mean pooling + - ``TextEmbeddingModel`` → ``transformer`` (BERT) + ``fc`` (projection) + + Algorithm + --------- + For each temporal field: + + 1. Route ``inputs[field]["value"]`` through a modality-specific encoder → + ``(B, N_i, E')`` per-event embeddings. + 2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours). + 3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or + ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if + token-level. + + Then: + + 4. Concatenate across all fields → ``(B, S_total, E')``. + 5. Sort events along dim=1 by timestamp (ascending). + 6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``. + 7. Return ``{"sequence", "time", "mask", "type_ids"}``. + + Args: + processors: ``dict[field_name, TemporalFeatureProcessor]``, the + processors for each temporal field in the dataset. Pass + ``dataset.input_processors`` directly. + embedding_dim: Shared embedding dimension ``E'``. + time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. + max_time_hours: Normalisation constant for the time embedding. + Defaults to 720 h (30 days). + image_size: Image size (H=W) assumed for IMAGE fields when using + PatchEmbedding. Defaults to 224. + image_channels: Number of input channels for IMAGE fields. Defaults to 3. + patch_size: Patch size for IMAGE PatchEmbedding encoder. Defaults to 16. + image_pool: Pooling strategy applied to IMAGE patch tokens to produce + one vector per image event. Only ``"mean"`` (global mean pooling) + is currently implemented. Defaults to ``"mean"``. + field_embeddings: Optional mapping of field names to pre-built unimodal + embedding models. Supported types: + + - :class:`EmbeddingModel` (codes / numeric) — extracts + ``embedding_layers[field_name]``. + - :class:`VisionEmbeddingModel` — extracts the backbone layer and + wraps it with global mean pooling. + - :class:`TextEmbeddingModel` — reuses ``transformer`` and ``fc`` + for BERT-based CLS extraction. + + Fields not present in this dict fall back to the default + internally-built encoders. + + Example:: + + model = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + ) + # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}} + out = model(inputs) + seq = out["sequence"] # (B, S_total, 128) + mask = out["mask"] # (B, S_total) float, 1=valid 0=pad + + # With pre-built unimodal models: + vision = VisionEmbeddingModel(dataset, embedding_dim=128) + model = UnifiedMultimodalEmbeddingModel( + processors=dataset.input_processors, + embedding_dim=128, + field_embeddings={"chest_xray": vision}, + ) + """ + + def __init__( + self, + processors: dict[str, Any], + embedding_dim: int = 128, + time_embedding: str = "sinusoidal", + max_time_hours: float = 720.0, + image_size: int = 224, + image_channels: int = 3, + patch_size: int = 16, + image_pool: str = "mean", + field_embeddings: Optional[dict[str, Any]] = None, + freeze_text_encoder: bool = False, + ): + super().__init__() + if image_pool != "mean": + raise NotImplementedError( + f"Only image_pool='mean' is implemented, got {image_pool!r}." + ) + self._embedding_dim = embedding_dim + self._freeze_text_encoder = freeze_text_encoder + self.image_pool = image_pool + _field_embeddings = field_embeddings or {} + + self.encoders: nn.ModuleDict = nn.ModuleDict() + self.projections: nn.ModuleDict = nn.ModuleDict() + self.modality_types: dict[str, ModalityType] = {} + self._shared_text_field_by_model: dict[str, str] = {} + self._text_canonical: dict[str, str] = {} # field → first field sharing the same tokenizer + + for field_name, processor in processors.items(): + if not isinstance(processor, TemporalFeatureProcessor): + raise TypeError( + f"UnifiedMultimodalEmbeddingModel requires every input processor " + f"to be a TemporalFeatureProcessor subclass, but '{field_name}' " + f"uses {type(processor).__name__}. For non-temporal fields use " + f"EmbeddingModel." + ) + + m = processor.modality() + self.modality_types[field_name] = m + pre_built = _field_embeddings.get(field_name) + + if m == ModalityType.CODE: + self.encoders[field_name] = self._build_code_encoder( + field_name, processor, pre_built, embedding_dim + ) + + elif m == ModalityType.TEXT: + self._build_text_encoder( + field_name, processor, pre_built, embedding_dim, + freeze=freeze_text_encoder, + ) + + elif m == ModalityType.IMAGE: + self.encoders[field_name] = self._build_image_encoder( + field_name, + processor, + pre_built, + embedding_dim, + image_size, + image_channels, + patch_size, + image_pool, + ) + + elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL): + self.encoders[field_name] = self._build_numeric_encoder( + field_name, processor, pre_built, embedding_dim + ) + + else: + raise NotImplementedError( + f"No encoder implemented for modality {m!r} (field '{field_name}')." + ) + + # Shared type embedding, one vector per unique modality in this dataset + unique_modalities = sorted(set(self.modality_types.values())) + self._modality_to_idx: dict[ModalityType, int] = { + mod: i for i, mod in enumerate(unique_modalities) + } + self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim) + self._warned_nested_code_flatten = False + + # Time embedding + if time_embedding == "sinusoidal": + self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours) + else: + raise NotImplementedError( + "Only 'sinusoidal' time embedding is implemented." + ) + + # ── Encoder builders ────────────────────────────────────────────────────── + + def _build_code_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + ) -> nn.Module: + """Build CODE encoder: nn.Embedding, optionally from a pre-built EmbeddingModel.""" + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + layer = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim)) + return layer + + vocab_size = processor.value_dim() + return nn.Embedding(vocab_size, embedding_dim, padding_idx=0) + + def _build_text_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + freeze: bool = False, + ) -> None: + """Build TEXT encoder: BERT + projection, optionally from TextEmbeddingModel.""" + + def _set_projection( + pre_dim: int, proj_source: Optional[nn.Module] = None + ) -> None: + if pre_dim != embedding_dim: + if proj_source is not None: + self.projections[field_name] = nn.Sequential( + proj_source, + nn.Linear(pre_dim, embedding_dim), + ) + else: + self.projections[field_name] = nn.Linear(pre_dim, embedding_dim) + elif proj_source is not None: + self.projections[field_name] = proj_source + + if ( + pre_built is not None + and hasattr(pre_built, "transformer") + and hasattr(pre_built, "fc") + ): + self.encoders[field_name] = pre_built.transformer + if freeze: + for p in pre_built.transformer.parameters(): + p.requires_grad = False + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + _set_projection(pre_dim, pre_built.fc) + return + + if processor.is_token(): + from transformers import AutoModel + + bert = AutoModel.from_pretrained(processor.tokenizer_model) + if freeze: + for p in bert.parameters(): + p.requires_grad = False + self.encoders[field_name] = bert + hidden = bert.config.hidden_size + if hidden != embedding_dim: + self.projections[field_name] = nn.Linear(hidden, embedding_dim) + else: + raise ValueError( + f"TEXT processor '{field_name}' must either supply a pre-built " + f"TextEmbeddingModel via field_embeddings or use a tokenizer " + f"(set tokenizer_model=...) to be used with " + f"UnifiedMultimodalEmbeddingModel." + ) + + def _build_image_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + image_size: int, + image_channels: int, + patch_size: int, + image_pool: str, + ) -> nn.Module: + """Build IMAGE encoder: backbone + pool, optionally from VisionEmbeddingModel.""" + pool_layers: dict[str, nn.Module] = {"mean": _MeanPool()} + pool_layer = pool_layers[image_pool] + + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + backbone = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential( + backbone, pool_layer, nn.Linear(pre_dim, embedding_dim) + ) + return nn.Sequential(backbone, pool_layer) + + _image_size = getattr(processor, "image_size", image_size) + _in_channels = getattr(processor, "in_channels", image_channels) + return nn.Sequential( + PatchEmbedding(_image_size, patch_size, _in_channels, embedding_dim), + pool_layer, + ) + + def _build_numeric_encoder( + self, + field_name: str, + processor: TemporalFeatureProcessor, + pre_built: Any, + embedding_dim: int, + ) -> nn.Module: + """Build NUMERIC/SIGNAL encoder: nn.Linear, optionally from EmbeddingModel.""" + if ( + pre_built is not None + and hasattr(pre_built, "embedding_layers") + and field_name in pre_built.embedding_layers + ): + layer = pre_built.embedding_layers[field_name] + pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) + if pre_dim != embedding_dim: + return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim)) + return layer + + in_features = processor.value_dim() + return nn.Linear(in_features, embedding_dim) + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + # ── Forward ─────────────────────────────────────────────────────────────── + + def forward( + self, + inputs: dict[str, dict[str, torch.Tensor]], + ) -> dict[str, torch.Tensor]: + """Encode and temporally align all temporal features. + + Args: + inputs: ``{field_name: {"value": Tensor, "time": Tensor, + "mask": Tensor (optional)}}`` + , one dict per temporal feature, exactly as produced by + ``collate_temporal``. + + Returns: + A dict with keys: + + * ``"sequence"``, ``(B, S_total, E')`` temporally-sorted events + (content + time + type embeddings) + * ``"time"`` , ``(B, S_total)`` timestamps (hours) + * ``"mask"`` , ``(B, S_total)`` 1=real event, 0=padding + * ``"type_ids"``, ``(B, S_total)`` modality index per event + * ``"token_emb"``, ``(B, S_total, E')`` content-only event embedding + (before time/type are added); the target for masked modeling. + """ + all_embeddings: list[torch.Tensor] = [] + all_times: list[torch.Tensor] = [] + all_masks: list[torch.Tensor] = [] + all_types: list[torch.Tensor] = [] + + for field_name, feat_dict in inputs.items(): + value = feat_dict["value"] # (B, N_i, ...) or (B, S, F) + time = feat_dict["time"] # (B, N_i) + mask = feat_dict.get("mask") + + if time is None: + # Fallback: treat every event as occurring at t=0 + time = torch.zeros(value.shape[:2], device=value.device) + + modality = self.modality_types[field_name] + encoder_key = self._text_canonical.get(field_name, field_name) + encoder = self.encoders[encoder_key] + + # ── Encode ──────────────────────────────────────────────────── + if modality == ModalityType.CODE: + # CODE values may be either: + # - flat indices: (B, S) + # - nested indices: (B, S, C) where C is codes-per-event + # For nested indices, flatten to (B, S*C, E') so code-level + # detail is preserved, and expand time/mask to match. + if value.dim() == 2: + emb = encoder(value) # (B, S, E') + elif value.dim() == 3: + bsz, seq_len, per_event_codes = value.shape + token_emb = encoder(value.long()) # (B, S, C, E') + emb = token_emb.reshape(bsz, seq_len * per_event_codes, -1) + + if not self._warned_nested_code_flatten: + warnings.warn( + ( + "UnifiedMultimodalEmbeddingModel detected " + f"nested CODE input for '{field_name}' with " + f"shape={tuple(value.shape)}. Flattening to " + f"(B, S*C, E) and repeating time along C." + ), + stacklevel=2, + ) + self._warned_nested_code_flatten = True + + if time is not None: + time = ( + time.unsqueeze(-1) + .expand(-1, -1, per_event_codes) + .reshape(bsz, seq_len * per_event_codes) + ) + + if mask is not None: + if mask.dim() == 2: + mask = ( + mask.unsqueeze(-1) + .expand(-1, -1, per_event_codes) + .reshape(bsz, seq_len * per_event_codes) + ) + elif mask.dim() == 3: + mask = mask.reshape(bsz, seq_len * per_event_codes) + else: + raise ValueError( + f"Unsupported CODE value rank for '{field_name}': " + f"shape={tuple(value.shape)}" + ) + + elif modality == ModalityType.TEXT: + b, n, l = value.shape + flat_ids = value.view(b * n, l) + flat_mask = mask.view(b * n, l) if mask is not None else None + out = encoder(input_ids=flat_ids, attention_mask=flat_mask) + cls_emb = out.last_hidden_state[:, 0, :] # (B*N, H) + if field_name in self.projections: + cls_emb = self.projections[field_name](cls_emb) + emb = cls_emb.view(b, n, -1) # (B, N, E') + + elif modality == ModalityType.IMAGE: + # encoder = Sequential(PatchEmbedding, _MeanPool) → (B*N, E') + b, n, c, h, w = value.shape + flat_imgs = value.view(b * n, c, h, w) + img_emb = encoder(flat_imgs) # (B*N, E') + emb = img_emb.view(b, n, -1) # (B, N, E') + + else: # NUMERIC / SIGNAL + emb = encoder(value) # (B, T, E') + + # ── Build event-level validity mask ─────────────────────────── + if mask is None: + event_mask = torch.ones(emb.shape[:2], device=emb.device) + else: + if mask.dim() > time.dim(): + # token-level (B, N, L) → event-level (B, N) + event_mask = (mask.sum(dim=-1) > 0).float() + else: + event_mask = mask.float() + + # ── Modality type indices ───────────────────────────────────── + type_idx = self._modality_to_idx[modality] + type_ids = torch.full( + emb.shape[:2], type_idx, dtype=torch.long, device=emb.device + ) + + all_embeddings.append(emb) + all_times.append(time) + all_masks.append(event_mask) + all_types.append(type_ids) + + # ── Concatenate across all fields ───────────────────────────────── + cat_emb = torch.cat(all_embeddings, dim=1) # (B, S_total, E') + cat_time = torch.cat(all_times, dim=1) # (B, S_total) + cat_mask = torch.cat(all_masks, dim=1) # (B, S_total) + cat_types = torch.cat(all_types, dim=1) # (B, S_total) + + # ── Sort by time ────────────────────────────────────────────────── + sort_idx = cat_time.argsort(dim=1) + cat_emb = cat_emb.gather(1, sort_idx.unsqueeze(-1).expand_as(cat_emb)) + cat_time = cat_time.gather(1, sort_idx) + cat_mask = cat_mask.gather(1, sort_idx) + cat_types = cat_types.gather(1, sort_idx) + + # ── Add time + type embeddings ──────────────────────────────────── + time_emb = self.time_embed(cat_time) # (B, S_total, E') + type_emb = self.type_embedding(cat_types) # (B, S_total, E') + final = cat_emb + time_emb + type_emb # (B, S_total, E') + + return { + "sequence": final, # (B, S_total, E') + "time": cat_time, # (B, S_total) + "mask": cat_mask, # (B, S_total) + "type_ids": cat_types, # (B, S_total) + # Per-event content embedding BEFORE time/type are added (same sort + # order as ``sequence``). Masked-modeling pretrainers should + # reconstruct THIS rather than ``sequence``: the time/type + # components are largely recoverable from event position, so + # including them in the target dilutes the content signal. + "token_emb": cat_emb, # (B, S_total, E') + } \ No newline at end of file diff --git a/pyhealth/models/embedding/vanilla.py b/pyhealth/models/embedding/vanilla.py new file mode 100644 index 000000000..9b684d0c0 --- /dev/null +++ b/pyhealth/models/embedding/vanilla.py @@ -0,0 +1,366 @@ +from __future__ import annotations + +from typing import Dict, Any, Optional, Union +import os + +import torch +import torch.nn as nn + +from ...datasets import SampleDataset +from ...processors import ( + MultiHotProcessor, + NestedFloatsProcessor, + NestedSequenceProcessor, + SequenceProcessor, + StageNetProcessor, + StageNetTensorProcessor, + TensorProcessor, + TimeseriesProcessor, + DeepNestedSequenceProcessor, + DeepNestedFloatsProcessor, +) +from ..base_model import BaseModel +from .base import BaseEmbeddingModel + + +def _iter_text_vectors( + path: str, + embedding_dim: int, + wanted_tokens: set[str], + encoding: str = "utf-8", +) -> Dict[str, torch.Tensor]: + """Loads word vectors from a text file (e.g., GloVe) for a subset of tokens. + + Expected format: one token per line followed by embedding_dim floats. + + This function reads the file line-by-line and only retains vectors for + tokens present in `wanted_tokens`. + """ + + if not os.path.exists(path): + raise FileNotFoundError(f"pretrained embedding file not found: {path}") + + vectors: Dict[str, torch.Tensor] = {} + with open(path, "r", encoding=encoding) as f: + for line in f: + line = line.strip() + if not line: + continue + parts = line.split() + # token + embedding_dim values + if len(parts) < embedding_dim + 1: + continue + token = parts[0] + if token not in wanted_tokens: + continue + try: + vec = torch.tensor( + [float(x) for x in parts[1 : embedding_dim + 1]], + dtype=torch.float, + ) + except ValueError: + continue + vectors[token] = vec + return vectors + + +def init_embedding_with_pretrained( + embedding: nn.Embedding, + code_vocab: Dict[Any, int], + pretrained_path: str, + embedding_dim: int, + pad_token: str = "", + unk_token: str = "", + normalize: bool = False, + freeze: bool = False, +) -> int: + """Initializes an nn.Embedding from a pretrained text-vector file. + + Tokens not found in the pretrained file are left as the module's existing + random initialization. + + Returns: + int: number of tokens successfully initialized from the file. + """ + + # Build wanted token set (stringified) + vocab_tokens = {str(t) for t in code_vocab.keys()} + vectors = _iter_text_vectors(pretrained_path, embedding_dim, vocab_tokens) + + loaded = 0 + with torch.no_grad(): + for tok, idx in code_vocab.items(): + tok_s = str(tok) + if tok_s in vectors: + vec = vectors[tok_s] + if normalize: + vec = vec / (vec.norm(p=2) + 1e-12) + embedding.weight[idx].copy_(vec) + loaded += 1 + + # Ensure pad row is zero + if pad_token in code_vocab: + embedding.weight[code_vocab[pad_token]].zero_() + # If embedding has a padding_idx, keep it consistent + if embedding.padding_idx is not None: + embedding.weight[embedding.padding_idx].zero_() + + if freeze: + embedding.weight.requires_grad_(False) + + return loaded + + +class EmbeddingModel(BaseModel): + """ + EmbeddingModel is responsible for creating embedding layers for different types of input data. + + This model automatically creates appropriate embedding transformations based on the processor type: + + - SequenceProcessor: nn.Embedding + Input: (batch, seq_len) + Output: (batch, seq_len, embedding_dim) + + - NestedSequenceProcessor: nn.Embedding + Input: (batch, num_visits, max_codes_per_visit) + Output: (batch, num_visits, max_codes_per_visit, embedding_dim) + + - DeepNestedSequenceProcessor: nn.Embedding + Input: (batch, num_groups, num_visits, max_codes_per_visit) + Output: (batch, num_groups, num_visits, max_codes_per_visit, embedding_dim) + + - TimeseriesProcessor / NestedFloatsProcessor / DeepNestedFloatsProcessor / StageNetTensorProcessor: + nn.Linear over the last dimension + Input: (..., size) + Output: (..., embedding_dim) + + - TensorProcessor: nn.Linear (size inferred from first sample) + + - MultiHotProcessor: nn.Linear over multi-hot vector + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + pretrained_emb_path: Optional[Union[str, Dict[str, str]]] = None, + freeze_pretrained: bool = False, + normalize_pretrained: bool = False, + ): + super().__init__(dataset) + # BaseEmbeddingModel declares `embedding_dim` as an abstract property, + # so we can't set self.embedding_dim directly (no setter). Use a + # private backing attribute and expose it through the property below. + self._embedding_dim = embedding_dim + self.embedding_layers = nn.ModuleDict() + + for field_name, processor in self.dataset.input_processors.items(): + # Deep categorical: use special module that collapses last dim to embedding_dim + + # Regular categorical sequences -> nn.Embedding (adds embedding dim) + if isinstance( + processor, + ( + SequenceProcessor, + StageNetProcessor, + NestedSequenceProcessor, + DeepNestedSequenceProcessor, + ), + ): + vocab_size = len(processor.code_vocab) + + # For NestedSequenceProcessor and DeepNestedSequenceProcessor, don't use padding_idx + # because empty visits/groups need non-zero embeddings. + if isinstance( + processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor) + ): + self.embedding_layers[field_name] = nn.Embedding( + num_embeddings=vocab_size, + embedding_dim=embedding_dim, + padding_idx=None, + ) + else: + self.embedding_layers[field_name] = nn.Embedding( + num_embeddings=vocab_size, + embedding_dim=embedding_dim, + padding_idx=0, + ) + + # Optional pretrained initialization (e.g., GloVe). + if pretrained_emb_path is not None: + if isinstance(pretrained_emb_path, str): + path = pretrained_emb_path + else: + path = pretrained_emb_path.get(field_name) + if path: + init_embedding_with_pretrained( + self.embedding_layers[field_name], + processor.code_vocab, + path, + embedding_dim=embedding_dim, + normalize=normalize_pretrained, + freeze=freeze_pretrained, + ) + + # Numeric features (including deep nested floats) -> nn.Linear over last dim + elif isinstance( + processor, + ( + TimeseriesProcessor, + StageNetTensorProcessor, + NestedFloatsProcessor, + DeepNestedFloatsProcessor, + ), + ): + # Assuming processor.size() returns the last-dim size + in_features = processor.size() + self.embedding_layers[field_name] = nn.Linear( + in_features=in_features, out_features=embedding_dim + ) + + elif isinstance(processor, TensorProcessor): + # Infer size from first sample + sample_tensor = None + for sample in dataset: + if field_name in sample: + sample_tensor = processor.process(sample[field_name]) + break + if sample_tensor is not None: + input_size = ( + sample_tensor.shape[-1] if sample_tensor.dim() > 0 else 1 + ) + self.embedding_layers[field_name] = nn.Linear( + in_features=input_size, out_features=embedding_dim + ) + + elif isinstance(processor, MultiHotProcessor): + num_categories = processor.size() + self.embedding_layers[field_name] = nn.Linear( + in_features=num_categories, out_features=embedding_dim + ) + + # Smart Processor (Token-based) -> Transformers + elif hasattr(processor, "is_token") and processor.is_token(): + try: + from transformers import AutoModel + except ImportError: + raise ImportError( + "Please install `transformers` to use token-based processors." + ) + + # Load the model + self.embedding_layers[field_name] = AutoModel.from_pretrained( + processor.tokenizer_model + ) + + # Check if we need projection + if ( + self.embedding_layers[field_name].config.hidden_size + != self.embedding_dim + ): + self.embedding_layers[f"{field_name}_proj"] = nn.Linear( + self.embedding_layers[field_name].config.hidden_size, + self.embedding_dim, + ) + + else: + print( + "Warning: No embedding created for field due to lack of compatible processor:", + field_name, + ) + + def forward( + self, + inputs: Dict[str, torch.Tensor], + masks: Dict[str, torch.Tensor] = None, + output_mask: bool = False, + ) -> ( + Dict[str, torch.Tensor] + | tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]] + ): + + embedded: Dict[str, torch.Tensor] = {} + out_masks: Dict[str, torch.Tensor] = {} if output_mask else None + + for field_name, tensor in inputs.items(): + processor = self.dataset.input_processors.get(field_name, None) + + if field_name not in self.embedding_layers: + # No embedding layer -> passthrough + embedded[field_name] = tensor + continue + + # Check if it's a transformer model + layer = self.embedding_layers[field_name] + + # Check for transformers.PreTrainedModel (but without importing if possible, use class name check) + # or check if it has 'config' attribute + if hasattr(layer, "config") and hasattr(layer, "forward"): + # It's likely a transformer + tensor = tensor.to(self.device).long() # Ensure LongTensor for IDs + + mask = None + if masks is not None and field_name in masks: + mask = masks[field_name].to(self.device) + + # Handle 3D input (Batch, Num_Notes, Seq_Len) + is_3d = inputs[field_name].dim() == 3 + + if is_3d: + b, n, l = inputs[field_name].shape + tensor = tensor.view(b * n, l) + if mask is not None: + mask = mask.view(b * n, l) + + # Forward pass through transformer + output = layer(input_ids=tensor, attention_mask=mask) + x = output.last_hidden_state # (Batch, Seq, Hidden) + + if is_3d: + # If we had 3D input, we MUST pool the sequence dim (L) to get one vector per note + # Resulting shape: (B, N, H) + + # Pool L dim -> (B*N, H) using CLS token (index 0) + x = x[:, 0, :] + + # Check projections + if f"{field_name}_proj" in self.embedding_layers: + x = self.embedding_layers[f"{field_name}_proj"](x) + + x = x.view(b, n, -1) + + else: + # 2D input (Batch, Seq) -> (Batch, Seq, Hidden) + # No pooling, treating as sequence of tokens (word embeddings) + if f"{field_name}_proj" in self.embedding_layers: + x = self.embedding_layers[f"{field_name}_proj"](x) + + embedded[field_name] = x + + else: + # Standard layers + tensor = tensor.to(self.device) + embedded[field_name] = layer(tensor) + + if output_mask: + # Generate a mask for this field + if masks is not None and field_name in masks: + out_masks[field_name] = masks[field_name].to(self.device) + elif hasattr(processor, "code_vocab"): + pad_idx = processor.code_vocab.get("", 0) + out_masks[field_name] = tensor != pad_idx + else: + # Default mask generation (e.g. for simple linear layers where 0 might be padding?) + # Be careful changing this behavior. + # Previous code: + # masks[field_name] = (tensor != pad_idx) -> where pad_idx was 0 default + pad_idx = 0 + out_masks[field_name] = tensor != pad_idx + + if output_mask: + return embedded, out_masks + else: + return embedded + + def __repr__(self) -> str: + return f"EmbeddingModel(embedding_layers={self.embedding_layers})" \ No newline at end of file diff --git a/pyhealth/models/embedding/vision.py b/pyhealth/models/embedding/vision.py new file mode 100644 index 000000000..57eedda22 --- /dev/null +++ b/pyhealth/models/embedding/vision.py @@ -0,0 +1,384 @@ +# Author: Josh Steier +# Description: Vision embedding model for medical imaging + +from typing import Any, Dict, Literal, Optional, Tuple, Union + +import torch +import torch.nn as nn +import shutil +from ...datasets import SampleDataset +from ..base_model import BaseModel +from ...processors import ImageProcessor +from .base import BaseEmbeddingModel + + +class Permute(nn.Module): + """Utility module to permute tensor dimensions in nn.Sequential. + + Args: + dims: Variable number of integers specifying the desired ordering + of dimensions. + + Example: + >>> permute = Permute(0, 2, 1) + >>> x = torch.randn(32, 256, 49) # (B, E, spatial) + >>> out = permute(x) # (B, spatial, E) + """ + + def __init__(self, *dims: int) -> None: + super().__init__() + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x.permute(*self.dims) + + +class PatchEmbedding(nn.Module): + """Convert images to patch embeddings using ViT-style projection. + + Splits an image into non-overlapping patches and projects each patch + to an embedding vector using a convolutional layer. + + Args: + image_size: Input image size (assumes square images). + patch_size: Size of each square patch. + in_channels: Number of input channels. + embedding_dim: Output embedding dimension for each patch. + + Example: + >>> patch_embed = PatchEmbedding(224, 16, 3, 256) + >>> x = torch.randn(4, 3, 224, 224) + >>> patches = patch_embed(x) # (4, 196, 256) + """ + + def __init__( + self, + image_size: int = 224, + patch_size: int = 16, + in_channels: int = 3, + embedding_dim: int = 128, + ) -> None: + super().__init__() + if image_size % patch_size != 0: + raise ValueError( + f"image_size ({image_size}) must be divisible by " + f"patch_size ({patch_size})" + ) + self.patch_size = patch_size + self.num_patches = (image_size // patch_size) ** 2 + self.proj = nn.Conv2d( + in_channels, embedding_dim, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # (B, C, H, W) -> (B, E, H/P, W/P) -> (B, num_patches, E) + x = self.proj(x) + x = x.flatten(2).transpose(1, 2) + return x + + +class VisionEmbeddingModel(BaseModel, BaseEmbeddingModel): + """Vision embedding model for medical image inputs. + + Converts medical images to sequences of patch embeddings suitable for + attention-based fusion with other modalities (EHR, text). + + Supports multiple backbone types: + - "patch": ViT-style patch projection (lightweight) + - "cnn": Small CNN encoder (good inductive bias) + - "resnet18"/"resnet50": Pretrained backbones + + Output shape: (batch, num_patches, embedding_dim) + + Args: + dataset: SampleDataset with ImageProcessor fields. + embedding_dim: Output embedding dimension. Default 128. + patch_size: Patch size for "patch" backbone. Default 16. + backbone: One of "patch", "cnn", "resnet18", "resnet50". + pretrained: Use ImageNet weights for ResNet. Default True. + freeze_backbone: Freeze pretrained weights. Default False. + dropout: Dropout rate. Default 0.0. + use_cls_token: Prepend learnable [CLS] token. Default False. + + Example: + >>> from pyhealth.datasets import create_sample_dataset + >>> model = VisionEmbeddingModel(dataset, embedding_dim=256) + >>> embeddings = model({"chest_xray": images}) + """ + + def __init__( + self, + dataset: SampleDataset, + embedding_dim: int = 128, + patch_size: int = 16, + backbone: Literal["patch", "cnn", "resnet18", "resnet50"] = "patch", + pretrained: bool = True, + freeze_backbone: bool = False, + dropout: float = 0.0, + use_cls_token: bool = False, + pool: Optional[Literal["mean"]] = None, + ) -> None: + super().__init__(dataset) + + self._embedding_dim = embedding_dim + self.patch_size = patch_size + self.pool = pool + self.backbone_type = backbone + self.use_cls_token = use_cls_token + + self.embedding_layers = nn.ModuleDict() + self.pos_embeddings = nn.ParameterDict() + self.cls_tokens = nn.ParameterDict() if use_cls_token else None + self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity() + + self._field_info: Dict[str, Dict[str, Any]] = {} + + for field_name, processor in self.dataset.input_processors.items(): + if not isinstance(processor, ImageProcessor): + continue + + image_size = processor.image_size + in_channels = self._infer_channels(processor) + + num_patches = self._build_embedding_layer( + field_name, image_size, in_channels, backbone, pretrained, freeze_backbone + ) + + num_positions = num_patches + 1 if use_cls_token else num_patches + self.pos_embeddings[field_name] = nn.Parameter( + torch.randn(1, num_positions, embedding_dim) * 0.02 + ) + + if use_cls_token: + self.cls_tokens[field_name] = nn.Parameter( + torch.randn(1, 1, embedding_dim) * 0.02 + ) + + self._field_info[field_name] = { + "num_patches": num_patches, + "image_size": image_size, + "in_channels": in_channels, + } + + @property + def embedding_dim(self) -> int: + return self._embedding_dim + + def _infer_channels(self, processor: ImageProcessor) -> int: + """Infer number of input channels from processor mode.""" + mode = getattr(processor, "mode", None) + if mode == "L": + return 1 + elif mode == "RGBA": + return 4 + return 3 + + def _build_embedding_layer( + self, + field_name: str, + image_size: int, + in_channels: int, + backbone: str, + pretrained: bool, + freeze_backbone: bool, + ) -> int: + """Build embedding layer and return number of output patches.""" + if backbone == "patch": + num_patches = (image_size // self.patch_size) ** 2 + self.embedding_layers[field_name] = PatchEmbedding( + image_size, self.patch_size, in_channels, self._embedding_dim + ) + + elif backbone == "cnn": + num_patches = 7 * 7 + self.embedding_layers[field_name] = nn.Sequential( + nn.Conv2d(in_channels, 64, 7, stride=2, padding=3), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, 128, 3, stride=2, padding=1), + nn.BatchNorm2d(128), + nn.ReLU(inplace=True), + nn.Conv2d(128, self._embedding_dim, 3, stride=2, padding=1), + nn.BatchNorm2d(self._embedding_dim), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d((7, 7)), + nn.Flatten(2), + Permute(0, 2, 1), + ) + + elif backbone in ("resnet18", "resnet50"): + num_patches = 7 * 7 + self.embedding_layers[field_name] = self._build_resnet_backbone( + backbone, in_channels, pretrained, freeze_backbone + ) + + else: + raise ValueError(f"Unknown backbone: {backbone}") + + return num_patches + + def _build_resnet_backbone( + self, backbone: str, in_channels: int, pretrained: bool, freeze: bool + ) -> nn.Module: + """Build pretrained ResNet backbone with spatial output.""" + try: + import torchvision.models as models + except ImportError as e: + raise ImportError("torchvision required for ResNet backbones") from e + + if backbone == "resnet18": + weights = models.ResNet18_Weights.DEFAULT if pretrained else None + resnet = models.resnet18(weights=weights) + feature_dim = 512 + else: + weights = models.ResNet50_Weights.DEFAULT if pretrained else None + resnet = models.resnet50(weights=weights) + feature_dim = 2048 + + if in_channels != 3: + resnet.conv1 = nn.Conv2d( + in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False + ) + + layers = list(resnet.children())[:-2] + backbone_net = nn.Sequential(*layers) + + if freeze: + for param in backbone_net.parameters(): + param.requires_grad = False + + return nn.Sequential( + backbone_net, + nn.Conv2d(feature_dim, self._embedding_dim, kernel_size=1), + nn.Flatten(2), + Permute(0, 2, 1), + ) + + def forward( + self, + inputs: Dict[str, torch.Tensor], + output_mask: bool = False, + ) -> Union[Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]]: + """Forward pass. + + Args: + inputs: Dict mapping field names to image tensors (B, C, H, W). + output_mask: If True, also return attention masks. + + Returns: + Dict of embeddings (B, num_patches, E), optionally with masks. + """ + embedded: Dict[str, torch.Tensor] = {} + masks: Dict[str, torch.Tensor] = {} if output_mask else None + + for field_name, tensor in inputs.items(): + if field_name not in self.embedding_layers: + embedded[field_name] = tensor + continue + + tensor = tensor.to(self.device) + batch_size = tensor.size(0) + + x = self.embedding_layers[field_name](tensor) + + if self.use_cls_token: + cls = self.cls_tokens[field_name].expand(batch_size, -1, -1) + x = torch.cat([cls, x], dim=1) + + x = x + self.pos_embeddings[field_name] + x = self.dropout(x) + + if self.pool == "mean": + x = x.mean(dim=1, keepdim=True) + + embedded[field_name] = x + + if output_mask: + masks[field_name] = torch.ones( + batch_size, x.size(1), dtype=torch.bool, device=x.device + ) + + return (embedded, masks) if output_mask else embedded + + def get_output_info(self, field_name: str) -> Dict[str, Any]: + """Get metadata about embedding output for a field.""" + if field_name not in self._field_info: + raise KeyError(f"Field '{field_name}' not found") + + info = self._field_info[field_name].copy() + info["embedding_dim"] = self._embedding_dim + info["has_cls_token"] = self.use_cls_token + if self.pool == "mean": + info["num_tokens"] = 1 + else: + info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0) + return info + + def __repr__(self) -> str: + fields = list(self.embedding_layers.keys()) + return ( + f"VisionEmbeddingModel(backbone={self.backbone_type!r}, " + f"embedding_dim={self._embedding_dim}, fields={fields})" + ) + + +if __name__ == "__main__": + from pyhealth.datasets import create_sample_dataset + from pyhealth.datasets.utils import get_dataloader + import tempfile + import os + from PIL import Image + import numpy as np + + # Create synthetic images + temp_dir = tempfile.mkdtemp() + samples = [] + for i in range(10): + img_path = os.path.join(temp_dir, f"img_{i}.png") + img = Image.fromarray(np.random.randint(0, 255, (224, 224), dtype=np.uint8), mode="L") + img.save(img_path) + samples.append({ + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "chest_xray": img_path, + "label": i % 2, + }) + + dataset = create_sample_dataset( + samples=samples, + input_schema={"chest_xray": ("image", {"image_size": 224, "mode": "L"})}, + output_schema={"label": "binary"}, + dataset_name="test_vision", + ) + + model = VisionEmbeddingModel( + dataset=dataset, + embedding_dim=128, + backbone="cnn", + use_cls_token=True, + ) + + model_pooled = VisionEmbeddingModel( + dataset=dataset, + embedding_dim=128, + backbone="cnn", + pool="mean", + ) + + + + loader = get_dataloader(dataset, batch_size=4, shuffle=False) + batch = next(iter(loader)) + + embeddings_pooled = model_pooled({"chest_xray": batch["chest_xray"]}) + print(f"Pooled output shape: {embeddings_pooled['chest_xray'].shape}") # expect (4, 1, 128) + print(f"Pooled output info: {model_pooled.get_output_info('chest_xray')}") # expect num_tokens=1 + + + embeddings = model({"chest_xray": batch["chest_xray"]}) + print(f"Input shape: {batch['chest_xray'].shape}") + print(f"Output shape: {embeddings['chest_xray'].shape}") + print(f"Output info: {model.get_output_info('chest_xray')}") + + # Cleanup + shutil.rmtree(temp_dir) diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py index fea902bd1..37d738f3f 100644 --- a/pyhealth/models/jamba_ehr.py +++ b/pyhealth/models/jamba_ehr.py @@ -15,6 +15,7 @@ from pyhealth.datasets import SampleDataset from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.models.transformer import TransformerBlock from pyhealth.models.ehrmamba import MambaBlock from pyhealth.models.utils import get_last_visit @@ -177,6 +178,11 @@ class JambaEHR(BaseModel): by an independent :class:`JambaLayer`. The resulting patient embeddings are concatenated and projected through a classification head. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + :class:`JambaLayer` rather than one layer per field. + Args: dataset (SampleDataset): Dataset providing processed inputs. embedding_dim (int): Embedding and hidden dimension. Default 128. @@ -186,6 +192,8 @@ class JambaEHR(BaseModel): dropout (float): Dropout rate. Default 0.3. state_size (int): SSM state size in Mamba blocks. Default 16. conv_kernel (int): Causal conv kernel in Mamba blocks. Default 4. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, enables unified multi-modal mode with a single JambaLayer. Examples: >>> from pyhealth.datasets import create_sample_dataset, get_dataloader @@ -234,6 +242,7 @@ def __init__( dropout: float = 0.3, state_size: int = 16, conv_kernel: int = 4, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super(JambaEHR, self).__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -243,6 +252,7 @@ def __init__( self.dropout_rate = dropout self.state_size = state_size self.conv_kernel = conv_kernel + self._use_unified = unified_embedding is not None assert ( len(self.label_keys) == 1 @@ -250,11 +260,12 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) + output_size = self.get_output_size() + self.dropout = nn.Dropout(dropout) - self.jamba: nn.ModuleDict = nn.ModuleDict() - for feature_key in self.feature_keys: - self.jamba[feature_key] = JambaLayer( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_jamba = JambaLayer( feature_size=embedding_dim, num_transformer_layers=num_transformer_layers, num_mamba_layers=num_mamba_layers, @@ -263,12 +274,66 @@ def __init__( state_size=state_size, conv_kernel=conv_kernel, ) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.jamba: nn.ModuleDict = nn.ModuleDict() + for feature_key in self.feature_keys: + self.jamba[feature_key] = JambaLayer( + feature_size=embedding_dim, + num_transformer_layers=num_transformer_layers, + num_mamba_layers=num_mamba_layers, + heads=heads, + dropout=dropout, + state_size=state_size, + conv_kernel=conv_kernel, + ) + self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Calls UnifiedMultimodalEmbeddingModel to produce a single + temporally-sorted event sequence, then encodes it with one shared + JambaLayer and pools to the last valid event. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S_total, E) + mask = out["mask"] # (B, S_total) float, 1=valid 0=pad - output_size = self.get_output_size() - self.dropout = nn.Dropout(dropout) - self.fc = nn.Linear( - len(self.feature_keys) * embedding_dim, output_size - ) + _, cls_emb = self._unified_jamba(sequence, mask) + logits = self.fc(self.dropout(cls_emb)) + y_prob = self.prepare_y_prob(logits) + + results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob} + if self.label_key in kwargs: + y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = cls_emb + return results @staticmethod def _pool_embedding(x: torch.Tensor) -> torch.Tensor: @@ -317,9 +382,10 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. - Embeds each feature stream, encodes through the hybrid - Transformer-Mamba stack, concatenates per-stream patient - representations, and projects to label space. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields and processes them with a + single JambaLayer backbone. Otherwise each field is embedded and + encoded independently. Args: **kwargs: Must include all feature keys (tensors or tuples @@ -330,6 +396,9 @@ def forward( ``y_prob``, ``y_true``, ``logit``, and optionally ``embed`` if ``kwargs["embed"] is True``. """ + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] for feature_key in self.feature_keys: diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 3393d7287..94fcea0ad 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional, Tuple +from typing import Any, Dict, Optional, Tuple import torch import torch.nn as nn @@ -20,6 +20,7 @@ ) from .embedding import EmbeddingModel +from .embedding.unified import UnifiedMultimodalEmbeddingModel class RNNLayer(nn.Module): @@ -92,9 +93,7 @@ def forward( Args: x: a tensor of shape [batch size, sequence len, input size]. mask: an optional tensor of shape [batch size, sequence len], where - 1 indicates valid and 0 indicates invalid. Samples with all-zero - masks are clamped to length 1 to prevent pack_padded_sequence - from receiving zero-length sequences. + 1 indicates valid and 0 indicates invalid. Returns: outputs: a tensor of shape [batch size, sequence len, hidden size], @@ -111,13 +110,10 @@ def forward( ) else: lengths = torch.sum(mask.int(), dim=-1).cpu() - # Clamp lengths to at least 1 to handle empty sequences, - # matching TCNLayer (tcn.py:186). - lengths = torch.clamp(lengths, min=1) # Ensure tensor is contiguous for cuDNN compatibility x = x.contiguous() x = rnn_utils.pack_padded_sequence( - x, lengths, batch_first=True, enforce_sorted=False + x.contiguous(), lengths, batch_first=True, enforce_sorted=False ) outputs, _ = self.rnn(x) outputs, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True) @@ -211,6 +207,7 @@ def __init__( dataset: SampleDataset, embedding_dim: int = 128, hidden_dim: int = 128, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, **kwargs ): super(RNN, self).__init__( @@ -218,6 +215,7 @@ def __init__( ) self.embedding_dim = embedding_dim self.hidden_dim = hidden_dim + self._use_unified = unified_embedding is not None # validate kwargs for RNN layer if "input_size" in kwargs: raise ValueError("input_size is determined by embedding_dim") @@ -227,20 +225,71 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - - self.rnn = nn.ModuleDict() - for feature_key in self.dataset.input_processors.keys(): - self.rnn[feature_key] = RNNLayer( - input_size=embedding_dim, hidden_size=hidden_dim, **kwargs - ) output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + if self._use_unified: + self.embedding_model = unified_embedding + self.rnn = nn.ModuleDict({ + "unified": RNNLayer(input_size=embedding_dim, hidden_size=hidden_dim, **kwargs) + }) + self.fc = nn.Linear(hidden_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.rnn = nn.ModuleDict() + for feature_key in self.dataset.input_processors.keys(): + self.rnn[feature_key] = RNNLayer( + input_size=embedding_dim, hidden_size=hidden_dim, **kwargs + ) + self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Embeds all temporal fields jointly as a single time-sorted sequence + and processes it with one RNN backbone. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + mask = out["mask"].int() # (B, S) + + _, last_hidden = self.rnn["unified"](sequence, mask) # (B, hidden_dim) + logits = self.fc(last_hidden) + y_true = kwargs[self.label_key].to(self.device) + loss = self.get_loss_function()(logits, y_true) + y_prob = self.prepare_y_prob(logits) + results = {"loss": loss, "y_prob": y_prob, "y_true": y_true, "logit": logits} + if kwargs.get("embed", False): + results["embed"] = last_hidden + return results def forward(self, **kwargs) -> Dict[str, torch.Tensor]: """Forward propagation. - The label `kwargs[self.label_key]` is a list of labels for each patient. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields as a single time-sorted + sequence and processes it with one RNN. Otherwise each field is + embedded and encoded independently. Args: **kwargs: keyword arguments for the model. The keys must contain @@ -254,6 +303,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: - logit: a tensor representing the logits. - embed (optional): a tensor representing the patient embeddings if requested. """ + if self._use_unified: + return self._forward_unified(**kwargs) + patient_emb = [] # We need to preprocess kwargs to extract values and masks for EmbeddingModel @@ -569,4 +621,4 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: } if kwargs.get("embed", False): results["embed"] = patient_emb - return results + return results \ No newline at end of file diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index cc0dfc5ca..5bb2bdfe3 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -12,6 +12,7 @@ from pyhealth.datasets import SampleDataset from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel from pyhealth.interpret.api import CheferInterpretable # VALID_OPERATION_LEVEL = ["visit", "event"] @@ -54,7 +55,7 @@ def forward( # Use -inf so softmax produces exact zeros on padded positions, # avoiding a second masked_fill after softmax (saves one full # [B, H, S, S] boolean allocation and an extra copy). - pad_mask = (mask == 0) + pad_mask = mask == 0 scores = scores.masked_fill(pad_mask, -1e9) p_attn = self.softmax(scores) if dropout is not None: @@ -164,7 +165,7 @@ def forward( self.attn_map = None # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) - + return self.output_linear(x) @@ -246,7 +247,7 @@ def set_activation_hooks(self, hooks) -> None: """Deprecated compatibility stub; no-op.""" return None - def forward(self, x, mask=None, register_hook = False): + def forward(self, x, mask=None, register_hook=False): """Forward propagation. Args: @@ -256,7 +257,12 @@ def forward(self, x, mask=None, register_hook = False): Returns: A tensor of shape [batch_size, seq_len, hidden] """ - x = self.input_sublayer(x, lambda _x: self.attention(_x, _x, _x, mask=mask, register_hook=register_hook)) + x = self.input_sublayer( + x, + lambda _x: self.attention( + _x, _x, _x, mask=mask, register_hook=register_hook + ), + ) x = self.output_sublayer(x, lambda _x: self.feed_forward(_x, mask=mask)) return self.dropout(x) @@ -297,7 +303,10 @@ def set_activation_hooks(self, hooks) -> None: return None def forward( - self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, register_hook: bool = False + self, + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + register_hook: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """Forward propagation. @@ -328,12 +337,21 @@ class Transformer(BaseModel, CheferInterpretable): an independent :class:`TransformerLayer`. The resulting [CLS]-style embeddings are concatenated and passed to a classification head. + When ``unified_embedding`` is supplied the model switches to **unified + mode**: all temporal fields are jointly embedded and time-sorted by + :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single* + :class:`TransformerLayer` rather than one layer per field. This allows + full cross-modal attention over the interleaved event sequence. + Args: dataset (SampleDataset): dataset providing processed inputs. embedding_dim (int): shared embedding dimension. heads (int): number of attention heads per transformer block. dropout (float): dropout rate applied inside transformer blocks. num_layers (int): number of transformer blocks per feature stream. + unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when + provided, the model uses a single backbone over the unified + multi-modal sequence instead of per-field transformers. Examples: >>> from pyhealth.datasets import create_sample_dataset, get_dataloader @@ -377,6 +395,7 @@ def __init__( dropout: float = 0.5, num_layers: int = 1, max_seq_len: int = 1024, + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, ): super().__init__(dataset=dataset) self.embedding_dim = embedding_dim @@ -385,6 +404,7 @@ def __init__( self.num_layers = num_layers self.max_seq_len = max_seq_len self._attention_hooks_enabled = False + self._use_unified = unified_embedding is not None assert ( len(self.label_keys) == 1 @@ -392,19 +412,28 @@ def __init__( self.label_key = self.label_keys[0] self.mode = self.dataset.output_schema[self.label_key] - self.embedding_model = EmbeddingModel(dataset, embedding_dim) + output_size = self.get_output_size() - self.transformer: nn.ModuleDict = nn.ModuleDict() - for feature_key in self.feature_keys: - self.transformer[feature_key] = TransformerLayer( + if self._use_unified: + self.embedding_model = unified_embedding + self._unified_backbone = TransformerLayer( feature_size=embedding_dim, heads=heads, dropout=dropout, num_layers=num_layers, ) - - output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) + self.fc = nn.Linear(embedding_dim, output_size) + else: + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + self.transformer: nn.ModuleDict = nn.ModuleDict() + for feature_key in self.feature_keys: + self.transformer[feature_key] = TransformerLayer( + feature_size=embedding_dim, + heads=heads, + dropout=dropout, + num_layers=num_layers, + ) + self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size) def _pool_embedding(self, x: torch.Tensor) -> torch.Tensor: """Pool nested embeddings to ``[batch, seq_len, hidden]`` format. @@ -443,6 +472,61 @@ def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor: mask[invalid_rows, 0] = True return mask.bool() + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build inputs expected by UnifiedMultimodalEmbeddingModel.""" + + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + inputs[field_name] = field_dict + + return inputs + + def _forward_unified( + self, + **kwargs: torch.Tensor | tuple[torch.Tensor, ...], + ) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode.""" + + register_hook = self._attention_hooks_enabled + inputs = self._build_unified_inputs(cast(Dict[str, Any], kwargs)) + out = self.embedding_model(inputs) + sequence = cast(torch.Tensor, out["sequence"]) + event_mask = cast(torch.Tensor, out["mask"]).bool() + + _, patient_emb = self._unified_backbone(sequence, event_mask, register_hook) + + logits = self.fc(patient_emb) + y_prob = self.prepare_y_prob(logits) + + results: Dict[str, torch.Tensor] = { + "logit": logits, + "y_prob": y_prob, + } + + if self.label_key in kwargs: + y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) + loss = self.get_loss_function()(logits, y_true) + results["loss"] = loss + results["y_true"] = y_true + + if kwargs.get("embed", False): + results["embed"] = patient_emb + return results + def forward_from_embedding( self, **kwargs: torch.Tensor | tuple[torch.Tensor, ...], @@ -500,8 +584,7 @@ def forward_from_embedding( if value is None: raise ValueError( - f"Feature '{feature_key}' must contain 'value' " - f"in the schema." + f"Feature '{feature_key}' must contain 'value' " f"in the schema." ) else: value = value.to(self.device) @@ -515,9 +598,7 @@ def forward_from_embedding( else: mask = self._mask_from_embeddings(value).to(self.device) - _, cls_emb = self.transformer[feature_key]( - value, mask, register_hook - ) + _, cls_emb = self.transformer[feature_key](value, mask, register_hook) patient_emb.append(cls_emb) patient_emb = torch.cat(patient_emb, dim=1) @@ -545,6 +626,11 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields and processes them with a + single transformer backbone. Otherwise each field is embedded and + encoded independently. + Args: **kwargs: keyword arguments for the model. @@ -562,6 +648,9 @@ def forward( logit: the raw logits before activation. embed: (if embed=True in kwargs) the patient embedding. """ + if self._use_unified: + return self._forward_unified(**kwargs) + for feature_key in self.feature_keys: feature = kwargs[feature_key] @@ -575,15 +664,16 @@ def forward( if value is None: raise ValueError( - f"Feature '{feature_key}' must contain 'value' " - f"in the schema." + f"Feature '{feature_key}' must contain 'value' " f"in the schema." ) else: value = value.to(self.device) if mask is not None: mask = mask.to(self.device) - value = self.embedding_model({feature_key: value}, masks={feature_key: mask})[feature_key] + value = self.embedding_model( + {feature_key: value}, masks={feature_key: mask} + )[feature_key] else: value = self.embedding_model({feature_key: value})[feature_key] @@ -591,9 +681,9 @@ def forward( # Reconstruct tuple with embedded value # Note: we need to handle list/tuple conversion carefully # feature is a tuple. - + # Simple slice reconstruction - kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1:] + kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1 :] return self.forward_from_embedding(**kwargs) @@ -621,9 +711,7 @@ def get_attention_layers( cast(TransformerBlock, blk).attention.get_attn_map(), cast(TransformerBlock, blk).attention.get_attn_grad(), ) - for blk in cast( - TransformerLayer, self.transformer[key] - ).transformer + for blk in cast(TransformerLayer, self.transformer[key]).transformer ] for key in self.feature_keys } @@ -683,4 +771,4 @@ def get_relevance_tensor( result = model(**data_batch) print(result) - result["loss"].backward() + result["loss"].backward() \ No newline at end of file diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index 9d313e6bc..205782c7f 100644 --- a/pyhealth/processors/time_image_processor.py +++ b/pyhealth/processors/time_image_processor.py @@ -75,6 +75,11 @@ class TimeImageProcessor(TemporalFeatureProcessor): patient has more images, the most recent (by timestamp) are kept. If None, all images are kept. Defaults to None. + padding: Sentinel string that marks a missing image. When + a path equals this value, a zero tensor of shape + (C, H, W) is returned instead of loading from disk. + If None, all paths are treated as real file paths. + Defaults to None. Raises: ValueError: If normalize is True but mean or std is missing. @@ -107,6 +112,7 @@ def __init__( std: Optional[List[float]] = None, mode: Optional[str] = None, max_images: Optional[int] = None, + padding: Optional[str] = None, ) -> None: self.image_size = image_size self.to_tensor = to_tensor @@ -115,18 +121,14 @@ def __init__( self.std = std self.mode = mode self.max_images = max_images + self.padding = padding self.n_channels = None - if self.normalize and ( - self.mean is None or self.std is None - ): + if self.normalize and (self.mean is None or self.std is None): raise ValueError( - "Normalization requires both mean and std to be " - "provided." + "Normalization requires both mean and std to be " "provided." ) - if not self.normalize and ( - self.mean is not None or self.std is not None - ): + if not self.normalize and (self.mean is not None or self.std is not None): raise ValueError( "Mean and std are provided but normalize is set " "to False. Either provide normalize=True, or " @@ -146,36 +148,49 @@ def _build_transform(self) -> transforms.Compose: transform_list = [] if self.mode is not None: transform_list.append( - transforms.Lambda( - partial(_convert_mode, mode=self.mode) - ) + transforms.Lambda(partial(_convert_mode, mode=self.mode)) ) if self.image_size is not None: - transform_list.append( - transforms.Resize( - (self.image_size, self.image_size) - ) - ) + transform_list.append(transforms.Resize((self.image_size, self.image_size))) if self.to_tensor: transform_list.append(transforms.ToTensor()) if self.normalize: - transform_list.append( - transforms.Normalize( - mean=self.mean, std=self.std - ) - ) + transform_list.append(transforms.Normalize(mean=self.mean, std=self.std)) return transforms.Compose(transform_list) - def _load_single_image( - self, path: Union[str, Path] - ) -> torch.Tensor: + def _zero_image_tensor(self) -> torch.Tensor: + """Return a zero tensor matching the expected image shape (C, H, W). + + Used as a placeholder when an image path is an empty string. + Channel count is inferred from self.n_channels if available, + otherwise derived from self.mode ("L"→1, "RGBA"→4, else 3). + + Returns: + Zero tensor of shape (C, image_size, image_size). + """ + if self.n_channels is not None: + c = self.n_channels + elif self.mode == "L": + c = 1 + elif self.mode == "RGBA": + c = 4 + else: + c = 3 + return torch.zeros(c, self.image_size, self.image_size) + + def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor: """Load and transform a single image from disk. + If path equals missing_path_token, returns a zero tensor of + the same shape as a normal image (C, H, W) via _zero_image_tensor. + Called internally by process() for each image path in the input list. Args: - path: Path to the image file. + path: Path to the image file. If this equals + missing_path_token, a zero-filled placeholder tensor + is returned instead. Returns: Transformed image tensor of shape (C, H, W). @@ -183,18 +198,16 @@ def _load_single_image( Raises: FileNotFoundError: If the image file does not exist. """ + if self.padding is not None and str(path) == self.padding: + return self._zero_image_tensor() image_path = Path(path) if not image_path.exists(): - raise FileNotFoundError( - f"Image file not found: {image_path}" - ) + raise FileNotFoundError(f"Image file not found: {image_path}") with Image.open(image_path) as img: img.load() return self.transform(img) - def fit( - self, samples: Iterable[Dict[str, Any]], field: str - ) -> None: + def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: """Fit the processor by inferring n_channels from data. Scans samples to find the first valid entry for the given @@ -214,8 +227,10 @@ def fit( for sample in samples: if field in sample and sample[field] is not None: image_paths, _ = sample[field] - if len(image_paths) > 0: - path = Path(image_paths[0]) + for raw_path in image_paths: + if self.padding is not None and str(raw_path) == self.padding: + continue + path = Path(raw_path) if path.exists(): with Image.open(path) as img: if img.mode == "L": @@ -225,14 +240,14 @@ def fit( else: self.n_channels = 3 break + if self.n_channels is not None: + break if self.n_channels is None: self.n_channels = 3 def process( self, - value: Tuple[ - List[Union[str, Path]], List[float] - ], + value: Tuple[List[Union[str, Path]], List[float]], ) -> Tuple[torch.Tensor, torch.Tensor, str]: """Process paired image paths and timestamps. @@ -278,15 +293,10 @@ def process( if len(image_paths) == 0: raise ValueError("image_paths must be non-empty.") - paired = sorted( - zip(time_diffs, image_paths), key=lambda x: x[0] - ) + paired = sorted(zip(time_diffs, image_paths), key=lambda x: x[0]) - if ( - self.max_images is not None - and len(paired) > self.max_images - ): - paired = paired[-self.max_images:] + if self.max_images is not None and len(paired) > self.max_images: + paired = paired[-self.max_images :] timestamps = [] image_tensors = [] @@ -295,9 +305,7 @@ def process( timestamps.append(t) images = torch.stack(image_tensors, dim=0) - timestamps = torch.tensor( - timestamps, dtype=torch.float32 - ) + timestamps = torch.tensor(timestamps, dtype=torch.float32) if self.n_channels is None: self.n_channels = images.shape[1] @@ -344,5 +352,6 @@ def __repr__(self) -> str: f"mean={self.mean}, " f"std={self.std}, " f"mode={self.mode}, " - f"max_images={self.max_images})" + f"max_images={self.max_images}, " + f"padding={self.padding!r})" ) \ No newline at end of file diff --git a/pyhealth/processors/tuple_time_text_processor.py b/pyhealth/processors/tuple_time_text_processor.py index bbe74c4e6..cb1fe99a0 100644 --- a/pyhealth/processors/tuple_time_text_processor.py +++ b/pyhealth/processors/tuple_time_text_processor.py @@ -5,6 +5,7 @@ from . import register_processor logger = logging.getLogger(__name__) +_MISSING_TEXT_TOKEN = "[MISSING_TEXT]" @register_processor("tuple_time_text") class TupleTimeTextProcessor(TemporalFeatureProcessor): @@ -81,6 +82,41 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str] - str: Type tag """ texts, time_diffs = value + texts = list(texts or []) + time_diffs = list(time_diffs or []) + + # Keep text/time aligned and filter malformed text entries. + pair_count = min(len(texts), len(time_diffs)) + cleaned_texts: List[str] = [] + cleaned_times: List[float] = [] + for i in range(pair_count): + raw_text = texts[i] + raw_time = time_diffs[i] + + # Normalize text; skip null/whitespace-only entries. + if raw_text is None: + continue + text = str(raw_text).strip() + if text == "": + continue + + # Best-effort float normalization; skip unparseable timestamps. + try: + t = float(raw_time) + except (TypeError, ValueError): + continue + + cleaned_texts.append(text) + cleaned_times.append(t) + + # Fast tokenizer path crashes on empty batches; force a single + # missingness token when all notes are empty/malformed. + if len(cleaned_texts) == 0: + cleaned_texts = [_MISSING_TEXT_TOKEN] + cleaned_times = [0.0] + + texts = cleaned_texts + time_diffs = cleaned_times time_tensor = torch.tensor(time_diffs, dtype=torch.float32) if self.tokenizer is not None: @@ -93,16 +129,18 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str] return_tensors="pt" ) - input_ids = encoded["input_ids"] - attention_mask = encoded["attention_mask"] + input_ids = encoded.get("input_ids") + if input_ids is None: + raise ValueError("Tokenizer output is missing required `input_ids`.") + + attention_mask = encoded.get("attention_mask") + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) # Not all tokenizers return token_type_ids (e.g. RoBERTa might not, BERT does) - if "token_type_ids" in encoded: - token_type_ids = encoded["token_type_ids"] - else: - # meaningful text usually 0, padding 0? BERT uses 0 for sent A. - # If not provided, we can just use zeros or omit. - # For consistency with schema, let's provide zeros if expected. + token_type_ids = encoded.get("token_type_ids") + if token_type_ids is None: + # Some tokenizers do not return token_type_ids. token_type_ids = torch.zeros_like(input_ids) return input_ids, attention_mask, token_type_ids, time_tensor, self.type_tag @@ -172,4 +210,4 @@ def process_temporal(self, value) -> dict: def __repr__(self): if self.tokenizer_model: return f"TupleTimeTextProcessor(type_tag='{self.type_tag}', tokenizer='{self.tokenizer_model}')" - return f"TupleTimeTextProcessor(type_tag='{self.type_tag}')" + return f"TupleTimeTextProcessor(type_tag='{self.type_tag}')" \ No newline at end of file diff --git a/pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub b/pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub new file mode 100644 index 000000000..b0768f201 --- /dev/null +++ b/pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub @@ -0,0 +1,45 @@ +# HTCondor submission — labs-only RNN mortality run +# +# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py: +# same task (labs), same model (rnn), same hyperparameters. Runs unattended +# instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_only_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_only_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/run_labs_only_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \ + CACHE_DIR=/shared/eng/wp14/pyhealth_cache_labs \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-only \ + WANDB_RUN_NAME=labs_rnn_seed$(seed)" + +output = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 4 +request_memory = 163840MB +request_disk = 20GB + +Rank = TARGET.CUDAGlobalMemoryMb + +queue seed from ( + 12 +) \ No newline at end of file diff --git a/pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh b/pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh new file mode 100644 index 000000000..649eaccb1 --- /dev/null +++ b/pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh @@ -0,0 +1,168 @@ +#!/usr/bin/env bash +# HTCondor executable — labs-only RNN mortality run. +# +# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py: +# same task (labs), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_only_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/eng/wp14/pyhealth_cache_labs/* +set -euo pipefail + +SEED="${1:?usage: run_labs_only_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-0}" +EMBEDDING_DIM="${EMBEDDING_DIM:-64}" +HIDDEN_DIM="${HIDDEN_DIM:-64}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-1}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Skip the distributed Dask cluster (falls back to the plain local scheduler) +# to avoid touching NVML during event-dataframe preprocessing. +export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-only}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs-only RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir: ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --ehr-root "${EHR_ROOT}" + --cache-dir "${CACHE_DIR}" + --task labs + --model rnn + --embedding-dim "${EMBEDDING_DIM}" + --hidden-dim "${HIDDEN_DIM}" + --rnn-type "${RNN_TYPE}" + --rnn-layers "${RNN_LAYERS}" + --dropout "${DROPOUT}" + --epochs "${EPOCHS}" + --batch-size "${BATCH_SIZE}" + --lr "${LR}" + --weight-decay "${WEIGHT_DECAY}" + --patience "${PATIENCE}" + --num-workers "${NUM_WORKERS}" + --seed "${SEED}" + --output-dir "${OUTPUT_DIR}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" \ No newline at end of file diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py new file mode 100644 index 000000000..c3b1308e9 --- /dev/null +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -0,0 +1,1064 @@ +import logging +import re +from datetime import datetime, timedelta +from typing import Any, Dict, List, Optional, Union, Tuple, ClassVar + +from pyhealth.tasks.base_task import BaseTask + +logger = logging.getLogger(__name__) + + +class BaseMultimodalMIMIC4Task(BaseTask): + """Base class for multimodal MIMIC-IV tasks. + + Provides shared constants and utility methods used across all multimodal + task variants (notes, ICD codes, lab values). + """ + + MISSING_TEXT_TOKEN: ClassVar[str] = "" + MISSING_CODE_TOKEN: ClassVar[str] = "" + MISSING_FLOAT_TOKEN: ClassVar[float] = 0.0 + + LAB_CATEGORIES: ClassVar[Dict[str, List[str]]] = { + "Sodium": ["50824", "52455", "50983", "52623"], + "Potassium": ["50822", "52452", "50971", "52610"], + "Chloride": ["50806", "52434", "50902", "52535"], + "Bicarbonate": ["50803", "50804"], + "Glucose": ["50809", "52027", "50931", "52569"], + "Calcium": ["50808", "51624"], + "Magnesium": ["50960"], + "Anion Gap": ["50868", "52500"], + "Osmolality": ["52031", "50964", "51701"], + "Phosphate": ["50970"], + } + + LAB_CATEGORY_NAMES: ClassVar[List[str]] = [ + "Sodium", + "Potassium", + "Chloride", + "Bicarbonate", + "Glucose", + "Calcium", + "Magnesium", + "Anion Gap", + "Osmolality", + "Phosphate", + ] + + LABITEMS: ClassVar[List[str]] = [ + item for itemids in LAB_CATEGORIES.values() for item in itemids + ] + + RADIOLOGY_CLINICAL_HEADERS: ClassVar[List[str]] = [ + "indication", + "impression", + # "findings", + # "clinical history", + # "history", + # "comparison", + # "technique", + # "conclusion", + # "summary" + ] + + DISCHARGE_CLINICAL_HEADERS: ClassVar[List[str]] = [ + "chief complaint", + # "history of present illness", + # "hpi", + # "past medical history", + # "past medical and surgical history", + # "past medical/surgical history", + # "past surgical history", + # "medications on admission", + # "admission medications", + # "home medications", + # "social history", + # "family history", + # "allergies", + # "review of systems", + ] + + def __init__( + self, + window_hours: Optional[float] = None, + ): + self.window_hours = window_hours + + @staticmethod + def _clean_text(text: Optional[str]) -> Optional[str]: + """Return text if non-empty, otherwise None.""" + return text if text else None + + @staticmethod + def _parse_note_sections(text: str, note_type: str) -> Dict[str, str]: + """Split a note into {lowercased_header: content_text} pairs.""" + ext_text = text + '\n\n' + if note_type == "radiology": + section_re = re.compile(r'([a-zA-Z ]+):[ \t\n]+(.+?)\n{2,}', re.DOTALL) + elif note_type == "discharge": + section_re = re.compile(r'([a-zA-Z ]+):\n+(.+?)\n{2,}', re.DOTALL) + else: + raise ValueError(f"Note Type '{note_type}' not supported.") + return { + m.group(1).strip().lower(): m.group(2).strip() + for m in section_re.finditer(ext_text) + if m.end() - m.start() > 0 + } + + @staticmethod + def _parse_datetime(value: Any) -> Optional[datetime]: + if isinstance(value, datetime): + return value + if isinstance(value, str): + for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"): + try: + return datetime.strptime(value, fmt) + except ValueError: + continue + return None + + @staticmethod + def _to_hours(delta_seconds: float) -> float: + return delta_seconds / 3600.0 + + def _compute_effective_window( + self, + admissions_to_process: List[Any], + ) -> Tuple[datetime, Optional[datetime]]: + """Compute effective start/end from the global span of processed admissions. + + Returns: + Tuple of (effective_start, effective_end). + """ + global_start = admissions_to_process[0].timestamp + global_end: Optional[datetime] = None + + for a in admissions_to_process: + dt = self._parse_datetime(getattr(a, "dischtime", None)) + if dt is not None and (global_end is None or dt > global_end): + global_end = dt + + if self.window_hours is not None: + effective_start = global_start + effective_end = effective_start + timedelta(hours=self.window_hours) + return effective_start, effective_end + + effective_start = global_start + effective_end = global_end + + return effective_start, effective_end + + def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]: + """Build admissions to process and derive mortality label. + + Includes all admissions up to and including the first death admission. + Patients who die in their first (and only) admission are included as + positives — previously they were dropped, which silently discarded most + ICU mortality positives and collapsed positive rate from ~10% to ~2.7%. + This now matches stagenet's semantics: use all available admission data + and label as positive if any admission has hospital_expire_flag=1. + """ + admissions = patient.get_events(event_type="admissions") + if len(admissions) == 0: + return [], 0 + + admissions_to_process: List[Any] = [] + mortality_label = 0 + + for admission in admissions: + admissions_to_process.append(admission) + if admission.hospital_expire_flag in [1, "1"]: + mortality_label = 1 + break + + return admissions_to_process, mortality_label + + def _collect_icd_codes(self, patient: Any, hadm_id: Any) -> List[str]: + """Collect ICD diagnosis and procedure codes for one admission. + + Returns: + List of ICD code strings, or an empty list if none found. + """ + diagnoses_icd = patient.get_events( + event_type="diagnoses_icd", filters=[("hadm_id", "==", hadm_id)] + ) + procedures_icd = patient.get_events( + event_type="procedures_icd", filters=[("hadm_id", "==", hadm_id)] + ) + return [ + e.icd_code for e in diagnoses_icd if hasattr(e, "icd_code") and e.icd_code + ] + [ + e.icd_code for e in procedures_icd if hasattr(e, "icd_code") and e.icd_code + ] + + def _collect_labs( + self, + patient: Any, + admission_time: datetime, + end_time: datetime, + ) -> Tuple[List[float], List[List[float]], List[List[bool]]]: + """Collect lab values and observation masks for one admission. + + Args: + patient: Patient object. + admission_time: Start of the window; times are relative to this. + end_time: End of the window (inclusive). + + Returns: + Tuple of (lab_times, lab_values, lab_masks). ``lab_masks`` is a + parallel boolean tensor where ``True`` means observed and ``False`` + means imputed with 0.0. Falls back to a single missing placeholder + row when no valid lab events are found. + """ + try: + import polars as pl + except ImportError as exc: + raise ImportError("Polars is required for lab collection.") from exc + + labevents_df = patient.get_events( + event_type="labevents", + start=admission_time, + end=end_time, + return_df=True, + ) + + lab_times: List[float] = [] + lab_values: List[List[float]] = [] + lab_masks: List[List[bool]] = [] + + labevents_df = labevents_df.filter( + pl.col("labevents/itemid").is_in(self.LABITEMS) + ) + if labevents_df.height > 0: + labevents_df = labevents_df.with_columns( + pl.col("labevents/storetime").str.strptime( + pl.Datetime, "%Y-%m-%d %H:%M:%S" + ) + ) + labevents_df = labevents_df.filter( + pl.col("labevents/storetime") <= end_time + ) + if labevents_df.height > 0: + labevents_df = labevents_df.select( + pl.col("timestamp"), + pl.col("labevents/itemid"), + pl.col("labevents/valuenum").cast(pl.Float64), + ) + for lab_ts in sorted(labevents_df["timestamp"].unique().to_list()): + ts_labs = labevents_df.filter(pl.col("timestamp") == lab_ts) + lab_vector: List[float] = [] + lab_mask: List[bool] = [] + for category_name in self.LAB_CATEGORY_NAMES: + category_value = self.MISSING_FLOAT_TOKEN + observed = False + for itemid in self.LAB_CATEGORIES[category_name]: + matching = ts_labs.filter( + pl.col("labevents/itemid") == itemid + ) + if matching.height > 0: + category_value = matching["labevents/valuenum"][0] + observed = True + break + lab_vector.append(category_value) + lab_mask.append(observed) + lab_times.append( + self._to_hours((lab_ts - admission_time).total_seconds()) + ) + lab_values.append(lab_vector) + lab_masks.append(lab_mask) + else: # If missing lab for a given admission + lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + lab_times.append(self.MISSING_FLOAT_TOKEN) + + if len(lab_values) == 0: # If missing lab for ALL admissions + lab_values.append([self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)) + lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + lab_times.append(self.MISSING_FLOAT_TOKEN) + return lab_times, lab_values, lab_masks + + def _collect_notes( + self, + patient: Any, + note_event_type: str, + hadm_id: Any, + admission_time: datetime, + start_time: Optional[datetime] = None, + end_time: Optional[datetime] = None, + section_headers: Optional[List[str]] = None, + fallback_to_full_note: bool = False, + ) -> Tuple[List[str], List[float]]: + """Collect notes of a given type for one admission. + + Args: + patient: Patient object. + note_event_type: Event type string (e.g. "discharge", "radiology"). + hadm_id: Admission ID to filter by. + admission_time: Admission start time; used to compute time offsets. + start_time: Optional start of the time window. + end_time: Optional end of the time window. + section_headers: When provided, extract only these named sections + from each note (lowercased match against parsed headers). + fallback_to_full_note: When True (default), falls back to the full + note text if no matching sections are found. When False, notes + with no matching sections are dropped entirely. + + Returns: + Tuple of (texts, hours_from_admission). Falls back to + ``([MISSING_TEXT_TOKEN], [MISSING_FLOAT_TOKEN])`` when the events + list is empty. + """ + notes = patient.get_events( + event_type=note_event_type, + start=start_time, + end=end_time, + filters=[("hadm_id", "==", hadm_id)], + ) + + texts: List[str] = [] + note_times: List[float] = [] + for note in notes: + try: + note_text = self._clean_text(note.text) + if note_text: + if section_headers is not None: + parsed = self._parse_note_sections(note_text, note_type=note_event_type) + extracted = [f"{k}: {v}" for k, v in parsed.items() if k in section_headers and v] + if extracted: + note_text = " [SEP] ".join(extracted) + elif not fallback_to_full_note: + continue + + time_from_admission = self._to_hours( + (note.timestamp - admission_time).total_seconds() + ) + texts.append(note_text) + note_times.append(time_from_admission) + except ( + AttributeError + ): # note object is missing .text or .timestamp attribute (e.g. malformed note) + pass + + return texts, note_times + + +class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): + """Task for ICD codes + lab values mortality prediction using MIMIC-IV. + + A notes-free structured-EHR task that uses only: + + - **ICD codes**: diagnosis and procedure codes per admission, processed by + ``StageNetProcessor`` with inter-admission time offsets. + - **Lab values**: 10-dimensional lab vectors (one per lab category) at each + measurement timestamp, processed by ``StageNetTensorProcessor``. + + Examples: + >>> from pyhealth.datasets import MIMIC4Dataset + >>> from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4 + >>> dataset = MIMIC4Dataset( + ... ehr_root="/path/to/mimic-iv/2.2", + ... ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + ... ) + >>> task = ICDLabsMIMIC4() + >>> samples = dataset.set_task(task) + """ + + PADDING: int = 0 + + task_name: str = "ICDLabsMIMIC4" + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { + "icd_codes": ("stagenet", {"padding": PADDING}), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + } + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + demographics = patient.get_events(event_type="patients") + if not demographics: + return [] + + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + + if len(admissions_to_process) == 0: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + previous_admission_time = None + + for admission in admissions_to_process: + admission_time = admission.timestamp + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + if visit_icd_codes: + if previous_admission_time is None: + time_from_previous = 0.0 + else: + time_from_previous = self._to_hours( + (admission_time - previous_admission_time).total_seconds() + ) + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + else: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + + previous_admission_time = admission_time + + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=admission_dischtime, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + if len(all_lab_values) == 0: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) + + if len(all_icd_codes) == 0: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "icd_codes": (all_icd_times, all_icd_codes), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + return [single_patient_longitudinal_record] + + +class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): + """Mortality prediction from admission-context notes and lab values. + + Follows the approach of Lee et al. (2023): use text that is clinically + available *at admission* rather than discharge notes. ICD codes are + excluded by default but can be re-enabled for ablation experiments. + + Text is extracted from the MIMIC-IV discharge note by parsing the Chief + Complaint, History of Present Illness, Past Medical History, and Medications + on Admission sections — all of which describe the patient's state at the + start of the stay. The extracted text is assigned timestamp 0.0. + + Radiology reports are also included, parsed for their Indication and + Impression sections and bounded to the same observation window as labs + (rather than timestamp 0.0), since — unlike the discharge summary — they + are written at exam time and describe findings from later in the stay. + + Fields: + admission_note_times: Admission-context discharge-note text at time + 0.0, plus in-window radiology note text at its exam-relative + timestamp. + labs: 10-dim lab vectors at each measurement timestamp. + labs_mask: Boolean observation mask parallel to ``labs``. + icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes + per admission with inter-admission time offsets. + + Args: + window_hours: Hours from admission for lab collection. ``None`` + collects for the full admission span. Default: 24. + include_icd: When ``True``, collect discharge-coded ICD codes and add + ``icd_codes`` to the sample dict / input schema. Default: ``False``. + """ + + PADDING: int = 0 + + task_name: str = "NotesLabsMIMIC4" + + _BASE_INPUT_SCHEMA: ClassVar[Dict] = { + "admission_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", + "type_tag": "note", + }, + ), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + } + + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __init__( + self, + window_hours: Optional[float] = None, + include_icd: bool = False, + ) -> None: + super().__init__(window_hours=window_hours) + self.include_icd = include_icd + schema = dict(self._BASE_INPUT_SCHEMA) + if include_icd: + schema["icd_codes"] = ("stagenet", {"padding": self.PADDING}) + self.input_schema = schema + logger.info( + "NotesLabsMIMIC4: filtering discharge notes to sections: %s; " + "radiology notes to sections: %s", + self.DISCHARGE_CLINICAL_HEADERS, + self.RADIOLOGY_CLINICAL_HEADERS, + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + if not patient.get_events(event_type="patients"): + return [] + + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_note_texts: List[str] = [] + all_note_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + previous_admission_time = None + + for admission in admissions_to_process: + admission_time = admission.timestamp + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + note_texts, note_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + section_headers=self.DISCHARGE_CLINICAL_HEADERS, + ) + all_note_texts.extend(note_texts) + all_note_times.extend(note_times) + + # Labs within the observation window + lab_end = ( + effective_end + if self.window_hours is not None + else admission_dischtime + ) + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=lab_end, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + # Radiology notes within the observation window. Unlike the + # discharge note (a retrospective summary parsed for its + # admission-context sections), radiology reports are written at + # exam time, so they're bounded to the same window as labs + # to avoid pulling in findings from later in the stay. + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=admission_time, + end_time=lab_end, + section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + ) + all_note_texts.extend(radiology_texts) + all_note_times.extend(radiology_times) + + if self.include_icd: + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + time_from_previous = ( + 0.0 + if previous_admission_time is None + else self._to_hours( + (admission_time - previous_admission_time).total_seconds() + ) + ) + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + else: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + previous_admission_time = admission_time + + if not all_lab_values: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) + + if not all_note_texts: + all_note_texts = [self.MISSING_TEXT_TOKEN] + all_note_times = [self.MISSING_FLOAT_TOKEN] + + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "admission_note_times": (all_note_texts, all_note_times), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + if self.include_icd: + if not all_icd_codes: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + record["icd_codes"] = (all_icd_times, all_icd_codes) + + return [record] + + +class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): + """Mortality prediction from admission-context notes, labs, and CXR. + + Extends ``NotesLabsMIMIC4`` with chest X-ray images: the same + admission-context discharge-note sections, in-window radiology reports, + and labs, plus CXR studies (StudyDate+StudyTime from the ``metadata`` + event table) bounded to the same observation window as labs/radiology. + ICD codes are excluded by default but can be re-enabled for ablation + experiments, same as ``NotesLabsMIMIC4``. + + Fields: + admission_note_times: Admission-context discharge-note text at time + 0.0, plus in-window radiology note text at its exam-relative + timestamp. + labs: 10-dim lab vectors at each measurement timestamp. + labs_mask: Boolean observation mask parallel to ``labs``. + cxr_image_times: In-window CXR image paths at their exam-relative + timestamp. + icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes + per admission with inter-admission time offsets. + + Args: + window_hours: Hours from admission for lab/CXR collection. + ``None`` (default) collects for the full admission span. + include_icd: When ``True``, collect discharge-coded ICD codes and add + ``icd_codes`` to the sample dict / input schema. Default: ``False``. + """ + + PADDING: int = 0 + + task_name: str = "NotesLabsCXRMIMIC4" + + _BASE_INPUT_SCHEMA: ClassVar[Dict] = { + "admission_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", + "type_tag": "note", + }, + ), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA + output_schema: Dict[str, str] = {"mortality": "binary"} + + def __init__( + self, + window_hours: Optional[float] = None, + include_icd: bool = False, + ) -> None: + super().__init__(window_hours=window_hours) + self.include_icd = include_icd + schema = dict(self._BASE_INPUT_SCHEMA) + if include_icd: + schema["icd_codes"] = ("stagenet", {"padding": self.PADDING}) + self.input_schema = schema + logger.info( + "NotesLabsCXRMIMIC4: filtering discharge notes to sections: %s; " + "radiology notes to sections: %s", + self.DISCHARGE_CLINICAL_HEADERS, + self.RADIOLOGY_CLINICAL_HEADERS, + ) + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: + if not patient.get_events(event_type="patients"): + return [] + + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_note_texts: List[str] = [] + all_note_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + all_lab_times: List[float] = [] + all_icd_codes: List[List[str]] = [] + all_icd_times: List[float] = [] + all_cxr_paths: List[str] = [] + all_cxr_times: List[float] = [] + previous_admission_time = None + + for admission in admissions_to_process: + admission_time = admission.timestamp + + # Skip admissions that start at or after the observation window + # closes, prevents Polars searchsorted OverflowError in CXR lookup. + if effective_end is not None and admission_time >= effective_end: + continue + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + note_texts, note_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + section_headers=self.DISCHARGE_CLINICAL_HEADERS, + ) + all_note_texts.extend(note_texts) + all_note_times.extend(note_times) + + # Labs within the observation window + lab_end = ( + effective_end + if self.window_hours is not None + else admission_dischtime + ) + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=lab_end, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + # Radiology notes within the observation window. Unlike the + # discharge note (a retrospective summary parsed for its + # admission-context sections), radiology reports are written at + # exam time, so they're bounded to the same window as labs/CXR + # to avoid pulling in findings from later in the stay. + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=admission_time, + end_time=lab_end, + section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + ) + all_note_texts.extend(radiology_texts) + all_note_times.extend(radiology_times) + + # CXR studies within the same observation window as labs/radiology. + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=lab_end, + ) + for event in metadata_events: + try: + if event.image_path: + all_cxr_paths.append(event.image_path) + all_cxr_times.append( + self._to_hours( + (event.timestamp - admission_time).total_seconds() + ) + ) + except AttributeError: + continue + + if self.include_icd: + visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) + time_from_previous = ( + 0.0 + if previous_admission_time is None + else self._to_hours( + (admission_time - previous_admission_time).total_seconds() + ) + ) + if visit_icd_codes: + all_icd_codes.append(visit_icd_codes) + all_icd_times.append(time_from_previous) + else: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + previous_admission_time = admission_time + + if not all_lab_values: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) + + if not all_note_texts: + all_note_texts = [self.MISSING_TEXT_TOKEN] + all_note_times = [self.MISSING_FLOAT_TOKEN] + + # time_image processor expects at least one path/time pair. + if len(all_cxr_paths) == 0: + all_cxr_paths = [self.MISSING_TEXT_TOKEN] + all_cxr_times = [self.MISSING_FLOAT_TOKEN] + + record: Dict[str, Any] = { + "patient_id": patient.patient_id, + "admission_note_times": (all_note_texts, all_note_times), + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "cxr_image_times": (all_cxr_paths, all_cxr_times), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + if self.include_icd: + if not all_icd_codes: + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) + all_icd_times.append(self.MISSING_FLOAT_TOKEN) + record["icd_codes"] = (all_icd_times, all_icd_codes) + + return [record] + + +class LabsMIMIC4(BaseMultimodalMIMIC4Task): + """EHR-only mortality prediction using lab values — no notes, no ICD codes. + + Serves as the structured-EHR reference baseline for multimodal ablations. + Collecting only ``labevents`` keeps the dataset loader fast and avoids any + leakage from discharge-coded ICD tables. + + Schema mirrors the ``labs`` / ``labs_mask`` fields from ``NotesLabsMIMIC4`` + so the same backbone models (MLP, RNN, Transformer, etc.) work unchanged. + + Args: + window_hours: Hours from admission to collect lab measurements. + ``None`` collects for the full admission span. Default: 24. + """ + + PADDING: int = 0 + + task_name: str = "LabsMIMIC4" + + input_schema: ClassVar[Dict] = { + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + } + output_schema: ClassVar[Dict] = {"mortality": "binary"} + + def __init__(self, window_hours: Optional[float] = 24) -> None: + super().__init__() + self.window_hours = window_hours + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_lab_times: List[float] = [] + all_lab_values: List[List[float]] = [] + all_lab_masks: List[List[bool]] = [] + + for admission in admissions_to_process: + admission_time = admission.timestamp + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=admission_dischtime, + ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) + + if len(all_lab_values) == 0: + all_lab_values.append( + [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) + ) + all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) + all_lab_times.append(self.MISSING_FLOAT_TOKEN) + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "labs": (all_lab_times, all_lab_values), + "labs_mask": (all_lab_times, all_lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + return [single_patient_longitudinal_record] + + +class CXRMIMIC4(BaseMultimodalMIMIC4Task): + """CXR-only mortality prediction using chest X-ray images. + + Serves as the imaging-only reference baseline for multimodal ablations — + no notes, no ICD codes, no labs — isolating the chest X-ray + modality the same way ``LabsMIMIC4`` isolates labs. + + CXR studies are filtered by timestamp (StudyDate+StudyTime, from the + ``metadata`` event table) within each admission's observation window. + + Args: + window_hours: Hours from admission to collect CXR studies. ``None`` + collects for the full admission span. Default: None. + """ + + task_name: str = "CXRMIMIC4" + + input_schema: ClassVar[Dict] = { + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + output_schema: ClassVar[Dict] = {"mortality": "binary"} + + def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + if not admissions_to_process: + return [] + + effective_start, effective_end = self._compute_effective_window( + admissions_to_process + ) + + all_cxr_paths: List[str] = [] + all_cxr_times: List[float] = [] + + for admission in admissions_to_process: + admission_time = admission.timestamp + + # Skip admissions that start at or after the observation window + # closes, prevents Polars searchsorted OverflowError. + if effective_end is not None and admission_time >= effective_end: + continue + + try: + admission_dischtime = datetime.strptime( + admission.dischtime, "%Y-%m-%d %H:%M:%S" + ) + except (ValueError, AttributeError): + admission_dischtime = admission_time + if admission_dischtime < admission_time: + admission_dischtime = admission_time + + admission_end = admission_dischtime + if effective_end is not None and effective_end < admission_end: + admission_end = effective_end + + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=admission_end, + ) + for event in metadata_events: + try: + if event.image_path: + all_cxr_paths.append(event.image_path) + all_cxr_times.append( + self._to_hours( + (event.timestamp - admission_time).total_seconds() + ) + ) + except AttributeError: + continue + + # time_image processor expects at least one path/time pair. + if len(all_cxr_paths) == 0: + all_cxr_paths = [self.MISSING_TEXT_TOKEN] + all_cxr_times = [self.MISSING_FLOAT_TOKEN] + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "cxr_image_times": (all_cxr_paths, all_cxr_times), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } + + return [single_patient_longitudinal_record] \ No newline at end of file diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index bc6a28677..8a86cc709 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -1,5 +1,7 @@ +import json import logging import os +import time from datetime import datetime from typing import Callable, Dict, List, Optional, Type @@ -37,6 +39,15 @@ def set_logger(log_path: str) -> None: return +def _vram_stats(device: str) -> Dict[str, float]: + """Returns current and peak VRAM usage in MB for a CUDA device.""" + if not torch.cuda.is_available() or not str(device).startswith("cuda"): + return {} + allocated = torch.cuda.memory_allocated(device) / 1024**2 + peak = torch.cuda.max_memory_allocated(device) / 1024**2 + return {"vram_allocated_mb": allocated, "vram_peak_mb": peak} + + def get_metrics_fn(mode: str) -> Callable: if mode == "binary": return binary_metrics_fn @@ -126,6 +137,9 @@ def train( monitor_criterion: str = "max", load_best_model_at_last: bool = True, patience=None, + accumulation_steps: int = 1, + use_amp: bool = False, + amp_dtype: str = "bf16", ): """Trains the model. @@ -145,10 +159,25 @@ def train( Default is True. patience: Number of epochs to wait for improvement before early stopping. Default is None, which means no early stopping. + accumulation_steps: Gradient accumulation steps to simulate a larger + effective batch size. Default is 1 (no accumulation). + use_amp: Whether to use automatic mixed precision. Default is False. + amp_dtype: AMP dtype — "bf16" (stable, recommended) or "fp16". + Default is "bf16". """ if optimizer_params is None: optimizer_params = {"lr": 1e-3} + _amp_dtype = ( + torch.bfloat16 if amp_dtype == "bf16" else torch.float16 + ) + # GradScaler only needed for fp16; bf16 has fp32 dynamic range + scaler = ( + torch.cuda.amp.GradScaler() + if (use_amp and _amp_dtype == torch.float16) + else None + ) + # logging logger.info("Training:") logger.info(f"Batch size: {train_dataloader.batch_size}") @@ -161,6 +190,8 @@ def train( logger.info(f"Monitor criterion: {monitor_criterion}") logger.info(f"Epochs: {epochs}") logger.info(f"Patience: {patience}") + logger.info(f"Accumulation steps: {accumulation_steps}") + logger.info(f"AMP: {use_amp} (dtype={amp_dtype})") # set optimizer param = list(self.model.named_parameters()) @@ -184,50 +215,129 @@ def train( steps_per_epoch = len(train_dataloader) global_step = 0 patience_counter = 0 + metrics_history: List[Dict] = [] + train_start = time.perf_counter() + total_skipped_steps = 0 # epoch training loop - for epoch in range(epochs): + epoch_iterator = tqdm(range(epochs), desc="Epochs", unit="epoch") + for epoch in epoch_iterator: + epoch_iterator.set_postfix_str(f"{epoch + 1}/{epochs}", refresh=False) training_loss = [] + epoch_skipped_steps = 0 self.model.zero_grad() self.model.train() + if torch.cuda.is_available() and str(self.device).startswith("cuda"): + torch.cuda.reset_peak_memory_stats(self.device) + epoch_start = time.perf_counter() # batch training loop logger.info("") - for _ in trange( + for step_idx in trange( steps_per_epoch, - desc=f"Epoch {epoch} / {epochs}", + desc=f"Epoch {epoch + 1}/{epochs}", smoothing=0.05, + leave=False, ): try: data = next(data_iterator) except StopIteration: data_iterator = iter(train_dataloader) data = next(data_iterator) - # forward - output = self.model(**data) - loss = output["loss"] + # forward (with optional AMP) + if use_amp: + with torch.autocast(device_type="cuda", dtype=_amp_dtype): + output = self.model(**data) + loss = output["loss"] / accumulation_steps + else: + output = self.model(**data) + loss = output["loss"] / accumulation_steps # backward - loss.backward() - if max_grad_norm is not None: - torch.nn.utils.clip_grad_norm_( - self.model.parameters(), max_grad_norm + if scaler is not None: + scaler.scale(loss).backward() + else: + loss.backward() + training_loss.append(loss.item() * accumulation_steps) + # optimizer step every accumulation_steps batches or epoch end + is_update_step = ( + (step_idx + 1) % accumulation_steps == 0 + or (step_idx + 1) == steps_per_epoch + ) + if is_update_step: + if scaler is not None: + scaler.unscale_(optimizer) + # Always compute the grad norm (even with no clipping + # configured) so non-finite gradients can be detected and + # skipped before they permanently poison the model with + # NaN weights. + grad_norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + max_grad_norm if max_grad_norm is not None else float("inf"), ) - # update - optimizer.step() - optimizer.zero_grad() - training_loss.append(loss.item()) - global_step += 1 + step_ok = bool(torch.isfinite(grad_norm)) + if not step_ok: + epoch_skipped_steps += 1 + total_skipped_steps += 1 + logger.warning( + f"epoch-{epoch} step-{global_step}: non-finite " + f"gradient norm ({grad_norm}); skipping optimizer " + f"step." + ) + if scaler is not None: + if step_ok: + scaler.step(optimizer) + scaler.update() + elif step_ok: + optimizer.step() + optimizer.zero_grad() + global_step += 1 + + epoch_time = time.perf_counter() - epoch_start + vram = _vram_stats(self.device) + + epochs_done = epoch + 1 + epochs_left = epochs - epochs_done + elapsed_total = time.perf_counter() - train_start + avg_epoch_time = elapsed_total / epochs_done + eta_s = avg_epoch_time * epochs_left + eta_h, eta_rem = divmod(int(eta_s), 3600) + eta_m = eta_rem // 60 + eta_str = f"{eta_h}h{eta_m:02d}m" + # log and save logger.info(f"--- Train epoch-{epoch}, step-{global_step} ---") logger.info(f"loss: {sum(training_loss) / len(training_loss):.4f}") + logger.info(f"epoch_time: {epoch_time:.2f}s elapsed: {elapsed_total:.0f}s ETA: {eta_str} ({epochs_done}/{epochs} epochs)") + print(f"[ETA] epoch {epochs_done}/{epochs} done in {epoch_time:.0f}s — ETA to finish: {eta_str}", flush=True) + if vram: + logger.info( + f"vram_peak: {vram['vram_peak_mb']:.1f} MB " + f"vram_current: {vram['vram_allocated_mb']:.1f} MB" + ) + if epoch_skipped_steps > 0: + logger.warning( + f"Skipped {epoch_skipped_steps} optimizer step(s) this " + f"epoch due to non-finite gradients (total so far: " + f"{total_skipped_steps})." + ) if self.exp_path is not None: self.save_ckpt(os.path.join(self.exp_path, "last.ckpt")) + epoch_record: Dict = { + "epoch": epoch, + "global_step": global_step, + "train_loss": sum(training_loss) / len(training_loss), + "epoch_time_s": round(epoch_time, 3), + "skipped_steps": epoch_skipped_steps, + **{f"train_{k}": v for k, v in vram.items()}, + } + # validation if val_dataloader is not None: scores = self.evaluate(val_dataloader) logger.info(f"--- Eval epoch-{epoch}, step-{global_step} ---") for key in scores.keys(): logger.info("{}: {:.4f}".format(key, scores[key])) + epoch_record.update({f"val_{k}": v for k, v in scores.items()}) # save best model if monitor is not None: score = scores[monitor] @@ -247,8 +357,21 @@ def train( logger.info( f"Early stopping at epoch-{epoch}, step-{global_step}" ) + metrics_history.append(epoch_record) break + metrics_history.append(epoch_record) + + total_time = time.perf_counter() - train_start + logger.info(f"--- Training complete: {total_time:.2f}s total ---") + + # persist metrics history + if self.exp_path is not None: + history_path = os.path.join(self.exp_path, "metrics_history.json") + with open(history_path, "w") as f: + json.dump(metrics_history, f, indent=2) + logger.info(f"Metrics history saved to {history_path}") + # load best model if load_best_model_at_last and self.exp_path is not None and os.path.isfile( os.path.join(self.exp_path, "best.ckpt")): @@ -262,7 +385,7 @@ def train( for key in scores.keys(): logger.info("{}: {:.4f}".format(key, scores[key])) - return + return metrics_history def inference(self, dataloader, additional_outputs=None, return_patient_ids=False) -> Dict[str, float]: @@ -430,4 +553,4 @@ def forward(self, x, y, **kwargs): monitor="accuracy", epochs=5, test_dataloader=val_dataloader, - ) + ) \ No newline at end of file From bc053040ec91e125c6a2a9782d5cd4384157d884 Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 01:45:04 -0700 Subject: [PATCH 02/27] Updates --- .../will/condor/labs_only/labs_only_rnn.sub | 0 .../will/condor/labs_only/run_labs_only_rnn.sh | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename {pyhealth/scripts_delete_me => scripts}/will/condor/labs_only/labs_only_rnn.sub (100%) rename {pyhealth/scripts_delete_me => scripts}/will/condor/labs_only/run_labs_only_rnn.sh (100%) diff --git a/pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub b/scripts/will/condor/labs_only/labs_only_rnn.sub similarity index 100% rename from pyhealth/scripts_delete_me/will/condor/labs_only/labs_only_rnn.sub rename to scripts/will/condor/labs_only/labs_only_rnn.sub diff --git a/pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh b/scripts/will/condor/labs_only/run_labs_only_rnn.sh similarity index 100% rename from pyhealth/scripts_delete_me/will/condor/labs_only/run_labs_only_rnn.sh rename to scripts/will/condor/labs_only/run_labs_only_rnn.sh From 5483dee99dd04302fd31ef5996cb18862538b38b Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 03:59:00 -0500 Subject: [PATCH 03/27] New Changes --- scripts/will/condor/labs_only/labs_only_rnn.sub | 4 ++-- scripts/will/condor/labs_only/run_labs_only_rnn.sh | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) mode change 100644 => 100755 scripts/will/condor/labs_only/run_labs_only_rnn.sh diff --git a/scripts/will/condor/labs_only/labs_only_rnn.sub b/scripts/will/condor/labs_only/labs_only_rnn.sub index b0768f201..e7d66a1b1 100644 --- a/scripts/will/condor/labs_only/labs_only_rnn.sub +++ b/scripts/will/condor/labs_only/labs_only_rnn.sub @@ -7,14 +7,14 @@ # # To submit (from the project root): # mkdir -p /home/wp14/logs/condor -# condor_submit scripts/will/condor/labs_only_rnn.sub +# condor_submit scripts/will/condor/labs_only/labs_only_rnn.sub # # Monitor: # condor_q # tail -f /home/wp14/logs/condor/labs_only_rnn__0.out initialdir = /home/wp14/PyHealth -executable = /home/wp14/PyHealth/scripts/will/condor/run_labs_only_rnn.sh +executable = /home/wp14/PyHealth/scripts/will/condor/labs_only/run_labs_only_rnn.sh transfer_executable = False arguments = $(seed) getenv = True diff --git a/scripts/will/condor/labs_only/run_labs_only_rnn.sh b/scripts/will/condor/labs_only/run_labs_only_rnn.sh old mode 100644 new mode 100755 index 649eaccb1..e05a8895f --- a/scripts/will/condor/labs_only/run_labs_only_rnn.sh +++ b/scripts/will/condor/labs_only/run_labs_only_rnn.sh @@ -20,7 +20,7 @@ CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs}" OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" CONDA_SH="${CONDA_SH:-}" -DEV_MODE="${DEV_MODE:-0}" +DEV_MODE="${DEV_MODE:-1}" EMBEDDING_DIM="${EMBEDDING_DIM:-64}" HIDDEN_DIM="${HIDDEN_DIM:-64}" RNN_TYPE="${RNN_TYPE:-GRU}" From a0f1422e4c885254ac2f9301c68f5689ac501b8e Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 04:21:41 -0500 Subject: [PATCH 04/27] New updates --- pyhealth/datasets/base_dataset.py | 187 ++++++++++---------------- pyhealth/models/embedding/__init__.py | 33 +++++ 2 files changed, 101 insertions(+), 119 deletions(-) create mode 100644 pyhealth/models/embedding/__init__.py diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 3d449d579..5618f4e9c 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -3,7 +3,7 @@ import pickle from abc import ABC from pathlib import Path -from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable +from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable, Union import functools import operator from urllib.parse import urlparse, urlunparse @@ -84,7 +84,13 @@ def path_exists(path: str) -> bool: except requests.RequestException: return False else: - return Path(path).exists() + try: + return Path(path).exists() + except OSError: + # Treat unreadable paths (e.g. stale/corrupted filesystem + # entries that raise I/O errors on stat) as non-existent so + # callers can fall back to an alternate extension. + return False def _csv_tsv_gz_path(path: str) -> str: @@ -314,16 +320,7 @@ class BaseDataset(ABC): dataset_name (str): Name of the dataset. config (dict): Configuration loaded from a YAML file. global_event_df (pl.LazyFrame): The global event data frame. - dev (bool): Whether to enable dev mode (limit to 1000 patients). - - Examples: - >>> from pyhealth.datasets import BaseDataset - >>> dataset = BaseDataset( - ... root="/path/to/source", - ... tables=["patients", "diagnoses"], - ... config_path="/path/to/config.yaml", - ... ) - >>> dataset.stats() + dev (Union[bool, int]): Whether to enable dev mode. If True, limit to 1000 patients. If an int, limit to that many patients. """ def __init__( @@ -334,7 +331,7 @@ def __init__( config_path: Optional[str] = None, cache_dir: str | Path | None = None, num_workers: int = 1, - dev: bool = False, + dev: Union[bool, int] = False, ): """Initializes the BaseDataset. @@ -351,7 +348,7 @@ def __init__( - **str** or **Path**: Used as the root cache directory path. A UUID is appended to the provided path to capture dataset configuration. num_workers (int): Number of worker processes for parallel operations. - dev (bool): Whether to run in dev mode (limits to 1000 patients). + dev (Union[bool, int]): Whether to run in dev mode. If True, limits to 1000 patients. If an int, limits to that many patients. """ if len(set(tables)) != len(tables): logger.warning("Duplicate table names in tables list. Removing duplicates.") @@ -434,68 +431,6 @@ def clean_tmpdir(self) -> None: if tmp_dir.exists(): shutil.rmtree(tmp_dir) - def _scan_table(self, source_path: str) -> dd.DataFrame: - """Routes a table source to the appropriate scanner based on its format. - - Parquet sources (``.parquet``/``.pq`` files, glob patterns targeting - such files, or directories of Parquet shards) are handled by - :meth:`_scan_parquet`. Any other source falls back to the existing - CSV/TSV(.gz) scanner, preserving prior behavior for all datasets. - - Args: - source_path (str): Path to the table source. - - Returns: - dd.DataFrame: The Dask DataFrame for the table source. - """ - stripped = source_path.rstrip("/") - if stripped.endswith((".parquet", ".pq")) or ( - not is_url(source_path) and Path(source_path).is_dir() - ): - return self._scan_parquet(source_path) - return self._scan_csv_tsv_gz(source_path) - - def _scan_parquet(self, source_path: str) -> dd.DataFrame: - """Scans a Parquet source and returns a Dask DataFrame. - - The source may be a single ``.parquet``/``.pq`` file, a glob pattern, - or a directory that is scanned recursively — which supports sharded - datasets such as MEDS, laid out as ``data//.parquet``. - - Unlike :meth:`_scan_csv_tsv_gz`, no all-string schema coercion is - applied: Parquet files embed their schema, so source dtypes (native - timestamps, numeric columns, nullable strings) are preserved and - handled downstream by :meth:`load_table`. - - Args: - source_path (str): Path to a Parquet file, directory, or glob. - - Returns: - dd.DataFrame: The Dask DataFrame backed by the Parquet source. - - Raises: - FileNotFoundError: If the source path does not exist, or if a - directory source contains no Parquet files. - """ - path = Path(source_path) - is_glob = any(ch in source_path for ch in "*?[") - if not is_glob: - if not path.exists(): - raise FileNotFoundError( - f"Parquet source does not exist: {source_path}" - ) - if path.is_dir() and not any( - itertools.chain(path.rglob("*.parquet"), path.rglob("*.pq")) - ): - raise FileNotFoundError( - f"Directory contains no Parquet files: {source_path}" - ) - return dd.read_parquet( - source_path, - split_row_groups=True, # type: ignore - blocksize="64MB", - ) - def _scan_csv_tsv_gz(self, source_path: str) -> dd.DataFrame: """Scans a CSV/TSV file (possibly gzipped) and returns a Dask DataFrame. @@ -571,30 +506,52 @@ def _event_transform(self, output_dir: Path) -> None: compute_ok = False try: df = self.load_data() - with DaskCluster( - n_workers=self.num_workers, - threads_per_worker=1, - processes=not in_notebook(), - # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory - local_directory=str(self.create_tmpdir()), - ) as cluster: - with DaskClient(cluster) as client: - if self.dev: - logger.info("Dev mode enabled: limiting to 1000 patients") - patients = df["patient_id"].unique().head(1000).tolist() - filter = df["patient_id"].isin(patients) - df = df[filter] - - logger.info(f"Caching event dataframe to {output_dir}...") - collection = df.sort_values("patient_id").to_parquet( - output_dir, - write_index=False, - compute=False, - ) - handle = client.compute(collection) - dask_progress(handle) - handle.result() # type: ignore - compute_ok = True # Data is fully written to disk + disable_distributed = os.environ.get( + "PYHEALTH_DISABLE_DASK_DISTRIBUTED", "0" + ) == "1" + + if disable_distributed: + logger.info( + "PYHEALTH_DISABLE_DASK_DISTRIBUTED=1 detected; using local dask scheduler." + ) + if self.dev: + n = 1000 if self.dev is True else int(self.dev) + logger.info(f"Dev mode enabled: limiting to {n} patients") + patients = df["patient_id"].unique().head(n, compute=True).tolist() + patient_filter = df["patient_id"].isin(patients) + df = df[patient_filter] + + logger.info(f"Caching event dataframe to {output_dir}...") + df.sort_values("patient_id").to_parquet( + output_dir, + write_index=False, + compute=True, + ) + else: + with DaskCluster( + n_workers=self.num_workers, + threads_per_worker=1, + processes=not in_notebook(), + # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory + local_directory=str(self.create_tmpdir()), + ) as cluster: + with DaskClient(cluster) as client: + if self.dev: + logger.info(f"Dev mode enabled: limiting to {1000 if self.dev is True else int(self.dev)} patients") + patients = df["patient_id"].unique().head(1000 if self.dev is True else int(self.dev)).tolist() + filter = df["patient_id"].isin(patients) + df = df[filter] + + logger.info(f"Caching event dataframe to {output_dir}...") + collection = df.sort_values("patient_id").to_parquet( + output_dir, + write_index=False, + compute=False, + ) + handle = client.compute(collection) + dask_progress(handle) + handle.result() # type: ignore + compute_ok = True # Data is fully written to disk except TimeoutError: if compute_ok: # Cluster shutdown timed out after successful compute — data is intact @@ -667,8 +624,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: Raises: ValueError: If the table is not found in the config. - FileNotFoundError: If the source file (CSV/TSV or Parquet) for the - table or join is not found. + FileNotFoundError: If the CSV file for the table or join is not found. """ assert self.config is not None, "Config must be provided to load tables" @@ -680,7 +636,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: csv_path = clean_path(csv_path) logger.info(f"Scanning table: {table_name} from {csv_path}") - df = self._scan_table(csv_path) + df = self._scan_csv_tsv_gz(csv_path) # Convert column names to lowercase before calling preprocess_func df = df.rename(columns=str.lower) @@ -699,7 +655,7 @@ def load_table(self, table_name: str) -> dd.DataFrame: other_csv_path = f"{self.root}/{join_cfg.file_path}" other_csv_path = clean_path(other_csv_path) logger.info(f"Joining with table: {other_csv_path}") - join_df = self._scan_table(other_csv_path) + join_df = self._scan_csv_tsv_gz(other_csv_path) join_df = join_df.rename(columns=str.lower) join_key = join_cfg.on columns = join_cfg.columns @@ -723,21 +679,14 @@ def load_table(self, table_name: str) -> dd.DataFrame: timestamp_series: dd.Series = functools.reduce( operator.add, (df[col].astype("string") for col in timestamp_col) ) - timestamp_series = dd.to_datetime( - timestamp_series, - format=timestamp_format, - errors="raise", - ) - elif pd.api.types.is_datetime64_any_dtype(df[timestamp_col].dtype): - # Typed sources (e.g. Parquet) already carry native timestamps: - # skip the string round-trip and only normalize the unit below. - timestamp_series: dd.Series = df[timestamp_col] else: - timestamp_series = dd.to_datetime( - df[timestamp_col].astype("string"), - format=timestamp_format, - errors="raise", - ) + timestamp_series: dd.Series = df[timestamp_col].astype("string") + + timestamp_series: dd.Series = dd.to_datetime( + timestamp_series, + format=timestamp_format, + errors="raise", + ) df: dd.DataFrame = df.assign( timestamp=timestamp_series.astype("datetime64[ms]") ) @@ -1165,4 +1114,4 @@ def _main_guard(self, func_name: str): f"{func_name} method accessed from a non-main process. This may lead to unexpected behavior.\n" + "Consider use __name__ == '__main__' guard when using multiprocessing." ) - exit(1) + exit(1) \ No newline at end of file diff --git a/pyhealth/models/embedding/__init__.py b/pyhealth/models/embedding/__init__.py new file mode 100644 index 000000000..2bc66d73b --- /dev/null +++ b/pyhealth/models/embedding/__init__.py @@ -0,0 +1,33 @@ +"""Embedding models for PyHealth multimodal pipelines. + +All embedding models share the :class:`BaseEmbeddingModel` interface: +they expose an ``embedding_dim`` property and a ``forward`` method that +transforms processor output tensors into dense vector embeddings. + +Available models: + +- :class:`EmbeddingModel` — generic encoder for codes, sequences, timeseries +- :class:`VisionEmbeddingModel` — ViT-style patch encoder for medical images (Josh) +- :class:`UnifiedMultimodalEmbeddingModel` — temporally-aligned multi-modal encoder + +Helper utilities: + +- :class:`SinusoidalTimeEmbedding` — continuous time positional encoding +- :func:`init_embedding_with_pretrained` — load GloVe-style pretrained vectors +""" + +from .base import BaseEmbeddingModel +from .vanilla import EmbeddingModel, init_embedding_with_pretrained +from .vision import VisionEmbeddingModel, PatchEmbedding, Permute +from .unified import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding + +__all__ = [ + "BaseEmbeddingModel", + "EmbeddingModel", + "VisionEmbeddingModel", + "PatchEmbedding", + "Permute", + "UnifiedMultimodalEmbeddingModel", + "SinusoidalTimeEmbedding", + "init_embedding_with_pretrained", +] \ No newline at end of file From b80f7efa29e1208715e6cba317927b359a2ceee6 Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 04:42:10 -0500 Subject: [PATCH 05/27] New Updates --- pyhealth/models/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py index 2f30ae673..7464f7a90 100644 --- a/pyhealth/models/__init__.py +++ b/pyhealth/models/__init__.py @@ -46,7 +46,7 @@ from .text_embedding import TextEmbedding from .sdoh import SdohClassifier from .medlink import MedLink -from .unified_embedding import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding +from .embedding import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding from .califorest import CaliForest from .generators.halo import HALO from .generators.gpt2 import GPT2 From f0187f929c7c5cce5b5861eb6d3a68aec9523abe Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 05:26:38 -0500 Subject: [PATCH 06/27] New Updates --- .../datasets/configs/mimic4_cxr_sunlab.yaml | 105 ++++++++++ pyhealth/datasets/mimic4.py | 138 ++++++++++-- .../labs_notes_cxr/labs_notes_cxr_rnn.sub | 61 ++++++ .../labs_notes_cxr/run_labs_notes_cxr_rnn.sh | 197 ++++++++++++++++++ 4 files changed, 489 insertions(+), 12 deletions(-) create mode 100644 pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml create mode 100644 scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub create mode 100755 scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh diff --git a/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml new file mode 100644 index 000000000..e631de4aa --- /dev/null +++ b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml @@ -0,0 +1,105 @@ +version: "2.1.0" +tables: + metadata: + file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + patient_id: "subject_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "image_path" + - "dicom_id" + - "study_id" + - "performedprocedurestepdescription" + - "viewposition" + - "rows" + - "columns" + - "procedurecodesequence_codemeaning" + - "viewcodesequence_codemeaning" + - "patientorientationcodesequence_codemeaning" + + chexpert: + file_path: "mimic-cxr-2.0.0-chexpert.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "study_id" + how: "inner" + columns: + - "studydate" + - "studytime" + - "dicom_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "atelectasis" + - "cardiomegaly" + - "consolidation" + - "edema" + - "enlarged cardiomediastinum" + - "fracture" + - "lung lesion" + - "lung opacity" + - "no finding" + - "pleural effusion" + - "pleural other" + - "pneumonia" + - "pneumothorax" + - "support devices" + + negbio: + file_path: "mimic-cxr-2.0.0-negbio.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "study_id" + how: "inner" + columns: + - "studydate" + - "studytime" + - "dicom_id" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "atelectasis" + - "cardiomegaly" + - "consolidation" + - "edema" + - "enlarged cardiomediastinum" + - "fracture" + - "lung lesion" + - "lung opacity" + - "no finding" + - "pleural effusion" + - "pleural other" + - "pneumonia" + - "pneumothorax" + - "support devices" + + split: + file_path: "mimic-cxr-2.0.0-split.csv" + patient_id: "subject_id" + join: + - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + "on": "dicom_id" + how: "inner" + columns: + - "studydate" + - "studytime" + timestamp: + - "studydate" + - "studytime" + timestamp_format: "%Y%m%d%H%M%S" + attributes: + - "dicom_id" + - "study_id" + - "split" \ No newline at end of file diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index 9d1aa55d8..f9a7da38c 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -1,7 +1,7 @@ import logging import os import warnings -from typing import List, Optional +from typing import List, Optional, Union import pandas as pd import dask.dataframe as dd @@ -223,6 +223,102 @@ def process_image_path(x): return +class MIMIC4CXRSunlabDataset(BaseDataset): + """ + Sunlab variant of the MIMIC-CXR Chest X-ray dataset. + + This variant uses the existing metadata CSV and derives flattened image + paths at ``images/{dicom_id}.jpg``. + """ + + def __init__( + self, + root: str, + tables: List[str], + dataset_name: str = "mimic4_cxr_sunlab", + config_path: Optional[str] = None, + cache_dir: Optional[str] = None, + **kwargs, + ): + if config_path is None: + config_path = os.path.join( + os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml" + ) + logger.info(f"Using default Sunlab CXR config: {config_path}") + self.prepare_metadata(root) + log_memory_usage(f"Before initializing {dataset_name}") + super().__init__( + root=root, + tables=tables, + dataset_name=dataset_name, + config_path=config_path, + cache_dir=cache_dir, + **kwargs, + ) + log_memory_usage(f"After initializing {dataset_name}") + + @staticmethod + def _resolve_column_name(columns: List[str], target: str) -> str: + lower_to_original = {col.lower(): col for col in columns} + resolved = lower_to_original.get(target.lower()) + if resolved is None: + raise ValueError( + f"Expected column '{target}' in metadata, available columns: {columns}" + ) + return resolved + + def prepare_metadata(self, root: str) -> None: + metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv") + if not os.path.exists(metadata_path): + raise FileNotFoundError( + f"Sunlab metadata file not found: {metadata_path}. " + "Expected existing metadata linked by dicom_id/subject_id/study_id." + ) + + images_dir = os.path.join(root, "images") + if not os.path.isdir(images_dir): + raise FileNotFoundError( + f"Sunlab images directory not found: {images_dir}. " + "Expected flattened image files at images/{dicom_id}.jpg." + ) + + metadata = pd.read_csv(metadata_path, dtype=str) + + dicom_col = self._resolve_column_name(metadata.columns.tolist(), "dicom_id") + study_time_col = self._resolve_column_name( + metadata.columns.tolist(), "studytime" + ) + + # Normalize StudyTime so timestamps parse with %Y%m%d%H%M%S in config. + def normalize_studytime(value: Optional[str]) -> str: + if value is None: + return "000000" + value_str = str(value).strip() + if value_str == "" or value_str.lower() == "nan": + return "000000" + try: + return f"{int(float(value_str)):06d}" + except Exception: + digits = "".join(ch for ch in value_str if ch.isdigit()) + if digits == "": + return "000000" + return digits[:6].zfill(6) + + metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime) + + metadata["image_path"] = metadata[dicom_col].apply( + lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg") + ) + + # Align with existing config conventions by using lowercase headers. + metadata.columns = [col.lower() for col in metadata.columns] + + metadata.to_csv( + os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"), + index=False, + ) + + class MIMIC4Dataset(BaseDataset): """ Unified MIMIC-IV dataset with support for EHR, clinical notes, and X-rays. @@ -242,6 +338,7 @@ class MIMIC4Dataset(BaseDataset): ehr_config_path: Path to the EHR config file note_config_path: Path to the note config file cxr_config_path: Path to the CXR config file + cxr_variant: Which CXR variant to load ("default" or "sunlab") dataset_name: Name of the dataset dev: Whether to enable dev mode (limit to 1000 patients) @@ -279,8 +376,9 @@ def __init__( ehr_config_path: Optional[str] = None, note_config_path: Optional[str] = None, cxr_config_path: Optional[str] = None, + cxr_variant: str = "default", dataset_name: str = "mimic4", - dev: bool = False, + dev: Union[bool, int] = False, cache_dir: Optional[str] = None, num_workers: int = 1, ): @@ -340,17 +438,33 @@ def __init__( # Initialize CXR dataset if root is provided if cxr_root is not None: + if cxr_variant not in {"default", "sunlab"}: + raise ValueError( + f"Unknown cxr_variant '{cxr_variant}'. " + "Expected one of {'default', 'sunlab'}." + ) + logger.info( - f"Initializing MIMIC4CXRDataset with tables: {cxr_tables} (dev mode: {dev})" - ) - self.sub_datasets["cxr"] = MIMIC4CXRDataset( - root=cxr_root, - tables=cxr_tables, - config_path=cxr_config_path, - cache_dir=str(self.cache_dir), - dev=dev, - num_workers=num_workers, + f"Initializing MIMIC4 CXR variant '{cxr_variant}' with tables: {cxr_tables} (dev mode: {dev})" ) + if cxr_variant == "sunlab": + self.sub_datasets["cxr"] = MIMIC4CXRSunlabDataset( + root=cxr_root, + tables=cxr_tables, + config_path=cxr_config_path, + cache_dir=str(self.cache_dir), + dev=dev, + num_workers=num_workers, + ) + else: + self.sub_datasets["cxr"] = MIMIC4CXRDataset( + root=cxr_root, + tables=cxr_tables, + config_path=cxr_config_path, + cache_dir=str(self.cache_dir), + dev=dev, + num_workers=num_workers, + ) log_memory_usage("After CXR dataset initialization") log_memory_usage("Completed MIMIC4Dataset init") @@ -374,4 +488,4 @@ def load_data(self) -> dd.DataFrame: if len(frames) == 1: return frames[0] else: - return dd.concat(frames, axis=0, join="outer") + return dd.concat(frames, axis=0, join="outer") \ No newline at end of file diff --git a/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub new file mode 100644 index 000000000..118fa8149 --- /dev/null +++ b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub @@ -0,0 +1,61 @@ +# HTCondor submission — labs+notes+CXR RNN mortality run +# +# Condor equivalent of +# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py: +# same task (notes_labs_cxr), same model (rnn), same hyperparameters. Runs +# unattended instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_notes_cxr_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \ + NOTE_ROOT=/shared/rsaas/physionet.org/files/mimic-note \ + CXR_ROOT=/shared/rsaas/physionet.org/files/MIMIC-CXR \ + CXR_VARIANT=sunlab \ + CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-notes-cxr \ + WANDB_RUN_NAME=labs_notes_cxr_rnn_seed$(seed) \ + FREEZE_ENCODER=1" + +output = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 8 +# Starting point copied from labs_notes_rnn.sub (400000MB was sized for the +# full-scale patient_id sort/shuffle there). Untested for this variant: the +# added CXR branch means more dataloader workers decoding JPEGs plus an +# image-metadata join during caching, so this may need to be raised further +# if the job gets cgroup-killed. +request_memory = 400000MB +request_disk = 20GB + +# Same reasoning as labs_notes_rnn.sub: FREEZE_ENCODER=1 (frozen +# Bio_ClinicalBERT text encoder) keeps VRAM down for the text branch, but the +# CXR image encoder is unfrozen here and adds its own footprint on top, so +# the GPU memory floor is a conservative starting guess pending a real OOM +# data point on this variant. +Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000) +Rank = TARGET.GPUs_GlobalMemoryMb + +queue seed from ( + 12 +) diff --git a/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh new file mode 100755 index 000000000..4f40d8801 --- /dev/null +++ b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh @@ -0,0 +1,197 @@ +#!/usr/bin/env bash +# HTCondor executable — labs+notes+CXR RNN mortality run. +# +# Condor equivalent of +# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py: +# same task (notes_labs_cxr), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_notes_cxr_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr/* +set -euo pipefail + +SEED="${1:?usage: run_labs_notes_cxr_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}" +CXR_ROOT="${CXR_ROOT:-/shared/rsaas/physionet.org/files/MIMIC-CXR}" +CXR_VARIANT="${CXR_VARIANT:-default}" +CACHE_DIR="${CACHE_DIR:-/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-1}" +EMBEDDING_DIM="${EMBEDDING_DIM:-128}" +HIDDEN_DIM="${HIDDEN_DIM:-128}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-2}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" +FREEZE_ENCODER="${FREEZE_ENCODER:-0}" +INCLUDE_VITALS="${INCLUDE_VITALS:-0}" +USE_AMP="${USE_AMP:-0}" +AMP_DTYPE="${AMP_DTYPE:-bf16}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1) +# also throws away its disk-spilling/memory limits, which OOM'd the full-scale +# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup +# limit hit). Scope the fix to just the NVML probe instead, keeping the real +# distributed cluster (with spilling) for the sort. +export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes-cxr}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_notes_cxr_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs+notes+CXR RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Note root : ${NOTE_ROOT}" +echo " CXR root : ${CXR_ROOT} (variant=${CXR_VARIANT})" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir : ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo " Use AMP : ${USE_AMP} (dtype=${AMP_DTYPE})" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --ehr-root "${EHR_ROOT}" + --note-root "${NOTE_ROOT}" + --cxr-root "${CXR_ROOT}" + --cxr-variant "${CXR_VARIANT}" + --cache-dir "${CACHE_DIR}" + --task notes_labs_cxr + --model rnn + --embedding-dim "${EMBEDDING_DIM}" + --hidden-dim "${HIDDEN_DIM}" + --rnn-type "${RNN_TYPE}" + --rnn-layers "${RNN_LAYERS}" + --dropout "${DROPOUT}" + --epochs "${EPOCHS}" + --batch-size "${BATCH_SIZE}" + --lr "${LR}" + --weight-decay "${WEIGHT_DECAY}" + --patience "${PATIENCE}" + --num-workers "${NUM_WORKERS}" + --seed "${SEED}" + --output-dir "${OUTPUT_DIR}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${INCLUDE_VITALS}" == "1" ]]; then + COMMON+=(--include-vitals) +fi + +if [[ "${USE_AMP}" == "1" ]]; then + COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}") +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" From b91ec9eaa59d9325f58b3f832c013b8d7343fb70 Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 05:46:58 -0500 Subject: [PATCH 07/27] New Updates --- .../will/condor/labs_notes/lab_notes_rnn.sub | 58 ++++++ .../condor/labs_notes/run_labs_notes_rnn.sh | 191 ++++++++++++++++++ 2 files changed, 249 insertions(+) create mode 100644 scripts/will/condor/labs_notes/lab_notes_rnn.sub create mode 100755 scripts/will/condor/labs_notes/run_labs_notes_rnn.sh diff --git a/scripts/will/condor/labs_notes/lab_notes_rnn.sub b/scripts/will/condor/labs_notes/lab_notes_rnn.sub new file mode 100644 index 000000000..a5ede861e --- /dev/null +++ b/scripts/will/condor/labs_notes/lab_notes_rnn.sub @@ -0,0 +1,58 @@ +# HTCondor submission — labs+notes RNN mortality run +# +# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py: +# same task (notes_labs), same model (rnn), same hyperparameters. Runs unattended +# instead of in a tmux session; Condor assigns the GPU (no manual +# nvidia-smi / CUDA_VISIBLE_DEVICES step needed). +# +# To submit (from the project root): +# mkdir -p /home/wp14/logs/condor +# condor_submit scripts/will/condor/labs_notes/labs_notes_rnn.sub +# +# Monitor: +# condor_q +# tail -f /home/wp14/logs/condor/labs_notes_rnn__0.out + +initialdir = /home/wp14/PyHealth +executable = /home/wp14/PyHealth/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh +transfer_executable = False +arguments = $(seed) +getenv = True + +environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \ + NOTE_ROOT=/shared/rsaas/physionet.org/files/mimic-note \ + CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes \ + OUTPUT_DIR=/home/wp14/output \ + USE_WANDB=1 \ + WANDB_PROJECT=pyhealth-multimodal-labs-notes \ + WANDB_RUN_NAME=labs_notes_rnn_seed$(seed)" + +output = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).out +error = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).err +log = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).log + +stream_output = False +stream_error = False + +request_gpus = 1 +request_cpus = 4 +# Was 163840MB (160GB) — cgroup-killed job 10821 at 162747MB during the +# full-scale patient_id sort/shuffle (see run_labs_notes_rnn.sh comment). +# Bumped for headroom now that the distributed cluster (with disk-spilling) +# is back in play; c02 has ~1TB total and is otherwise idle. +request_memory = 400000MB +request_disk = 20GB + +# Previously hardcoded to sunlab-c01 (A100 80GB) because the previous run +# OOM'd on a 47GB card with ~46.8GB resident. sunlab-c01's condor_startd is +# currently down (master alive, STARTD_StartTime=0), so it never matches — +# the only machine in the pool is sunlab-c02 (8x RTX 6000 Ada, 48509MB each). +# Match any GPU with enough headroom and prefer the biggest; FREEZE_ENCODER=1 +# below (frozen Bio_ClinicalBERT text encoder, ~50% less VRAM for the text +# branch) is what actually keeps this under 48GB instead of the hostname pin. +Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000) +Rank = TARGET.GPUs_GlobalMemoryMb + +queue seed from ( + 12 +) \ No newline at end of file diff --git a/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh new file mode 100755 index 000000000..0e26d3fc1 --- /dev/null +++ b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh @@ -0,0 +1,191 @@ +#!/usr/bin/env bash +# HTCondor executable — labs+notes RNN mortality run. +# +# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py: +# same task (notes_labs), same model (rnn), same hyperparameter defaults. GPU +# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor +# assigns the GPU via request_gpus / cgroups. +# +# usage: run_labs_notes_rnn.sh +# to remove logs: rm -rf logs/condor/* +# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes/* +set -euo pipefail + +SEED="${1:?usage: run_labs_notes_rnn.sh }" + +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}" +EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}" +NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}" +CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs_notes}" +OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}" +CONDA_SH="${CONDA_SH:-}" + +DEV_MODE="${DEV_MODE:-1}" +EMBEDDING_DIM="${EMBEDDING_DIM:-128}" +HIDDEN_DIM="${HIDDEN_DIM:-128}" +RNN_TYPE="${RNN_TYPE:-GRU}" +RNN_LAYERS="${RNN_LAYERS:-2}" +DROPOUT="${DROPOUT:-0.1}" +EPOCHS="${EPOCHS:-15}" +BATCH_SIZE="${BATCH_SIZE:-32}" +LR="${LR:-1e-3}" +WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" +PATIENCE="${PATIENCE:-5}" +NUM_WORKERS="${NUM_WORKERS:-4}" +FREEZE_ENCODER="${FREEZE_ENCODER:-0}" +INCLUDE_VITALS="${INCLUDE_VITALS:-0}" +USE_AMP="${USE_AMP:-0}" +AMP_DTYPE="${AMP_DTYPE:-bf16}" + +# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that +# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup. +# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1) +# also throws away its disk-spilling/memory limits, which OOM'd the full-scale +# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup +# limit hit). Scope the fix to just the NVML probe instead, keeping the real +# distributed cluster (with spilling) for the sort. +export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}" + +USE_WANDB="${USE_WANDB:-0}" +WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes}" +WANDB_RUN_NAME="${WANDB_RUN_NAME:-}" + +resolve_conda_sh() { + if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then + echo "${CONDA_SH}" + return 0 + fi + if command -v conda >/dev/null 2>&1; then + local base + base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then + echo "${base}/etc/profile.d/conda.sh" + return 0 + fi + fi + if [ -f /etc/profile.d/modules.sh ]; then + source /etc/profile.d/modules.sh >/dev/null 2>&1 || true + if command -v module >/dev/null 2>&1; then + module load miniconda3 >/dev/null 2>&1 || true + module load anaconda3 >/dev/null 2>&1 || true + if command -v conda >/dev/null 2>&1; then + local mod_base + mod_base="$(conda info --base 2>/dev/null || true)" + if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then + echo "${mod_base}/etc/profile.d/conda.sh" + return 0 + fi + fi + fi + fi + local user_name home_dir + user_name="${USER:-$(id -un 2>/dev/null || true)}" + home_dir="${HOME:-/home/${user_name}}" + local candidates=( + "${home_dir}/miniconda3/etc/profile.d/conda.sh" + "/home/${user_name}/miniconda3/etc/profile.d/conda.sh" + "${home_dir}/anaconda3/etc/profile.d/conda.sh" + "/home/${user_name}/anaconda3/etc/profile.d/conda.sh" + "/opt/miniconda3/etc/profile.d/conda.sh" + "/opt/anaconda3/etc/profile.d/conda.sh" + "/opt/conda/etc/profile.d/conda.sh" + ) + local c + for c in "${candidates[@]}"; do + if [[ -f "${c}" ]]; then + echo "${c}" + return 0 + fi + done + local found="" + found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)" + if [[ -n "${found}" && -f "${found}" ]]; then + echo "${found}" + return 0 + fi + return 1 +} + +CONDA_SH="$(resolve_conda_sh || true)" +if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then + echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2 + exit 1 +fi +source "${CONDA_SH}" +eval "$(conda shell.bash hook)" +conda activate "${CONDA_ENV}" + +cd "${PROJECT_DIR}" +export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}" + +JOB_TAG="rnn_labs_notes_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}" + +echo "========================================================" +echo " Labs+notes RNN run | ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo " Conda env : ${CONDA_ENV}" +echo " EHR root : ${EHR_ROOT}" +echo " Note root : ${NOTE_ROOT}" +echo " Cache dir : ${CACHE_DIR}" +echo " Output dir: ${OUTPUT_DIR}" +echo " Seed : ${SEED}" +echo " Dev mode : ${DEV_MODE}" +echo " Use AMP : ${USE_AMP} (dtype=${AMP_DTYPE})" +echo "========================================================" + +if ! python -c "import pyhealth" >/dev/null 2>&1; then + echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2 + exit 1 +fi + +COMMON=( + --ehr-root "${EHR_ROOT}" + --note-root "${NOTE_ROOT}" + --cache-dir "${CACHE_DIR}" + --task notes_labs + --model rnn + --embedding-dim "${EMBEDDING_DIM}" + --hidden-dim "${HIDDEN_DIM}" + --rnn-type "${RNN_TYPE}" + --rnn-layers "${RNN_LAYERS}" + --dropout "${DROPOUT}" + --epochs "${EPOCHS}" + --batch-size "${BATCH_SIZE}" + --lr "${LR}" + --weight-decay "${WEIGHT_DECAY}" + --patience "${PATIENCE}" + --num-workers "${NUM_WORKERS}" + --seed "${SEED}" + --output-dir "${OUTPUT_DIR}" +) + +if [[ "${DEV_MODE}" == "1" ]]; then + COMMON+=(--dev) +fi + +if [[ "${FREEZE_ENCODER}" == "1" ]]; then + COMMON+=(--freeze-encoder) +fi + +if [[ "${INCLUDE_VITALS}" == "1" ]]; then + COMMON+=(--include-vitals) +fi + +if [[ "${USE_AMP}" == "1" ]]; then + COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}") +fi + +if [[ "${USE_WANDB}" == "1" ]]; then + COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}") + if [[ -n "${WANDB_RUN_NAME}" ]]; then + COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}") + fi +fi + +python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" + +echo "========================================================" +echo " Completed ${JOB_TAG}" +echo " $(date -u '+%Y-%m-%d %H:%M:%S UTC')" +echo "========================================================" From 2b7b2c5d4fc0ef0d7fe19032aee0a8cc99a9a495 Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 11 Aug 2026 16:24:07 -0500 Subject: [PATCH 08/27] New Updates --- pyhealth/data/data.py | 2 +- pyhealth/datasets/utils.py | 66 +++++++++---------- .../condor/labs_notes/run_labs_notes_rnn.sh | 2 +- .../labs_notes_cxr/run_labs_notes_cxr_rnn.sh | 2 +- 4 files changed, 35 insertions(+), 37 deletions(-) diff --git a/pyhealth/data/data.py b/pyhealth/data/data.py index 14b1b526c..39dd1206f 100644 --- a/pyhealth/data/data.py +++ b/pyhealth/data/data.py @@ -155,7 +155,7 @@ def _filter_by_time_range_fast(self, df: pl.DataFrame, start: Optional[datetime] start_idx = np.searchsorted(ts_col, np.datetime64(start, "ms"), side="left") if end is not None: end_idx = np.searchsorted(ts_col, np.datetime64(end, "ms"), side="right") - return df.slice(start_idx, end_idx - start_idx) + return df.slice(start_idx, max(0, end_idx - start_idx)) def _filter_by_event_type_regular(self, df: pl.DataFrame, event_type: Optional[str]) -> pl.DataFrame: """Regular filtering by event type. Time complexity: O(n).""" diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 24c87a1d5..0125ab4f9 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -15,9 +15,10 @@ MODULE_CACHE_PATH = os.path.join(BASE_CACHE_PATH, "datasets") create_directory(MODULE_CACHE_PATH) -#PyG import for graph-based models +# PyG import for graph-based models try: from torch_geometric.data import Data as PyGData, Batch as PyGBatch + HAS_PYG = True except ImportError: HAS_PYG = False @@ -267,40 +268,37 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: for key in keys: values = [sample[key] for sample in batch] - # Check if this is a temporal feature tuple (time, values) - if isinstance(values[0], tuple) and len(values[0]) == 2: - # Handle (time, values) tuples from processors - time_tensors = [v[0] for v in values] - value_tensors = [v[1] for v in values] - - # Collate values - if value_tensors[0].dim() == 0: - # Scalars - collated_values = torch.stack(value_tensors) - elif all(v.shape == value_tensors[0].shape for v in value_tensors): - # All same shape - collated_values = torch.stack(value_tensors) - else: - # Variable shapes, use pad_sequence - collated_values = pad_sequence( - value_tensors, batch_first=True, padding_value=0 - ) - - # Collate times (if present) - collated_times = None - # Check if ALL samples have time (not just some) - if all(t is not None for t in time_tensors): - time_tensors_all = [t for t in time_tensors if t is not None] - if all(t.shape == time_tensors_all[0].shape for t in time_tensors_all): - collated_times = torch.stack(time_tensors_all) + if isinstance(values[0], tuple): + # Generic tuple collation for processor outputs, e.g. + # - (time, value) from StageNet processors + # - (value, mask, token_type_ids, time, type_tag) + # from TupleTimeTextProcessor with tokenizer. + transposed = list(zip(*values)) + collated_elems = [] + + for elem_vals in transposed: + first = elem_vals[0] + + if first is None and all(v is None for v in elem_vals): + collated_elems.append(None) + elif isinstance(first, torch.Tensor): + tensor_vals = list(elem_vals) + if all(v.shape == tensor_vals[0].shape for v in tensor_vals): + collated_elems.append(torch.stack(tensor_vals)) + else: + collated_elems.append( + pad_sequence( + tensor_vals, + batch_first=True, + padding_value=0, + ) + ) else: - collated_times = pad_sequence( - time_tensors_all, batch_first=True, padding_value=0 - ) + collated_elems.append(list(elem_vals)) + + collated[key] = tuple(collated_elems) - # Return as tuple (time, values) - collated[key] = (collated_times, collated_values) - # PyG Data objects (graph processor output) + # PyG Data objects (graph processor output) elif HAS_PYG and isinstance(values[0], PyGData): collated[key] = PyGBatch.from_data_list(values) @@ -453,4 +451,4 @@ def load_processors(processor_dir: str) -> Tuple[Dict, Dict]: print(list_nested_levels([[1, [2], [[3]]]])) print(is_homo_list([1, 2, 3])) print(is_homo_list([1, 2, [3]])) - print(is_homo_list([1, 2.0])) + print(is_homo_list([1, 2.0])) \ No newline at end of file diff --git a/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh index 0e26d3fc1..33661ff11 100755 --- a/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh +++ b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh @@ -33,7 +33,7 @@ LR="${LR:-1e-3}" WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" PATIENCE="${PATIENCE:-5}" NUM_WORKERS="${NUM_WORKERS:-4}" -FREEZE_ENCODER="${FREEZE_ENCODER:-0}" +FREEZE_ENCODER="${FREEZE_ENCODER:-1}" INCLUDE_VITALS="${INCLUDE_VITALS:-0}" USE_AMP="${USE_AMP:-0}" AMP_DTYPE="${AMP_DTYPE:-bf16}" diff --git a/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh index 4f40d8801..aa9037ae7 100755 --- a/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh +++ b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh @@ -36,7 +36,7 @@ LR="${LR:-1e-3}" WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}" PATIENCE="${PATIENCE:-5}" NUM_WORKERS="${NUM_WORKERS:-4}" -FREEZE_ENCODER="${FREEZE_ENCODER:-0}" +FREEZE_ENCODER="${FREEZE_ENCODER:-1}" INCLUDE_VITALS="${INCLUDE_VITALS:-0}" USE_AMP="${USE_AMP:-0}" AMP_DTYPE="${AMP_DTYPE:-bf16}" From 9782aca82bf139707dc6b5a2eb9dc4ac019c27d8 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 07:20:22 -0400 Subject: [PATCH 09/27] Stop emitting fake missing-event placeholders. Empty notes, labs, CXR, and ICD are now zero-length sequences instead of a constant [MISSING_TEXT] row, black image, or pad visit. The tokenizer crash on empty batches is handled by constructing empty tensors, so BERT cannot treat note presence as a free mortality feature. Co-authored-by: Cursor --- pyhealth/processors/stagenet_processor.py | 40 +++-- pyhealth/processors/time_image_processor.py | 15 +- .../processors/tuple_time_text_processor.py | 24 +-- pyhealth/tasks/multimodal_mimic4.py | 149 ++++++------------ tests/core/test_stagenet_processor.py | 11 +- tests/test_tuple_time_text_processor.py | 6 + 6 files changed, 109 insertions(+), 136 deletions(-) diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py index 604376ec1..557fc4a7b 100644 --- a/pyhealth/processors/stagenet_processor.py +++ b/pyhealth/processors/stagenet_processor.py @@ -86,7 +86,7 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: if len(first_elem) > 0 and isinstance(first_elem[0], str): # Case 2: [["A", "B"], ["C"], ...] self._is_nested = True - break + break # Build vocabulary for codes and find max nested length max_inner_len = 0 @@ -178,9 +178,9 @@ def process( def _encode_codes(self, codes: List[str]) -> torch.Tensor: """Encode flat code list to indices.""" - # Handle empty code list - return single padding token + # Handle empty code list — zero events, not a fake pad token. if len(codes) == 0: - return torch.tensor([self.code_vocab[""]], dtype=torch.long) + return torch.zeros((0,), dtype=torch.long) indices = [] for code in codes: @@ -198,10 +198,9 @@ def _encode_nested_codes(self, nested_codes: List[List[str]]) -> torch.Tensor: assert self._max_nested_len is not None, "Max nested length must be set during fit()" # Handle empty nested codes (no visits/events) - # Return single padding token with shape (1, max_len) if len(nested_codes) == 0: - pad_token = self.code_vocab[""] - return torch.tensor([[pad_token] * self._max_nested_len], dtype=torch.long) + max_len = self._max_nested_len if self._max_nested_len is not None else 1 + return torch.zeros((0, max_len), dtype=torch.long) encoded_sequences = [] # Use global max length determined during fit @@ -345,9 +344,10 @@ class StageNetTensorProcessor(TemporalFeatureProcessor): >>> time.shape # (3,) """ - def __init__(self): + def __init__(self, forward_fill: bool = True): self._size = None # Feature dimension (set during fit) self._is_nested = None + self.forward_fill = forward_fill def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: """Determine input structure. @@ -370,13 +370,14 @@ def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None: # Flat numeric: [1.5, 2.0, ...] self._is_nested = False self._size = 1 + break elif isinstance(first_elem, list): if len(first_elem) > 0: if isinstance(first_elem[0], (int, float)): # Nested numerics: [[1.0, 2.0], [3.0, 4.0]] self._is_nested = True self._size = len(first_elem) - break + break def process( self, value: Tuple[Optional[List], List] @@ -395,14 +396,31 @@ def process( """ # Unpack tuple: (time, values) time_data, value_data = value + value_data = list(value_data or []) - # Convert to numpy for easier imputation handling import numpy as np + if len(value_data) == 0: + n_feat = self._size if self._size is not None else 1 + nested = True if self._is_nested is None else self._is_nested + if nested: + value_tensor = torch.zeros((0, n_feat), dtype=torch.float) + else: + value_tensor = torch.zeros((0,), dtype=torch.float) + time_tensor = ( + torch.zeros((0,), dtype=torch.float) if time_data is not None else None + ) + return time_tensor, value_tensor + + # Convert to numpy for easier imputation handling value_array = np.array(value_data, dtype=float) - # Apply forward-fill imputation - if value_array.ndim == 1: + # Observation-mask fields must preserve false after a true observation; + # forward-filling a 0/1 mask would turn later missing labs into observed + # values. Ordinary numeric time-series retain the historical behaviour. + if not self.forward_fill: + value_array = np.nan_to_num(value_array, nan=0.0, posinf=0.0, neginf=0.0) + elif value_array.ndim == 1: # Flat numeric: [1.5, 2.0, nan, 3.0, ...] last_value = 0.0 for i in range(len(value_array)): diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index 205782c7f..449fff729 100644 --- a/pyhealth/processors/time_image_processor.py +++ b/pyhealth/processors/time_image_processor.py @@ -279,7 +279,6 @@ def process( Raises: ValueError: If image_paths and time_diffs have different lengths. - ValueError: If image_paths is empty. FileNotFoundError: If any image file does not exist. """ image_paths, time_diffs = value @@ -291,7 +290,19 @@ def process( f"match." ) if len(image_paths) == 0: - raise ValueError("image_paths must be non-empty.") + if self.n_channels is not None: + c = self.n_channels + elif self.mode == "L": + c = 1 + elif self.mode == "RGBA": + c = 4 + else: + c = 3 + images = torch.zeros( + (0, c, self.image_size, self.image_size), dtype=torch.float32 + ) + timestamps = torch.zeros((0,), dtype=torch.float32) + return images, timestamps, "image" paired = sorted(zip(time_diffs, image_paths), key=lambda x: x[0]) diff --git a/pyhealth/processors/tuple_time_text_processor.py b/pyhealth/processors/tuple_time_text_processor.py index cb1fe99a0..7f1fe6a6f 100644 --- a/pyhealth/processors/tuple_time_text_processor.py +++ b/pyhealth/processors/tuple_time_text_processor.py @@ -5,7 +5,6 @@ from . import register_processor logger = logging.getLogger(__name__) -_MISSING_TEXT_TOKEN = "[MISSING_TEXT]" @register_processor("tuple_time_text") class TupleTimeTextProcessor(TemporalFeatureProcessor): @@ -22,7 +21,7 @@ def __init__( self, type_tag: str = "note", tokenizer_model: Optional[str] = None, - max_length: int = 128, + max_length: int = 512, padding: bool = True, truncation: bool = True, ): @@ -32,8 +31,9 @@ def __init__( type_tag: Modality identifier for automatic routing. Default: "note" tokenizer_model: Name or path of the HuggingFace tokenizer to use. If None, texts are returned as raw strings. Default: None - max_length: Maximum sequence length for tokenization. Default: 128 - padding: Whether to pad sequences to max_length. Default: True + max_length: Maximum sequence length for tokenization. Default: 512 + padding: Whether to pad sequences to the longest note in the sample. + Default: True truncation: Whether to truncate sequences to max_length. Default: True """ super().__init__() @@ -109,21 +109,21 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str] cleaned_texts.append(text) cleaned_times.append(t) - # Fast tokenizer path crashes on empty batches; force a single - # missingness token when all notes are empty/malformed. - if len(cleaned_texts) == 0: - cleaned_texts = [_MISSING_TEXT_TOKEN] - cleaned_times = [0.0] - texts = cleaned_texts time_diffs = cleaned_times time_tensor = torch.tensor(time_diffs, dtype=torch.float32) if self.tokenizer is not None: - # Tokenize the list of texts + # Fast tokenizers crash on tokenizer([]). Build empty tensors + # ourselves so a patient with no notes is zero events, not a + # fake "[MISSING_TEXT]" row whose BERT embedding is a constant + # the classifier can use as a mortality feature. + if len(texts) == 0: + empty = torch.zeros((0, 1), dtype=torch.long) + return empty, empty.clone(), empty.clone(), time_tensor, self.type_tag encoded = self.tokenizer( texts, - padding="max_length" if self.padding else False, + padding="longest" if self.padding else False, truncation=self.truncation, max_length=self.max_length, return_tensors="pt" diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index c3b1308e9..768c75819 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -83,6 +83,9 @@ def __init__( window_hours: Optional[float] = None, ): self.window_hours = window_hours + # Task cache key is uuid5 over {**vars(task), schemas}. Bump when + # emitted data changes so leaky caches cannot be reused. + self.emitted_data_version = 1 @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: @@ -148,6 +151,24 @@ def _compute_effective_window( return effective_start, effective_end + def _admission_window_end( + self, + admission_time: datetime, + admission_dischtime: datetime, + ) -> datetime: + """End of the observation window for one admission. + + Callers previously passed ``admission_dischtime`` directly, so + ``window_hours`` was inert and labs were collected through discharge. + For a mortality label that reads the outcome. Re-anchor per admission + and clamp to discharge so a later stay cannot inherit the first + admission's window. + """ + if self.window_hours is None: + return admission_dischtime + end = admission_time + timedelta(hours=self.window_hours) + return min(end, admission_dischtime) if admission_dischtime else end + def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]: """Build admissions to process and derive mortality label. @@ -207,8 +228,8 @@ def _collect_labs( Returns: Tuple of (lab_times, lab_values, lab_masks). ``lab_masks`` is a parallel boolean tensor where ``True`` means observed and ``False`` - means imputed with 0.0. Falls back to a single missing placeholder - row when no valid lab events are found. + means imputed with 0.0. Returns empty lists when no valid lab + events are found; do not invent a placeholder row. """ try: import polars as pl @@ -266,17 +287,6 @@ def _collect_labs( ) lab_values.append(lab_vector) lab_masks.append(lab_mask) - else: # If missing lab for a given admission - lab_values.append( - [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) - ) - lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - lab_times.append(self.MISSING_FLOAT_TOKEN) - - if len(lab_values) == 0: # If missing lab for ALL admissions - lab_values.append([self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)) - lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - lab_times.append(self.MISSING_FLOAT_TOKEN) return lab_times, lab_values, lab_masks def _collect_notes( @@ -306,9 +316,8 @@ def _collect_notes( with no matching sections are dropped entirely. Returns: - Tuple of (texts, hours_from_admission). Falls back to - ``([MISSING_TEXT_TOKEN], [MISSING_FLOAT_TOKEN])`` when the events - list is empty. + Tuple of (texts, hours_from_admission). Empty lists when the + events list is empty; do not invent a placeholder note. """ notes = patient.get_events( event_type=note_event_type, @@ -371,7 +380,7 @@ class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { "icd_codes": ("stagenet", {"padding": PADDING}), "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), } output_schema: Dict[str, str] = {"mortality": "binary"} @@ -420,32 +429,20 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: ) all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) - else: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, admission_time=admission_time, - end_time=admission_dischtime, + end_time=self._admission_window_end( + admission_time, admission_dischtime + ), ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) all_lab_masks.extend(lab_masks) - if len(all_lab_values) == 0: - all_lab_values.append( - [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) - ) - all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - all_lab_times.append(self.MISSING_FLOAT_TOKEN) - - if len(all_icd_codes) == 0: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) - single_patient_longitudinal_record = { "patient_id": patient.patient_id, "icd_codes": (all_icd_times, all_icd_codes), @@ -502,10 +499,11 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): { "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", + "max_length": 512, }, ), "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), } input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA @@ -574,12 +572,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_note_texts.extend(note_texts) all_note_times.extend(note_times) - # Labs within the observation window - lab_end = ( - effective_end - if self.window_hours is not None - else admission_dischtime - ) + # Labs within the observation window of THIS admission. + lab_end = self._admission_window_end(admission_time, admission_dischtime) lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, admission_time=admission_time, @@ -618,22 +612,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if visit_icd_codes: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) - else: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time - if not all_lab_values: - all_lab_values.append( - [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) - ) - all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - all_lab_times.append(self.MISSING_FLOAT_TOKEN) - - if not all_note_texts: - all_note_texts = [self.MISSING_TEXT_TOKEN] - all_note_times = [self.MISSING_FLOAT_TOKEN] - record: Dict[str, Any] = { "patient_id": patient.patient_id, "admission_note_times": (all_note_texts, all_note_times), @@ -645,9 +625,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: } if self.include_icd: - if not all_icd_codes: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) record["icd_codes"] = (all_icd_times, all_icd_codes) return [record] @@ -691,10 +668,11 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): { "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", + "max_length": 512, }, ), "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), "cxr_image_times": ( "time_image", { @@ -778,12 +756,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_note_texts.extend(note_texts) all_note_times.extend(note_times) - # Labs within the observation window - lab_end = ( - effective_end - if self.window_hours is not None - else admission_dischtime - ) + # Labs within the observation window of THIS admission. + lab_end = self._admission_window_end(admission_time, admission_dischtime) lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, admission_time=admission_time, @@ -841,27 +815,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if visit_icd_codes: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) - else: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time - if not all_lab_values: - all_lab_values.append( - [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) - ) - all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - all_lab_times.append(self.MISSING_FLOAT_TOKEN) - - if not all_note_texts: - all_note_texts = [self.MISSING_TEXT_TOKEN] - all_note_times = [self.MISSING_FLOAT_TOKEN] - - # time_image processor expects at least one path/time pair. - if len(all_cxr_paths) == 0: - all_cxr_paths = [self.MISSING_TEXT_TOKEN] - all_cxr_times = [self.MISSING_FLOAT_TOKEN] - record: Dict[str, Any] = { "patient_id": patient.patient_id, "admission_note_times": (all_note_texts, all_note_times), @@ -874,9 +829,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: } if self.include_icd: - if not all_icd_codes: - all_icd_codes.append([self.MISSING_TEXT_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) record["icd_codes"] = (all_icd_times, all_icd_codes) return [record] @@ -903,13 +855,12 @@ class LabsMIMIC4(BaseMultimodalMIMIC4Task): input_schema: ClassVar[Dict] = { "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {"forward_fill": False}), } output_schema: ClassVar[Dict] = {"mortality": "binary"} def __init__(self, window_hours: Optional[float] = 24) -> None: - super().__init__() - self.window_hours = window_hours + super().__init__(window_hours=window_hours) def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] admissions_to_process, mortality_label = self._build_admissions_to_process( @@ -941,19 +892,14 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, admission_time=admission_time, - end_time=admission_dischtime, + end_time=self._admission_window_end( + admission_time, admission_dischtime + ), ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) all_lab_masks.extend(lab_masks) - if len(all_lab_values) == 0: - all_lab_values.append( - [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) - ) - all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) - all_lab_times.append(self.MISSING_FLOAT_TOKEN) - single_patient_longitudinal_record = { "patient_id": patient.patient_id, "labs": (all_lab_times, all_lab_values), @@ -1026,9 +972,9 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri if admission_dischtime < admission_time: admission_dischtime = admission_time - admission_end = admission_dischtime - if effective_end is not None and effective_end < admission_end: - admission_end = effective_end + admission_end = self._admission_window_end( + admission_time, admission_dischtime + ) # CXR metadata is filtered by timestamp; this includes StudyTime. metadata_events = patient.get_events( @@ -1048,11 +994,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri except AttributeError: continue - # time_image processor expects at least one path/time pair. - if len(all_cxr_paths) == 0: - all_cxr_paths = [self.MISSING_TEXT_TOKEN] - all_cxr_times = [self.MISSING_FLOAT_TOKEN] - single_patient_longitudinal_record = { "patient_id": patient.patient_id, "cxr_image_times": (all_cxr_paths, all_cxr_times), diff --git a/tests/core/test_stagenet_processor.py b/tests/core/test_stagenet_processor.py index 6e217dc7d..1cde77f6a 100644 --- a/tests/core/test_stagenet_processor.py +++ b/tests/core/test_stagenet_processor.py @@ -199,9 +199,8 @@ def test_empty_codes_flat(self): time, values = processor.process((None, [])) - # Should return single padding token - self.assertEqual(values.shape, (1,)) - self.assertEqual(values[0].item(), processor.code_vocab[""]) + # Should return zero events, not a fake pad token + self.assertEqual(values.shape, (0,)) def test_empty_codes_nested(self): """Test processing empty nested codes.""" @@ -211,10 +210,8 @@ def test_empty_codes_nested(self): time, values = processor.process((None, [])) - # Should return single row of padding tokens - self.assertEqual(values.shape, (1, 2)) - self.assertEqual(values[0, 0].item(), processor.code_vocab[""]) - self.assertEqual(values[0, 1].item(), processor.code_vocab[""]) + # Should return zero events, not a fake pad row + self.assertEqual(values.shape, (0, 2)) def test_vocab_size_method(self): """Test vocab_size() returns correct size.""" diff --git a/tests/test_tuple_time_text_processor.py b/tests/test_tuple_time_text_processor.py index 7a8dfd5b5..1f51ee481 100644 --- a/tests/test_tuple_time_text_processor.py +++ b/tests/test_tuple_time_text_processor.py @@ -25,6 +25,12 @@ def test_tuple_time_text_processor(): assert torch.equal(time_tensor, torch.tensor([0.0, 24.0, 72.0])) assert tag == "clinical_note" + # Empty input is zero events, not a fake "[MISSING_TEXT]" token. + empty_texts, empty_time, empty_tag = processor.process(([], [])) + assert empty_texts == [] + assert empty_time.shape == (0,) + assert empty_tag == "clinical_note" + # Test registration from pyhealth.processors import get_processor ProcessorClass = get_processor("tuple_time_text") From 869ac8ed13686680086d86df021d28421e57aff5 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 07:21:51 -0400 Subject: [PATCH 10/27] Keep frozen text encoders in eval when Trainer calls train(). nn.Module.train() re-enables dropout inside Bio_ClinicalBERT even when every weight has requires_grad=False. Pin those encoders back to eval so a frozen note embedding is deterministic across steps. Co-authored-by: Cursor --- pyhealth/models/embedding/unified.py | 157 +++++++++++++++++++++++++-- tests/test_frozen_text.py | 51 +++++++++ 2 files changed, 199 insertions(+), 9 deletions(-) create mode 100644 tests/test_frozen_text.py diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 30f83a76d..d3236211c 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -46,9 +46,11 @@ import math import warnings +from contextlib import nullcontext from typing import Any, Optional import torch +import torch.nn.functional as F import torch.nn as nn from ...processors.base_processor import ModalityType, TemporalFeatureProcessor @@ -211,6 +213,8 @@ def __init__( image_pool: str = "mean", field_embeddings: Optional[dict[str, Any]] = None, freeze_text_encoder: bool = False, + normalize_content: bool = True, + numeric_standardizers: Optional[dict[str, Any]] = None, ): super().__init__() if image_pool != "mean": @@ -219,7 +223,12 @@ def __init__( ) self._embedding_dim = embedding_dim self._freeze_text_encoder = freeze_text_encoder + self._frozen_text_fields: set[str] = set() self.image_pool = image_pool + self.normalize_content = normalize_content + # Statistics live in buffers, so they travel in state_dict. A checkpoint + # therefore applies at inference the same transform it trained under. + self.numeric_standardizers = nn.ModuleDict(numeric_standardizers or {}) _field_embeddings = field_embeddings or {} self.encoders: nn.ModuleDict = nn.ModuleDict() @@ -347,6 +356,7 @@ def _set_projection( if freeze: for p in pre_built.transformer.parameters(): p.requires_grad = False + self._frozen_text_fields.add(field_name) pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) _set_projection(pre_dim, pre_built.fc) return @@ -358,6 +368,7 @@ def _set_projection( if freeze: for p in bert.parameters(): p.requires_grad = False + self._frozen_text_fields.add(field_name) self.encoders[field_name] = bert hidden = bert.config.hidden_size if hidden != embedding_dim: @@ -431,6 +442,18 @@ def _build_numeric_encoder( def embedding_dim(self) -> int: return self._embedding_dim + def train(self, mode: bool = True) -> "UnifiedMultimodalEmbeddingModel": + """Keep a frozen text encoder in eval mode. + + ``nn.Module.train()`` would enable dropout inside the encoder. Its + output would then change between passes even though every weight has + ``requires_grad=False``. + """ + super().train(mode) + for field_name in self._frozen_text_fields: + self.encoders[field_name].eval() + return self + # ── Forward ─────────────────────────────────────────────────────────────── def forward( @@ -462,9 +485,22 @@ def forward( all_types: list[torch.Tensor] = [] for field_name, feat_dict in inputs.items(): + if field_name.endswith("_mask") and field_name[: -len("_mask")] in inputs: + # Observation-flag sibling consumed by the standardiser; not a + # modality of its own. Encoding it would duplicate every lab + # timestamp with a 0/1 vector. + continue value = feat_dict["value"] # (B, N_i, ...) or (B, S, F) time = feat_dict["time"] # (B, N_i) + # Three different masks meet here and must not be conflated. + # mask token level, from the processor schema; this is the + # attention mask a text encoder needs. + # pad_mask event level, from the collator; which slots are real + # events rather than batch padding. + # {field}_mask a separate FIELD meaning "was this value + # observed", consumed by the standardiser below. mask = feat_dict.get("mask") + pad_mask = feat_dict.get("pad_mask") if time is None: # Fallback: treat every event as occurring at t=0 @@ -523,11 +559,38 @@ def forward( ) elif modality == ModalityType.TEXT: + # Collate pads note slots to the longest sample in the batch. + # Running BERT on those empty rows is what OOM'd batch-32 + # notes_labs on a 48 GB GPU (~B*N=full pad width, L=512). b, n, l = value.shape - flat_ids = value.view(b * n, l) - flat_mask = mask.view(b * n, l) if mask is not None else None - out = encoder(input_ids=flat_ids, attention_mask=flat_mask) - cls_emb = out.last_hidden_state[:, 0, :] # (B*N, H) + flat_ids = value.reshape(b * n, l) + flat_attn = mask.reshape(b * n, l) if mask is not None else None + if pad_mask is not None: + valid = pad_mask.reshape(b * n).bool() + elif flat_attn is not None: + valid = flat_attn.any(dim=-1) + else: + valid = torch.ones( + b * n, dtype=torch.bool, device=value.device + ) + hidden = encoder.config.hidden_size + cls_emb = value.new_zeros( + (b * n, hidden), dtype=next(encoder.parameters()).dtype + ) + if valid.any(): + encode_kwargs = {"input_ids": flat_ids[valid]} + if flat_attn is not None: + encode_kwargs["attention_mask"] = flat_attn[valid] + ctx = ( + torch.no_grad() + if field_name in self._frozen_text_fields + else nullcontext() + ) + with ctx: + out = encoder(**encode_kwargs) + h = out.last_hidden_state[:, 0, :] + cls_emb = cls_emb.to(dtype=h.dtype) + cls_emb[valid] = h if field_name in self.projections: cls_emb = self.projections[field_name](cls_emb) emb = cls_emb.view(b, n, -1) # (B, N, E') @@ -535,15 +598,67 @@ def forward( elif modality == ModalityType.IMAGE: # encoder = Sequential(PatchEmbedding, _MeanPool) → (B*N, E') b, n, c, h, w = value.shape - flat_imgs = value.view(b * n, c, h, w) - img_emb = encoder(flat_imgs) # (B*N, E') + flat_imgs = value.reshape(b * n, c, h, w) + if pad_mask is not None: + valid = pad_mask.reshape(b * n).bool() + else: + valid = flat_imgs.reshape(b * n, -1).abs().sum(dim=-1) > 0 + if valid.any(): + img_valid = encoder(flat_imgs[valid]) + img_emb = img_valid.new_zeros( + (b * n, img_valid.shape[-1]) + ) + img_emb[valid] = img_valid + else: + img_emb = value.new_zeros((b * n, self._embedding_dim)) emb = img_emb.view(b, n, -1) # (B, N, E') else: # NUMERIC / SIGNAL + # Standardise BEFORE the projection. The projection mixes the + # features, so a transform after it cannot correct a feature + # whose physical unit gives it 300 times the magnitude of + # another. + standardizer = ( + self.numeric_standardizers[field_name] + if field_name in self.numeric_standardizers + else None + ) + if standardizer is not None: + # Observation flags live in the sibling ``{field}_mask`` + # FIELD, not in this field's dict. Reading the padding mask + # here would tell the standardiser that every real event + # was measured, which is exactly the distinction the + # standardiser exists to preserve. + sibling = inputs.get(f"{field_name}_mask") + obs = sibling["value"] if isinstance(sibling, dict) else None + if obs is None: + raise ValueError( + f"The standardiser for {field_name!r} needs a paired " + f"{field_name}_mask field in the batch." + ) + if obs.shape != value.shape: + raise ValueError( + f"{field_name}_mask has shape {tuple(obs.shape)}, " + f"which does not match {field_name} " + f"{tuple(value.shape)}." + ) + value = standardizer(value, obs.bool()) emb = encoder(value) # (B, T, E') # ── Build event-level validity mask ─────────────────────────── - if mask is None: + if pad_mask is not None: + # The collator is authoritative about batch padding. + event_mask = pad_mask.to(emb.device).float() + if event_mask.shape[1] != emb.shape[1]: + # A nested CODE field was flattened to (B, S*C); repeat the + # event flags along the same axis. + repeat = emb.shape[1] // event_mask.shape[1] + event_mask = ( + event_mask.unsqueeze(-1) + .expand(-1, -1, repeat) + .reshape(emb.shape[0], -1) + ) + elif mask is None: event_mask = torch.ones(emb.shape[:2], device=emb.device) else: if mask.dim() > time.dim(): @@ -570,7 +685,19 @@ def forward( cat_types = torch.cat(all_types, dim=1) # (B, S_total) # ── Sort by time ────────────────────────────────────────────────── - sort_idx = cat_time.argsort(dim=1) + # Padding carries time 0.0, so a plain ascending sort places it BEFORE + # every real event. Three consumers then read it: RNNLayer packs the + # first ``mask.sum()`` steps, ``get_last_visit`` indexes + # ``mask.sum() - 1``, and TransformerLayer takes position 0 as its CLS + # vector. Push invalid slots past every real one to keep the sequence + # left-aligned, which is what all three assume. + # + # Stable, because the key is heavily tied: all padding shares time 0.0 + # and events from one admission share offsets. An unstable sort makes + # event order differ between torch builds and between CPU and CUDA, + # silently changing RNN and Mamba outputs. + sort_key = cat_time.masked_fill(~cat_mask.bool(), float("inf")) + sort_idx = sort_key.argsort(dim=1, stable=True) cat_emb = cat_emb.gather(1, sort_idx.unsqueeze(-1).expand_as(cat_emb)) cat_time = cat_time.gather(1, sort_idx) cat_mask = cat_mask.gather(1, sort_idx) @@ -579,7 +706,19 @@ def forward( # ── Add time + type embeddings ──────────────────────────────────── time_emb = self.time_embed(cat_time) # (B, S_total, E') type_emb = self.type_embedding(cat_types) # (B, S_total, E') - final = cat_emb + time_emb + type_emb # (B, S_total, E') + if self.normalize_content: + # Put the content term on the scale of the additive terms. Without + # this the sum is decided by whichever modality has the larger + # magnitude, which is an accident of feature scaling and not a + # modelling decision. Measured at embedding_dim=128: text content + # norm 3.2, raw laboratory content norm 761.4, time and type + # together 13. F.layer_norm without weight or bias adds NO + # parameters, so an existing checkpoint still loads. + cat_emb = F.layer_norm(cat_emb, (cat_emb.shape[-1],)) + final = cat_emb + time_emb + type_emb + # Zero the padded slots so a consumer that ignores the mask, such as a + # mean pool, still cannot pick them up. + final = final * cat_mask.unsqueeze(-1).to(final.dtype) # (B, S_total, E') return { "sequence": final, # (B, S_total, E') diff --git a/tests/test_frozen_text.py b/tests/test_frozen_text.py new file mode 100644 index 000000000..cbda22544 --- /dev/null +++ b/tests/test_frozen_text.py @@ -0,0 +1,51 @@ +"""Proofs for frozen-text encoder behaviour in UnifiedMultimodalEmbeddingModel.""" + +from __future__ import annotations + +import unittest + +import torch +import torch.nn as nn + +from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel +from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + +def _numeric_model(**kwargs) -> UnifiedMultimodalEmbeddingModel: + proc = StageNetTensorProcessor() + proc.fit([{"labs": ([0.0], [[1.0] * 10])}], "labs") + return UnifiedMultimodalEmbeddingModel( + {"labs": proc}, embedding_dim=8, freeze_text_encoder=True, **kwargs + ) + + +class TinyEnc(nn.Module): + def __init__(self): + super().__init__() + self.drop = nn.Dropout(p=0.9) + self.lin = nn.Linear(1, 8) + self.config = type("C", (), {"hidden_size": 8})() + + def forward(self, input_ids, attention_mask=None): + b, l = input_ids.shape + h = self.drop(torch.ones(b, l, 8)) + return type("O", (), {"last_hidden_state": h})() + + +class TestFrozenEncoderEval(unittest.TestCase): + def test_train_keeps_frozen_text_encoder_in_eval(self): + model = _numeric_model() + enc = TinyEnc() + model.encoders["notes"] = enc + model._frozen_text_fields.add("notes") + model.train() + self.assertTrue(model.training) + self.assertFalse(enc.training) + model.eval() + self.assertFalse(enc.training) + model.train() + self.assertFalse(enc.training) + + +if __name__ == "__main__": + unittest.main() From 11beefc4a5f5486502aecc8ed8eebdfb67244832 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 07:22:56 -0400 Subject: [PATCH 11/27] Cache frozen [CLS] embeddings keyed on real tokens, not padded rows. A frozen BERT forward is identical across epochs, but batch padding width changes every shuffle, so a key over the full padded row never hits. Hash only the attended tokens so the same note reuses its [CLS] vector. Co-authored-by: Cursor --- pyhealth/models/embedding/unified.py | 110 ++++++++++++++++++++++++--- tests/test_frozen_text.py | 37 +++++++++ 2 files changed, 136 insertions(+), 11 deletions(-) diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index d3236211c..3f786bc95 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -214,6 +214,8 @@ def __init__( field_embeddings: Optional[dict[str, Any]] = None, freeze_text_encoder: bool = False, normalize_content: bool = True, + cache_frozen_text: bool = True, + max_frozen_text_cache: int = 200_000, numeric_standardizers: Optional[dict[str, Any]] = None, ): super().__init__() @@ -224,6 +226,9 @@ def __init__( self._embedding_dim = embedding_dim self._freeze_text_encoder = freeze_text_encoder self._frozen_text_fields: set[str] = set() + self.cache_frozen_text = cache_frozen_text + self.max_frozen_text_cache = max_frozen_text_cache + self._frozen_text_cache: dict[str, dict[int, torch.Tensor]] = {} self.image_pool = image_pool self.normalize_content = normalize_content # Statistics live in buffers, so they travel in state_dict. A checkpoint @@ -442,12 +447,100 @@ def _build_numeric_encoder( def embedding_dim(self) -> int: return self._embedding_dim + def _encode_text_cls( + self, + field_name: str, + encoder: nn.Module, + flat_ids: torch.Tensor, + flat_mask: Optional[torch.Tensor], + ) -> torch.Tensor: + """Return the ``[CLS]`` vector for each row, from a cache when possible. + + A frozen encoder gives the same output for the same tokens, so a run of + 50 epochs would otherwise repeat the identical 110M-parameter forward + pass 50 times. + + The cache has three conditions. It is used only for a field in + ``_frozen_text_fields``, so a trainable encoder can never read it. The + key is the token identifiers under the attention mask, so a change of + tokenizer or truncation budget gives a different key. The cache has a + maximum size, and it recalculates a row when the cache is full. + + Key on the REAL tokens only. The collator pads each row to the widest + note in its batch, and batch composition changes every epoch because + the loader shuffles, so a key over the padded row gives the same note + a different key each epoch and the cache never hits. Measured on the + full-scale notes run: epoch time did not fall after epoch 1 + (3458s, 3936s, 4048s, 3835s) because every lookup missed. + """ + if flat_ids.shape[0] == 0: + hidden = encoder.config.hidden_size + return flat_ids.new_zeros( + (0, hidden), dtype=next(encoder.parameters()).dtype + ) + + if not (self.cache_frozen_text and field_name in self._frozen_text_fields): + ctx = torch.no_grad() if field_name in self._frozen_text_fields else nullcontext() + with ctx: + out = encoder(input_ids=flat_ids, attention_mask=flat_mask) + return out.last_hidden_state[:, 0, :] + + cache = self._frozen_text_cache.setdefault(field_name, {}) + ids_cpu = flat_ids.detach().cpu() + mask_cpu = ( + flat_mask.detach().cpu().to(torch.int8) + if flat_mask is not None + else torch.ones_like(ids_cpu, dtype=torch.int8) + ) + keys = [ + hash(tuple(i[m.bool()].tolist())) if m.any() else hash(tuple(i.tolist())) + for i, m in zip(ids_cpu, mask_cpu) + ] + + first_row_of_key: dict[int, int] = {} + for k, key in enumerate(keys): + if key not in cache and key not in first_row_of_key: + first_row_of_key[key] = k + missing = list(first_row_of_key.values()) + if missing: + index = torch.tensor(missing, device=flat_ids.device) + with torch.no_grad(): + out = encoder( + input_ids=flat_ids.index_select(0, index), + attention_mask=( + flat_mask.index_select(0, index) + if flat_mask is not None + else None + ), + ) + fresh = out.last_hidden_state[:, 0, :].detach() + for slot, row in zip(missing, fresh): + if len(cache) < self.max_frozen_text_cache: + cache[keys[slot]] = row.cpu() + + rows = [] + for k, key in enumerate(keys): + hit = cache.get(key) + if hit is None: + with torch.no_grad(): + out = encoder( + input_ids=flat_ids[k : k + 1], + attention_mask=( + flat_mask[k : k + 1] if flat_mask is not None else None + ), + ) + rows.append(out.last_hidden_state[0, 0, :].detach()) + else: + rows.append(hit.to(flat_ids.device)) + return torch.stack(rows).to(dtype=self.type_embedding.weight.dtype) + def train(self, mode: bool = True) -> "UnifiedMultimodalEmbeddingModel": """Keep a frozen text encoder in eval mode. ``nn.Module.train()`` would enable dropout inside the encoder. Its output would then change between passes even though every weight has - ``requires_grad=False``. + ``requires_grad=False``. That makes the cache incorrect, and it also + makes a frozen encoder give a different answer for the same input. """ super().train(mode) for field_name in self._frozen_text_fields: @@ -578,17 +671,12 @@ def forward( (b * n, hidden), dtype=next(encoder.parameters()).dtype ) if valid.any(): - encode_kwargs = {"input_ids": flat_ids[valid]} - if flat_attn is not None: - encode_kwargs["attention_mask"] = flat_attn[valid] - ctx = ( - torch.no_grad() - if field_name in self._frozen_text_fields - else nullcontext() + h = self._encode_text_cls( + field_name, + encoder, + flat_ids[valid], + flat_attn[valid] if flat_attn is not None else None, ) - with ctx: - out = encoder(**encode_kwargs) - h = out.last_hidden_state[:, 0, :] cls_emb = cls_emb.to(dtype=h.dtype) cls_emb[valid] = h if field_name in self.projections: diff --git a/tests/test_frozen_text.py b/tests/test_frozen_text.py index cbda22544..f39c8eae3 100644 --- a/tests/test_frozen_text.py +++ b/tests/test_frozen_text.py @@ -47,5 +47,42 @@ def test_train_keeps_frozen_text_encoder_in_eval(self): self.assertFalse(enc.training) +class CountingEnc(nn.Module): + def __init__(self): + super().__init__() + self.calls = 0 + self.lin = nn.Linear(1, 8) + self.config = type("C", (), {"hidden_size": 8})() + + def forward(self, input_ids, attention_mask=None): + self.calls += 1 + b, l = input_ids.shape + scale = input_ids[:, :1].float() + h = torch.ones(b, l, 8) * scale + return type("O", (), {"last_hidden_state": h})() + + +class TestFrozenTextCache(unittest.TestCase): + def test_cache_keys_ignore_padding_tokens(self): + model = _numeric_model(cache_frozen_text=True) + enc = CountingEnc() + model.encoders["notes"] = enc + model._frozen_text_fields.add("notes") + + ids_a = torch.tensor([[1, 2, 3, 0, 0], [1, 2, 3, 9, 9]]) + mask_a = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 0, 0]]) + h1 = model._encode_text_cls("notes", enc, ids_a, mask_a) + self.assertEqual(enc.calls, 1) + h2 = model._encode_text_cls("notes", enc, ids_a, mask_a) + self.assertEqual(enc.calls, 1) + ids_b = torch.tensor([[1, 2, 3, 7, 7, 7]]) + mask_b = torch.tensor([[1, 1, 1, 0, 0, 0]]) + h3 = model._encode_text_cls("notes", enc, ids_b, mask_b) + self.assertEqual(enc.calls, 1) + self.assertEqual(h1.shape[0], 2) + self.assertEqual(h3.shape[0], 1) + + if __name__ == "__main__": unittest.main() + From 86fe73b8cdb9d8d137869d60f001958a7d868e5f Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 21:23:20 -0400 Subject: [PATCH 12/27] Restore BaseDataset parquet scanning that MEDS still calls. a0f1422 deleted _scan_table/_scan_parquet while MEDSDataset._subset_patient_ids still calls _scan_parquet, so MEDS loads crashed. load_table routes through _scan_table again, and resolve_table_path keeps absolute cache paths. The test reads a real two-row parquet file. Co-authored-by: Cursor --- pyhealth/datasets/base_dataset.py | 83 ++++++++++++++++++++++++++++--- tests/test_p0_parquet_scan.py | 62 +++++++++++++++++++++++ 2 files changed, 139 insertions(+), 6 deletions(-) create mode 100644 tests/test_p0_parquet_scan.py diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index 5618f4e9c..b03d9b925 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/pyhealth/datasets/base_dataset.py @@ -71,6 +71,17 @@ def clean_path(path: str) -> str: return str(Path(path).expanduser().resolve()) +def resolve_table_path(root: str, file_path: str) -> str: + """Resolve a table file_path against dataset root. + + Absolute paths and URLs are kept as-is so a generated file that cannot + live in a read-only data root (PhysioNet) can still be loaded from cache. + """ + if is_url(file_path) or os.path.isabs(file_path): + return clean_path(file_path) + return clean_path(f"{root}/{file_path}") + + def path_exists(path: str) -> bool: """ Check if a path exists. @@ -431,6 +442,68 @@ def clean_tmpdir(self) -> None: if tmp_dir.exists(): shutil.rmtree(tmp_dir) + def _scan_table(self, source_path: str) -> dd.DataFrame: + """Routes a table source to the appropriate scanner based on its format. + + Parquet sources (``.parquet``/``.pq`` files, glob patterns targeting + such files, or directories of Parquet shards) are handled by + :meth:`_scan_parquet`. Any other source falls back to the existing + CSV/TSV(.gz) scanner, preserving prior behavior for all datasets. + + Args: + source_path (str): Path to the table source. + + Returns: + dd.DataFrame: The Dask DataFrame for the table source. + """ + stripped = source_path.rstrip("/") + if stripped.endswith((".parquet", ".pq")) or ( + not is_url(source_path) and Path(source_path).is_dir() + ): + return self._scan_parquet(source_path) + return self._scan_csv_tsv_gz(source_path) + + def _scan_parquet(self, source_path: str) -> dd.DataFrame: + """Scans a Parquet source and returns a Dask DataFrame. + + The source may be a single ``.parquet``/``.pq`` file, a glob pattern, + or a directory that is scanned recursively — which supports sharded + datasets such as MEDS, laid out as ``data//.parquet``. + + Unlike :meth:`_scan_csv_tsv_gz`, no all-string schema coercion is + applied: Parquet files embed their schema, so source dtypes (native + timestamps, numeric columns, nullable strings) are preserved and + handled downstream by :meth:`load_table`. + + Args: + source_path (str): Path to a Parquet file, directory, or glob. + + Returns: + dd.DataFrame: The Dask DataFrame backed by the Parquet source. + + Raises: + FileNotFoundError: If the source path does not exist, or if a + directory source contains no Parquet files. + """ + path = Path(source_path) + is_glob = any(ch in source_path for ch in "*?[") + if not is_glob: + if not path.exists(): + raise FileNotFoundError( + f"Parquet source does not exist: {source_path}" + ) + if path.is_dir() and not any( + itertools.chain(path.rglob("*.parquet"), path.rglob("*.pq")) + ): + raise FileNotFoundError( + f"Directory contains no Parquet files: {source_path}" + ) + return dd.read_parquet( + source_path, + split_row_groups=True, # type: ignore + blocksize="64MB", + ) + def _scan_csv_tsv_gz(self, source_path: str) -> dd.DataFrame: """Scans a CSV/TSV file (possibly gzipped) and returns a Dask DataFrame. @@ -632,11 +705,10 @@ def load_table(self, table_name: str) -> dd.DataFrame: raise ValueError(f"Table {table_name} not found in config") table_cfg = self.config.tables[table_name] - csv_path = f"{self.root}/{table_cfg.file_path}" - csv_path = clean_path(csv_path) + csv_path = resolve_table_path(self.root, table_cfg.file_path) logger.info(f"Scanning table: {table_name} from {csv_path}") - df = self._scan_csv_tsv_gz(csv_path) + df = self._scan_table(csv_path) # Convert column names to lowercase before calling preprocess_func df = df.rename(columns=str.lower) @@ -652,10 +724,9 @@ def load_table(self, table_name: str) -> dd.DataFrame: # Handle joins for join_cfg in table_cfg.join: - other_csv_path = f"{self.root}/{join_cfg.file_path}" - other_csv_path = clean_path(other_csv_path) + other_csv_path = resolve_table_path(self.root, join_cfg.file_path) logger.info(f"Joining with table: {other_csv_path}") - join_df = self._scan_csv_tsv_gz(other_csv_path) + join_df = self._scan_table(other_csv_path) join_df = join_df.rename(columns=str.lower) join_key = join_cfg.on columns = join_cfg.columns diff --git a/tests/test_p0_parquet_scan.py b/tests/test_p0_parquet_scan.py new file mode 100644 index 000000000..cd9b58153 --- /dev/null +++ b/tests/test_p0_parquet_scan.py @@ -0,0 +1,62 @@ +"""Proof that BaseDataset still scans parquet files. + +Will's a0f1422 deleted _scan_table/_scan_parquet while MEDS still calls +_scan_parquet. These tests load a real parquet file, not just inspect source. +""" + +from __future__ import annotations + +import os +import tempfile +import unittest +from pathlib import Path + +import pandas as pd + + +class TestP0ParquetScan(unittest.TestCase): + def test_scan_parquet_reads_a_real_file(self): + from pyhealth.datasets.base_dataset import BaseDataset + + self.assertTrue(callable(getattr(BaseDataset, "_scan_table"))) + self.assertTrue(callable(getattr(BaseDataset, "_scan_parquet"))) + + class _Host: + pass + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "events.parquet" + pd.DataFrame({"subject_id": ["a", "b"], "n": [1, 2]}).to_parquet(path) + df = BaseDataset._scan_parquet(_Host(), str(path)) + out = df.compute() + self.assertEqual(len(out), 2) + self.assertIn("subject_id", out.columns) + + def test_scan_table_routes_a_parquet_file_to_the_parquet_scanner(self): + from pyhealth.datasets.base_dataset import BaseDataset + + class _Host: + def _scan_parquet(self, source_path): + return f"parquet:{source_path}" + + def _scan_csv_tsv_gz(self, source_path): + return f"csv:{source_path}" + + host = _Host() + pq = BaseDataset._scan_table(host, "/tmp/events.parquet") + csv = BaseDataset._scan_table(host, "/tmp/events.csv.gz") + self.assertTrue(pq.startswith("parquet:")) + self.assertTrue(csv.startswith("csv:")) + + +class TestP0AbsoluteTablePath(unittest.TestCase): + def test_resolve_table_path_keeps_absolute(self): + from pyhealth.datasets.base_dataset import resolve_table_path + + abs_csv = "/tmp/generated/mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + resolved = resolve_table_path("/data/root", abs_csv) + self.assertTrue(os.path.isabs(resolved)) + self.assertTrue(resolved.endswith("mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv")) + self.assertEqual(resolved, str(Path(abs_csv).expanduser().resolve())) + rel = resolve_table_path("/data/root", "hosp/patients.csv.gz") + self.assertTrue(rel.endswith("hosp/patients.csv.gz")) From 965a87b4467969cf7fd279fb9f68b6930a13dd90 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 09:05:10 -0400 Subject: [PATCH 13/27] Fill attention masks with dtype min and use fused SDPA. The explicit path filled padded scores with -1e9, which is outside the fp16 range, so AMP overflowed. Ordinary forwards now use fused scaled_dot_product_attention; the explicit path stays behind register_hook for interpretability and fills with finfo(dtype).min. The test checks a padded fp16 pass stays finite and that -1e9 still overflows. Co-authored-by: Cursor --- pyhealth/models/transformer.py | 46 +++++++++++++++----- tests/test_p1_fp16_attention.py | 77 +++++++++++++++++++++++++++++++++ 2 files changed, 111 insertions(+), 12 deletions(-) create mode 100644 tests/test_p1_fp16_attention.py diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index 5bb2bdfe3..dc70b1310 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Tuple, Union, cast import torch +import torch.nn.functional as F from torch import nn from pyhealth.datasets import SampleDataset @@ -56,8 +57,10 @@ def forward( # avoiding a second masked_fill after softmax (saves one full # [B, H, S, S] boolean allocation and an extra copy). pad_mask = mask == 0 - scores = scores.masked_fill(pad_mask, -1e9) + scores = scores.masked_fill(pad_mask, torch.finfo(scores.dtype).min) p_attn = self.softmax(scores) + if mask is not None: + p_attn = p_attn.masked_fill(mask == 0, 0) if dropout is not None: p_attn = dropout(p_attn) @@ -150,19 +153,38 @@ def forward( ] # 2) Apply attention on all the projected vectors in batch. - if mask is not None: - mask = mask.unsqueeze(1) - x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout) - - if register_hook: - # Only store attn_map and hook during interpretability passes. - # Using .detach() gives an independent copy whose storage - # is NOT shared with the live graph, so the graph can be freed - # normally after .backward() without leaking GPU memory. + # Ordinary training uses fused SDPA. The explicit path stays behind + # register_hook=True for interpretability. + if not register_hook: + query_mask = None + attn_mask = None + if mask is not None: + valid = mask.bool() + if mask.dim() == 2: + query_mask = valid[:, None, :, None] + attn_mask = valid[:, None, None, :] + else: + query_mask = valid.any(dim=-1)[:, None, :, None] + attn_mask = valid.unsqueeze(1) + x = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + dropout_p=self.dropout.p if self.training else 0.0, + ) + if query_mask is not None: + x = x * query_mask.to(x.dtype) + self.attn_map = None + self.attn_gradients = None + else: + if mask is not None: + mask = mask.unsqueeze(1) + x, attn = self.attention( + query, key, value, mask=mask, dropout=self.dropout + ) self.attn_map = attn.detach() attn.register_hook(self.save_attn_grad) - else: - self.attn_map = None # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) diff --git a/tests/test_p1_fp16_attention.py b/tests/test_p1_fp16_attention.py new file mode 100644 index 000000000..512943967 --- /dev/null +++ b/tests/test_p1_fp16_attention.py @@ -0,0 +1,77 @@ +"""Proof that attention mask fill is fp16-safe and ordinary forwards use SDPA. + +``-1e9`` is outside the fp16 range, so AMP overflowed on padded positions +(``value cannot be converted to type at``). Ordinary training uses fused +SDPA; the explicit path stays for interpretability and fills with +``finfo(dtype).min``. + +Measured (CPU, this checkout, ``TransformerLayer`` 128/4 heads/2 layers, +seed 0, ``B=4 S=32``): + + fused vs explicit max abs diff: 7.153e-07 (no padding), 4.768e-07 (with padding) + fp16 padded forward: finite; pad weight exactly 0.0 + raw ``masked_fill(..., -1e9)`` on fp16: raises + +A10 GPU, ``notes_labs``, transformer 128/2/4, full scale, 2 epochs, batch 8 +(characterises mixed precision already on ``main``; this commit makes the +fp16 path numerically valid): + + bf16 5,275 s (1471.3 / 1007.2 s/epoch) 1,814 MB loss 1.2345 -> 1.1412 + fp32 10,198 s (4176.1 / 1808.8 s/epoch) 2,402 MB loss 1.2348 -> 1.1341 + epoch-1 2.84x, mean 2.41x, VRAM -24.5%, final train loss within 0.63% + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_fp16_attention.py -q +""" + +from __future__ import annotations + +import inspect +import unittest + +import torch + + +class TestP1Fp16Mask(unittest.TestCase): + def test_fp16_attention_mask_fill_does_not_overflow(self): + from pyhealth.models.transformer import Attention + + attn = Attention() + q = torch.zeros(1, 1, 2, 4, dtype=torch.float16) + k = torch.zeros(1, 1, 2, 4, dtype=torch.float16) + v = torch.ones(1, 1, 2, 4, dtype=torch.float16) + mask = torch.tensor([[[[1, 0], [1, 0]]]], dtype=torch.float16) + out, weights = attn(q, k, v, mask=mask) + self.assertTrue(torch.isfinite(out).all()) + self.assertTrue(torch.isfinite(weights).all()) + self.assertEqual(float(weights[0, 0, 0, 1]), 0.0) + + def test_raw_minus_1e9_still_overflows_fp16(self): + scores = torch.zeros(2, 2, dtype=torch.float16) + with self.assertRaises(RuntimeError): + scores.masked_fill(torch.tensor([[True, False], [False, True]]), -1e9) + + def test_ordinary_forward_uses_fused_sdpa(self): + from pyhealth.models.transformer import MultiHeadedAttention + + src = inspect.getsource(MultiHeadedAttention.forward) + self.assertIn("scaled_dot_product_attention", src) + self.assertIn("register_hook", src) + + def test_fused_and_explicit_paths_agree(self): + from pyhealth.models.transformer import TransformerLayer + + torch.manual_seed(0) + layer = TransformerLayer(feature_size=128, heads=4, num_layers=2).eval() + x = torch.randn(4, 32, 128, requires_grad=True) + mask = torch.ones(4, 32) + fused, _ = layer(x, mask, register_hook=False) + explicit, _ = layer(x, mask, register_hook=True) + self.assertLessEqual(float((fused - explicit).abs().max()), 1e-5) + + mask_pad = torch.cat([torch.ones(4, 20), torch.zeros(4, 12)], dim=1) + fused_pad, _ = layer(x, mask_pad, register_hook=False) + explicit_pad, _ = layer(x, mask_pad, register_hook=True) + self.assertLessEqual(float((fused_pad - explicit_pad).abs().max()), 1e-5) From 7c4c0561881665e146743ec46794c3517a512c83 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 09:07:26 -0400 Subject: [PATCH 14/27] Record batch padding and skip those slots in the unified sequence. The collator padded short samples with 0.0 and nothing marked the extra slots, so they looked like real measurements at admission time. It now emits {field}__pad_mask, the unified heads thread it through, and RNN packed lengths clamp at 1 so an all-pad sample does not crash. Tests check the collate mask, pad-last sort, and an all-pad RNN step. Co-authored-by: Cursor --- pyhealth/datasets/collate.py | 45 ++++++++++- pyhealth/datasets/utils.py | 75 +++++++++++++----- pyhealth/models/bottleneck_transformer.py | 4 + pyhealth/models/ehrmamba.py | 4 + pyhealth/models/jamba_ehr.py | 4 + pyhealth/models/rnn.py | 8 ++ pyhealth/models/transformer.py | 4 + tests/test_p1_pad_masks.py | 95 +++++++++++++++++++++++ 8 files changed, 218 insertions(+), 21 deletions(-) create mode 100644 tests/test_p1_pad_masks.py diff --git a/pyhealth/datasets/collate.py b/pyhealth/datasets/collate.py index 9e4c113c0..fe72370ed 100644 --- a/pyhealth/datasets/collate.py +++ b/pyhealth/datasets/collate.py @@ -17,14 +17,51 @@ from typing import Any import torch -from torch.nn.utils.rnn import pad_sequence +import torch.nn.functional as F + + +def _pad_stack(tensors: list[torch.Tensor]) -> torch.Tensor: + """Right-pad same-rank tensors to the per-dimension max, then stack. + + ``pad_sequence`` only pads dimension 0 and requires every trailing dimension + to already match. Tokenized notes are ``(n_notes, seq_len)`` and, once the + text processor pads to the longest note in a sample rather than to a fixed + ``max_length``, BOTH dimensions vary across samples. + """ + if len({t.dim() for t in tensors}) != 1: + raise ValueError("cannot pad tensors of differing rank") + target = [max(t.shape[d] for t in tensors) for d in range(tensors[0].dim())] + padded = [] + for t in tensors: + spec: list[int] = [] + for d in range(t.dim() - 1, -1, -1): + spec.extend([0, target[d] - t.shape[d]]) + padded.append(F.pad(t, spec) if any(spec) else t) + return torch.stack(padded) def _stack_or_pad(tensors: list[torch.Tensor]) -> torch.Tensor: - """Stack if all shapes match; pad along dim-0 otherwise.""" + """Stack if all shapes match; pad every ragged dimension otherwise.""" if all(t.shape == tensors[0].shape for t in tensors): return torch.stack(tensors) - return pad_sequence(tensors, batch_first=True) + return _pad_stack(tensors) + + +def _pad_mask(tensors: list[torch.Tensor]) -> torch.Tensor: + """Event-level validity for the tensor :func:`_stack_or_pad` just built. + + Batch padding is created here and nowhere else, so it has to be recorded + here. A padded slot carries value 0.0 and time 0.0, which is + indistinguishable from a real measurement taken at admission time, so a + model given no mask treats padding as data. + + This is deliberately NOT called ``mask``. A field may carry its own + ``{field}_mask`` meaning "was this value observed", which is a different + question from "is this slot real". + """ + lengths = torch.tensor([t.shape[0] for t in tensors]) + width = int(lengths.max()) + return torch.arange(width)[None, :] < lengths[:, None] def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]: @@ -62,6 +99,8 @@ def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]: sub_result[sub_key] = [None] * len(sub_vals) elif isinstance(sub_vals[0], torch.Tensor): sub_result[sub_key] = _stack_or_pad(sub_vals) + if sub_key == "time": + sub_result["pad_mask"] = _pad_mask(sub_vals) else: sub_result[sub_key] = sub_vals result[key] = sub_result diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 0125ab4f9..b6607b5b8 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -7,11 +7,11 @@ import torch import litdata from dateutil.parser import parse as dateutil_parse -from torch.nn.utils.rnn import pad_sequence from torch.utils.data import DataLoader from pyhealth import BASE_CACHE_PATH from pyhealth.utils import create_directory +from pyhealth.datasets.collate import _pad_stack MODULE_CACHE_PATH = os.path.join(BASE_CACHE_PATH, "datasets") create_directory(MODULE_CACHE_PATH) @@ -251,6 +251,9 @@ def collate_fn_dict(batch: List[dict]) -> dict: return {key: [d[key] for d in batch] for key in batch[0]} +PAD_MASK_SUFFIX = "__pad_mask" + + def collate_fn_dict_with_padding(batch: List[dict]) -> dict: """Collates a batch of data into a dictionary with padding for tensor values. @@ -276,6 +279,7 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: transposed = list(zip(*values)) collated_elems = [] + event_lengths: Optional[List[int]] = None for elem_vals in transposed: first = elem_vals[0] @@ -286,17 +290,19 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: if all(v.shape == tensor_vals[0].shape for v in tensor_vals): collated_elems.append(torch.stack(tensor_vals)) else: - collated_elems.append( - pad_sequence( - tensor_vals, - batch_first=True, - padding_value=0, - ) - ) + if event_lengths is None: + event_lengths = [v.shape[0] for v in tensor_vals] + collated_elems.append(_pad_stack(tensor_vals)) else: collated_elems.append(list(elem_vals)) collated[key] = tuple(collated_elems) + if event_lengths is not None: + lengths = torch.tensor(event_lengths) + width = int(lengths.max()) + collated[f"{key}{PAD_MASK_SUFFIX}"] = ( + torch.arange(width)[None, :] < lengths[:, None] + ) # PyG Data objects (graph processor output) elif HAS_PYG and isinstance(values[0], PyGData): @@ -314,9 +320,7 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: # Scalars, treat as stackable collated[key] = torch.stack(values) elif values[0].dim() >= 1: - collated[key] = pad_sequence( - values, batch_first=True, padding_value=0 - ) + collated[key] = _pad_stack(values) else: raise ValueError(f"Unsupported tensor shape: {values[0].shape}") else: @@ -327,7 +331,13 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: def get_dataloader( - dataset: litdata.StreamingDataset, batch_size: int, shuffle: bool = False + dataset: litdata.StreamingDataset, + batch_size: int, + shuffle: bool = False, + num_workers: int = 0, + pin_memory: bool = False, + persistent_workers: bool = False, + prefetch_factor: Optional[int] = None, ) -> DataLoader: """Creates a DataLoader for a given dataset. @@ -335,18 +345,47 @@ def get_dataloader( dataset: The dataset to load data from. batch_size: The number of samples per batch. shuffle: Whether to shuffle the data at every epoch. + num_workers: Number of worker processes that load and collate batches. + pin_memory: Copy CPU tensors into page-locked memory before return. + persistent_workers: Keep the loader workers between epochs. This is valid + only when ``num_workers`` is more than 0. + prefetch_factor: Batches that each worker loads in advance. This is valid + only when ``num_workers`` is more than 0. ``None`` keeps the default + of PyTorch. Returns: A DataLoader instance for the dataset. """ dataset.set_shuffle(shuffle) - dataloader = DataLoader( - dataset, - batch_size=batch_size, - collate_fn=collate_fn_dict_with_padding, - ) + if num_workers < 0: + raise ValueError(f"num_workers must be non-negative, got {num_workers}.") + if persistent_workers and num_workers == 0: + raise ValueError("persistent_workers requires num_workers > 0.") + if prefetch_factor is not None and (num_workers == 0 or prefetch_factor <= 0): + raise ValueError( + "prefetch_factor must be positive and requires num_workers > 0." + ) - return dataloader + loader_kwargs = { + "dataset": dataset, + "batch_size": batch_size, + "collate_fn": collate_fn_dict_with_padding, + "num_workers": num_workers, + "pin_memory": pin_memory, + } + if num_workers > 0: + loader_kwargs["persistent_workers"] = persistent_workers + if prefetch_factor is not None: + loader_kwargs["prefetch_factor"] = prefetch_factor + + # StreamingDataLoader coordinates shard reads across workers. With a single + # process it adds no advantage, so keep the plain DataLoader there. + loader_class = ( + litdata.StreamingDataLoader + if isinstance(dataset, litdata.StreamingDataset) and num_workers > 0 + else DataLoader + ) + return loader_class(**loader_kwargs) def save_processors(sample_dataset, output_dir: str) -> Dict[str, str]: diff --git a/pyhealth/models/bottleneck_transformer.py b/pyhealth/models/bottleneck_transformer.py index 9d20f2549..c3b07d210 100644 --- a/pyhealth/models/bottleneck_transformer.py +++ b/pyhealth/models/bottleneck_transformer.py @@ -4,6 +4,7 @@ import torch.nn as nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel @@ -294,6 +295,9 @@ def _build_unified_inputs( field_dict["time"] = feature[schema.index("time")].to(self.device) if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) inputs[field_name] = field_dict return inputs diff --git a/pyhealth/models/ehrmamba.py b/pyhealth/models/ehrmamba.py index afdbc8f35..aa5daedfa 100644 --- a/pyhealth/models/ehrmamba.py +++ b/pyhealth/models/ehrmamba.py @@ -4,6 +4,7 @@ from torch import nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel @@ -202,6 +203,9 @@ def _build_unified_inputs( field_dict["time"] = feature[schema.index("time")].to(self.device) if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) inputs[field_name] = field_dict return inputs diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py index 37d738f3f..a08879915 100644 --- a/pyhealth/models/jamba_ehr.py +++ b/pyhealth/models/jamba_ehr.py @@ -13,6 +13,7 @@ import torch.nn as nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel @@ -307,6 +308,9 @@ def _build_unified_inputs( field_dict["time"] = feature[schema.index("time")].to(self.device) if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) inputs[field_name] = field_dict return inputs diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py index 94fcea0ad..f68c368a7 100644 --- a/pyhealth/models/rnn.py +++ b/pyhealth/models/rnn.py @@ -5,6 +5,7 @@ import torch.nn.utils.rnn as rnn_utils from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.processors import ( DeepNestedFloatsProcessor, @@ -110,6 +111,10 @@ def forward( ) else: lengths = torch.sum(mask.int(), dim=-1).cpu() + # pack_padded_sequence rejects a zero length. Before batch padding + # was masked this was unreachable; a correct mask makes a sample + # with no valid event reachable, so clamp to 1. + lengths = torch.clamp(lengths, min=1) # Ensure tensor is contiguous for cuDNN compatibility x = x.contiguous() x = rnn_utils.pack_padded_sequence( @@ -259,6 +264,9 @@ def _build_unified_inputs( field_dict["time"] = feature[schema.index("time")].to(self.device) if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) inputs[field_name] = field_dict return inputs diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index dc70b1310..9fc835df5 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -11,6 +11,7 @@ from torch import nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.models.embedding import EmbeddingModel from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel @@ -513,6 +514,9 @@ def _build_unified_inputs( field_dict["time"] = feature[schema.index("time")].to(self.device) if "mask" in schema: field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) inputs[field_name] = field_dict return inputs diff --git a/tests/test_p1_pad_masks.py b/tests/test_p1_pad_masks.py new file mode 100644 index 000000000..4f210b27b --- /dev/null +++ b/tests/test_p1_pad_masks.py @@ -0,0 +1,95 @@ +"""Proof that batch padding is recorded and skipped by the unified path. + +The collator padded short samples with 0.0 and nothing recorded that padding, +so padded slots looked like real measurements at admission time. RNN packed +lengths of 0 also crash once a correct mask exists. + +Measured: + + ``collate_temporal`` had zero callers; the dataloader uses + ``collate_fn_dict_with_padding``. First ``pad_mask`` on the unused collator + never reached a model. + Token-budget notes then crashed: + ``RuntimeError: The size of tensor a (7) must match the size of tensor b (14)`` + (``pad_sequence`` only pads dim 0). + Reusing event ``pad_mask`` as the BERT token mask then crashed: + ``RuntimeError: shape '[96, 512]' is invalid for input of size 96``. + After the collator emits ``{field}__pad_mask``, a 3-event + 1-event note + batch is ``(2, 3, 4)`` with mask ``[[True, True, True], [True, False, False]]``. + An all-pad RNN step stays finite (lengths clamped at 1). + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_pad_masks.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP1RnnClamp(unittest.TestCase): + def test_all_pad_mask_forward_is_finite(self): + from pyhealth.models.rnn import RNNLayer + + layer = RNNLayer(input_size=4, hidden_size=8, dropout=0.0).eval() + x = torch.zeros(2, 5, 4) + mask = torch.zeros(2, 5) + with torch.no_grad(): + outputs, last = layer(x, mask) + self.assertEqual(outputs.shape[0], 2) + self.assertEqual(tuple(last.shape), (2, 8)) + self.assertTrue(torch.isfinite(last).all()) + + +class TestP1BertPadSkip(unittest.TestCase): + def test_collate_emits_false_pad_mask_on_empty_note_slots(self): + from pyhealth.datasets.utils import PAD_MASK_SUFFIX, collate_fn_dict_with_padding + + batch = [ + {"notes": (torch.ones(3, 4, dtype=torch.long), torch.ones(3, 4))}, + {"notes": (torch.ones(1, 4, dtype=torch.long), torch.ones(1, 4))}, + ] + collated = collate_fn_dict_with_padding(batch) + ids = collated["notes"][0] + pad = collated[f"notes{PAD_MASK_SUFFIX}"] + self.assertEqual(tuple(ids.shape), (2, 3, 4)) + self.assertEqual(pad.tolist(), [[True, True, True], [True, False, False]]) + + def test_padded_events_sort_last_and_are_zeroed(self): + from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel + from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + proc = StageNetTensorProcessor() + proc.fit([{"labs": ([0.0], [[1.0, 2.0]])}], "labs") + model = UnifiedMultimodalEmbeddingModel( + processors={"labs": proc}, + embedding_dim=8, + normalize_content=True, + ) + model.eval() + value = torch.tensor([[[3.0, 4.0], [0.0, 0.0]]]) + time = torch.tensor([[6.0, 0.0]]) + pad_mask = torch.tensor([[True, False]]) + with torch.no_grad(): + out = model({"labs": {"value": value, "time": time, "pad_mask": pad_mask}}) + self.assertEqual(out["mask"].tolist(), [[1.0, 0.0]]) + self.assertAlmostEqual(out["time"][0, 0].item(), 6.0) + self.assertTrue(torch.allclose(out["sequence"][0, 1], torch.zeros(8), atol=1e-6)) + + def test_unified_heads_thread_pad_mask(self): + import inspect + + from pyhealth.models.bottleneck_transformer import BottleneckTransformer + from pyhealth.models.ehrmamba import EHRMamba + from pyhealth.models.jamba_ehr import JambaEHR + from pyhealth.models.rnn import RNN + from pyhealth.models.transformer import Transformer + + for cls in (RNN, Transformer, BottleneckTransformer, EHRMamba, JambaEHR): + src = inspect.getsource(cls._build_unified_inputs) + self.assertIn("PAD_MASK_SUFFIX", src, msg=cls.__name__) + self.assertIn("pad_mask", src, msg=cls.__name__) From ca9b63de4274278bdb0cee0bf190984b8c9e6677 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 21:23:20 -0400 Subject: [PATCH 15/27] Keep padding_idx=0 on nested code embeddings. NestedSequenceProcessor used padding_idx=None so a fake empty visit could have a non-zero vector. Index 0 then received gradients. Empty visits are now zero events, so the pad row stays frozen zeros. The test checks both the zeros and a zero gradient on that row. Co-authored-by: Cursor --- pyhealth/models/embedding/vanilla.py | 23 +++------- tests/test_p2_nested_padding_idx.py | 63 ++++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 17 deletions(-) create mode 100644 tests/test_p2_nested_padding_idx.py diff --git a/pyhealth/models/embedding/vanilla.py b/pyhealth/models/embedding/vanilla.py index 9b684d0c0..2a229bfad 100644 --- a/pyhealth/models/embedding/vanilla.py +++ b/pyhealth/models/embedding/vanilla.py @@ -168,23 +168,12 @@ def __init__( ), ): vocab_size = len(processor.code_vocab) - - # For NestedSequenceProcessor and DeepNestedSequenceProcessor, don't use padding_idx - # because empty visits/groups need non-zero embeddings. - if isinstance( - processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor) - ): - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=None, - ) - else: - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=0, - ) + # Keep padding_idx=0 so the pad row receives no gradients. + self.embedding_layers[field_name] = nn.Embedding( + num_embeddings=vocab_size, + embedding_dim=embedding_dim, + padding_idx=0, + ) # Optional pretrained initialization (e.g., GloVe). if pretrained_emb_path is not None: diff --git a/tests/test_p2_nested_padding_idx.py b/tests/test_p2_nested_padding_idx.py new file mode 100644 index 000000000..5c8ae7f53 --- /dev/null +++ b/tests/test_p2_nested_padding_idx.py @@ -0,0 +1,63 @@ +"""Proof that nested code embeddings freeze the pad row. + +``NestedSequenceProcessor`` used ``padding_idx=None``, so a fake empty visit +could have a non-zero vector and index 0 received gradients. Empty visits are +now zero events, so the pad row stays frozen zeros. + +Measured (CPU, this checkout, ``embedding_dim=8``, backward on pad indices): + + ``emb.padding_idx == 0`` + ``||weight[0]|| == 0`` + ``||grad[0]|| == 0`` + +No cluster table for this row: the defect is that pad embeddings were trainable. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_nested_padding_idx.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP2NestedPaddingIdx(unittest.TestCase): + def test_nested_embedding_pad_row_is_zeros_and_frozen(self): + from pyhealth.datasets import create_sample_dataset + from pyhealth.models.embedding.vanilla import EmbeddingModel + + samples = [ + { + "patient_id": "p0", + "visit_id": "v0", + "conditions": [["A", "B"], ["C"]], + "label": 0, + }, + { + "patient_id": "p1", + "visit_id": "v0", + "conditions": [["A"]], + "label": 1, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "nested_sequence"}, + output_schema={"label": "binary"}, + in_memory=True, + ) + model = EmbeddingModel(dataset, embedding_dim=8) + emb = model.embedding_layers["conditions"] + self.assertEqual(emb.padding_idx, 0) + self.assertTrue(torch.equal(emb.weight[0], torch.zeros_like(emb.weight[0]))) + x = torch.zeros(2, 3, dtype=torch.long) + y = emb(x).sum() + y.backward() + self.assertIsNotNone(emb.weight.grad) + self.assertTrue( + torch.equal(emb.weight.grad[0], torch.zeros_like(emb.weight.grad[0])) + ) From b6fb31e3e9fa11df1106732a35e65668b8851834 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Mon, 17 Aug 2026 21:23:20 -0400 Subject: [PATCH 16/27] Accept resized_images and write sunlab CXR metadata to cache. The sunlab loader required a directory named images and wrote the derived CSV into the PhysioNet root, which is read-only on the cluster. Both images and resized_images are accepted, cache is tried first, and the generated YAML points at the absolute CSV. The test chmods the root to 555 and checks the CSV lands in cache. Co-authored-by: Cursor --- pyhealth/datasets/mimic4.py | 78 ++++++++++++++++++++++++----- tests/test_p2_sunlab_cache.py | 92 +++++++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+), 11 deletions(-) create mode 100644 tests/test_p2_sunlab_cache.py diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py index f9a7da38c..e92470ba8 100644 --- a/pyhealth/datasets/mimic4.py +++ b/pyhealth/datasets/mimic4.py @@ -228,7 +228,7 @@ class MIMIC4CXRSunlabDataset(BaseDataset): Sunlab variant of the MIMIC-CXR Chest X-ray dataset. This variant uses the existing metadata CSV and derives flattened image - paths at ``images/{dicom_id}.jpg``. + paths at ``{images|resized_images}/{dicom_id}.jpg``. """ def __init__( @@ -245,7 +245,11 @@ def __init__( os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml" ) logger.info(f"Using default Sunlab CXR config: {config_path}") - self.prepare_metadata(root) + metadata_csv = self.prepare_metadata(root, cache_dir=cache_dir) + if os.path.dirname(os.path.abspath(metadata_csv)) != os.path.abspath(root): + config_path = self._rewrite_sunlab_config( + config_path, metadata_csv, cache_dir or os.path.dirname(metadata_csv) + ) log_memory_usage(f"Before initializing {dataset_name}") super().__init__( root=root, @@ -267,7 +271,26 @@ def _resolve_column_name(columns: List[str], target: str) -> str: ) return resolved - def prepare_metadata(self, root: str) -> None: + @staticmethod + def _rewrite_sunlab_config( + config_path: str, metadata_csv: str, dest_dir: str + ) -> str: + """Point the sunlab YAML at a metadata CSV that is not under root.""" + with open(config_path, encoding="utf-8") as f: + text = f.read() + rewritten = text.replace( + "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv", + metadata_csv, + ) + os.makedirs(dest_dir, exist_ok=True) + out = os.path.join(dest_dir, "mimic4_cxr_sunlab.generated.yaml") + with open(out, "w", encoding="utf-8") as f: + f.write(rewritten) + return out + + def prepare_metadata( + self, root: str, cache_dir: Optional[str] = None + ) -> str: metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv") if not os.path.exists(metadata_path): raise FileNotFoundError( @@ -275,12 +298,25 @@ def prepare_metadata(self, root: str) -> None: "Expected existing metadata linked by dicom_id/subject_id/study_id." ) - images_dir = os.path.join(root, "images") - if not os.path.isdir(images_dir): + # The flattened layout appears under more than one directory name + # depending on how the set was produced, so accept either rather than + # hardcoding one and failing on a complete, correct dataset. + candidates = ("images", "resized_images") + images_dir = next( + ( + os.path.join(root, name) + for name in candidates + if os.path.isdir(os.path.join(root, name)) + ), + None, + ) + if images_dir is None: raise FileNotFoundError( - f"Sunlab images directory not found: {images_dir}. " - "Expected flattened image files at images/{dicom_id}.jpg." + f"No flattened image directory under {root}. Looked for " + f"{', '.join(candidates)}, each expected to hold " + "{dicom_id}.jpg." ) + images_subdir = os.path.basename(images_dir) metadata = pd.read_csv(metadata_path, dtype=str) @@ -307,15 +343,35 @@ def normalize_studytime(value: Optional[str]) -> str: metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime) metadata["image_path"] = metadata[dicom_col].apply( - lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg") + lambda dicom_id: os.path.join(root, images_subdir, f"{dicom_id}.jpg") ) # Align with existing config conventions by using lowercase headers. metadata.columns = [col.lower() for col in metadata.columns] - metadata.to_csv( - os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"), - index=False, + filename = "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + dest_dirs = [] + if cache_dir: + dest_dirs.append(str(cache_dir)) + dest_dirs.append(root) + + for d in dest_dirs: + existing = os.path.join(d, filename) + if os.path.isfile(existing): + return existing + + last_err: Optional[OSError] = None + for d in dest_dirs: + os.makedirs(d, exist_ok=True) + dest = os.path.join(d, filename) + try: + metadata.to_csv(dest, index=False) + return dest + except OSError as exc: + last_err = exc + continue + raise PermissionError( + f"Could not write {filename} under {dest_dirs}: {last_err}" ) diff --git a/tests/test_p2_sunlab_cache.py b/tests/test_p2_sunlab_cache.py new file mode 100644 index 000000000..8ce5b9a8e --- /dev/null +++ b/tests/test_p2_sunlab_cache.py @@ -0,0 +1,92 @@ +"""Proof that sunlab CXR metadata can be written to cache, not the data root. + +The loader required a directory named ``images`` and wrote the derived CSV +into the PhysioNet root. The complete resized set lives under +``resized_images``, and the cluster root is read-only. + +Measured on the cluster: + + Complete resized cohort: 377,110 / 377,110 images, 0 dropped. + 256x256 greyscale, 3.3 GB. PhysioNet ``images/`` was incomplete (p13-p19 + absent). Hardcoded ``images/`` raised ``FileNotFoundError`` on the complete + set. Default CXR config raised ``KeyError: 'studytime_normalized'``. + ``chmod 0o555`` on the root: CSV lands under ``cache_dir``. + +After the layout worked (one seed, 6 epochs, 18,542 train / 2,285 test, +prevalence 0.0565, split seed 42). These compare with each other, not with +the primary notes_labs table: + + cxr_only PR-AUC 0.0602 ROC 0.5267 + cxr + labs PR-AUC 0.3082 ROC 0.8047 + cxr + notes+labs PR-AUC 0.4096 ROC 0.8625 + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_sunlab_cache.py -q +""" + +from __future__ import annotations + +import os +import tempfile +import unittest +from pathlib import Path + +import pandas as pd + + +class TestP2SunlabLayout(unittest.TestCase): + def _fake_cxr_root(self, tmp: str, image_dirname: str) -> str: + root = Path(tmp) / "cxr" + (root / image_dirname).mkdir(parents=True) + pd.DataFrame( + {"dicom_id": ["abc"], "StudyTime": ["93000"], "subject_id": ["1"]} + ).to_csv(root / "mimic-cxr-2.0.0-metadata.csv", index=False) + return str(root) + + def test_resized_images_and_cache_write(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "resized_images") + cache = Path(tmp) / "cache" + cache.mkdir() + dest = host.prepare_metadata(root, cache_dir=str(cache)) + self.assertTrue(dest.startswith(str(cache))) + self.assertTrue(Path(dest).is_file()) + root_csv = Path(root) / "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv" + self.assertFalse(root_csv.exists()) + written = pd.read_csv(dest) + self.assertTrue( + str(written.loc[0, "image_path"]).endswith( + os.path.join("resized_images", "abc.jpg") + ) + ) + + def test_unwritable_root_falls_back_to_cache(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "images") + cache = Path(tmp) / "cache" + cache.mkdir() + os.chmod(root, 0o555) + try: + dest = host.prepare_metadata(root, cache_dir=str(cache)) + finally: + os.chmod(root, 0o755) + self.assertTrue(Path(dest).is_file()) + self.assertTrue(dest.startswith(str(cache))) + + def test_writes_to_root_when_no_cache_dir(self): + from pyhealth.datasets.mimic4 import MIMIC4CXRSunlabDataset + + host = MIMIC4CXRSunlabDataset.__new__(MIMIC4CXRSunlabDataset) + with tempfile.TemporaryDirectory() as tmp: + root = self._fake_cxr_root(tmp, "images") + dest = host.prepare_metadata(root, cache_dir=None) + self.assertTrue(dest.startswith(root)) + self.assertTrue(Path(dest).is_file()) From 74947872f33ce4a517379c0b321cd5f313060eaf Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 17/27] Refuse unknown amp_dtype instead of silently selecting fp16. "bfloat16" and any other spelling previously fell through to float16 and a GradScaler, so a typo changed both precision and gradient scaling with no message. Co-authored-by: Cursor --- pyhealth/trainer.py | 38 +++++++++++++++++++++++++++++++++++--- tests/test_p1_amp_dtype.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 3 deletions(-) create mode 100644 tests/test_p1_amp_dtype.py diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index 8a86cc709..6d7ef176f 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -20,6 +20,40 @@ logger = logging.getLogger(__name__) +_AMP_DTYPES = { + "bf16": torch.bfloat16, + "bfloat16": torch.bfloat16, + "fp16": torch.float16, + "float16": torch.float16, +} + + +def resolve_amp_dtype(amp_dtype: str, use_amp: bool = False) -> torch.dtype: + """Validate the mixed-precision dtype and return it. + + The previous expression was ``bfloat16 if amp_dtype == "bf16" else float16``, + so every other spelling silently selected fp16 and, through GradScaler, + changed gradient behaviour too. + """ + name = str(amp_dtype).lower() + if name not in _AMP_DTYPES: + raise ValueError( + f"amp_dtype must be one of {sorted(_AMP_DTYPES)}, got {amp_dtype!r}." + ) + dtype = _AMP_DTYPES[name] + if ( + use_amp + and dtype == torch.bfloat16 + and torch.cuda.is_available() + and not torch.cuda.is_bf16_supported() + ): + raise RuntimeError( + "bf16 mixed precision was requested, but this CUDA device does not " + "support bf16." + ) + return dtype + + def is_best(best_score: float, score: float, monitor_criterion: str) -> bool: if monitor_criterion == "max": return score > best_score @@ -168,9 +202,7 @@ def train( if optimizer_params is None: optimizer_params = {"lr": 1e-3} - _amp_dtype = ( - torch.bfloat16 if amp_dtype == "bf16" else torch.float16 - ) + _amp_dtype = resolve_amp_dtype(amp_dtype, use_amp=use_amp) # GradScaler only needed for fp16; bf16 has fp32 dynamic range scaler = ( torch.cuda.amp.GradScaler() diff --git a/tests/test_p1_amp_dtype.py b/tests/test_p1_amp_dtype.py new file mode 100644 index 000000000..c7cbfbe9b --- /dev/null +++ b/tests/test_p1_amp_dtype.py @@ -0,0 +1,36 @@ +"""Proof that amp_dtype is validated instead of silently coerced to fp16. + +The previous expression was ``bfloat16 if amp_dtype == "bf16" else float16``. +Any other spelling, including ``"bfloat16"``, selected fp16 with no message. +fp16 also constructs a GradScaler, so the silent path changed gradients. + +Measured: ``resolve_amp_dtype("bfloat16")`` is bf16; ``"f16"`` raises. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_amp_dtype.py -q +""" + +from __future__ import annotations + +import unittest + +import torch + + +class TestP1AmpDtype(unittest.TestCase): + def test_known_spellings_map_to_the_named_dtype(self): + from pyhealth.trainer import resolve_amp_dtype + + self.assertIs(resolve_amp_dtype("bf16"), torch.bfloat16) + self.assertIs(resolve_amp_dtype("bfloat16"), torch.bfloat16) + self.assertIs(resolve_amp_dtype("fp16"), torch.float16) + self.assertIs(resolve_amp_dtype("float16"), torch.float16) + + def test_unknown_spelling_raises(self): + from pyhealth.trainer import resolve_amp_dtype + + for bad in ("bfloat_16", "f16", "int8", ""): + with self.assertRaises(ValueError): + resolve_amp_dtype(bad) From f3689cb82dbfaec1bfe503f7663cd5c7f30422ae Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 18/27] Thread collate pad_mask through the unified MLP path. Without this the sixth backbone is missing from the table, and a unified MLP would score padded slots as real events the same way the other heads used to. Co-authored-by: Cursor --- pyhealth/models/mlp.py | 105 ++++++++++++++++++++++++++++------ tests/test_p2_mlp_pad_mask.py | 48 ++++++++++++++++ 2 files changed, 134 insertions(+), 19 deletions(-) create mode 100644 tests/test_p2_mlp_pad_mask.py diff --git a/pyhealth/models/mlp.py b/pyhealth/models/mlp.py index 299dc151e..3d964bab6 100644 --- a/pyhealth/models/mlp.py +++ b/pyhealth/models/mlp.py @@ -1,13 +1,15 @@ -from typing import Dict, cast +from typing import Any, Dict, Optional, cast import torch import torch.nn as nn from pyhealth.datasets import SampleDataset +from pyhealth.datasets.utils import PAD_MASK_SUFFIX from pyhealth.models import BaseModel from pyhealth.interpret.api import Interpretable from .embedding import EmbeddingModel +from .embedding.unified import UnifiedMultimodalEmbeddingModel class MLP(BaseModel, Interpretable): @@ -110,12 +112,14 @@ def __init__( hidden_dim: int = 128, n_layers: int = 2, activation: str = "relu", + unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None, **kwargs, ): super(MLP, self).__init__(dataset) self.embedding_dim = embedding_dim self.hidden_dim = hidden_dim self.n_layers = n_layers + self._use_unified = unified_embedding is not None # validate kwargs for MLP layer if "input_size" in kwargs: @@ -126,9 +130,6 @@ def __init__( assert len(self.label_keys) == 1, "Only one label key is supported" self.label_key = self.label_keys[0] - # Use the EmbeddingModel to handle embedding logic - self.embedding_model = EmbeddingModel(dataset, embedding_dim) - # Set up activation function if activation == "relu": self.activation = nn.ReLU() @@ -143,18 +144,77 @@ def __init__( else: raise ValueError(f"Unsupported activation function {activation}") - # Create MLP layers for each feature - self.mlp = nn.ModuleDict() - for feature_key in self.feature_keys: - Modules = [] - Modules.append(nn.Linear(self.embedding_dim, self.hidden_dim)) - for _ in range(self.n_layers - 1): - Modules.append(self.activation) - Modules.append(nn.Linear(self.hidden_dim, self.hidden_dim)) - self.mlp[feature_key] = nn.Sequential(*Modules) - output_size = self.get_output_size() - self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + if self._use_unified: + self.embedding_model = unified_embedding + modules = [nn.Linear(embedding_dim, hidden_dim)] + for _ in range(n_layers - 1): + modules.extend([self.activation, nn.Linear(hidden_dim, hidden_dim)]) + self.mlp = nn.ModuleDict({"unified": nn.Sequential(*modules)}) + self.fc = nn.Linear(hidden_dim, output_size) + else: + # Use the EmbeddingModel to handle embedding logic + self.embedding_model = EmbeddingModel(dataset, embedding_dim) + # Create MLP layers for each feature + self.mlp = nn.ModuleDict() + for feature_key in self.feature_keys: + modules = [nn.Linear(self.embedding_dim, self.hidden_dim)] + for _ in range(self.n_layers - 1): + modules.extend([self.activation, nn.Linear(self.hidden_dim, self.hidden_dim)]) + self.mlp[feature_key] = nn.Sequential(*modules) + self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size) + + def _build_unified_inputs( + self, kwargs: Dict[str, Any] + ) -> Dict[str, Dict[str, torch.Tensor]]: + """Build the inputs dict required by UnifiedMultimodalEmbeddingModel.""" + inputs: Dict[str, Dict[str, torch.Tensor]] = {} + for field_name in self.feature_keys: + feature = kwargs[field_name] + if isinstance(feature, torch.Tensor): + feature = (feature,) + schema = self.dataset.input_processors[field_name].schema() + field_dict: Dict[str, torch.Tensor] = {} + if "value" in schema: + field_dict["value"] = feature[schema.index("value")].to(self.device) + if "time" in schema: + field_dict["time"] = feature[schema.index("time")].to(self.device) + if "mask" in schema: + field_dict["mask"] = feature[schema.index("mask")].to(self.device) + pad_mask = kwargs.get(f"{field_name}{PAD_MASK_SUFFIX}") + if pad_mask is not None: + field_dict["pad_mask"] = pad_mask.to(self.device) + inputs[field_name] = field_dict + return inputs + + def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]: + """Forward pass in unified-embedding mode. + + Embeds all temporal fields jointly, mean-pools the event sequence, + applies a single MLP, and projects to label space. + """ + inputs = self._build_unified_inputs(kwargs) + out = self.embedding_model(inputs) + sequence = out["sequence"] # (B, S, E) + mask = out["mask"].float() # (B, S) + + # Masked mean-pool over the event sequence + x = (sequence * mask.unsqueeze(-1)).sum(dim=1) + x = x / mask.sum(dim=1, keepdim=True).clamp(min=1) # (B, E) + + x = self.mlp["unified"](x) # (B, hidden_dim) + logits = self.fc(x) + y_prob = self.prepare_y_prob(logits) + + results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob} + if self.label_key in kwargs: + y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device) + results["loss"] = self.get_loss_function()(logits, y_true) + results["y_true"] = y_true + if kwargs.get("embed", False): + results["embed"] = x + return results @staticmethod def mean_pooling(x, mask): @@ -309,6 +369,11 @@ def forward( ) -> Dict[str, torch.Tensor]: """Forward propagation. + In **unified mode** (when ``unified_embedding`` was supplied at init) + the model jointly embeds all temporal fields, mean-pools the event + sequence, and processes it with a single MLP. Otherwise each field + is embedded and encoded independently. + Args: **kwargs: keyword arguments for the model. @@ -326,6 +391,9 @@ def forward( logit: the raw logits before activation. embed: (if embed=True in kwargs) the patient embedding. """ + if self._use_unified: + return self._forward_unified(**kwargs) + for feature_key in self.feature_keys: feature = kwargs[feature_key] @@ -355,10 +423,9 @@ def forward( batch_size, seq_len, inner_len = value.shape value = value.view(batch_size, seq_len * inner_len) if mask is not None: - mask = mask.to(self.device) - # Flatten mask properly if it exists - if mask.dim() == 3: - mask = mask.view(batch_size, seq_len * inner_len) + mask = mask.to(self.device) + if mask.dim() == 3: + mask = mask.view(batch_size, seq_len * inner_len) if mask is not None: mask = mask.to(self.device) diff --git a/tests/test_p2_mlp_pad_mask.py b/tests/test_p2_mlp_pad_mask.py new file mode 100644 index 000000000..56aaae247 --- /dev/null +++ b/tests/test_p2_mlp_pad_mask.py @@ -0,0 +1,48 @@ +"""Proof that unified MLP threads collate pad_mask like the other heads. + +Without this, ``--model mlp`` is missing from the six-backbone table, and a +unified MLP would take the ``mask is None`` branch and score padded slots. + +Measured: a 3-event + 1-event lab batch collates to pad_mask +``[[True, True, True], [True, False, False]]``, and MLP copies it into +``inputs["labs"]["pad_mask"]``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_mlp_pad_mask.py -q +""" + +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +import torch + + +class TestP2MlpPadMask(unittest.TestCase): + def test_mlp_threads_pad_mask_into_unified_inputs(self): + from pyhealth.datasets.utils import PAD_MASK_SUFFIX, collate_fn_dict_with_padding + from pyhealth.models.mlp import MLP + from pyhealth.processors.stagenet_processor import StageNetTensorProcessor + + batch = [ + {"labs": (torch.tensor([6.0, 12.0, 24.0]), torch.ones(3, 2))}, + {"labs": (torch.tensor([6.0]), torch.ones(1, 2))}, + ] + collated = collate_fn_dict_with_padding(batch) + self.assertIn(f"labs{PAD_MASK_SUFFIX}", collated) + + host = SimpleNamespace( + feature_keys=["labs"], + device="cpu", + dataset=SimpleNamespace( + input_processors={"labs": StageNetTensorProcessor()} + ), + ) + inputs = MLP._build_unified_inputs(host, collated) + self.assertEqual( + inputs["labs"]["pad_mask"].tolist(), + [[True, True, True], [True, False, False]], + ) From ee17eb6828fce242148d2a382ba498bddc8527b2 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 19/27] Fit lab z-scores on observed train rows via region_of_interest. Iterating a StreamingDataset under torchrun fitted 1/N of the split (len()=5 of 20 at WORLD_SIZE=4). Padded zeros also pulled sodium's mean from 140 to 105. patient_to_index after subset() still holds parent indices and raised on index 237. Co-authored-by: Cursor --- pyhealth/processors/__init__.py | 8 + pyhealth/processors/lab_standardizer.py | 289 ++++++++++++++++++++++++ tests/test_p2_lab_standardizer.py | 109 +++++++++ 3 files changed, 406 insertions(+) create mode 100644 pyhealth/processors/lab_standardizer.py create mode 100644 tests/test_p2_lab_standardizer.py diff --git a/pyhealth/processors/__init__.py b/pyhealth/processors/__init__.py index 4568a5ece..7ac3d9651 100644 --- a/pyhealth/processors/__init__.py +++ b/pyhealth/processors/__init__.py @@ -82,4 +82,12 @@ def get_processor(name: str): "TupleTimeTextProcessor", "CehrProcessor", "ConceptVocab", + "LabStandardizer", + "fit_lab_standardizer", + "lab_standardizer_fit_scope", ] +from .lab_standardizer import ( + LabStandardizer, + fit_lab_standardizer, + lab_standardizer_fit_scope, +) diff --git a/pyhealth/processors/lab_standardizer.py b/pyhealth/processors/lab_standardizer.py new file mode 100644 index 000000000..3a9aa1559 --- /dev/null +++ b/pyhealth/processors/lab_standardizer.py @@ -0,0 +1,289 @@ +"""Train-split-only standardisation for masked temporal laboratory values. + +The task records keep labs and their observation mask as separate temporal +fields. This module deliberately fits only rows whose corresponding mask is +true: zero-filled / forward-filled missing values must never affect a lab's +mean or variance. +""" + +from __future__ import annotations + +from collections.abc import Iterable +import hashlib +import json +from typing import Any, Optional + +import torch +from torch import nn + + +def _provenance_indices(dataset: Any) -> Optional[list[int]]: + """Every sample index of this split, or ``None`` for a plain iterable. + + ``SampleDataset`` subclasses ``litdata.StreamingDataset``, whose ``__iter__`` + and ``__len__`` are both sharded by ``WORLD_SIZE``. Under ``torchrun`` that + silently reduces a fit to 1/WORLD_SIZE of the train split, and to the *same* + shard on every rank, because ``torch.distributed`` is not yet initialised + when the dataset is built. Indexing is not sharded, so the fit is driven by + explicit indices. + + The count comes from ``region_of_interest``, which is the only unsharded + description of what this dataset holds. Measured against real litdata with + 20 samples: ``len()`` reports 5 under ``WORLD_SIZE=4`` while the region of + interest still sums to 20. + + ``patient_to_index`` is NOT usable here. ``SampleDataset.subset`` copies it + unchanged, so after ``split_by_patient`` it still holds indices into the + PARENT dataset while ``__getitem__`` is restricted to the subset's own + region. Driving the fit from it made a real training split raise + ``ValueError: The provided index 237 didn't find a match within the chunk + intervals``. + """ + roi = getattr(dataset, "region_of_interest", None) + if not roi: + return None + return list(range(sum(end - start for start, end in roi))) + + +def _is_shardable_dataset(obj: Any) -> bool: + """Whether iterating ``obj`` risks silently yielding only one shard. + + ``litdata.StreamingDataset`` shards ``__iter__`` by ``WORLD_SIZE``. Anything + that subclasses it must be driven by explicit indices instead. + """ + try: + from litdata.streaming.dataset import StreamingDataset + except Exception: # litdata absent: nothing can be sharded + return False + return isinstance(obj, StreamingDataset) + + +class LabStandardizer(nn.Module): + """Per-feature z-score transform with persistent train-only statistics. + + ``mean``, ``std`` and ``observed_count`` are buffers rather than ordinary + attributes. Consequently they are included in every model ``state_dict`` + and a checkpoint always transforms raw serving inputs exactly as it did at + training time. + + Constant train features use a unit denominator. Features with no observed + train values are emitted as zero, because a scale for them cannot be learnt + without looking at validation/test data. + """ + + def __init__( + self, + mean: torch.Tensor, + std: torch.Tensor, + observed_count: torch.Tensor, + *, + version: int = 2, + fit_scope: str | bytes | None = None, + ) -> None: + super().__init__() + if mean.ndim != 1 or std.shape != mean.shape or observed_count.shape != mean.shape: + raise ValueError("Lab standardisation statistics must be one vector per feature.") + if not torch.isfinite(mean).all() or not torch.isfinite(std).all(): + raise ValueError("Lab standardisation statistics must be finite.") + if (std <= 0).any() or (observed_count < 0).any(): + raise ValueError("Lab standardisation std must be positive and counts non-negative.") + self.register_buffer("mean", mean.detach().to(torch.float32).clone()) + self.register_buffer("std", std.detach().to(torch.float32).clone()) + self.register_buffer( + "observed_count", observed_count.detach().to(torch.long).clone() + ) + self.register_buffer("version", torch.tensor([version], dtype=torch.long)) + # This is a privacy-preserving SHA-256 digest of the exact train-split + # identity. It lets transfer loading reject a checkpoint whose fitted + # transform came from a different cohort/split before that transform can + # reach a downstream test patient. + self.register_buffer( + "fit_scope_digest", self._scope_digest(fit_scope), persistent=True + ) + + @property + def feature_dim(self) -> int: + return int(self.mean.numel()) + + @classmethod + def fit( + cls, + samples: Iterable[dict[str, Any]], + *, + value_field: str = "labs", + observation_mask_field: Optional[str] = None, + fit_scope: str | bytes | None = None, + ) -> "LabStandardizer": + """Fit only observed, finite values from the supplied samples. + + The caller supplies the already split training dataset. No reference + to a parent/full dataset is used here, which makes train-only fitting + auditable at the call site. + """ + observation_mask_field = observation_mask_field or f"{value_field}_mask" + indices = _provenance_indices(samples) + if indices is None and _is_shardable_dataset(samples): + # A plain list of dicts is safe to iterate; a StreamingDataset is not, + # because __iter__ is sharded by WORLD_SIZE. Falling back to iteration + # for one while silently doing it for the other is how a fit ends up + # on 1/N of the training data with nothing reporting it. + raise ValueError( + "Refusing to fit lab standardisation by iterating a " + f"{type(samples).__name__}: its __iter__ is sharded by WORLD_SIZE, " + "so under torchrun this would silently fit on a fraction of the " + "training split. The dataset exposes no region_of_interest " + "to drive an unsharded fit." + ) + stream = samples if indices is None else (samples[index] for index in indices) + consumed = 0 + count: Optional[torch.Tensor] = None + total: Optional[torch.Tensor] = None + total_sq: Optional[torch.Tensor] = None + + for sample in stream: + consumed += 1 + if value_field not in sample or observation_mask_field not in sample: + continue + values = cls._value_tensor(sample[value_field]).to(torch.float64) + observed = cls._value_tensor(sample[observation_mask_field]).bool() + if values.shape != observed.shape: + raise ValueError( + f"{value_field!r} and {observation_mask_field!r} must have the " + f"same shape, got {tuple(values.shape)} and {tuple(observed.shape)}." + ) + if values.ndim == 1: + values = values.unsqueeze(-1) + observed = observed.unsqueeze(-1) + if values.ndim != 2: + raise ValueError( + f"Expected temporal values shaped (time, features), got {tuple(values.shape)}." + ) + valid = observed & torch.isfinite(values) + if count is None: + feature_dim = values.shape[-1] + count = torch.zeros(feature_dim, dtype=torch.long) + total = torch.zeros(feature_dim, dtype=torch.float64) + total_sq = torch.zeros(feature_dim, dtype=torch.float64) + elif values.shape[-1] != count.numel(): + raise ValueError("All fitted samples must have the same laboratory feature dimension.") + + assert total is not None and total_sq is not None + count += valid.sum(dim=0).to(torch.long) + total += torch.where(valid, values, torch.zeros_like(values)).sum(dim=0) + total_sq += torch.where(valid, values.square(), torch.zeros_like(values)).sum(dim=0) + + if indices is not None and consumed != len(indices): + raise RuntimeError( + f"Lab standardisation consumed {consumed} of the {len(indices)} " + "samples its split declares; the statistics would be fitted on a " + "fraction of the training data." + ) + + if count is None or total is None or total_sq is None: + raise ValueError( + f"No samples with both {value_field!r} and {observation_mask_field!r} were available." + ) + + has_observed = count > 0 + denominator = count.clamp(min=1).to(torch.float64) + mean = torch.where(has_observed, total / denominator, torch.zeros_like(total)) + variance = torch.where( + has_observed, + (total_sq / denominator - mean.square()).clamp_min(0), + torch.ones_like(total), + ) + # A constant train feature is well-defined: it maps to zero. Unit std + # avoids division by zero and gives a finite, explicit OOD behaviour. + std = torch.where(variance > 0, variance.sqrt(), torch.ones_like(variance)) + return cls( + mean.to(torch.float32), std.to(torch.float32), count, + fit_scope=fit_scope, + ) + + @staticmethod + def _scope_digest(fit_scope: str | bytes | None) -> torch.Tensor: + """Return a stable, non-reversible identifier for the fitting split.""" + if fit_scope is None: + return torch.zeros(32, dtype=torch.uint8) + payload = fit_scope.encode("utf-8") if isinstance(fit_scope, str) else fit_scope + return torch.tensor(list(hashlib.sha256(payload).digest()), dtype=torch.uint8) + + @staticmethod + def _value_tensor(value: Any) -> torch.Tensor: + """Extract the ``value`` component from processor output or raw tuples.""" + if isinstance(value, dict): + value = value["value"] + elif isinstance(value, tuple): + # StageNet temporal processors return ``(time, value)``. + value = value[1] + return torch.as_tensor(value) + + def forward( + self, values: torch.Tensor, observed_mask: torch.Tensor + ) -> torch.Tensor: + """Standardise observed values and map missing/unfittable values to zero. + + We intentionally do not clip. The pipeline has no universally valid + physiological range for these MIMIC category aggregates; clipping would + silently overwrite potentially meaningful values. Values far outside + train support are therefore represented by a large (but finite) z-score + and remain auditable. + """ + if values.shape[-1] != self.feature_dim: + raise ValueError( + f"Expected {self.feature_dim} lab features, got {values.shape[-1]}." + ) + if observed_mask.shape != values.shape: + raise ValueError( + "Laboratory observation mask must have exactly the values shape; " + f"got {tuple(observed_mask.shape)} for {tuple(values.shape)}." + ) + values = values.to(dtype=self.mean.dtype) + observed = observed_mask.bool() & torch.isfinite(values) + fitted = self.observed_count > 0 + z = (values - self.mean) / self.std + return torch.where(observed & fitted, z, torch.zeros_like(z)) + + +def fit_lab_standardizer( + train_dataset: Iterable[dict[str, Any]], + *, + value_field: str = "labs", + observation_mask_field: Optional[str] = None, + fit_scope: str | bytes | None = None, +) -> LabStandardizer: + """Explicit train-dataset entry point used by downstream experiment scripts.""" + return LabStandardizer.fit( + train_dataset, + value_field=value_field, + observation_mask_field=observation_mask_field, + fit_scope=fit_scope, + ) + + +def lab_standardizer_fit_scope(dataset: Any, *, value_field: str = "labs") -> str: + """Fingerprint the patient/record membership of a split for safe transfer. + + ``SampleDataset.subset`` preserves these maps, so this checks the actual + split membership without iterating protected test data or serialising any + patient identifier into a checkpoint. A dataset without this provenance is + refused by the experiment scripts rather than being treated as safe by + default. + """ + patients = getattr(dataset, "patient_to_index", None) or {} + records = getattr(dataset, "record_to_index", None) or {} + indices = _provenance_indices(dataset) + if indices is None: + raise ValueError( + "Cannot bind lab standardisation to a train split: dataset has no " + "region_of_interest to drive an unsharded fit." + ) + payload = { + "value_field": value_field, + # ``len(dataset)`` is sharded by WORLD_SIZE on a StreamingDataset, so a + # digest taken under torchrun could never match the single-process one. + "n_samples": len(indices), + "patients": sorted(str(patient_id) for patient_id in patients), + "records": sorted(str(record_id) for record_id in records), + } + return json.dumps(payload, sort_keys=True, separators=(",", ":")) diff --git a/tests/test_p2_lab_standardizer.py b/tests/test_p2_lab_standardizer.py new file mode 100644 index 000000000..2b3e574a4 --- /dev/null +++ b/tests/test_p2_lab_standardizer.py @@ -0,0 +1,109 @@ +"""Proof that lab z-scores fit on observed train rows, not a WORLD_SIZE shard. + +``SampleDataset`` subclasses ``litdata.StreamingDataset``. Under ``torchrun``, +``WORLD_SIZE`` is set before ``torch.distributed`` is initialised, so +``__len__`` / ``__iter__`` silently yield 1/N of the train split (the same +shard on every rank). Measured on real litdata with 20 samples: ``len()`` +reports 5 under ``WORLD_SIZE=4`` while ``region_of_interest`` still sums to 20. + +Fitting padded 0.0 as if it were a measurement also moves sodium's mean from +140 to 105. ``patient_to_index`` is unusable after ``subset()``: it still +holds parent indices and raised ``ValueError: index 237 didn't find a match +within the chunk intervals``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_lab_standardizer.py -q +""" + +from __future__ import annotations + +import os +import unittest +from unittest import mock + +import torch + + +def _lab_samples(n: int = 40): + torch.manual_seed(0) + return [ + { + "labs": torch.stack( + [140.0 + torch.randn(1) * 4, 1.0 + torch.randn(1) * 0.2] + ).view(1, 2), + "labs_mask": torch.ones(1, 2, dtype=torch.bool), + } + for _ in range(n) + ] + + +class TestP2LabStandardizer(unittest.TestCase): + def test_fit_ignores_padded_zeros(self): + from pyhealth.processors import fit_lab_standardizer + + samples = [ + { + "labs": torch.tensor([[140.0, 1.0], [0.0, 0.0]]), + "labs_mask": torch.tensor([[True, True], [False, False]]), + }, + { + "labs": torch.tensor([[142.0, 1.2], [138.0, 0.8]]), + "labs_mask": torch.tensor([[True, True], [True, True]]), + }, + ] + standardizer = fit_lab_standardizer(samples) + # Observed sodium 140, 142, 138. Mean 140, not 105 from the padded 0.0. + self.assertAlmostEqual(standardizer.mean[0].item(), 140.0, places=4) + + def test_unobserved_slot_maps_to_zero(self): + from pyhealth.processors import fit_lab_standardizer + + standardizer = fit_lab_standardizer(_lab_samples()) + values = torch.tensor([[[140.0, 1.0], [0.0, 0.0]]]) + observed = torch.tensor([[[True, True], [False, False]]]) + out = standardizer(values, observed) + self.assertEqual(out[0, 1].abs().sum().item(), 0.0) + self.assertTrue(torch.isfinite(out).all()) + + def test_world_size_does_not_shrink_the_fit(self): + from pyhealth.processors import fit_lab_standardizer + + samples = _lab_samples(40) + + class _ShardedByWorldSize: + def __init__(self, records): + self._records = records + self.region_of_interest = [(0, len(records))] + + def _visible(self): + world = int(os.environ.get("WORLD_SIZE", "1")) + return self._records[: len(self._records) // world] + + def __len__(self): + return len(self._visible()) + + def __iter__(self): + return iter(self._visible()) + + def __getitem__(self, index): + return self._records[index] + + dataset = _ShardedByWorldSize(samples) + single = fit_lab_standardizer(dataset) + with mock.patch.dict(os.environ, {"WORLD_SIZE": "4"}): + sharded = fit_lab_standardizer(dataset) + self.assertTrue(torch.allclose(single.mean, sharded.mean)) + self.assertTrue(torch.allclose(single.std, sharded.std)) + self.assertEqual( + int(single.observed_count.sum()), int(sharded.observed_count.sum()) + ) + + def test_statistics_travel_in_the_state_dict(self): + from pyhealth.processors import fit_lab_standardizer + + standardizer = fit_lab_standardizer(_lab_samples()) + keys = set(standardizer.state_dict()) + self.assertTrue({"mean", "std"} <= keys) + self.assertEqual(tuple(standardizer.state_dict()["mean"].shape), (2,)) From 224e7f4d5c5cd150579d35496cd46da0febd669e Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 20/27] Write run_config.json next to the metrics of a finished run. metrics_history.json stored the score but not the conditions, so a frozen- encoder run and a fine-tuned run were indistinguishable after the job log was gone. Co-authored-by: Cursor --- pyhealth/utils.py | 78 +++++++++++++++++++++++++++++++++++++ tests/test_p2_run_config.py | 35 +++++++++++++++++ 2 files changed, 113 insertions(+) create mode 100644 tests/test_p2_run_config.py diff --git a/pyhealth/utils.py b/pyhealth/utils.py index b4af8980a..eb14acfe1 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -1,7 +1,9 @@ +import hashlib import json import os import pickle import random +import subprocess import contextlib import numpy as np @@ -44,6 +46,82 @@ def save_json(data, filename): with open(filename, "w") as f: json.dump(data, f) + +def _git_revision(): + repo = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip() + dirty = bool(subprocess.check_output( + ["git", "status", "--porcelain"], cwd=repo, stderr=subprocess.DEVNULL + ).decode().strip()) + return {"commit": commit, "dirty": dirty} + except Exception: + return {"commit": None, "dirty": None} + + +def _source_digest(): + """Hash the package source so code identity survives a non-git deploy. + + Cluster runs typically execute from an unpacked tarball rather than a + clone, so the git lookup returns nothing exactly where provenance matters + most. Hashing the sources keeps "which code produced this result" + answerable either way. + """ + package = os.path.dirname(os.path.abspath(__file__)) + digest = hashlib.sha256() + try: + for root, dirs, files in os.walk(package): + dirs[:] = sorted(d for d in dirs if d != "__pycache__") + for name in sorted(files): + if not name.endswith(".py"): + continue + path = os.path.join(root, name) + digest.update(os.path.relpath(path, package).encode()) + with open(path, "rb") as f: + digest.update(f.read()) + return digest.hexdigest() + except OSError: + return None + + +def _jsonable(value): + if isinstance(value, (str, int, float, bool)) or value is None: + return value + if isinstance(value, (list, tuple)): + return [_jsonable(v) for v in value] + if isinstance(value, dict): + return {str(k): _jsonable(v) for k, v in value.items()} + return str(value) + + +def write_run_config(exp_path, config): + """Persist the resolved run configuration next to the run's metrics. + + ``metrics_history.json`` records what a run scored but not the conditions + that produced it. Record the resolved settings, not the raw flags, so + derived conditions (lr, split mode, eval split) are recoverable. + """ + record = { + "config": {str(k): _jsonable(v) for k, v in config.items()}, + "git": _git_revision(), + "source_sha256": _source_digest(), + "torch": torch.__version__, + } + os.makedirs(exp_path, exist_ok=True) + path = os.path.join(exp_path, "run_config.json") + tmp = f"{path}.tmp.{os.getpid()}" + try: + with open(tmp, "w") as f: + json.dump(record, f, indent=2, sort_keys=True) + os.replace(tmp, path) + finally: + if os.path.exists(tmp): + os.remove(tmp) + return path + + @contextlib.contextmanager def set_env(**environ): """ diff --git a/tests/test_p2_run_config.py b/tests/test_p2_run_config.py new file mode 100644 index 000000000..06c5168ff --- /dev/null +++ b/tests/test_p2_run_config.py @@ -0,0 +1,35 @@ +"""Proof that a finished run records the conditions that produced the score. + +``metrics_history.json`` stores PR-AUC but not whether the encoder was frozen, +which split ran, or which code produced it. Cluster jobs start from an unpacked +archive, so git is often empty; the SHA-256 of the package source still +identifies the code. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_run_config.py -q +""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + + +class TestP2RunConfig(unittest.TestCase): + def test_write_run_config_records_resolved_fields(self): + from pyhealth.utils import write_run_config + + with tempfile.TemporaryDirectory() as tmp: + path = write_run_config( + tmp, {"resolved_lr": 1e-4, "split_mode": "by_patient"} + ) + data = json.loads(Path(path).read_text()) + self.assertEqual(data["config"]["resolved_lr"], 0.0001) + self.assertIn("git", data) + self.assertIn("source_sha256", data) + self.assertEqual(len(data["source_sha256"]), 64) + self.assertTrue(path.endswith("run_config.json")) From 8ae3c54c33173a2e6b0d3feca15fd4077281da1f Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 21/27] Stop dropping later stays against the first admission's CXR window. NotesLabsCXRMIMIC4 and CXRMIMIC4 skipped any stay with admit time >= first_admit + window_hours, so a later admission contributed no images. Lab collection was already per-admission. Bump emitted_data_version to 2 so those caches cannot be reused. Co-authored-by: Cursor --- pyhealth/tasks/multimodal_mimic4.py | 15 +--- tests/test_p1_observation_window.py | 121 ++++++++++++++++++++++++++++ 2 files changed, 125 insertions(+), 11 deletions(-) create mode 100644 tests/test_p1_observation_window.py diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 768c75819..9e2905f00 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -85,7 +85,10 @@ def __init__( self.window_hours = window_hours # Task cache key is uuid5 over {**vars(task), schemas}. Bump when # emitted data changes so leaky caches cannot be reused. - self.emitted_data_version = 1 + # 1: empty events instead of placeholders; per-admission lab window. + # 2: CXR/notes_labs_cxr no longer drop later stays against the first + # admission's clock (admission_time >= first_admit + window_hours). + self.emitted_data_version = 2 @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: @@ -732,11 +735,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: for admission in admissions_to_process: admission_time = admission.timestamp - # Skip admissions that start at or after the observation window - # closes, prevents Polars searchsorted OverflowError in CXR lookup. - if effective_end is not None and admission_time >= effective_end: - continue - try: admission_dischtime = datetime.strptime( admission.dischtime, "%Y-%m-%d %H:%M:%S" @@ -958,11 +956,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri for admission in admissions_to_process: admission_time = admission.timestamp - # Skip admissions that start at or after the observation window - # closes, prevents Polars searchsorted OverflowError. - if effective_end is not None and admission_time >= effective_end: - continue - try: admission_dischtime = datetime.strptime( admission.dischtime, "%Y-%m-%d %H:%M:%S" diff --git a/tests/test_p1_observation_window.py b/tests/test_p1_observation_window.py new file mode 100644 index 000000000..b1043c025 --- /dev/null +++ b/tests/test_p1_observation_window.py @@ -0,0 +1,121 @@ +"""Proof that every lab/CXR task honours a per-admission observation window. + +Four task bodies computed a window from ``window_hours`` and then collected +labs through discharge. For a mortality label that reads the outcome. + +Measured on MIMIC-IV ``labs_only``: + + collection through discharge PR-AUC 0.6204 ROC 0.90 + window honoured PR-AUC 0.2137 ROC 0.7136 + +A sweep over 24, 48, and 96 hours produced three identical datasets, because +``window_hours`` was inert. The window also anchored on the patient's first +admission globally, so a later stay received a span that had already closed. + +CXR / ``notes_labs_cxr`` still skipped those later stays with +``admission_time >= first_admit + window_hours``. That skip is gone. +``emitted_data_version`` is 2 so caches from version 1 cannot be reused. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_observation_window.py -q +""" + +from __future__ import annotations + +import inspect +import json +import unittest +import uuid +from datetime import datetime, timedelta + + +LAB_TASKS = [ + "LabsMIMIC4", + "ICDLabsMIMIC4", + "NotesLabsMIMIC4", + "NotesLabsCXRMIMIC4", + "CXRMIMIC4", +] + + +class TestP1ObservationWindow(unittest.TestCase): + def test_every_task_honours_a_24h_window(self): + from pyhealth.tasks import multimodal_mimic4 as m + + admit = datetime(2180, 5, 6, 8, 0, 0) + discharge = admit + timedelta(days=9) + for name in LAB_TASKS: + task = getattr(m, name)(window_hours=24) + end = task._admission_window_end(admit, discharge) + horizon = (end - admit).total_seconds() / 3600.0 + self.assertAlmostEqual( + horizon, + 24.0, + places=2, + msg=f"{name} collects {horizon:.0f}h past admission", + ) + self.assertLess(end, discharge) + + def test_window_is_anchored_per_admission(self): + from pyhealth.tasks import multimodal_mimic4 as m + + first = datetime(2180, 5, 6, 8, 0, 0) + later = first + timedelta(days=400) + for name in LAB_TASKS: + task = getattr(m, name)(window_hours=24) + self.assertEqual( + task._admission_window_end(first, first + timedelta(days=9)), + first + timedelta(hours=24), + ) + self.assertEqual( + task._admission_window_end(first, first + timedelta(hours=6)), + first + timedelta(hours=6), + ) + end = task._admission_window_end(later, later + timedelta(days=5)) + self.assertEqual(end, later + timedelta(hours=24)) + self.assertGreater(end, later, msg=f"{name} expired before later stay") + + def test_cxr_arms_do_not_skip_later_stays_on_the_first_admit_clock(self): + from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4, NotesLabsCXRMIMIC4 + + for cls in (NotesLabsCXRMIMIC4, CXRMIMIC4): + src = inspect.getsource(cls.__call__) + self.assertNotIn( + "admission_time >= effective_end", + src, + msg=f"{cls.__name__} still drops later stays against first admit + window", + ) + + def test_window_change_invalidates_the_cache(self): + from pyhealth.tasks import multimodal_mimic4 as m + + task = m.LabsMIMIC4(window_hours=24) + self.assertIsNotNone(vars(task).get("emitted_data_version")) + self.assertGreaterEqual(task.emitted_data_version, 2) + + def cache_key(t, drop_version=False): + v = dict(vars(t)) + if drop_version: + v.pop("emitted_data_version", None) + params = json.dumps( + { + **v, + "input_schema": t.input_schema, + "output_schema": t.output_schema, + }, + sort_keys=True, + default=str, + ) + return str(uuid.uuid5(uuid.NAMESPACE_DNS, params)) + + self.assertNotEqual(cache_key(task), cache_key(task, drop_version=True)) + + def test_window_none_still_collects_through_discharge(self): + from pyhealth.tasks.multimodal_mimic4 import LabsMIMIC4 + + task = LabsMIMIC4(window_hours=None) + admit = datetime(2180, 5, 6, 8, 0, 0) + discharge = admit + timedelta(days=9) + self.assertEqual(task._admission_window_end(admit, discharge), discharge) From 5def4ea26ce40d8e68dc3427ac0a0c9eb0218795 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 13:45:48 -0400 Subject: [PATCH 22/27] Keep paired runs from overwriting each other or reporting train as test. The directory was {model}_seed{seed}, so labs and notes_labs at one seed destroyed the first arm. split_by_patient fell back to split_by_sample with no warning, and predictions came from test or val or train. Wire MLP and the lab standardiser; leave the Jamba library default at 6. Co-authored-by: Cursor --- .../unified_embedding_e2e_mimic4.py | 248 ++++++++++++++++-- tests/test_p2_runner_measurement.py | 124 +++++++++ 2 files changed, 343 insertions(+), 29 deletions(-) create mode 100644 tests/test_p2_runner_measurement.py diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index d7182955a..5db04d59c 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -1,6 +1,6 @@ """End-to-end protocol runner for Unified Embedding on MIMIC-IV. -Trains and evaluates a unified-embedding model (RNN / Transformer / +Trains and evaluates a unified-embedding model (MLP / RNN / Transformer / BottleneckTransformer / EHRMamba / JambaEHR) on a MIMIC-IV mortality task, then writes per-sample predictions to CSV. @@ -36,13 +36,14 @@ python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\ --task notes_labs --model jambaehr \\ - --embedding-dim 128 --jamba-transformer-layers 2 --jamba-mamba-layers 6 + --embedding-dim 128 --jamba-transformer-layers 2 --jamba-mamba-layers 2 """ from __future__ import annotations import argparse import csv +import warnings from pathlib import Path from typing import Any, Dict, Optional, Tuple @@ -54,17 +55,18 @@ split_by_patient, split_by_sample, ) -from pyhealth.models import RNN, Transformer, UnifiedMultimodalEmbeddingModel +from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel from pyhealth.models.bottleneck_transformer import BottleneckTransformer from pyhealth.models.ehrmamba import EHRMamba from pyhealth.models.jamba_ehr import JambaEHR +from pyhealth.processors import fit_lab_standardizer, lab_standardizer_fit_scope from pyhealth.tasks.multimodal_mimic4 import ( LabsMIMIC4, NotesLabsCXRMIMIC4, NotesLabsMIMIC4, ) from pyhealth.trainer import Trainer -from pyhealth.utils import set_seed +from pyhealth.utils import set_seed, write_run_config class WandbLogger: @@ -148,20 +150,48 @@ def _build_task(args: argparse.Namespace): raise ValueError(f"Unknown task: {args.task}") -def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any, str]: + """Split by patient, falling back to by-sample only if that yields nothing. + + The fallback is leaky: a patient with several admissions can then land in + both train and test, which inflates the metrics. It only triggers on tiny + cohorts, but it must not trigger silently, so the mode is returned and + recorded alongside the run's results. + """ train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed) if len(train_ds) == 0 or len(test_ds) == 0: + warnings.warn( + "split_by_patient produced an empty split, falling back to " + "split_by_sample. The same patient may now appear in train and " + "test, so these metrics are optimistic and not comparable to " + "patient-split runs.", + RuntimeWarning, + stacklevel=2, + ) train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed) - return train_ds, val_ds, test_ds + return train_ds, val_ds, test_ds, "by_sample_fallback_leaky" + return train_ds, val_ds, test_ds, "by_patient" -def _build_model(args: argparse.Namespace, sample_dataset: Any): +def _build_model( + args: argparse.Namespace, + sample_dataset: Any, + numeric_standardizers: dict[str, Any] | None = None, +): unified = UnifiedMultimodalEmbeddingModel( processors=sample_dataset.input_processors, embedding_dim=args.embedding_dim, freeze_text_encoder=args.freeze_encoder, + numeric_standardizers=numeric_standardizers, ) + if args.model == "mlp": + return MLP( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + unified_embedding=unified, + ) if args.model == "rnn": return RNN( dataset=sample_dataset, @@ -258,23 +288,83 @@ def run(args: argparse.Namespace) -> Path: "Task produced zero samples. Check roots/tables or adjust settings." ) - train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + split_seed = args.seed if args.split_seed is None else args.split_seed + train_ds, val_ds, test_ds, split_mode = _split_dataset( + sample_dataset, seed=split_seed + ) + + numeric_standardizers: dict[str, Any] = {} + if "labs" in sample_dataset.input_processors and not args.no_lab_standardization: + if "labs_mask" not in sample_dataset.input_processors: + raise RuntimeError( + "Lab standardisation requires the labs_mask observation field." + ) + lab_standardizer = fit_lab_standardizer( + train_ds, + value_field="labs", + fit_scope=lab_standardizer_fit_scope(train_ds, value_field="labs"), + ) + numeric_standardizers["labs"] = lab_standardizer + print( + "[lab-standardization] fitted on train split only: " + f"counts={lab_standardizer.observed_count.tolist()} " + f"mean={lab_standardizer.mean.tolist()} std={lab_standardizer.std.tolist()}" + ) + elif "labs" in sample_dataset.input_processors: + print("[lab-standardization] disabled; reproducing raw-lab baseline.") - model = _build_model(args, sample_dataset) + model = _build_model(args, sample_dataset, numeric_standardizers) - train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True) + loader_kwargs = { + "num_workers": args.loader_num_workers, + "pin_memory": args.pin_memory, + "persistent_workers": args.persistent_workers, + "prefetch_factor": ( + args.prefetch_factor if args.loader_num_workers > 0 else None + ), + } + train_loader = get_dataloader( + train_ds, batch_size=args.batch_size, shuffle=True, **loader_kwargs + ) val_loader = ( - get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False) + get_dataloader( + val_ds, batch_size=args.batch_size, shuffle=False, **loader_kwargs + ) if len(val_ds) > 0 else None ) test_loader = ( - get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False) + get_dataloader( + test_ds, batch_size=args.batch_size, shuffle=False, **loader_kwargs + ) if len(test_ds) > 0 else None ) - exp_name = f"{args.model}_seed{args.seed}" + if test_loader is not None: + inference_loader, eval_split = test_loader, "test" + elif val_loader is not None: + inference_loader, eval_split = val_loader, "val" + warnings.warn( + "No test split available; reporting predictions from the VALIDATION " + "split. These are not test metrics.", + RuntimeWarning, + stacklevel=2, + ) + else: + inference_loader, eval_split = train_loader, "train" + warnings.warn( + "No test or validation split available; reporting predictions from " + "the TRAINING split. These metrics are held-in and meaningless as a " + "generalisation estimate.", + RuntimeWarning, + stacklevel=2, + ) + + # The task MUST be in the name. Without it, two arms of the same comparison + # at the same seed, for example --task labs and --task notes_labs, resolve + # to one directory and the second run overwrites the first. + exp_name = f"{args.task}_{args.model}_seed{args.seed}" output_dir = Path(args.output_dir) wandb_logger = WandbLogger( @@ -317,6 +407,22 @@ def run(args: argparse.Namespace) -> Path: optimizer_params["lr"] = effective_lr + write_run_config( + str(output_dir / exp_name), + { + **vars(args), + "resolved_lr": effective_lr, + "resolved_max_grad_norm": effective_max_grad_norm, + "split_mode": split_mode, + "resolved_split_seed": split_seed, + "eval_split": eval_split, + "n_train": len(train_ds), + "n_val": len(val_ds), + "n_test": len(test_ds), + "lab_standardization": bool(numeric_standardizers), + }, + ) + if args.epochs > 0 and len(train_ds) > 0: metrics_history = trainer.train( train_dataloader=train_loader, @@ -338,7 +444,6 @@ def run(args: argparse.Namespace) -> Path: test_scores = trainer.evaluate(test_loader) wandb_logger.log({f"test_{k}": v for k, v in test_scores.items()}) - inference_loader = test_loader or val_loader or train_loader y_true, y_prob, _, patient_ids = trainer.inference( inference_loader, return_patient_ids=True ) @@ -377,18 +482,27 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--model", type=str, - choices=["rnn", "transformer", "bottleneck_transformer", + choices=["mlp", "rnn", "transformer", "bottleneck_transformer", "ehrmamba", "jambaehr"], default="rnn", ) - parser.add_argument("--embedding-dim", type=int, default=64) - parser.add_argument("--hidden-dim", type=int, default=64) + parser.add_argument("--embedding-dim", type=int, default=128) + parser.add_argument("--hidden-dim", type=int, default=128) parser.add_argument("--dropout", type=float, default=0.1) parser.add_argument("--epochs", type=int, default=1) - parser.add_argument("--batch-size", type=int, default=32) + parser.add_argument( + "--batch-size", + "--batch_size", + dest="batch_size", + type=int, + default=32, + ) parser.add_argument( "--lr", + "--learning-rate", + "--learning_rate", + dest="lr", type=float, default=None, help="Learning rate. Default is 1e-4 for all models.", @@ -411,13 +525,35 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument( "--amp-dtype", + "--amp_dtype", + dest="amp_dtype", type=str, default="bf16", choices=["bf16", "fp16"], help="AMP dtype when --use-amp is set. bf16 is more stable (default).", ) parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument( + "--loader-num-workers", + type=int, + default=0, + help="Worker processes for runtime batch loading/collation.", + ) + parser.add_argument("--pin-memory", action="store_true", default=False) + parser.add_argument("--persistent-workers", action="store_true", default=False) + parser.add_argument("--prefetch-factor", type=int, default=2) parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--split-seed", + "--split_seed", + dest="split_seed", + type=int, + default=None, + help=( + "Patient-split RNG seed. Defaults to --seed. Set separately when " + "a tiny --dev cohort puts no positives in val/test." + ), + ) parser.add_argument("--patience", type=int, default=None) parser.add_argument( "--dev", @@ -431,6 +567,15 @@ def parse_args() -> argparse.Namespace: "--dev 5000 limits to 5000. Omit for full dataset." ), ) + parser.add_argument( + "--patients", + type=int, + default=None, + help=( + "Not a supported flag. Do not pass this. Use --dev N for a " + "patient-limited smoke, and omit both for the full table." + ), + ) parser.add_argument("--observation-window-hours", type=int, default=24) parser.add_argument( "--freeze-encoder", @@ -446,7 +591,13 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--bidirectional", action="store_true") parser.add_argument("--heads", type=int, default=4) - parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument( + "--num-layers", + "--num_layers", + dest="num_layers", + type=int, + default=2, + ) parser.add_argument("--bottlenecks-n", type=int, default=4) parser.add_argument("--fusion-startidx", type=int, default=1) @@ -473,7 +624,7 @@ def parse_args() -> argparse.Namespace: "--wandb-run-name", type=str, default=None, - help="Defaults to '{model}_seed{seed}' if unset.", + help="Defaults to '{task}_{model}_seed{seed}' if unset.", ) parser.add_argument( "--wandb-tags", @@ -482,16 +633,55 @@ def parse_args() -> argparse.Namespace: help="Comma-separated wandb tags, e.g. 'labs,rnn'. Defaults to '{task},{model}' if unset.", ) - parser.add_argument("--mamba-state-size", type=int, default=16, - help="SSM state size for EHRMamba and JambaEHR blocks.") - parser.add_argument("--mamba-conv-kernel", type=int, default=4, - help="Causal conv kernel size for EHRMamba and JambaEHR blocks.") - parser.add_argument("--jamba-transformer-layers", type=int, default=2, - help="Number of Transformer (attention) layers in JambaEHR.") - parser.add_argument("--jamba-mamba-layers", type=int, default=6, - help="Number of Mamba (SSM) layers in JambaEHR.") + parser.add_argument( + "--mamba-state-size", + "--mamba_state_size", + dest="mamba_state_size", + type=int, + default=16, + help="SSM state size for EHRMamba and JambaEHR blocks.", + ) + parser.add_argument( + "--mamba-conv-kernel", + "--mamba_conv_kernel", + "--conv-kernel", + "--conv_kernel", + dest="mamba_conv_kernel", + type=int, + default=4, + help="Causal conv kernel size for EHRMamba and JambaEHR blocks.", + ) + parser.add_argument( + "--jamba-transformer-layers", + "--jamba_transformer_layers", + dest="jamba_transformer_layers", + type=int, + default=2, + help="Number of Transformer (attention) layers in JambaEHR.", + ) + parser.add_argument( + "--jamba-mamba-layers", + "--jamba_mamba_layers", + dest="jamba_mamba_layers", + type=int, + default=6, + help="Number of Mamba (SSM) layers in JambaEHR. Library default is 6; " + "pass 2 for a depth-matched comparison against --num-layers 2.", + ) + parser.add_argument( + "--no-lab-standardization", + action="store_true", + default=False, + help="Disable train-split lab z-scoring (raw-lab ablation).", + ) - return parser.parse_args() + args = parser.parse_args() + if args.patients is not None: + parser.error( + "--patients is not a supported flag and is not a 5-patient smoke. " + "Omit it for the full table; use --dev N for a patient-limited run." + ) + return args if __name__ == "__main__": diff --git a/tests/test_p2_runner_measurement.py b/tests/test_p2_runner_measurement.py new file mode 100644 index 000000000..fb25ec4c7 --- /dev/null +++ b/tests/test_p2_runner_measurement.py @@ -0,0 +1,124 @@ +"""Proof that paired runs cannot overwrite each other or report train as test. + +The run directory was ``{model}_seed{seed}``. ``--task labs`` and +``--task notes_labs`` at seed 42 resolved to one path and the second run +destroyed the first. ``split_by_patient`` fell back to ``split_by_sample`` +with no warning (patient leak). Predictions came from +``test_loader or val_loader or train_loader``. + +The sixth backbone was also missing: ``--model`` had no ``mlp``. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p2_runner_measurement.py -q +""" + +from __future__ import annotations + +import importlib.util +import sys +import unittest +import warnings +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + + +RUNNER = ( + Path(__file__).resolve().parents[1] + / "examples" + / "mortality_prediction" + / "unified_embedding_e2e_mimic4.py" +) + + +def _load_runner(): + spec = importlib.util.spec_from_file_location("e2e_runner_measurement", RUNNER) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def _parse(mod, *argv): + with mock.patch.object(sys, "argv", ["e2e.py", "--ehr-root", "/tmp", *argv]): + return mod.parse_args() + + +def _exp_name_assignment() -> str: + for line in RUNNER.read_text().splitlines(): + stripped = line.strip() + if stripped.startswith("exp_name = "): + return stripped + raise AssertionError("exp_name is no longer assigned in the runner") + + +class TestP2RunnerMeasurement(unittest.TestCase): + def test_run_directory_includes_the_task(self): + assignment = _exp_name_assignment() + self.assertIn("args.task", assignment) + template = assignment.split("=", 1)[1].strip() + args = SimpleNamespace(task="labs", model="transformer", seed=42) + labs = eval(template, {"args": args}) # noqa: S307 + args.task = "notes_labs" + notes = eval(template, {"args": args}) # noqa: S307 + self.assertEqual(labs, "labs_transformer_seed42") + self.assertEqual(notes, "notes_labs_transformer_seed42") + self.assertNotEqual(labs, notes) + + def test_cli_accepts_mlp(self): + from pyhealth.models import MLP + + self.assertTrue(hasattr(MLP, "_forward_unified")) + mod = _load_runner() + args = _parse(mod, "--model", "mlp") + self.assertEqual(args.model, "mlp") + + def test_split_warns_and_labels_leaky_fallback(self): + from pyhealth.datasets import create_sample_dataset + + mod = _load_runner() + samples = [ + { + "patient_id": "only", + "visit_id": "v0", + "labs": [1.0, 2.0], + "label": 0, + }, + { + "patient_id": "only", + "visit_id": "v1", + "labs": [3.0, 4.0], + "label": 1, + }, + ] + dataset = create_sample_dataset( + samples=samples, + input_schema={"labs": "tensor"}, + output_schema={"label": "binary"}, + in_memory=True, + ) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _train, _val, _test, mode = mod._split_dataset(dataset, seed=1) + self.assertEqual(mode, "by_sample_fallback_leaky") + self.assertTrue(any("split_by_sample" in str(w.message) for w in caught)) + + def test_eval_split_is_named_and_recorded(self): + src = RUNNER.read_text() + self.assertIn('inference_loader, eval_split = test_loader, "test"', src) + self.assertIn("write_run_config", src) + self.assertIn("These are not test metrics", src) + self.assertNotIn( + "inference_loader = test_loader or val_loader or train_loader", + src, + ) + + def test_jamba_cli_default_matches_the_library(self): + from pyhealth.models.jamba_ehr import JambaLayer + + self.assertEqual(JambaLayer.__init__.__defaults__[1], 6) + mod = _load_runner() + args = _parse(mod, "--model", "jambaehr") + self.assertEqual(args.jamba_transformer_layers, 2) + self.assertEqual(args.jamba_mamba_layers, 6) From 5be22e4de8d95de0a578e7f548117b68cb921dc9 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Tue, 18 Aug 2026 16:26:56 -0400 Subject: [PATCH 23/27] Give NotesLabs a 24h window and put concatenated stays on one timeline. The class default was None (labs through discharge) while the docstring, LabsMIMIC4, and the runner all used 24. Event times were hours from each stay's own admit, so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay; times are hours from the first stay in the sample. Bump emitted_data_version to 3. Co-authored-by: Cursor --- pyhealth/tasks/multimodal_mimic4.py | 107 ++++++++++++++-------------- tests/test_p1_observation_window.py | 15 +++- tests/test_p1_time_axis.py | 71 ++++++++++++++++++ 3 files changed, 138 insertions(+), 55 deletions(-) create mode 100644 tests/test_p1_time_axis.py diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 9e2905f00..33a303c7e 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -88,7 +88,10 @@ def __init__( # 1: empty events instead of placeholders; per-admission lab window. # 2: CXR/notes_labs_cxr no longer drop later stays against the first # admission's clock (admission_time >= first_admit + window_hours). - self.emitted_data_version = 2 + # 3: event times are hours from the first stay in the sample, not + # reset per admission (reset times made stay 2 at +6h sort with + # stay 1 at +6h). + self.emitted_data_version = 3 @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: @@ -127,6 +130,16 @@ def _parse_datetime(value: Any) -> Optional[datetime]: def _to_hours(delta_seconds: float) -> float: return delta_seconds / 3600.0 + @classmethod + def _hours_since(cls, timestamp: datetime, origin: datetime) -> float: + """Hours from ``origin`` to ``timestamp``. + + Collection windows stay per admission. The value written onto the + unified timeline is hours from the first stay in this sample, so a + later stay at +6h does not sort with the first stay at +6h. + """ + return cls._to_hours((timestamp - origin).total_seconds()) + def _compute_effective_window( self, admissions_to_process: List[Any], @@ -220,13 +233,16 @@ def _collect_labs( patient: Any, admission_time: datetime, end_time: datetime, + time_origin: datetime, ) -> Tuple[List[float], List[List[float]], List[List[bool]]]: """Collect lab values and observation masks for one admission. Args: patient: Patient object. - admission_time: Start of the window; times are relative to this. + admission_time: Start of this stay's collection window. end_time: End of the window (inclusive). + time_origin: First stay in this sample. Event times are hours + from here, not from ``admission_time``. Returns: Tuple of (lab_times, lab_values, lab_masks). ``lab_masks`` is a @@ -285,9 +301,7 @@ def _collect_labs( break lab_vector.append(category_value) lab_mask.append(observed) - lab_times.append( - self._to_hours((lab_ts - admission_time).total_seconds()) - ) + lab_times.append(self._hours_since(lab_ts, time_origin)) lab_values.append(lab_vector) lab_masks.append(lab_mask) return lab_times, lab_values, lab_masks @@ -302,6 +316,7 @@ def _collect_notes( end_time: Optional[datetime] = None, section_headers: Optional[List[str]] = None, fallback_to_full_note: bool = False, + time_origin: Optional[datetime] = None, ) -> Tuple[List[str], List[float]]: """Collect notes of a given type for one admission. @@ -309,7 +324,9 @@ def _collect_notes( patient: Patient object. note_event_type: Event type string (e.g. "discharge", "radiology"). hadm_id: Admission ID to filter by. - admission_time: Admission start time; used to compute time offsets. + admission_time: This stay's admit time (unused for the timeline + once ``time_origin`` is set; kept so existing call sites that + pass it positionally stay valid). start_time: Optional start of the time window. end_time: Optional end of the time window. section_headers: When provided, extract only these named sections @@ -319,8 +336,8 @@ def _collect_notes( with no matching sections are dropped entirely. Returns: - Tuple of (texts, hours_from_admission). Empty lists when the - events list is empty; do not invent a placeholder note. + Tuple of (texts, hours from the sample's first stay). Empty lists + when the events list is empty; do not invent a placeholder note. """ notes = patient.get_events( event_type=note_event_type, @@ -343,11 +360,9 @@ def _collect_notes( elif not fallback_to_full_note: continue - time_from_admission = self._to_hours( - (note.timestamp - admission_time).total_seconds() - ) + origin = time_origin if time_origin is not None else admission_time texts.append(note_text) - note_times.append(time_from_admission) + note_times.append(self._hours_since(note.timestamp, origin)) except ( AttributeError ): # note object is missing .text or .timestamp attribute (e.g. malformed note) @@ -361,8 +376,8 @@ class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): A notes-free structured-EHR task that uses only: - - **ICD codes**: diagnosis and procedure codes per admission, processed by - ``StageNetProcessor`` with inter-admission time offsets. + - **ICD codes**: diagnosis and procedure codes per admission, placed on + the same hours-from-first-stay timeline as labs. - **Lab values**: 10-dimensional lab vectors (one per lab category) at each measurement timestamp, processed by ``StageNetTensorProcessor``. @@ -402,13 +417,13 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: effective_start, effective_end = self._compute_effective_window( admissions_to_process ) + time_origin = admissions_to_process[0].timestamp all_icd_codes: List[List[str]] = [] all_icd_times: List[float] = [] all_lab_values: List[List[float]] = [] all_lab_masks: List[List[bool]] = [] all_lab_times: List[float] = [] - previous_admission_time = None for admission in admissions_to_process: admission_time = admission.timestamp @@ -424,16 +439,8 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) if visit_icd_codes: - if previous_admission_time is None: - time_from_previous = 0.0 - else: - time_from_previous = self._to_hours( - (admission_time - previous_admission_time).total_seconds() - ) all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) - - previous_admission_time = admission_time + all_icd_times.append(self._hours_since(admission_time, time_origin)) lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, @@ -441,6 +448,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: end_time=self._admission_window_end( admission_time, admission_dischtime ), + time_origin=time_origin, ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) @@ -514,7 +522,7 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): def __init__( self, - window_hours: Optional[float] = None, + window_hours: Optional[float] = 24, include_icd: bool = False, ) -> None: super().__init__(window_hours=window_hours) @@ -543,6 +551,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: effective_start, effective_end = self._compute_effective_window( admissions_to_process ) + time_origin = admissions_to_process[0].timestamp all_note_texts: List[str] = [] all_note_times: List[float] = [] @@ -551,7 +560,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_times: List[float] = [] all_icd_codes: List[List[str]] = [] all_icd_times: List[float] = [] - previous_admission_time = None for admission in admissions_to_process: admission_time = admission.timestamp @@ -571,6 +579,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.hadm_id, admission_time, section_headers=self.DISCHARGE_CLINICAL_HEADERS, + time_origin=time_origin, ) all_note_texts.extend(note_texts) all_note_times.extend(note_times) @@ -581,6 +590,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: patient=patient, admission_time=admission_time, end_time=lab_end, + time_origin=time_origin, ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) @@ -599,23 +609,18 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: start_time=admission_time, end_time=lab_end, section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + time_origin=time_origin, ) all_note_texts.extend(radiology_texts) all_note_times.extend(radiology_times) if self.include_icd: visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) - time_from_previous = ( - 0.0 - if previous_admission_time is None - else self._to_hours( - (admission_time - previous_admission_time).total_seconds() - ) - ) if visit_icd_codes: all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) - previous_admission_time = admission_time + all_icd_times.append( + self._hours_since(admission_time, time_origin) + ) record: Dict[str, Any] = { "patient_id": patient.patient_id, @@ -656,7 +661,7 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): Args: window_hours: Hours from admission for lab/CXR collection. - ``None`` (default) collects for the full admission span. + ``None`` collects for the full admission span. Default: 24. include_icd: When ``True``, collect discharge-coded ICD codes and add ``icd_codes`` to the sample dict / input schema. Default: ``False``. """ @@ -691,7 +696,7 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): def __init__( self, - window_hours: Optional[float] = None, + window_hours: Optional[float] = 24, include_icd: bool = False, ) -> None: super().__init__(window_hours=window_hours) @@ -720,6 +725,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: effective_start, effective_end = self._compute_effective_window( admissions_to_process ) + time_origin = admissions_to_process[0].timestamp all_note_texts: List[str] = [] all_note_times: List[float] = [] @@ -730,7 +736,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_icd_times: List[float] = [] all_cxr_paths: List[str] = [] all_cxr_times: List[float] = [] - previous_admission_time = None for admission in admissions_to_process: admission_time = admission.timestamp @@ -750,6 +755,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.hadm_id, admission_time, section_headers=self.DISCHARGE_CLINICAL_HEADERS, + time_origin=time_origin, ) all_note_texts.extend(note_texts) all_note_times.extend(note_times) @@ -760,6 +766,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: patient=patient, admission_time=admission_time, end_time=lab_end, + time_origin=time_origin, ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) @@ -778,6 +785,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: start_time=admission_time, end_time=lab_end, section_headers=self.RADIOLOGY_CLINICAL_HEADERS, + time_origin=time_origin, ) all_note_texts.extend(radiology_texts) all_note_times.extend(radiology_times) @@ -794,26 +802,18 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: if event.image_path: all_cxr_paths.append(event.image_path) all_cxr_times.append( - self._to_hours( - (event.timestamp - admission_time).total_seconds() - ) + self._hours_since(event.timestamp, time_origin) ) except AttributeError: continue if self.include_icd: visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) - time_from_previous = ( - 0.0 - if previous_admission_time is None - else self._to_hours( - (admission_time - previous_admission_time).total_seconds() - ) - ) if visit_icd_codes: all_icd_codes.append(visit_icd_codes) - all_icd_times.append(time_from_previous) - previous_admission_time = admission_time + all_icd_times.append( + self._hours_since(admission_time, time_origin) + ) record: Dict[str, Any] = { "patient_id": patient.patient_id, @@ -870,6 +870,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri effective_start, effective_end = self._compute_effective_window( admissions_to_process ) + time_origin = admissions_to_process[0].timestamp all_lab_times: List[float] = [] all_lab_values: List[List[float]] = [] @@ -893,6 +894,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri end_time=self._admission_window_end( admission_time, admission_dischtime ), + time_origin=time_origin, ) all_lab_times.extend(lab_times) all_lab_values.extend(lab_values) @@ -949,6 +951,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri effective_start, effective_end = self._compute_effective_window( admissions_to_process ) + time_origin = admissions_to_process[0].timestamp all_cxr_paths: List[str] = [] all_cxr_times: List[float] = [] @@ -980,9 +983,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri if event.image_path: all_cxr_paths.append(event.image_path) all_cxr_times.append( - self._to_hours( - (event.timestamp - admission_time).total_seconds() - ) + self._hours_since(event.timestamp, time_origin) ) except AttributeError: continue diff --git a/tests/test_p1_observation_window.py b/tests/test_p1_observation_window.py index b1043c025..4c0dea61e 100644 --- a/tests/test_p1_observation_window.py +++ b/tests/test_p1_observation_window.py @@ -14,7 +14,7 @@ CXR / ``notes_labs_cxr`` still skipped those later stays with ``admission_time >= first_admit + window_hours``. That skip is gone. -``emitted_data_version`` is 2 so caches from version 1 cannot be reused. +``emitted_data_version`` is 3 so caches from version 1-2 cannot be reused. Repro:: @@ -93,7 +93,7 @@ def test_window_change_invalidates_the_cache(self): task = m.LabsMIMIC4(window_hours=24) self.assertIsNotNone(vars(task).get("emitted_data_version")) - self.assertGreaterEqual(task.emitted_data_version, 2) + self.assertGreaterEqual(task.emitted_data_version, 3) def cache_key(t, drop_version=False): v = dict(vars(t)) @@ -119,3 +119,14 @@ def test_window_none_still_collects_through_discharge(self): admit = datetime(2180, 5, 6, 8, 0, 0) discharge = admit + timedelta(days=9) self.assertEqual(task._admission_window_end(admit, discharge), discharge) + + def test_notes_labs_defaults_to_a_24h_window(self): + from pyhealth.tasks.multimodal_mimic4 import ( + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, + ) + + self.assertEqual(NotesLabsMIMIC4().window_hours, 24) + self.assertEqual(NotesLabsCXRMIMIC4().window_hours, 24) + self.assertEqual(LabsMIMIC4().window_hours, 24) diff --git a/tests/test_p1_time_axis.py b/tests/test_p1_time_axis.py new file mode 100644 index 000000000..61023885c --- /dev/null +++ b/tests/test_p1_time_axis.py @@ -0,0 +1,71 @@ +"""Proof that concatenated stays share one time origin. + +The sample is patient-level: every admission up to the first death is one +sequence. Event times were hours from *that stay's* admit, then concatenated, +so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay +(admit, admit+window]. Times are hours from the first stay in the sample. + +Single-stay patients are unchanged. The sinusoid still wraps at +``max_time_hours=720``; that affects the embedding of multi-year gaps, not +the sort order. + +Repro:: + + PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=. \\ + python -m pytest tests/test_p1_time_axis.py -q +""" + +from __future__ import annotations + +import inspect +import unittest +from datetime import datetime, timedelta + + +class TestP1TimeAxis(unittest.TestCase): + def test_hours_since_does_not_reset_per_stay(self): + from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task + + first = datetime(2180, 5, 6, 8, 0, 0) + later = first + timedelta(days=400) + stay1_event = first + timedelta(hours=6) + stay2_event = later + timedelta(hours=6) + + self.assertEqual(BaseMultimodalMIMIC4Task._hours_since(stay1_event, first), 6.0) + # Old convention: both events were 6.0 and collided after concat. + self.assertEqual(BaseMultimodalMIMIC4Task._hours_since(stay2_event, later), 6.0) + self.assertAlmostEqual( + BaseMultimodalMIMIC4Task._hours_since(stay2_event, first), + 400 * 24 + 6.0, + ) + self.assertGreater( + BaseMultimodalMIMIC4Task._hours_since(stay2_event, first), + BaseMultimodalMIMIC4Task._hours_since(stay1_event, first), + ) + + def test_collectors_write_hours_from_the_sample_origin(self): + from pyhealth.tasks.multimodal_mimic4 import ( + BaseMultimodalMIMIC4Task, + CXRMIMIC4, + LabsMIMIC4, + NotesLabsCXRMIMIC4, + NotesLabsMIMIC4, + ) + + for cls in (LabsMIMIC4, NotesLabsMIMIC4, NotesLabsCXRMIMIC4, CXRMIMIC4): + src = inspect.getsource(cls.__call__) + self.assertIn( + "time_origin = admissions_to_process[0].timestamp", + src, + msg=cls.__name__, + ) + self.assertNotIn( + "event.timestamp - admission_time", + src, + msg=f"{cls.__name__} still resets times per stay", + ) + + labs_src = inspect.getsource(BaseMultimodalMIMIC4Task._collect_labs) + self.assertIn("time_origin", labs_src) + self.assertIn("_hours_since", labs_src) + self.assertNotIn("lab_ts - admission_time", labs_src) From 9f50b763dff22c3d8a70c4c48af086dec483277c Mon Sep 17 00:00:00 2001 From: Rian354 Date: Wed, 19 Aug 2026 14:32:22 -0400 Subject: [PATCH 24/27] Give time embeddings a 10-year span and drop ICDLabsMIMIC4. The old sinusoid wrapped every 720 hours, so later stays aliased with the first. MIMIC timestamps ICD at discharge, which leaks the in-hospital mortality label. Co-authored-by: Cursor --- pyhealth/models/embedding/unified.py | 49 ++++++++----- pyhealth/models/unified_embedding.py | 46 +++++++----- pyhealth/tasks/multimodal_mimic4.py | 102 ++------------------------- tests/test_p1_observation_window.py | 8 ++- tests/test_p1_time_axis.py | 18 ++++- tests/test_unified_multimodal.py | 10 ++- 6 files changed, 100 insertions(+), 133 deletions(-) diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 3f786bc95..bf8050088 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -62,37 +62,54 @@ class SinusoidalTimeEmbedding(nn.Module): - """Continuous sinusoidal embedding for scalar time values (in hours). + """Multi-scale sinusoidal embedding for times in hours. - Identical in spirit to the positional encoding in "Attention is All You - Need" but operating on real-valued timestamps rather than integer positions. + Wavelengths are spaced geometrically from ``min_hours`` (within-stay + resolution) to ``max_hours`` (longitudinal span). The previous encoding + mapped ``t / 720 * 2π``, so every frequency wrapped every 30 days: a later + stay at +9606h produced the same embedding as +6h even after the task + stopped resetting the clock. Args: dim: Output embedding dimension (must be even). - max_hours: Maximum expected time value in hours. Values are normalised - to ``[0, 2π]`` before the sin/cos projection. Default 720 (30 days). + max_hours: Longest wavelength in hours. Default 87600 (10 years). + min_hours: Shortest wavelength in hours. Default 1.0. Shape: Input: ``(*, )`` float tensor of times in hours Output: ``(*, dim)`` """ - def __init__(self, dim: int, max_hours: float = 720.0): + def __init__( + self, + dim: int, + max_hours: float = 87600.0, + min_hours: float = 1.0, + ): super().__init__() assert dim % 2 == 0, f"dim must be even, got {dim}" + if min_hours <= 0 or max_hours <= min_hours: + raise ValueError( + f"need 0 < min_hours < max_hours, got min={min_hours}, max={max_hours}" + ) self.dim = dim - self.max_hours = max_hours + self.max_hours = float(max_hours) + self.min_hours = float(min_hours) half = dim // 2 - freqs = torch.exp( - -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1) + periods = torch.exp( + torch.linspace( + math.log(self.min_hours), + math.log(self.max_hours), + half, + dtype=torch.float32, + ) ) - self.register_buffer("freqs", freqs) # (dim//2,) + self.register_buffer("freqs", 2 * math.pi / periods) def forward(self, t: torch.Tensor) -> torch.Tensor: """:param t: ``(...,)`` float, times in hours.""" - t_norm = t / self.max_hours * 2 * math.pi # (...,) - args = t_norm.unsqueeze(-1) * self.freqs # (..., dim//2) - return torch.cat([args.sin(), args.cos()], dim=-1) # (..., dim) + args = t.unsqueeze(-1).to(dtype=self.freqs.dtype) * self.freqs + return torch.cat([args.sin(), args.cos()], dim=-1) class _MeanPool(nn.Module): @@ -159,8 +176,8 @@ class UnifiedMultimodalEmbeddingModel(nn.Module, BaseEmbeddingModel): ``dataset.input_processors`` directly. embedding_dim: Shared embedding dimension ``E'``. time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. - max_time_hours: Normalisation constant for the time embedding. - Defaults to 720 h (30 days). + max_time_hours: Longest wavelength of the time embedding, in hours. + Defaults to 87600 (10 years). Shortest wavelength is 1 hour. image_size: Image size (H=W) assumed for IMAGE fields when using PatchEmbedding. Defaults to 224. image_channels: Number of input channels for IMAGE fields. Defaults to 3. @@ -206,7 +223,7 @@ def __init__( processors: dict[str, Any], embedding_dim: int = 128, time_embedding: str = "sinusoidal", - max_time_hours: float = 720.0, + max_time_hours: float = 87600.0, image_size: int = 224, image_channels: int = 3, patch_size: int = 16, diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py index 014326b41..3784bc0c0 100644 --- a/pyhealth/models/unified_embedding.py +++ b/pyhealth/models/unified_embedding.py @@ -35,37 +35,51 @@ class SinusoidalTimeEmbedding(nn.Module): - """Continuous sinusoidal embedding for scalar time values (in hours). + """Multi-scale sinusoidal embedding for times in hours. - Identical in spirit to the positional encoding in "Attention is All You - Need" but operating on real-valued timestamps rather than integer positions. + Wavelengths are spaced geometrically from ``min_hours`` to ``max_hours``. + The previous encoding wrapped every 720 hours. Args: dim: Output embedding dimension (must be even). - max_hours: Maximum expected time value in hours. Values are normalised - to ``[0, 2π]`` before the sin/cos projection. Default 720 (30 days). + max_hours: Longest wavelength in hours. Default 87600 (10 years). + min_hours: Shortest wavelength in hours. Default 1.0. Shape: Input: ``(*, )`` float tensor of times in hours Output: ``(*, dim)`` """ - def __init__(self, dim: int, max_hours: float = 720.0): + def __init__( + self, + dim: int, + max_hours: float = 87600.0, + min_hours: float = 1.0, + ): super().__init__() assert dim % 2 == 0, f"dim must be even, got {dim}" + if min_hours <= 0 or max_hours <= min_hours: + raise ValueError( + f"need 0 < min_hours < max_hours, got min={min_hours}, max={max_hours}" + ) self.dim = dim - self.max_hours = max_hours + self.max_hours = float(max_hours) + self.min_hours = float(min_hours) half = dim // 2 - freqs = torch.exp( - -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1) + periods = torch.exp( + torch.linspace( + math.log(self.min_hours), + math.log(self.max_hours), + half, + dtype=torch.float32, + ) ) - self.register_buffer("freqs", freqs) # (dim//2,) + self.register_buffer("freqs", 2 * math.pi / periods) def forward(self, t: torch.Tensor) -> torch.Tensor: """:param t: ``(...,)`` float, times in hours.""" - t_norm = t / self.max_hours * 2 * math.pi # (...,) - args = t_norm.unsqueeze(-1) * self.freqs # (..., dim//2) - return torch.cat([args.sin(), args.cos()], dim=-1) # (..., dim) + args = t.unsqueeze(-1).to(dtype=self.freqs.dtype) * self.freqs + return torch.cat([args.sin(), args.cos()], dim=-1) def _build_image_encoder(embedding_dim: int) -> nn.Module: @@ -128,8 +142,8 @@ class UnifiedMultimodalEmbeddingModel(nn.Module): ``dataset.input_processors`` directly. embedding_dim: Shared embedding dimension ``E'``. time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. - max_time_hours: Normalisation constant for the time embedding. - Defaults to 720 h (30 days). + max_time_hours: Longest wavelength of the time embedding, in hours. + Defaults to 87600 (10 years). Example:: @@ -148,7 +162,7 @@ def __init__( processors: dict[str, Any], embedding_dim: int = 128, time_embedding: str = "sinusoidal", - max_time_hours: float = 720.0, + max_time_hours: float = 87600.0, ): super().__init__() self.embedding_dim = embedding_dim diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 33a303c7e..f0b8e6011 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -371,102 +371,6 @@ def _collect_notes( return texts, note_times -class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): - """Task for ICD codes + lab values mortality prediction using MIMIC-IV. - - A notes-free structured-EHR task that uses only: - - - **ICD codes**: diagnosis and procedure codes per admission, placed on - the same hours-from-first-stay timeline as labs. - - **Lab values**: 10-dimensional lab vectors (one per lab category) at each - measurement timestamp, processed by ``StageNetTensorProcessor``. - - Examples: - >>> from pyhealth.datasets import MIMIC4Dataset - >>> from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4 - >>> dataset = MIMIC4Dataset( - ... ehr_root="/path/to/mimic-iv/2.2", - ... ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], - ... ) - >>> task = ICDLabsMIMIC4() - >>> samples = dataset.set_task(task) - """ - - PADDING: int = 0 - - task_name: str = "ICDLabsMIMIC4" - input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "icd_codes": ("stagenet", {"padding": PADDING}), - "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {"forward_fill": False}), - } - output_schema: Dict[str, str] = {"mortality": "binary"} - - def __call__(self, patient: Any) -> List[Dict[str, Any]]: - demographics = patient.get_events(event_type="patients") - if not demographics: - return [] - - admissions_to_process, mortality_label = self._build_admissions_to_process( - patient - ) - - if len(admissions_to_process) == 0: - return [] - - effective_start, effective_end = self._compute_effective_window( - admissions_to_process - ) - time_origin = admissions_to_process[0].timestamp - - all_icd_codes: List[List[str]] = [] - all_icd_times: List[float] = [] - all_lab_values: List[List[float]] = [] - all_lab_masks: List[List[bool]] = [] - all_lab_times: List[float] = [] - - for admission in admissions_to_process: - admission_time = admission.timestamp - - try: - admission_dischtime = datetime.strptime( - admission.dischtime, "%Y-%m-%d %H:%M:%S" - ) - except (ValueError, AttributeError): - admission_dischtime = admission_time - if admission_dischtime < admission_time: - admission_dischtime = admission_time - - visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) - if visit_icd_codes: - all_icd_codes.append(visit_icd_codes) - all_icd_times.append(self._hours_since(admission_time, time_origin)) - - lab_times, lab_values, lab_masks = self._collect_labs( - patient=patient, - admission_time=admission_time, - end_time=self._admission_window_end( - admission_time, admission_dischtime - ), - time_origin=time_origin, - ) - all_lab_times.extend(lab_times) - all_lab_values.extend(lab_values) - all_lab_masks.extend(lab_masks) - - single_patient_longitudinal_record = { - "patient_id": patient.patient_id, - "icd_codes": (all_icd_times, all_icd_codes), - "labs": (all_lab_times, all_lab_values), - "labs_mask": (all_lab_times, all_lab_masks), - "mortality": mortality_label, - "window_start": effective_start, - "window_end": effective_end, - } - - return [single_patient_longitudinal_record] - - class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): """Mortality prediction from admission-context notes and lab values. @@ -498,6 +402,9 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): collects for the full admission span. Default: 24. include_icd: When ``True``, collect discharge-coded ICD codes and add ``icd_codes`` to the sample dict / input schema. Default: ``False``. + MIMIC-IV timestamps those codes at ``dischtime``, so this leaks + the in-hospital mortality label. Keep it off except as an + explicit leaky ablation. """ PADDING: int = 0 @@ -664,6 +571,9 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): ``None`` collects for the full admission span. Default: 24. include_icd: When ``True``, collect discharge-coded ICD codes and add ``icd_codes`` to the sample dict / input schema. Default: ``False``. + MIMIC-IV timestamps those codes at ``dischtime``, so this leaks + the in-hospital mortality label. Keep it off except as an + explicit leaky ablation. """ PADDING: int = 0 diff --git a/tests/test_p1_observation_window.py b/tests/test_p1_observation_window.py index 4c0dea61e..aa59a682c 100644 --- a/tests/test_p1_observation_window.py +++ b/tests/test_p1_observation_window.py @@ -33,7 +33,6 @@ LAB_TASKS = [ "LabsMIMIC4", - "ICDLabsMIMIC4", "NotesLabsMIMIC4", "NotesLabsCXRMIMIC4", "CXRMIMIC4", @@ -130,3 +129,10 @@ def test_notes_labs_defaults_to_a_24h_window(self): self.assertEqual(NotesLabsMIMIC4().window_hours, 24) self.assertEqual(NotesLabsCXRMIMIC4().window_hours, 24) self.assertEqual(LabsMIMIC4().window_hours, 24) + + def test_discharge_coded_icd_is_not_a_mortality_task(self): + from pyhealth.tasks import multimodal_mimic4 as m + from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4 + + self.assertFalse(hasattr(m, "ICDLabsMIMIC4")) + self.assertFalse(NotesLabsMIMIC4().include_icd) diff --git a/tests/test_p1_time_axis.py b/tests/test_p1_time_axis.py index 61023885c..521687ad3 100644 --- a/tests/test_p1_time_axis.py +++ b/tests/test_p1_time_axis.py @@ -5,9 +5,9 @@ so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay (admit, admit+window]. Times are hours from the first stay in the sample. -Single-stay patients are unchanged. The sinusoid still wraps at -``max_time_hours=720``; that affects the embedding of multi-year gaps, not -the sort order. +Single-stay patients are unchanged. Time embeddings use geometric +wavelengths from 1 hour to 10 years so a later stay at +9606h does not +alias with +6h. Repro:: @@ -21,6 +21,8 @@ import unittest from datetime import datetime, timedelta +import torch + class TestP1TimeAxis(unittest.TestCase): def test_hours_since_does_not_reset_per_stay(self): @@ -69,3 +71,13 @@ def test_collectors_write_hours_from_the_sample_origin(self): self.assertIn("time_origin", labs_src) self.assertIn("_hours_since", labs_src) self.assertNotIn("lab_ts - admission_time", labs_src) + + def test_sinusoid_does_not_wrap_every_thirty_days(self): + from pyhealth.models.embedding.unified import SinusoidalTimeEmbedding + + emb = SinusoidalTimeEmbedding(dim=32) + t6 = emb(torch.tensor([6.0])) + t726 = emb(torch.tensor([6.0 + 720.0])) + t9606 = emb(torch.tensor([9606.0])) + self.assertFalse(torch.allclose(t6, t726, atol=1e-5)) + self.assertFalse(torch.allclose(t6, t9606, atol=1e-5)) diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 01bc194e6..99b78cbae 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -170,7 +170,7 @@ def test_collate_temporal_variable_length(): def test_sinusoidal_time_embedding_shape(): from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding - emb = SinusoidalTimeEmbedding(dim=64, max_hours=720.0) + emb = SinusoidalTimeEmbedding(dim=64) t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]]) # (2, 3) out = emb(t) assert out.shape == (2, 3, 64) @@ -184,6 +184,14 @@ def test_sinusoidal_different_times_differ(): assert not torch.allclose(t0, t1) +def test_sinusoidal_does_not_alias_every_720h(): + from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=32) + t6 = emb(torch.tensor([6.0])) + t726 = emb(torch.tensor([726.0])) + assert not torch.allclose(t6, t726, atol=1e-5) + + # ── 6. UnifiedMultimodalEmbeddingModel — code-only smoke test ───────────────── def _make_code_processors_and_inputs(batch_size=2, seq_len=5): From 89f7bc742ec5e415dff59b64e2455ef2308803f1 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Wed, 19 Aug 2026 16:05:55 -0400 Subject: [PATCH 25/27] Point the old unified_embedding import at the package copy. Tests and docs imported pyhealth.models.unified_embedding while runtime used embedding.unified, so the two files could drift. embedding.py next to the embedding/ package was unused. Co-authored-by: Cursor --- pyhealth/models/embedding.py | 361 --------------------------- pyhealth/models/unified_embedding.py | 349 +------------------------- tests/test_unified_multimodal.py | 28 ++- 3 files changed, 33 insertions(+), 705 deletions(-) delete mode 100644 pyhealth/models/embedding.py diff --git a/pyhealth/models/embedding.py b/pyhealth/models/embedding.py deleted file mode 100644 index 4232b2788..000000000 --- a/pyhealth/models/embedding.py +++ /dev/null @@ -1,361 +0,0 @@ -from __future__ import annotations - -from typing import Dict, Any, Optional, Union -import os - -import torch -import torch.nn as nn - -from ..datasets import SampleDataset -from ..processors import ( - MultiHotProcessor, - NestedFloatsProcessor, - NestedSequenceProcessor, - SequenceProcessor, - StageNetProcessor, - StageNetTensorProcessor, - TensorProcessor, - TimeseriesProcessor, - DeepNestedSequenceProcessor, - DeepNestedFloatsProcessor, -) -from .base_model import BaseModel - - -def _iter_text_vectors( - path: str, - embedding_dim: int, - wanted_tokens: set[str], - encoding: str = "utf-8", -) -> Dict[str, torch.Tensor]: - """Loads word vectors from a text file (e.g., GloVe) for a subset of tokens. - - Expected format: one token per line followed by embedding_dim floats. - - This function reads the file line-by-line and only retains vectors for - tokens present in `wanted_tokens`. - """ - - if not os.path.exists(path): - raise FileNotFoundError(f"pretrained embedding file not found: {path}") - - vectors: Dict[str, torch.Tensor] = {} - with open(path, "r", encoding=encoding) as f: - for line in f: - line = line.strip() - if not line: - continue - parts = line.split() - # token + embedding_dim values - if len(parts) < embedding_dim + 1: - continue - token = parts[0] - if token not in wanted_tokens: - continue - try: - vec = torch.tensor( - [float(x) for x in parts[1 : embedding_dim + 1]], - dtype=torch.float, - ) - except ValueError: - continue - vectors[token] = vec - return vectors - - -def init_embedding_with_pretrained( - embedding: nn.Embedding, - code_vocab: Dict[Any, int], - pretrained_path: str, - embedding_dim: int, - pad_token: str = "", - unk_token: str = "", - normalize: bool = False, - freeze: bool = False, -) -> int: - """Initializes an nn.Embedding from a pretrained text-vector file. - - Tokens not found in the pretrained file are left as the module's existing - random initialization. - - Returns: - int: number of tokens successfully initialized from the file. - """ - - # Build wanted token set (stringified) - vocab_tokens = {str(t) for t in code_vocab.keys()} - vectors = _iter_text_vectors(pretrained_path, embedding_dim, vocab_tokens) - - loaded = 0 - with torch.no_grad(): - for tok, idx in code_vocab.items(): - tok_s = str(tok) - if tok_s in vectors: - vec = vectors[tok_s] - if normalize: - vec = vec / (vec.norm(p=2) + 1e-12) - embedding.weight[idx].copy_(vec) - loaded += 1 - - # Ensure pad row is zero - if pad_token in code_vocab: - embedding.weight[code_vocab[pad_token]].zero_() - # If embedding has a padding_idx, keep it consistent - if embedding.padding_idx is not None: - embedding.weight[embedding.padding_idx].zero_() - - if freeze: - embedding.weight.requires_grad_(False) - - return loaded - - -class EmbeddingModel(BaseModel): - """ - EmbeddingModel is responsible for creating embedding layers for different types of input data. - - This model automatically creates appropriate embedding transformations based on the processor type: - - - SequenceProcessor: nn.Embedding - Input: (batch, seq_len) - Output: (batch, seq_len, embedding_dim) - - - NestedSequenceProcessor: nn.Embedding - Input: (batch, num_visits, max_codes_per_visit) - Output: (batch, num_visits, max_codes_per_visit, embedding_dim) - - - DeepNestedSequenceProcessor: nn.Embedding - Input: (batch, num_groups, num_visits, max_codes_per_visit) - Output: (batch, num_groups, num_visits, max_codes_per_visit, embedding_dim) - - - TimeseriesProcessor / NestedFloatsProcessor / DeepNestedFloatsProcessor / StageNetTensorProcessor: - nn.Linear over the last dimension - Input: (..., size) - Output: (..., embedding_dim) - - - TensorProcessor: nn.Linear (size inferred from first sample) - - - MultiHotProcessor: nn.Linear over multi-hot vector - """ - - def __init__( - self, - dataset: SampleDataset, - embedding_dim: int = 128, - pretrained_emb_path: Optional[Union[str, Dict[str, str]]] = None, - freeze_pretrained: bool = False, - normalize_pretrained: bool = False, - ): - super().__init__(dataset) - self.embedding_dim = embedding_dim - self.embedding_layers = nn.ModuleDict() - - for field_name, processor in self.dataset.input_processors.items(): - # Deep categorical: use special module that collapses last dim to embedding_dim - - # Regular categorical sequences -> nn.Embedding (adds embedding dim) - if isinstance( - processor, - ( - SequenceProcessor, - StageNetProcessor, - NestedSequenceProcessor, - DeepNestedSequenceProcessor, - ), - ): - vocab_size = len(processor.code_vocab) - - if isinstance( - processor, (NestedSequenceProcessor, DeepNestedSequenceProcessor) - ): - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=0, - ) - else: - self.embedding_layers[field_name] = nn.Embedding( - num_embeddings=vocab_size, - embedding_dim=embedding_dim, - padding_idx=0, - ) - - # Optional pretrained initialization (e.g., GloVe). - if pretrained_emb_path is not None: - if isinstance(pretrained_emb_path, str): - path = pretrained_emb_path - else: - path = pretrained_emb_path.get(field_name) - if path: - init_embedding_with_pretrained( - self.embedding_layers[field_name], - processor.code_vocab, - path, - embedding_dim=embedding_dim, - normalize=normalize_pretrained, - freeze=freeze_pretrained, - ) - - # Numeric features (including deep nested floats) -> nn.Linear over last dim - elif isinstance( - processor, - ( - TimeseriesProcessor, - StageNetTensorProcessor, - NestedFloatsProcessor, - DeepNestedFloatsProcessor, - ), - ): - # Assuming processor.size() returns the last-dim size - in_features = processor.size() - self.embedding_layers[field_name] = nn.Linear( - in_features=in_features, out_features=embedding_dim - ) - - elif isinstance(processor, TensorProcessor): - # Infer size from first sample - sample_tensor = None - for sample in dataset: - if field_name in sample: - sample_tensor = processor.process(sample[field_name]) - break - if sample_tensor is not None: - input_size = ( - sample_tensor.shape[-1] if sample_tensor.dim() > 0 else 1 - ) - self.embedding_layers[field_name] = nn.Linear( - in_features=input_size, out_features=embedding_dim - ) - - elif isinstance(processor, MultiHotProcessor): - num_categories = processor.size() - self.embedding_layers[field_name] = nn.Linear( - in_features=num_categories, out_features=embedding_dim - ) - - # Smart Processor (Token-based) -> Transformers - elif hasattr(processor, "is_token") and processor.is_token(): - try: - from transformers import AutoModel - except ImportError: - raise ImportError( - "Please install `transformers` to use token-based processors." - ) - - # Load the model - self.embedding_layers[field_name] = AutoModel.from_pretrained( - processor.tokenizer_model - ) - - # Check if we need projection - if ( - self.embedding_layers[field_name].config.hidden_size - != self.embedding_dim - ): - self.embedding_layers[f"{field_name}_proj"] = nn.Linear( - self.embedding_layers[field_name].config.hidden_size, - self.embedding_dim, - ) - - else: - print( - "Warning: No embedding created for field due to lack of compatible processor:", - field_name, - ) - - def forward( - self, - inputs: Dict[str, torch.Tensor], - masks: Dict[str, torch.Tensor] = None, - output_mask: bool = False, - ) -> ( - Dict[str, torch.Tensor] - | tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]] - ): - - embedded: Dict[str, torch.Tensor] = {} - out_masks: Dict[str, torch.Tensor] = {} if output_mask else None - - for field_name, tensor in inputs.items(): - processor = self.dataset.input_processors.get(field_name, None) - - if field_name not in self.embedding_layers: - # No embedding layer -> passthrough - embedded[field_name] = tensor - continue - - # Check if it's a transformer model - layer = self.embedding_layers[field_name] - - # Check for transformers.PreTrainedModel (but without importing if possible, use class name check) - # or check if it has 'config' attribute - if hasattr(layer, "config") and hasattr(layer, "forward"): - # It's likely a transformer - tensor = tensor.to(self.device).long() # Ensure LongTensor for IDs - - mask = None - if masks is not None and field_name in masks: - mask = masks[field_name].to(self.device) - - # Handle 3D input (Batch, Num_Notes, Seq_Len) - is_3d = inputs[field_name].dim() == 3 - - if is_3d: - b, n, l = inputs[field_name].shape - tensor = tensor.view(b * n, l) - if mask is not None: - mask = mask.view(b * n, l) - - # Forward pass through transformer - output = layer(input_ids=tensor, attention_mask=mask) - x = output.last_hidden_state # (Batch, Seq, Hidden) - - if is_3d: - # If we had 3D input, we MUST pool the sequence dim (L) to get one vector per note - # Resulting shape: (B, N, H) - - # Pool L dim -> (B*N, H) using CLS token (index 0) - x = x[:, 0, :] - - # Check projections - if f"{field_name}_proj" in self.embedding_layers: - x = self.embedding_layers[f"{field_name}_proj"](x) - - x = x.view(b, n, -1) - - else: - # 2D input (Batch, Seq) -> (Batch, Seq, Hidden) - # No pooling, treating as sequence of tokens (word embeddings) - if f"{field_name}_proj" in self.embedding_layers: - x = self.embedding_layers[f"{field_name}_proj"](x) - - embedded[field_name] = x - - else: - # Standard layers - tensor = tensor.to(self.device) - embedded[field_name] = layer(tensor) - - if output_mask: - # Generate a mask for this field - # For transformers, we might already have a mask, or use pad token - if masks is not None and field_name in masks: - out_masks[field_name] = masks[field_name].to(self.device) - elif hasattr(processor, "code_vocab"): - pad_idx = processor.code_vocab.get("", 0) - out_masks[field_name] = tensor != pad_idx - else: - # Default mask generation (e.g. for simple linear layers where 0 might be padding?) - # Be careful changing this behavior. - # Previous code: - # masks[field_name] = (tensor != pad_idx) -> where pad_idx was 0 default - pad_idx = 0 - out_masks[field_name] = tensor != pad_idx - - if output_mask: - return embedded, out_masks - else: - return embedded - - def __repr__(self) -> str: - return f"EmbeddingModel(embedding_layers={self.embedding_layers})" diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py index 3784bc0c0..0b2771989 100644 --- a/pyhealth/models/unified_embedding.py +++ b/pyhealth/models/unified_embedding.py @@ -1,341 +1,16 @@ -"""UnifiedMultimodalEmbeddingModel — temporally aligned multimodal embedding. +"""Deprecated import path for the unified multimodal encoder. -Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor`` -subclasses ), embeds each event with a modality-specific encoder, then -interleaves all events on a shared timeline by sorting on timestamp and adding -sinusoidal time embeddings + learned modality-type embeddings. - -Output shape: ``(B, S_total, E')`` — a single sequence of events usable by -any downstream sequence model (Transformer, Mamba, RNN, …). - -Quickstart:: - - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel - from pyhealth.datasets.collate import collate_temporal - model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128) - # inside forward: - # inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...} - out = model(inputs) - # out["sequence"]: (B, S_total, 128) - # out["mask"]: (B, S_total) — 1 = real event, 0 = padding - # out["time"]: (B, S_total) — hours from first event +Use ``pyhealth.models.embedding`` instead. This module re-exports the live +classes so older ``from pyhealth.models.unified_embedding import ...`` call +sites keep working against one implementation. """ -from __future__ import annotations - -import math -from typing import Any - -import torch -import torch.nn as nn - -from pyhealth.processors.base_processor import ModalityType, TemporalFeatureProcessor - - -# ── Helpers ─────────────────────────────────────────────────────────────────── - - -class SinusoidalTimeEmbedding(nn.Module): - """Multi-scale sinusoidal embedding for times in hours. - - Wavelengths are spaced geometrically from ``min_hours`` to ``max_hours``. - The previous encoding wrapped every 720 hours. - - Args: - dim: Output embedding dimension (must be even). - max_hours: Longest wavelength in hours. Default 87600 (10 years). - min_hours: Shortest wavelength in hours. Default 1.0. - - Shape: - Input: ``(*, )`` float tensor of times in hours - Output: ``(*, dim)`` - """ - - def __init__( - self, - dim: int, - max_hours: float = 87600.0, - min_hours: float = 1.0, - ): - super().__init__() - assert dim % 2 == 0, f"dim must be even, got {dim}" - if min_hours <= 0 or max_hours <= min_hours: - raise ValueError( - f"need 0 < min_hours < max_hours, got min={min_hours}, max={max_hours}" - ) - self.dim = dim - self.max_hours = float(max_hours) - self.min_hours = float(min_hours) - half = dim // 2 - periods = torch.exp( - torch.linspace( - math.log(self.min_hours), - math.log(self.max_hours), - half, - dtype=torch.float32, - ) - ) - self.register_buffer("freqs", 2 * math.pi / periods) - - def forward(self, t: torch.Tensor) -> torch.Tensor: - """:param t: ``(...,)`` float, times in hours.""" - args = t.unsqueeze(-1).to(dtype=self.freqs.dtype) * self.freqs - return torch.cat([args.sin(), args.cos()], dim=-1) - - -def _build_image_encoder(embedding_dim: int) -> nn.Module: - """Lightweight 5-layer CNN encoder: C × H × W → embedding_dim. - - Uses ``torchvision.models.resnet18`` pre-trained backbone, strips the - final FC layer, and adds a projection to ``embedding_dim``. Falls back to - a toy Conv-pool-flatten network if torchvision is not installed. - """ - try: - import torchvision.models as tv - - backbone = tv.resnet18(weights=None) - in_features = backbone.fc.in_features - backbone.fc = nn.Linear(in_features, embedding_dim) - return backbone - except ImportError: - # Minimal fallback: single conv → global avg pool → linear - return nn.Sequential( - nn.Conv2d(3, 32, 3, padding=1), - nn.ReLU(), - nn.AdaptiveAvgPool2d(1), - nn.Flatten(), - nn.Linear(32, embedding_dim), - ) - - -# ── Main model ─────────────────────────────────────────────────────────────── - - -class UnifiedMultimodalEmbeddingModel(nn.Module): - """Embed heterogeneous temporal features into a single aligned sequence. - - **All** input processors must be ``TemporalFeatureProcessor`` subclasses. - Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``) - are rejected with a clear error — use the existing ``EmbeddingModel`` for - those fields. - - Algorithm - --------- - For each temporal field: - - 1. Route ``inputs[field]["value"]`` through a modality-specific encoder → - ``(B, N_i, E')`` per-event embeddings. - 2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours). - 3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or - ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if - token-level. - - Then: - - 4. Concatenate across all fields → ``(B, S_total, E')``. - 5. Sort events along dim=1 by timestamp (ascending). - 6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``. - 7. Return ``{"sequence", "time", "mask", "type_ids"}``. - - Args: - processors: ``dict[field_name, TemporalFeatureProcessor]`` — the - processors for each temporal field in the dataset. Pass - ``dataset.input_processors`` directly. - embedding_dim: Shared embedding dimension ``E'``. - time_embedding: ``"sinusoidal"`` (default) or ``"learned"``. - max_time_hours: Longest wavelength of the time embedding, in hours. - Defaults to 87600 (10 years). - - Example:: - - model = UnifiedMultimodalEmbeddingModel( - processors=dataset.input_processors, - embedding_dim=128, - ) - # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}} - out = model(inputs) - seq = out["sequence"] # (B, S_total, 128) - mask = out["mask"] # (B, S_total) float, 1=valid 0=pad - """ - - def __init__( - self, - processors: dict[str, Any], - embedding_dim: int = 128, - time_embedding: str = "sinusoidal", - max_time_hours: float = 87600.0, - ): - super().__init__() - self.embedding_dim = embedding_dim - - self.encoders: nn.ModuleDict = nn.ModuleDict() - self.projections: nn.ModuleDict = nn.ModuleDict() - self.modality_types: dict[str, ModalityType] = {} - - for field_name, processor in processors.items(): - if not isinstance(processor, TemporalFeatureProcessor): - raise TypeError( - f"UnifiedMultimodalEmbeddingModel requires every input processor " - f"to be a TemporalFeatureProcessor subclass, but '{field_name}' " - f"uses {type(processor).__name__}. For non-temporal fields use " - f"the existing EmbeddingModel." - ) - - m = processor.modality() - self.modality_types[field_name] = m - - if m == ModalityType.CODE: - vocab_size = processor.value_dim() - self.encoders[field_name] = nn.Embedding( - vocab_size, embedding_dim, padding_idx=0 - ) - - elif m == ModalityType.TEXT: - if processor.is_token(): - from transformers import AutoModel - - bert = AutoModel.from_pretrained(processor.tokenizer_model) - self.encoders[field_name] = bert - hidden = bert.config.hidden_size - if hidden != embedding_dim: - self.projections[field_name] = nn.Linear(hidden, embedding_dim) - else: - raise ValueError( - f"TEXT processor '{field_name}' must use a tokenizer " - f"(set tokenizer_model=...) to be used with " - f"UnifiedMultimodalEmbeddingModel." - ) - - elif m == ModalityType.IMAGE: - self.encoders[field_name] = _build_image_encoder(embedding_dim) - - elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL): - in_features = processor.value_dim() - self.encoders[field_name] = nn.Linear(in_features, embedding_dim) - - else: - raise NotImplementedError( - f"No encoder implemented for modality {m!r} (field '{field_name}')." - ) - - # Shared type embedding — one vector per unique modality in this dataset - unique_modalities = sorted(set(self.modality_types.values())) - self._modality_to_idx: dict[ModalityType, int] = { - mod: i for i, mod in enumerate(unique_modalities) - } - self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim) - - # Time embedding - if time_embedding == "sinusoidal": - self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours) - else: - raise NotImplementedError("Only 'sinusoidal' time embedding is implemented.") - - # ── Forward ─────────────────────────────────────────────────────────────── - - def forward( - self, - inputs: dict[str, dict[str, torch.Tensor]], - ) -> dict[str, torch.Tensor]: - """Encode and temporally align all temporal features. - - Args: - inputs: ``{field_name: {"value": Tensor, "time": Tensor, - "mask": Tensor (optional)}}`` - — one dict per temporal feature, exactly as produced by - ``collate_temporal``. - - Returns: - A dict with keys: - - * ``"sequence"`` — ``(B, S_total, E')`` temporally-sorted events - * ``"time"`` — ``(B, S_total)`` timestamps (hours) - * ``"mask"`` — ``(B, S_total)`` 1=real event, 0=padding - * ``"type_ids"`` — ``(B, S_total)`` modality index per event - """ - all_embeddings: list[torch.Tensor] = [] - all_times: list[torch.Tensor] = [] - all_masks: list[torch.Tensor] = [] - all_types: list[torch.Tensor] = [] - - for field_name, feat_dict in inputs.items(): - value = feat_dict["value"] # (B, N_i, ...) or (B, S, F) - time = feat_dict["time"] # (B, N_i) - mask = feat_dict.get("mask") - - if time is None: - # Fallback: treat every event as occurring at t=0 - time = torch.zeros(value.shape[:2], device=value.device) - - modality = self.modality_types[field_name] - encoder = self.encoders[field_name] - - # ── Encode ──────────────────────────────────────────────────── - if modality == ModalityType.CODE: - emb = encoder(value) # (B, S, E') - - elif modality == ModalityType.TEXT: - b, n, l = value.shape - flat_ids = value.view(b * n, l) - flat_mask = mask.view(b * n, l) if mask is not None else None - out = encoder(input_ids=flat_ids, attention_mask=flat_mask) - cls_emb = out.last_hidden_state[:, 0, :] # (B*N, H) - if field_name in self.projections: - cls_emb = self.projections[field_name](cls_emb) - emb = cls_emb.view(b, n, -1) # (B, N, E') - - elif modality == ModalityType.IMAGE: - b, n, c, h, w = value.shape - flat_imgs = value.view(b * n, c, h, w) - img_emb = encoder(flat_imgs) # (B*N, E') - emb = img_emb.view(b, n, -1) - - else: # NUMERIC / SIGNAL - emb = encoder(value) # (B, T, E') - - # ── Build event-level validity mask ─────────────────────────── - if mask is None: - event_mask = torch.ones(emb.shape[:2], device=emb.device) - else: - if mask.dim() > time.dim(): - # token-level (B, N, L) → event-level (B, N) - event_mask = (mask.sum(dim=-1) > 0).float() - else: - event_mask = mask.float() - - # ── Modality type indices ───────────────────────────────────── - type_idx = self._modality_to_idx[modality] - type_ids = torch.full( - emb.shape[:2], type_idx, dtype=torch.long, device=emb.device - ) - - all_embeddings.append(emb) - all_times.append(time) - all_masks.append(event_mask) - all_types.append(type_ids) - - # ── Concatenate across all fields ───────────────────────────────── - cat_emb = torch.cat(all_embeddings, dim=1) # (B, S_total, E') - cat_time = torch.cat(all_times, dim=1) # (B, S_total) - cat_mask = torch.cat(all_masks, dim=1) # (B, S_total) - cat_types = torch.cat(all_types, dim=1) # (B, S_total) - - # ── Sort by time ────────────────────────────────────────────────── - sort_idx = cat_time.argsort(dim=1) - cat_emb = cat_emb.gather( - 1, sort_idx.unsqueeze(-1).expand_as(cat_emb) - ) - cat_time = cat_time.gather(1, sort_idx) - cat_mask = cat_mask.gather(1, sort_idx) - cat_types = cat_types.gather(1, sort_idx) - # ── Add time + type embeddings ──────────────────────────────────── - time_emb = self.time_embed(cat_time) # (B, S_total, E') - type_emb = self.type_embedding(cat_types) # (B, S_total, E') - final = cat_emb + time_emb + type_emb # (B, S_total, E') +from pyhealth.models.embedding.unified import ( + SinusoidalTimeEmbedding, + UnifiedMultimodalEmbeddingModel, +) - return { - "sequence": final, # (B, S_total, E') - "time": cat_time, # (B, S_total) - "mask": cat_mask, # (B, S_total) - "type_ids": cat_types, # (B, S_total) - } +__all__ = [ + "SinusoidalTimeEmbedding", + "UnifiedMultimodalEmbeddingModel", +] diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 99b78cbae..f2e3f5fac 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -168,8 +168,22 @@ def test_collate_temporal_variable_length(): # ── 5. SinusoidalTimeEmbedding ──────────────────────────────────────────────── +def test_legacy_unified_import_is_the_package(): + from pyhealth.models.embedding.unified import ( + SinusoidalTimeEmbedding as PackageTime, + UnifiedMultimodalEmbeddingModel as PackageModel, + ) + from pyhealth.models.unified_embedding import ( + SinusoidalTimeEmbedding as LegacyTime, + UnifiedMultimodalEmbeddingModel as LegacyModel, + ) + + assert PackageTime is LegacyTime + assert PackageModel is LegacyModel + + def test_sinusoidal_time_embedding_shape(): - from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding + from pyhealth.models.embedding import SinusoidalTimeEmbedding emb = SinusoidalTimeEmbedding(dim=64) t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]]) # (2, 3) out = emb(t) @@ -177,7 +191,7 @@ def test_sinusoidal_time_embedding_shape(): def test_sinusoidal_different_times_differ(): - from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding + from pyhealth.models.embedding import SinusoidalTimeEmbedding emb = SinusoidalTimeEmbedding(dim=32) t0 = emb(torch.tensor([0.0])) t1 = emb(torch.tensor([24.0])) @@ -185,7 +199,7 @@ def test_sinusoidal_different_times_differ(): def test_sinusoidal_does_not_alias_every_720h(): - from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding + from pyhealth.models.embedding import SinusoidalTimeEmbedding emb = SinusoidalTimeEmbedding(dim=32) t6 = emb(torch.tensor([6.0])) t726 = emb(torch.tensor([726.0])) @@ -214,7 +228,7 @@ def _make_code_processors_and_inputs(batch_size=2, seq_len=5): def test_unified_model_code_only(): - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64) @@ -232,7 +246,7 @@ def test_unified_model_code_only(): def test_unified_model_rejects_non_temporal(): - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel from pyhealth.processors import SequenceProcessor bad_proc = SequenceProcessor() @@ -242,7 +256,7 @@ def test_unified_model_rejects_non_temporal(): def test_unified_model_gradient_flow(): """Loss.backward() should propagate through time + type embeddings.""" - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=32) @@ -258,7 +272,7 @@ def test_unified_model_gradient_flow(): def test_unified_model_time_sort(): """Events should be sorted by time ascending in the output.""" - from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel from pyhealth.processors import StageNetProcessor samples = [{"c": (None, ["a", "b"])}] From b50bf2f1efd067706a67ea08abfae03f528e30dc Mon Sep 17 00:00:00 2001 From: Rian354 Date: Wed, 19 Aug 2026 16:06:00 -0400 Subject: [PATCH 26/27] Point AMP autocast at the trainer device instead of hardcoding CUDA. use_amp=True on CPU or MPS warned that CUDA was unavailable and skipped mixed precision. Unknown amp_dtype spellings were already rejected. Co-authored-by: Cursor --- pyhealth/trainer.py | 13 ++++++++++++- tests/test_p1_amp_dtype.py | 7 +++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index 6d7ef176f..2085221b3 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -54,6 +54,14 @@ def resolve_amp_dtype(amp_dtype: str, use_amp: bool = False) -> torch.dtype: return dtype +def autocast_device_type(device) -> str: + """Device type string for ``torch.autocast``, from the trainer device.""" + name = str(device).split(":", 1)[0].lower() + if name in ("cuda", "cpu", "mps"): + return name + return "cpu" + + def is_best(best_score: float, score: float, monitor_criterion: str) -> bool: if monitor_criterion == "max": return score > best_score @@ -277,7 +285,10 @@ def train( data = next(data_iterator) # forward (with optional AMP) if use_amp: - with torch.autocast(device_type="cuda", dtype=_amp_dtype): + with torch.autocast( + device_type=autocast_device_type(self.device), + dtype=_amp_dtype, + ): output = self.model(**data) loss = output["loss"] / accumulation_steps else: diff --git a/tests/test_p1_amp_dtype.py b/tests/test_p1_amp_dtype.py index c7cbfbe9b..57fb1b959 100644 --- a/tests/test_p1_amp_dtype.py +++ b/tests/test_p1_amp_dtype.py @@ -28,6 +28,13 @@ def test_known_spellings_map_to_the_named_dtype(self): self.assertIs(resolve_amp_dtype("fp16"), torch.float16) self.assertIs(resolve_amp_dtype("float16"), torch.float16) + def test_autocast_device_follows_the_trainer_device(self): + from pyhealth.trainer import autocast_device_type + + self.assertEqual(autocast_device_type("cuda:0"), "cuda") + self.assertEqual(autocast_device_type("cpu"), "cpu") + self.assertEqual(autocast_device_type("mps"), "mps") + def test_unknown_spelling_raises(self): from pyhealth.trainer import resolve_amp_dtype From 8d4a4c9dcd01044a368a4b9ee8cff23846b001b7 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Wed, 19 Aug 2026 16:06:04 -0400 Subject: [PATCH 27/27] Default to a full stay and stamp admission-context notes at admit. Will's protocol is through discharge, so the 24h class/runner default is gone. Discharge-section text is still written at the end of the stay; stamping it at charttime leaked length of stay. Radiology stays at exam time. Cache version 4. Co-authored-by: Cursor --- .../unified_embedding_e2e_mimic4.py | 10 ++- pyhealth/tasks/multimodal_mimic4.py | 65 +++++++++++-------- tests/test_p1_observation_window.py | 17 +++-- tests/test_p1_time_axis.py | 51 ++++++++++++++- tests/test_p2_runner_measurement.py | 7 ++ 5 files changed, 116 insertions(+), 34 deletions(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 5db04d59c..d0806c8e7 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -576,7 +576,15 @@ def parse_args() -> argparse.Namespace: "patient-limited smoke, and omit both for the full table." ), ) - parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument( + "--observation-window-hours", + type=int, + default=None, + help=( + "If set, collect labs/CXR/radiology only this many hours from each " + "admission. Default: full stay (through discharge)." + ), + ) parser.add_argument( "--freeze-encoder", action="store_true", diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index f0b8e6011..ff9022cbe 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -85,13 +85,17 @@ def __init__( self.window_hours = window_hours # Task cache key is uuid5 over {**vars(task), schemas}. Bump when # emitted data changes so leaky caches cannot be reused. - # 1: empty events instead of placeholders; per-admission lab window. + # 1: empty events instead of placeholders; per-admission collection span. # 2: CXR/notes_labs_cxr no longer drop later stays against the first # admission's clock (admission_time >= first_admit + window_hours). # 3: event times are hours from the first stay in the sample, not # reset per admission (reset times made stay 2 at +6h sort with # stay 1 at +6h). - self.emitted_data_version = 3 + # 4: class/runner default is full stay (window_hours=None); + # admission-context discharge sections stamped at admit, not + # charttime. v3 caches used a 24h class default and discharge + # charttime on those notes. + self.emitted_data_version = 4 @staticmethod def _clean_text(text: Optional[str]) -> Optional[str]: @@ -134,9 +138,10 @@ def _to_hours(delta_seconds: float) -> float: def _hours_since(cls, timestamp: datetime, origin: datetime) -> float: """Hours from ``origin`` to ``timestamp``. - Collection windows stay per admission. The value written onto the - unified timeline is hours from the first stay in this sample, so a - later stay at +6h does not sort with the first stay at +6h. + Collection windows stay per admission (full stay, or admit+window + when ``window_hours`` is set). The value written onto the unified + timeline is hours from the first stay in this sample, so a later + stay at +6h does not sort with the first stay at +6h. """ return cls._to_hours((timestamp - origin).total_seconds()) @@ -317,6 +322,7 @@ def _collect_notes( section_headers: Optional[List[str]] = None, fallback_to_full_note: bool = False, time_origin: Optional[datetime] = None, + event_time: Optional[datetime] = None, ) -> Tuple[List[str], List[float]]: """Collect notes of a given type for one admission. @@ -324,16 +330,20 @@ def _collect_notes( patient: Patient object. note_event_type: Event type string (e.g. "discharge", "radiology"). hadm_id: Admission ID to filter by. - admission_time: This stay's admit time (unused for the timeline - once ``time_origin`` is set; kept so existing call sites that - pass it positionally stay valid). + admission_time: This stay's admit time. Used as the timeline + origin when ``time_origin`` is omitted. start_time: Optional start of the time window. end_time: Optional end of the time window. section_headers: When provided, extract only these named sections from each note (lowercased match against parsed headers). - fallback_to_full_note: When True (default), falls back to the full - note text if no matching sections are found. When False, notes - with no matching sections are dropped entirely. + fallback_to_full_note: When True, falls back to the full note text + if no matching sections are found. When False, notes with no + matching sections are dropped entirely. + time_origin: Timeline origin. Hours are measured from here. + event_time: If set, stamp every collected note at this instant. + Discharge-section text is admission-context and should pass + this stay's admit. Radiology omits it and uses exam + ``charttime``. Returns: Tuple of (texts, hours from the sample's first stay). Empty lists @@ -361,8 +371,9 @@ def _collect_notes( continue origin = time_origin if time_origin is not None else admission_time + stamp = event_time if event_time is not None else note.timestamp texts.append(note_text) - note_times.append(self._hours_since(note.timestamp, origin)) + note_times.append(self._hours_since(stamp, origin)) except ( AttributeError ): # note object is missing .text or .timestamp attribute (e.g. malformed note) @@ -381,17 +392,17 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): Text is extracted from the MIMIC-IV discharge note by parsing the Chief Complaint, History of Present Illness, Past Medical History, and Medications on Admission sections — all of which describe the patient's state at the - start of the stay. The extracted text is assigned timestamp 0.0. + start of the stay. The extracted text is stamped at that stay's admit. Radiology reports are also included, parsed for their Indication and Impression sections and bounded to the same observation window as labs - (rather than timestamp 0.0), since — unlike the discharge summary — they + (rather than that stay's admit), since — unlike the discharge summary — they are written at exam time and describe findings from later in the stay. Fields: - admission_note_times: Admission-context discharge-note text at time - 0.0, plus in-window radiology note text at its exam-relative - timestamp. + admission_note_times: Admission-context discharge-note text stamped + at that stay's admit (hours from the first stay), plus in-window + radiology note text at exam time. labs: 10-dim lab vectors at each measurement timestamp. labs_mask: Boolean observation mask parallel to ``labs``. icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes @@ -399,7 +410,7 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): Args: window_hours: Hours from admission for lab collection. ``None`` - collects for the full admission span. Default: 24. + collects for the full admission span. Default: ``None``. include_icd: When ``True``, collect discharge-coded ICD codes and add ``icd_codes`` to the sample dict / input schema. Default: ``False``. MIMIC-IV timestamps those codes at ``dischtime``, so this leaks @@ -429,7 +440,7 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): def __init__( self, - window_hours: Optional[float] = 24, + window_hours: Optional[float] = None, include_icd: bool = False, ) -> None: super().__init__(window_hours=window_hours) @@ -487,6 +498,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission_time, section_headers=self.DISCHARGE_CLINICAL_HEADERS, time_origin=time_origin, + event_time=admission_time, ) all_note_texts.extend(note_texts) all_note_times.extend(note_times) @@ -556,9 +568,9 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): experiments, same as ``NotesLabsMIMIC4``. Fields: - admission_note_times: Admission-context discharge-note text at time - 0.0, plus in-window radiology note text at its exam-relative - timestamp. + admission_note_times: Admission-context discharge-note text stamped + at that stay's admit (hours from the first stay), plus in-window + radiology note text at exam time. labs: 10-dim lab vectors at each measurement timestamp. labs_mask: Boolean observation mask parallel to ``labs``. cxr_image_times: In-window CXR image paths at their exam-relative @@ -568,7 +580,7 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): Args: window_hours: Hours from admission for lab/CXR collection. - ``None`` collects for the full admission span. Default: 24. + ``None`` collects for the full admission span. Default: ``None``. include_icd: When ``True``, collect discharge-coded ICD codes and add ``icd_codes`` to the sample dict / input schema. Default: ``False``. MIMIC-IV timestamps those codes at ``dischtime``, so this leaks @@ -606,7 +618,7 @@ class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): def __init__( self, - window_hours: Optional[float] = 24, + window_hours: Optional[float] = None, include_icd: bool = False, ) -> None: super().__init__(window_hours=window_hours) @@ -666,6 +678,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission_time, section_headers=self.DISCHARGE_CLINICAL_HEADERS, time_origin=time_origin, + event_time=admission_time, ) all_note_texts.extend(note_texts) all_note_times.extend(note_times) @@ -754,7 +767,7 @@ class LabsMIMIC4(BaseMultimodalMIMIC4Task): Args: window_hours: Hours from admission to collect lab measurements. - ``None`` collects for the full admission span. Default: 24. + ``None`` collects for the full admission span. Default: ``None``. """ PADDING: int = 0 @@ -767,7 +780,7 @@ class LabsMIMIC4(BaseMultimodalMIMIC4Task): } output_schema: ClassVar[Dict] = {"mortality": "binary"} - def __init__(self, window_hours: Optional[float] = 24) -> None: + def __init__(self, window_hours: Optional[float] = None) -> None: super().__init__(window_hours=window_hours) def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] diff --git a/tests/test_p1_observation_window.py b/tests/test_p1_observation_window.py index aa59a682c..6cb694ec4 100644 --- a/tests/test_p1_observation_window.py +++ b/tests/test_p1_observation_window.py @@ -14,7 +14,10 @@ CXR / ``notes_labs_cxr`` still skipped those later stays with ``admission_time >= first_admit + window_hours``. That skip is gone. -``emitted_data_version`` is 3 so caches from version 1-2 cannot be reused. +``emitted_data_version`` is 4 so caches from version 1-3 cannot be reused. +The table protocol default is full stay (``window_hours=None``). Passing +``window_hours=24`` still caps collection per admission; CXR arms must not +skip later stays against the first admission's clock. Repro:: @@ -92,7 +95,7 @@ def test_window_change_invalidates_the_cache(self): task = m.LabsMIMIC4(window_hours=24) self.assertIsNotNone(vars(task).get("emitted_data_version")) - self.assertGreaterEqual(task.emitted_data_version, 3) + self.assertGreaterEqual(task.emitted_data_version, 4) def cache_key(t, drop_version=False): v = dict(vars(t)) @@ -119,16 +122,18 @@ def test_window_none_still_collects_through_discharge(self): discharge = admit + timedelta(days=9) self.assertEqual(task._admission_window_end(admit, discharge), discharge) - def test_notes_labs_defaults_to_a_24h_window(self): + def test_protocol_default_is_full_stay(self): from pyhealth.tasks.multimodal_mimic4 import ( + CXRMIMIC4, LabsMIMIC4, NotesLabsCXRMIMIC4, NotesLabsMIMIC4, ) - self.assertEqual(NotesLabsMIMIC4().window_hours, 24) - self.assertEqual(NotesLabsCXRMIMIC4().window_hours, 24) - self.assertEqual(LabsMIMIC4().window_hours, 24) + self.assertIsNone(NotesLabsMIMIC4().window_hours) + self.assertIsNone(NotesLabsCXRMIMIC4().window_hours) + self.assertIsNone(LabsMIMIC4().window_hours) + self.assertIsNone(CXRMIMIC4().window_hours) def test_discharge_coded_icd_is_not_a_mortality_task(self): from pyhealth.tasks import multimodal_mimic4 as m diff --git a/tests/test_p1_time_axis.py b/tests/test_p1_time_axis.py index 521687ad3..04f21f137 100644 --- a/tests/test_p1_time_axis.py +++ b/tests/test_p1_time_axis.py @@ -3,7 +3,11 @@ The sample is patient-level: every admission up to the first death is one sequence. Event times were hours from *that stay's* admit, then concatenated, so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay -(admit, admit+window]. Times are hours from the first stay in the sample. +(admit through discharge, or admit+window if ``window_hours`` is set). Times +are hours from the first stay in the sample. + +Admission-context discharge sections are stamped at that stay's admit, not +at the discharge note's ``charttime``. Radiology stays at exam time. Single-stay patients are unchanged. Time embeddings use geometric wavelengths from 1 hour to 10 years so a later stay at +9606h does not @@ -81,3 +85,48 @@ def test_sinusoid_does_not_wrap_every_thirty_days(self): t9606 = emb(torch.tensor([9606.0])) self.assertFalse(torch.allclose(t6, t726, atol=1e-5)) self.assertFalse(torch.allclose(t6, t9606, atol=1e-5)) + + def test_discharge_sections_are_stamped_at_admit_not_charttime(self): + from types import SimpleNamespace + + from pyhealth.tasks.multimodal_mimic4 import NotesLabsCXRMIMIC4, NotesLabsMIMIC4 + + first = datetime(2180, 5, 6, 8, 0, 0) + stay2_admit = first + timedelta(days=400) + charttime = stay2_admit + timedelta(days=9) + note = SimpleNamespace( + text="Chief Complaint:\n\nshortness of breath\n\n", + timestamp=charttime, + ) + + class FakePatient: + def get_events(self, event_type, start=None, end=None, filters=None): + return [note] + + task = NotesLabsMIMIC4() + _texts, times = task._collect_notes( + FakePatient(), + "discharge", + hadm_id=1, + admission_time=stay2_admit, + section_headers=task.DISCHARGE_CLINICAL_HEADERS, + time_origin=first, + event_time=stay2_admit, + ) + self.assertEqual(times, [400 * 24.0]) + + _texts, rad_times = task._collect_notes( + FakePatient(), + "radiology", + hadm_id=1, + admission_time=stay2_admit, + time_origin=first, + ) + self.assertAlmostEqual(rad_times[0], 400 * 24.0 + 9 * 24.0) + + for cls in (NotesLabsMIMIC4, NotesLabsCXRMIMIC4): + self.assertIn( + "event_time=admission_time", + inspect.getsource(cls.__call__), + msg=cls.__name__, + ) diff --git a/tests/test_p2_runner_measurement.py b/tests/test_p2_runner_measurement.py index fb25ec4c7..84658451d 100644 --- a/tests/test_p2_runner_measurement.py +++ b/tests/test_p2_runner_measurement.py @@ -122,3 +122,10 @@ def test_jamba_cli_default_matches_the_library(self): args = _parse(mod, "--model", "jambaehr") self.assertEqual(args.jamba_transformer_layers, 2) self.assertEqual(args.jamba_mamba_layers, 6) + + def test_observation_window_defaults_to_full_stay(self): + mod = _load_runner() + args = _parse(mod) + self.assertIsNone(args.observation_window_hours) + args_24 = _parse(mod, "--observation-window-hours", "24") + self.assertEqual(args_24.observation_window_hours, 24)