Upload folder using huggingface_hub

ESM_per_token.py

@@ -0,0 +1,150 @@
+from transformers.models.esm.modeling_esm import EsmPreTrainedModel, EsmModel
+import torch
+import torch.nn as nn
+from typing import Optional, Union, Tuple
+from transformers.models.auto.modeling_auto import AutoModel
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+from torch.nn import MSELoss
+from transformers.modeling_outputs import TokenClassifierOutput
+import numpy as np
+import re
+from utils.lora_utils import LoRAConfig, modify_with_lora
+from models.enm_adaptor_heads import (
+    ENMAdaptedAttentionClassifier, ENMAdaptedDirectClassifier, 
+    ENMAdaptedConvClassifier, ENMNoAdaptorClassifier
+)
+from peft import LoraConfig, inject_adapter_in_model
+
+class EsmForTokenRegression(EsmPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+    _keys_to_ignore_on_load_missing = [r"position_ids"]
+
+    def __init__(self, config, class_config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.add_pearson_loss = class_config.add_pearson_loss
+        self.add_sse_loss = class_config.add_sse_loss
+
+        self.esm = EsmModel(config, add_pooling_layer=False)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        
+        if class_config.adaptor_architecture == 'attention':
+            self.classifier = ENMAdaptedAttentionClassifier(
+                config.hidden_size, 
+                class_config.num_labels, 
+                class_config.enm_embed_dim, 
+                class_config.enm_att_heads
+            )
+        elif class_config.adaptor_architecture == 'direct':
+            self.classifier = ENMAdaptedDirectClassifier(
+                config.hidden_size, 
+                class_config.num_labels
+            )
+        elif class_config.adaptor_architecture == 'conv':
+            self.classifier = ENMAdaptedConvClassifier(
+                config.hidden_size, 
+                class_config.num_labels, 
+                class_config.kernel_size, 
+                class_config.enm_embed_dim, 
+                class_config.num_layers
+            )
+        elif class_config.adaptor_architecture == 'no-adaptor':
+            self.classifier = ENMNoAdaptorClassifier(
+                config.hidden_size, 
+                class_config.num_labels
+            )
+        else:
+            raise ValueError('Only attention, direct, conv and no-adaptor architectures are supported.')
+
+        self.init_weights()
+
+    def forward(
+        self,
+        enm_vals=None,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, TokenClassifierOutput]:
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.esm(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        
+        logits = self.classifier(sequence_output, enm_vals, attention_mask)
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return output
+
+        return TokenClassifierOutput(
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+def ESM_classification_model(half_precision, class_config, lora_config):
+    # Load ESM and tokenizer
+    if not half_precision:
+        model = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D")
+        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
+    elif half_precision and torch.cuda.is_available():
+        model = EsmModel.from_pretrained("facebook/esm2_t36_3B_UR50D", torch_dtype=torch.float16).to(torch.device('cuda'))
+        tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
+    else:
+        raise ValueError('Half precision can be run on GPU only.')
+    
+    # Create new Classifier model with ESM dimensions
+    class_model = EsmForTokenRegression(model.config, class_config)
+    
+    # Set encoder weights to checkpoint weights
+    class_model.esm = model
+    
+    # Delete the checkpoint model
+    del model
+    
+    # Print number of trainable parameters
+    model_parameters = filter(lambda p: p.requires_grad, class_model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ESM_Classifier\nTrainable Parameter: " + str(params))
+    
+    # Add model modification lora
+    esm_lora_peft_config = LoraConfig(
+        r=4, lora_alpha=1, bias="all", target_modules=["query","key","value","dense"]
+    )
+    
+    # Add LoRA layers
+    class_model.esm = inject_adapter_in_model(esm_lora_peft_config, class_model.esm) 
+    
+    # Freeze Encoder (except LoRA)
+    for (param_name, param) in class_model.esm.named_parameters():
+        param.requires_grad = False
+    
+    for (param_name, param) in class_model.esm.named_parameters():
+        if re.fullmatch(".*lora.*", param_name): #".*layer_norm.*|.*lora_[ab].*"
+            param.requires_grad = True
+        if re.fullmatch(".*layer_norm.*", param_name): #".*layer_norm.*|.*lora_[ab].*"
+            param.requires_grad = True
+    # Print trainable Parameter
+    model_parameters = filter(lambda p: p.requires_grad, class_model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ESM_LoRA_Classifier\nTrainable Parameter: " + str(params) + "\n")
+    
+    return class_model, tokenizer

T5_encoder_per_token.py

@@ -0,0 +1,172 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import copy
+import re
+from transformers import T5Config, T5PreTrainedModel, T5EncoderModel, T5Tokenizer
+from transformers.models.t5.modeling_t5 import T5Stack
+from transformers.modeling_outputs import TokenClassifierOutput
+from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+from models.enm_adaptor_heads import ENMAdaptedAttentionClassifier, ENMAdaptedDirectClassifier, ENMAdaptedConvClassifier, ENMNoAdaptorClassifier
+from utils.lora_utils import LoRAConfig, modify_with_lora
+
+class T5EncoderForTokenClassification(T5PreTrainedModel):
+
+    def __init__(self, config: T5Config, class_config):
+        super().__init__(config)
+        self.num_labels = class_config.num_labels
+        self.config = config
+        self.add_pearson_loss = class_config.add_pearson_loss
+        self.add_sse_loss = class_config.add_sse_loss
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = T5Stack(encoder_config, self.shared)
+
+        self.dropout = nn.Dropout(class_config.dropout_rate)
+        if class_config.adaptor_architecture == 'attention':
+            self.classifier = ENMAdaptedAttentionClassifier(config.hidden_size, class_config.num_labels, class_config.enm_embed_dim, class_config.enm_att_heads) #nn.Linear(config.hidden_size, class_config.num_labels)
+        elif class_config.adaptor_architecture == 'direct':
+            self.classifier = ENMAdaptedDirectClassifier(config.hidden_size, class_config.num_labels)
+        elif class_config.adaptor_architecture == 'conv':
+            self.classifier = ENMAdaptedConvClassifier(config.hidden_size, class_config.num_labels, class_config.kernel_size, class_config.enm_embed_dim, class_config.num_layers)
+        elif class_config.adaptor_architecture == 'no-adaptor':
+            self.classifier = ENMNoAdaptorClassifier(config.hidden_size, class_config.num_labels)
+        else:
+            raise ValueError('Only attention, direct, conv and no-adaptor architectures are supported for the adaptor.')
+
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.encoder.block), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.encoder.block))
+        self.encoder.parallelize(self.device_map)
+        self.classifier = self.classifier.to(self.encoder.first_device)
+        self.model_parallel = True
+
+    def deparallelize(self):
+        self.encoder.deparallelize()
+        self.encoder = self.encoder.to("cpu")
+        self.model_parallel = False
+        self.device_map = None
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    def get_encoder(self):
+        return self.encoder
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def forward(
+        self,
+        enm_vals = None,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        # import pdb; pdb.set_trace()
+        outputs = self.encoder(input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        sequence_output = self.dropout(sequence_output)
+        #TODO: check the enm_vals are padded properly and check that the sequence limit (in the transformer) is indeed 512
+        logits = self.classifier(sequence_output, enm_vals, attention_mask)
+        
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            #loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+def PT5_classification_model(half_precision, class_config):
+    # Load PT5 and tokenizer
+    # possible to load the half preciion model (thanks to @pawel-rezo for pointing that out)
+    if not half_precision:
+        model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50", local_files_only=True)
+        tokenizer = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50", local_files_only=True)
+    elif half_precision and torch.cuda.is_available(): 
+        tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False, local_files_only=True)
+        model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc", torch_dtype=torch.float16, local_files_only=True).to(torch.device('cuda'))
+    else:
+          raise ValueError('Half precision can be run on GPU only.')
+    
+    # Create new Classifier model with PT5 dimensions
+    class_model=T5EncoderForTokenClassification(model.config,class_config)
+    
+    # Set encoder and embedding weights to checkpoint weights
+    class_model.shared=model.shared
+    class_model.encoder=model.encoder    
+    
+    # Delete the checkpoint model
+    model=class_model
+    del class_model
+    
+    # Print number of trainable parameters
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ProtT5_Classfier\nTrainable Parameter: "+ str(params))    
+ 
+    # Add model modification lora
+    config = LoRAConfig('configs/lora_config.yaml')
+    
+    # Add LoRA layers
+    model = modify_with_lora(model, config)
+    
+    # Freeze Embeddings and Encoder (except LoRA)
+    for (param_name, param) in model.shared.named_parameters():
+                param.requires_grad = False
+    for (param_name, param) in model.encoder.named_parameters():
+                param.requires_grad = False       
+
+    for (param_name, param) in model.named_parameters():
+            if re.fullmatch(config.trainable_param_names, param_name):
+                param.requires_grad = True
+
+    # Print trainable Parameter          
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ProtT5_LoRA_Classfier\nTrainable Parameter: "+ str(params) + "\n")
+    
+    return model, tokenizer

enm_adaptor_heads.py

@@ -0,0 +1,85 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ENMAdaptedAttentionClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim, enm_embed_dim, num_att_heads):
+        super(ENMAdaptedAttentionClassifier, self).__init__()
+        self.embedding = nn.Linear(1, enm_embed_dim)
+        self.enm_attention = nn.MultiheadAttention(enm_embed_dim, num_att_heads)
+        self.layer_norm = nn.LayerNorm(enm_embed_dim)
+        self.enm_adaptor = nn.Linear(enm_embed_dim, seq_embedding_dim)
+        self.adapted_classifier = nn.Linear(2*seq_embedding_dim, out_dim)
+    
+    def forward(self, seq_embedding, enm_input):
+        enm_input = enm_input.transpose(0, 1)  # Transpose to shape (N, B, E) for MultiheadAttention
+        enm_input = enm_input.unsqueeze(-1)  # Add a dimension for the embedding
+        enm_input_embedded = self.embedding(enm_input)
+        enm_att, _ = self.enm_attention(enm_input_embedded, enm_input_embedded, enm_input_embedded)
+        enm_att = enm_att.transpose(0, 1)  # Transpose back to shape (B, N, E)
+        enm_att = self.layer_norm(enm_att + enm_input.transpose(0, 1))
+        enm_embedding = self.enm_adaptor(enm_att)
+        # import pdb; pdb.set_trace()
+        combined_embedding = torch.cat((seq_embedding, enm_embedding), dim=-1)
+        logits = self.adapted_classifier(combined_embedding)
+        return logits
+    
+class ENMAdaptedConvClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim, kernel_size, enm_embedding_dim, num_layers):
+        super(ENMAdaptedConvClassifier, self).__init__()
+        layers = []
+        self.conv1 = nn.Conv1d(1, enm_embedding_dim, kernel_size=kernel_size, padding=(kernel_size-1)//2)
+        layers.append(self.conv1)
+        layers.append(nn.ReLU())
+        for i in range(num_layers-1):
+            layers.append(nn.Conv1d(enm_embedding_dim, enm_embedding_dim, kernel_size=kernel_size, padding=(kernel_size-1)//2))
+            layers.append(nn.ReLU())
+        self.conv_net = nn.Sequential(*layers)
+        self.adapted_classifier = nn.Linear(seq_embedding_dim+1, out_dim)
+
+    def forward(self, seq_embedding, enm_input, attention_mask=None):
+        enm_input = torch.nan_to_num(enm_input, nan=0.0)
+        enm_input = enm_input.unsqueeze(1)
+        enm_input = enm_input.to(seq_embedding.device)
+        conv_out = self.conv_net(enm_input)
+        enm_embedding = conv_out.transpose(1,2)
+        
+        if attention_mask is not None:
+            # Use attention_mask to ignore padded elements
+            mask = attention_mask.unsqueeze(-1).float()
+            enm_embedding = enm_embedding * mask
+            # Compute mean over non-padded elements
+            
+            enm_embedding = enm_embedding.mean(dim=-1).unsqueeze(-1)
+            # enm_embedding = enm_embedding.sum(dim=2)/ mask.sum(dim=2).clamp(min=1e-9)
+        else:
+            raise ValueError('We actually want to provide the mask.')
+            enm_embedding = torch.mean(enm_embedding, dim=1)
+            
+        # enm_embedding = enm_embedding.unsqueeze(1).expand(-1, seq_embedding.size(1), -1)
+        combined_embedding = torch.cat((seq_embedding, enm_embedding), dim=-1)
+        logits = self.adapted_classifier(combined_embedding)
+        return logits
+    
+
+        
+class ENMAdaptedDirectClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim):
+        super(ENMAdaptedDirectClassifier, self).__init__()
+        self.adapted_classifier = nn.Linear(seq_embedding_dim+1, out_dim)
+
+    def forward(self, seq_embedding, enm_input):
+            enm_input = enm_input.unsqueeze(-1)
+            combined_embedding = torch.cat((seq_embedding, enm_input), dim=-1)
+            logits = self.adapted_classifier(combined_embedding)
+            return logits
+
+class ENMNoAdaptorClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim):
+        super(ENMNoAdaptorClassifier, self).__init__()
+        self.adapted_classifier = nn.Linear(seq_embedding_dim, out_dim)
+
+    def forward(self, seq_embedding, enm_input, attention_mask=None):
+            _ = enm_input #ignoring enm_input
+            logits = self.adapted_classifier(seq_embedding)
+            return logits

weights/flexpert_3d_weights.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3cbc7a6bed15e92cc6b5f65b947c3c838e46e5815f7cbd57f54bbc19741558e6
+size 4843266070

weights/flexpert_seq_weights.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca9ffd974154597e372c30faa728e4c61c5811fc98a148af66d31dfe2b5c0061
+size 4842603885