跳转至

ESM Cambrian (ESMC)

Pre-trained model on protein sequences using a masked language modeling (MLM) objective.

Disclaimer

This is an UNOFFICIAL implementation of the Language Modeling Materializes a World Model of Protein Biology by Salvatore Candido, Thomas Hayes, Alexander Derry, Roshan Rao, Zeming Lin, Alexander Rives, et al.

The OFFICIAL repository of ESMC is at Biohub/esm.

Tip

The MultiMolecule team has confirmed that the provided model and checkpoints are producing the same intermediate representations as the original implementation.

The team releasing ESMC did not write this model card for this model so this model card has been written by the MultiMolecule team.

Model Details

ESM Cambrian is an encoder-only protein language model trained on billions of protein sequences in a self-supervised fashion. The model learns protein representations by predicting masked amino-acid tokens from sequence context. Please refer to the Training Details section for more information on the training process.

Variants

Model Specification

Variants Num Layers Hidden Size Num Heads Intermediate Size Num Parameters (M) FLOPs (G) MACs (G) Max Num Tokens
ESMC-6B 80 2560 40 6912 6351.87 6716.35 3355.44 2048
ESMC-600M 36 1152 18 3072 574.97 631.66 315.28
ESMC-300M 30 960 15 2560 332.95 370.71 184.97
  • Code: multimolecule.esmc
  • Data: UniRef, MGnify, and Joint Genome Institute
  • Paper: Language Modeling Materializes a World Model of Protein Biology
  • Developed by: Salvatore Candido, Thomas Hayes, Alexander Derry, Roshan Rao, Zeming Lin, Robert Verkuil, Bryan Wu, Jin Sub Lee, Elise S. Bruguera, Jehan A. Keval, Mykhailo Kopylov, John E. Pak, Wesley Wu, Neil Thomas, Samson Mataraso, Alvin Hsu, Ashton C. Trotman-Grant, Kilian Fatras, Allan dos Santos Costa, Rohil Badkundri, Halil Akin, Deniz Oktay, Jonathan Deaton, Elizabeth Montabana, Hrishita Sitwala, Yue Yu, Marius Wiggert, Dylan Alexander Carlin, Anthony W. Goering, Tomasz Blazejewski, McCullen Sandora, Michael Hla, Tina Z. Jia, Leon H. Kloker, Nicholas J. Sofroniew, Masatoshi Uehara, Jassi Pannu, Sharrol Bachas, Daniel S. Liu, Tom Sercu, Alexander Rives
  • Model type: Encoder-only Transformer with rotary position embeddings, QK layer normalization, and SwiGLU feed-forward blocks
  • Original Repository: Biohub/esm

Usage

The model file depends on the multimolecule library. You can install it using pip:

Bash
1
pip install multimolecule

Direct Use

Masked Language Modeling

You can use this model directly with a pipeline for masked language modeling:

Python
1
2
3
4
5
import multimolecule  # you must import multimolecule to register models
from transformers import pipeline

predictor = pipeline("fill-mask", model="multimolecule/esmc-300m")
output = predictor("MSK<mask>EELFTGVVPILVELDGDVNGHK")

Downstream Use

Extract Features

Here is how to use this model to get the features of a given sequence in PyTorch:

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
from multimolecule import EsmCModel, ProteinTokenizer


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/esmc-300m")
model = EsmCModel.from_pretrained("multimolecule/esmc-300m")

text = "MSKGEELFTGVVPILVELDGDVNGHK"
input = tokenizer(text, return_tensors="pt")

output = model(**input)

Sequence Classification / Regression

Note

This model is not fine-tuned for any specific task. You will need to fine-tune the model on a downstream task to use it for sequence classification or regression.

Here is how to use this model as backbone to fine-tune for a sequence-level task in PyTorch:

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import torch
from multimolecule import EsmCForSequencePrediction, ProteinTokenizer


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/esmc-300m")
model = EsmCForSequencePrediction.from_pretrained("multimolecule/esmc-300m")

