From 86d2e7545187ea3a93d7f95a9da3f4f40b7ec5df Mon Sep 17 00:00:00 2001 From: John Wu Date: Wed, 15 Apr 2026 18:33:41 -0500 Subject: [PATCH 01/30] fixes --- .../multimodal_embedding_jamba_mimic4_cxr.py | 326 ++++++++++++++++++ .../multimodal_embedding_mamba_mimic4_cxr.py | 15 +- .../multimodal_embedding_mlp_mimic4_cxr.py | 321 +++++++++++++++++ .../multimodal_embedding_rnn_mimic4_cxr.py | 322 +++++++++++++++++ ...imodal_embedding_transformer_mimic4_cxr.py | 314 +++++++++++++++++ examples/mortality_prediction/readme.md | 11 +- pyhealth/models/transformer.py | 96 +++++- pyhealth/processors/time_image_processor.py | 69 +--- 8 files changed, 1400 insertions(+), 74 deletions(-) create mode 100644 examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py create mode 100644 examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py create mode 100644 examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py create mode 100644 examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py diff --git a/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py new file mode 100644 index 000000000..3fc268825 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py @@ -0,0 +1,326 @@ +"""Unified multimodal embedding + JambaEHR runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +JambaEHR is a hybrid Transformer + Mamba architecture that interleaves +attention and SSM (state-space model) blocks. + +Default roots are set to shared PhysioNet mounts: +- 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 + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_jamba_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_jamba_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import JambaEHR, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +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_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_attn{args.num_transformer_layers}" + f"_mamba{args.num_mamba_layers}" + f"_heads{args.heads}" + f"_state{args.state_size}" + f"_conv{args.conv_kernel}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_jamba_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + 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 use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = JambaEHR( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + num_transformer_layers=args.num_transformer_layers, + num_mamba_layers=args.num_mamba_layers, + heads=args.heads, + dropout=args.dropout, + state_size=args.state_size, + conv_kernel=args.conv_kernel, + unified_embedding=unified, + ) + print(f"JambaEHR unified mode: {model._use_unified}") + + 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 + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + JambaEHR on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-transformer-layers", type=int, default=2) + parser.add_argument("--num-mamba-layers", type=int, default=6) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--state-size", type=int, default=16) + parser.add_argument("--conv-kernel", type=int, default=4) + parser.add_argument("--dropout", type=float, default=0.3) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py index e5a0672af..03a42cfe8 100644 --- a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py @@ -185,7 +185,7 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: trainer = Trainer( model=model, - metrics=["accuracy"], + metrics=["pr_auc", "roc_auc"], device=args.device, enable_logging=True, output_path=_build_run_output_path(args), @@ -220,6 +220,7 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: y_true, y_prob, _, patient_ids = trainer.inference( inference_loader, return_patient_ids=True ) + scores = trainer.evaluate(inference_loader) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -236,6 +237,9 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: print(" peak_train_vram_mb: N/A (non-CUDA device)") else: print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") return len(patient_ids), y_true.shape[0] @@ -281,7 +285,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--observation-window-hours", type=int, default=24) parser.add_argument("--epochs", type=int, default=1) - parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--batch-size", type=int, default=1) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--device", type=str, default="cuda:1") parser.add_argument("--num-workers", type=int, default=16) @@ -290,8 +294,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dev", action="store_true") parser.add_argument("--quick-test", action="store_true") parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) args = parser.parse_args() + if args.condor: + args.device = "cuda" if args.quick_test: args.dev = True args.epochs = 1 diff --git a/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py new file mode 100644 index 000000000..9f37d7153 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py @@ -0,0 +1,321 @@ +"""Unified multimodal embedding + MLP runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- 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 + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_mlp_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_mlp_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import MLP, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +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_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_hidden{args.hidden_dim}" + f"_layers{args.num_layers}" + f"_act{args.activation}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_mlp_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + 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 use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = MLP( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + hidden_dim=args.hidden_dim, + n_layers=args.num_layers, + activation=args.activation, + unified_embedding=unified, + ) + print(f"MLP unified mode: {model._use_unified}") + + 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 + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + MLP on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument( + "--activation", + type=str, + default="relu", + choices=["relu", "tanh", "sigmoid", "leaky_relu", "elu"], + ) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py new file mode 100644 index 000000000..e58ba7243 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py @@ -0,0 +1,322 @@ +"""Unified multimodal embedding + RNN runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- 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 + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_rnn_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_rnn_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import RNN, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +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_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_hidden{args.hidden_dim}" + f"_{args.rnn_type.lower()}" + f"_layers{args.num_layers}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_rnn_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + 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 use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = 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.num_layers, + dropout=args.dropout, + ) + print(f"RNN unified mode: {model._use_unified}") + + 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 + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + RNN on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--hidden-dim", type=int, default=128) + parser.add_argument( + "--rnn-type", + type=str, + default="GRU", + choices=["GRU", "LSTM", "RNN"], + ) + parser.add_argument("--num-layers", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py new file mode 100644 index 000000000..6a1727571 --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py @@ -0,0 +1,314 @@ +"""Unified multimodal embedding + Transformer runner for MIMIC-IV + CXR. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- 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 + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_transformer_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_transformer_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import Transformer, UnifiedMultimodalEmbeddingModel +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +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_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_layers{args.num_layers}" + f"_heads{args.heads}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), "output", "multimodal_embedding_transformer_mimic4_cxr", run_tag + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + 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 use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " f"schema={processor.schema()}") + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = Transformer( + dataset=sample_dataset, + embedding_dim=args.embedding_dim, + heads=args.heads, + dropout=args.dropout, + num_layers=args.num_layers, + unified_embedding=unified, + ) + print(f"Transformer unified mode: {model._use_unified}") + + 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 + ) + + print( + "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run unified multimodal embedding + Transformer on " + "MIMIC-IV mortality with CXR." + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=1) + parser.add_argument("--heads", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).") diff --git a/examples/mortality_prediction/readme.md b/examples/mortality_prediction/readme.md index 5335ffd9a..fb82f291b 100644 --- a/examples/mortality_prediction/readme.md +++ b/examples/mortality_prediction/readme.md @@ -5,5 +5,14 @@ # Multimodality +## Model Variants -nohup python examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py --batch-size 1 > ../logs/multimodal_embedding_mamba_mimic4_cxr_b1.log & \ No newline at end of file +nohup python examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mamba_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mlp_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_rnn_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_transformer_mimic4_cxr_b1.log & + +nohup python examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_jamba_mimic4_cxr_b1.log & \ No newline at end of file diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py index 2bed5a3d4..f678403b2 100644 --- a/pyhealth/models/transformer.py +++ b/pyhealth/models/transformer.py @@ -55,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: @@ -165,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) @@ -247,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: @@ -257,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) @@ -298,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. @@ -387,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 @@ -463,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, ...], @@ -520,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) @@ -535,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) @@ -603,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] @@ -619,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) @@ -649,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 } diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py index ad53fa350..421998d07 100644 --- a/pyhealth/processors/time_image_processor.py +++ b/pyhealth/processors/time_image_processor.py @@ -124,16 +124,11 @@ def __init__( 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 " @@ -153,24 +148,14 @@ 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 _zero_image_tensor(self) -> torch.Tensor: @@ -193,9 +178,7 @@ def _zero_image_tensor(self) -> torch.Tensor: c = 3 return torch.zeros(c, self.image_size, self.image_size) - def _load_single_image( - self, path: Union[str, Path] - ) -> torch.Tensor: + 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 @@ -215,23 +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 - ): + 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 @@ -271,9 +247,7 @@ def fit( 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. @@ -319,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 = [] @@ -336,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] @@ -386,9 +353,5 @@ def __repr__(self) -> str: f"std={self.std}, " f"mode={self.mode}, " f"max_images={self.max_images}, " -<<<<<<< HEAD f"padding={self.padding!r})" -======= - f"padding={self.padding})" ->>>>>>> main - ) \ No newline at end of file + ) From 6eeda743a26c9597d136bce80b3457581dee657a Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 15 Apr 2026 22:46:46 -0500 Subject: [PATCH 02/30] Add script to run on sunlab --- .../unified_embedding_e2e_mimic4.py | 2 +- scripts/sunlab/run_ehrmamba_timewindow.sh | 38 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 scripts/sunlab/run_ehrmamba_timewindow.sh diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 0db4b0d97..f72e41f8e 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -299,7 +299,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, - choices=["stagenet", "icd_labs", "clinical_notes_icd_labs"], + choices=["icd_labs", "clinical_notes_icd_labs"], default="stagenet", ) parser.add_argument( diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh new file mode 100644 index 000000000..fc4c49c32 --- /dev/null +++ b/scripts/sunlab/run_ehrmamba_timewindow.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# Run ehrmamba jobs for multiple seeds in tmux sessions +# Usage: bash scripts/run_ehrmamba.sh +set -euo pipefail + +SEEDS=(267573289 1872967241 706384748) +TIME_WINDOWS=(24 48 96) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="" + +cd "${PROJECT_DIR}" +mkdir -p logs + +for seed in "${SEEDS[@]}"; do + for tw in "${TIME_WINDOWS[@]}"; do + echo " Launching ehrmamba seed=${seed} time_window=${tw} in tmux" + tmux new -s "ehrmamba_${seed}_observationtimewindow_${tw}" -d "bash -lc ' + conda activate pyhealth2 && \ + CUDA_VISIBLE_DEVICES=4 python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \ + --cache-dir /home/${USER}/pyhealth_cache \ + --task icd_labs \ + --model ehrmamba \ + --embedding-dim 128 \ + --num-layers 2 \ + --observation-window-hours ${tw} \ + --mamba-state-size 16 \ + --mamba-conv-kernel 4 \ + --epochs 20 \ + --batch-size 8 \ + --seed ${seed} \ + --output-dir /home/${USER}/output/table2/ehrmamba_seed${seed}_observationtimewindow_${tw} \ + 2>&1 | tee logs/ehrmamba_seed${seed}_observationtimewindow_${tw}.log; exec bash'" + echo " Launched ehrmamba seed=${seed} time_window=${tw} → tmux session: ehrmamba_${seed}_observationtimewindow_${tw}" + done +done + +tmux attach -t "ehrmamba_${seed}_observationtimewindow_${tw}" \ No newline at end of file From 63718f1255c77ce3a6ae730cd970132f81a3ee6a Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 15 Apr 2026 22:49:44 -0500 Subject: [PATCH 03/30] Update script --- scripts/sunlab/run_ehrmamba_timewindow.sh | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh index fc4c49c32..6d9faf8aa 100644 --- a/scripts/sunlab/run_ehrmamba_timewindow.sh +++ b/scripts/sunlab/run_ehrmamba_timewindow.sh @@ -6,7 +6,7 @@ set -euo pipefail SEEDS=(267573289 1872967241 706384748) TIME_WINDOWS=(24 48 96) PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" -USER="" +USER="wp14" cd "${PROJECT_DIR}" mkdir -p logs @@ -33,6 +33,4 @@ for seed in "${SEEDS[@]}"; do 2>&1 | tee logs/ehrmamba_seed${seed}_observationtimewindow_${tw}.log; exec bash'" echo " Launched ehrmamba seed=${seed} time_window=${tw} → tmux session: ehrmamba_${seed}_observationtimewindow_${tw}" done -done - -tmux attach -t "ehrmamba_${seed}_observationtimewindow_${tw}" \ No newline at end of file +done \ No newline at end of file From 71a45ed005c1adf93cc807de5a0ca497acfc7080 Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 15 Apr 2026 22:51:58 -0500 Subject: [PATCH 04/30] small changes --- scripts/sunlab/run_ehrmamba_timewindow.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh index 6d9faf8aa..f50b7588f 100644 --- a/scripts/sunlab/run_ehrmamba_timewindow.sh +++ b/scripts/sunlab/run_ehrmamba_timewindow.sh @@ -6,7 +6,7 @@ set -euo pipefail SEEDS=(267573289 1872967241 706384748) TIME_WINDOWS=(24 48 96) PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" -USER="wp14" +USER="" cd "${PROJECT_DIR}" mkdir -p logs From 131ca34a092316a9c66c66d7307b626bf95661da Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 15 Apr 2026 21:03:13 -0700 Subject: [PATCH 05/30] Update run_ehrmamba_timewindow.sh --- scripts/sunlab/run_ehrmamba_timewindow.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh index f50b7588f..2ea666be9 100644 --- a/scripts/sunlab/run_ehrmamba_timewindow.sh +++ b/scripts/sunlab/run_ehrmamba_timewindow.sh @@ -14,7 +14,7 @@ mkdir -p logs for seed in "${SEEDS[@]}"; do for tw in "${TIME_WINDOWS[@]}"; do echo " Launching ehrmamba seed=${seed} time_window=${tw} in tmux" - tmux new -s "ehrmamba_${seed}_observationtimewindow_${tw}" -d "bash -lc ' + tmux new -s "ehrmamba_${seed}_observationtimewindow_${tw}" "bash -lc ' conda activate pyhealth2 && \ CUDA_VISIBLE_DEVICES=4 python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \ @@ -33,4 +33,4 @@ for seed in "${SEEDS[@]}"; do 2>&1 | tee logs/ehrmamba_seed${seed}_observationtimewindow_${tw}.log; exec bash'" echo " Launched ehrmamba seed=${seed} time_window=${tw} → tmux session: ehrmamba_${seed}_observationtimewindow_${tw}" done -done \ No newline at end of file +done From 02443ddb3a44075cf9e9953ad0ce5958900c52af Mon Sep 17 00:00:00 2001 From: John Wu Date: Thu, 16 Apr 2026 10:59:38 -0500 Subject: [PATCH 06/30] enable shared encoder usage for same tokenizers --- pyhealth/models/embedding/unified.py | 46 +++-- tests/test_unified_multimodal.py | 266 +++++++++++++++++++++------ 2 files changed, 242 insertions(+), 70 deletions(-) diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 5164d9955..2e5e9560d 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -214,6 +214,7 @@ def __init__( 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] = {} for field_name, processor in processors.items(): if not isinstance(processor, TemporalFeatureProcessor): @@ -307,6 +308,21 @@ def _build_text_encoder( embedding_dim: int, ) -> 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") @@ -314,23 +330,29 @@ def _build_text_encoder( ): self.encoders[field_name] = pre_built.transformer pre_dim = getattr(pre_built, "embedding_dim", embedding_dim) - if pre_dim != embedding_dim: - self.projections[field_name] = nn.Sequential( - pre_built.fc, - nn.Linear(pre_dim, embedding_dim), - ) - else: - self.projections[field_name] = pre_built.fc + _set_projection(pre_dim, pre_built.fc) return 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) + tokenizer_model = getattr(processor, "tokenizer_model", None) + if not tokenizer_model: + raise ValueError( + f"TEXT processor '{field_name}' is token-based but does not " + "define tokenizer_model." + ) + + shared_field = self._shared_text_field_by_model.get(tokenizer_model) + if shared_field is not None: + shared_encoder = self.encoders[shared_field] + else: + shared_encoder = AutoModel.from_pretrained(tokenizer_model) + self._shared_text_field_by_model[tokenizer_model] = field_name + + self.encoders[field_name] = shared_encoder + hidden = shared_encoder.config.hidden_size + _set_projection(hidden) else: raise ValueError( f"TEXT processor '{field_name}' must either supply a pre-built " diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 22c88785c..f17e0f20e 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -2,49 +2,73 @@ collate_temporal helper, and UnifiedMultimodalEmbeddingModel. Run with: - TOKENIZERS_PARALLELISM=false pytest tests/test_unified_multimodal.py -v + TOKENIZERS_PARALLELISM=false python tests/test_unified_multimodal.py """ + import math +import unittest from datetime import datetime, timedelta +from unittest.mock import patch -import pytest import torch import numpy as np # ── 1. TemporalFeatureProcessor ABC & ModalityType ──────────────────────────── + def test_modality_type_values(): from pyhealth.processors import ModalityType - assert ModalityType.CODE == "code" - assert ModalityType.TEXT == "text" - assert ModalityType.IMAGE == "image" + + assert ModalityType.CODE == "code" + assert ModalityType.TEXT == "text" + assert ModalityType.IMAGE == "image" assert ModalityType.NUMERIC == "numeric" def test_stagenet_is_temporal(): - from pyhealth.processors import StageNetProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + StageNetProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = StageNetProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.CODE def test_stagenet_tensor_is_temporal(): - from pyhealth.processors import StageNetTensorProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + StageNetTensorProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = StageNetTensorProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.NUMERIC def test_tuple_time_text_is_temporal(): - from pyhealth.processors import TupleTimeTextProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + TupleTimeTextProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = TupleTimeTextProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.TEXT def test_time_image_is_temporal(): - from pyhealth.processors import TimeImageProcessor, TemporalFeatureProcessor, ModalityType + from pyhealth.processors import ( + TimeImageProcessor, + TemporalFeatureProcessor, + ModalityType, + ) + p = TimeImageProcessor() assert isinstance(p, TemporalFeatureProcessor) assert p.modality() == ModalityType.IMAGE @@ -52,8 +76,10 @@ def test_time_image_is_temporal(): # ── 2. StageNetProcessor.process_temporal() ─────────────────────────────────── + def test_stagenet_process_temporal(): from pyhealth.processors import StageNetProcessor + samples = [{"codes": (None, ["A", "B", "C"])}] p = StageNetProcessor() p.fit(samples, "codes") @@ -63,13 +89,14 @@ def test_stagenet_process_temporal(): assert set(out.keys()) == {"value", "time"} assert out["value"].dtype == torch.long - assert out["time"].dtype == torch.float32 + assert out["time"].dtype == torch.float32 assert out["value"].shape == (3,) - assert out["time"].shape == (3,) + assert out["time"].shape == (3,) def test_stagenet_tensor_process_temporal(): from pyhealth.processors import StageNetTensorProcessor + samples = [{"vitals": ([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]])}] p = StageNetTensorProcessor() p.fit(samples, "vitals") @@ -77,35 +104,37 @@ def test_stagenet_tensor_process_temporal(): out = p.process_temporal(([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]])) assert set(out.keys()) == {"value", "time"} assert out["value"].shape == (2, 2) - assert out["time"].shape == (2,) + assert out["time"].shape == (2,) assert p.value_dim() == 2 assert p.modality().value == "numeric" # ── 3. TemporalTimeseriesProcessor ──────────────────────────────────────────── + def test_temporal_timeseries_basic(): from pyhealth.processors import TemporalTimeseriesProcessor p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=2)) ts = [ - datetime(2023, 1, 1, 0), - datetime(2023, 1, 1, 4), - datetime(2023, 1, 1, 8), + datetime(2023, 1, 1, 0), + datetime(2023, 1, 1, 4), + datetime(2023, 1, 1, 8), ] val = np.array([[120.0, 80.0], [115.0, 78.0], [118.0, 82.0]]) out = p.process((ts, val)) # 8 h window / 2 h step + 1 = 5 steps assert out["value"].shape == (5, 2) - assert out["time"].shape == (5,) + assert out["time"].shape == (5,) # Times should be [0, 2, 4, 6, 8] - expected_times = torch.tensor([0., 2., 4., 6., 8.]) + expected_times = torch.tensor([0.0, 2.0, 4.0, 6.0, 8.0]) assert torch.allclose(out["time"], expected_times) def test_temporal_timeseries_fit(): from pyhealth.processors import TemporalTimeseriesProcessor + p = TemporalTimeseriesProcessor() ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 1)] val = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) @@ -117,29 +146,35 @@ def test_temporal_timeseries_fit(): def test_temporal_timeseries_imputation(): from pyhealth.processors import TemporalTimeseriesProcessor + p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=1)) ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 2)] # gap at h=1 val = np.array([[10.0], [20.0]]) out = p.process((ts, val)) # 3 steps: h=0 → 10, h=1 → forward-filled to 10, h=2 → 20 assert out["value"].shape == (3, 1) - assert float(out["value"][1, 0]) == pytest.approx(10.0) + assert math.isclose(float(out["value"][1, 0]), 10.0, rel_tol=1e-6) # ── 4. collate_temporal ─────────────────────────────────────────────────────── + def test_collate_temporal_basic(): from pyhealth.datasets.collate import collate_temporal batch = [ { - "codes": {"value": torch.tensor([1, 2, 3], dtype=torch.long), - "time": torch.tensor([0., 1., 2.])}, + "codes": { + "value": torch.tensor([1, 2, 3], dtype=torch.long), + "time": torch.tensor([0.0, 1.0, 2.0]), + }, "label": torch.tensor(1), }, { - "codes": {"value": torch.tensor([4, 5, 3], dtype=torch.long), - "time": torch.tensor([0.5, 1.5, 2.5])}, + "codes": { + "value": torch.tensor([4, 5, 3], dtype=torch.long), + "time": torch.tensor([0.5, 1.5, 2.5]), + }, "label": torch.tensor(0), }, ] @@ -147,7 +182,7 @@ def test_collate_temporal_basic(): collated = collate_temporal(batch) assert collated["codes"]["value"].shape == (2, 3) - assert collated["codes"]["time"].shape == (2, 3) + assert collated["codes"]["time"].shape == (2, 3) assert collated["label"].shape == (2,) @@ -156,10 +191,18 @@ def test_collate_temporal_variable_length(): from pyhealth.datasets.collate import collate_temporal batch = [ - {"codes": {"value": torch.tensor([1, 2], dtype=torch.long), - "time": torch.tensor([0., 1.])}}, - {"codes": {"value": torch.tensor([3, 4, 5], dtype=torch.long), - "time": torch.tensor([0., 1., 2.])}}, + { + "codes": { + "value": torch.tensor([1, 2], dtype=torch.long), + "time": torch.tensor([0.0, 1.0]), + } + }, + { + "codes": { + "value": torch.tensor([3, 4, 5], dtype=torch.long), + "time": torch.tensor([0.0, 1.0, 2.0]), + } + }, ] collated = collate_temporal(batch) # Padded to length 3 @@ -168,24 +211,28 @@ def test_collate_temporal_variable_length(): # ── 5. SinusoidalTimeEmbedding ──────────────────────────────────────────────── + def test_sinusoidal_time_embedding_shape(): from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=64, max_hours=720.0) - t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]]) # (2, 3) + 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) def test_sinusoidal_different_times_differ(): from pyhealth.models.embedding import SinusoidalTimeEmbedding + emb = SinusoidalTimeEmbedding(dim=32) - t0 = emb(torch.tensor([0.0])) - t1 = emb(torch.tensor([24.0])) + t0 = emb(torch.tensor([0.0])) + t1 = emb(torch.tensor([24.0])) assert not torch.allclose(t0, t1) # ── 6. UnifiedMultimodalEmbeddingModel, code-only smoke test ───────────────── + def _make_code_processors_and_inputs(batch_size=2, seq_len=5): """Build a minimal dataset mock with a single CODE-modality field.""" from pyhealth.processors import StageNetProcessor @@ -199,7 +246,9 @@ def _make_code_processors_and_inputs(batch_size=2, seq_len=5): # Fake batch dict (as produced by collate_temporal) value = torch.randint(1, vocab_size, (batch_size, seq_len)) - time = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) + time = ( + torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1) + ) inputs = {"codes": {"value": value, "time": time}} return processors, inputs @@ -214,8 +263,8 @@ def test_unified_model_code_only(): out = model(inputs) assert "sequence" in out - assert "time" in out - assert "mask" in out + assert "time" in out + assert "mask" in out B, S, E = out["sequence"].shape assert B == 2 @@ -228,8 +277,15 @@ def test_unified_model_rejects_non_temporal(): from pyhealth.processors import SequenceProcessor bad_proc = SequenceProcessor() - with pytest.raises(TypeError, match="TemporalFeatureProcessor"): - UnifiedMultimodalEmbeddingModel(processors={"field": bad_proc}, embedding_dim=64) + try: + UnifiedMultimodalEmbeddingModel( + processors={"field": bad_proc}, + embedding_dim=64, + ) + except TypeError as exc: + assert "TemporalFeatureProcessor" in str(exc) + else: + raise AssertionError("Expected TypeError for non-temporal processor") def test_unified_model_gradient_flow(): @@ -239,7 +295,7 @@ def test_unified_model_gradient_flow(): processors, inputs = _make_code_processors_and_inputs() model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=32) - out = model(inputs) + out = model(inputs) loss = out["sequence"].mean() loss.backward() @@ -257,15 +313,13 @@ def test_unified_model_time_sort(): proc = StageNetProcessor() proc.fit(samples, "c") - model = UnifiedMultimodalEmbeddingModel( - processors={"c": proc}, embedding_dim=16 - ) + model = UnifiedMultimodalEmbeddingModel(processors={"c": proc}, embedding_dim=16) # Reverse-order times - value = torch.tensor([[2, 1]]) # (1, 2) - time = torch.tensor([[10.0, 0.0]]) # t=10 then t=0 → should sort to [0, 10] - out = model({"c": {"value": value, "time": time}}) - assert out["time"][0, 0].item() == pytest.approx(0.0) - assert out["time"][0, 1].item() == pytest.approx(10.0) + value = torch.tensor([[2, 1]]) # (1, 2) + time = torch.tensor([[10.0, 0.0]]) # t=10 then t=0 → should sort to [0, 10] + out = model({"c": {"value": value, "time": time}}) + assert math.isclose(out["time"][0, 0].item(), 0.0, rel_tol=1e-6) + assert math.isclose(out["time"][0, 1].item(), 10.0, rel_tol=1e-6) # ── 7. field_embeddings: reuse pre-built unimodal encoder ───────────────────── @@ -322,8 +376,8 @@ class _MockEmbedModel: assert isinstance(model.encoders["codes"], nn.Sequential) # Forward should produce embedding_dim=32 value = torch.randint(1, vocab_size, (2, 3)) - time = torch.arange(3, dtype=torch.float32).unsqueeze(0).expand(2, -1) - out = model({"codes": {"value": value, "time": time}}) + time = torch.arange(3, dtype=torch.float32).unsqueeze(0).expand(2, -1) + out = model({"codes": {"value": value, "time": time}}) assert out["sequence"].shape[-1] == 32 @@ -349,13 +403,90 @@ class _MockEmbedModel: field_embeddings={"codes": _MockEmbedModel()}, ) value = torch.randint(1, vocab_size, (3, 4)) - time = torch.arange(4, dtype=torch.float32).unsqueeze(0).expand(3, -1) - out = model({"codes": {"value": value, "time": time}}) + time = torch.arange(4, dtype=torch.float32).unsqueeze(0).expand(3, -1) + out = model({"codes": {"value": value, "time": time}}) assert out["sequence"].shape == (3, 4, 64) assert out["mask"].shape == (3, 4) +def test_unified_text_encoder_shared_by_tokenizer(): + """Token-based TEXT fields with the same tokenizer share one encoder.""" + from types import SimpleNamespace + import torch.nn as nn + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + from pyhealth.processors import ModalityType, TemporalFeatureProcessor + + class _DummyTemporalTextProcessor(TemporalFeatureProcessor): + def __init__(self, tokenizer_model: str): + self.tokenizer_model = tokenizer_model + + def process(self, value): + return value + + def modality(self): + return ModalityType.TEXT + + def value_dim(self): + return 0 + + def is_token(self): + return True + + def schema(self): + return ("value", "mask", "time") + + def dim(self): + return (2, 2, 1) + + def spatial(self): + return (False, False) + + class _DummyBert(nn.Module): + def __init__(self, hidden_size: int): + super().__init__() + self.config = SimpleNamespace(hidden_size=hidden_size) + + def forward(self, input_ids=None, attention_mask=None): + if input_ids is None: + raise ValueError("input_ids is required") + b, l = input_ids.shape + hidden = self.config.hidden_size + out = torch.zeros(b, l, hidden) + return SimpleNamespace(last_hidden_state=out) + + call_count = {"n": 0} + + def _fake_from_pretrained(_name): + call_count["n"] += 1 + return _DummyBert(hidden_size=48) + + with patch("transformers.AutoModel.from_pretrained", _fake_from_pretrained): + processors = { + "discharge_note_times": _DummyTemporalTextProcessor("bert-base-uncased"), + "radiology_note_times": _DummyTemporalTextProcessor("bert-base-uncased"), + } + + model = UnifiedMultimodalEmbeddingModel( + processors=processors, + embedding_dim=32, + ) + + # Both text fields reuse the same encoder instance. + assert ( + model.encoders["discharge_note_times"] is model.encoders["radiology_note_times"] + ) + assert call_count["n"] == 1 + + # Projections remain field-specific. + assert "discharge_note_times" in model.projections + assert "radiology_note_times" in model.projections + assert ( + model.projections["discharge_note_times"] + is not model.projections["radiology_note_times"] + ) + + # ── 8. Downstream models in unified mode ────────────────────────────────────── @@ -408,8 +539,8 @@ def test_transformer_unified_mode(): model = Transformer(dataset=dataset, embedding_dim=32, unified_embedding=unified) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out and "logit" in out out["loss"].backward() @@ -436,8 +567,8 @@ def test_ehrmamba_unified_mode(): ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out out["loss"].backward() @@ -465,8 +596,8 @@ def test_jamba_ehr_unified_mode(): ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out and "y_prob" in out out["loss"].backward() @@ -484,7 +615,9 @@ def test_mlp_unified_mode(): processors=dataset.input_processors, embedding_dim=32, ) - model = MLP(dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified) + model = MLP( + dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified + ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) @@ -507,7 +640,9 @@ def test_rnn_unified_mode(): processors=dataset.input_processors, embedding_dim=32, ) - model = RNN(dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified) + model = RNN( + dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified + ) loader = get_dataloader(dataset, batch_size=2, shuffle=False) batch = next(iter(loader)) @@ -560,6 +695,7 @@ def test_unified_per_field_backward_compat(): dataset = _make_stagenet_dataset() # Uses SequenceProcessor-style input_schema for per-field mode from pyhealth.datasets import create_sample_dataset + samples = [ {"patient_id": "p0", "visit_id": "v0", "codes": ["A", "B", "C"], "label": 1}, {"patient_id": "p1", "visit_id": "v1", "codes": ["D", "E"], "label": 0}, @@ -572,7 +708,21 @@ def test_unified_per_field_backward_compat(): ) model = Transformer(dataset=ds, embedding_dim=32) loader = get_dataloader(ds, batch_size=2, shuffle=False) - batch = next(iter(loader)) - out = model(**batch) + batch = next(iter(loader)) + out = model(**batch) assert "loss" in out out["loss"].backward() + + +def load_tests(loader, tests, pattern): + """Expose top-level test_ functions to unittest discovery.""" + suite = unittest.TestSuite() + namespace = globals() + for name in sorted(namespace): + if name.startswith("test_") and callable(namespace[name]): + suite.addTest(unittest.FunctionTestCase(namespace[name])) + return suite + + +if __name__ == "__main__": + unittest.main() From 77d098b698de030f7797ab2a1e5f6e8c733933f1 Mon Sep 17 00:00:00 2001 From: John Wu Date: Sun, 19 Apr 2026 15:38:38 -0500 Subject: [PATCH 07/30] attempts at memory optimization --- ...timodal_embedding_bottleneck_mimic4_cxr.py | 360 ++++++++++++++++++ .../multimodal_embedding_jamba_mimic4_cxr.py | 13 + .../multimodal_embedding_mamba_mimic4_cxr.py | 13 + .../multimodal_embedding_mlp_mimic4_cxr.py | 13 + .../multimodal_embedding_rnn_mimic4_cxr.py | 13 + ...imodal_embedding_transformer_mimic4_cxr.py | 13 + pyhealth/models/embedding/unified.py | 9 +- pyhealth/trainer.py | 67 +++- 8 files changed, 485 insertions(+), 16 deletions(-) create mode 100644 examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py diff --git a/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py new file mode 100644 index 000000000..e807848af --- /dev/null +++ b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py @@ -0,0 +1,360 @@ +"""Unified multimodal embedding + BottleneckTransformer runner. + +Runs EHR + notes + X-ray (MIMIC-IV + CXR) with ClinicalNotesICDLabsCXRMIMIC4. + +This script runs EHR + notes + X-ray (metadata/negbio) with the +ClinicalNotesICDLabsCXRMIMIC4 task. + +Default roots are set to shared PhysioNet mounts: +- 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 + +Quick start: + python examples/mortality_prediction/ + multimodal_embedding_bottleneck_mimic4_cxr.py \ + --quick-test + +Smoke test (single forward + inference, no train): + python examples/mortality_prediction/ + multimodal_embedding_bottleneck_mimic4_cxr.py \ + --smoke-forward +""" + +from __future__ import annotations + +import argparse +import os +import time +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +from pyhealth.models import UnifiedMultimodalEmbeddingModel +from pyhealth.models.bottleneck_transformer import BottleneckTransformer +from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsCXRMIMIC4 +from pyhealth.trainer import Trainer + + +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_run_output_path(args: argparse.Namespace) -> str: + device_tag = args.device.replace(":", "") + lr_tag = f"{args.lr:g}".replace(".", "p") + dropout_tag = f"{args.dropout:g}".replace(".", "p") + run_tag = ( + f"cxr{args.cxr_variant}" + f"_emb{args.embedding_dim}" + f"_layers{args.num_layers}" + f"_heads{args.heads}" + f"_bottleneck{args.bottlenecks_n}" + f"_fuse{args.fusion_startidx}" + f"_drop{dropout_tag}" + f"_win{args.observation_window_hours}" + f"_ep{args.epochs}" + f"_bs{args.batch_size}" + f"_lr{lr_tag}" + f"_{device_tag}" + f"_nw{args.num_workers}" + f"_seed{args.seed}" + f"_dev{int(args.dev)}" + f"_quick{int(args.quick_test)}" + f"_smoke{int(args.smoke_forward)}" + ) + return os.path.join( + os.getcwd(), + "output", + "multimodal_embedding_bottleneck_mimic4_cxr", + run_tag, + ) + + +def run(args: argparse.Namespace) -> Tuple[int, int]: + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + cuda_device_index = None + if args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + print("Using dataset roots:") + print(f" ehr_root: {args.ehr_root}") + print(f" note_root: {args.note_root}") + print(f" cxr_root: {args.cxr_root}") + print(f" cxr_variant: {args.cxr_variant}") + print(f" cache_dir: {args.cache_dir}") + print(f" num_workers: {args.num_workers}") + + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root, + cxr_root=args.cxr_root, + cxr_variant=args.cxr_variant, + ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"], + note_tables=["discharge", "radiology"], + cxr_tables=["metadata", "negbio", "chexpert", "split"], + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = ClinicalNotesICDLabsCXRMIMIC4( + window_hours=args.observation_window_hours, + ) + 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 use " + "--quick-test first." + ) + + print(f"Task sample count: {len(sample_dataset)}") + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print( + f" - {key}: {type(processor).__name__}, " + f"schema={processor.schema()}" + ) + + train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) + + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + model = 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, + ) + print( + f"BottleneckTransformer bottlenecks_n={args.bottlenecks_n}, " + f"fusion_startidx={args.fusion_startidx}" + ) + + 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 + ) + + print( + "Split sizes: " + f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}" + ) + + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print( + f" tuple[{i}] type={type(elem).__name__} " + f"shape={shape}" + ) + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print( + f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}" + ) + + trainer = Trainer( + model=model, + metrics=["pr_auc", "roc_auc"], + device=args.device, + enable_logging=True, + output_path=_build_run_output_path(args), + ) + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + trainer.train( + train_dataloader=train_loader, + val_dataloader=val_loader, + epochs=args.epochs, + optimizer_params={"lr": args.lr}, + monitor=None, + load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, + ) + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + else: + peak_train_vram_mb = None + train_runtime_sec = time.perf_counter() - train_start + else: + peak_train_vram_mb = None + train_runtime_sec = None + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + scores = trainer.evaluate(inference_loader) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + print("Evaluation metrics:") + for metric_name, metric_value in scores.items(): + print(f" {metric_name}: {metric_value:.4f}") + + return len(patient_ids), y_true.shape[0] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Run unified multimodal embedding + BottleneckTransformer " + "on MIMIC-IV mortality with CXR." + ) + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", + type=str, + default="/shared/eng/pyhealth", + ) + + parser.add_argument("--embedding-dim", type=int, default=64) + parser.add_argument("--num-layers", type=int, default=2) + parser.add_argument("--heads", type=int, default=1) + parser.add_argument("--dropout", type=float, default=0.1) + parser.add_argument("--bottlenecks-n", type=int, default=4) + parser.add_argument("--fusion-startidx", type=int, default=1) + + parser.add_argument("--observation-window-hours", type=int, default=24) + + parser.add_argument("--epochs", type=int, default=1) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--device", type=str, default="cuda:1") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--seed", type=int, default=42) + + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--smoke-forward", action="store_true") + parser.add_argument( + "--condor", + action="store_true", + help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", + ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) + args = parser.parse_args() + + if args.condor: + args.device = "cuda" + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return args + + +if __name__ == "__main__": + cli_args = parse_args() + num_patient_ids, num_rows = run(cli_args) + print( + f"Inference completed (patient_ids={num_patient_ids}, " + f"rows={num_rows})." + ) diff --git a/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py index 3fc268825..4746eb342 100644 --- a/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py @@ -211,6 +211,10 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: optimizer_params={"lr": args.lr}, monitor=None, load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, ) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -308,6 +312,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) args = parser.parse_args() if args.condor: diff --git a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py index 03a42cfe8..4896b037d 100644 --- a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py @@ -204,6 +204,10 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: optimizer_params={"lr": args.lr}, monitor=None, load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, ) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -299,6 +303,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) args = parser.parse_args() if args.condor: diff --git a/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py index 9f37d7153..ac09f2c61 100644 --- a/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py @@ -203,6 +203,10 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: optimizer_params={"lr": args.lr}, monitor=None, load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, ) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -303,6 +307,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) args = parser.parse_args() if args.condor: diff --git a/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py index e58ba7243..4aeeed9ce 100644 --- a/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py @@ -204,6 +204,10 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: optimizer_params={"lr": args.lr}, monitor=None, load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, ) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -304,6 +308,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) args = parser.parse_args() if args.condor: diff --git a/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py index 6a1727571..69942738c 100644 --- a/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py +++ b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py @@ -202,6 +202,10 @@ def run(args: argparse.Namespace) -> Tuple[int, int]: optimizer_params={"lr": args.lr}, monitor=None, load_best_model_at_last=False, + max_grad_norm=args.max_grad_norm, + accumulation_steps=args.grad_accum, + use_amp=args.amp, + amp_dtype=args.amp_dtype, ) if cuda_device_index is not None: torch.cuda.synchronize(cuda_device_index) @@ -296,6 +300,15 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Set device to 'cuda' (no index) for HTCondor GPU jobs.", ) + parser.add_argument("--grad-accum", type=int, default=1) + parser.add_argument("--max-grad-norm", type=float, default=None) + parser.add_argument("--amp", action="store_true") + parser.add_argument( + "--amp-dtype", + type=str, + default="bf16", + choices=["bf16", "fp16"], + ) args = parser.parse_args() if args.condor: diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 2e5e9560d..d70992e03 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -215,6 +215,7 @@ def __init__( 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): @@ -345,12 +346,15 @@ def _set_projection( shared_field = self._shared_text_field_by_model.get(tokenizer_model) if shared_field is not None: + # Second+ field with same tokenizer: reuse existing encoder, do NOT + # register under a new key (avoids duplicate parameter registration). + self._text_canonical[field_name] = shared_field shared_encoder = self.encoders[shared_field] else: shared_encoder = AutoModel.from_pretrained(tokenizer_model) self._shared_text_field_by_model[tokenizer_model] = field_name + self.encoders[field_name] = shared_encoder - self.encoders[field_name] = shared_encoder hidden = shared_encoder.config.hidden_size _set_projection(hidden) else: @@ -455,7 +459,8 @@ def forward( time = torch.zeros(value.shape[:2], device=value.device) modality = self.modality_types[field_name] - encoder = self.encoders[field_name] + encoder_key = self._text_canonical.get(field_name, field_name) + encoder = self.encoders[encoder_key] # ── Encode ──────────────────────────────────────────────────── if modality == ModalityType.CODE: diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index 78d15c479..34c32dc61 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -137,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. @@ -156,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}") @@ -172,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()) @@ -208,7 +228,7 @@ def train( 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}", smoothing=0.05, @@ -218,20 +238,39 @@ def train( 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 - ) - # update - optimizer.step() - optimizer.zero_grad() - training_loss.append(loss.item()) - global_step += 1 + 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 max_grad_norm is not None: + if scaler is not None: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), max_grad_norm + ) + if scaler is not None: + scaler.step(optimizer) + scaler.update() + else: + optimizer.step() + optimizer.zero_grad() + global_step += 1 epoch_time = time.perf_counter() - epoch_start vram = _vram_stats(self.device) From 4cd3def02055244d3cb326b5b1035ff46ba18086 Mon Sep 17 00:00:00 2001 From: John Wu Date: Tue, 21 Apr 2026 09:53:56 -0500 Subject: [PATCH 08/30] add stats file to share --- .../multimodal_dataset_stats.py | 391 ++++++++++++++++++ 1 file changed, 391 insertions(+) create mode 100644 examples/mortality_prediction/multimodal_dataset_stats.py diff --git a/examples/mortality_prediction/multimodal_dataset_stats.py b/examples/mortality_prediction/multimodal_dataset_stats.py new file mode 100644 index 000000000..ceaaf7765 --- /dev/null +++ b/examples/mortality_prediction/multimodal_dataset_stats.py @@ -0,0 +1,391 @@ +"""Standalone multimodal dataset statistics auditor. + +Loads a MIMIC-IV multimodal task dataset and reports per-modality missingness +rates and token/element counts — no model, no GPU, no trainer required. + +Works directly on the processed SampleDataset. Processed schema (single sample, +no batch dim): + discharge_note_times / radiology_note_times: + (input_ids, attn_mask, token_type_ids, time, type_tag) + input_ids shape: (N_notes, 128) attn_mask shape: (N_notes, 128) + attn_mask.sum() gives real (non-padding) tokens seen by the model. + Note: each note is independently truncated to 128 wordpieces — no chunking. + icd_codes: + (time, value) value shape: (N_visits, vocab_size) [multi-hot] + labs_mask: + (time, value) value shape: (N_timesteps, 10) [bool/float] + cxr_image_times: + (image, time, paths) image shape: (N_images, 3, 224, 224) + +Usage: + python examples/mortality_prediction/multimodal_dataset_stats.py \\ + --dev --quick-test + + python examples/mortality_prediction/multimodal_dataset_stats.py \\ + --task ClinicalNotesICDLabsCXRMIMIC4 --output-csv /tmp/stats.csv +""" + +from __future__ import annotations + +import argparse +import csv +from typing import Any, Dict, List, Tuple + +import numpy as np + +from pyhealth.datasets import MIMIC4Dataset +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsCXRMIMIC4, + ClinicalNotesICDLabsMIMIC4, + ClinicalNotesMIMIC4, + ICDLabsMIMIC4, +) + +TASK_MAP = { + "ClinicalNotesICDLabsCXRMIMIC4": ClinicalNotesICDLabsCXRMIMIC4, + "ClinicalNotesICDLabsMIMIC4": ClinicalNotesICDLabsMIMIC4, + "ClinicalNotesMIMIC4": ClinicalNotesMIMIC4, + "ICDLabsMIMIC4": ICDLabsMIMIC4, +} + +# bert-base-uncased encodes "[MISSING_TEXT]" to ~7 tokens with padding. +# Real notes are always longer. Used to detect the missingness sentinel. +_MISSING_NOTE_TOKEN_THRESHOLD = 15 + + +def _note_stats(tup: tuple, note_max_len: int) -> Tuple[bool, int, int]: + """Stats from a processed note tuple. + + Schema: (input_ids, attn_mask, token_type_ids, time, type_tag) + input_ids shape: (N_notes, note_max_len) + attn_mask shape: (N_notes, note_max_len) + + Returns (is_missing, n_notes, seen_tokens_total). + seen_tokens = attn_mask.sum() — real non-padding tokens, already capped + at note_max_len by the processor's truncation. + """ + input_ids, attn_mask = tup[0], tup[1] + n_notes = input_ids.shape[0] + seen = int(attn_mask.sum().item()) + is_missing = (n_notes == 1) and (seen < _MISSING_NOTE_TOKEN_THRESHOLD) + if is_missing: + return True, 0, 0 + return False, n_notes, seen + + +def _icd_stats(tup: tuple) -> Tuple[bool, int, int]: + """Stats from a processed icd_codes tuple. + + Schema: (time, value) value shape: (N_visits, vocab_size) [multi-hot] + + Returns (is_missing, n_visits, n_code_activations). + """ + value = tup[1] + n_visits = value.shape[0] + n_act = int(value.sum().item()) + return n_act == 0, n_visits, n_act + + +def _labs_stats(tup: tuple) -> Tuple[bool, int]: + """Stats from a processed labs_mask tuple. + + Schema: (time, value) value shape: (N_timesteps, 10) [bool/float] + + Returns (is_missing, n_observed). + """ + value = tup[1] + n_obs = int(value.sum().item()) + return n_obs == 0, n_obs + + +def _cxr_stats(tup: tuple, patch_count: int) -> Tuple[bool, int, int]: + """Stats from a processed cxr_image_times tuple. + + Schema: (image, time, paths) image shape: (N_images, 3, H, W) + + Returns (is_missing, n_images, cxr_tokens). + A zero-valued image tensor indicates the missing-image sentinel. + """ + image = tup[0] + n_images = image.shape[0] + if n_images == 0 or float(image.sum()) < 1e-8: + return True, 0, 0 + return False, n_images, n_images * patch_count + + +def audit_sample( + sample: Dict[str, Any], + note_max_len: int, + cxr_patch_count: int, +) -> Dict[str, Any]: + """Audit one processed SampleDataset sample dict.""" + row: Dict[str, Any] = {} + + for note_key in ("discharge_note_times", "radiology_note_times"): + if note_key not in sample: + continue + miss, n_notes, seen_tok = _note_stats( + sample[note_key], note_max_len + ) + row[f"{note_key}__missing"] = int(miss) + row[f"{note_key}__n_notes"] = n_notes + row[f"{note_key}__seen_tokens"] = seen_tok + + if "icd_codes" in sample: + miss, n_visits, n_act = _icd_stats(sample["icd_codes"]) + row["icd_codes__missing"] = int(miss) + row["icd_codes__n_visits"] = n_visits + row["icd_codes__n_activations"] = n_act + + if "labs_mask" in sample: + miss, n_obs = _labs_stats(sample["labs_mask"]) + row["labs__missing"] = int(miss) + row["labs__n_obs"] = n_obs + + if "cxr_image_times" in sample: + miss, n_img, cxr_tok = _cxr_stats( + sample["cxr_image_times"], cxr_patch_count + ) + row["cxr__missing"] = int(miss) + row["cxr__n_images"] = n_img + row["cxr__tokens"] = cxr_tok + + note_tok = sum( + row.get(f"{k}__seen_tokens", 0) + for k in ("discharge_note_times", "radiology_note_times") + ) + row["total_tokens"] = ( + note_tok + + row.get("icd_codes__n_activations", 0) + + row.get("labs__n_obs", 0) + + row.get("cxr__tokens", 0) + ) + return row + + +def _stats(arr: np.ndarray) -> str: + if len(arr) == 0: + return "N/A" + return ( + f"mean={arr.mean():.1f} median={np.median(arr):.0f}" + f" p90={np.percentile(arr, 90):.0f} max={arr.max():.0f}" + ) + + +def print_report(rows: List[Dict], args: argparse.Namespace) -> None: + n = len(rows) + print() + print(f"Task: {args.task} Samples: {n:,} dev={args.dev}") + print() + + def _arr(key: str) -> np.ndarray: + return np.array([r[key] for r in rows if key in r], dtype=float) + + hdr = ( + f"{'Modality':<28} {'missing%':>10}" + f" {'mean':>8} {'median':>8} {'p90':>8} {'max':>8}" + ) + print(hdr) + print("-" * len(hdr)) + + modality_specs = [] + for note_key, label in [ + ("discharge_note_times", "discharge notes"), + ("radiology_note_times", "radiology notes"), + ]: + miss = _arr(f"{note_key}__missing") + if len(miss): + modality_specs.append( + (label, miss, _arr(f"{note_key}__n_notes")) + ) + + if rows and "icd_codes__missing" in rows[0]: + modality_specs.append(( + "icd codes", + _arr("icd_codes__missing"), + _arr("icd_codes__n_activations"), + )) + if rows and "labs__missing" in rows[0]: + modality_specs.append(( + "labs (observations)", + _arr("labs__missing"), + _arr("labs__n_obs"), + )) + if rows and "cxr__missing" in rows[0]: + modality_specs.append(( + "cxr images", + _arr("cxr__missing"), + _arr("cxr__n_images"), + )) + + for label, miss, counts in modality_specs: + miss_pct = f"{miss.mean() * 100:.1f}%" + print( + f"{label:<28} {miss_pct:>10} " + f"{counts.mean():>8.1f} {np.median(counts):>8.0f} " + f"{np.percentile(counts, 90):>8.0f} {counts.max():>8.0f}" + ) + + print() + print(f"Note seen-tokens (cap@{args.note_max_len}, from attn_mask):") + for note_key, label in [ + ("discharge_note_times", " discharge"), + ("radiology_note_times", " radiology"), + ]: + seen = _arr(f"{note_key}__seen_tokens") + if len(seen) == 0: + continue + pct_full = (seen >= args.note_max_len).mean() * 100 + print( + f"{label}: {_stats(seen)}" + f" % hitting cap={pct_full:.0f}%" + ) + + if rows and "cxr__tokens" in rows[0]: + cxr_tok = _arr("cxr__tokens") + print() + print(f"CXR tokens (@{args.cxr_patch_count} patches/image):") + print(f" {_stats(cxr_tok)}") + + total = _arr("total_tokens") + print() + print("Aggregate tokens/sample:") + print(f" {_stats(total)}") + print() + + +def run(args: argparse.Namespace) -> None: + task_cls = TASK_MAP[args.task] + needs_notes = args.task in ( + "ClinicalNotesMIMIC4", + "ClinicalNotesICDLabsMIMIC4", + "ClinicalNotesICDLabsCXRMIMIC4", + ) + needs_cxr = args.task == "ClinicalNotesICDLabsCXRMIMIC4" + needs_icd = args.task in ( + "ClinicalNotesICDLabsMIMIC4", + "ICDLabsMIMIC4", + "ClinicalNotesICDLabsCXRMIMIC4", + ) + + ehr_tables = ( + ["diagnoses_icd", "procedures_icd", "labevents"] + if needs_icd + else [] + ) + note_tables = ["discharge", "radiology"] if needs_notes else [] + cxr_tables = ( + ["metadata", "negbio", "chexpert", "split"] if needs_cxr else [] + ) + + print("Loading MIMIC4Dataset ...") + base_dataset = MIMIC4Dataset( + ehr_root=args.ehr_root, + note_root=args.note_root if needs_notes else None, + cxr_root=args.cxr_root if needs_cxr else None, + cxr_variant=args.cxr_variant, + ehr_tables=ehr_tables, + note_tables=note_tables, + cxr_tables=cxr_tables, + cache_dir=args.cache_dir, + dev=args.dev, + num_workers=args.num_workers, + ) + + task = task_cls(window_hours=args.observation_window_hours) + print("Running set_task ...") + sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers) + total = len(sample_dataset) + print(f"Total samples: {total:,}") + + limit = ( + min(args.sample_limit, total) + if args.sample_limit and args.sample_limit > 0 + else total + ) + print(f"Auditing {limit:,} samples ...") + + rows: List[Dict] = [] + for i in range(limit): + if i % 5000 == 0 and i > 0: + print(f" {i:,} / {limit:,} ...") + rows.append( + audit_sample( + sample_dataset[i], + args.note_max_len, + args.cxr_patch_count, + ) + ) + + if not rows: + print("No samples. Check roots/tables/task combination.") + return + + print_report(rows, args) + + if args.output_csv: + all_keys = list(rows[0].keys()) + with open(args.output_csv, "w", newline="") as f: + writer = csv.DictWriter(f, fieldnames=all_keys) + writer.writeheader() + writer.writerows(rows) + print(f"Per-sample CSV written to: {args.output_csv}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Audit modality missingness and token counts " + "for MIMIC-IV multimodal tasks." + ) + ) + parser.add_argument( + "--ehr-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimiciv/2.2", + ) + parser.add_argument( + "--note-root", + type=str, + default="/shared/rsaas/physionet.org/files/mimic-note", + ) + parser.add_argument( + "--cxr-root", + type=str, + default="/shared/rsaas/physionet.org/files/MIMIC-CXR", + ) + parser.add_argument( + "--cxr-variant", + type=str, + default="sunlab", + choices=["default", "sunlab"], + ) + parser.add_argument( + "--cache-dir", type=str, default="/shared/eng/pyhealth" + ) + parser.add_argument( + "--task", + type=str, + default="ClinicalNotesICDLabsCXRMIMIC4", + choices=list(TASK_MAP.keys()), + ) + parser.add_argument("--observation-window-hours", type=int, default=24) + parser.add_argument("--dev", action="store_true") + parser.add_argument("--quick-test", action="store_true") + parser.add_argument("--num-workers", type=int, default=16) + parser.add_argument("--note-max-len", type=int, default=128) + parser.add_argument("--cxr-patch-count", type=int, default=196) + parser.add_argument("--sample-limit", type=int, default=None) + parser.add_argument("--output-csv", type=str, default=None) + + args = parser.parse_args() + if args.quick_test: + args.dev = True + if args.sample_limit is None: + args.sample_limit = 50 + return args + + +if __name__ == "__main__": + args = parse_args() + run(args) From ce630744cccc734f18af5d1b88359d01cca7a7ab Mon Sep 17 00:00:00 2001 From: William Pang Date: Thu, 23 Apr 2026 08:02:55 -0700 Subject: [PATCH 09/30] Add ability to specify task flags --- scripts/slurm/run_table2.sh | 2 +- scripts/slurm/submit_rnn.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/slurm/run_table2.sh b/scripts/slurm/run_table2.sh index c0f294733..dcd66769b 100755 --- a/scripts/slurm/run_table2.sh +++ b/scripts/slurm/run_table2.sh @@ -69,7 +69,7 @@ COMMON=( --ehr-root "${EHR_ROOT}" --note-root "${NOTE_ROOT}" --cache-dir "${CACHE_DIR}" - --task clinical_notes_icd_labs + --task "${TABLE2_TASK:-clinical_notes_icd_labs}" --model "${MODEL}" --embedding-dim 128 --hidden-dim 128 diff --git a/scripts/slurm/submit_rnn.sh b/scripts/slurm/submit_rnn.sh index 900dbe51b..847f8210d 100755 --- a/scripts/slurm/submit_rnn.sh +++ b/scripts/slurm/submit_rnn.sh @@ -21,7 +21,7 @@ for seed in "${SEEDS[@]}"; do --mem=32G --gres=gpu:1 --time=6:00:00 \ --output="logs/slurm/table2_rnn_seed${seed}_%j.out" \ --error="logs/slurm/table2_rnn_seed${seed}_%j.err" \ - --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=rnn,SEED="${seed}",TABLE2_BS_RNN=16 \ + --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=rnn,SEED="${seed}",TABLE2_BS_RNN=16,TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}" \ scripts/slurm/run_table2.sh | awk '{print $NF}') echo " Submitted rnn seed=${seed} → ${job}" done From 41a15e5a2b04af462e24e0c067b43d264fa55938 Mon Sep 17 00:00:00 2001 From: Joshua Chen Date: Thu, 23 Apr 2026 20:07:06 -0500 Subject: [PATCH 10/30] Add CXR experiment script and fix negative observation window --- .../unified_embedding_e2e_mimic4_cxr.py | 552 ++++++++++++++++++ pyhealth/tasks/multimodal_mimic4.py | 4 + 2 files changed, 556 insertions(+) create mode 100644 examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py new file mode 100644 index 000000000..95259915f --- /dev/null +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py @@ -0,0 +1,552 @@ +"""End-to-end protocol runner for Unified Embedding on MIMIC-IV + CXR. + +Trains and evaluates a unified-embedding model (MLP / RNN / Transformer / +BottleneckTransformer / EHRMamba / JambaEHR) on MIMIC-IV mortality using +all four modalities: clinical notes, ICD codes, lab values, and chest X-rays. + +This script is the CXR-extended version of unified_embedding_e2e_mimic4.py. +It adds --cxr-root, --cxr-variant, and the clinical_notes_icd_labs_cxr task +on top of the existing non-CXR tasks. + +VRAM note: CXR embeddings push peak VRAM to ~40 GB. Use a small batch size +(2-4) and target the largest available GPU (A6000 / H100). + +Example +------- + # Quick sanity check (--dev + 1 epoch + small batch): + python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ + --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 \\ + --task clinical_notes_icd_labs_cxr \\ + --model mlp --quick-test + + # Smoke test (single forward + inference, no training): + python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ + --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 \\ + --task clinical_notes_icd_labs_cxr \\ + --model mlp --smoke-forward --dev + + # Full training with Transformer + CXR: + python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ + --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 \\ + --task clinical_notes_icd_labs_cxr \\ + --model transformer --heads 4 --num-layers 2 \\ + --epochs 10 --batch-size 4 --device cuda:0 + + # JambaEHR + CXR: + python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ + --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 \\ + --task clinical_notes_icd_labs_cxr \\ + --model jambaehr --embedding-dim 128 \\ + --jamba-transformer-layers 2 --jamba-mamba-layers 6 +""" + +from __future__ import annotations + +import argparse +import csv +import time +from pathlib import Path +from typing import Any, Tuple + +import numpy as np +import torch + +from pyhealth.datasets import ( + MIMIC4Dataset, + get_dataloader, + split_by_patient, + split_by_sample, +) +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.tasks import MortalityPredictionStageNetMIMIC4 +from pyhealth.tasks.multimodal_mimic4 import ( + ClinicalNotesICDLabsMIMIC4, + ClinicalNotesICDLabsCXRMIMIC4, +) +from pyhealth.trainer import Trainer + + +# --------------------------------------------------------------------------- +# 1. Dataset construction +# --------------------------------------------------------------------------- + +def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: + """Load the MIMIC-IV base dataset with the appropriate tables. + + For non-CXR tasks this behaves identically to Rian's runner. + For the CXR task it additionally passes cxr_root / cxr_variant / + cxr_tables so the dataset loads chest X-ray metadata. + """ + ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] + + # Notes are needed for any task that includes clinical text. + note_tables = None + if args.task in ("clinical_notes_icd_labs", "clinical_notes_icd_labs_cxr"): + if not args.note_root: + raise ValueError( + f"--task {args.task} requires --note-root." + ) + note_tables = ["discharge", "radiology"] + + # CXR tables are only loaded for the CXR task. + cxr_kwargs = {} + if args.task == "clinical_notes_icd_labs_cxr": + if not args.cxr_root: + raise ValueError( + "--task clinical_notes_icd_labs_cxr requires --cxr-root." + ) + 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, + num_workers=args.num_workers, + **cxr_kwargs, + ) + + +# --------------------------------------------------------------------------- +# 2. Task construction +# --------------------------------------------------------------------------- + +def _build_task(args: argparse.Namespace): + """Return the task object matching --task.""" + if args.task == "stagenet": + return MortalityPredictionStageNetMIMIC4() + if args.task == "clinical_notes_icd_labs": + return ClinicalNotesICDLabsMIMIC4(window_hours=args.observation_window_hours) + if args.task == "clinical_notes_icd_labs_cxr": + return ClinicalNotesICDLabsCXRMIMIC4(window_hours=args.observation_window_hours) + raise ValueError(f"Unknown task: {args.task}") + + +# --------------------------------------------------------------------------- +# 3. Train / val / test split +# --------------------------------------------------------------------------- + +def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: + """Split by patient first; fall back to split by sample if empty.""" + 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 + + +# --------------------------------------------------------------------------- +# 4. Model construction +# --------------------------------------------------------------------------- + +def _build_model(args: argparse.Namespace, sample_dataset: Any): + """Instantiate UnifiedMultimodalEmbeddingModel + the chosen backbone.""" + unified = UnifiedMultimodalEmbeddingModel( + processors=sample_dataset.input_processors, + embedding_dim=args.embedding_dim, + ) + + 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, + 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}") + + +# --------------------------------------------------------------------------- +# 5. Prediction CSV writer +# --------------------------------------------------------------------------- + +def _write_predictions( + output_csv: Path, + patient_ids: list[str], + y_true: np.ndarray, + y_prob: np.ndarray, +) -> None: + """Write per-sample predictions to a CSV file.""" + 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), + } + ) + + +# --------------------------------------------------------------------------- +# 6. Main training + evaluation loop +# --------------------------------------------------------------------------- + +def run(args: argparse.Namespace) -> Path: + """Execute the full pipeline: load → task → split → train → evaluate.""" + torch.manual_seed(args.seed) + np.random.seed(args.seed) + total_start = time.perf_counter() + + # Resolve CUDA device index for VRAM tracking. + cuda_device_index = None + if args.device and args.device.startswith("cuda"): + device_index = torch.device(args.device).index + cuda_device_index = 0 if device_index is None else device_index + + 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." + ) + + print(f"Task sample count: {len(sample_dataset)}") + + # Print processor schemas so mismatches are caught early. + print("Input processor schemas:") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + if processor is None: + print(f" - {key}: ") + continue + print(f" - {key}: {type(processor).__name__}, " + f"schema={processor.schema()}") + + 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 + ) + + print(f"Split sizes: train={len(train_ds)}, val={len(val_ds)}, " + f"test={len(test_ds)}") + + # Debug batch diagnostics: print types, shapes, and schema mappings + # for the first training batch so processor issues surface immediately. + debug_batch = next(iter(train_loader)) + print("Batch field diagnostics (train batch 0):") + for key in sample_dataset.input_schema.keys(): + processor = sample_dataset.input_processors.get(key) + feature = debug_batch.get(key) + schema = processor.schema() if processor is not None else () + print(f" - {key}: type={type(feature).__name__}, schema={schema}") + + if isinstance(feature, tuple): + for i, elem in enumerate(feature): + shape = getattr(elem, "shape", None) + print(f" tuple[{i}] type={type(elem).__name__} " + f"shape={shape}") + + if processor is not None and isinstance(feature, tuple): + for field_name in ("value", "time", "mask"): + if field_name in schema: + idx = schema.index(field_name) + if idx < len(feature): + selected = feature[idx] + shape = getattr(selected, "shape", None) + print(f" schema['{field_name}'] -> tuple[{idx}] " + f"type={type(selected).__name__} shape={shape}") + + exp_name = f"{args.model}_seed{args.seed}" + output_dir = Path(args.output_dir) + + 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, + ) + + # Model-specific optimizer defaults (from Rian's runner). + 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-3 + if args.adam_eps is not None: + optimizer_params["eps"] = args.adam_eps + + optimizer_params["lr"] = effective_lr + + # --smoke-forward: skip training, only run inference to verify the + # pipeline works end-to-end without waiting for a full epoch. + peak_train_vram_mb = None + train_runtime_sec = None + + if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: + if cuda_device_index is not None: + torch.cuda.reset_peak_memory_stats(cuda_device_index) + torch.cuda.synchronize(cuda_device_index) + + train_start = time.perf_counter() + 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, + ) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) + peak_train_vram_mb = peak_train_bytes / (1024**2) + + train_runtime_sec = time.perf_counter() - train_start + + inference_loader = test_loader or val_loader or train_loader + y_true, y_prob, _, patient_ids = trainer.inference( + inference_loader, return_patient_ids=True + ) + + if cuda_device_index is not None: + torch.cuda.synchronize(cuda_device_index) + + total_runtime_sec = time.perf_counter() - total_start + + # Benchmark summary — matches the Mamba CXR script's output format. + print("Benchmark summary:") + print(f" total_runtime_sec: {total_runtime_sec:.2f}") + if train_runtime_sec is None: + print(" training_runtime_sec: N/A (training skipped)") + print(" peak_train_vram_mb: N/A (training skipped)") + else: + print(f" training_runtime_sec: {train_runtime_sec:.2f}") + if peak_train_vram_mb is None: + print(" peak_train_vram_mb: N/A (non-CUDA device)") + else: + print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") + + output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" + _write_predictions(output_csv, patient_ids, y_true, y_prob) + return output_csv + + +# --------------------------------------------------------------------------- +# 7. CLI argument parser +# --------------------------------------------------------------------------- + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run E2E unified embedding on MIMIC-IV (+ CXR) with " + "any of six backbone models." + ) + + # --- Data paths --- + 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, + help="Root directory for MIMIC-CXR images. Required for " + "clinical_notes_icd_labs_cxr task.", + ) + parser.add_argument( + "--cxr-variant", type=str, default="sunlab", + choices=["default", "sunlab"], + help="CXR directory layout variant.", + ) + parser.add_argument("--cache-dir", type=str, default=None) + parser.add_argument("--output-dir", type=str, default="./output/unified_e2e_cxr") + + # --- Task and model selection --- + parser.add_argument( + "--task", type=str, default="clinical_notes_icd_labs_cxr", + choices=[ + "stagenet", + "clinical_notes_icd_labs", + "clinical_notes_icd_labs_cxr", + ], + ) + parser.add_argument( + "--model", type=str, default="rnn", + choices=[ + "mlp", "rnn", "transformer", "bottleneck_transformer", + "ehrmamba", "jambaehr", + ], + ) + + # --- Shared embedding / training hyperparameters --- + 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=4, + help="Batch size. Default is 4 (CXR is VRAM-heavy, ~40 GB peak).", + ) + parser.add_argument( + "--lr", type=float, default=None, + help="Learning rate. Default: 1e-3 (1e-4 for bottleneck_transformer).", + ) + parser.add_argument( + "--adam-eps", type=float, default=None, + help="Adam epsilon. Default: 1e-8 (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("--num-workers", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--dev", action="store_true") + parser.add_argument( + "--quick-test", action="store_true", + help="Force --dev, 1 epoch, small batch for a quick sanity check.", + ) + parser.add_argument( + "--smoke-forward", action="store_true", + help="Skip training; only run a single forward pass + inference " + "to verify the pipeline end-to-end.", + ) + + # --- Task-specific --- + parser.add_argument("--observation-window-hours", type=int, default=24) + + # --- RNN-specific --- + 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") + + # --- Transformer / BottleneckTransformer --- + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--num-layers", type=int, default=2) + + # --- BottleneckTransformer-specific --- + parser.add_argument("--bottlenecks-n", type=int, default=4) + parser.add_argument("--fusion-startidx", type=int, default=1) + + # --- Training stability --- + parser.add_argument( + "--max-grad-norm", type=float, default=None, + help="Gradient clip norm. Default: None (0.5 for bottleneck_transformer).", + ) + + # --- Mamba / JambaEHR-specific --- + parser.add_argument("--mamba-state-size", type=int, default=16) + parser.add_argument("--mamba-conv-kernel", type=int, default=4) + parser.add_argument("--jamba-transformer-layers", type=int, default=2) + parser.add_argument("--jamba-mamba-layers", type=int, default=6) + + args = parser.parse_args() + + if args.quick_test: + args.dev = True + args.epochs = 1 + args.batch_size = min(args.batch_size, 4) + + return 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/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index f4f09744f..efc6787c1 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -812,6 +812,10 @@ 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. + 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" From f6eabc8e97f6ca5d7d274b314ea3f63a4be79aff Mon Sep 17 00:00:00 2001 From: Joshua Chen Date: Thu, 23 Apr 2026 20:40:59 -0500 Subject: [PATCH 11/30] tidy up --- .../unified_embedding_e2e_mimic4_cxr.py | 280 +++--------------- 1 file changed, 37 insertions(+), 243 deletions(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py index 95259915f..a4e0f7277 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py @@ -2,59 +2,25 @@ Trains and evaluates a unified-embedding model (MLP / RNN / Transformer / BottleneckTransformer / EHRMamba / JambaEHR) on MIMIC-IV mortality using -all four modalities: clinical notes, ICD codes, lab values, and chest X-rays. - -This script is the CXR-extended version of unified_embedding_e2e_mimic4.py. -It adds --cxr-root, --cxr-variant, and the clinical_notes_icd_labs_cxr task -on top of the existing non-CXR tasks. - -VRAM note: CXR embeddings push peak VRAM to ~40 GB. Use a small batch size -(2-4) and target the largest available GPU (A6000 / H100). +clinical notes, ICD codes, lab values, and chest X-rays. Example ------- - # Quick sanity check (--dev + 1 epoch + small batch): - python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ - --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 \\ - --task clinical_notes_icd_labs_cxr \\ - --model mlp --quick-test - - # Smoke test (single forward + inference, no training): - python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ - --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 \\ - --task clinical_notes_icd_labs_cxr \\ - --model mlp --smoke-forward --dev - - # Full training with Transformer + CXR: - python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ - --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 \\ - --task clinical_notes_icd_labs_cxr \\ - --model transformer --heads 4 --num-layers 2 \\ - --epochs 10 --batch-size 4 --device cuda:0 - - # JambaEHR + CXR: python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\ --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 \\ --task clinical_notes_icd_labs_cxr \\ - --model jambaehr --embedding-dim 128 \\ - --jamba-transformer-layers 2 --jamba-mamba-layers 6 + --model ehrmamba --num-layers 2 \\ + --epochs 10 --batch-size 8 --device cuda:0 """ from __future__ import annotations import argparse import csv -import time from pathlib import Path -from typing import Any, Tuple +from typing import Any, Optional, Tuple import numpy as np import torch @@ -77,35 +43,21 @@ from pyhealth.trainer import Trainer -# --------------------------------------------------------------------------- -# 1. Dataset construction -# --------------------------------------------------------------------------- - def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: - """Load the MIMIC-IV base dataset with the appropriate tables. - - For non-CXR tasks this behaves identically to Rian's runner. - For the CXR task it additionally passes cxr_root / cxr_variant / - cxr_tables so the dataset loads chest X-ray metadata. - """ ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] - - # Notes are needed for any task that includes clinical text. note_tables = None + + # Notes required for any task that includes clinical text if args.task in ("clinical_notes_icd_labs", "clinical_notes_icd_labs_cxr"): if not args.note_root: - raise ValueError( - f"--task {args.task} requires --note-root." - ) + raise ValueError(f"--task {args.task} requires --note-root.") note_tables = ["discharge", "radiology"] - # CXR tables are only loaded for the CXR task. + # CXR metadata tables only loaded for the CXR task cxr_kwargs = {} if args.task == "clinical_notes_icd_labs_cxr": if not args.cxr_root: - raise ValueError( - "--task clinical_notes_icd_labs_cxr requires --cxr-root." - ) + raise ValueError("--task clinical_notes_icd_labs_cxr requires --cxr-root.") cxr_kwargs = dict( cxr_root=args.cxr_root, cxr_variant=args.cxr_variant, @@ -124,12 +76,7 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: ) -# --------------------------------------------------------------------------- -# 2. Task construction -# --------------------------------------------------------------------------- - def _build_task(args: argparse.Namespace): - """Return the task object matching --task.""" if args.task == "stagenet": return MortalityPredictionStageNetMIMIC4() if args.task == "clinical_notes_icd_labs": @@ -139,28 +86,14 @@ def _build_task(args: argparse.Namespace): raise ValueError(f"Unknown task: {args.task}") -# --------------------------------------------------------------------------- -# 3. Train / val / test split -# --------------------------------------------------------------------------- - def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]: - """Split by patient first; fall back to split by sample if empty.""" - train_ds, val_ds, test_ds = split_by_patient( - dataset, [0.8, 0.1, 0.1], seed=seed - ) + 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 - ) + 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 -# --------------------------------------------------------------------------- -# 4. Model construction -# --------------------------------------------------------------------------- - def _build_model(args: argparse.Namespace, sample_dataset: Any): - """Instantiate UnifiedMultimodalEmbeddingModel + the chosen backbone.""" unified = UnifiedMultimodalEmbeddingModel( processors=sample_dataset.input_processors, embedding_dim=args.embedding_dim, @@ -229,17 +162,12 @@ def _build_model(args: argparse.Namespace, sample_dataset: Any): raise ValueError(f"Unknown model: {args.model}") -# --------------------------------------------------------------------------- -# 5. Prediction CSV writer -# --------------------------------------------------------------------------- - def _write_predictions( output_csv: Path, patient_ids: list[str], y_true: np.ndarray, y_prob: np.ndarray, ) -> None: - """Write per-sample predictions to a CSV file.""" output_csv.parent.mkdir(parents=True, exist_ok=True) y_true_flat = y_true.reshape(-1).tolist() @@ -262,21 +190,9 @@ def _write_predictions( ) -# --------------------------------------------------------------------------- -# 6. Main training + evaluation loop -# --------------------------------------------------------------------------- - def run(args: argparse.Namespace) -> Path: - """Execute the full pipeline: load → task → split → train → evaluate.""" torch.manual_seed(args.seed) np.random.seed(args.seed) - total_start = time.perf_counter() - - # Resolve CUDA device index for VRAM tracking. - cuda_device_index = None - if args.device and args.device.startswith("cuda"): - device_index = torch.device(args.device).index - cuda_device_index = 0 if device_index is None else device_index base_dataset = _build_base_dataset(args) task = _build_task(args) @@ -287,18 +203,6 @@ def run(args: argparse.Namespace) -> Path: "Task produced zero samples. Check roots/tables or adjust settings." ) - print(f"Task sample count: {len(sample_dataset)}") - - # Print processor schemas so mismatches are caught early. - print("Input processor schemas:") - for key in sample_dataset.input_schema.keys(): - processor = sample_dataset.input_processors.get(key) - if processor is None: - print(f" - {key}: ") - continue - print(f" - {key}: {type(processor).__name__}, " - f"schema={processor.schema()}") - train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed) model = _build_model(args, sample_dataset) @@ -314,35 +218,6 @@ def run(args: argparse.Namespace) -> Path: else None ) - print(f"Split sizes: train={len(train_ds)}, val={len(val_ds)}, " - f"test={len(test_ds)}") - - # Debug batch diagnostics: print types, shapes, and schema mappings - # for the first training batch so processor issues surface immediately. - debug_batch = next(iter(train_loader)) - print("Batch field diagnostics (train batch 0):") - for key in sample_dataset.input_schema.keys(): - processor = sample_dataset.input_processors.get(key) - feature = debug_batch.get(key) - schema = processor.schema() if processor is not None else () - print(f" - {key}: type={type(feature).__name__}, schema={schema}") - - if isinstance(feature, tuple): - for i, elem in enumerate(feature): - shape = getattr(elem, "shape", None) - print(f" tuple[{i}] type={type(elem).__name__} " - f"shape={shape}") - - if processor is not None and isinstance(feature, tuple): - for field_name in ("value", "time", "mask"): - if field_name in schema: - idx = schema.index(field_name) - if idx < len(feature): - selected = feature[idx] - shape = getattr(selected, "shape", None) - print(f" schema['{field_name}'] -> tuple[{idx}] " - f"type={type(selected).__name__} shape={shape}") - exp_name = f"{args.model}_seed{args.seed}" output_dir = Path(args.output_dir) @@ -355,7 +230,7 @@ def run(args: argparse.Namespace) -> Path: exp_name=exp_name, ) - # Model-specific optimizer defaults (from Rian's runner). + # BT uses tighter optimizer settings to stabilize training on full MIMIC-IV effective_lr = args.lr effective_max_grad_norm = args.max_grad_norm optimizer_params = {} @@ -365,9 +240,7 @@ def run(args: argparse.Namespace) -> Path: 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 - ) + 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-3 @@ -376,17 +249,7 @@ def run(args: argparse.Namespace) -> Path: optimizer_params["lr"] = effective_lr - # --smoke-forward: skip training, only run inference to verify the - # pipeline works end-to-end without waiting for a full epoch. - peak_train_vram_mb = None - train_runtime_sec = None - - if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0: - if cuda_device_index is not None: - torch.cuda.reset_peak_memory_stats(cuda_device_index) - torch.cuda.synchronize(cuda_device_index) - - train_start = time.perf_counter() + if args.epochs > 0 and len(train_ds) > 0: trainer.train( train_dataloader=train_loader, val_dataloader=val_loader, @@ -398,155 +261,86 @@ def run(args: argparse.Namespace) -> Path: load_best_model_at_last=True, ) - if cuda_device_index is not None: - torch.cuda.synchronize(cuda_device_index) - peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index) - peak_train_vram_mb = peak_train_bytes / (1024**2) - - train_runtime_sec = time.perf_counter() - train_start - inference_loader = test_loader or val_loader or train_loader y_true, y_prob, _, patient_ids = trainer.inference( inference_loader, return_patient_ids=True ) - if cuda_device_index is not None: - torch.cuda.synchronize(cuda_device_index) - - total_runtime_sec = time.perf_counter() - total_start - - # Benchmark summary — matches the Mamba CXR script's output format. - print("Benchmark summary:") - print(f" total_runtime_sec: {total_runtime_sec:.2f}") - if train_runtime_sec is None: - print(" training_runtime_sec: N/A (training skipped)") - print(" peak_train_vram_mb: N/A (training skipped)") - else: - print(f" training_runtime_sec: {train_runtime_sec:.2f}") - if peak_train_vram_mb is None: - print(" peak_train_vram_mb: N/A (non-CUDA device)") - else: - print(f" peak_train_vram_mb: {peak_train_vram_mb:.2f}") - output_csv = output_dir / exp_name / f"predictions_{args.model}.csv" _write_predictions(output_csv, patient_ids, y_true, y_prob) return output_csv -# --------------------------------------------------------------------------- -# 7. CLI argument parser -# --------------------------------------------------------------------------- - def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Run E2E unified embedding on MIMIC-IV (+ CXR) with " - "any of six backbone models." + description="Run E2E unified embedding on MIMIC-IV + CXR with any of six backbone models." ) - - # --- Data paths --- 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, - help="Root directory for MIMIC-CXR images. Required for " - "clinical_notes_icd_labs_cxr task.", - ) - parser.add_argument( - "--cxr-variant", type=str, default="sunlab", - choices=["default", "sunlab"], - help="CXR directory layout variant.", - ) + 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_cxr") - # --- Task and model selection --- parser.add_argument( "--task", type=str, default="clinical_notes_icd_labs_cxr", - choices=[ - "stagenet", - "clinical_notes_icd_labs", - "clinical_notes_icd_labs_cxr", - ], + choices=["stagenet", "clinical_notes_icd_labs", "clinical_notes_icd_labs_cxr"], ) parser.add_argument( "--model", type=str, default="rnn", - choices=[ - "mlp", "rnn", "transformer", "bottleneck_transformer", - "ehrmamba", "jambaehr", - ], + choices=["mlp", "rnn", "transformer", "bottleneck_transformer", "ehrmamba", "jambaehr"], ) - # --- Shared embedding / training hyperparameters --- 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=4, - help="Batch size. Default is 4 (CXR is VRAM-heavy, ~40 GB peak).", - ) - parser.add_argument( - "--lr", type=float, default=None, - help="Learning rate. Default: 1e-3 (1e-4 for bottleneck_transformer).", - ) - parser.add_argument( - "--adam-eps", type=float, default=None, - help="Adam epsilon. Default: 1e-8 (1e-6 for bottleneck_transformer).", - ) + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--lr", type=float, default=None) + parser.add_argument("--adam-eps", type=float, default=None) parser.add_argument("--weight-decay", type=float, default=0.0) parser.add_argument("--device", type=str, default=None) parser.add_argument("--num-workers", type=int, default=1) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--dev", action="store_true") - parser.add_argument( - "--quick-test", action="store_true", - help="Force --dev, 1 epoch, small batch for a quick sanity check.", - ) - parser.add_argument( - "--smoke-forward", action="store_true", - help="Skip training; only run a single forward pass + inference " - "to verify the pipeline end-to-end.", - ) - # --- Task-specific --- parser.add_argument("--observation-window-hours", type=int, default=24) - # --- RNN-specific --- 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") - # --- Transformer / BottleneckTransformer --- parser.add_argument("--heads", type=int, default=4) parser.add_argument("--num-layers", type=int, default=2) - # --- BottleneckTransformer-specific --- parser.add_argument("--bottlenecks-n", type=int, default=4) parser.add_argument("--fusion-startidx", type=int, default=1) - # --- Training stability --- - parser.add_argument( - "--max-grad-norm", type=float, default=None, - help="Gradient clip norm. Default: None (0.5 for bottleneck_transformer).", - ) + parser.add_argument("--max-grad-norm", type=float, default=None) - # --- Mamba / JambaEHR-specific --- parser.add_argument("--mamba-state-size", type=int, default=16) parser.add_argument("--mamba-conv-kernel", type=int, default=4) parser.add_argument("--jamba-transformer-layers", type=int, default=2) parser.add_argument("--jamba-mamba-layers", type=int, default=6) - args = parser.parse_args() + return parser.parse_args() - if args.quick_test: - args.dev = True - args.epochs = 1 - args.batch_size = min(args.batch_size, 4) - return args +def print_vram_summary(device: Optional[str] = None) -> None: + """Print peak allocated and total VRAM for the given CUDA device.""" + if not torch.cuda.is_available(): + print("VRAM summary: CUDA not available.") + return + idx = torch.device(device).index if device and device.startswith("cuda") else 0 + if idx is None: + idx = 0 + allocated_mb = torch.cuda.max_memory_allocated(idx) / (1024 ** 2) + total_mb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 2) + print(f"VRAM summary (cuda:{idx}): peak_allocated={allocated_mb:.0f} MB / total={total_mb:.0f} MB") 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 + print(f"Saved predictions to: {output_csv_path}") + print_vram_summary(args.device) From f528f18108cc303f7ee4a2dfdfe399fa09133aea Mon Sep 17 00:00:00 2001 From: William Pang Date: Fri, 24 Apr 2026 19:34:28 -0500 Subject: [PATCH 12/30] Add RNN Script for Sunlab --- scripts/sunlab/run_rnn.sh | 44 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100755 scripts/sunlab/run_rnn.sh diff --git a/scripts/sunlab/run_rnn.sh b/scripts/sunlab/run_rnn.sh new file mode 100755 index 000000000..f47238805 --- /dev/null +++ b/scripts/sunlab/run_rnn.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit Table 2 jobs for rnn — 3 fixed seeds, sunlab cluster (tmux sessions). +# Seeds are shared across all lab members for aligned comparisons. +# Usage: bash scripts/sunlab/run_rnn.sh +set -euo pipefail + +# SEEDS=(267573289 1872967241 706384748) +SEEDS=(42) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="${USER:-wp14}" +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +CUDA_DEVICE="${CUDA_DEVICE:-0}" +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:-/home/${USER}/pyhealth_cache}" +OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}" + +cd "${PROJECT_DIR}" +mkdir -p logs/sunlab + +for seed in "${SEEDS[@]}"; do + session="rnn_seed${seed}" + echo " Launching rnn seed=${seed} in tmux session: ${session}" + tmux new-session -d -s "${session}" "bash -lc ' + conda activate ${CONDA_ENV} && \ + CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root ${EHR_ROOT} \ + --cache-dir ${CACHE_DIR} \ + --task icd_labs \ + --model rnn \ + --embedding-dim 128 \ + --hidden-dim 128 \ + --heads 4 \ + --num-layers 2 \ + --dropout 0.1 \ + --epochs 20 \ + --batch-size 16 \ + --weight-decay 1e-5 \ + --patience 5 \ + --seed ${seed} \ + --output-dir ${OUTPUT_BASE}/rnn_seed${seed} \ + 2>&1 | tee logs/sunlab/rnn_seed${seed}.log; exec bash'" + echo " Launched rnn seed=${seed} → tmux session: ${session}" +done From 26754ee60a9aebe0e64d7aa0201d16c7a85ba48f Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 28 Apr 2026 10:22:47 -0500 Subject: [PATCH 13/30] Add run transfomrer --- scripts/sunlab/run_transformer.sh | 44 +++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 scripts/sunlab/run_transformer.sh diff --git a/scripts/sunlab/run_transformer.sh b/scripts/sunlab/run_transformer.sh new file mode 100644 index 000000000..0a2aacaff --- /dev/null +++ b/scripts/sunlab/run_transformer.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Submit Table 2 jobs for transformer — 3 fixed seeds, sunlab cluster (tmux sessions). +# Seeds are shared across all lab members for aligned comparisons. +# Usage: bash scripts/sunlab/run_transformer.sh +set -euo pipefail + +# SEEDS=(267573289 1872967241 706384748) +SEEDS=(44) +PROJECT_DIR="${PROJECT_DIR:-$(pwd)}" +USER="${USER:-wp14}" +CONDA_ENV="${CONDA_ENV:-pyhealth2}" +CUDA_DEVICE="${CUDA_DEVICE:-4}" +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:-/home/${USER}/pyhealth_cache}" +OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}" + +cd "${PROJECT_DIR}" + +for seed in "${SEEDS[@]}"; do + session="transformer_seed${seed}" + log_dir="logs/sunlab/$(date +%Y%m%d)" + mkdir -p "${log_dir}" + tmux new-session -d -s "${session}" "bash -lc ' + conda activate ${CONDA_ENV} && \ + CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \ + --ehr-root ${EHR_ROOT} \ + --cache-dir ${CACHE_DIR} \ + --task icd_labs \ + --model transformer \ + --embedding-dim 64 \ + --hidden-dim 64 \ + --heads 2 \ + --num-layers 1 \ + --dropout 0.1 \ + --epochs 20 \ + --batch-size 4 \ + --weight-decay 1e-5 \ + --patience 5 \ + --seed ${seed} \ + --output-dir ${OUTPUT_BASE}/$(date +%Y%m%d) \ + 2>&1 | tee ${log_dir}/transformer_seed${seed}.log; exec bash'" + echo " Launched transformer seed=${seed} → tmux session: ${session}, log: ${log_dir}/transformer_seed${seed}.log" +done From 4e56e2801ea4ccb7c12f2e18b6bfd30f41061116 Mon Sep 17 00:00:00 2001 From: William Pang Date: Sat, 2 May 2026 21:53:08 -0700 Subject: [PATCH 14/30] Update discharge.csv --- test-resources/core/mimic4demo/note/discharge.csv | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test-resources/core/mimic4demo/note/discharge.csv b/test-resources/core/mimic4demo/note/discharge.csv index 116c6d69e..a021ce257 100644 --- a/test-resources/core/mimic4demo/note/discharge.csv +++ b/test-resources/core/mimic4demo/note/discharge.csv @@ -6,7 +6,7 @@ d4,10001,20002,DS,1,2150-06-25 12:00:00,2150-06-25 13:00:00,Patient 10001 admitt d5,10002,20003,DS,1,2151-01-12 12:00:00,2151-01-12 13:00:00,Patient 10002 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d6,10002,20004,DS,1,2151-04-20 12:00:00,2151-04-20 13:00:00,Patient 10002 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d7,10003,20005,DS,1,2152-03-05 12:00:00,2152-03-05 13:00:00,Patient 10003 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. -d8,10003,20006,DS,1,2152-08-15 12:00:00,2152-08-15 13:00:00,Patient 10003 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. +d8,10003,20006,DS,1,2152-08-15 12:00:00,2152-08-15 13:00:00,Patient 10003 admitted for evaluation. Vitals unstable on presentation. Clinical status deteriorated despite aggressive management. Patient expired during this admission. Death pronounced on 2152-08-15. d9,10004,20007,DS,1,2150-05-02 12:00:00,2150-05-02 13:00:00,Patient 10004 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d10,10005,20008,DS,1,2151-07-22 12:00:00,2151-07-22 13:00:00,Patient 10005 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. d11,10006,20009,DS,1,2152-09-08 12:00:00,2152-09-08 13:00:00,Patient 10006 admitted for evaluation. Vitals stable. Labs reviewed. Discharged in good condition. From 6b9d54283fb8544edb328645cd6a82483371ce86 Mon Sep 17 00:00:00 2001 From: William Pang Date: Sat, 2 May 2026 22:11:21 -0700 Subject: [PATCH 15/30] Update diagnoses_icd.csv --- test-resources/core/mimic4demo/hosp/diagnoses_icd.csv | 1 + 1 file changed, 1 insertion(+) diff --git a/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv index 804e2ae40..21511f015 100644 --- a/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv +++ b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv @@ -19,6 +19,7 @@ subject_id,hadm_id,seq_num,icd_code,icd_version 10003,20006,2,E1065,10 10003,20006,3,N170,10 10003,20006,4,I509,10 +10003,20006,5,R99,10 10004,20007,1,K219,10 10004,20007,2,I10,10 10005,20008,1,25001,9 From 8e3e1b4b9787c22dc3df22e41ab1d2bec815b7b7 Mon Sep 17 00:00:00 2001 From: William Pang Date: Sat, 2 May 2026 23:15:10 -0700 Subject: [PATCH 16/30] Update multimodal_mimic4_task_tutorial.ipynb --- .../multimodal_mimic4_task_tutorial.ipynb | 121 +++++++++++------- 1 file changed, 78 insertions(+), 43 deletions(-) diff --git a/examples/multimodal_mimic4_task_tutorial.ipynb b/examples/multimodal_mimic4_task_tutorial.ipynb index da288e2fb..87b12d4db 100644 --- a/examples/multimodal_mimic4_task_tutorial.ipynb +++ b/examples/multimodal_mimic4_task_tutorial.ipynb @@ -92,7 +92,7 @@ "metadata": {}, "outputs": [], "source": [ - "SEED = 25\n", + "SEED = 22\n", "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)" @@ -113,7 +113,7 @@ "metadata": {}, "outputs": [], "source": [ - "MORTALITY_LABEL = [1]" + "MORTALITY_LABEL = 0" ] }, { @@ -144,8 +144,8 @@ " return [s for s in samples if s['mortality'].item() in labels]\n", "\n", "\n", - "def print_patient_admission_info(sample_patient_id):\n", - " patient = dataset.get_patient(sample_patient_id)\n", + "def print_patient_admission_info(sample):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", " admissions = patient.get_events(\n", " event_type=\"admissions\",\n", " start=sample['window_start'],\n", @@ -158,8 +158,8 @@ " print(f\" dischtime: {adm.dischtime}\")\n", "\n", "\n", - "def print_patient_note_info(sample_patient_id, note_type=\"discharge\", char_limit=80):\n", - " patient = dataset.get_patient(sample_patient_id)\n", + "def print_patient_note_info(sample, note_type=\"discharge\", char_limit=80):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", " notes = patient.get_events(\n", " event_type=note_type,\n", " start=sample['window_start'],\n", @@ -181,8 +181,8 @@ "}\n", "\n", "\n", - "def print_patient_lab_info(sample_patient_id):\n", - " patient = dataset.get_patient(sample_patient_id)\n", + "def print_patient_lab_info(sample):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", " labs = patient.get_events(\n", " event_type=\"labevents\",\n", " start=sample['window_start'],\n", @@ -196,6 +196,31 @@ " print(f\" storetime: {lab['storetime']}\")\n", " print(f\" valuenum: {lab['valuenum']}\")\n", "\n", + "\n", + "def print_patient_icd_info(sample):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", + " diagnoses = patient.get_events(\n", + " event_type=\"diagnoses_icd\",\n", + " start=sample['window_start'],\n", + " end=sample['window_end'],\n", + " )\n", + " procedures = patient.get_events(\n", + " event_type=\"procedures_icd\",\n", + " start=sample['window_start'],\n", + " end=sample['window_end'],\n", + " )\n", + " print(f\"\\ndiagnoses_icd: {len(diagnoses)}\")\n", + " for dx in diagnoses:\n", + " print(f\" hadm_id: {dx.hadm_id}\")\n", + " print(f\" seq_num: {dx.seq_num}\")\n", + " print(f\" icd_code: {dx.icd_code} (ICD-{dx.icd_version})\\n\")\n", + " print(f\"\\nprocedures_icd: {len(procedures)}\")\n", + " for px in procedures:\n", + " print(f\" hadm_id: {px.hadm_id}\")\n", + " print(f\" seq_num: {px.seq_num}\")\n", + " print(f\" icd_code: {px.icd_code} (ICD-{px.icd_version})\\n\")\n", + "\n", + "\n", "def clear_cache_directory(path):\n", " for item in os.listdir(path):\n", " item_path = os.path.join(path, item)\n", @@ -275,20 +300,22 @@ "outputs": [], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", - "sample = samples[0]\n", - "sample_patient_id = sample['patient_id']\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", "# Admissions\n", - "print_patient_admission_info(sample_patient_id)\n", + "print_patient_admission_info(sample)\n", "# Discharge Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"discharge\")\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"radiology\")\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", "# Labs\n", - "print_patient_lab_info(sample_patient_id)" + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" ] }, { @@ -335,20 +362,22 @@ "outputs": [], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", - "sample = samples[0]\n", - "sample_patient_id = sample['patient_id']\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", "# Admissions\n", - "print_patient_admission_info(sample_patient_id)\n", + "print_patient_admission_info(sample)\n", "# Discharge Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"discharge\")\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"radiology\")\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", "# Labs\n", - "print_patient_lab_info(sample_patient_id)" + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" ] }, { @@ -397,20 +426,22 @@ "outputs": [], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", - "sample = samples[0]\n", - "sample_patient_id = sample['patient_id']\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", "# Admissions\n", - "print_patient_admission_info(sample_patient_id)\n", + "print_patient_admission_info(sample)\n", "# Discharge Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"discharge\")\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"radiology\")\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", "# Labs\n", - "print_patient_lab_info(sample_patient_id)" + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" ] }, { @@ -468,20 +499,22 @@ "outputs": [], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", - "sample = samples[0]\n", - "sample_patient_id = sample['patient_id']\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "\n", "# Admissions\n", - "print_patient_admission_info(sample_patient_id)\n", + "print_patient_admission_info(sample)\n", "# Discharge Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"discharge\")\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"radiology\")\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", "# Labs\n", - "print_patient_lab_info(sample_patient_id)" + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" ] }, { @@ -541,22 +574,24 @@ "outputs": [], "source": [ "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n", - "sample = samples[0]\n", - "sample_patient_id = sample['patient_id']\n", + "random.seed(SEED)\n", + "sample = random.choice(samples)\n", "\n", - "print(f\"\\nPatient ID: {sample_patient_id}\")\n", + "print(f\"\\nPatient ID: {sample['patient_id']}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", "print(f\"Window start: {sample['window_start']}\")\n", "print(f\"Window end: {sample['window_end']}\")\n", "\n", "# Admissions\n", - "print_patient_admission_info(sample_patient_id)\n", + "print_patient_admission_info(sample)\n", "# Discharge Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"discharge\")\n", + "print_patient_note_info(sample, note_type=\"discharge\")\n", "# Radiology Notes\n", - "print_patient_note_info(sample_patient_id, note_type=\"radiology\")\n", + "print_patient_note_info(sample, note_type=\"radiology\")\n", "# Labs\n", - "print_patient_lab_info(sample_patient_id)" + "print_patient_lab_info(sample)\n", + "# ICD Codes\n", + "print_patient_icd_info(sample)" ] }, { From d1ddce5c27271850152dc2a02b5ac85d294a690d Mon Sep 17 00:00:00 2001 From: Rian354 Date: Thu, 14 May 2026 05:39:45 -0400 Subject: [PATCH 17/30] feat(scripts): YAML-driven unified training config and CLI wrapper Adds configs/train/ hierarchy and scripts/train_unified.py so any experiment can be launched from a single reproducible config file instead of long bash one-liners or manual env-var edits. Co-Authored-By: Claude Sonnet 4.6 --- configs/train/base.yaml | 80 +++++++++++++++ configs/train/e2e_balanced.yaml | 8 ++ configs/train/e2e_baseline.yaml | 11 +++ configs/train/e2e_icd_on.yaml | 7 ++ configs/train/e2e_labs_only.yaml | 6 ++ configs/train/smoke.yaml | 14 +++ scripts/train_unified.py | 163 +++++++++++++++++++++++++++++++ 7 files changed, 289 insertions(+) create mode 100644 configs/train/base.yaml create mode 100644 configs/train/e2e_balanced.yaml create mode 100644 configs/train/e2e_baseline.yaml create mode 100644 configs/train/e2e_icd_on.yaml create mode 100644 configs/train/e2e_labs_only.yaml create mode 100644 configs/train/smoke.yaml create mode 100644 scripts/train_unified.py diff --git a/configs/train/base.yaml b/configs/train/base.yaml new file mode 100644 index 000000000..70c7f2c7c --- /dev/null +++ b/configs/train/base.yaml @@ -0,0 +1,80 @@ +# PyHealth unified training config — base defaults +# All values here are the OOM-safe defaults validated on devsplit. +# Override per-condition and per-model in condition/model-specific configs. + +# ── paths ────────────────────────────────────────────────────────────────────── +ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2 +note_root: /shared/rsaas/physionet.org/files/mimic-note +cache_dir: /home/rianatri/pyhealth_cache +output_dir: output/unified + +# ── task ────────────────────────────────────────────────────────────────────── +# One of: notes_labs | labs_only | icd_labs | clinical_notes_icd_labs | stagenet +task: notes_labs +observation_window_hours: 24 + +# Task flags (notes_labs only) +icd_codes: false +include_vitals: false +balanced_sampling: false +balanced_ratio: 1.0 + +# ── model ───────────────────────────────────────────────────────────────────── +# One of: mlp | rnn | transformer | bottleneck_transformer | ehrmamba | jambaehr +model: mlp +freeze_encoder: false + +# Shared embedding dims — safe defaults across all models on 24 GB GPU (frozen) +# or 80 GB A100 (full BERT). +embedding_dim: 128 +hidden_dim: 128 +heads: 4 +num_layers: 2 +dropout: 0.1 + +# ── training ────────────────────────────────────────────────────────────────── +epochs: 50 +batch_size: 16 +lr: null # null = model-specific default (1e-4 for all) +weight_decay: 1.0e-5 +patience: 10 +seed: 42 +num_workers: 2 + +# ── dev mode ────────────────────────────────────────────────────────────────── +# 0 = full dataset; N > 0 = limit to N patients (devsplit) +dev: 0 + +# ── model-specific overrides ────────────────────────────────────────────────── +# These are merged at runtime based on `model` value. +# Keys match CLI args without the leading `--`. +_model_overrides: + transformer: + batch_size: 2 + embedding_dim: 64 + hidden_dim: 64 + heads: 2 + num_layers: 1 + bottleneck_transformer: + batch_size: 2 + embedding_dim: 96 + hidden_dim: 96 + heads: 2 + num_layers: 1 + max_grad_norm: 0.5 + bottlenecks_n: 4 + fusion_startidx: 1 + ehrmamba: + batch_size: 2 + embedding_dim: 96 + hidden_dim: 96 + mamba_state_size: 16 + mamba_conv_kernel: 4 + jambaehr: + batch_size: 2 + embedding_dim: 64 + hidden_dim: 64 + jamba_transformer_layers: 1 + jamba_mamba_layers: 2 + mamba_state_size: 16 + mamba_conv_kernel: 4 diff --git a/configs/train/e2e_balanced.yaml b/configs/train/e2e_balanced.yaml new file mode 100644 index 000000000..74aaedcb9 --- /dev/null +++ b/configs/train/e2e_balanced.yaml @@ -0,0 +1,8 @@ +# PyHealth E2E balanced condition — notes+labs with 1:1 pos:neg undersampling +_inherit: base.yaml + +task: notes_labs +balanced_sampling: true +balanced_ratio: 1.0 +epochs: 50 +output_dir: output/e2e_full/balanced diff --git a/configs/train/e2e_baseline.yaml b/configs/train/e2e_baseline.yaml new file mode 100644 index 000000000..df1a6243e --- /dev/null +++ b/configs/train/e2e_baseline.yaml @@ -0,0 +1,11 @@ +# PyHealth E2E baseline condition — notes+labs, ICD off, no balancing +# Full dataset, 50 epochs, patience=10. + +_inherit: base.yaml + +task: notes_labs +icd_codes: false +include_vitals: false +balanced_sampling: false +epochs: 50 +output_dir: output/e2e_full/baseline diff --git a/configs/train/e2e_icd_on.yaml b/configs/train/e2e_icd_on.yaml new file mode 100644 index 000000000..13f817fd3 --- /dev/null +++ b/configs/train/e2e_icd_on.yaml @@ -0,0 +1,7 @@ +# PyHealth E2E ICD-on condition — notes+labs+ICD (discharge-coded leakage ablation) +_inherit: base.yaml + +task: notes_labs +icd_codes: true +epochs: 50 +output_dir: output/e2e_full/icd_on diff --git a/configs/train/e2e_labs_only.yaml b/configs/train/e2e_labs_only.yaml new file mode 100644 index 000000000..d0ca5ade8 --- /dev/null +++ b/configs/train/e2e_labs_only.yaml @@ -0,0 +1,6 @@ +# PyHealth E2E labs-only condition — EHR-only reference baseline +_inherit: base.yaml + +task: labs_only +epochs: 50 +output_dir: output/e2e_full/labs_only diff --git a/configs/train/smoke.yaml b/configs/train/smoke.yaml new file mode 100644 index 000000000..af0079770 --- /dev/null +++ b/configs/train/smoke.yaml @@ -0,0 +1,14 @@ +# PyHealth smoke test config — devsplit, fast iteration +# Inherits base.yaml defaults; overrides for fast validation. + +_inherit: base.yaml + +task: notes_labs +model: mlp +dev: 1000 +epochs: 5 +patience: 5 +batch_size: 16 +embedding_dim: 64 +hidden_dim: 64 +output_dir: output/smoke diff --git a/scripts/train_unified.py b/scripts/train_unified.py new file mode 100644 index 000000000..6d908b491 --- /dev/null +++ b/scripts/train_unified.py @@ -0,0 +1,163 @@ +"""Unified training CLI for PyHealth multimodal experiments. + +Reads a YAML config (with optional _inherit chain) then fires the E2E script +with all resolved arguments. CLI flags override YAML values. + +Usage: + python scripts/train_unified.py --config configs/train/e2e_baseline.yaml \ + --model transformer --seed 0 + + # Smoke test + python scripts/train_unified.py --config configs/train/smoke.yaml + + # Full run with frozen encoder + python scripts/train_unified.py --config configs/train/e2e_balanced.yaml \ + --model ehrmamba --freeze-encoder --seed 42 + +The script resolves _model_overrides from the YAML, merges CLI overrides, then +calls unified_embedding_e2e_mimic4.py via subprocess so it can be used from +condor jobs without reimplementing the full arg-build logic. +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict + + +def _load_yaml(path: Path) -> Dict[str, Any]: + try: + import yaml # type: ignore + except ImportError: + raise SystemExit("PyYAML not installed. Run: pip install pyyaml") + with open(path) as f: + return yaml.safe_load(f) or {} + + +def _resolve_config(config_path: Path) -> Dict[str, Any]: + """Load config and recursively merge _inherit chain.""" + cfg = _load_yaml(config_path) + inherit = cfg.pop("_inherit", None) + if inherit: + parent_path = config_path.parent / inherit + parent = _resolve_config(parent_path) + # parent values are defaults; child values win + merged = {**parent, **cfg} + return merged + return cfg + + +def _apply_model_overrides(cfg: Dict[str, Any]) -> Dict[str, Any]: + """Merge _model_overrides[model] into cfg if present.""" + overrides = cfg.pop("_model_overrides", {}) + model = cfg.get("model", "mlp") + if model in overrides: + cfg = {**cfg, **overrides[model]} + return cfg + + +def _to_args(cfg: Dict[str, Any]) -> list[str]: + """Convert resolved config dict to CLI args for unified_embedding_e2e_mimic4.py.""" + # Map config keys to CLI flags (underscores → hyphens, bool flags → store_true) + bool_flags = { + "icd_codes", "freeze_encoder", "include_vitals", + "balanced_sampling", "bidirectional", + } + skip_keys = {"_inherit", "_model_overrides"} + + args = [] + for key, val in cfg.items(): + if key in skip_keys or val is None: + continue + flag = "--" + key.replace("_", "-") + if key in bool_flags: + if val: + args.append(flag) + else: + args += [flag, str(val)] + return args + + +def _parse_cli() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Unified PyHealth training CLI — loads a YAML config then fires the E2E script." + ) + parser.add_argument( + "--config", + required=True, + help="Path to a YAML config file (e.g. configs/train/e2e_baseline.yaml).", + ) + # Pass-through overrides — any key from the YAML can be overridden here. + parser.add_argument("--model", default=None) + parser.add_argument("--seed", type=int, default=None) + parser.add_argument("--task", default=None) + parser.add_argument("--epochs", type=int, default=None) + parser.add_argument("--batch-size", type=int, default=None) + parser.add_argument("--embedding-dim", type=int, default=None) + parser.add_argument("--dev", type=int, default=None) + parser.add_argument("--freeze-encoder", action="store_true", default=None) + parser.add_argument("--balanced-sampling", action="store_true", default=None) + parser.add_argument("--icd-codes", action="store_true", default=None) + parser.add_argument("--output-dir", default=None) + parser.add_argument("--ehr-root", default=None) + parser.add_argument("--note-root", default=None) + parser.add_argument("--cache-dir", default=None) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the resolved command without running it.", + ) + return parser.parse_args() + + +def main() -> None: + cli = _parse_cli() + + config_path = Path(cli.config) + if not config_path.exists(): + raise SystemExit(f"Config not found: {config_path}") + + cfg = _resolve_config(config_path) + cfg = _apply_model_overrides(cfg) + + # Apply CLI overrides (non-None values win over YAML) + cli_overrides = { + "model": cli.model, + "seed": cli.seed, + "task": cli.task, + "epochs": cli.epochs, + "batch_size": cli.batch_size, + "embedding_dim": cli.embedding_dim, + "dev": cli.dev, + "freeze_encoder": cli.freeze_encoder if cli.freeze_encoder else None, + "balanced_sampling": cli.balanced_sampling if cli.balanced_sampling else None, + "icd_codes": cli.icd_codes if cli.icd_codes else None, + "output_dir": cli.output_dir, + "ehr_root": cli.ehr_root, + "note_root": cli.note_root, + "cache_dir": cli.cache_dir, + } + for k, v in cli_overrides.items(): + if v is not None: + cfg[k] = v + + script = Path(__file__).parent.parent / "examples" / "mortality_prediction" / "unified_embedding_e2e_mimic4.py" + cmd = [sys.executable, str(script)] + _to_args(cfg) + + print("Resolved command:") + print(" " + " \\\n ".join(cmd)) + + if cli.dry_run: + print("\n[dry-run] Not executing.") + return + + result = subprocess.run(cmd, env={**os.environ, "PYTHONPATH": str(script.parent.parent.parent)}) + sys.exit(result.returncode) + + +if __name__ == "__main__": + main() From 98ddfe095d605758e8c4379ecaaccd099359a6bc Mon Sep 17 00:00:00 2001 From: Rian354 Date: Thu, 14 May 2026 05:40:11 -0400 Subject: [PATCH 18/30] fix(tasks): standardize tokenizer to Bio_ClinicalBERT across all task classes ClinicalNotesMIMIC4 and ClinicalNotesICDLabsMIMIC4 were using bert-base-uncased while the runtime encoder loaded Bio_ClinicalBERT, causing a silent token-ID mismatch on clinical terminology. Co-Authored-By: Claude Sonnet 4.6 --- pyhealth/tasks/multimodal_mimic4.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index dbd643bf0..a87628c83 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -35,14 +35,14 @@ class ClinicalNotesMIMIC4(BaseTask): "discharge_note_times": ( "tuple_time_text", { - "tokenizer_model": "bert-base-uncased", + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", }, ), "radiology_note_times": ( "tuple_time_text", { - "tokenizer_model": "bert-base-uncased", + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", }, ) @@ -200,14 +200,14 @@ class ClinicalNotesICDLabsMIMIC4(BaseTask): "discharge_note_times": ( "tuple_time_text", { - "tokenizer_model": "bert-base-uncased", + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", }, ), "radiology_note_times": ( "tuple_time_text", { - "tokenizer_model": "bert-base-uncased", + "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", "type_tag": "note", }, ), From 5b12bc6581590a190dca7c84f42c35d2a46274d9 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Thu, 14 May 2026 05:44:55 -0400 Subject: [PATCH 19/30] feat(tasks): add LabsOnlyMIMIC4 EHR-only reference baseline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pure lab-value mortality prediction (no notes, no ICD codes). Mirrors the labs/labs_mask schema from NotesLabsMIMIC4 so the same backbones work unchanged. Full-scale results: transformer 0.496/0.888 PR/ROC-AUC at ep9, MLP 0.399/0.880 at ep38 — best non-leaky result across 18 runs. Co-Authored-By: Claude Sonnet 4.6 --- .../unified_embedding_e2e_mimic4.py | 10 ++- pyhealth/tasks/__init__.py | 2 + pyhealth/tasks/multimodal_mimic4.py | 63 +++++++++++++++++++ 3 files changed, 74 insertions(+), 1 deletion(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 0a2dfce91..e520f181d 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -72,6 +72,7 @@ from pyhealth.tasks.multimodal_mimic4 import ( ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, + LabsOnlyMIMIC4, NotesLabsMIMIC4, ) from pyhealth.trainer import Trainer @@ -104,6 +105,11 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if "chartevents" not in ehr_tables: ehr_tables.append("chartevents") + if args.task == "labs_only": + # Pure EHR baseline: only labevents, no notes, no ICD codes. + ehr_tables = ["labevents"] + note_tables = None + return MIMIC4Dataset( ehr_root=args.ehr_root, ehr_tables=ehr_tables, @@ -128,6 +134,8 @@ def _build_task(args: argparse.Namespace): include_icd=args.icd_codes, include_vitals=args.include_vitals, ) + if args.task == "labs_only": + return LabsOnlyMIMIC4(window_hours=args.observation_window_hours) raise ValueError(f"Unknown task: {args.task}") @@ -379,7 +387,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, - choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs"], + choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs_only"], default="stagenet", help=( "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index be6ec9fca..569b035f0 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -49,6 +49,8 @@ from .multimodal_mimic4 import ( ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, + ICDLabsMIMIC4, + LabsOnlyMIMIC4, NotesLabsMIMIC4, ) from .patient_linkage import patient_linkage_mimic3_fn diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 248377322..56ce98f69 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -1077,3 +1077,66 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: record["icd_codes"] = (all_icd_times, all_icd_codes) return [record] + + +class LabsOnlyMIMIC4(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 = "LabsOnlyMIMIC4" + + 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 = self._build_admissions_to_process(patient) + if not admissions: + return [] + + for admission_time, admission_dischtime, admission in admissions: + mortality_label = self._get_mortality_label(patient, admission) + effective_start, effective_end = self._compute_window( + admission_time, admission_dischtime, self.window_hours + ) + + lab_times, lab_values, lab_masks = self._collect_labs( + patient=patient, + admission_time=admission_time, + end_time=effective_end, + ) + + if len(lab_values) == 0: + 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) + + yield { + "patient_id": patient.patient_id, + "labs": (lab_times, lab_values), + "labs_mask": (lab_times, lab_masks), + "mortality": mortality_label, + "window_start": effective_start, + "window_end": effective_end, + } From 3d14ccedfcc3cb4c76e504ac12924f216fbe1c85 Mon Sep 17 00:00:00 2001 From: Joshua Chen Date: Sun, 17 May 2026 14:22:34 -0500 Subject: [PATCH 20/30] Add pool='mean' option to VisionEmbeddingModel --- pyhealth/models/embedding/vision.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/pyhealth/models/embedding/vision.py b/pyhealth/models/embedding/vision.py index b1fef865e..7ce8d9bb3 100644 --- a/pyhealth/models/embedding/vision.py +++ b/pyhealth/models/embedding/vision.py @@ -116,11 +116,13 @@ def __init__( 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 @@ -286,6 +288,9 @@ def forward( 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: @@ -303,7 +308,10 @@ def get_output_info(self, field_name: str) -> Dict[str, Any]: info = self._field_info[field_name].copy() info["embedding_dim"] = self._embedding_dim info["has_cls_token"] = self.use_cls_token - info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0) + 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: From af66cef34f43b116c4a5a3cadaa88fdefcb3c655 Mon Sep 17 00:00:00 2001 From: Joshua Chen Date: Sun, 17 May 2026 14:38:18 -0500 Subject: [PATCH 21/30] tests --- pyhealth/models/embedding/vision.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pyhealth/models/embedding/vision.py b/pyhealth/models/embedding/vision.py index 7ce8d9bb3..57eedda22 100644 --- a/pyhealth/models/embedding/vision.py +++ b/pyhealth/models/embedding/vision.py @@ -358,9 +358,23 @@ def __repr__(self) -> str: 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}") From ac476708eb5ead84ce5804237b95bf9fc3d151f3 Mon Sep 17 00:00:00 2001 From: William Pang Date: Sun, 24 May 2026 18:43:40 -0700 Subject: [PATCH 22/30] Update vision_embedding_tutorial.ipynb --- examples/vision_embedding_tutorial.ipynb | 581 +++-------------------- 1 file changed, 76 insertions(+), 505 deletions(-) diff --git a/examples/vision_embedding_tutorial.ipynb b/examples/vision_embedding_tutorial.ipynb index 913422105..3d62b68bf 100644 --- a/examples/vision_embedding_tutorial.ipynb +++ b/examples/vision_embedding_tutorial.ipynb @@ -9,7 +9,7 @@ "\n", "This notebook demonstrates how to use the `VisionEmbeddingModel` for medical imaging tasks in PyHealth.\n", "\n", - "**Contributors:** Josh Steier \n", + "**Contributors:** Josh Steier, Joshua Chen, William Pang\n", "\n", "\n", "**Overview:**\n", @@ -31,18 +31,10 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "272cc5bb", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Running on device: cpu\n" - ] - } - ], + "outputs": [], "source": [ "import os\n", "import random\n", @@ -84,18 +76,10 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "f07f439c", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Created 200 synthetic images in C:\\Users\\637682\\AppData\\Local\\Temp\\chest_xray_jzp9zbiz\n" - ] - } - ], + "outputs": [], "source": [ "USE_SYNTHETIC = True\n", "MIMIC_CXR_ROOT = \"/path/to/physionet.org/files/mimic-cxr-jpg/2.0.0\"\n", @@ -171,22 +155,10 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "c3f49ff5", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Label label vocab: {0: 0, 1: 1}\n", - "Total task samples: 200\n", - "Input schema: {'image': 'image'}\n", - "Output schema: {'label': 'binary'}\n", - "Train/Val/Test sizes: 140, 28, 32\n" - ] - } - ], + "outputs": [], "source": [ "image_mode = \"L\" if USE_SYNTHETIC else \"RGB\"\n", "\n", @@ -218,19 +190,10 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "4684f6c1", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "{'patient_id': 'list(len=16)', 'visit_id': 'list(len=16)', 'image': 'Tensor(shape=(16, 1, 224, 224))', 'label': 'Tensor(shape=(16, 1))'}\n", - "Sample labels: [[0.0], [0.0], [0.0], [1.0], [1.0]]\n" - ] - } - ], + "outputs": [], "source": [ "BATCH_SIZE = 16\n", "\n", @@ -260,21 +223,10 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "a47ebd7e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "patient_id: list(len=16)\n", - "visit_id: list(len=16)\n", - "image: shape=torch.Size([16, 1, 224, 224]), dtype=torch.float32\n", - "label: shape=torch.Size([16, 1]), dtype=torch.float32\n" - ] - } - ], + "outputs": [], "source": [ "batch = next(iter(train_loader))\n", "\n", @@ -304,52 +256,10 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "c4d35bc6", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "C:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth\\sampler\\sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n", - " import pkg_resources\n", - "c:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth-env\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionEmbeddingModel Output Shape Verification\n", - "============================================================\n", - "\n", - "Input shape: torch.Size([16, 1, 224, 224]) # (B, C, H, W)\n", - "\n", - "patch:\n", - " Output shape: (16, 197, 128) # (B, num_tokens, E)\n", - " Tokens: 197 = 196 patches + 1 CLS\n", - " Parameters: 58,240\n", - "\n", - "cnn:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 231,808\n", - "\n", - "resnet18:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 11,242,432\n", - "\n", - "resnet50:\n", - " Output shape: (16, 50, 128) # (B, num_tokens, E)\n", - " Tokens: 50 = 49 patches + 1 CLS\n", - " Parameters: 23,770,560\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.models import VisionEmbeddingModel\n", "\n", @@ -389,6 +299,52 @@ " print()" ] }, + { + "cell_type": "markdown", + "id": "604e21d3", + "metadata": {}, + "source": [ + "## 4.6. Mean Pooling: Compact Single-Vector Embeddings\n", + "\n", + "Pass `pool='mean'` to collapse all patch tokens into a single averaged vector per image:\n", + "\n", + "```\n", + "With pool='mean': (B, 1, E) → e.g., (16, 1, 128)\n", + "```" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3be8a484", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Mean Pooling vs Token Sequence Output\")\n", + "print(\"=\" * 60)\n", + "print(f\"\\nInput shape: {batch['image'].shape} # (B, C, H, W)\")\n", + "print()\n", + "\n", + "# With mean pooling: single averaged vector\n", + "model_pooled = VisionEmbeddingModel(\n", + " dataset=sample_dataset,\n", + " embedding_dim=128,\n", + " backbone=\"cnn\",\n", + " pool=\"mean\",\n", + ")\n", + "\n", + "with torch.no_grad():\n", + " out_pooled = model_pooled({\"image\": batch[\"image\"]})\n", + "\n", + "info_seq = model_seq.get_output_info(\"image\")\n", + "info_pooled = model_pooled.get_output_info(\"image\")\n", + "\n", + "print(\"With mean pooling (pool='mean'):\")\n", + "print(f\" Output shape: {tuple(out_pooled['image'].shape)} # (B, 1, E)\")\n", + "print(f\" num_tokens: {info_pooled['num_tokens']}\")\n", + "print()" + ] + }, { "cell_type": "markdown", "id": "4ce223e3", @@ -401,49 +357,10 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "344eac50", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionEmbeddingModel Backbone Comparison\n", - "==================================================\n", - "\n", - "patch:\n", - " Tokens: 197 (196 patches + CLS)\n", - " Parameters: 58,240\n", - "\n", - "cnn:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 231,808\n", - "\n", - "resnet18:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 11,242,432\n", - "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to C:\\Users\\637682/.cache\\torch\\hub\\checkpoints\\resnet50-11ad3fa6.pth\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "100%|██████████| 97.8M/97.8M [00:22<00:00, 4.63MB/s]\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "resnet50:\n", - " Tokens: 50 (49 patches + CLS)\n", - " Parameters: 23,770,560\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.models import VisionEmbeddingModel\n", "\n", @@ -476,18 +393,7 @@ "execution_count": null, "id": "91b5f6a6", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Feature keys: ['image']\n", - "Label key: label\n", - "Mode: binary\n", - "Total parameters: 11,259,329\n" - ] - } - ], + "outputs": [], "source": [ "class VisionClassifier(nn.Module):\n", " \"\"\"End-to-end classifier using VisionEmbeddingModel.\"\"\"\n", @@ -575,30 +481,10 @@ }, { "cell_type": "code", - "execution_count": 25, + "execution_count": null, "id": "699c7e03", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "VisionClassifier(\n", - " (vision_encoder): VisionEmbeddingModel(backbone='resnet18', embedding_dim=128, fields=['image'])\n", - " (classifier): Sequential(\n", - " (0): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n", - " (1): Linear(in_features=128, out_features=128, bias=True)\n", - " (2): GELU(approximate='none')\n", - " (3): Dropout(p=0.1, inplace=False)\n", - " (4): Linear(in_features=128, out_features=1, bias=True)\n", - " )\n", - ")\n", - "Metrics: ['roc_auc']\n", - "Device: cpu\n", - "\n" - ] - } - ], + "outputs": [], "source": [ "from pyhealth.trainer import Trainer\n", "\n", @@ -628,259 +514,10 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "id": "b12a766e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Training:\n", - "Batch size: 16\n", - "Optimizer: \n", - "Optimizer params: {'lr': 0.0001}\n", - "Weight decay: 0.0\n", - "Max grad norm: 1.0\n", - "Val dataloader: \n", - "Monitor: roc_auc\n", - "Monitor criterion: max\n", - "Epochs: 5\n", - "Patience: None\n", - "\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "3cadbf78b48241199f56d04123c7adb0", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Epoch 0 / 5: 0%| | 0/9 [00:00 Date: Sun, 24 May 2026 18:52:00 -0700 Subject: [PATCH 23/30] Update vision_embedding_tutorial.ipynb --- examples/vision_embedding_tutorial.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/vision_embedding_tutorial.ipynb b/examples/vision_embedding_tutorial.ipynb index 3d62b68bf..d241358c8 100644 --- a/examples/vision_embedding_tutorial.ipynb +++ b/examples/vision_embedding_tutorial.ipynb @@ -309,7 +309,8 @@ "Pass `pool='mean'` to collapse all patch tokens into a single averaged vector per image:\n", "\n", "```\n", - "With pool='mean': (B, 1, E) → e.g., (16, 1, 128)\n", + "Input: (B, C, H, W)\n", + "Output with pool='mean': (B, 1, E)\n", "```" ] }, From 0a67c764f0891dc947b2b8698ef73c84c7fe59d0 Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 3 Jun 2026 01:36:08 -0700 Subject: [PATCH 24/30] Add suggestions to Rian's PR --- .../unified_embedding_e2e_mimic4.py | 12 +++++------- pyhealth/tasks/__init__.py | 2 +- pyhealth/tasks/multimodal_mimic4.py | 8 +++++--- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index e520f181d..5447d7077 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -72,7 +72,7 @@ from pyhealth.tasks.multimodal_mimic4 import ( ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, - LabsOnlyMIMIC4, + LabsMIMIC4, NotesLabsMIMIC4, ) from pyhealth.trainer import Trainer @@ -105,10 +105,8 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if "chartevents" not in ehr_tables: ehr_tables.append("chartevents") - if args.task == "labs_only": - # Pure EHR baseline: only labevents, no notes, no ICD codes. + if args.task == "labs": ehr_tables = ["labevents"] - note_tables = None return MIMIC4Dataset( ehr_root=args.ehr_root, @@ -134,8 +132,8 @@ def _build_task(args: argparse.Namespace): include_icd=args.icd_codes, include_vitals=args.include_vitals, ) - if args.task == "labs_only": - return LabsOnlyMIMIC4(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}") @@ -387,7 +385,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, - choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs_only"], + choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs"], default="stagenet", help=( "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index 569b035f0..f026289bf 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -50,7 +50,7 @@ ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, - LabsOnlyMIMIC4, + LabsMIMIC4, NotesLabsMIMIC4, ) from .patient_linkage import patient_linkage_mimic3_fn diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 56ce98f69..61f72a69e 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -1079,7 +1079,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: return [record] -class LabsOnlyMIMIC4(BaseMultimodalMIMIC4Task): +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. @@ -1096,7 +1096,7 @@ class LabsOnlyMIMIC4(BaseMultimodalMIMIC4Task): PADDING: int = 0 - task_name: str = "LabsOnlyMIMIC4" + task_name: str = "LabsMIMIC4" input_schema: ClassVar[Dict] = { "labs": ("stagenet_tensor", {}), @@ -1132,7 +1132,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) lab_times.append(self.MISSING_FLOAT_TOKEN) - yield { + single_patient_longitudinal_record = { "patient_id": patient.patient_id, "labs": (lab_times, lab_values), "labs_mask": (lab_times, lab_masks), @@ -1140,3 +1140,5 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri "window_start": effective_start, "window_end": effective_end, } + + return [single_patient_longitudinal_record] From 11f04a3b7a73f40823ae0df5ae7a0836d4cbcd3b Mon Sep 17 00:00:00 2001 From: Pang Date: Wed, 3 Jun 2026 19:14:59 -0500 Subject: [PATCH 25/30] Minor change --- examples/mortality_prediction/multimodal_mimic4.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/mortality_prediction/multimodal_mimic4.py b/examples/mortality_prediction/multimodal_mimic4.py index a4c9d6e1c..eb16b005c 100644 --- a/examples/mortality_prediction/multimodal_mimic4.py +++ b/examples/mortality_prediction/multimodal_mimic4.py @@ -18,7 +18,7 @@ TASK = "ClinicalNotesICDLabsMIMIC4" # Options: ClinicalNotesMIMIC4, ICDLabsMIMIC4, ClinicalNotesICDLabsMIMIC4, ClinicalNotesICDLabsCXRMIMIC4 # The idea here is that we want additive tasks so we can evaluate the value in adding more modalities DEV_MODE = True -ENVIRONMENT = "SunLabCluster" # Either 'Local' or 'Cluster' or "SunLabCluster" +ENVIRONMENT = "SunLabCluster" # Either 'Local' or 'CampusCluster' or "SunLabCluster" NETID = "wp14" # For personal cache if ENVIRONMENT == "Local": @@ -38,7 +38,7 @@ cache_dir = os.path.join( pyhealth_repo_root, "local_data/local/data/wp/pyhealth_cache" ) -elif ENVIRONMENT == "Cluster": +elif ENVIRONMENT == "CampusCluster": ehr_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2" note_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note" From afda6eb7d6603f44358e5abbdeb8199edbdacf26 Mon Sep 17 00:00:00 2001 From: Pang Date: Wed, 3 Jun 2026 19:21:39 -0500 Subject: [PATCH 26/30] Testing main branch protection rules --- examples/mortality_prediction/multimodal_mimic4.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/mortality_prediction/multimodal_mimic4.py b/examples/mortality_prediction/multimodal_mimic4.py index eb16b005c..4261c25b4 100644 --- a/examples/mortality_prediction/multimodal_mimic4.py +++ b/examples/mortality_prediction/multimodal_mimic4.py @@ -18,7 +18,7 @@ TASK = "ClinicalNotesICDLabsMIMIC4" # Options: ClinicalNotesMIMIC4, ICDLabsMIMIC4, ClinicalNotesICDLabsMIMIC4, ClinicalNotesICDLabsCXRMIMIC4 # The idea here is that we want additive tasks so we can evaluate the value in adding more modalities DEV_MODE = True -ENVIRONMENT = "SunLabCluster" # Either 'Local' or 'CampusCluster' or "SunLabCluster" +ENVIRONMENT = "CampusCluster" # Either 'Local' or 'CampusCluster' or "SunLabCluster" NETID = "wp14" # For personal cache if ENVIRONMENT == "Local": From 4e5824ab966aa4b4ec56cf179c3c76ed69ce3a10 Mon Sep 17 00:00:00 2001 From: Pang Date: Tue, 16 Jun 2026 18:26:17 -0500 Subject: [PATCH 27/30] Added updates --- .../unified_embedding_e2e_mimic4.py | 2 +- .../multimodal_mimic4_task_tutorial.ipynb | 288 ++++++++++++++++-- pyhealth/tasks/multimodal_mimic4.py | 65 ++-- scripts/will/run_labs.py | 116 +++++++ 4 files changed, 425 insertions(+), 46 deletions(-) create mode 100644 scripts/will/run_labs.py diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 5447d7077..b57df5a8b 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -316,7 +316,7 @@ def run(args: argparse.Namespace) -> Path: trainer = Trainer( model=model, - metrics=["pr_auc", "roc_auc", "f1", "f1_opt", "accuracy"], + metrics=["pr_auc", "roc_auc", "f1", "accuracy"], device=args.device, enable_logging=True, output_path=str(output_dir), diff --git a/examples/multimodal_mimic4_task_tutorial.ipynb b/examples/multimodal_mimic4_task_tutorial.ipynb index 3fda88002..77763540c 100644 --- a/examples/multimodal_mimic4_task_tutorial.ipynb +++ b/examples/multimodal_mimic4_task_tutorial.ipynb @@ -25,7 +25,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "id": "8477e2ac-322f-49e9-819f-ce7621b214c6", "metadata": {}, "outputs": [], @@ -33,6 +33,7 @@ "from datetime import datetime\n", "from typing import Any, Dict, List, Optional\n", "import os\n", + "from pathlib import Path\n", "import tempfile\n", "import shutil \n", "import random\n", @@ -50,7 +51,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "id": "201c3a62-2a2e-47ab-9ab8-eea77812f349", "metadata": {}, "outputs": [], @@ -69,12 +70,23 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 3, "id": "9a466c4b-2283-443f-94a3-5768babddf50", "metadata": {}, - "outputs": [], + "outputs": [ + { + "data": { + "text/plain": [ + "" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], "source": [ - "SEED = 42\n", + "SEED = 30\n", "random.seed(SEED)\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)" @@ -90,7 +102,7 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 4, "id": "05ec2dc9-1a0c-4360-a028-d93173dcba65", "metadata": {}, "outputs": [], @@ -121,7 +133,30 @@ " print(f\" hadm_id: {note.hadm_id}\")\n", " print(f\" charttime: {note.timestamp}\")\n", " print(f\" storetime: {note.storetime}\")\n", - " print(f\" text: {note.text[:char_limit]}..(Limited to {char_limit} Characters).\")" + " print(f\" text: {note.text[:char_limit]}..(Limited to {char_limit} Characters).\")\n", + "\n", + "def print_patient_lab_info(sample, event_limit=5):\n", + " patient = dataset.get_patient(sample['patient_id'])\n", + " lab_events = patient.get_events(\n", + " event_type=\"labevents\",\n", + " start=sample['window_start'],\n", + " end=sample['window_end'],\n", + " )\n", + " print(f\"\\nlab events: {len(lab_events)}\")\n", + " for event in lab_events[:event_limit]:\n", + " print(f\" itemid: {event.itemid}\")\n", + " print(f\" charttime: {event.timestamp}\")\n", + " print(f\" storetime: {event.storetime}\")\n", + " print(f\" valuenum: {event.valuenum}\")\n", + " if len(lab_events) > event_limit:\n", + " print(f\" ... ({len(lab_events) - event_limit} more events)\")\n", + "\n", + "def get_project_root(marker=\"PyHealth\"):\n", + " path = Path(os.getcwd())\n", + " for parent in [path, *path.parents]:\n", + " if parent.name == marker:\n", + " return str(parent)\n", + " raise ValueError(f\"'{marker}' not found in path\")" ] }, { @@ -139,12 +174,12 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 5, "id": "tv8ir9mp9t", "metadata": {}, "outputs": [], "source": [ - "PYHEALTH_REPO_ROOT = '/Users/wpang/Desktop/PyHealth'\n", + "PYHEALTH_REPO_ROOT = get_project_root()\n", "\n", "EHR_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n", "NOTE_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n", @@ -153,10 +188,43 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 6, "id": "2f0f9ac3-1e3e-438b-b463-32e839941ce0", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Memory usage Starting MIMIC4Dataset init: 576.3 MB\n", + "Initializing mimic4 dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo|/u/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132\n", + "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents'] (dev mode: False)\n", + "Using default EHR config: /u/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n", + "Memory usage Before initializing mimic4_ehr: 576.3 MB\n", + "Initializing mimic4_ehr dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/1756cd94-34cf-553f-8bbf-732cb4253291\n", + "Memory usage After initializing mimic4_ehr: 576.3 MB\n", + "Memory usage After EHR dataset initialization: 576.3 MB\n", + "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n", + "Using default note config: /u/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n", + "Memory usage Before initializing mimic4_note: 576.3 MB\n", + "Initializing mimic4_note dataset from /u/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n", + "Using provided cache_dir: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/7ddd05e8-68d1-57f8-a246-264b9bfb6f49\n", + "Memory usage After initializing mimic4_note: 576.3 MB\n", + "Memory usage After Note dataset initialization: 576.3 MB\n", + "Memory usage Completed MIMIC4Dataset init: 576.3 MB\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/u/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n", + " warnings.warn(\n" + ] + } + ], "source": [ "dataset = MIMIC4Dataset(\n", " ehr_root=EHR_ROOT,\n", @@ -187,10 +255,49 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 7, "id": "d3b027ab-c814-49cd-8542-885d76b2401b", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n", + "Task cache paths: task_df=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/task_df.ld, samples=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n", + "Applying task transformations on data with 1 workers...\n", + "No cached event dataframe found. Creating: /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/global_event_df.parquet\n", + "Combining data from ehr dataset\n", + "Scanning table: diagnoses_icd from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: procedures_icd from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: labevents from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n", + "Joining with table: /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n", + "Scanning table: patients from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n", + "Scanning table: admissions from /u/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n", + "Scanning table: icustays from /u/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n", + "Combining data from note dataset\n", + "Scanning table: discharge from /u/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n", + "Scanning table: radiology from /u/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n", + "Creating combined dataframe\n", + "Caching event dataframe to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/global_event_df.parquet...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 patients. (Polars threads: 6)\n", + "Worker 0 finished processing patients.\n", + "Fitting processors on the dataset...\n", + "Label mortality vocab: {0: 0, 1: 1}\n", + "Processing samples and saving to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n", + "Applying processors on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 samples. (0 to 13)\n", + "Worker 0 finished processing samples.\n", + "Cached processed samples to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_493ab707-d120-581c-a795-319ffd5c468e/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n" + ] + } + ], "source": [ "%%capture\n", "task = ClinicalNotesICDLabsMIMIC4()\n", @@ -199,10 +306,89 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 8, "id": "6b70acb2-bdea-4a43-8b9a-7725c34fcfce", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Patient ID: 10001\n", + "Mortality Flag: tensor([0.])\n", + "Found 13 unique patient IDs\n", + "\n", + "num_admissions: 3\n", + " hadm_id: 19999\n", + " admittime: 2150-02-15 08:00:00\n", + " dischtime: 2150-02-18 14:00:00\n", + " hadm_id: 20001\n", + " admittime: 2150-03-15 08:00:00\n", + " dischtime: 2150-03-18 14:00:00\n", + " hadm_id: 20002\n", + " admittime: 2150-06-20 10:30:00\n", + " dischtime: 2150-06-25 16:00:00\n", + "\n", + "discharge notes: 3\n", + " note_id: d1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-18 12:00:00\n", + " storetime: 2150-02-18 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + " note_id: d3\n", + " hadm_id: 20001\n", + " charttime: 2150-03-18 12:00:00\n", + " storetime: 2150-03-18 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + " note_id: d4\n", + " hadm_id: 20002\n", + " charttime: 2150-06-25 12:00:00\n", + " storetime: 2150-06-25 13:00:00\n", + " text: Patient 10001 admitted for evaluation. Vitals stable. Labs reviewed. Discharged ..(Limited to 80 Characters).\n", + "\n", + "radiology notes: 3\n", + " note_id: r1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-15 14:00:00\n", + " storetime: 2150-02-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + " note_id: r3\n", + " hadm_id: 20001\n", + " charttime: 2150-03-15 14:00:00\n", + " storetime: 2150-03-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + " note_id: r4\n", + " hadm_id: 20002\n", + " charttime: 2150-06-20 14:00:00\n", + " storetime: 2150-06-20 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + "\n", + "lab events: 12\n", + " itemid: 50806\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 102.0\n", + " itemid: 50809\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 285.0\n", + " itemid: 50803\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 18.0\n", + " itemid: 50868\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 16.0\n", + " itemid: 50822\n", + " charttime: 2150-03-15 10:00:00\n", + " storetime: 2150-03-15 10:30:00\n", + " valuenum: 4.8\n", + " ... (7 more events)\n" + ] + } + ], "source": [ "random_patient_record = np.random.randint(0, len(samples)-1)\n", "sample = samples[random_patient_record]\n", @@ -218,7 +404,10 @@ "print_patient_note_info(sample, note_type=\"discharge\")\n", "\n", "# Radiology Notes\n", - "print_patient_note_info(sample, note_type=\"radiology\")" + "print_patient_note_info(sample, note_type=\"radiology\")\n", + "\n", + "# Labs\n", + "print_patient_lab_info(sample)" ] }, { @@ -231,10 +420,33 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 9, "id": "16121916-e52e-4f66-8463-ff0f931cf416", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n", + "Task cache paths: task_df=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/task_df.ld, samples=/tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n", + "Applying task transformations on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 patients. (Polars threads: 6)\n", + "Worker 0 finished processing patients.\n", + "Fitting processors on the dataset...\n", + "Label mortality vocab: {0: 0, 1: 1}\n", + "Processing samples and saving to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n", + "Applying processors on data with 1 workers...\n", + "Detected Jupyter notebook environment, setting num_workers to 1\n", + "Single worker mode, processing sequentially\n", + "Worker 0 started processing 13 samples. (0 to 13)\n", + "Worker 0 finished processing samples.\n", + "Cached processed samples to /tmp/tmp6vbrmqeu/db0f934a-30cc-5e3e-b373-0057d5a4b132/tasks/ClinicalNotesICDLabsMIMIC4_6c674c02-3635-50c0-bdde-3a45cb1e6f31/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n" + ] + } + ], "source": [ "%%capture\n", "task = ClinicalNotesICDLabsMIMIC4(window_hours=12)\n", @@ -245,10 +457,38 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 10, "id": "42df85ab-2649-497b-92a2-415dc169798e", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "Patient ID: 10001\n", + "Mortality Flag: tensor([0.])\n", + "Window start: 2150-02-15 08:00:00\n", + "Window end: 2150-02-15 20:00:00\n", + "\n", + "num_admissions: 1\n", + " hadm_id: 19999\n", + " admittime: 2150-02-15 08:00:00\n", + " dischtime: 2150-02-18 14:00:00\n", + "\n", + "discharge notes: 0\n", + "\n", + "radiology notes: 1\n", + " note_id: r1\n", + " hadm_id: 19999\n", + " charttime: 2150-02-15 14:00:00\n", + " storetime: 2150-02-15 15:00:00\n", + " text: Chest X-ray for patient 10001. No acute cardiopulmonary process identified. Hear..(Limited to 80 Characters).\n", + "\n", + "lab events: 0\n" + ] + } + ], "source": [ "print(f\"\\nPatient ID: {sample_patient_id}\")\n", "print(f\"Mortality Flag: {sample['mortality']}\")\n", @@ -262,7 +502,10 @@ "print_patient_note_info(sample, note_type=\"discharge\")\n", "\n", "# Radiology Notes\n", - "print_patient_note_info(sample, note_type=\"radiology\")" + "print_patient_note_info(sample, note_type=\"radiology\")\n", + "\n", + "# Labs\n", + "print_patient_lab_info(sample)" ] }, { @@ -307,7 +550,8 @@ "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", - "pygments_lexer": "ipython3" + "pygments_lexer": "ipython3", + "version": "3.12.13" } }, "nbformat": 4, diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 61f72a69e..24f7db005 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -1109,36 +1109,55 @@ def __init__(self, window_hours: Optional[float] = 24) -> None: self.window_hours = window_hours def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[override] - admissions = self._build_admissions_to_process(patient) - if not admissions: + admissions_to_process, mortality_label = self._build_admissions_to_process( + patient + ) + if not admissions_to_process: return [] - for admission_time, admission_dischtime, admission in admissions: - mortality_label = self._get_mortality_label(patient, admission) - effective_start, effective_end = self._compute_window( - admission_time, admission_dischtime, self.window_hours - ) + 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=effective_end, + end_time=admission_dischtime, ) + all_lab_times.extend(lab_times) + all_lab_values.extend(lab_values) + all_lab_masks.extend(lab_masks) - if len(lab_values) == 0: - 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(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": (lab_times, lab_values), - "labs_mask": (lab_times, lab_masks), - "mortality": mortality_label, - "window_start": effective_start, - "window_end": effective_end, - } + 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] + return [single_patient_longitudinal_record] diff --git a/scripts/will/run_labs.py b/scripts/will/run_labs.py new file mode 100644 index 000000000..2bc674b49 --- /dev/null +++ b/scripts/will/run_labs.py @@ -0,0 +1,116 @@ +project_dir = "/u/wp14/PyHealth" +seed = 12 +account = "jimeng-ic" +partition = "eng-research-gpu" +time = "24:00:00" +mem = "64G" +gres = "gpu:A10:1" +conda_env = "pyhealth2" +ehr_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2" +cache_dir = "/u/wp14/pyhealth_cache" +output_dir = "output/rnn_labs" +embedding_dim = 128 +hidden_dim = 128 +rnn_type = "GRU" +rnn_layers = 2 +dropout = 0.1 +epochs = 5 +batch_size = 32 +lr = 1e-3 +weight_decay = 1e-5 +patience = 5 +num_workers = 4 +dev = False + +# ── Step 0: Clean logs and cache ──────────────────────────────────────────── +print(f""" +### STEP 0 (optional): Clean logs and cache + +rm -rf {project_dir}/logs/* +rm -rf {cache_dir}/* +""") + +# ── Step 1: Reserve resources ───────────────────────────────────────────────── +if dev: + print(f""" +### STEP 1: Reserve an interactive node + +srun \\ + --account={account} \\ + --partition={partition} \\ + --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\ + --mem={mem} --gres={gres} --time={time} \\ + --pty bash +""") +else: + print(f""" +### STEP 1: Submit batch job (includes step 2 automatically) + +mkdir -p {project_dir}/logs && sbatch \\ + --job-name=rnn_labs_s{seed} \\ + --account={account} \\ + --partition={partition} \\ + --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\ + --mem={mem} --gres={gres} --time={time} \\ + --output={project_dir}/logs/rnn_labs_s{seed}_%j.out \\ + --error={project_dir}/logs/rnn_labs_s{seed}_%j.err \\ + --wrap=' + module load miniconda3/24.9.2 && + eval "$(conda shell.bash hook)" && + conda activate {conda_env} && + cd {project_dir} && + export PYTHONPATH={project_dir}:$PYTHONPATH && + python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --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} + ' +""") + +# ── Step 2: Run (only needed for dev/interactive) ───────────────────────────── +print("\n" + "=" * 60 + "\n") +if dev: + print(f""" +### STEP 2: Once on the node, run + +module load miniconda3/24.9.2 && +eval "$(conda shell.bash hook)" && +conda activate {conda_env} && +cd {project_dir} && +export PYTHONPATH={project_dir}:$PYTHONPATH && +mkdir -p {project_dir}/logs/dev && +python {project_dir}/examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\ + --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} --dev \\ + > >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.out) \\ + 2> >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.err >&2) +""") From 45b7b6b097a27b290e3cc408816068435f5dc988 Mon Sep 17 00:00:00 2001 From: William Pang Date: Tue, 16 Jun 2026 18:42:18 -0500 Subject: [PATCH 28/30] Change File Name --- scripts/will/{run_labs.py => slurm_run_labs_only_variant.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename scripts/will/{run_labs.py => slurm_run_labs_only_variant.py} (100%) diff --git a/scripts/will/run_labs.py b/scripts/will/slurm_run_labs_only_variant.py similarity index 100% rename from scripts/will/run_labs.py rename to scripts/will/slurm_run_labs_only_variant.py From 5eb65f32d3710809cee353152a34caf1a21993a3 Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 17 Jun 2026 14:57:56 -0500 Subject: [PATCH 29/30] Clean markdown --- .../unified_embedding_e2e_mimic4.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 9275a90d4..a572dd1b6 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -73,16 +73,12 @@ from pyhealth.models.ehrmamba import EHRMamba from pyhealth.models.jamba_ehr import JambaEHR from pyhealth.tasks import MortalityPredictionStageNetMIMIC4 -<<<<<<< HEAD from pyhealth.tasks.multimodal_mimic4 import ( ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4, LabsMIMIC4, NotesLabsMIMIC4, ) -======= -from pyhealth.tasks.multimodal_mimic4 import ClinicalNotesICDLabsMIMIC4, ICDLabsMIMIC4 ->>>>>>> origin/main from pyhealth.trainer import Trainer from pyhealth.utils import set_seed @@ -99,7 +95,6 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if args.task == "icd_labs": ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] -<<<<<<< HEAD if args.task == "notes_labs": if not args.note_root: raise ValueError("--task notes_labs requires --note-root.") @@ -117,8 +112,6 @@ def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset: if args.task == "labs": ehr_tables = ["labevents"] -======= ->>>>>>> origin/main return MIMIC4Dataset( ehr_root=args.ehr_root, ehr_tables=ehr_tables, @@ -396,11 +389,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--task", type=str, -<<<<<<< HEAD choices=["stagenet", "icd_labs", "clinical_notes_icd_labs", "notes_labs", "labs"], -======= - choices=["icd_labs", "clinical_notes_icd_labs"], ->>>>>>> origin/main default="stagenet", help=( "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. " From 0542dee53cb1cfbda0429979076799293174e734 Mon Sep 17 00:00:00 2001 From: William Pang Date: Wed, 17 Jun 2026 18:48:13 -0500 Subject: [PATCH 30/30] Updates to markdown --- pyhealth/models/embedding/unified.py | 14 +- pyhealth/tasks/__init__.py | 6 +- pyhealth/tasks/multimodal_mimic4.py | 495 ++++++++++++--------------- 3 files changed, 226 insertions(+), 289 deletions(-) diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 787a4183f..703f77630 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -344,16 +344,6 @@ def _set_projection( if processor.is_token(): from transformers import AutoModel -<<<<<<< HEAD - 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) -======= tokenizer_model = getattr(processor, "tokenizer_model", None) if not tokenizer_model: raise ValueError( @@ -369,12 +359,14 @@ def _set_projection( shared_encoder = self.encoders[shared_field] else: shared_encoder = AutoModel.from_pretrained(tokenizer_model) + if freeze: + for p in shared_encoder.parameters(): + p.requires_grad = False self._shared_text_field_by_model[tokenizer_model] = field_name self.encoders[field_name] = shared_encoder hidden = shared_encoder.config.hidden_size _set_projection(hidden) ->>>>>>> origin/main else: raise ValueError( f"TEXT processor '{field_name}' must either supply a pre-built " diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py index ecd5cd1bb..e98d98db6 100644 --- a/pyhealth/tasks/__init__.py +++ b/pyhealth/tasks/__init__.py @@ -47,6 +47,7 @@ from .multimodal_mimic4 import ( ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, + ClinicalNotesICDLabsCXRMIMIC4, ICDLabsMIMIC4, LabsMIMIC4, NotesLabsMIMIC4, @@ -72,9 +73,4 @@ MutationPathogenicityPrediction, VariantClassificationClinVar, ) -from .multimodal_mimic4 import ( - ClinicalNotesMIMIC4, - ClinicalNotesICDLabsMIMIC4, - ClinicalNotesICDLabsCXRMIMIC4, -) from .patient_linkage_mimic3 import PatientLinkageMIMIC3Task diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py index 418177078..71f8d3f87 100644 --- a/pyhealth/tasks/multimodal_mimic4.py +++ b/pyhealth/tasks/multimodal_mimic4.py @@ -497,21 +497,21 @@ class ClinicalNotesMIMIC4(BaseMultimodalMIMIC4Task): task_name: str = "ClinicalNotesMIMIC4" input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "radiology_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ) - } + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + } output_schema: Dict[str, str] = {"mortality": "binary"} def __call__(self, patient: Any) -> List[Dict[str, Any]]: @@ -574,17 +574,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_radiology_times_from_admission, ) - if len(all_discharge_texts) == 0: - discharge_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - if len(all_radiology_texts) == 0: - radiology_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - single_patient_longitudinal_record = { "patient_id": patient.patient_id, "discharge_note_times": discharge_note_times_from_admission, @@ -634,37 +623,23 @@ class ClinicalNotesICDLabsMIMIC4(BaseMultimodalMIMIC4Task): task_name: str = "ClinicalNotesICDLabsMIMIC4" input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "radiology_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT", - "type_tag": "note", - }, - ), - "icd_codes": ("stagenet", {"padding": PADDING}), - "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), - } - output_schema: Dict[str, str] = {"mortality": "binary"} - - 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"], + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "icd_codes": ("stagenet", {"padding": PADDING}), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), } output_schema: Dict[str, str] = {"mortality": "binary"} @@ -711,12 +686,10 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): - # Skip malformed admissions; note collectors are not called for these. - continue + admission_dischtime = admission_time if admission_dischtime < admission_time: - # Guard against invalid chronology in source records. - continue + admission_dischtime = admission_time discharge_texts, discharge_times = self._collect_notes( patient, @@ -752,7 +725,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) else: # Add missingness token if there are no ICD diagnosis/inpatient procedure codes - all_icd_codes.append([self.MISSING_CODE_TOKEN]) + all_icd_codes.append([self.MISSING_TEXT_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time @@ -774,12 +747,195 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) all_lab_times.append(self.MISSING_FLOAT_TOKEN) - # If all admissions were skipped before ICD collection, ensure a - # single placeholder step so StageNetProcessor does not emit None time. + discharge_note_times_from_admission = ( + all_discharge_texts, + all_discharge_times_from_admission, + ) + radiology_note_times_from_admission = ( + all_radiology_texts, + all_radiology_times_from_admission, + ) + + single_patient_longitudinal_record = { + "patient_id": patient.patient_id, + "discharge_note_times": discharge_note_times_from_admission, + "radiology_note_times": radiology_note_times_from_admission, + "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 ClinicalNotesICDLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): + """Task combining notes, ICD, labs, and CXR for MIMIC-IV mortality. + + Adds temporally filtered CXR signals on top of ``ClinicalNotesICDLabsMIMIC4``: + - ``cxr_image_times``: ``(image_paths, hours_from_admission)`` processed by + ``TimeImageProcessor``. + + CXR filtering uses event timestamps from the metadata table. Since + timestamps are built from StudyDate+StudyTime in dataset configs, studytime + is naturally respected in temporal windows. + """ + + PADDING: int = 0 + + task_name: str = "ClinicalNotesICDLabsCXRMIMIC4" + input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { + "discharge_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "radiology_note_times": ( + "tuple_time_text", + { + "tokenizer_model": "bert-base-uncased", + "type_tag": "note", + }, + ), + "icd_codes": ("stagenet", {"padding": PADDING}), + "labs": ("stagenet_tensor", {}), + "labs_mask": ("stagenet_tensor", {}), + "cxr_image_times": ( + "time_image", + { + "image_size": 224, + "mode": "RGB", + "padding": "", + }, + ), + } + 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_discharge_texts: List[str] = [] + all_discharge_times_from_admission: List[float] = [] + all_radiology_texts: List[str] = [] + all_radiology_times_from_admission: List[float] = [] + 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] = [] + all_cxr_paths: List[str] = [] + all_cxr_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 + + discharge_texts, discharge_times = self._collect_notes( + patient, + "discharge", + admission.hadm_id, + admission_time, + start_time=effective_start, + end_time=effective_end, + ) + all_discharge_texts.extend(discharge_texts) + all_discharge_times_from_admission.extend(discharge_times) + + radiology_texts, radiology_times = self._collect_notes( + patient, + "radiology", + admission.hadm_id, + admission_time, + start_time=effective_start, + end_time=effective_end, + ) + all_radiology_texts.extend(radiology_texts) + all_radiology_times_from_admission.extend(radiology_times) + + 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_CODE_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) + + # CXR metadata is filtered by timestamp; this includes StudyTime. + metadata_events = patient.get_events( + event_type="metadata", + start=admission_time, + end=admission_dischtime, + ) + 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 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_CODE_TOKEN]) all_icd_times.append(self.MISSING_FLOAT_TOKEN) + if len(all_cxr_paths) == 0: + all_cxr_paths = [self.MISSING_TEXT_TOKEN] + all_cxr_times = [self.MISSING_FLOAT_TOKEN] + discharge_note_times_from_admission = ( all_discharge_texts, all_discharge_times_from_admission, @@ -789,9 +945,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_radiology_times_from_admission, ) - # Per-admission note fallback happens inside _collect_notes(). - # This final guard handles the edge case where every admission was - # skipped before _collect_notes() was reached. if len(all_discharge_texts) == 0: discharge_note_times_from_admission = ( [self.MISSING_TEXT_TOKEN], @@ -810,6 +963,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: "icd_codes": (all_icd_times, all_icd_codes), "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, @@ -818,6 +972,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: return [single_patient_longitudinal_record] + class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task): """Task for ICD codes + lab values mortality prediction using MIMIC-IV. @@ -880,16 +1035,9 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): -<<<<<<< HEAD admission_dischtime = admission_time if admission_dischtime < admission_time: admission_dischtime = admission_time -======= - continue - - if admission_dischtime < admission_time: - continue ->>>>>>> origin/main visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id) if visit_icd_codes: @@ -902,11 +1050,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_icd_codes.append(visit_icd_codes) all_icd_times.append(time_from_previous) else: -<<<<<<< HEAD all_icd_codes.append([self.MISSING_TEXT_TOKEN]) -======= - all_icd_codes.append([self.MISSING_CODE_TOKEN]) ->>>>>>> origin/main all_icd_times.append(self.MISSING_FLOAT_TOKEN) previous_admission_time = admission_time @@ -928,11 +1072,7 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_lab_times.append(self.MISSING_FLOAT_TOKEN) if len(all_icd_codes) == 0: -<<<<<<< HEAD all_icd_codes.append([self.MISSING_TEXT_TOKEN]) -======= - all_icd_codes.append([self.MISSING_CODE_TOKEN]) ->>>>>>> origin/main all_icd_times.append(self.MISSING_FLOAT_TOKEN) single_patient_longitudinal_record = { @@ -948,7 +1088,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: return [single_patient_longitudinal_record] -<<<<<<< HEAD class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): """Mortality prediction from admission-context notes and lab values. @@ -980,23 +1119,10 @@ class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task): include_vitals: When ``True``, collect ICU vital signs from chartevents and add ``vitals`` and ``vitals_mask`` to the sample dict / input schema. Default: ``False``. -======= -class ClinicalNotesICDLabsCXRMIMIC4(BaseMultimodalMIMIC4Task): - """Task combining notes, ICD, labs, and CXR for MIMIC-IV mortality. - - Adds temporally filtered CXR signals on top of ``ClinicalNotesICDLabsMIMIC4``: - - ``cxr_image_times``: ``(image_paths, hours_from_admission)`` processed by - ``TimeImageProcessor``. - - CXR filtering uses event timestamps from the metadata table. Since - timestamps are built from StudyDate+StudyTime in dataset configs, studytime - is naturally respected in temporal windows. ->>>>>>> origin/main """ PADDING: int = 0 -<<<<<<< HEAD task_name: str = "NotesLabsMIMIC4" _BASE_INPUT_SCHEMA: ClassVar[Dict] = { @@ -1033,58 +1159,18 @@ def __init__( def __call__(self, patient: Any) -> List[Dict[str, Any]]: if not patient.get_events(event_type="patients"): -======= - task_name: str = "ClinicalNotesICDLabsCXRMIMIC4" - input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = { - "discharge_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "bert-base-uncased", - "type_tag": "note", - }, - ), - "radiology_note_times": ( - "tuple_time_text", - { - "tokenizer_model": "bert-base-uncased", - "type_tag": "note", - }, - ), - "icd_codes": ("stagenet", {"padding": PADDING}), - "labs": ("stagenet_tensor", {}), - "labs_mask": ("stagenet_tensor", {}), - "cxr_image_times": ( - "time_image", - { - "image_size": 224, - "mode": "RGB", - "padding": "", - }, - ), - } - 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: ->>>>>>> origin/main return [] admissions_to_process, mortality_label = self._build_admissions_to_process( patient ) -<<<<<<< HEAD if not admissions_to_process: -======= - if len(admissions_to_process) == 0: ->>>>>>> origin/main return [] effective_start, effective_end = self._compute_effective_window( admissions_to_process ) -<<<<<<< HEAD all_note_texts: List[str] = [] all_note_times: List[float] = [] all_lab_values: List[List[float]] = [] @@ -1095,37 +1181,16 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: all_vital_times: List[float] = [] all_icd_codes: List[List[str]] = [] all_icd_times: List[float] = [] -======= - all_discharge_texts: List[str] = [] - all_discharge_times_from_admission: List[float] = [] - all_radiology_texts: List[str] = [] - all_radiology_times_from_admission: List[float] = [] - 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] = [] - all_cxr_paths: List[str] = [] - all_cxr_times: List[float] = [] ->>>>>>> origin/main previous_admission_time = None for admission in admissions_to_process: admission_time = admission.timestamp -<<<<<<< HEAD -======= - # 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 - ->>>>>>> origin/main try: admission_dischtime = datetime.strptime( admission.dischtime, "%Y-%m-%d %H:%M:%S" ) except (ValueError, AttributeError): -<<<<<<< HEAD admission_dischtime = admission_time if admission_dischtime < admission_time: admission_dischtime = admission_time @@ -1272,56 +1337,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri admission_dischtime = admission_time if admission_dischtime < admission_time: admission_dischtime = admission_time -======= - # Skip malformed admissions; note collectors are not called for these. - continue - - if admission_dischtime < admission_time: - # Guard against invalid chronology in source records. - continue - - admission_end = admission_dischtime - if effective_end is not None and effective_end < admission_end: - admission_end = effective_end - - discharge_texts, discharge_times = self._collect_notes( - patient, - "discharge", - admission.hadm_id, - admission_time, - start_time=effective_start, - end_time=admission_end, - ) - all_discharge_texts.extend(discharge_texts) - all_discharge_times_from_admission.extend(discharge_times) - - radiology_texts, radiology_times = self._collect_notes( - patient, - "radiology", - admission.hadm_id, - admission_time, - start_time=effective_start, - end_time=admission_end, - ) - all_radiology_texts.extend(radiology_texts) - all_radiology_times_from_admission.extend(radiology_times) - - 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_CODE_TOKEN]) - all_icd_times.append(self.MISSING_FLOAT_TOKEN) - - previous_admission_time = admission_time ->>>>>>> origin/main lab_times, lab_values, lab_masks = self._collect_labs( patient=patient, @@ -1332,27 +1347,6 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri all_lab_values.extend(lab_values) all_lab_masks.extend(lab_masks) -<<<<<<< HEAD -======= - # 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 - ->>>>>>> origin/main if len(all_lab_values) == 0: all_lab_values.append( [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES) @@ -1360,55 +1354,10 @@ def __call__(self, patient: Any) -> List[Dict[str, Any]]: # type: ignore[overri all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES)) all_lab_times.append(self.MISSING_FLOAT_TOKEN) -<<<<<<< HEAD single_patient_longitudinal_record = { "patient_id": patient.patient_id, "labs": (all_lab_times, all_lab_values), "labs_mask": (all_lab_times, all_lab_masks), -======= - # If all admissions were skipped before ICD collection, ensure a - # single placeholder step so StageNetProcessor does not emit None time. - if len(all_icd_codes) == 0: - all_icd_codes.append([self.MISSING_CODE_TOKEN]) - all_icd_times.append(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] - - discharge_note_times_from_admission = ( - all_discharge_texts, - all_discharge_times_from_admission, - ) - radiology_note_times_from_admission = ( - all_radiology_texts, - all_radiology_times_from_admission, - ) - - # Per-admission note fallback happens inside _collect_notes(). - # This final guard handles the edge case where every admission was - # skipped before _collect_notes() was reached. - if len(all_discharge_texts) == 0: - discharge_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - if len(all_radiology_texts) == 0: - radiology_note_times_from_admission = ( - [self.MISSING_TEXT_TOKEN], - [self.MISSING_FLOAT_TOKEN], - ) - - single_patient_longitudinal_record = { - "patient_id": patient.patient_id, - "discharge_note_times": discharge_note_times_from_admission, - "radiology_note_times": radiology_note_times_from_admission, - "icd_codes": (all_icd_times, all_icd_codes), - "labs": (all_lab_times, all_lab_values), - "labs_mask": (all_lab_times, all_lab_masks), - "cxr_image_times": (all_cxr_paths, all_cxr_times), ->>>>>>> origin/main "mortality": mortality_label, "window_start": effective_start, "window_end": effective_end,