Add multi-task learning functionality#99
Conversation
… to _architecture and register legacy module shim The bundled multitask_model.pt was serialised when MultitaskDeepLCModel and BatchedHeads lived in a top-level module called `multitask_model`. That module no longer exists, so torch.load raised ModuleNotFoundError for any user who tried to load the default model. Fix: - Add BatchedHeads and MultitaskDeepLCModel to deeplc/_architecture.py, where the rest of the model architecture already lives. - Add _patch_legacy_multitask_module() to deeplc/_model_ops.py, which registers a sys.modules shim mapping the old import path to the new classes before torch.load is called. The shim is a no-op if the module is already registered. - Add test_load_multitask_model_without_prior_shim to tests/test_model_ops.py to prevent regression. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Fix multitask_model.pt loading (ModuleNotFoundError on import)
…d add annotations - Move MultitaskAdapter from multitask.py into _architecture.py alongside BatchedHeads and MultitaskDeepLCModel; delete multitask.py - Fix MultitaskAdapter.forward: unsqueeze(-1) -> unsqueeze(0); the guard branch produces shape [n_heads, 1] which is incompatible with nn.Linear(n_heads, ...); unsqueeze(0) gives the correct [1, n_heads]. Branch is dead code in normal inference since BatchedHeads always returns [batch, n_heads] (ndim=2). - Add type annotations to MultitaskAdapter.forward, freeze_backbone, unfreeze_backbone
_best_correlating_head() and _is_multitask_output() are pure numpy helpers specific to multitask calibration; they do not belong in core.py. Move them to a dedicated _multitask_utils.py and import from there.
…test - core.py: predict_and_calibrate used psm_list_reference (original param, may be None for auto-calibration) instead of parsed_psm_list_ref; fix both occurrences - core.py: multitask_model.pt is always the default and is always packaged; remove DEFAULT_MODEL_FALLBACK, DEFAULT_MODEL_NAME, DEFAULT_MULTITASK_MODEL_WORKSPACE and the three-way conditional — set DEFAULT_MODEL = DEFAULT_MULTITASK_MODEL_PACKAGED - test_model_ops.py: skip multitask model load test when model file is not bundled
predict() now returns [n, n_heads] with the multitask default model. Select first head before comparing raw vs calibrated predictions.
Replace MultitaskDeepLCModel + MultitaskAdapter with a single DeepLCModel that holds an optional fine-tuning MLP (attached via add_adapter()). All models now return [batch, n_heads]; core.predict() normalises to 1D at the public boundary. Checkpoints are plain state dicts loaded with weights_only=True, removing the class-path dependency on the legacy pickle format. Migrate bundled multitask_model.pt in-place to the new format. Co-authored-by: Copilot <copilot@github.com>
Adds a typed `selected_model_head: int | None` field to the `Calibration` base class to record which model head the calibration was fitted to. This replaces the ad-hoc `getattr`-based monkey-patch previously used in core.py to communicate head selection between `calibrate()` and `predict_and_calibrate()`. `predict_and_calibrate()` now raises a clear `ValueError` when a pre-fitted calibration is passed without `selected_model_head` set, instead of silently re-deriving it via an extra predict call on the reference data. Also adds missing docstrings and Pylance type narrowing fixes to calibration.py. And update ruff settings. Co-authored-by: Copilot <copilot@github.com>
|
Changes from my side:
Other |
2 participants
Adds a bundled multitask pretrained model as the new default and refactors the model architecture to support it cleanly.
Bundled multitask model
Packages a pretrained multitask model (
multitask_model.pt) trained jointly across multiple LC setups as the new default. Zero-shot predictions are immediately reasonable without retraining; fine-tuning starts from a well-initialised backbone.Unified model architecture
Consolidates
DeepLCModel,MultitaskDeepLCModel, and related legacy classes into a singleDeepLCModelthat always outputs[batch, n_heads]. An optional adapter MLP (add_adapter()) attaches on top of the head output for fine-tuning and maps to[batch, 1].freeze_backbone()/unfreeze_backbone()control which parameters train during fine-tuning. Checkpoint loading infersn_headsandfinal_num_layersfrom the state dict; legacy full-module pickles still load via a shim.Calibration and head selection
calibrate()automatically selects the output head with the highest Pearson correlation to observed retention times. The selected index is stored asCalibration.selected_model_head— a typed field on the base class — so the same head is used consistently when the fitted calibration is reused.