text = "MSKGEELFTGVVPILVELDGDVNGHK"
input = tokenizer(text, return_tensors="pt")
label = torch.tensor([1])

output = model(**input, labels=label)

Token Classification / Regression

Note

This model is not fine-tuned for any specific task. You will need to fine-tune the model on a downstream task to use it for token classification or regression.

Here is how to use this model as backbone to fine-tune for a residue-level task in PyTorch:

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import torch
from multimolecule import EsmCForTokenPrediction, ProteinTokenizer


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/esmc-300m")
model = EsmCForTokenPrediction.from_pretrained("multimolecule/esmc-300m")

text = "MSKGEELFTGVVPILVELDGDVNGHK"
input = tokenizer(text, return_tensors="pt")
label = torch.randint(2, (len(text), ))

output = model(**input, labels=label)

Contact Classification / Regression

Note

This model is not fine-tuned for any specific task. You will need to fine-tune the model on a downstream task to use it for contact classification or regression.

Here is how to use this model as backbone to fine-tune for a contact-level task in PyTorch:

Python
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
import torch
from multimolecule import EsmCForContactPrediction, ProteinTokenizer


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/esmc-300m")
model = EsmCForContactPrediction.from_pretrained("multimolecule/esmc-300m")

text = "MSKGEELFTGVVPILVELDGDVNGHK"
input = tokenizer(text, return_tensors="pt")
label = torch.randint(2, (len(text), len(text)))

output = model(**input, labels=label)

Training Details

ESMC used Masked Language Modeling (MLM) as the pre-training objective: taking a protein sequence, the model predicts amino-acid tokens that have been masked from the input sequence.

Training Data

ESMC was trained on protein sequences from UniRef, MGnify, and the Joint Genome Institute. The released materials describe ESMC as trained on approximately 2.8 billion protein sequences drawn from across life.

The upstream model card reports that sequence data was clustered at 70% sequence identity, resulting in 83 million UniRef clusters, 372 million MGnify clusters, and 2 billion JGI clusters.

Note that ProteinTokenizer will tokenize protein sequences using the MultiMolecule protein vocabulary for you.

Training Procedure

Preprocessing

ESMC tokenizes protein chains as amino-acid sequences and uses masked amino-acid prediction from surrounding sequence context. The released models use rotary position embeddings and were trained with context length up to 2048 tokens.

Pre-training

The upstream model card describes a two-stage training schedule.

  • Stage 1: 1,000,000 steps with context length 512 and 64% metagenomic data.
  • Stage 2: 500,000 steps with context length 2048 and 37.5% metagenomic data.
  • Training FLOPs: 1.26e22 for ESMC-300M, 2.17e22 for ESMC-600M, and 2.37e23 for ESMC-6B.

Citation

