Add model_embedding_dim argument to Dataset constructors

[willdumm]

This PR makes DXSM datasets specific to the embedding dimension of the model that the dataset is intended for.
Minimal changes in dnsm-experiments-1 were required, and these are in https://github.com/matsengrp/dnsm-experiments-1/pull/78.

Here are some discussion points:

• Should the token insertion steps be moved from the load_pcp_df implementation to the dataset constructor? Right now, we do any token insertion (and heavy/light splicing, etc) any time we load a pcp dataframe, then the Dataset constructor strips out any unrecognized tokens.
• Should we do anything to ensure that sequences containing X’s aren’t presented to old models that don’t support them? (models have always had an embedding dimension dedicated to X’s)
• How do we prevent old models that don’t support them from being presented with paired data? If they were, the new dataset construction code would just strip out the separator token.
• If we give a crepe an aa sequence directly, it may not have the scaffolding tokens that the model has been trained to expect, unless the sequence has been moved through the load_pcp_df pipeline
• What does the src_key_padding_mask actually do? Where is the documentation for the torch Embedding forward function? It does not seem to exempt sites from being looked up in the embedding, anyway. I just want to make sure that we shouldn’t be unmasking ambiguous sites before presenting them to the model, like we do with other tokens.

[willdumm]

To recap our conversation just now about this:

• If we give a crepe an aa sequence directly, it may not have the scaffolding tokens that the model has been trained to expect, unless the sequence has been moved through the load_pcp_df pipeline
• Should the token insertion steps be moved from the load_pcp_df implementation to the dataset constructor? Right now, we do any token insertion (and heavy/light splicing, etc) any time we load a pcp dataframe, then the Dataset constructor strips out any unrecognized tokens.

Yes, we've decided so. I'll backup the current working status of this branch (and the corresponding dnsm-experiments-1 branch) to wd-token-fix-bu, then implement the following here:

• load_pcp_df and associated functions will load pcp_df without inserting any special token scaffolding into the sequences. We will always maintain separate heavy and light-chain columns wherever applicable.
• There will be a free function that takes perhaps pairs of heavy- and light- chain sequences and scaffolds them with special tokens so they can be presented to the model. This function will take a known_token_count so that it knows how to process the sequences.
• This free function will be called by the Dataset constructor, which will also still need to accept the known_token_count parameter.
• Calling the model forward/represent functions directly (or through model.call) will require the user to do any sequence token scaffolding on their own, perhaps using the free function mentioned above
• Calling the Crepe.__call__ function on sequences will do the required scaffolding automatically. Perhaps we'll have to strip out the model predictions for special token sites, so the outputs match the input sequence lengths? Erick suggests input format of something like crepe([(heavy, None), (heavy1, light1), ...]) to allow heavy and light chain sequences to be passed to the crepe for proper scaffolding.

• Should we do anything to ensure that sequences containing X’s aren’t presented to old models that don’t support them? (models have always had an embedding dimension dedicated to X’s)
• How do we prevent old models that don’t support them from being presented with paired data? If they were, the new dataset construction code would just strip out the separator token.

We'll just use pretrained crepe names to make this clear.

• What does the src_key_padding_mask actually do? Where is the documentation for the torch Embedding forward function? It does not seem to exempt sites from being looked up in the embedding, anyway. I just want to make sure that we shouldn’t be unmasking ambiguous sites before presenting them to the model, like we do with other tokens.

We understand that although all sites in the input are included in the embedding, the model ignores sites masked sites while keeping track of the fact that they take up space. We may want to try unmasking X's that aren't just for padding here, as a separate issue.

[matsen]

I was able to take a quick look (half-attention as attending conference)...

[matsen]

Is it not used for branch length optimization?

If not, let's cut it!

[willdumm]

It's used in a couple of tests, so I prepended an underscore to its name instead of removing it.

[matsen]

I think that we decided on OR.

[willdumm]

I am still waiting for all the notebooks to run, and I expect that changes will be needed to make some of them run. However, all tests pass, and the test snakemake configs run for both dnsm and dasm.

[matsen]

🎉