dl4h final

[raymondcoding15]

Contributor: Andrew Raymond (andrew46@illinois.edu), CS 598 DL4H SP26 final project.
Type of contribution: Option 2, Model.
Paper: Kim et al., "Integrating ChatGPT into Secure Hospital Networks: A Case Study on Improving Radiology Report Analysis." CHIL 2024. PMLR 248:72-87.

What this PR does

This PR adds SentenceKDTransformer to pyhealth.models. It's the student-side model from Kim et al.: a BERT-family encoder trained on sentence-level labels with a combined cross-entropy plus supervised-contrastive objective (paper Eq. 5, contrastive term from Khosla et al., NeurIPS 2020). A small helper aggregates per-sentence probabilities into a document-level anomaly score the way the paper does in Eq. 4 (max over sentences), and also exposes two other aggregation modes I found useful during ablations.

No new dataset or task is introduced. The example exercises the model on the existing MedicalTranscriptionsDataset so reviewers don't need a DUA, Kaggle key, or anything else.

Files to review

The diff is limited to seven files:

• pyhealth/models/sentence_kd_transformer.py: the model and a module-level supervised_contrastive_loss helper. Start here.
• pyhealth/models/__init__.py: one-line export of both symbols.
• docs/api/models/pyhealth.models.SentenceKDTransformer.rst: Sphinx autodoc page.
• docs/api/models.rst: toctree entry.
• tests/core/test_sentence_kd_transformer.py: 16 unit tests, all synthetic.
• examples/medical_transcriptions_classification_sentence_kd_transformer.py: the headless ablation suite.
• examples/medical_transcriptions_classification_sentence_kd_transformer.ipynb: Colab-runnable walkthrough of the same ablations.

What I replicated vs. what I extended

Faithful to the paper: BERT student with a [CLS] classifier head, the combined CE + λ·SupCon loss, RadBERT as the default backbone, and max-pool document scoring.

Things the paper doesn't do that this PR adds:

• A full λ sweep over {0, 0.1, 0.5, 1, 2, 5}. The paper only compares λ=0 to λ=1 in Table 4, so it can't tell you where the contrastive term starts to dominate cross-entropy; the sweep does.
• A temperature sweep over τ ∈ {0.05, 0.1, 0.2, 0.5, 1.0}. The paper inherits Khosla et al.'s default of 0.07 without evaluating it, so I treated τ as a free knob.
• A learning-rate × dropout × batch-size grid. Mostly to rule out "the lr did all the work" when a contrastive configuration wins.
• Backbone comparison (BERT-base vs Bio_ClinicalBERT vs RadBERT) under λ=1. The paper compares backbones only under plain CE (Table 2).
• Two aggregation modes besides max: a topk_mean that's less sensitive to a single high-scoring sentence, and an attention-pooled variant.

A note on the teacher step

Kim et al. use ChatGPT as a cloud teacher to produce ternary sentence labels (normal/abnormal/uncertain) for MIMIC-CXR, and then distill a BERT student against those labels. I deliberately did not replicate that step here, for three reasons:

1. The example dataset (mtsamples) ships with ground-truth specialty labels, so an LLM teacher isn't needed to train the student.
2. An OpenAI API call would block "Run All" reproducibility for anyone without a paid API key.
3. For the paper's actual target data (credentialed MIMIC-CXR text), most institutional DUAs don't allow sending report text to a third-party cloud API without a BAA, so I didn't want to model that pattern in an example other people will copy.

The model itself is teacher-agnostic: whatever source produces the sentence labels (a human annotator, a rule-based labeler, or an on-prem model like Llama-3-8B) drops in without code changes. What this PR covers is the student side of the pipeline (the loss, the forward pass, and the aggregation), which is the portion of the paper PyHealth can actually ship.

How to test locally

Unit tests run on CPU and use prajjwal1/bert-tiny (~17 MB, one-time HuggingFace cache download). Full suite finishes in roughly 6 seconds, with each individual test under 2 seconds:

python -m pytest tests/core/test_sentence_kd_transformer.py -v

For a fast end-to-end sanity check without touching real data:

python examples/medical_transcriptions_classification_sentence_kd_transformer.py --quick --epochs 1

For the real mtsamples run (same code path the notebook uses):

python examples/medical_transcriptions_classification_sentence_kd_transformer.py \
    --data_root /path/to/mtsamples --epochs 3

If you're compute-constrained, cap the dataset with --max_samples:

python examples/medical_transcriptions_classification_sentence_kd_transformer.py \
    --data_root /path/to/mtsamples --epochs 2 --max_samples 1500

The notebook is set up for Colab with a T4 and takes about 3 to 4 minutes end to end with the default 1500-row subsample; flip MAX_SAMPLES = None in the data cell to run on the full ~5,000 mtsamples corpus instead, which takes about 10 to 12 minutes. The notebook auto-downloads mtsamples from harishnair04/mtsamples on HuggingFace (CC0, ~17 MB), so there's no credential setup. Subsampling is seeded, so every ablation run sees the same data and trends are preserved at 1500 rows.

[github-actions]

This PR has been automatically marked as stale because it has not had recent activity. It will be closed in 7 days if no further activity occurs.