AbLang2

Pre-trained model on paired and unpaired antibody sequences using a modified masked language modeling objective.

Disclaimer

This is an UNOFFICIAL implementation of Addressing the antibody germline bias and its effect on language models for improved antibody design by Tobias H. Olsen, et al.

The OFFICIAL repository of AbLang2 is at oxpig/AbLang2.

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

Model Details

AbLang2 is an antibody-specific encoder-only protein language model trained to reduce antibody germline bias in masked residue prediction. It uses multi-head self-attention with rotary position embeddings and SwiGLU feed-forward blocks. The released paired model is trained on paired and unpaired antibody sequence data and is optimized for non-germline residue prediction.

Model Specification

Num Layers	Hidden Size	Num Heads	Intermediate Size	Num Parameters (M)	FLOPs (G)	MACs (G)	Max Num Tokens
12	480	20	1920	44.82	24.48	12.20	256

Note

Max Num Tokens reflects the training sequence length of the released checkpoint. AbLang2 uses rotary position embeddings and has no max_position_embeddings field, so the architecture itself does not impose a hard length limit.

Links

• Code: multimolecule.ablang2
• Data: Observed Antibody Space
• Paper: Addressing the antibody germline bias and its effect on language models for improved antibody design
• Developed by: Tobias H. Olsen, Iain H. Moal, Charlotte M. Deane
• Model type: Encoder-only antibody language model with rotary position embeddings and SwiGLU feed-forward blocks
• Original Repository: oxpig/AbLang2

Usage

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

pip install multimolecule

Direct Use

Masked Language Modeling

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

import multimolecule  # you must import multimolecule to register models
from transformers import pipeline

predictor = pipeline("fill-mask", model="multimolecule/ablang2")
output = predictor("EVQLVESGGGLVQPGGSLRLSCAAS<mask>FTFSSYAMSWVRQAPGKGLEWV")

Downstream Use

Extract Features

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

from multimolecule import ProteinTokenizer, AbLang2Model


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/ablang2")
model = AbLang2Model.from_pretrained("multimolecule/ablang2")

text = "EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMSWVRQAPGKGLEWV"
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:

import torch
from multimolecule import ProteinTokenizer, AbLang2ForSequencePrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/ablang2")
model = AbLang2ForSequencePrediction.from_pretrained("multimolecule/ablang2")

text = "EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMSWVRQAPGKGLEWV"
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:

import torch
from multimolecule import ProteinTokenizer, AbLang2ForTokenPrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/ablang2")
model = AbLang2ForTokenPrediction.from_pretrained("multimolecule/ablang2")

