Use ESM2 computed representations as starting point for ESM Fold

[kentwait]

TL;DR I already have ESM2 embeddings. Can you use ESM2 embeddings as input to ESM Fold structure prediction instead of sequence?

In the sample code for using esmfold, a string of the amino acid sequence was used as input to predict the structure.

...
with torch.no_grad():
    output = model.infer_pdb(sequence)
...

I looked into the code for infer_pdb which lead me ultimately to the forward method where the sequence is converted to embeddings and then used for structure prediction.

...
# === ESM ===
esmaa = self._af2_idx_to_esm_idx(aa, mask)

if masking_pattern is not None:
    esmaa = self._mask_inputs_to_esm(esmaa, masking_pattern)

esm_s = self._compute_language_model_representations(esmaa)

# Convert esm_s to the precision used by the trunk and
# the structure module. These tensors may be a lower precision if, for example,
# we're running the language model in fp16 precision.
esm_s = esm_s.to(self.esm_s_combine.dtype)
...

Since self._compute_language_model_representations computes the embeddings, can I hijack this function and rewrite it to use per-site ESM2 embeddings I already have precomputed instead of starting from the amino acid sequence? Please let me know if I missed something.

Thanks

[tomsercu]

Yes sounds like you should absolutely monkeypatch that function to load up your embeddings. (will disable the masking_pattern option though to randomize the embeddings a bit).
Just FYI the ESM forward pass is actually contributing significantely less to the overall compute than the folding trunk and structure module.

[kentwait]

Got it. Thanks for the help