Feat/vals - #346
Conversation
|
Warning Review limit reached
Next review available in: 27 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughAdds optional per-token numeric values across data preparation, dataset transformations, embeddings, masked and autoregressive pretraining, numeric prediction, and weighted loss computation. Training entrypoints and YAML examples expose numeric-column, embedding-mode, and value-loss settings. ChangesNumeric value modeling pipeline
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant DataPreparation
participant SubjectData
participant PretrainDataset
participant EhrValueEmbeddings
participant BonsaiValuePretrain
participant ValuePretrainModule
DataPreparation->>SubjectData: numeric_value column
SubjectData->>PretrainDataset: per-token numeric_value
PretrainDataset->>EhrValueEmbeddings: numeric_value and numeric_target
EhrValueEmbeddings->>BonsaiValuePretrain: fused embeddings
BonsaiValuePretrain->>ValuePretrainModule: concept and value logits/labels
ValuePretrainModule->>ValuePretrainModule: weighted concept and value loss
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 2📝 Generate docstrings 💡
🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
bonsai/modules/networks/bonsai_nets.py (1)
28-39: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winPersist
value_embedding_modein checkpoint hyperparameters.
BonsaiBase.hparamsomits this new architecture setting. Consequently, finetuning cannot reliably reconstruct the pretrained embedding layout;strict=Falsecan silently skip or randomly initialize numeric-embedding weights when runtime config differs. Add it toself.hparams, then restore or validate it against the pretraining checkpoint before loading weights.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@bonsai/modules/networks/bonsai_nets.py` around lines 28 - 39, Persist value_embedding_mode in BonsaiBase.hparams during initialization, and ensure checkpoint loading restores or validates this setting against the pretrained checkpoint before weights are loaded. Use the existing hparams/checkpoint-loading flow and preserve the embedding layout so mismatches cannot silently proceed under strict=False.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@bonsai/modules/networks/components/embeddings.py`:
- Around line 137-141: Update the "film" branch of the value embedding flow to
derive gamma and beta from value_embed, then apply those FiLM parameters to
concept_embeds. Preserve the existing dtype conversion and mask-based selection
so numeric positions use the modulated concept representation while other
positions remain unchanged.
---
Outside diff comments:
In `@bonsai/modules/networks/bonsai_nets.py`:
- Around line 28-39: Persist value_embedding_mode in BonsaiBase.hparams during
initialization, and ensure checkpoint loading restores or validates this setting
against the pretrained checkpoint before weights are loaded. Use the existing
hparams/checkpoint-loading flow and preserve the embedding layout so mismatches
cannot silently proceed under strict=False.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ba97667d-a4f9-463d-8604-2e97748dda53
⛔ Files ignored due to path filters (33)
example_data/example_MEDS_data/held_out/0.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data/train/0.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data/tuning/0.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/0.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/1.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/2.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/3.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/4.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/5.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/6.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/train/7.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/tuning/0.parquetis excluded by!**/*.parquetexample_data/example_MEDS_data_wo_held_out/tuning/1.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/held_out/0.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/held_out/1.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/0.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/1.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/10.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/11.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/12.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/13.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/14.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/15.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/2.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/3.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/4.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/5.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/6.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/7.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/8.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/train/9.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/tuning/0.parquetis excluded by!**/*.parquetexample_data/lab_correlated_MEDS_data/tuning/1.parquetis excluded by!**/*.parquet
📒 Files selected for processing (25)
bonsai/functional/censoring.pybonsai/functional/collate.pybonsai/functional/create_data.pybonsai/functional/features.pybonsai/functional/subject_data.pybonsai/functional/truncation.pybonsai/modules/datasets/FinetuneDataset.pybonsai/modules/datasets/PretrainDataset.pybonsai/modules/lightningmodules/PretrainModule.pybonsai/modules/networks/bonsai_nets.pybonsai/modules/networks/components/embeddings.pybonsai/run/create_data.pybonsai/run/finetune.pybonsai/run/pretrain.pybonsai/run/train.pyconfigs/data_creation/default_create_data.yamlconfigs/examples/example_data.yamlconfigs/examples/example_finetune.yamlconfigs/examples/example_finetune_val.yamlconfigs/examples/example_outcome1.yamlconfigs/examples/example_outcome2.yamlconfigs/examples/example_value_data.yamlconfigs/finetune.yamlconfigs/pretrain.yamlexample_data/lab_correlated_MEDS_data/hash_to_integer_map.pt
There was a problem hiding this comment.
Overall, I think the code looks good from a functionality standpoint. However, there's one crucial thing that cements itself into two parts and that's all the if statement and optionals present.
- @Sllambias would probably agree that this amount of toggling and optionals is a "code smell". I don't think we can fully avoid it, but neither do I think we should fully embrace it
- Having
ifstatements present in any model architecture (inside our model and nn.Modules), lead to graph breaks when using torch.compile. I.e. we're losing optimization when having this.
My suggestions would be:
- We get @Sllambias input and see if we can minimize the amount of optionals. This is mainly for all our functionals and data components. I can't really see many potential drops, but perhaps we can design something to avoid it being a general pattern
- Make a seperate bonsai_net for anything to do with values. We cannot have
ifstatements inside anynn.Module
|
|
||
| embeddings = self.code_embedding(code) | ||
| if self.value_embedding_mode is not None and numeric_value is not None: | ||
| embeddings = self.numeric_value_embedding(numeric_value, embeddings) |
There was a problem hiding this comment.
Might be due to FiLM, so this is potentially incorrect, but I would've thought that it would be embeddings += self.numeric_value_embedding(numeric_value)
There was a problem hiding this comment.
Looking further, looks indeed to be FiLM-specific, so I guess this is of no concern. It's only a concern if we expect to support additional non-FiLM value embeddings; then this pattern breaks
Yeah I agree with the many if statements. Another alternative would be to add a nan numeric column, and then the code just runs as usual, since nan numerics are ignored. But then we would be carrying nan values throughout the model, which seems sub-optimal. I think that the separate model might make more sense |
Sllambias
left a comment
Sllambias
left a comment
There was a problem hiding this comment.
I agree with Mikkel and I see two options (but probably others are also possible):
- Split into value specific funtions/classes that inherit the base and build upon. Downside of this is inheritance also becomes messy but if it's limited to 1 layer I think it can be workable.
- Make values always included, e.g. as an empty list per subject. This would solve most of the conditionals and then only be required to split into separate classes at the model class, rather than having to make parallel classes from the dataset and forward. Would maybe also introduce some issues to have these, but is probably the simplest codewise if it is feasible.
Functionality looks good but I think it's better to do more proper review when the format is nailed down as things will probably change
| ], | ||
| ) | ||
|
|
||
| def compute_combined_loss(self, batch, concept_loss_fn): |
There was a problem hiding this comment.
I think this should be a loss function specific to training with values that that is only instantiated with values so no conditionals are necessary
There was a problem hiding this comment.
I think (2) is definitely feasible as that's just what I ended up doing in the original codebase, but would we prefer this? Or something like (1) instead?
There was a problem hiding this comment.
(2) is probably the cleanest code-wise, but it's hacky, and (1) is probably the cleanest logic-wise. For this, I would try out (1) first and see if this resolves things. I don't mind having some if statements (e.g. if "numeric_value" in subject, then add it to the returned object).
So we could do either:
- If logic is completely new, make a separate function/class for it
- If this just extends functionality with another few lines of code, use inheritance or keep it as if statements. Here we could use inheritance for all classes and just if statements for functions.
Hard to tell if this also becomes messy in its own way, but I'm not the biggest fan of (2) unless alternatives don't work
There was a problem hiding this comment.
I think it's a bad idea to go forward with a maybe it's there maybe it's not approach. That forces these lines in all functions, which is an eyesore in itself, and it also extends horribly:
if "numeric_value" in subject:
subject["numeric_value"] = subject["numeric_value"].float()
I think either commit to it always being there, with the option of it being empty (solution 2), or commit to having specific functionality that should be invoked when values are there (solution 1)
There was a problem hiding this comment.
So I split up the bonsai_nets to value and non-values, but there are still quite a lot of if-statements during data creation. Are we content with that, or would be also like to split this up? I'm not what the most clean way of doing this would be, as it is mostly just a few lines in each function.
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
bonsai/modules/networks/components/embeddings.py (1)
150-163: 🎯 Functional Correctness | 🔴 Critical | ⚡ Quick winFiLM parameters still derived from the wrong tensor — looks like a regression of a previously "addressed" bug.
gamma/betaare computed fromconcept_embedsand then used to modulatevalue_embed, and the result replacesconcept_embedsat numeric positions. This is backwards: FiLM should derivegamma/betafrom the conditioning signal (value_embed) and apply them to modulateconcept_embeds, preserving the code identity. A prior review on this exact code (same line numbers, 157-161) flagged this precise pattern and is marked "Addressed in commits f54b019 to df801ef" — but the code shown here is identical to the pre-fix version, so the fix does not appear to be present on this branch.🐛 Proposed fix (matches the previously accepted suggestion)
if self.value_embedding_mode == "film": - gamma = self.gamma_layer(concept_embeds) - beta = self.beta_layer(concept_embeds) - fused = (gamma * value_embed + beta).to(dtype=concept_embeds.dtype) + gamma = self.gamma_layer(value_embed) + beta = self.beta_layer(value_embed) + fused = (gamma * concept_embeds + beta).to(dtype=concept_embeds.dtype) return fused * mask + concept_embeds * (1 - mask)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@bonsai/modules/networks/components/embeddings.py` around lines 150 - 163, Update the film branch in forward so gamma_layer and beta_layer derive their FiLM parameters from value_embed, then apply them to concept_embeds instead of modulating value_embed. Preserve concept_embeds for masked/non-numeric positions and ensure the output retains the existing dtype conversion and mask behavior.
🧹 Nitpick comments (1)
bonsai/modules/networks/bonsai_nets.py (1)
146-308: 📐 Maintainability & Code Quality | 🔵 TrivialDuplicated pretrain-head/finetune-forward logic across base and value variants.
pretrain_head/weight-tying construction is copied identically inBonsaiPretrain.__init__andBonsaiValuePretrain.__init__(lines 172-174, 215-216), andforward()inBonsaiFinetuneandBonsaiValueFinetune(265-269, 304-308) are byte-for-byte identical. A shared mixin/base holding the common head-construction and forward logic (leaving onlyembed()to vary) would remove this duplication and reduce the risk of divergent fixes. This lines up with the PR discussion's own direction toward separating value-related functionality into a dedicated module.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@bonsai/modules/networks/bonsai_nets.py` around lines 146 - 308, Extract the identical pretrain-head construction and weight tying from BonsaiPretrain and BonsaiValuePretrain into a shared mixin or base, and reuse it in both classes. Likewise, move the common prediction-token selection and finetune-head forward logic from BonsaiFinetune and BonsaiValueFinetune into the shared abstraction, while preserving each variant’s existing embed/value behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@bonsai/modules/networks/bonsai_nets.py`:
- Around line 219-232: Update BonsaiValuePretrain.forward to access
batch["numeric_target"] optionally before indexing with mask, and raise a clear
configuration error when numeric_target is absent while value pretraining is
enabled. Preserve the existing val_mask filtering and value-logit computation
when numeric data is present.
---
Duplicate comments:
In `@bonsai/modules/networks/components/embeddings.py`:
- Around line 150-163: Update the film branch in forward so gamma_layer and
beta_layer derive their FiLM parameters from value_embed, then apply them to
concept_embeds instead of modulating value_embed. Preserve concept_embeds for
masked/non-numeric positions and ensure the output retains the existing dtype
conversion and mask behavior.
---
Nitpick comments:
In `@bonsai/modules/networks/bonsai_nets.py`:
- Around line 146-308: Extract the identical pretrain-head construction and
weight tying from BonsaiPretrain and BonsaiValuePretrain into a shared mixin or
base, and reuse it in both classes. Likewise, move the common prediction-token
selection and finetune-head forward logic from BonsaiFinetune and
BonsaiValueFinetune into the shared abstraction, while preserving each variant’s
existing embed/value behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: ac493e62-338a-4f12-8028-4bd83d7ee587
📒 Files selected for processing (12)
bonsai/functional/create_data.pybonsai/modules/datasets/PretrainDataset.pybonsai/modules/lightningmodules/PretrainModule.pybonsai/modules/networks/bonsai_nets.pybonsai/modules/networks/components/embeddings.pybonsai/run/finetune.pybonsai/run/pretrain.pybonsai/run/train.pyconfigs/examples/example_finetune_val.yamlconfigs/examples/example_pretrain_val.yamlconfigs/finetune.yamlconfigs/pretrain.yaml
| def forward(self, batch: dict): | ||
| last_hidden_state = super().forward(batch) | ||
| labels = batch["target"] | ||
|
|
||
| mask = labels != -100 | ||
| last_hidden_state = last_hidden_state[mask] | ||
| labels = labels[mask] | ||
| logits = self.pretrain_head(last_hidden_state) | ||
|
|
||
| val_labels = batch["numeric_target"][mask] | ||
| val_mask = ~torch.isnan(val_labels) | ||
| val_logits = self.pretrain_head_value(last_hidden_state[val_mask]).squeeze(-1) | ||
| val_labels = val_labels[val_mask] | ||
| return logits, labels, val_logits, val_labels |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== candidate files =="
fd -a bonsai_nets.py . | sed 's#^\./##'
echo "== outline =="
ast-grep outline bonsai/modules/networks/bonsai_nets.py --view expanded || true
echo "== relevant source =="
sed -n '1,280p' bonsai/modules/networks/bonsai_nets.py
echo "== search datasets / numeric_target =="
rg -n "numeric_target|numeric_value|PretrainDataset|value_embedding_mode|BonsaiValuePretrain" .Repository: FGA-DIKU/BONSAI
Length of output: 18792
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== PretrainDataset.py =="
sed -n '1,220p' bonsai/modules/datasets/PretrainDataset.py
echo "== value config snippets =="
sed -n '40,65p' configs/pretrain.yaml
sed -n '130,175p' bonsai/modules/datamodules/PretrainDataModule.py
echo "== deterministic batch-key behavior probe, as data =="
python3 - <<'PY'
from pathlib import Path
text = Path("bonsai/modules/datasets/PretrainDataset.py").read_text()
blocks = text.split("\n\n")
by_key = {}
for i, block in enumerate(blocks,1):
if '"numeric_target"' in block:
by_key[f'block_{i}'] = [j for j,line in enumerate(block.splitlines(),1) if '"numeric_target"' in line or "numeric_column" in line]
for k, lines in by_key.items():
print(f"\n{ k } line refs:", lines)
for line_no in [73,74,78,79,80,81,82,156,157,159,160,161]:
line = text.splitlines()[line_no-1]
print(f"{line_no:3}: {line}")
PYRepository: FGA-DIKU/BONSAI
Length of output: 7487
Guard numeric_target before indexing it in BonsaiValuePretrain.
PretrainDataset.__getitem__ only adds numeric_target when the subject contains numeric_value; without that data key, forward() raises an opaque KeyError at batch["numeric_target"]. Use an optional access path and throw a clear configuration error when value pretraining is enabled but numeric data is missing.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@bonsai/modules/networks/bonsai_nets.py` around lines 219 - 232, Update
BonsaiValuePretrain.forward to access batch["numeric_target"] optionally before
indexing with mask, and raise a clear configuration error when numeric_target is
absent while value pretraining is enabled. Preserve the existing val_mask
filtering and value-logit computation when numeric data is present.
3 participants
Summary by CodeRabbit
numeric_column.value_embedding_mode.numeric_columnand numeric-training options such asvalue_loss_weight.