[DUMMY PR] Wp dev/ra/labs only variant - #1184
Draft
will-pang wants to merge 232 commits into
Draft
Conversation
Introduces a shared BaseEmbeddingModel ABC (embedding_dim property +
forward) and reorganises the four embedding models into a unified
pyhealth/models/embedding/ package:
base.py - BaseEmbeddingModel abstract base class
vanilla.py - EmbeddingModel (codes, sequences, timeseries)
vision.py - VisionEmbeddingModel (patch/CNN/ResNet)
text.py - TextEmbeddingModel (renamed from TextEmbedding,
TextEmbedding kept as backward compatability alias)
unified.py - UnifiedMultimodalEmbeddingModel, IMAGE encoding now
delegates to PatchEmbedding from vision.py + mean pool
instead of an inline _build_image_encoder helper
All existing public exports in pyhealth/models/__init__.py are unchanged, w/ test imports updated to new paths.
…tream models
UnifiedMultimodalEmbeddingModel now takes prebuilt unimodal embedding
models via the new field_embeddings parameter, pulling their trained
encoder weights instead of building from scratch:
- EmbeddingModel -> embedding_layers[field] (nn.Embedding / nn.Linear)
- VisionEmbeddingModel -> embedding_layers[field] backbone + mean pool
- TextEmbeddingModel -> transformer (BERT) + fc (projection)
Dims mismatches are handled automatically with an nn.Linear bridge.
Transformer, EHRMamba, and JambaEHR have an optional
unified_embedding parameter (UnifiedMultimodalEmbeddingModel). When
provided the model switches to unified mode:
- All temporal fields are jointly embedded and time-sorted by
UnifiedMultimodalEmbeddingModel -> single (B, S_total, E) sequence
- A single backbone (TransformerLayer / MambaBlock stack / JambaLayer)
processes the interleaved sequence instead of one backbone per field
- The fc head takes embedding_dim inputs (not n_fields * embedding_dim)
- Per-field mode is unchanged; backward compat guaranteed
_forward_unified and _build_unified_inputs helpers are added to each
downstream model. 7 new tests cover field_embeddings reuse, projection
bridging, and unified-mode forward + backward through all three models.
…xes ModuleNotFoundError for users who don't have polars installed - Track per category observation status when building lab vectors: True = value was observed, False = imputed 0.0, emitted as labs_mask field with the same (times, values) structure as labs - Add labs_mask to input_schema (stagenet_tensor) so it is batched and padded in sync w/ labs - Propagate all false mask rows through both the inner fallback path and the post-loop all-admissions guard
forward() now also returns "token_emb": the per-event content embedding before the time and type embeddings are added, in the same temporally-sorted order as "sequence". This is the correct reconstruction target for masked-modeling pretrainers (MAE/SimMIM). The time and type components of "sequence" are largely recoverable from event position, so reconstructing the composed "sequence" dilutes the content learning signal. The change is additive and backward-compatible (a new output key only; all existing keys are unchanged). Adds a test asserting sequence == token_emb + time_embed(time) + type_embedding(type_ids).
…into ra/labs-only-variant
…kage, Balanced Sampling, ICU Vitals (#28)
…into ra/labs-only-variant # Conflicts: # examples/mortality_prediction/unified_embedding_e2e_mimic4.py # pyhealth/models/embedding/unified.py # pyhealth/tasks/multimodal_mimic4.py # scripts/slurm/run_table2.sh
will-pang
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.