BibTeX
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
@UNPUBLISHED{Candido2026-wk,
  title       = "Language modeling materializes a world model of protein
                 biology",
  author      = "Candido, Salvatore and Hayes, Thomas and Derry, Alexander and
                 Rao, Roshan and Lin, Zeming and Verkuil, Robert and Wu, Bryan
                 Z and Lee, Jin Sub and Bruguera, Elise S and Keval, Jehan A
                 and Kopylov, Mykhailo and Pak, John E and Wu, Wesley and
                 Thomas, Neil and Mataraso, Samson and Hsu, Alvin and
                 Trotman-Grant, Ashton C and Fatras, Kilian and dos Santos
                 Costa, Allan and Badkundri, Rohil and Ak{\i}n, Halil and
                 Oktay, Deniz and Deaton, Jonathan and Montabana, Elizabeth and
                 Sitwala, Hrishita and Yu, Yue and Wiggert, Marius and Carlin,
                 Dylan Alexander and Goering, Anthony W and Blazejewski, Tomasz
                 and Sandora, Mccullen and Hla, Michael and Jia, Tina Z and
                 Kloker, Leon H and Sofroniew, Nicholas J and Uehara, Masatoshi
                 and Pannu, Jassi and Bachas, Sharrol and Liu, Daniel S and
                 Sercu, Tom and Rives, Alexander",
  abstract    = "Abstract Proteins are fundamental to life. The full extent of
                 their biology is beyond our ability to characterize with
                 experimental approaches in the physical laboratory. Accurate
                 digital representations could accelerate the discovery of
                 protein biology through virtual experiments. We propose
                 language modeling to learn unified and general representations
                 that can be scaled to all of protein biology. Building on
                 these representations, we develop a structure prediction model
                 that exceeds the performance of established methods for
                 biomolecular complex prediction across benchmarks, including
                 for the interactions of antibodies with their targets. A
                 simple search procedure yields high experimental success rates
                 for the discovery of proteins with nanomolar binding
                 affinities for both miniproteins and single-chain antibodies,
                 a modality critical for therapeutic design. Study of the
                 concepts in the language model's representation space reveals
                 a systematic organization aligned with the reductionist
                 understanding of proteins developed through empirical science.
                 Leveraging this organization, we generate a comprehensive map
                 of protein biology encompassing over 6.8 billion sequences and
                 1.1 billion predicted structures, identifying connections
                 across known and unknown biology. As a whole, this shows
                 language modeling as a powerful substrate for representing the
                 biology of proteins, operating across scales from the
                 prediction and design of protein interactions at the atomic
                 level, to identifying properties of proteins at different
                 levels of granularity and abstraction, to the scale of mapping
                 connections between proteins across billions of years of
                 evolution.",
  journal     = "bioRxiv",
  institution = "bioRxiv",
  month       =  jun,
  year        =  2026,
  copyright   = "http://creativecommons.org/licenses/by/4.0/"
}

Note

The artifacts distributed in this repository are part of the MultiMolecule project. If MultiMolecule supports your research, please cite the MultiMolecule project as follows:

