Add missing citations in PyHealth code (#1181)

* Add missing citations for TFMTokenizer, EHRMambaCEHR, CEHR embeddings, comprehensiveness/sufficiency metrics, MLE, and NNAAR

* to pass the pr checks

Author: lehendo

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
 

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

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::
 

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

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

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

pyhealth/models/cehr_embeddings.py

@@ -48,7 +48,32 @@ 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).
+
+    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__(
         self,

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``).
@@ -27,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__(

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)