faisalashraf/abaffinity

1# Copyright (c) Meta Platforms, Inc. and affiliates.2#3# This source code is licensed under the MIT license found in the4# LICENSE file in the root directory of this source tree.5 6from typing import Union7import torch8import torch.nn as nn9 10from .modules import ContactPredictionHead, ESM1bLayerNorm, RobertaLMHead, TransformerLayer11from .data import Alphabet12 13class ESM2(nn.Module):14    def __init__(15        self,16        num_layers: int = 33,17        embed_dim: int = 1280,18        attention_heads: int = 20,19        alphabet: Union[Alphabet, str] = "ESM-1b",20        token_dropout: bool = True,21    ):22        super().__init__()23        self.num_layers = num_layers24        self.embed_dim = embed_dim25        self.attention_heads = attention_heads26        if not isinstance(alphabet, Alphabet):27            alphabet = Alphabet.from_architecture(alphabet) 28        self.alphabet = alphabet29        self.alphabet_size = len(alphabet)30        self.padding_idx = alphabet.padding_idx31        self.mask_idx = alphabet.mask_idx32        self.cls_idx = alphabet.cls_idx33        self.eos_idx = alphabet.eos_idx34        self.prepend_bos = alphabet.prepend_bos35        self.append_eos = alphabet.append_eos36        self.token_dropout = token_dropout37 38        self._init_submodules()39 40    def _init_submodules(self):41        self.embed_scale = 142        self.embed_tokens = nn.Embedding(43            self.alphabet_size,44            self.embed_dim,45            padding_idx=self.padding_idx,46        )47 48        self.layers = nn.ModuleList(49            [50                TransformerLayer(51                    self.embed_dim,52                    4 * self.embed_dim,53                    self.attention_heads,54                    add_bias_kv=False,55                    use_esm1b_layer_norm=True,56                    use_rotary_embeddings=True,57                )58                for _ in range(self.num_layers)59            ]60        )61 62        self.contact_head = ContactPredictionHead(63            self.num_layers * self.attention_heads,64            self.prepend_bos,65            self.append_eos,66            eos_idx=self.eos_idx,67        )68        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)69 70        self.lm_head = RobertaLMHead(71            embed_dim=self.embed_dim,72            output_dim=self.alphabet_size,73            weight=self.embed_tokens.weight,74        )75 76    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):77        if return_contacts:78            need_head_weights = True79 80        assert tokens.ndim == 281        padding_mask = tokens.eq(self.padding_idx)  # B, T82 83        x = self.embed_scale * self.embed_tokens(tokens)84 85        if self.token_dropout:86            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)87            # x: B x T x C88            mask_ratio_train = 0.15 * 0.889            src_lengths = (~padding_mask).sum(-1)90            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths91            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]92 93        if padding_mask is not None:94            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))95 96        repr_layers = set(repr_layers)97        hidden_representations = {}98        if 0 in repr_layers:99            hidden_representations[0] = x100 101        if need_head_weights:102            attn_weights = []103 104        # (B, T, E) => (T, B, E)105        x = x.transpose(0, 1)106 107        if not padding_mask.any():108            padding_mask = None109 110        for layer_idx, layer in enumerate(self.layers):111            x, attn = layer(112                x,113                self_attn_padding_mask=padding_mask,114                need_head_weights=need_head_weights,115            )116            if (layer_idx + 1) in repr_layers:117                hidden_representations[layer_idx + 1] = x.transpose(0, 1)118            if need_head_weights:119                # (H, B, T, T) => (B, H, T, T)120                attn_weights.append(attn.transpose(1, 0))121 122        x = self.emb_layer_norm_after(x)123        x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)124 125        # last hidden representation should have layer norm applied126        if (layer_idx + 1) in repr_layers:127            hidden_representations[layer_idx + 1] = x128        x = self.lm_head(x)129 130        result = {"logits": x, "representations": hidden_representations}131        if need_head_weights:132            # attentions: B x L x H x T x T133            attentions = torch.stack(attn_weights, 1)134            if padding_mask is not None:135                attention_mask = 1 - padding_mask.type_as(attentions)136                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)137                attentions = attentions * attention_mask[:, None, None, :, :]138            result["attentions"] = attentions139            if return_contacts:140                contacts = self.contact_head(tokens, attentions)141                result["contacts"] = contacts142 143        return result144 145    def predict_contacts(self, tokens):146        return self(tokens, return_contacts=True)["contacts"]147