BibTeX
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
@software{chen_2024_12638419,
  author    = {Chen, Zhiyuan and Zhu, Sophia Y.},
  title     = {MultiMolecule},
  doi       = {10.5281/zenodo.12638419},
  publisher = {Zenodo},
  url       = {https://doi.org/10.5281/zenodo.12638419},
  year      = 2024,
  month     = may,
  day       = 4
}

Contact

Please use GitHub issues of MultiMolecule for any questions or comments on the model card.

Please contact the authors of the ESMC paper for questions or comments on the paper/model.

License

This model implementation is licensed under the GNU Affero General Public License.

For additional terms and clarifications, please refer to our License FAQ.

Text Only
1
SPDX-License-Identifier: AGPL-3.0-or-later

API Reference

EsmCConfig

Bases: PreTrainedConfig

This is the configuration class to store the configuration of an EsmCModel. It is used to instantiate an ESMC model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to Biohub’s biohub/ESMC-300M architecture after conversion to the MultiMolecule protein vocabulary.

Parameters:

Name Type Description Default

vocab_size

int

Vocabulary size of the ESMC model. MultiMolecule checkpoints use the standard ProteinTokenizer vocabulary.

 37

hidden_size

int | None

Dimensionality of the encoder layers.

 None

num_hidden_layers

int | None

Number of hidden layers in the Transformer encoder.

 None

num_attention_heads

int | None

Number of attention heads for each attention layer.

 None

intermediate_size

int | None

Dimensionality of the SwiGLU hidden projection after splitting the packed upstream projection. If None, it is computed from hidden_size using ESMC’s 8 / 3 expansion rounded up to a multiple of 256.

 None

hidden_act

str

The non-linear activation function used in the feed-forward layer. ESMC uses "swiglu".

 'swiglu'

hidden_dropout

float

Dropout probability applied to residual branches.

 0.0

attention_dropout

float

Dropout probability applied to attention probabilities.

 0.0

max_position_embeddings

int

Maximum sequence length used for documentation and tokenizer compatibility. ESMC uses rotary embeddings, so this does not allocate learned position embeddings.

 2048

initializer_range

float

Standard deviation of the initializer for weight matrices.

 0.02

layer_norm_eps

float

Epsilon used by layer normalization.

 1e-05

attention_bias

bool

Whether to use bias terms in attention projection layers.

 False

feedforward_bias

bool

Whether to use bias terms in feed-forward projection layers.

 False

attention_layer_norm_bias

bool

Whether to use a bias in the pre-attention layer normalization.

 True

qk_layer_norm

bool

Whether to apply layer normalization to query and key projections before rotary embeddings.

 True

qk_layer_norm_bias

bool

Whether to use a bias in query/key layer normalization.

 False

final_layer_norm_bias

bool

Whether to use a bias in the final encoder layer normalization.

 False

residue_scaling_factor

float | None

Residual branch divisor used by ESMC. If None, defaults to sqrt(num_hidden_layers / 36).

 None

head

HeadConfig | None

The configuration of the downstream prediction heads.

 None

lm_head

MaskedLMHeadConfig | None

The configuration of the masked language model head.

 None
Source code in multimolecule/models/esmc/configuration_esmc.py
Python
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
class EsmCConfig(PreTrainedConfig):
    r"""
    This is the configuration class to store the configuration of an
    [`EsmCModel`][multimolecule.models.EsmCModel]. It is used to instantiate an ESMC model according to the specified
    arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar
    configuration to Biohub's [biohub/ESMC-300M](https://huggingface.co/biohub/ESMC-300M) architecture after conversion
    to the MultiMolecule protein vocabulary.

    Args:
        vocab_size:
            Vocabulary size of the ESMC model. MultiMolecule checkpoints use the standard
            [`ProteinTokenizer`][multimolecule.ProteinTokenizer] vocabulary.
        hidden_size:
            Dimensionality of the encoder layers.
        num_hidden_layers:
            Number of hidden layers in the Transformer encoder.
        num_attention_heads:
            Number of attention heads for each attention layer.
        intermediate_size:
            Dimensionality of the SwiGLU hidden projection after splitting the packed upstream projection. If `None`,
            it is computed from `hidden_size` using ESMC's `8 / 3` expansion rounded up to a multiple of 256.
        hidden_act:
            The non-linear activation function used in the feed-forward layer. ESMC uses `"swiglu"`.
        hidden_dropout:
            Dropout probability applied to residual branches.
        attention_dropout:
            Dropout probability applied to attention probabilities.
        max_position_embeddings:
            Maximum sequence length used for documentation and tokenizer compatibility. ESMC uses rotary embeddings,
            so this does not allocate learned position embeddings.
        initializer_range:
            Standard deviation of the initializer for weight matrices.
        layer_norm_eps:
            Epsilon used by layer normalization.
        attention_bias:
            Whether to use bias terms in attention projection layers.
        feedforward_bias:
            Whether to use bias terms in feed-forward projection layers.
        attention_layer_norm_bias:
            Whether to use a bias in the pre-attention layer normalization.
        qk_layer_norm:
            Whether to apply layer normalization to query and key projections before rotary embeddings.
        qk_layer_norm_bias:
            Whether to use a bias in query/key layer normalization.
        final_layer_norm_bias:
            Whether to use a bias in the final encoder layer normalization.
        residue_scaling_factor:
            Residual branch divisor used by ESMC. If `None`, defaults to `sqrt(num_hidden_layers / 36)`.
        head:
            The configuration of the downstream prediction heads.
        lm_head:
            The configuration of the masked language model head.
    """

    model_type = "esmc"

    def __init__(
        self,
        vocab_size: int = 37,
        hidden_size: int | None = None,
        num_hidden_layers: int | None = None,
        num_attention_heads: int | None = None,
        intermediate_size: int | None = None,
        hidden_act: str = "swiglu",
        hidden_dropout: float = 0.0,
        attention_dropout: float = 0.0,
        max_position_embeddings: int = 2048,
        initializer_range: float = 0.02,
        layer_norm_eps: float = 1.0e-5,
        attention_bias: bool = False,
        feedforward_bias: bool = False,
        attention_layer_norm_bias: bool = True,
        qk_layer_norm: bool = True,
        qk_layer_norm_bias: bool = False,
        final_layer_norm_bias: bool = False,
        residue_scaling_factor: float | None = None,
        pad_token_id: int = 0,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        unk_token_id: int = 3,
        mask_token_id: int = 4,
        null_token_id: int = 5,
        head: HeadConfig | None = None,
        lm_head: MaskedLMHeadConfig | None = None,
        **kwargs,
    ):
        hidden_size = int(kwargs.pop("d_model", 960) if hidden_size is None else hidden_size)
        num_hidden_layers = int(kwargs.pop("n_layers", 30) if num_hidden_layers is None else num_hidden_layers)
        num_attention_heads = int(kwargs.pop("n_heads", 15) if num_attention_heads is None else num_attention_heads)
        intermediate_size = (
            _swiglu_intermediate_size(hidden_size) if intermediate_size is None else int(intermediate_size)
        )
        residue_scaling_factor = (
            math.sqrt(num_hidden_layers / 36) if residue_scaling_factor is None else residue_scaling_factor
        )
        # ESMC stores the token embedding and MLM decoder as independent parameters.
        kwargs.setdefault("tie_word_embeddings", False)
        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            unk_token_id=unk_token_id,
            mask_token_id=mask_token_id,
            null_token_id=null_token_id,
            **kwargs,
        )
        validate_attention_dimensions(hidden_size, num_attention_heads)
        if (hidden_size // num_attention_heads) % 2 != 0:
            raise ValueError(
                "Rotary embeddings require an even head dimension; got "
                f"hidden_size={hidden_size}, num_attention_heads={num_attention_heads}."
            )
        hidden_act = hidden_act.lower()
        if hidden_act != "swiglu":
            raise ValueError(f"ESMC only supports hidden_act='swiglu'; got {hidden_act!r}.")
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.intermediate_size = intermediate_size
        self.hidden_act = hidden_act
        self.hidden_dropout = hidden_dropout
        self.attention_dropout = attention_dropout
        self.max_position_embeddings = max_position_embeddings
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.attention_bias = attention_bias
        self.feedforward_bias = feedforward_bias
        self.attention_layer_norm_bias = attention_layer_norm_bias
        self.qk_layer_norm = qk_layer_norm
        self.qk_layer_norm_bias = qk_layer_norm_bias
        self.final_layer_norm_bias = final_layer_norm_bias
        self.residue_scaling_factor = residue_scaling_factor
        self.head = HeadConfig(**head) if head is not None else None
        self.lm_head = (
            MaskedLMHeadConfig(**lm_head)
            if lm_head is not None
            else MaskedLMHeadConfig(
                transform="nonlinear",
                transform_act="gelu",
                bias=True,
                layer_norm_eps=layer_norm_eps,
            )
        )

EsmCForContactPrediction

Bases: EsmCPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> import torch
>>> from multimolecule import EsmCConfig, EsmCForContactPrediction, ProteinTokenizer
>>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
>>> model = EsmCForContactPrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**input, labels=torch.randint(2, (1, 9, 9)))
>>> output["logits"].shape
torch.Size([1, 9, 9, 1])
>>> output["loss"]
tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
Source code in multimolecule/models/esmc/modeling_esmc.py
Python
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
class EsmCForContactPrediction(EsmCPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import EsmCConfig, EsmCForContactPrediction, ProteinTokenizer
        >>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
        >>> model = EsmCForContactPrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**input, labels=torch.randint(2, (1, 9, 9)))
        >>> output["logits"].shape
        torch.Size([1, 9, 9, 1])
        >>> output["loss"]  # doctest:+ELLIPSIS
        tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
    """

    def __init__(self, config: EsmCConfig):
        super().__init__(config)
        self.model = EsmCModel(config, add_pooling_layer=False)
        self.contact_head = ContactPredictionHead(config)
        self.head_config = self.contact_head.config
        self.require_attentions = self.contact_head.require_attentions

        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        labels: Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | ContactPredictorOutput:
        if self.require_attentions:
            output_attentions = kwargs.get("output_attentions", self.config.output_attentions)
            if output_attentions is False:
                warn("output_attentions must be True since prediction head requires attentions.")
            kwargs["output_attentions"] = True
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            return_dict=True,
            **kwargs,
        )

        output = self.contact_head(outputs, attention_mask, input_ids, labels)
        logits, loss = output.logits, output.loss

        return ContactPredictorOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

EsmCForMaskedLM

Bases: EsmCPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> import torch
>>> from multimolecule import EsmCConfig, EsmCForMaskedLM, ProteinTokenizer
>>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
>>> model = EsmCForMaskedLM(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**input, labels=input["input_ids"])
>>> output["logits"].shape
torch.Size([1, 11, 37])
>>> output["loss"]
tensor(..., grad_fn=<NllLossBackward0>)
Source code in multimolecule/models/esmc/modeling_esmc.py
Python
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
class EsmCForMaskedLM(EsmCPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import EsmCConfig, EsmCForMaskedLM, ProteinTokenizer
        >>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
        >>> model = EsmCForMaskedLM(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**input, labels=input["input_ids"])
        >>> output["logits"].shape
        torch.Size([1, 11, 37])
        >>> output["loss"]  # doctest:+ELLIPSIS
        tensor(..., grad_fn=<NllLossBackward0>)
    """

    # ESMC does not tie input and output embeddings. Only the decoder bias alias
    # follows the Transformers convention used by ``MaskedLMHead``.
    _tied_weights_keys = {
        "lm_head.decoder.bias": "lm_head.bias",
    }

    def __init__(self, config: EsmCConfig):
        super().__init__(config)
        self.model = EsmCModel(config, add_pooling_layer=False)
        self.lm_head = MaskedLMHead(config)

        # Initialize weights and apply final processing
        self.post_init()

    def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict:
        tied_weights = super().get_expanded_tied_weights_keys(all_submodels=all_submodels)
        if all_submodels:
            return tied_weights
        return tied_weights | self._tied_weights_keys

    def get_output_embeddings(self):
        return self.lm_head.decoder

    def set_output_embeddings(self, embeddings):
        self.lm_head.decoder = embeddings
        if hasattr(self.lm_head, "bias"):
            self.lm_head.bias = embeddings.bias

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        labels: Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | MaskedLMOutput:
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            return_dict=True,
            **kwargs,
        )

        output = self.lm_head(outputs, labels)
        logits, loss = output.logits, output.loss

        return MaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

