Batch Support for Obtaining Residue Embeddings

[Junseok0207]

I am currently trying to obtain residue embeddings for protein sequences. The typical workflow involves the following steps:

protein = ESMProtein(sequence=sequence)
protein_tensor = self.model.encode(protein)
config = SamplingConfig(return_per_residue_embeddings=True)
output = client.forward_and_sample(protein_tensor, config)
embeddings = output.per_residue_embedding

However, I don't know how to get embeddings in batch mode. I checked the example in esm/examples/local_generate.py (lines 129-135), but it only shows the batch_generate function, which does not include a way to obtain embeddings. How can I achieve embeddings with batch?