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 6import math7from typing import Optional8 9import torch10import torch.nn as nn11import torch.nn.functional as F12 13from .multihead_attention import MultiheadAttention  # noqa14from .axial_attention import ColumnSelfAttention, RowSelfAttention15 16 17def gelu(x):18    """Implementation of the gelu activation function.19 20    For information: OpenAI GPT's gelu is slightly different21    (and gives slightly different results):22    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))23    """24    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))25 26 27def symmetrize(x):28    "Make layer symmetric in final two dimensions, used for contact prediction."29    return x + x.transpose(-1, -2)30 31 32def apc(x):33    "Perform average product correct, used for contact prediction."34    a1 = x.sum(-1, keepdims=True)35    a2 = x.sum(-2, keepdims=True)36    a12 = x.sum((-1, -2), keepdims=True)37 38    avg = a1 * a239    avg.div_(a12)  # in-place to reduce memory40    normalized = x - avg41    return normalized42 43 44class ESM1LayerNorm(nn.Module):45    def __init__(self, hidden_size, eps=1e-12, affine=True):46        """Construct a layernorm layer in the TF style (eps inside the sqrt)."""47        super().__init__()48        self.hidden_size = (hidden_size,) if isinstance(hidden_size, int) else tuple(hidden_size)49        self.eps = eps50        self.affine = bool(affine)51        if self.affine:52            self.weight = nn.Parameter(torch.ones(hidden_size))53            self.bias = nn.Parameter(torch.zeros(hidden_size))54        else:55            self.weight, self.bias = None, None56 57    def forward(self, x):58        dims = tuple(-(i + 1) for i in range(len(self.hidden_size)))59        means = x.mean(dims, keepdim=True)60        x_zeromean = x - means61        variances = x_zeromean.pow(2).mean(dims, keepdim=True)62        x = x_zeromean / torch.sqrt(variances + self.eps)63        if self.affine:64            x = (self.weight * x) + self.bias65        return x66 67 68try:69    from apex.normalization import FusedLayerNorm as _FusedLayerNorm70 71    class ESM1bLayerNorm(_FusedLayerNorm):72        @torch.jit.unused73        def forward(self, x):74            if not x.is_cuda:75                return super().forward(x)76            else:77                with torch.cuda.device(x.device):78                    return super().forward(x)79 80except ImportError:81    from torch.nn import LayerNorm as ESM1bLayerNorm82 83 84class TransformerLayer(nn.Module):85    """Transformer layer block."""86 87    def __init__(88        self,89        embed_dim,90        ffn_embed_dim,91        attention_heads,92        add_bias_kv=True,93        use_esm1b_layer_norm=False,94        use_rotary_embeddings: bool = False,95    ):96        super().__init__()97        self.embed_dim = embed_dim98        self.ffn_embed_dim = ffn_embed_dim99        self.attention_heads = attention_heads100        self.use_rotary_embeddings = use_rotary_embeddings101        self._init_submodules(add_bias_kv, use_esm1b_layer_norm)102 103    def _init_submodules(self, add_bias_kv, use_esm1b_layer_norm):104        BertLayerNorm = ESM1bLayerNorm if use_esm1b_layer_norm else ESM1LayerNorm105 106        self.self_attn = MultiheadAttention(107            self.embed_dim,108            self.attention_heads,109            add_bias_kv=add_bias_kv,110            add_zero_attn=False,111            use_rotary_embeddings=self.use_rotary_embeddings,112        )113        self.self_attn_layer_norm = BertLayerNorm(self.embed_dim)114 115        self.fc1 = nn.Linear(self.embed_dim, self.ffn_embed_dim)116        self.fc2 = nn.Linear(self.ffn_embed_dim, self.embed_dim)117 118        self.final_layer_norm = BertLayerNorm(self.embed_dim)119 120    def forward(121        self, x, self_attn_mask=None, self_attn_padding_mask=None, need_head_weights=False122    ):123        residual = x124        x = self.self_attn_layer_norm(x)125        x, attn = self.self_attn(126            query=x,127            key=x,128            value=x,129            key_padding_mask=self_attn_padding_mask,130            need_weights=True,131            need_head_weights=need_head_weights,132            attn_mask=self_attn_mask,133        )134        x = residual + x135 136        residual = x137        x = self.final_layer_norm(x)138        x = gelu(self.fc1(x))139        x = self.fc2(x)140        x = residual + x141 142        return x, attn143 144 145class AxialTransformerLayer(nn.Module):146    """Implements an Axial MSA Transformer block."""147 148    def __init__(149        self,150        embedding_dim: int = 768,151        ffn_embedding_dim: int = 3072,152        num_attention_heads: int = 8,153        dropout: float = 0.1,154        attention_dropout: float = 0.1,155        activation_dropout: float = 0.1,156        max_tokens_per_msa: int = 2**14,157    ) -> None:158        super().__init__()159 160        # Initialize parameters161        self.embedding_dim = embedding_dim162        self.dropout_prob = dropout163 164        row_self_attention = RowSelfAttention(165            embedding_dim,166            num_attention_heads,167            dropout=dropout,168            max_tokens_per_msa=max_tokens_per_msa,169        )170 171        column_self_attention = ColumnSelfAttention(172            embedding_dim,173            num_attention_heads,174            dropout=dropout,175            max_tokens_per_msa=max_tokens_per_msa,176        )177 178        feed_forward_layer = FeedForwardNetwork(179            embedding_dim,180            ffn_embedding_dim,181            activation_dropout=activation_dropout,182            max_tokens_per_msa=max_tokens_per_msa,183        )184 185        self.row_self_attention = self.build_residual(row_self_attention)186        self.column_self_attention = self.build_residual(column_self_attention)187        self.feed_forward_layer = self.build_residual(feed_forward_layer)188 189    def build_residual(self, layer: nn.Module):190        return NormalizedResidualBlock(191            layer,192            self.embedding_dim,193            self.dropout_prob,194        )195 196    def forward(197        self,198        x: torch.Tensor,199        self_attn_mask: Optional[torch.Tensor] = None,200        self_attn_padding_mask: Optional[torch.Tensor] = None,201        need_head_weights: bool = False,202    ):203        """204        LayerNorm is applied either before or after the self-attention/ffn205        modules similar to the original Transformer implementation.206        """207        x, row_attn = self.row_self_attention(208            x,209            self_attn_mask=self_attn_mask,210            self_attn_padding_mask=self_attn_padding_mask,211        )212        x, column_attn = self.column_self_attention(213            x,214            self_attn_mask=self_attn_mask,215            self_attn_padding_mask=self_attn_padding_mask,216        )217        x = self.feed_forward_layer(x)218        if need_head_weights:219            return x, column_attn, row_attn220        else:221            return x222 223 224class LearnedPositionalEmbedding(nn.Embedding):225    """226    This module learns positional embeddings up to a fixed maximum size.227    Padding ids are ignored by either offsetting based on padding_idx228    or by setting padding_idx to None and ensuring that the appropriate229    position ids are passed to the forward function.230    """231 232    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int):233        if padding_idx is not None:234            num_embeddings_ = num_embeddings + padding_idx + 1235        else:236            num_embeddings_ = num_embeddings237        super().__init__(num_embeddings_, embedding_dim, padding_idx)238        self.max_positions = num_embeddings239 240    def forward(self, input: torch.Tensor):241        """Input is expected to be of size [bsz x seqlen]."""242        if input.size(1) > self.max_positions:243            raise ValueError(244                f"Sequence length {input.size(1)} above maximum "245                f" sequence length of {self.max_positions}"246            )247        mask = input.ne(self.padding_idx).int()248        positions = (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + self.padding_idx249        return F.embedding(250            positions,251            self.weight,252            self.padding_idx,253            self.max_norm,254            self.norm_type,255            self.scale_grad_by_freq,256            self.sparse,257        )258 259 260class SinusoidalPositionalEmbedding(nn.Module):261    def __init__(self, embed_dim, padding_idx, learned=False):262        super().__init__()263        self.embed_dim = embed_dim264        self.padding_idx = padding_idx265        self.register_buffer("_float_tensor", torch.FloatTensor(1))266        self.weights = None267 268    def forward(self, x):269        bsz, seq_len = x.shape270        max_pos = self.padding_idx + 1 + seq_len271        if self.weights is None or max_pos > self.weights.size(0):272            self.weights = self.get_embedding(max_pos)273        self.weights = self.weights.type_as(self._float_tensor)274 275        positions = self.make_positions(x)276        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()277 278    def make_positions(self, x):279        mask = x.ne(self.padding_idx)280        range_buf = torch.arange(x.size(1), device=x.device).expand_as(x) + self.padding_idx + 1281        positions = range_buf.expand_as(x)282        return positions * mask.long() + self.padding_idx * (1 - mask.long())283 284    def get_embedding(self, num_embeddings):285        half_dim = self.embed_dim // 2286        emb = math.log(10000) / (half_dim - 1)287        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)288        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)289        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)290        if self.embed_dim % 2 == 1:291            # zero pad292            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)293        if self.padding_idx is not None:294            emb[self.padding_idx, :] = 0295        return emb296 297 298class RobertaLMHead(nn.Module):299    """Head for masked language modeling."""300 301    def __init__(self, embed_dim, output_dim, weight):302        super().__init__()303        self.dense = nn.Linear(embed_dim, embed_dim)304        self.layer_norm = ESM1bLayerNorm(embed_dim)305        self.weight = weight306        self.bias = nn.Parameter(torch.zeros(output_dim))307 308    def forward(self, features):309        x = self.dense(features)310        x = gelu(x)311        x = self.layer_norm(x)312        # project back to size of vocabulary with bias313        x = F.linear(x, self.weight) + self.bias314        return x315 316 317class ContactPredictionHead(nn.Module):318    """Performs symmetrization, apc, and computes a logistic regression on the output features"""319 320    def __init__(321        self,322        in_features: int,323        prepend_bos: bool,324        append_eos: bool,325        bias=True,326        eos_idx: Optional[int] = None,327    ):328        super().__init__()329        self.in_features = in_features330        self.prepend_bos = prepend_bos331        self.append_eos = append_eos332        if append_eos and eos_idx is None:333            raise ValueError("Using an alphabet with eos token, but no eos token was passed in.")334        self.eos_idx = eos_idx335        self.regression = nn.Linear(in_features, 1, bias)336        self.activation = nn.Sigmoid()337 338    def forward(self, tokens, attentions):339        # remove eos token attentions340        if self.append_eos:341            eos_mask = tokens.ne(self.eos_idx).to(attentions)342            eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2)343            attentions = attentions * eos_mask[:, None, None, :, :]344            attentions = attentions[..., :-1, :-1]345        # remove cls token attentions346        if self.prepend_bos:347            attentions = attentions[..., 1:, 1:]348        batch_size, layers, heads, seqlen, _ = attentions.size()349        attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen)350 351        # features: B x C x T x T352        attentions = attentions.to(353            self.regression.weight.device354        )  # attentions always float32, may need to convert to float16355        attentions = apc(symmetrize(attentions))356        attentions = attentions.permute(0, 2, 3, 1)357        return self.activation(self.regression(attentions).squeeze(3))358 359 360class NormalizedResidualBlock(nn.Module):361    def __init__(362        self,363        layer: nn.Module,364        embedding_dim: int,365        dropout: float = 0.1,366    ):367        super().__init__()368        self.embedding_dim = embedding_dim369 370        self.layer = layer371        self.dropout_module = nn.Dropout(372            dropout,373        )374        self.layer_norm = ESM1bLayerNorm(self.embedding_dim)375 376    def forward(self, x, *args, **kwargs):377        residual = x378        x = self.layer_norm(x)379        outputs = self.layer(x, *args, **kwargs)380        if isinstance(outputs, tuple):381            x, *out = outputs382        else:383            x = outputs384            out = None385 386        x = self.dropout_module(x)387        x = residual + x388 389        if out is not None:390            return (x,) + tuple(out)391        else:392            return x393 394 395class FeedForwardNetwork(nn.Module):396    def __init__(397        self,398        embedding_dim: int,399        ffn_embedding_dim: int,400        activation_dropout: float = 0.1,401        max_tokens_per_msa: int = 2**14,402    ):403        super().__init__()404        self.embedding_dim = embedding_dim405        self.ffn_embedding_dim = ffn_embedding_dim406        self.max_tokens_per_msa = max_tokens_per_msa407        self.activation_fn = nn.GELU()408        self.activation_dropout_module = nn.Dropout(409            activation_dropout,410        )411        self.fc1 = nn.Linear(embedding_dim, ffn_embedding_dim)412        self.fc2 = nn.Linear(ffn_embedding_dim, embedding_dim)413 414    def forward(self, x):415        x = self.activation_fn(self.fc1(x))416        x = self.activation_dropout_module(x)417        x = self.fc2(x)418        return x419