EsmCForSequencePrediction

Bases: EsmCPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> import torch
>>> from multimolecule import EsmCConfig, EsmCForSequencePrediction, ProteinTokenizer
>>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
>>> model = EsmCForSequencePrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**input, labels=torch.tensor([[1]]))
>>> output["logits"].shape
torch.Size([1, 1])
>>> output["loss"]
tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
Source code in multimolecule/models/esmc/modeling_esmc.py
Python
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class EsmCForSequencePrediction(EsmCPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import EsmCConfig, EsmCForSequencePrediction, ProteinTokenizer
        >>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
        >>> model = EsmCForSequencePrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**input, labels=torch.tensor([[1]]))
        >>> output["logits"].shape
        torch.Size([1, 1])
        >>> output["loss"]  # doctest:+ELLIPSIS
        tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
    """

    def __init__(self, config: EsmCConfig):
        super().__init__(config)
        self.model = EsmCModel(config)
        self.sequence_head = SequencePredictionHead(config)
        self.head_config = self.sequence_head.config

        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        labels: Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | SequencePredictorOutput:
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            return_dict=True,
            **kwargs,
        )

        output = self.sequence_head(outputs, labels)
        logits, loss = output.logits, output.loss

        return SequencePredictorOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

EsmCForTokenPrediction

Bases: EsmCPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> import torch
>>> from multimolecule import EsmCConfig, EsmCForTokenPrediction, ProteinTokenizer
>>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
>>> model = EsmCForTokenPrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**input, labels=torch.randint(2, (1, 9)))
>>> output["logits"].shape
torch.Size([1, 9, 1])
>>> output["loss"]
tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
Source code in multimolecule/models/esmc/modeling_esmc.py
Python
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
class EsmCForTokenPrediction(EsmCPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import EsmCConfig, EsmCForTokenPrediction, ProteinTokenizer
        >>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
        >>> model = EsmCForTokenPrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**input, labels=torch.randint(2, (1, 9)))
        >>> output["logits"].shape
        torch.Size([1, 9, 1])
        >>> output["loss"]  # doctest:+ELLIPSIS
        tensor(..., grad_fn=<BinaryCrossEntropyWithLogitsBackward0>)
    """

    def __init__(self, config: EsmCConfig):
        super().__init__(config)
        self.model = EsmCModel(config, add_pooling_layer=False)
        self.token_head = TokenPredictionHead(config)
        self.head_config = self.token_head.config

        # Initialize weights and apply final processing
        self.post_init()

    @can_return_tuple
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        labels: Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | TokenPredictorOutput:
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            return_dict=True,
            **kwargs,
        )

        output = self.token_head(outputs, attention_mask, input_ids, labels)
        logits, loss = output.logits, output.loss

        return TokenPredictorOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

