ML4H merge: tranche 1 - #1185
Conversation
Empty notes, labs, CXR, and ICD are now zero-length sequences instead of a constant [MISSING_TEXT] row, black image, or pad visit. The tokenizer crash on empty batches is handled by constructing empty tensors, so BERT cannot treat note presence as a free mortality feature. Co-authored-by: Cursor <cursoragent@cursor.com>
nn.Module.train() re-enables dropout inside Bio_ClinicalBERT even when every weight has requires_grad=False. Pin those encoders back to eval so a frozen note embedding is deterministic across steps. Co-authored-by: Cursor <cursoragent@cursor.com>
A frozen BERT forward is identical across epochs, but batch padding width changes every shuffle, so a key over the full padded row never hits. Hash only the attended tokens so the same note reuses its [CLS] vector. Co-authored-by: Cursor <cursoragent@cursor.com>
a0f1422 deleted _scan_table/_scan_parquet while MEDSDataset._subset_patient_ids still calls _scan_parquet, so MEDS loads crashed. load_table routes through _scan_table again, and resolve_table_path keeps absolute cache paths. The test reads a real two-row parquet file. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Empty notes, labs, and CXR are now zero-length sequences instead of a fake Frozen BERT stays in eval when Trainer calls
|
The explicit path filled padded scores with -1e9, which is outside the fp16 range, so AMP overflowed. Ordinary forwards now use fused scaled_dot_product_attention; the explicit path stays behind register_hook for interpretability and fills with finfo(dtype).min. The test checks a padded fp16 pass stays finite and that -1e9 still overflows. Co-authored-by: Cursor <cursoragent@cursor.com>
The collator padded short samples with 0.0 and nothing marked the extra slots, so they looked like real measurements at admission time. It now emits {field}__pad_mask, the unified heads thread it through, and RNN packed lengths clamp at 1 so an all-pad sample does not crash. Tests check the collate mask, pad-last sort, and an all-pad RNN step.
Co-authored-by: Cursor <cursoragent@cursor.com>
NestedSequenceProcessor used padding_idx=None so a fake empty visit could have a non-zero vector. Index 0 then received gradients. Empty visits are now zero events, so the pad row stays frozen zeros. The test checks both the zeros and a zero gradient on that row. Co-authored-by: Cursor <cursoragent@cursor.com>
The sunlab loader required a directory named images and wrote the derived CSV into the PhysioNet root, which is read-only on the cluster. Both images and resized_images are accepted, cache is tried first, and the generated YAML points at the absolute CSV. The test chmods the root to 555 and checks the CSV lands in cache. Co-authored-by: Cursor <cursoragent@cursor.com>
|
A few more measurement fixes: The attention mask fill is The collator now emits Nested code embeddings keep The sunlab CXR loader accepts |
"bfloat16" and any other spelling previously fell through to float16 and a GradScaler, so a typo changed both precision and gradient scaling with no message. Co-authored-by: Cursor <cursoragent@cursor.com>
Without this the sixth backbone is missing from the table, and a unified MLP would score padded slots as real events the same way the other heads used to. Co-authored-by: Cursor <cursoragent@cursor.com>
Iterating a StreamingDataset under torchrun fitted 1/N of the split (len()=5 of 20 at WORLD_SIZE=4). Padded zeros also pulled sodium's mean from 140 to 105. patient_to_index after subset() still holds parent indices and raised on index 237. Co-authored-by: Cursor <cursoragent@cursor.com>
metrics_history.json stored the score but not the conditions, so a frozen- encoder run and a fine-tuned run were indistinguishable after the job log was gone. Co-authored-by: Cursor <cursoragent@cursor.com>
NotesLabsCXRMIMIC4 and CXRMIMIC4 skipped any stay with admit time >= first_admit + window_hours, so a later admission contributed no images. Lab collection was already per-admission. Bump emitted_data_version to 2 so those caches cannot be reused. Co-authored-by: Cursor <cursoragent@cursor.com>
The directory was {model}_seed{seed}, so labs and notes_labs at one seed
destroyed the first arm. split_by_patient fell back to split_by_sample with
no warning, and predictions came from test or val or train. Wire MLP and the
lab standardiser; leave the Jamba library default at 6.
Co-authored-by: Cursor <cursoragent@cursor.com>
The class default was None (labs through discharge) while the docstring, LabsMIMIC4, and the runner all used 24. Event times were hours from each stay's own admit, so stay 2 at +6h sorted with stay 1 at +6h. Collection is still per stay; times are hours from the first stay in the sample. Bump emitted_data_version to 3. Co-authored-by: Cursor <cursoragent@cursor.com>
The old sinusoid wrapped every 720 hours, so later stays aliased with the first. MIMIC timestamps ICD at discharge, which leaks the in-hospital mortality label. Co-authored-by: Cursor <cursoragent@cursor.com>
Tests and docs imported pyhealth.models.unified_embedding while runtime used embedding.unified, so the two files could drift. embedding.py next to the embedding/ package was unused. Co-authored-by: Cursor <cursoragent@cursor.com>
use_amp=True on CPU or MPS warned that CUDA was unavailable and skipped mixed precision. Unknown amp_dtype spellings were already rejected. Co-authored-by: Cursor <cursoragent@cursor.com>
Will's protocol is through discharge, so the 24h class/runner default is gone. Discharge-section text is still written at the end of the stay; stamping it at charttime leaked length of stay. Radiology stays at exam time. Cache version 4. Co-authored-by: Cursor <cursoragent@cursor.com>
|
Pushed the rest of the measurement fixes onto this branch (fast-forward to Protocol matches the opening post: full stay (no 24h cap), no ICD, no re-weighting. Event times are hours from the first stay in the sample, so later admissions do not land on top of the first. Admission-context discharge sections are stamped at that stay's admit; radiology stays at exam time. Empty notes are still empty sequences, not a fake Also: AMP dtype is validated and autocast follows the trainer device; unified MLP gets the collate pad mask; lab z-scores fit on observed train rows; run directories are Task cache version is 4. Rebuild caches before comparing numbers. |
There was a problem hiding this comment.
Leaving feedback/thoughts here as I work through the commits:
- Stop emitting fake missing-event placeholders (9782aca): ✅ Reviewed and looks good to me
- Keep frozen text encoders in eval when Trainer calls train()(869ac8e): ✅ Reviewed and looks good to me
- Cache frozen [CLS] embeddings keyed on real tokens, not padded rows.: ✅ Reviewed and looks good to me
- Restore BaseDataset parquet scanning that MEDS still calls.: Do we really need this? Are we working with MEDS at the moment?
- Fill attention masks with dtype min and use fused SDPA: ⏸ Changes are beyond my domain knowledge, might want to loop someone else in to review (maybe @joshuasteier?)
- Record batch padding and skip those slots in the unified sequence: ⏸ Changes are beyond my domain knowledge, might want to loop someone else in to review (maybe @joshuasteier?)
- Keep padding_idx=0 on nested code embeddings: 🛑 Might want to double check this update @Rian354
- Accept resized_images and write sunlab CXR metadata to cache.: Tranche 2?
- Refuse unknown amp_dtype instead of silently selecting fp16.: ✅ Reviewed and looks good to me
- Thread collate pad_mask through the unified MLP path.: ✅ Reviewed and looks good to me
- Fit lab z-scores on observed train rows via region_of_interest.: I think I sort of get what this is trying to do (normalize lab values), but maybe not necessary for now?
- Write run_config.json next to the metrics of a finished run.: Useful, but maybe not appropriate for the full repo but instead just for our own multimodal-pyhealth-repo?
- Stop dropping later stays against the first admission's CXR window.: I thought we're dropping observation widnows for now?
- Keep paired runs from overwriting each other or reporting train as test.: 🤔 Let me give this a deeper review, I think this is a crucial issue to look into
- Give NotesLabs a 24h window and put concatenated stays on one timeline.: I thought we're dropping observation windows for now?
- Give time embeddings a 10-year span and drop ICDLabsMIMIC4.: ✅ Reviewed and looks good to me
- Point the old unified_embedding import at the package copy.: ✅ Reviewed and looks good to me
- Point AMP autocast at the trainer device instead of hardcoding CUDA.:✅ Reviewed and looks good to me
- Default to a full stay and stamp admission-context notes at admit.: 🤔 Let me give this a deeper review, I think this is a crucial issue to look into
There was a problem hiding this comment.
This looks good to me
There was a problem hiding this comment.
This looks good to me. I think I didn't port this update from a previous PR.
There was a problem hiding this comment.
Looks good to me as well. Seems like what this is doing is making the longest_length variable per sample, bounded by max_length.
There was a problem hiding this comment.
Looks good to me as well.
There was a problem hiding this comment.
Looks good to me. A lot of changes here, but to summarize:
(1) Removing encoding missingness with a token. Now it does:
Field: ICD codes (nested)
Tensor shape: (num_visits, max_len)
Which dim is data-dependent: num_visits
What's 0 when missing: num_visits = 0 → shape (0, max_len).
max_len is fixed (set once during fit(), same for every
patient)
────────────────────────────────────────
Field: Labs
Tensor shape: (num_timesteps, 10)
Which dim is data-dependent: num_timesteps
What's 0 when missing: num_timesteps = 0 → shape (0, 10).
The
10 (lab categories) is always fixed
────────────────────────────────────────
Field: Notes (tokenizer path)
Tensor shape: (num_notes, seq_len)
Which dim is data-dependent: both num_notes and seq_len
What's 0 when missing: num_notes = 0 → shape (0, 1)
specifically (the processor's empty-guard hardcodes width 1
rather than computing a "longest note" from an empty list,
since there's nothing to tokenize)
────────────────────────────────────────
Field: Notes (raw-text path)
Tensor shape: Python list + (num_notes,) time tensor
Which dim is data-dependent: num_notes
What's 0 when missing: num_notes = 0 → empty list,
torch.tensor([])
────────────────────────────────────────
(2) Creates a admission_window_end() helper function that standardizes windows across the board. I don't think we will be using it for now (i.e., observation windows), but very useful to keep in the future.
|
LGTM. Thank you! |
There was a problem hiding this comment.
- Looks good to me too. I think I just need to spend more time understanding how masking is passed based on different modalities, but no conceptual red flags at the moment.
- Looks good to me. Probably need to dig deeper into how we are cache-ing text, but can do that some other time
Give time embeddings a 10-year span and drop ICDLabsMIMIC4: Changes to TimeEmbedding looks good to me as well
There was a problem hiding this comment.
Changes are a bit out of my domain knowledge here - perhaps @jhnwu3 might have thoughts here.
3 participants
A first pass at consolidating all the work we did so that we can merge to main. I'll probably need everyone's help in reviewing this (@joshuasteier, @Rian354, @joshuxchn, as well as @jhnwu3), but in short the updates here add in unified embedding support and tested across the different backbones (RNN, transformer, bottleneck
transformer, EHRMamba, JambaEHR) that we're interested in presenting results. In terms of modality, this focuses only on labs, labs + notes, labs + notes + cxr.
Three big updates made
What to remove
scripts/will/*which I've used to submit to condor and can be removed.