跳转至

CARP

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

Disclaimer

This is an UNOFFICIAL implementation of Convolutions are competitive with transformers for protein sequence pretraining by Kevin K. Yang, et al.

The OFFICIAL repository of CARP is at microsoft/protein-sequence-models.

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 CARP did not write this model card for this model so this model card has been written by the MultiMolecule team.

Model Details

CARP is a family of ByteNet-style convolutional protein language models. It uses learned token embeddings, a stack of residual dilated 1D convolution blocks, and a final layer normalization before the masked-language-model decoder. The models were pre-trained on the March 2020 release of UniRef50 using the same masked language modeling task as BERT and ESM-1b.

Variants

Model Specification

Variant Num Layers Hidden Size Intermediate Size Num Parameters (M) FLOPs (G) MACs (G) Max Num Tokens
CARP-600k 16 128 64 0.61 1.25 0.61 1024
CARP-38M 16 1024 512 37.90 77.68 38.70
CARP-76M 32 75.74 155.26 77.36
CARP-640M 56 1280 1280 642.96 1317.22 657.73

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/carp-600k")
output = predictor("MVLSPADKTNVKAAW<mask>KVGAHAGEYGAEALER")

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 ProteinTokenizer, CarpModel


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/carp-600k")
model = CarpModel.from_pretrained("multimolecule/carp-600k")

text = "MVLSPADKTNVKAAWGKVGAHAGEYGAEALER"
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 ProteinTokenizer, CarpForSequencePrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/carp-600k")
model = CarpForSequencePrediction.from_pretrained("multimolecule/carp-600k")

text = "MVLSPADKTNVKAAWGKVGAHAGEYGAEALER"
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 ProteinTokenizer, CarpForTokenPrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/carp-600k")
model = CarpForTokenPrediction.from_pretrained("multimolecule/carp-600k")

text = "MVLSPADKTNVKAAWGKVGAHAGEYGAEALER"
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 ProteinTokenizer, CarpForContactPrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/carp-600k")
model = CarpForContactPrediction.from_pretrained("multimolecule/carp-600k")

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

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

Training Details

CARP was trained with Masked Language Modeling (MLM) as the pre-training objective. Masked residues are predicted from the surrounding protein sequence using bidirectional dilated convolution blocks rather than self-attention layers.

Training Data

CARP was pre-trained on the March 2020 release of UniRef50.

Training Procedure

Preprocessing

The released CARP checkpoints use the protein alphabet from the official sequence_models package. During conversion, equivalent amino-acid and special-token rows are mapped into the MultiMolecule protein tokenizer vocabulary.

Pre-training

The model was trained with masked language modeling over a ByteNet-style residual dilated convolution stack. Please refer to the original paper for details on the training setup.

Citation