EsmCModel

Bases: EsmCPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> import torch
>>> from multimolecule import EsmCConfig, EsmCModel, ProteinTokenizer
>>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
>>> model = EsmCModel(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**input)
>>> output["last_hidden_state"].shape
torch.Size([1, 11, 60])
>>> output["pooler_output"].shape
torch.Size([1, 60])
Source code in multimolecule/models/esmc/modeling_esmc.py
Python
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
class EsmCModel(EsmCPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import EsmCConfig, EsmCModel, ProteinTokenizer
        >>> config = EsmCConfig(hidden_size=60, num_hidden_layers=2, num_attention_heads=3, intermediate_size=128)
        >>> model = EsmCModel(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**input)
        >>> output["last_hidden_state"].shape
        torch.Size([1, 11, 60])
        >>> output["pooler_output"].shape
        torch.Size([1, 60])
    """

    def __init__(self, config: EsmCConfig, add_pooling_layer: bool = True):
        super().__init__(config)
        self.pad_token_id = config.pad_token_id
        self.gradient_checkpointing = False
        self.embeddings = EsmCEmbeddings(config)
        self.encoder = EsmCEncoder(config)
        self.pooler = EsmCPooler(config) if add_pooling_layer else None

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    @merge_with_config_defaults
    @capture_outputs(tie_last_hidden_states=False)
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        sequence_id: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | BaseModelOutputWithPoolingAndCrossAttentions:
        if (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if (
            attention_mask is None
            and input_ids is not None
            and not isinstance(input_ids, NestedTensor)
            and self.pad_token_id is not None
        ):
            attention_mask = input_ids.ne(self.pad_token_id)

        embedding_output = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)

        if sequence_id is not None:
            attn_mask = (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze(1)
        elif isinstance(embedding_output, NestedTensor) and attention_mask is None:
            attn_mask = None
        else:
            attn_mask = create_bidirectional_mask(
                config=self.config,
                inputs_embeds=embedding_output,
                attention_mask=attention_mask,
            )

        encoder_outputs = self.encoder(embedding_output, attention_mask=attn_mask, **kwargs)
        sequence_output = encoder_outputs.last_hidden_state
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )

EsmCPreTrainedModel

Bases: PreTrainedModel

An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models.

Source code in multimolecule/models/esmc/modeling_esmc.py
Python
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
class EsmCPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = EsmCConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _supports_flash_attn = True
    _supports_sdpa = True
    _supports_flex_attn = True
    _supports_attention_backend = True
    _can_record_outputs: dict[str, Any] | None = None
    _no_split_modules = ["EsmCLayer"]