torchcell.models.esm2.Esm2

class Esm2(model_name: str)[source]

Bases: PeptideModel

Facebook ESM-2 masked-LM wrapper that embeds protein sequences.

property max_sequence_size: int

Return the maximum supported residue sequence length.

load_model(model_name: str) → None[source]

Download if needed, then load the tokenizer and model onto the device.

embed(sequences: str | list[str], mean_embedding: bool = False) → Tensor[source]

Return last-hidden-state embeddings for one or more sequences.

Parameters:

  • sequences – A single sequence or list of sequences to embed.

  • mean_embedding – If True, return attention-masked mean-pooled embeddings instead of per-residue hidden states.

Returns:

Tensor of embeddings on the model device.