text = "EVQLVESGGGLVQPGGSLRLSCAASGFTFSSYAMSWVRQAPGKGLEWV"
input = tokenizer(text, return_tensors="pt")
label = torch.randint(2, (1, 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:

import torch
from multimolecule import ProteinTokenizer, AbLang2ForContactPrediction


tokenizer = ProteinTokenizer.from_pretrained("multimolecule/ablang2")
model = AbLang2ForContactPrediction.from_pretrained("multimolecule/ablang2")

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

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

Training Details

AbLang2 was trained with masked language modeling as the pre-training objective. The model is bidirectional, so each masked position attends to surrounding residues on both sides.

Training Data

AbLang2 is trained on sequences derived from the Observed Antibody Space (OAS), including 35.6 million unpaired heavy/light-chain sequences and 1.26 million paired antibody sequences for the final released model.

Training Procedure

The AbLang2 paper focuses on reducing antibody germline bias in residue prediction and model-guided antibody design. Please refer to the original paper for details on the training setup.

Citation

@article{olsen2024ablang2,
  title   = {Addressing the antibody germline bias and its effect on language models for improved antibody design},
  author  = {Olsen, Tobias H. and Moal, Iain H. and Deane, Charlotte M.},
  year    = {2024},
  journal = {Bioinformatics},
  volume  = {40},
  number  = {11},
  pages   = {btae618},
  doi     = {10.1093/bioinformatics/btae618},
  url     = {https://doi.org/10.1093/bioinformatics/btae618},
}

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:

@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 AbLang2 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.

SPDX-License-Identifier: AGPL-3.0-or-later

API Reference

AbLang2Config

Bases: PreTrainedConfig

This is the configuration class to store the configuration of an AbLang2Model. It is used to instantiate an AbLang2 model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to the official AbLang2 paired-antibody checkpoint.

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 AbLang2 model. Defines the number of different tokens that can be represented by the input_ids passed when calling [AbLang2Model].	37
hidden_size	int	Dimensionality of the encoder layers and the pooler layer.	480
num_hidden_layers	int	Number of hidden layers in the Transformer encoder.	12
num_attention_heads	int	Number of attention heads for each attention layer in the Transformer encoder.	20
intermediate_size	int	Dimensionality of the feed-forward hidden layer after the activation. For SwiGLU, the first feed-forward projection has twice this size.	1920
hidden_act	str	The non-linear activation function used in the feed-forward layers. AbLang2 uses "swiglu".	'swiglu'
hidden_dropout	float	The dropout probability applied to residual and feed-forward outputs.	0.0
attention_dropout	float	The dropout ratio applied to attention probabilities.	0.0
initializer_range	float	Standard deviation used by default initialization for embeddings and linear layers.	0.02
layer_norm_eps	float	The epsilon used by the layer normalization layers.	1e-12
rotary_base	float	Base used for rotary position embeddings.	10000.0
attention_bias	bool	Whether to use bias terms in the attention projections.	True
feedforward_bias	bool	Whether to use bias terms in the feed-forward projections.	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:

>>> from multimolecule import AbLang2Config, AbLang2Model
>>> # Initializing an AbLang2 multimolecule/ablang2 style configuration
>>> configuration = AbLang2Config()
>>> # Initializing a model (with random weights) from the multimolecule/ablang2 style configuration
>>> model = AbLang2Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config

Source code in multimolecule/models/ablang2/configuration_ablang2.py

class AbLang2Config(PreTrainedConfig):
    r"""
    This is the configuration class to store the configuration of an
    [`AbLang2Model`][multimolecule.models.AbLang2Model]. It is used to instantiate an AbLang2 model according to the
    specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a
    similar configuration to the official AbLang2 paired-antibody checkpoint.

    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 AbLang2 model. Defines the number of different tokens that can be represented by the
            `input_ids` passed when calling [`AbLang2Model`].
        hidden_size:
            Dimensionality of the encoder layers and the pooler layer.
        num_hidden_layers:
            Number of hidden layers in the Transformer encoder.
        num_attention_heads:
            Number of attention heads for each attention layer in the Transformer encoder.
        intermediate_size:
            Dimensionality of the feed-forward hidden layer after the activation. For SwiGLU, the first feed-forward
            projection has twice this size.
        hidden_act:
            The non-linear activation function used in the feed-forward layers. AbLang2 uses `"swiglu"`.
        hidden_dropout:
            The dropout probability applied to residual and feed-forward outputs.
        attention_dropout:
            The dropout ratio applied to attention probabilities.
        initializer_range:
            Standard deviation used by default initialization for embeddings and linear layers.
        layer_norm_eps:
            The epsilon used by the layer normalization layers.
        rotary_base:
            Base used for rotary position embeddings.
        attention_bias:
            Whether to use bias terms in the attention projections.
        feedforward_bias:
            Whether to use bias terms in the feed-forward projections.
        head:
            The configuration of the downstream prediction head.
        lm_head:
            The configuration of the masked language model head.

    Examples:
        >>> from multimolecule import AbLang2Config, AbLang2Model
        >>> # Initializing an AbLang2 multimolecule/ablang2 style configuration
        >>> configuration = AbLang2Config()
        >>> # Initializing a model (with random weights) from the multimolecule/ablang2 style configuration
        >>> model = AbLang2Model(configuration)
        >>> # Accessing the model configuration
        >>> configuration = model.config
    """

    model_type = "ablang2"

    def __init__(
        self,
        vocab_size: int = 37,
        hidden_size: int = 480,
        num_hidden_layers: int = 12,
        num_attention_heads: int = 20,
        intermediate_size: int = 1920,
        hidden_act: str = "swiglu",
        hidden_dropout: float = 0.0,
        attention_dropout: float = 0.0,
        initializer_range: float = 0.02,
        layer_norm_eps: float = 1.0e-12,
        rotary_base: float = 10000.0,
        attention_bias: bool = True,
        feedforward_bias: 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,
        sep_token_id: int = 32,
        head: HeadConfig | None = None,
        lm_head: MaskedLMHeadConfig | None = None,
        **kwargs,
    ):
        kwargs.setdefault("tie_word_embeddings", True)
        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,
            sep_token_id=sep_token_id,
            **kwargs,
        )
        validate_attention_dimensions(hidden_size, num_attention_heads)
        hidden_act = hidden_act.lower()
        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.initializer_range = initializer_range
        self.layer_norm_eps = layer_norm_eps
        self.rotary_base = rotary_base
        self.attention_bias = attention_bias
        self.feedforward_bias = feedforward_bias
        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(
                hidden_size=hidden_size,
                transform=None,
                transform_act=hidden_act,
                bias=True,
                layer_norm_eps=layer_norm_eps,
            )
        )

AbLang2ForContactPrediction

Bases: AbLang2PreTrainedModel

Examples:

>>> import torch
>>> from multimolecule import AbLang2Config, AbLang2ForContactPrediction, ProteinTokenizer
>>> config = AbLang2Config()
>>> model = AbLang2ForContactPrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
>>> output = model(**input, labels=torch.randint(2, (1, 25, 25)))
>>> output["logits"].shape
torch.Size([1, 25, 25, 1])

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2ForContactPrediction(AbLang2PreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import AbLang2Config, AbLang2ForContactPrediction, ProteinTokenizer
        >>> config = AbLang2Config()
        >>> model = AbLang2ForContactPrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
        >>> output = model(**input, labels=torch.randint(2, (1, 25, 25)))
        >>> output["logits"].shape
        torch.Size([1, 25, 25, 1])
    """

    def __init__(self, config: AbLang2Config):
        super().__init__(config)
        self.model = AbLang2Model(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,
        )

AbLang2ForMaskedLM

Bases: AbLang2PreTrainedModel

Examples:

>>> import torch
>>> from multimolecule import AbLang2Config, AbLang2ForMaskedLM, ProteinTokenizer
>>> config = AbLang2Config()
>>> model = AbLang2ForMaskedLM(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
>>> output = model(**input, labels=input["input_ids"])
>>> output["logits"].shape
torch.Size([1, 27, 37])

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2ForMaskedLM(AbLang2PreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import AbLang2Config, AbLang2ForMaskedLM, ProteinTokenizer
        >>> config = AbLang2Config()
        >>> model = AbLang2ForMaskedLM(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
        >>> output = model(**input, labels=input["input_ids"])
        >>> output["logits"].shape
        torch.Size([1, 27, 37])
    """

    _tied_weights_keys = {
        "lm_head.decoder.bias": "lm_head.bias",
        "lm_head.decoder.weight": "model.embeddings.word_embeddings.weight",
    }

    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 __init__(self, config: AbLang2Config):
        super().__init__(config)
        self.model = AbLang2Model(config, add_pooling_layer=False)
        self.lm_head = AbLang2LMHead(config, weight=self.model.embeddings.word_embeddings.weight)

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

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

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

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

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

AbLang2ForSequencePrediction

Bases: AbLang2PreTrainedModel

Examples:

>>> import torch
>>> from multimolecule import AbLang2Config, AbLang2ForSequencePrediction, ProteinTokenizer
>>> config = AbLang2Config()
>>> model = AbLang2ForSequencePrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
>>> output = model(**input, labels=torch.tensor([[1]]))
>>> output["logits"].shape
torch.Size([1, 1])

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2ForSequencePrediction(AbLang2PreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import AbLang2Config, AbLang2ForSequencePrediction, ProteinTokenizer
        >>> config = AbLang2Config()
        >>> model = AbLang2ForSequencePrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
        >>> output = model(**input, labels=torch.tensor([[1]]))
        >>> output["logits"].shape
        torch.Size([1, 1])
    """

    def __init__(self, config: AbLang2Config):
        super().__init__(config)
        self.model = AbLang2Model(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,
        )

AbLang2ForTokenPrediction

Bases: AbLang2PreTrainedModel

Examples:

>>> import torch
>>> from multimolecule import AbLang2Config, AbLang2ForTokenPrediction, ProteinTokenizer
>>> config = AbLang2Config()
>>> model = AbLang2ForTokenPrediction(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
>>> output = model(**input, labels=torch.randint(2, (1, 25)))
>>> output["logits"].shape
torch.Size([1, 25, 1])

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2ForTokenPrediction(AbLang2PreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import AbLang2Config, AbLang2ForTokenPrediction, ProteinTokenizer
        >>> config = AbLang2Config()
        >>> model = AbLang2ForTokenPrediction(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
        >>> output = model(**input, labels=torch.randint(2, (1, 25)))
        >>> output["logits"].shape
        torch.Size([1, 25, 1])
    """

    def __init__(self, config: AbLang2Config):
        super().__init__(config)
        self.model = AbLang2Model(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,
        )

AbLang2Model

Bases: AbLang2PreTrainedModel

Examples:

>>> import torch
>>> from multimolecule import AbLang2Config, AbLang2Model, ProteinTokenizer
>>> config = AbLang2Config()
>>> model = AbLang2Model(config)
>>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
>>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
>>> output = model(**input)
>>> output["last_hidden_state"].shape
torch.Size([1, 27, 480])
>>> output["pooler_output"].shape
torch.Size([1, 480])

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2Model(AbLang2PreTrainedModel):
    """
    Examples:
        >>> import torch
        >>> from multimolecule import AbLang2Config, AbLang2Model, ProteinTokenizer
        >>> config = AbLang2Config()
        >>> model = AbLang2Model(config)
        >>> tokenizer = ProteinTokenizer.from_pretrained("multimolecule/protein")
        >>> input = tokenizer("EVQLVESGGGLVQPGGSLRLSCAAS", return_tensors="pt")
        >>> output = model(**input)
        >>> output["last_hidden_state"].shape
        torch.Size([1, 27, 480])
        >>> output["pooler_output"].shape
        torch.Size([1, 480])
    """

    def __init__(self, config: AbLang2Config, add_pooling_layer: bool = True):
        super().__init__(config)
        self.pad_token_id = config.pad_token_id
        self.gradient_checkpointing = False
        self.embeddings = AbLang2Embeddings(config)
        self.encoder = AbLang2Encoder(config)
        self.pooler = AbLang2Pooler(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,
        inputs_embeds: Tensor | NestedTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[Tensor, ...] | BaseModelOutputWithPoolingAndCrossAttentions:
        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)

        embedding_output = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)
        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,
        )

AbLang2PreTrainedModel

Bases: PreTrainedModel

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

Source code in multimolecule/models/ablang2/modeling_ablang2.py

class AbLang2PreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = AbLang2Config
    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 = ["AbLang2Layer"]

    @torch.no_grad()
    def _init_weights(self, module: nn.Module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            init.normal_(module.weight, mean=0.0, std=std)
            if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            init.ones_(module.weight)
            init.zeros_(module.bias)