Upload 3 files

__init__.py

@@ -0,0 +1,7 @@
+from transformers import AutoConfig, AutoModel
+
+from .configuration import MetaLATTEConfig
+from .model import MultitaskProteinModel
+
+AutoConfig.register("metalatte", MetaLATTEConfig)
+AutoModel.register(MetaLATTEConfig, MultitaskProteinModel)

configuration.py

@@ -0,0 +1,47 @@
+from transformers import PretrainedConfig
+
+class MetaLATTEConfig(PretrainedConfig):
+    model_type = "metalatte"
+
+    def __init__(
+        self,
+        num_labels=15,
+        hidden_size=1280,
+        num_hidden_layers=33,
+        num_attention_heads=20,
+        intermediate_size=5120,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.0,
+        attention_probs_dropout_prob=0.0,
+        max_position_embeddings=1026,
+        initializer_range=0.02,
+        layer_norm_eps=1e-5,
+        esm_model_name="facebook/esm2_t33_650M_UR50D",
+        num_layers_to_finetune=2,
+        num_linear_layers=3,
+        hidden_dim=512,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.num_labels = num_labels
+        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_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.esm_model_name = esm_model_name
+        self.num_layers_to_finetune = num_layers_to_finetune
+        self.num_linear_layers = num_linear_layers
+        self.hidden_dim = hidden_dim
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        return super().from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+    def save_pretrained(self, save_directory):
+        super().save_pretrained(save_directory)

model.py

@@ -0,0 +1,275 @@
+import os
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:4096' # do this before importing pytorch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers import EsmModel
+import torch
+import numpy as np
+from lightning.pytorch import seed_everything
+from typing import Tuple
+import torch
+import gc
+from torch.optim.lr_scheduler import _LRScheduler
+from transformers import EsmModel, PreTrainedModel
+from configuration import MetaLATTEConfig
+
+seed_everything(42)
+   
+class GELU(nn.Module):
+    """Implementation of the gelu activation function.
+
+    For information: OpenAI GPT's gelu is slightly different
+    (and gives slightly different results):
+    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+    """
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1) # x: B, L, H, hidden # x1: B, L, H, hidden // 2
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(x, cos, sin):
+    # Assuming x has shape (B, L, H, HIDDEN_DIM)
+    # cos and sin have shape (1, L, HIDDEN_DIM)
+    cos = cos.unsqueeze(2)  # (1, L, 1, HIDDEN_DIM)
+    sin = sin.unsqueeze(2)  # (1, L, 1, HIDDEN_DIM)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    """
+    The rotary position embeddings from RoFormer_ (Su et. al).
+    A crucial insight from the method is that the query and keys are
+    transformed by rotation matrices which depend on the relative positions.
+    Other implementations are available in the Rotary Transformer repo_ and in
+    GPT-NeoX_, GPT-NeoX was an inspiration
+    .. _RoFormer: https://arxiv.org/abs/2104.09864
+    .. _repo: https://github.com/ZhuiyiTechnology/roformer
+    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
+    .. warning: Please note that this embedding is not registered on purpose, as it is transformative
+        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis
+    """
+
+    def __init__(self, dim: int, *_, **__):
+        super().__init__()
+        # Generate and save the inverse frequency buffer (non trainable)
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(self, x, seq_dimension=1):
+        seq_len = x.shape[seq_dimension]
+
+        # Reset the tables if the sequence length has changed,
+        # or if we're on a new device (possibly due to tracing for instance)
+        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
+            self._seq_len_cached = seq_len
+            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq) # L， 256
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device) # L, 512
+
+            self._cos_cached = emb.cos()[None, :, :] # 1, L, 512
+            self._sin_cached = emb.sin()[None, :, :] # 1, L, 512
+
+        return self._cos_cached, self._sin_cached
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k)
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached), # B, L, H, hidden
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )
+
+
+def macro_f1(y_true, y_pred, thresholds):
+    y_pred_binary = (y_pred >= thresholds).float()
+    tp = (y_true * y_pred_binary).sum(dim=0)
+    fp = ((1 - y_true) * y_pred_binary).sum(dim=0)
+    fn = (y_true * (1 - y_pred_binary)).sum(dim=0)
+    precision = tp / (tp + fp + 1e-7)
+    recall = tp / (tp + fn + 1e-7)
+    f1 = 2 * precision * recall / (precision + recall + 1e-7)
+    macro_f1 = f1.mean()
+    return macro_f1
+
+def safeguard_softmax(logits, dim=-1):
+    # remove max number to prevent exp() to be INF
+    max_logits, _ = logits.max(dim=dim, keepdim=True)
+    exp_logits = torch.exp(logits - max_logits)
+    exp_sum = exp_logits.sum(dim=dim, keepdim=True)
+    probs = exp_logits / (exp_sum + 1e-7)  # Adding a small epsilon to prevent division by zero
+    return probs
+
+class PositionalAttentionHead(nn.Module):
+    def __init__(self, hidden_dim, n_heads):
+        super(PositionalAttentionHead, self).__init__()
+        self.n_heads = n_heads
+        self.hidden_dim = hidden_dim
+        self.head_dim = hidden_dim // n_heads
+        self.preattn_ln = nn.LayerNorm(self.head_dim)
+        self.Q = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.K = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.V = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.rot_emb = RotaryEmbedding(self.head_dim)
+
+    def forward(self, x, attention_mask):
+        batch_size, seq_len, _ = x.size() # B, L, H
+        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        x = self.preattn_ln(x)
+
+        q = self.Q(x)
+        k = self.K(x)
+        v = self.V(x)
+
+        q, k = self.rot_emb(q, k)
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        attn_scores = torch.einsum('bqhd,bkhd->bhqk', q, k) / math.sqrt(self.head_dim)
+        #print(attention_mask.unsqueeze(1).shape)
+        #print(attention_mask.unsqueeze(1).unsqueeze(1).shape)
+        attn_scores = attn_scores.masked_fill(torch.logical_not(attention_mask.unsqueeze(1).unsqueeze(1)), float("-inf")) # B, H, L, L
+
+        attn_probs = safeguard_softmax(attn_scores, dim=-1)
+
+        x = torch.einsum('bhqk,bkhd->bqhd', attn_probs, v)
+        x = x.reshape(batch_size, seq_len, self.hidden_dim)  # B, L, H
+        gc.collect()
+        torch.cuda.empty_cache()
+        return x, attn_probs
+
+class CosineAnnealingWithWarmup(_LRScheduler):
+    # Implement based on Llama paper's description
+    # https://arxiv.org/abs/2302.13971
+    def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.eta_ratio = eta_ratio  # The ratio of minimum to maximum learning rate
+        super(CosineAnnealingWithWarmup, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs]
+
+        progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+        cosine_decay = 0.5 * (1 + np.cos(np.pi * progress))
+        decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio
+
+        return [decayed_lr * base_lr for base_lr in self.base_lrs]
+    
+class RobertaLMHead(nn.Module):
+    """Head for masked language modeling."""
+    def __init__(self, embed_dim, output_dim, weight):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+        self.weight = weight
+        self.gelu = GELU()
+        self.bias = nn.Parameter(torch.zeros(output_dim))
+    def forward(self, features):
+        x = self.dense(features)
+        x = self.gelu(x)
+        x = self.layer_norm(x)
+        # project back to size of vocabulary with bias
+        x = F.linear(x, self.weight) + self.bias
+        return x
+
+       
+class MultitaskProteinModel(PreTrainedModel):
+    config_class = MetaLATTEConfig  
+    base_model_prefix = "metalatte"
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.esm_model = EsmModel.from_pretrained(self.config.esm_model_name)      
+        # layer freezing for the original esm model
+        # first freeze all
+        for param in self.esm_model.parameters():
+            param.requires_grad = False
+        # unfreeze the required layers
+        for i in range(config.num_layers_to_finetune):
+            for param in self.esm_model.encoder.layer[-i-1].parameters():
+                param.requires_grad = True
+        self.lm_head = RobertaLMHead(embed_dim = 1280, output_dim=33, weight=self.esm_model.embeddings.word_embeddings.weight)
+        # esm_dim should be 1280
+        self.attn_head = PositionalAttentionHead(self.config.hidden_size, self.config.num_attention_heads)
+        self.attn_ln = nn.LayerNorm(self.config.hidden_size)
+        self.attn_skip = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+        self.linear_layers = nn.ModuleList()
+        # Add linear layers after the attention head
+        for _ in range(self.config.num_linear_layers):
+            self.linear_layers.append(nn.Linear(self.config.hidden_size, self.config.hidden_size))
+        self.reduction_layers = nn.Sequential(
+            nn.Linear(self.config.hidden_size, self.config.hidden_dim),
+            GELU(),
+            nn.Linear(self.config.hidden_dim, self.config.num_labels)
+        )
+        self.clf_ln = nn.LayerNorm(self.config.hidden_size)
+        self.classification_thresholds = nn.Parameter(torch.tensor([0.5]*self.config.num_labels))
+
+        # Initialize weights and apply final processing
+        self.post_init()
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        if config is None:
+            config = MetaLATTEConfig.from_pretrained(pretrained_model_name_or_path)
+        
+        model = cls(config)        
+        state_dict = torch.load(f"{pretrained_model_name_or_path}/pytorch_model.bin", map_location=torch.device('cpu'))['state_dict']
+        model.load_state_dict(state_dict, strict=False)
+        return model
+        
+    def forward(self, input_ids, attention_mask=None):
+        outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
+        embeddings = outputs.last_hidden_state
+        attention_masks = attention_mask
+
+        x_pool, x_attns = self.attn_head(embeddings, attention_masks)
+        x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool))  # Added skip connection for the attention layer
+
+        for linear_layer in self.linear_layers:
+            residue = x_pool
+            x_pool = linear_layer(x_pool)  # 1280 -> 1280
+            x_pool = F.silu(x_pool)
+            x_pool = x_pool + residue  # Skip connection
+
+        x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool)  # (B, H, L, 1280)
+        x_combined = x_weighted.mean(dim=1)  # Average over heads: (B, L, 1280)
+        x_combined = self.clf_ln(x_combined)
+
+        mlm_logits = self.lm_head(x_combined)
+        attention_masks = attention_masks.unsqueeze(-1).float()  # (B, L, 1)
+        attention_sum = attention_masks.sum(dim=1, keepdim=True)  # (B, 1, 1)
+        x_combined_masked = (x_combined * attention_masks).sum(dim=1) / attention_sum.squeeze(1)  # (B, 1280)
+
+        # Compute classification logits
+        x_pred = self.reduction_layers(x_combined_masked)
+        gc.collect()
+        torch.cuda.empty_cache()
+        return x_pred, x_attns, x_combined_masked, mlm_logits
+
+    def predict(self, input_ids, attention_mask=None):
+        x_pred, _, _, _ = self.forward(input_ids, attention_mask)
+        classification_output = torch.sigmoid(x_pred)
+        predictions = (classification_output >= self.classification_thresholds).float()
+
+        for i, pred in enumerate(predictions):
+            if pred.sum() == 0:
+                weighted_probs = classification_output[i]
+                max_class = torch.argmax(weighted_probs)
+                predictions[i, max_class] = 1.0
+
+        return classification_output, predictions
+
+