BibTeX
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
@article{yang2024convolutions,
  author  = {Yang, Kevin K. and Fusi, Nicolo and Lu, Alex X.},
  title   = {Convolutions are competitive with transformers for protein sequence pretraining},
  journal = {Cell Systems},
  volume  = {15},
  number  = {3},
  pages   = {286--294.e2},
  year    = {2024},
  doi     = {10.1016/j.cels.2024.01.008},
  url     = {https://doi.org/10.1016/j.cels.2024.01.008},
}

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 CARP 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

CarpConfig

Bases: PreTrainedConfig

This is the configuration class to store the configuration of a CarpModel. It is used to instantiate a CARP model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to the official carp_600k checkpoint from Microsoft Research.

Configuration objects inherit from PreTrainedConfig and can be used to control the model outputs. Read the documentation from PreTrainedConfig for more information.

Parameters:

Name Type Description Default

vocab_size

int

Vocabulary size of the CARP model. Defines the number of different tokens that can be represented by the input_ids passed when calling [CarpModel].

 37

embedding_size

int

Dimensionality of the token embeddings before the ByteNet projection.

 8

hidden_size

int

Dimensionality of the ByteNet residual stream.

 128

intermediate_size

int | None

Dimensionality used inside each dilated convolution block. If None, CARP uses hidden_size // 2 when slim=True, otherwise hidden_size.

 None

num_hidden_layers

int

Number of ByteNet residual convolution blocks.

 16

kernel_size

int

Width of the dilated one-dimensional convolution kernels.

 5

max_dilation

int

Largest dilation factor in the cyclic ByteNet dilation schedule.

 128

hidden_act

str

Non-linear activation function used in each residual convolution block. CARP checkpoints use "gelu".

 'gelu'

hidden_dropout

float

Dropout probability applied after each residual convolution block.

 0.0

initializer_range

float

Standard deviation used for newly initialized vocabulary rows during conversion.

 0.02

layer_norm_eps

float

Epsilon used by layer normalization layers.

 1e-05

slim

bool

Whether the checkpoint uses the half-width hidden channel inside residual convolution blocks.

 True

head

HeadConfig | None

The configuration of the downstream prediction head.

 None

lm_head

MaskedLMHeadConfig | None

The configuration of the masked language model head.

 None

Examples:

Python Console Session
1
2
3
4
5
6
7
>>> from multimolecule import CarpConfig, CarpModel
>>> # Initializing a CARP multimolecule/carp style configuration
>>> configuration = CarpConfig()
>>> # Initializing a model (with random weights) from the multimolecule/carp style configuration
>>> model = CarpModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
Source code in multimolecule/models/carp/configuration_carp.py
Python
 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
 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
class CarpConfig(PreTrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`CarpModel`][multimolecule.models.CarpModel].
    It is used to instantiate a CARP model according to the specified arguments, defining the model architecture.
    Instantiating a configuration with the defaults will yield a similar configuration to the official `carp_600k`
    checkpoint from Microsoft Research.

    Configuration objects inherit from [`PreTrainedConfig`][multimolecule.models.PreTrainedConfig] and can be used to
    control the model outputs. Read the documentation from [`PreTrainedConfig`][multimolecule.models.PreTrainedConfig]
    for more information.

    Args:
        vocab_size:
            Vocabulary size of the CARP model. Defines the number of different tokens that can be represented by the
            `input_ids` passed when calling [`CarpModel`].
        embedding_size:
            Dimensionality of the token embeddings before the ByteNet projection.
        hidden_size:
            Dimensionality of the ByteNet residual stream.
        intermediate_size:
            Dimensionality used inside each dilated convolution block. If `None`, CARP uses `hidden_size // 2` when
            `slim=True`, otherwise `hidden_size`.
        num_hidden_layers:
            Number of ByteNet residual convolution blocks.
        kernel_size:
            Width of the dilated one-dimensional convolution kernels.
        max_dilation:
            Largest dilation factor in the cyclic ByteNet dilation schedule.
        hidden_act:
            Non-linear activation function used in each residual convolution block. CARP checkpoints use `"gelu"`.
        hidden_dropout:
            Dropout probability applied after each residual convolution block.
        initializer_range:
            Standard deviation used for newly initialized vocabulary rows during conversion.
        layer_norm_eps:
            Epsilon used by layer normalization layers.
        slim:
            Whether the checkpoint uses the half-width hidden channel inside residual convolution blocks.
        head:
            The configuration of the downstream prediction head.
        lm_head:
            The configuration of the masked language model head.

    Examples:
        >>> from multimolecule import CarpConfig, CarpModel
        >>> # Initializing a CARP multimolecule/carp style configuration
        >>> configuration = CarpConfig()
        >>> # Initializing a model (with random weights) from the multimolecule/carp style configuration
        >>> model = CarpModel(configuration)
        >>> # Accessing the model configuration
        >>> configuration = model.config
    """

    model_type = "carp"

    def __init__(
        self,
        vocab_size: int = 37,
        embedding_size: int = 8,
        hidden_size: int = 128,
        intermediate_size: int | None = None,
        num_hidden_layers: int = 16,
        kernel_size: int = 5,
        max_dilation: int = 128,
        hidden_act: str = "gelu",
        hidden_dropout: float = 0.0,
        initializer_range: float = 0.02,
        layer_norm_eps: float = 1.0e-5,
        slim: bool = True,
        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,
    ):
        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,
        )

        hidden_act = hidden_act.lower()
        if kernel_size <= 0 or kernel_size % 2 == 0:
            raise ValueError(f"kernel_size must be a positive odd integer; got {kernel_size}.")
        if max_dilation < 1:
            raise ValueError(f"max_dilation must be positive; got {max_dilation}.")
        if intermediate_size is None:
            intermediate_size = hidden_size // 2 if slim else hidden_size
        for name, value in {
            "vocab_size": vocab_size,
            "embedding_size": embedding_size,
            "hidden_size": hidden_size,
            "intermediate_size": intermediate_size,
            "num_hidden_layers": num_hidden_layers,
        }.items():
            if value <= 0:
                raise ValueError(f"{name} must be positive; got {value}.")

        self.vocab_size = vocab_size
        self.embedding_size = embedding_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.kernel_size = kernel_size
        self.max_dilation = max_dilation
        self.hidden_act = hidden_act
        self.hidden_dropout = hidden_dropout
        self.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.slim = slim
        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=None, transform_act=None, bias=True)
        )

CarpEncoderOutput dataclass

Bases: ModelOutput

Base class for outputs of the CARP convolutional encoder.

Parameters:

Name Type Description Default

last_hidden_state

`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`

Sequence of hidden states at the output of the last encoder layer.

 None

attentions

tuple[FloatTensor, ...] | None

Always None; CARP is a convolutional model and has no attention layers. Provided for compatibility with the Transformers output convention.

 None
Source code in multimolecule/models/carp/modeling_carp.py
Python
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
@dataclass
class CarpEncoderOutput(ModelOutput):
    """
    Base class for outputs of the CARP convolutional encoder.

    Args:
        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden states at the output of the last encoder layer.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or
            when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the embedding output plus one after each encoder layer) of shape
            `(batch_size, sequence_length, hidden_size)`.
        attentions:
            Always `None`; CARP is a convolutional model and has no attention layers. Provided for compatibility with
            the Transformers output convention.
    """

    last_hidden_state: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None

CarpForContactPrediction

Bases: CarpPreTrainedModel

Examples:

Python Console Session
1
2
3
4
5
6
7
8
>>> import torch
>>> from multimolecule.models.carp import CarpConfig, CarpForContactPrediction
>>> config = CarpConfig()
>>> model = CarpForContactPrediction(config)
>>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
>>> output = model(inputs, labels=torch.randint(2, (1, 9, 9)))
>>> output["logits"].shape
torch.Size([1, 9, 9, 1])
Source code in multimolecule/models/carp/modeling_carp.py
Python
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
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
class CarpForContactPrediction(CarpPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule.models.carp import CarpConfig, CarpForContactPrediction
        >>> config = CarpConfig()
        >>> model = CarpForContactPrediction(config)
        >>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
        >>> output = model(inputs, labels=torch.randint(2, (1, 9, 9)))
        >>> output["logits"].shape
        torch.Size([1, 9, 9, 1])
    """

    def __init__(self, config: CarpConfig):
        super().__init__(config)
        self.model = CarpModel(config, add_pooling_layer=False)
        self.num_labels = config.num_labels
        self.contact_head = ContactPredictionHead(config)
        self.head_config = self.contact_head.config
        if self.contact_head.require_attentions:
            raise ValueError("CARP does not expose attention maps; use a representation-based contact head.")

        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:
        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,
        )

CarpForMaskedLM

Bases: CarpPreTrainedModel

Examples:

Python Console Session
1
2
3
4
5
6
7
8
>>> import torch
>>> from multimolecule.models.carp import CarpConfig, CarpForMaskedLM
>>> config = CarpConfig()
>>> model = CarpForMaskedLM(config)
>>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
>>> output = model(inputs, labels=inputs)
>>> output["logits"].shape
torch.Size([1, 11, 37])
Source code in multimolecule/models/carp/modeling_carp.py
Python
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
317
318
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
class CarpForMaskedLM(CarpPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule.models.carp import CarpConfig, CarpForMaskedLM
        >>> config = CarpConfig()
        >>> model = CarpForMaskedLM(config)
        >>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
        >>> output = model(inputs, labels=inputs)
        >>> output["logits"].shape
        torch.Size([1, 11, 37])
    """

    def __init__(self, config: CarpConfig):
        super().__init__(config)
        self.model = CarpModel(config, add_pooling_layer=False)
        self.lm_head = CarpLMHead(config)

        self.post_init()

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

    def set_output_embeddings(self, embeddings):
        self.lm_head.decoder = embeddings

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

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

    @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.last_hidden_state, labels)
        logits, loss = output.logits, output.loss

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

CarpForSequencePrediction

Bases: CarpPreTrainedModel

Examples:

Python Console Session
1
2
3
4
5
6
7
8
>>> import torch
>>> from multimolecule.models.carp import CarpConfig, CarpForSequencePrediction
>>> config = CarpConfig()
>>> model = CarpForSequencePrediction(config)
>>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
>>> output = model(inputs, labels=torch.tensor([[1]]))
>>> output["logits"].shape
torch.Size([1, 1])
Source code in multimolecule/models/carp/modeling_carp.py
Python
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
174
175
176
177
178
179
180
181
182
183
184
185
class CarpForSequencePrediction(CarpPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule.models.carp import CarpConfig, CarpForSequencePrediction
        >>> config = CarpConfig()
        >>> model = CarpForSequencePrediction(config)
        >>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
        >>> output = model(inputs, labels=torch.tensor([[1]]))
        >>> output["logits"].shape
        torch.Size([1, 1])
    """

    def __init__(self, config: CarpConfig):
        super().__init__(config)
        self.model = CarpModel(config)
        self.num_labels = config.num_labels
        self.sequence_head = SequencePredictionHead(config)
        self.head_config = self.sequence_head.config

        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,
        )

