Skip to content

ESMFold2: fold() cannot take precomputed LM hidden states #420

Open

Description

@daylight-00

EsmFold2Model.forward accepts lm_hidden_states, but ESMFold2InputBuilder.fold() has a fixed keyword list and never forwards them. To reuse one ESMC pass across several folds (seeds, sampler settings, ensembles), or to fold with a trunk whose ESMC lives on another device or process, callers have to drop to prepare_input + forward + decode by hand.

A related gap: when the model has no ESMC attached (model.esmc is None, e.g. from_pretrained(..., load_esmc=False)) and no states are given, fold() skips the LM pathway without any signal and returns the prediction of a different computation than the trained model.

PR #415 adds lm_hidden_states to fold() with input validation, makes fold() raise when it has neither an ESMC nor states, and adds a public compute_lm_hidden_states() for computing them in a separate step. The hard error is a behaviour change; an explicit opt-in is easy to add if you would prefer that. I would be glad to reshape the API to whatever fits best.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions