From 6220fb4f2ab384f0e99294b1bf71879d5ad3e78a Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 21 Jul 2026 23:24:30 -0700 Subject: [PATCH 1/2] Add missing citations for TFMTokenizer, EHRMambaCEHR, CEHR embeddings, comprehensiveness/sufficiency metrics, MLE, and NNAAR --- docs/api/metrics/pyhealth.metrics.generative.rst | 3 ++- examples/conformal_eeg/test_tfm_tuev_inference.py | 2 ++ pyhealth/metrics/generative/privacy.py | 6 ++++++ pyhealth/metrics/generative/utility.py | 7 +++++++ pyhealth/metrics/interpretability/comprehensiveness.py | 5 +++++ pyhealth/metrics/interpretability/sufficiency.py | 5 +++++ pyhealth/models/cehr_embeddings.py | 9 ++++++++- pyhealth/models/ehrmamba_cehr.py | 7 +++++++ pyhealth/models/tfm_tokenizer.py | 10 ++++++++-- 9 files changed, 50 insertions(+), 4 deletions(-) diff --git a/docs/api/metrics/pyhealth.metrics.generative.rst b/docs/api/metrics/pyhealth.metrics.generative.rst index 85e448a52..2116903d2 100644 --- a/docs/api/metrics/pyhealth.metrics.generative.rst +++ b/docs/api/metrics/pyhealth.metrics.generative.rst @@ -2,7 +2,8 @@ pyhealth.metrics.generative =================================== Evaluation metrics for synthetic (generative) EHR data, covering privacy, -utility, and statistical fidelity. +utility, and statistical fidelity. See each function's docstring for the +paper it implements. .. currentmodule:: pyhealth.metrics.generative diff --git a/examples/conformal_eeg/test_tfm_tuev_inference.py b/examples/conformal_eeg/test_tfm_tuev_inference.py index fede48d4d..d25eff7ac 100644 --- a/examples/conformal_eeg/test_tfm_tuev_inference.py +++ b/examples/conformal_eeg/test_tfm_tuev_inference.py @@ -5,6 +5,8 @@ present in this repo. No training — pure inference to verify weights and normalization are correct. +Model: TFMTokenizer, see pyhealth.models.tfm_tokenizer for the paper citation. + Usage: python examples/conformal_eeg/test_tfm_tuev_inference.py python examples/conformal_eeg/test_tfm_tuev_inference.py --gpu_id 1 diff --git a/pyhealth/metrics/generative/privacy.py b/pyhealth/metrics/generative/privacy.py index 70edc4957..9f9696caa 100644 --- a/pyhealth/metrics/generative/privacy.py +++ b/pyhealth/metrics/generative/privacy.py @@ -55,6 +55,12 @@ def calc_nnaar( ) -> Dict[str, Tuple[float, float]]: """Computes the Nearest Neighbor Adversarial Accuracy Risk (NNAAR). + Paper: + Yale, Andrew, Saloni Dash, Ritik Dutta, Isabelle Guyon, Adrien Pavao, + and Kristin P. Bennett. "Generation and Evaluation of Privacy + Preserving Synthetic Health Data." Neurocomputing 416 (2020): + 244-255. https://doi.org/10.1016/j.neucom.2019.12.136 + NNAAR measures whether the synthetic data sits closer to the real training data than to held-out test data, which would indicate memorization:: diff --git a/pyhealth/metrics/generative/utility.py b/pyhealth/metrics/generative/utility.py index d36ffcacd..3d10778eb 100644 --- a/pyhealth/metrics/generative/utility.py +++ b/pyhealth/metrics/generative/utility.py @@ -53,6 +53,13 @@ def compute_mle( (Train-Synthetic-Test-Real, TSTR). Both are evaluated on the same real test set. Synthetic accuracy/F1 close to real accuracy/F1 indicates high utility. + Paper: + Esteban, Cristobal, Stephanie L. Hyland, and Gunnar Ratsch. + "Real-valued (Medical) Time Series Generation with Recurrent + Conditional GANs." arXiv:1706.02633 (2017). Introduces and names + the "Train on Synthetic, Test on Real" (TSTR) evaluation protocol + this function implements. + Note: The current implementation hard-codes the downstream task to next-visit prediction (built via diff --git a/pyhealth/metrics/interpretability/comprehensiveness.py b/pyhealth/metrics/interpretability/comprehensiveness.py index 755daf47d..5374c622b 100644 --- a/pyhealth/metrics/interpretability/comprehensiveness.py +++ b/pyhealth/metrics/interpretability/comprehensiveness.py @@ -18,6 +18,11 @@ class ComprehensivenessMetric(RemovalBasedMetric): are REMOVED (ablated). Higher scores indicate more faithful interpretations. + Paper: + DeYoung, Jay, Sarthak Jain, Nazneen Fatema Rajani, Eric Lehman, + Caiming Xiong, Richard Socher, and Byron C. Wallace. + "ERASER: A Benchmark to Evaluate Rationalized NLP Models." ACL 2020. + The metric is computed as: COMP = (1/|B|) × Σ[p_c(x)(x) - p_c(x)(x \\ x:q%)] q∈B diff --git a/pyhealth/metrics/interpretability/sufficiency.py b/pyhealth/metrics/interpretability/sufficiency.py index c6c6a8fa5..cb6ab0efc 100644 --- a/pyhealth/metrics/interpretability/sufficiency.py +++ b/pyhealth/metrics/interpretability/sufficiency.py @@ -18,6 +18,11 @@ class SufficiencyMetric(RemovalBasedMetric): features are KEPT (all others removed). Lower scores indicate more faithful interpretations. + Paper: + DeYoung, Jay, Sarthak Jain, Nazneen Fatema Rajani, Eric Lehman, + Caiming Xiong, Richard Socher, and Byron C. Wallace. + "ERASER: A Benchmark to Evaluate Rationalized NLP Models." ACL 2020. + The metric is computed as: SUFF = (1/|B|) × Σ[p_c(x)(x) - p_c(x)(x:q%)] q∈B diff --git a/pyhealth/models/cehr_embeddings.py b/pyhealth/models/cehr_embeddings.py index 7974a699e..25c162f4f 100644 --- a/pyhealth/models/cehr_embeddings.py +++ b/pyhealth/models/cehr_embeddings.py @@ -48,7 +48,14 @@ def forward(self, visit_segments: torch.Tensor) -> torch.Tensor: class MambaEmbeddingsForCEHR(nn.Module): - """CEHR-style combined embeddings for Mamba (concept + type + time + age + visit).""" + """CEHR-style combined embeddings for Mamba (concept + type + time + age + visit). + + Paper: Same paper as :class:`~pyhealth.models.ehrmamba.EHRMamba` -- + EHRMAMBA: Towards Generalizable and Scalable Foundation Models for + Electronic Health Records (arxiv 2405.14567). This embedding scheme is + part of that paper's own Odyssey toolkit (see module header for the + code source). + """ def __init__( self, diff --git a/pyhealth/models/ehrmamba_cehr.py b/pyhealth/models/ehrmamba_cehr.py index cd555629c..a1a8e3c04 100644 --- a/pyhealth/models/ehrmamba_cehr.py +++ b/pyhealth/models/ehrmamba_cehr.py @@ -18,6 +18,13 @@ class EHRMambaCEHR(BaseModel): """Mamba backbone over CEHR embeddings (FHIR / MPF pipeline). + Paper: Same paper as :class:`~pyhealth.models.ehrmamba.EHRMamba` -- + EHRMAMBA: Towards Generalizable and Scalable Foundation Models for + Electronic Health Records (arxiv 2405.14567). This class combines that + paper's Mamba backbone (:class:`~pyhealth.models.ehrmamba.MambaBlock`) + with CEHR-style embeddings (see + :class:`~pyhealth.models.cehr_embeddings.MambaEmbeddingsForCEHR`). + Args: dataset: Fitted :class:`~pyhealth.datasets.SampleDataset` with MPF task schema. vocab_size: Concept embedding vocabulary size (typically ``task.vocab.vocab_size``). diff --git a/pyhealth/models/tfm_tokenizer.py b/pyhealth/models/tfm_tokenizer.py index 38bd6de04..04f83cc57 100644 --- a/pyhealth/models/tfm_tokenizer.py +++ b/pyhealth/models/tfm_tokenizer.py @@ -714,10 +714,16 @@ def load_embedding_weights(source_model, target_model): class TFMTokenizer(BaseModel): """TFM-Tokenizer model. - + This model uses VQ-VAE with transformers to tokenize EEG signals. It can extract discrete tokens and continuous embeddings for downstream tasks. - + + Paper: + Pradeepkumar, Jathurshan, Xihao Piao, Zheng Chen, and Jimeng Sun. + "Tokenizing Single-Channel EEG with Time-Frequency Motif Learning." + ICLR 2026. https://arxiv.org/abs/2502.16060 + Code: https://github.com/Jathurshan0330/TFM-Tokenizer + The model expects two inputs: - STFT spectrogram: shape (batch, n_freq, n_time) - Raw temporal signal: shape (batch, n_samples) From d21b68364fdd233e6837b01d9dfce5238f08bef3 Mon Sep 17 00:00:00 2001 From: lehendo Date: Tue, 21 Jul 2026 23:33:17 -0700 Subject: [PATCH 2/2] to pass the pr checks --- pyhealth/models/cehr_embeddings.py | 18 ++++++++++++++++++ pyhealth/models/ehrmamba_cehr.py | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/pyhealth/models/cehr_embeddings.py b/pyhealth/models/cehr_embeddings.py index 25c162f4f..c8ed4e88f 100644 --- a/pyhealth/models/cehr_embeddings.py +++ b/pyhealth/models/cehr_embeddings.py @@ -55,6 +55,24 @@ class MambaEmbeddingsForCEHR(nn.Module): Electronic Health Records (arxiv 2405.14567). This embedding scheme is part of that paper's own Odyssey toolkit (see module header for the code source). + + Examples: + >>> import torch + >>> from pyhealth.models.cehr_embeddings import MambaEmbeddingsForCEHR + >>> embeddings = MambaEmbeddingsForCEHR(vocab_size=100, hidden_size=32) + >>> batch_size, seq_len = 2, 5 + >>> input_ids = torch.randint(0, 100, (batch_size, seq_len)) + >>> token_type_ids = torch.zeros(batch_size, seq_len, dtype=torch.long) + >>> time_stamps = torch.zeros(batch_size, seq_len) + >>> ages = torch.zeros(batch_size, seq_len) + >>> visit_orders = torch.zeros(batch_size, seq_len, dtype=torch.long) + >>> visit_segments = torch.zeros(batch_size, seq_len, dtype=torch.long) + >>> out = embeddings( + ... input_ids, token_type_ids, time_stamps, ages, + ... visit_orders, visit_segments, + ... ) + >>> out.shape + torch.Size([2, 5, 32]) """ def __init__( diff --git a/pyhealth/models/ehrmamba_cehr.py b/pyhealth/models/ehrmamba_cehr.py index a1a8e3c04..711b032c6 100644 --- a/pyhealth/models/ehrmamba_cehr.py +++ b/pyhealth/models/ehrmamba_cehr.py @@ -34,6 +34,24 @@ class EHRMambaCEHR(BaseModel): state_size: SSM state size per channel. conv_kernel: Causal conv kernel in each block. dropout: Dropout before classifier. + + Examples: + >>> from pyhealth.datasets import MIMIC4FHIR, split_by_patient + >>> from pyhealth.tasks.mpf_clinical_prediction import ( + ... MPFClinicalPredictionTask, + ... ) + >>> from pyhealth.models import EHRMambaCEHR + >>> dataset = MIMIC4FHIR(root="/path/to/mimic-iv-fhir-demo") + >>> sample_dataset = dataset.set_task(MPFClinicalPredictionTask()) + >>> train_ds, val_ds, test_ds = split_by_patient( + ... sample_dataset, [0.7, 0.1, 0.2] + ... ) + >>> vocab_size = ( + ... sample_dataset.input_processors["concept_ids"].vocab.vocab_size + ... ) + >>> model = EHRMambaCEHR( + ... dataset=sample_dataset, vocab_size=vocab_size, embedding_dim=32 + ... ) """ def __init__(