CarpForTokenPrediction

Bases: CarpPreTrainedModel

Examples:

Python Console Session
1
2
3
4
5
6
7
8
>>> import torch
>>> from multimolecule.models.carp import CarpConfig, CarpForTokenPrediction
>>> config = CarpConfig()
>>> model = CarpForTokenPrediction(config)
>>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
>>> output = model(inputs, labels=torch.randint(2, (1, 9)))
>>> output["logits"].shape
torch.Size([1, 9, 1])
Source code in multimolecule/models/carp/modeling_carp.py
Python
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
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
class CarpForTokenPrediction(CarpPreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule.models.carp import CarpConfig, CarpForTokenPrediction
        >>> config = CarpConfig()
        >>> model = CarpForTokenPrediction(config)
        >>> inputs = torch.tensor([[1, 6, 23, 15, 21, 18, 6, 8, 14, 22, 2]])
        >>> output = model(inputs, labels=torch.randint(2, (1, 9)))
        >>> output["logits"].shape
        torch.Size([1, 9, 1])
    """

    def __init__(self, config: CarpConfig):
        super().__init__(config)
        self.model = CarpModel(config, add_pooling_layer=False)
        self.num_labels = config.num_labels
        self.token_head = TokenPredictionHead(config)
        self.head_config = self.token_head.config

        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,
        )

CarpModel

Bases: CarpPreTrainedModel

Examples:

Python Console Session
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
>>> from multimolecule import ProteinTokenizer
>>> from multimolecule.models.carp import CarpConfig, CarpModel
>>> config = CarpConfig()
>>> model = CarpModel(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> inputs = tokenizer("MVLSPADKT", return_tensors="pt")
>>> output = model(**inputs)
>>> output["last_hidden_state"].shape
torch.Size([1, 11, 128])
>>> output["pooler_output"].shape
torch.Size([1, 128])
Source code in multimolecule/models/carp/modeling_carp.py
Python
 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
class CarpModel(CarpPreTrainedModel):
    """
    Examples:
        >>> from multimolecule import ProteinTokenizer
        >>> from multimolecule.models.carp import CarpConfig, CarpModel
        >>> config = CarpConfig()
        >>> model = CarpModel(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> inputs = tokenizer("MVLSPADKT", return_tensors="pt")
        >>> output = model(**inputs)
        >>> output["last_hidden_state"].shape
        torch.Size([1, 11, 128])
        >>> output["pooler_output"].shape
        torch.Size([1, 128])
    """

    def __init__(self, config: CarpConfig, add_pooling_layer: bool = True):
        super().__init__(config)
        self.pad_token_id = config.pad_token_id
        self.embeddings = CarpEmbeddings(config)
        self.encoder = CarpEncoder(config)
        self.pooler = CarpPooler() if add_pooling_layer else None

        self.post_init()

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

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

    @can_return_tuple
    @merge_with_config_defaults
    def forward(
        self,
        input_ids: Tensor | NestedTensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | NestedTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | CarpModelOutput:
        if isinstance(input_ids, NestedTensor):
            if attention_mask is None:
                attention_mask = input_ids.mask
            input_ids = input_ids.tensor
        if isinstance(inputs_embeds, NestedTensor):
            if attention_mask is None:
                attention_mask = inputs_embeds.mask
            inputs_embeds = inputs_embeds.tensor
        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 self.pad_token_id is not None:
            attention_mask = input_ids.ne(self.pad_token_id)

        hidden_states = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)
        if attention_mask is None:
            attention_mask = torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)

        encoder_outputs = self.encoder(
            hidden_states,
            attention_mask=attention_mask,
            output_hidden_states=kwargs.get("output_hidden_states", self.config.output_hidden_states),
        )
        sequence_output = encoder_outputs.last_hidden_state
        pooled_output = self.pooler(sequence_output, attention_mask) if self.pooler is not None else None

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

CarpModelOutput dataclass

Bases: ModelOutput

Base class for outputs of the CARP backbone.

Parameters:

Name Type Description Default

last_hidden_state

`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`

Sequence of hidden states at the output of the last encoder layer.

 None

pooler_output

`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*

Mean-pooled sequence representation over unmasked tokens.

 None

attentions

tuple[FloatTensor, ...] | None

Always None; CARP is a convolutional model and has no attention layers. Provided for compatibility with the Transformers output convention.

 None
Source code in multimolecule/models/carp/modeling_carp.py
Python
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
@dataclass
class CarpModelOutput(ModelOutput):
    """
    Base class for outputs of the CARP backbone.

    Args:
        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden states at the output of the last encoder layer.
        pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*):
            Mean-pooled sequence representation over unmasked tokens.
        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or
            when `config.output_hidden_states=True`):
            Tuple of `torch.FloatTensor` (one for the embedding output plus one after each encoder layer) of shape
            `(batch_size, sequence_length, hidden_size)`.
        attentions:
            Always `None`; CARP is a convolutional model and has no attention layers. Provided for compatibility with
            the Transformers output convention.
    """

    last_hidden_state: torch.FloatTensor | None = None
    pooler_output: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None

CarpPreTrainedModel

Bases: PreTrainedModel

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

Source code in multimolecule/models/carp/modeling_carp.py
Python
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
class CarpPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = CarpConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["CarpLayer"]

    @torch.no_grad()
    def _init_weights(self, module: nn.Module):
        # CARP's reference implementation uses PyTorch module defaults. Keep
        # those constructor initializers instead of applying a Transformers
        # normal initialization pass.
        return