allow partial loading for pre trained ms2 models

[mo-sameh]

This PR include two main changes:

• Allow to dynamically change the charged_frag_types used for ms2 hence changing the dimensions of the output layer for the ms2 model.
• Allow for partially loading the ms2 model so that a user can benefit from a pretrained model backbone when using a different prediction head such as changing the fragment types. With this feature users can finetune/train a model with more fragment types than the ones used from for pretraining. From my experiments using a pretrained backbone had a significant performance gap compared to training fully from scratch.

From an alphaDia finetuning experiment where I want to predict the following fragment types ['b','y','c','a','x','z']

• Performance when training for 50 epochs from scratch:

Model tested on test dataset with the following metrics:

l1_loss                       : 0.0228
PCC-mean                      : 0.8015
COS-mean                      : 0.8140
SA-mean                       : 0.6054
SPC-mean                      : -0.2333

• Performance when training for 50 epochs when loading a pretrained backbone:

Model tested on test dataset with the following metrics:

l1_loss                       : 0.0122
PCC-mean                      : 0.9462
COS-mean                      : 0.9491
SA-mean                       : 0.7961
SPC-mean                      : -0.3072

[jalew188]

There may be two cases:

1. we save partial model params into the file and this method allows us to load this partial model.
2. we save the full model but we only need to load partial params by specifying param names or first Kth layers.

[mo-sameh]

Hi, I’m not sure I fully understand your concerns. The current interface saves the complete model weights, so I’m unclear about why users would want to save only partial models. Are you suggesting we add this functionality?

As for loading partial weights, the current implementation(in this PR) should already handle this automatically. It matches parameter keys and sizes, loading the matching weights while initializing the remaining parameters from scratch. Are you suggesting we modify the interface to let users explicitly specify which layers to load?

[jalew188]

No, I mean the use cases.

[mschwoer]

(nit) This more concise format might be easier on the eye:

    warning_msg += "".join(
        [
            f"\nKeys with size mismatches: {size_mismatches}" if size_mismatches else "",
            f"\nUnexpected keys: {unexpected_keys}" if unexpected_keys else "",
            f"\nMissing keys: {missing_keys}" if missing_keys else "",
        ]
    )

[mschwoer]

we still have requires-python = ">=3.8.0" .. so Optional[List[str]] it is ..

@jalew188 should we drop support for 3.8?

[jalew188]

We can, in the next release.