add functionality for continuous fine-tuning on subpopulation - #172
add functionality for continuous fine-tuning on subpopulation#172kirilklein wants to merge 12 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a subpopulation fine-tuning flow: new Azure component and pipeline, CLI entrypoint and config, dataset filtering and fold creation, trainer freeze-on-init option, directory setup flag, dataset utility to drop constant outcomes, and tests for folds and freezing behavior. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Pipeline as SUBPOP_FINETUNE_CALIBRATE_ESTIMATE<br/>Pipeline
participant Finetune as finetune_subpop<br/>Component
participant Calibrate as calibrate_exp_y<br/>Component
participant Estimate as estimate<br/>Component
participant GetStats as get_stats<br/>Component
Client->>Pipeline: trigger(prepared_data, finetune_model, subpopulation_pids, optional: counterfactual_outcomes, secondary_cohort_config)
rect rgba(100,150,200,0.5)
Note over Pipeline: select DSL variant based on optional inputs
end
Pipeline->>Finetune: run(prepared_data, restart_model=finetune_model, subpopulation_pids)
Finetune->>Finetune: create folds, drop constant outcomes, cv_loop, accumulate predictions
Finetune-->>Pipeline: model
Pipeline->>Calibrate: run(model, predictions)
Calibrate-->>Pipeline: calibrated_predictions
alt counterfactual_outcomes present
Pipeline->>Estimate: run(calibrated_predictions, counterfactual_outcomes)
else
Pipeline->>Estimate: run(calibrated_predictions)
end
Estimate-->>Pipeline: ps_calibrated_predictions
alt secondary_cohort_config present
Pipeline->>GetStats: run(ps_calibrated_predictions, secondary_cohort_config)
else
Pipeline->>GetStats: run(ps_calibrated_predictions)
end
GetStats-->>Pipeline: statistics
Pipeline-->>Client: results
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 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: 4
🧹 Nitpick comments (5)
corebehrt/main_causal/finetune_subpop.py (3)
67-68: Emptytest_datais intentional but worth documenting.An empty
CausalPatientDatasetis passed astest_datatocv_loop. Perfinetune_fold(context snippet 2), this is handled correctly viaif test_data and len(test_data) > 0. However, a brief comment would clarify the intent.📝 Proposed clarification
- # Run CV loop (loads per-fold checkpoints via restart_model) - test_data = CausalPatientDataset([], vocab) + # Run CV loop (loads per-fold checkpoints via restart_model) + # Note: No separate test set for subpop fine-tuning - only train/val splits + test_data = CausalPatientDataset([], vocab) # Empty placeholder🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/main_causal/finetune_subpop.py` around lines 67 - 68, Document that passing an empty CausalPatientDataset to cv_loop is intentional: add a brief inline comment where test_data = CausalPatientDataset([], vocab) is created explaining that test_data is intentionally empty because finetune_fold checks "if test_data and len(test_data) > 0" before using it, and cv_loop relies on that behavior; reference CausalPatientDataset, cv_loop, and finetune_fold in the comment so future readers understand the defensive check and intent.
54-58: Inconsistent config access:bootstrapat root vsn_folds/seedunderdata.
n_foldsandseedare accessed fromcfg.data, butbootstrapis accessed from the rootcfg. Looking at the YAML config (ft_subpop.yaml),bootstrap: trueis indeed at the root level whilen_folds/seedare underdata:. This works but is inconsistent.Consider moving
bootstrapunder thedata:section in the YAML for consistency, then access viadata_cfg.get("bootstrap", True).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/main_causal/finetune_subpop.py` around lines 54 - 58, The code accesses n_folds and seed from data_cfg but reads bootstrap from the root cfg, which is inconsistent; change the access to read bootstrap from the same data_cfg (use data_cfg.get("bootstrap", True)) so all data-related options use the same source; update any references to the bootstrap variable in this module (e.g., where bootstrap is used after it's defined) to rely on the new data_cfg-derived value and ensure the YAML moves bootstrap under data if you also update configuration files.
43-47: Addweights_only=Truetotorch.loadcalls for security.PyTorch 2.5.1 (your current version) emits a
FutureWarningwhentorch.loadis called withoutweights_only, as arbitrary pickle loading is a security risk—malicious pickles can execute arbitrary code during unpickling. Since these files contain only tensors/dicts, explicitly setweights_only=True.🛡️ Proposed fix
- loaded_data = torch.load(join(cfg.paths.prepared_data, PREPARED_ALL_PATIENTS)) + loaded_data = torch.load(join(cfg.paths.prepared_data, PREPARED_ALL_PATIENTS), weights_only=True) vocab = load_vocabulary(cfg.paths.prepared_data) data = CausalPatientDataset(loaded_data, vocab) - subpop_pids = torch.load(cfg.paths.subpopulation_pids) + subpop_pids = torch.load(cfg.paths.subpopulation_pids, weights_only=True)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/main_causal/finetune_subpop.py` around lines 43 - 47, The two torch.load calls that populate loaded_data and subpop_pids should pass weights_only=True to avoid unsafe pickle deserialization; update the calls that set loaded_data = torch.load(join(cfg.paths.prepared_data, PREPARED_ALL_PATIENTS)) and subpop_pids = torch.load(cfg.paths.subpopulation_pids) to include the weights_only=True kwarg so only tensor weights are loaded (no arbitrary pickles) while keeping the rest of the logic (vocab = load_vocabulary(...) and data = CausalPatientDataset(...)) unchanged.corebehrt/modules/trainer/causal/trainer.py (1)
63-64: Frozen parameters remain registered with optimizer.When
freeze_encoder_at_initis enabled,_freeze_encoder()setsrequires_grad=Falseon encoder parameters but doesn't remove them from the optimizer'sparam_groups. This means the optimizer still tracks these parameters, wasting memory and potentially causing confusion during debugging.Consider either:
- Recreating the optimizer with only unfrozen parameters, or
- Documenting this as intentional (if you want to support unfreezing later)
♻️ Option: Filter frozen params from optimizer after freezing
if self.args.get("freeze_encoder_at_init", False): self._freeze_encoder() + # Optionally rebuild optimizer with only trainable params + # self._rebuild_optimizer_for_frozen_params()Alternatively, add a comment explaining the design choice:
if self.args.get("freeze_encoder_at_init", False): + # Note: Frozen params remain in optimizer to support potential unfreezing self._freeze_encoder()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/modules/trainer/causal/trainer.py` around lines 63 - 64, The freeze_encoder_at_init path calls _freeze_encoder() which sets encoder parameters to requires_grad=False but leaves them in the optimizer.param_groups; update the code immediately after _freeze_encoder() to either recreate the optimizer with only parameters where p.requires_grad is True (e.g., rebuild optimizer from model.parameters() filtered by p.requires_grad) or explicitly remove frozen params from existing optimizer.param_groups by filtering each group to keep only params with p.requires_grad==True; reference the symbols freeze_encoder_at_init, _freeze_encoder, and optimizer when making this change so the optimizer no longer tracks frozen encoder parameters.tests/test_main_causal/test_finetune_subpop.py (1)
46-90: Tests manually replicate__init__logic instead of exercising it.These tests patch
EHRTrainer.__init__to returnNone, then manually recreate the freeze logic at lines 73-74 and 87-88. This doesn't verify that the actual__init__calls_freeze_encoder— it only tests the conditional expression in isolation.If someone refactors
__init__(e.g., moves the freeze check), these tests would still pass while the real behavior breaks.Consider testing the actual initialization path, perhaps by mocking the dependencies that
__init__requires rather than bypassing__init__entirely.♻️ Alternative test approach
class TestFreezeEncoderAtInit(unittest.TestCase): `@patch`("corebehrt.modules.trainer.causal.trainer.CausalEHRTrainer._freeze_encoder") def test_freeze_called_when_flag_set(self, mock_freeze): """Verify _freeze_encoder is called when freeze_encoder_at_init=True.""" # Create minimal mocks for actual __init__ dependencies mock_model = MagicMock() mock_model.config.outcome_names = ["outcome_1"] # Actually instantiate and let __init__ run trainer = CausalEHRTrainer( model=mock_model, optimizer=MagicMock(), train_dataset=MagicMock(), val_dataset=None, test_dataset=None, args={"freeze_encoder_at_init": True, "use_pcgrad": False, ...}, # ... other required args ) mock_freeze.assert_called_once()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_main_causal/test_finetune_subpop.py` around lines 46 - 90, The tests patch out EHRTrainer.__init__ and then duplicate the freeze logic, so they don't verify that CausalEHRTrainer.__init__ actually calls _freeze_encoder; update the tests (TestFreezeEncoderAtInit.test_freeze_called_when_flag_set and test_freeze_not_called_when_flag_not_set) to let CausalEHRTrainer.__init__ run instead of stubbing it: remove the `@patch` of EHRTrainer.__init__, supply minimal mocked dependencies (e.g., a MagicMock model with config.outcome_names, mock optimizer/datasets, and args containing or omitting "freeze_encoder_at_init") when instantiating CausalEHRTrainer so the real __init__ executes, and delete the manual if trainer.args... trainer._freeze_encoder() lines; keep patching CausalEHRTrainer._freeze_encoder to assert it was (or was not) called.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@corebehrt/azure/components/finetune_subpop.py`:
- Line 1: The file fails formatting (ruff) — reformat
corebehrt/azure/components/finetune_subpop.py so it passes ruff format checks;
specifically ensure the import line "from corebehrt.azure.util import job" and
the file overall follow project style (run `ruff format` on the file or apply
your formatter) and then re-run linting to confirm the file is no longer
reported as "would be reformatted".
In `@corebehrt/configs/causal/finetune/ft_subpop.yaml`:
- Line 34: The config has mismatched epoch counts: the top-level key epochs
(currently 3) does not match num_training_epochs (currently 5) used by the
scheduler; update one so both match (either set epochs to 5 to match
num_training_epochs or change num_training_epochs to 3) — edit the keys "epochs"
and/or "num_training_epochs" in ft_subpop.yaml so they are identical and
consistent with the intended training length.
In `@corebehrt/main_causal/finetune_subpop.py`:
- Line 1: The file failed ruff formatting; run the code formatter on
corebehrt/main_causal/finetune_subpop.py (e.g., run `ruff format` or your repo's
formatter) and commit the reformat; ensure any top-level docstring/quotes and
surrounding imports/whitespace are normalized so `ruff format --check` passes;
no logic changes needed—only apply the automatic formatting to
finetune_subpop.py.
In `@tests/test_main_causal/test_finetune_subpop.py`:
- Line 1: The file contains an unformatted import statement ("import unittest")
causing ruff format --check to fail; fix by running ruff format on that test
file (or applying ruff's suggested changes) so whitespace, import formatting and
trailing newline are corrected and the file passes ruff format checks.
---
Nitpick comments:
In `@corebehrt/main_causal/finetune_subpop.py`:
- Around line 67-68: Document that passing an empty CausalPatientDataset to
cv_loop is intentional: add a brief inline comment where test_data =
CausalPatientDataset([], vocab) is created explaining that test_data is
intentionally empty because finetune_fold checks "if test_data and
len(test_data) > 0" before using it, and cv_loop relies on that behavior;
reference CausalPatientDataset, cv_loop, and finetune_fold in the comment so
future readers understand the defensive check and intent.
- Around line 54-58: The code accesses n_folds and seed from data_cfg but reads
bootstrap from the root cfg, which is inconsistent; change the access to read
bootstrap from the same data_cfg (use data_cfg.get("bootstrap", True)) so all
data-related options use the same source; update any references to the bootstrap
variable in this module (e.g., where bootstrap is used after it's defined) to
rely on the new data_cfg-derived value and ensure the YAML moves bootstrap under
data if you also update configuration files.
- Around line 43-47: The two torch.load calls that populate loaded_data and
subpop_pids should pass weights_only=True to avoid unsafe pickle
deserialization; update the calls that set loaded_data =
torch.load(join(cfg.paths.prepared_data, PREPARED_ALL_PATIENTS)) and subpop_pids
= torch.load(cfg.paths.subpopulation_pids) to include the weights_only=True
kwarg so only tensor weights are loaded (no arbitrary pickles) while keeping the
rest of the logic (vocab = load_vocabulary(...) and data =
CausalPatientDataset(...)) unchanged.
In `@corebehrt/modules/trainer/causal/trainer.py`:
- Around line 63-64: The freeze_encoder_at_init path calls _freeze_encoder()
which sets encoder parameters to requires_grad=False but leaves them in the
optimizer.param_groups; update the code immediately after _freeze_encoder() to
either recreate the optimizer with only parameters where p.requires_grad is True
(e.g., rebuild optimizer from model.parameters() filtered by p.requires_grad) or
explicitly remove frozen params from existing optimizer.param_groups by
filtering each group to keep only params with p.requires_grad==True; reference
the symbols freeze_encoder_at_init, _freeze_encoder, and optimizer when making
this change so the optimizer no longer tracks frozen encoder parameters.
In `@tests/test_main_causal/test_finetune_subpop.py`:
- Around line 46-90: The tests patch out EHRTrainer.__init__ and then duplicate
the freeze logic, so they don't verify that CausalEHRTrainer.__init__ actually
calls _freeze_encoder; update the tests
(TestFreezeEncoderAtInit.test_freeze_called_when_flag_set and
test_freeze_not_called_when_flag_not_set) to let CausalEHRTrainer.__init__ run
instead of stubbing it: remove the `@patch` of EHRTrainer.__init__, supply minimal
mocked dependencies (e.g., a MagicMock model with config.outcome_names, mock
optimizer/datasets, and args containing or omitting "freeze_encoder_at_init")
when instantiating CausalEHRTrainer so the real __init__ executes, and delete
the manual if trainer.args... trainer._freeze_encoder() lines; keep patching
CausalEHRTrainer._freeze_encoder to assert it was (or was not) called.
🪄 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
Run ID: 012652ca-45b9-4ca7-a9da-072ca6916aaa
📒 Files selected for processing (7)
corebehrt/azure/components/finetune_subpop.pycorebehrt/azure/pipelines/SUBPOP_FINETUNE_CALIBRATE_ESTIMATE.pycorebehrt/azure/pipelines/__init__.pycorebehrt/configs/causal/finetune/ft_subpop.yamlcorebehrt/main_causal/finetune_subpop.pycorebehrt/modules/trainer/causal/trainer.pytests/test_main_causal/test_finetune_subpop.py
| batch_size: 128 | ||
| val_batch_size: 256 | ||
| effective_batch_size: 128 | ||
| epochs: 3 |
There was a problem hiding this comment.
Mismatch: epochs: 3 but num_training_epochs: 5.
The trainer will run for 3 epochs (line 34), but the scheduler is configured for 5 epochs (line 59). This means the learning rate schedule won't complete — training ends before reaching the final scheduled LR.
🐛 Proposed fix: Align epoch counts
trainer_args:
# ...
- epochs: 3
+ epochs: 5
Or alternatively:
scheduler:
_target_: transformers.get_linear_schedule_with_warmup
- num_training_epochs: 5
+ num_training_epochs: 3
num_warmup_epochs: 1Also applies to: 57-60
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@corebehrt/configs/causal/finetune/ft_subpop.yaml` at line 34, The config has
mismatched epoch counts: the top-level key epochs (currently 3) does not match
num_training_epochs (currently 5) used by the scheduler; update one so both
match (either set epochs to 5 to match num_training_epochs or change
num_training_epochs to 3) — edit the keys "epochs" and/or "num_training_epochs"
in ft_subpop.yaml so they are identical and consistent with the intended
training length.
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/format.yml:
- Around line 15-16: The workflow uses outdated action versions; update the
actions referenced in the workflow by replacing the usages of
actions/checkout@v3 and actions/setup-python@v4 with the current stable releases
(e.g., actions/checkout@v6.0.2 and actions/setup-python@v6.2.0) so the job
definitions that call these actions use the newer tags; ensure any inputs or
behavior differences are accounted for by verifying the checkout and
setup-python steps still pass after updating.
In `@corebehrt/main_causal/finetune_subpop.py`:
- Around line 55-60: The call to create_folds ignores data.val_ratio in
single-fold mode, so when n_folds == 1 the configured validation split isn't
passed through; update the code around data_cfg, n_folds, seed, bootstrap to
read val_ratio = data_cfg.get("val_ratio", <default>) and pass
val_ratio=val_ratio into create_folds (i.e., create_folds(train_val_pids,
n_folds, seed, bootstrap=bootstrap, val_ratio=val_ratio)) so create_folds
receives the intended validation proportion in single-fold mode.
- Around line 50-61: Before calling create_folds/validate_folds add explicit
precondition checks on the filtered subpopulation: after computing
train_val_pids (from data.filter_by_pids) check that len(train_val_pids) > 0 and
that len(train_val_pids) >= n_folds (n_folds comes from cfg.get("data") as used
above); if not, raise a clear ValueError or log an error via logger with a
descriptive message including len(train_val_pids) and n_folds (and mention
bootstrap mode if cfg.get("bootstrap") affects behavior), so create_folds and
validate_folds are only called when the subpopulation size meets fold
requirements.
🪄 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
Run ID: f3333f4a-22e0-4f47-adbf-387be7d22bf9
📒 Files selected for processing (4)
.github/workflows/format.ymlcorebehrt/azure/components/finetune_subpop.pycorebehrt/main_causal/finetune_subpop.pytests/test_main_causal/test_finetune_subpop.py
✅ Files skipped from review due to trivial changes (1)
- corebehrt/azure/components/finetune_subpop.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tests/test_main_causal/test_finetune_subpop.py
| - uses: actions/checkout@v3 | ||
| - uses: actions/setup-python@v4 |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check the latest releases for GitHub Actions used in this workflow
echo "Checking latest release for actions/checkout:"
gh api repos/actions/checkout/releases/latest --jq '.tag_name'
echo -e "\nChecking latest release for actions/setup-python:"
gh api repos/actions/setup-python/releases/latest --jq '.tag_name'Repository: kirilklein/PHAIR_EHR
Length of output: 464
🌐 Web query:
actions/checkout latest version 2025 2026
💡 Result:
The latest version of the GitHub Actions actions/checkout is v6.0.2, released on 2026-01-09. This is the current version as of 2026-04-01. Key recent releases: - v6.0.2 (2026-01-09): Fix tag handling: preserve annotations and explicit fetch-tags. - v6.0.1 (late 2025): Minor updates including worktree support and README clarifications. - v6.0.0 (late 2025): Improved credential security (persist-credentials to separate file), Node.js 24 support. - v5.0.0 (2025-08): Updated to Node.js 24 runtime. Usage example: - uses: actions/checkout@v6 Official sources confirm v6.0.2 as latest, with last repo push on 2026-02-03 but no newer release by April 2026.
Citations:
- 1: https://github.com/actions/checkout/releases
- 2: https://github.com/actions/checkout/blob/main/CHANGELOG.md
- 3: https://github.com/actions/checkout
- 4: actions/checkout@v5.0.1...v6.0.2
- 5: https://newreleases.io/project/github/actions/checkout/release/v6.0.1
🌐 Web query:
actions/setup-python latest version 2025 2026
💡 Result:
The latest version of the actions/setup-python GitHub Action is v6.2.0, released on January 22, 2026. This is confirmed directly from the official repository homepage, which states "v6.2.0 Latest Jan 22, 2026". Usage examples in the docs and marketplace also reference uses: actions/setup-python@v6. Recent activity includes commits up to January 22, 2026, such as dependency updates. Earlier 2025 releases include v6.0.0 in September 2025 with breaking changes (Node 20 to Node 24 upgrade, requiring runner v2.327.1+), and various enhancements like pip-version support.
Citations:
- 1: https://github.com/actions/setup-python
- 2: https://github.com/actions/setup-python/releases
- 3: https://newreleases.io/project/github/actions/setup-python/release/v6.0.0
- 4: actions/setup-python@a26af69...a309ff8
- 5: https://github.com/marketplace/actions/setup-python
Update GitHub Actions to current stable versions.
The workflow uses outdated action versions. actions/checkout@v3 (released in 2022) and actions/setup-python@v4 (released in 2023) are significantly behind current stable versions and may cause compatibility issues with modern GitHub Actions runners.
Current stable versions as of April 2026:
actions/checkout: v6.0.2 (latest)actions/setup-python: v6.2.0 (latest)
Update to the current stable versions:
Proposed fix
- - uses: actions/checkout@v3
- - uses: actions/setup-python@v4
+ - uses: actions/checkout@v6
+ - uses: actions/setup-python@v6📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| - uses: actions/checkout@v3 | |
| - uses: actions/setup-python@v4 | |
| - uses: actions/checkout@v6 | |
| - uses: actions/setup-python@v6 |
🧰 Tools
🪛 actionlint (1.7.11)
[error] 15-15: the runner of "actions/checkout@v3" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
[error] 16-16: the runner of "actions/setup-python@v4" action is too old to run on GitHub Actions. update the action's version to fix this issue
(action)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In @.github/workflows/format.yml around lines 15 - 16, The workflow uses
outdated action versions; update the actions referenced in the workflow by
replacing the usages of actions/checkout@v3 and actions/setup-python@v4 with the
current stable releases (e.g., actions/checkout@v6.0.2 and
actions/setup-python@v6.2.0) so the job definitions that call these actions use
the newer tags; ensure any inputs or behavior differences are accounted for by
verifying the checkout and setup-python steps still pass after updating.
| data = data.filter_by_pids(subpop_pids) | ||
| train_val_pids = data.get_pids() | ||
| logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data") | ||
|
|
||
| # Create fresh bootstrap folds from subpopulation | ||
| data_cfg = cfg.get("data", {}) | ||
| n_folds = data_cfg.get("n_folds", 5) | ||
| seed = data_cfg.get("seed", 42) | ||
| bootstrap = cfg.get("bootstrap", True) | ||
|
|
||
| folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) | ||
| validate_folds(folds, set(train_val_pids), logger, bootstrap=bootstrap) |
There was a problem hiding this comment.
Add explicit subpopulation/fold precondition checks before creating folds.
If filtering leaves too few patients, fold creation can fail at runtime (e.g., KFold with n_folds > len(train_val_pids)), and the failure message will be less actionable than a direct validation error.
Proposed fix
data = data.filter_by_pids(subpop_pids)
train_val_pids = data.get_pids()
logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data")
# Create fresh bootstrap folds from subpopulation
data_cfg = cfg.get("data", {})
n_folds = data_cfg.get("n_folds", 5)
seed = data_cfg.get("seed", 42)
bootstrap = cfg.get("bootstrap", True)
+
+ if not train_val_pids:
+ raise ValueError(
+ "No matching subpopulation patients found in prepared data. "
+ "Check cfg.paths.subpopulation_pids and prepared dataset alignment."
+ )
+ if n_folds < 1:
+ raise ValueError(f"n_folds must be >= 1, got {n_folds}")
+ if n_folds > len(train_val_pids):
+ raise ValueError(
+ f"n_folds ({n_folds}) cannot exceed number of filtered patients "
+ f"({len(train_val_pids)})."
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| data = data.filter_by_pids(subpop_pids) | |
| train_val_pids = data.get_pids() | |
| logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data") | |
| # Create fresh bootstrap folds from subpopulation | |
| data_cfg = cfg.get("data", {}) | |
| n_folds = data_cfg.get("n_folds", 5) | |
| seed = data_cfg.get("seed", 42) | |
| bootstrap = cfg.get("bootstrap", True) | |
| folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) | |
| validate_folds(folds, set(train_val_pids), logger, bootstrap=bootstrap) | |
| data = data.filter_by_pids(subpop_pids) | |
| train_val_pids = data.get_pids() | |
| logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data") | |
| # Create fresh bootstrap folds from subpopulation | |
| data_cfg = cfg.get("data", {}) | |
| n_folds = data_cfg.get("n_folds", 5) | |
| seed = data_cfg.get("seed", 42) | |
| bootstrap = cfg.get("bootstrap", True) | |
| if not train_val_pids: | |
| raise ValueError( | |
| "No matching subpopulation patients found in prepared data. " | |
| "Check cfg.paths.subpopulation_pids and prepared dataset alignment." | |
| ) | |
| if n_folds < 1: | |
| raise ValueError(f"n_folds must be >= 1, got {n_folds}") | |
| if n_folds > len(train_val_pids): | |
| raise ValueError( | |
| f"n_folds ({n_folds}) cannot exceed number of filtered patients " | |
| f"({len(train_val_pids)})." | |
| ) | |
| folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) | |
| validate_folds(folds, set(train_val_pids), logger, bootstrap=bootstrap) |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@corebehrt/main_causal/finetune_subpop.py` around lines 50 - 61, Before
calling create_folds/validate_folds add explicit precondition checks on the
filtered subpopulation: after computing train_val_pids (from
data.filter_by_pids) check that len(train_val_pids) > 0 and that
len(train_val_pids) >= n_folds (n_folds comes from cfg.get("data") as used
above); if not, raise a clear ValueError or log an error via logger with a
descriptive message including len(train_val_pids) and n_folds (and mention
bootstrap mode if cfg.get("bootstrap") affects behavior), so create_folds and
validate_folds are only called when the subpopulation size meets fold
requirements.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
corebehrt/modules/setup/directory.py (1)
391-397:⚠️ Potential issue | 🟠 MajorInconsistency:
write_configstill referencespretrain_modelwhen directory check is disabled.When
check_pretrain=False, line 396 still attempts to copy the config from thepretrain_modeldirectory viawrite_config("model", source="pretrain_model", ...). Thecheck_directorycall at line 391 is skipped, butwrite_configwill subsequently callcopyfile()on a source that may not exist, causing a runtime failure.Make the
write_configcall conditional on the same flag:Proposed fix
if check_pretrain: self.check_directory("pretrain_model") self.create_run_directory("model", base="runs") # Write config in output directory. - self.write_config("model", source="pretrain_model", name=PRETRAIN_CFG) + if check_pretrain: + self.write_config("model", source="pretrain_model", name=PRETRAIN_CFG) self.write_config("model", name=FINETUNE_CFG)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/modules/setup/directory.py` around lines 391 - 397, The write_config call currently always tries to copy PRETRAIN_CFG from the "pretrain_model" source even when check_pretrain is False; update the logic so that after calling create_run_directory("model", base="runs") you only call write_config("model", source="pretrain_model", name=PRETRAIN_CFG) when check_pretrain is True (the same condition used for calling check_directory("pretrain_model")), otherwise skip that source-based write and only call write_config("model", name=FINETUNE_CFG); this ensures check_pretrain, check_directory, write_config, PRETRAIN_CFG and "pretrain_model" are consistent and avoids copyfile on a missing source.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@corebehrt/modules/setup/directory.py`:
- Around line 391-397: The write_config call currently always tries to copy
PRETRAIN_CFG from the "pretrain_model" source even when check_pretrain is False;
update the logic so that after calling create_run_directory("model",
base="runs") you only call write_config("model", source="pretrain_model",
name=PRETRAIN_CFG) when check_pretrain is True (the same condition used for
calling check_directory("pretrain_model")), otherwise skip that source-based
write and only call write_config("model", name=FINETUNE_CFG); this ensures
check_pretrain, check_directory, write_config, PRETRAIN_CFG and "pretrain_model"
are consistent and avoids copyfile on a missing source.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d1d175be-6bce-43c7-92e2-d1354d42e13d
📒 Files selected for processing (1)
corebehrt/modules/setup/directory.py
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
corebehrt/main_causal/finetune_subpop.py (2)
62-67:⚠️ Potential issue | 🟡 MinorPass
data.val_ratiothrough tocreate_foldsto preserve configured split behavior.On Line 67,
create_folds(...)currently ignoresdata.val_ratio, which can diverge from config whenn_folds == 1.Proposed fix
data_cfg = cfg.get("data", {}) n_folds = data_cfg.get("n_folds", 5) seed = data_cfg.get("seed", 42) + val_ratio = data_cfg.get("val_ratio", 0.8) bootstrap = cfg.get("bootstrap", True) - folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) + folds = create_folds( + train_val_pids, + n_folds, + seed, + val_ratio=val_ratio, + bootstrap=bootstrap, + )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/main_causal/finetune_subpop.py` around lines 62 - 67, The call to create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap) is ignoring the configured validation split; read data_cfg.get("val_ratio", <default>) from the same cfg block (data_cfg variable) and pass it into create_folds (e.g., include val_ratio=val_ratio) so that when n_folds == 1 the configured validation proportion is preserved; update the invocation near create_folds and ensure the val_ratio symbol is defined from data_cfg before the call.
58-68:⚠️ Potential issue | 🟠 MajorAdd explicit filtered-subpopulation/fold precondition checks before fold creation.
On Line 67, fold creation can fail with unclear errors when
train_val_pidsis empty orn_foldsis invalid relative to filtered size. Add explicit checks right beforecreate_folds(...)for clearer failures.Proposed fix
train_val_pids = data.get_pids() logger.info(f"Filtered to {len(train_val_pids)} patients in prepared data") # Create fresh bootstrap folds from subpopulation data_cfg = cfg.get("data", {}) n_folds = data_cfg.get("n_folds", 5) seed = data_cfg.get("seed", 42) bootstrap = cfg.get("bootstrap", True) + + if not train_val_pids: + raise ValueError( + "No matching subpopulation patients found in prepared data. " + "Check cfg.paths.subpopulation_pids and prepared dataset alignment." + ) + if n_folds < 1: + raise ValueError(f"n_folds must be >= 1, got {n_folds}") + if n_folds > len(train_val_pids): + raise ValueError( + f"n_folds ({n_folds}) cannot exceed number of filtered patients " + f"({len(train_val_pids)})." + ) folds = create_folds(train_val_pids, n_folds, seed, bootstrap=bootstrap)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@corebehrt/main_causal/finetune_subpop.py` around lines 58 - 68, The code may call create_folds with an empty or inappropriate fold count causing unclear errors; before calling create_folds(validate_folds), check that train_val_pids (from data.get_pids()) is not empty and that n_folds is a positive integer and reasonably <= len(train_val_pids) (or if cfg.get("bootstrap") is True, document/validate the allowed relation) — if these preconditions fail, log a clear error via logger.error (including values of len(train_val_pids), n_folds, and bootstrap) and raise a ValueError with a descriptive message so callers see an explicit failure instead of a downstream cryptic one; update the block that sets data_cfg, n_folds, seed, bootstrap and calls create_folds/validate_folds to perform these checks first.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@corebehrt/modules/preparation/causal/dataset.py`:
- Around line 87-111: The method drop_constant_outcomes has formatting issues
causing ruff format --check to fail; run an auto-formatter (e.g., ruff format)
on the file or at least reformat this method so spacing and line wrapping meet
project style (ensure consistent indentation and wrapping around the list
comprehension for varying, the f-string in the ValueError, and the for-loop that
rebuilds p.outcomes for patients), then re-run ruff format --check to confirm;
key symbols to touch: drop_constant_outcomes, get_outcomes, outcomes, varying,
dropped, and p.outcomes.
---
Duplicate comments:
In `@corebehrt/main_causal/finetune_subpop.py`:
- Around line 62-67: The call to create_folds(train_val_pids, n_folds, seed,
bootstrap=bootstrap) is ignoring the configured validation split; read
data_cfg.get("val_ratio", <default>) from the same cfg block (data_cfg variable)
and pass it into create_folds (e.g., include val_ratio=val_ratio) so that when
n_folds == 1 the configured validation proportion is preserved; update the
invocation near create_folds and ensure the val_ratio symbol is defined from
data_cfg before the call.
- Around line 58-68: The code may call create_folds with an empty or
inappropriate fold count causing unclear errors; before calling
create_folds(validate_folds), check that train_val_pids (from data.get_pids())
is not empty and that n_folds is a positive integer and reasonably <=
len(train_val_pids) (or if cfg.get("bootstrap") is True, document/validate the
allowed relation) — if these preconditions fail, log a clear error via
logger.error (including values of len(train_val_pids), n_folds, and bootstrap)
and raise a ValueError with a descriptive message so callers see an explicit
failure instead of a downstream cryptic one; update the block that sets
data_cfg, n_folds, seed, bootstrap and calls create_folds/validate_folds to
perform these checks first.
🪄 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
Run ID: c70c83fb-2869-460b-a84c-a5abb9c82d37
📒 Files selected for processing (2)
corebehrt/main_causal/finetune_subpop.pycorebehrt/modules/preparation/causal/dataset.py
| def drop_constant_outcomes(self) -> List[str]: | ||
| """Remove outcome columns that take only one value across all patients. | ||
|
|
||
| Returns: | ||
| Names of outcomes that were dropped (empty if none). | ||
|
|
||
| Raises: | ||
| ValueError: If every outcome would be dropped (no varying labels left). | ||
| """ | ||
| if not self.patients: | ||
| return [] | ||
| outcomes = self.get_outcomes() | ||
| varying = [ | ||
| name for name, vals in outcomes.items() if len(set(vals)) > 1 | ||
| ] | ||
| dropped = [name for name in outcomes if name not in varying] | ||
| if not varying: | ||
| raise ValueError( | ||
| "After filtering, no outcome has both classes (0 and 1). " | ||
| f"All outcomes are constant: dropped candidates {dropped}" | ||
| ) | ||
| if dropped: | ||
| for p in self.patients: | ||
| p.outcomes = {k: p.outcomes[k] for k in varying} | ||
| return dropped |
There was a problem hiding this comment.
Fix formatting in this changed block to unblock CI.
ruff format --check is failing for this file, so this method block should be auto-formatted before merge (e.g., run ruff format corebehrt/modules/preparation/causal/dataset.py).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@corebehrt/modules/preparation/causal/dataset.py` around lines 87 - 111, The
method drop_constant_outcomes has formatting issues causing ruff format --check
to fail; run an auto-formatter (e.g., ruff format) on the file or at least
reformat this method so spacing and line wrapping meet project style (ensure
consistent indentation and wrapping around the list comprehension for varying,
the f-string in the ValueError, and the for-loop that rebuilds p.outcomes for
patients), then re-run ruff format --check to confirm; key symbols to touch:
drop_constant_outcomes, get_outcomes, outcomes, varying, dropped, and
p.outcomes.
|
@Montgomeryyyy we can skip calibration for outcomes entirely, only ps needs to be calibrated. I can implement this quickly. Depending how many jobs you have run already, it could just be a separate script |
I have mostly run the cvd and diab sub cohorts, and this does not take long. So if you can implement if quickly that would be nice:) |
2 participants
Summary by CodeRabbit
New Features
Tests
Chores