Upload code

modeling_prot2text.py

@@ -1,88 +1,16 @@
 from transformers import GPT2Config, AutoTokenizer, GPT2Config
 from transformers import PretrainedConfig, PreTrainedModel
 import transformers
-from typing import Optional, Tuple, Callable
+from typing import Optional, Tuple, Callable, List
 import torch
 import torch.nn as nn
 from transformers.modeling_utils import PreTrainedModel, PretrainedConfig
 from .utils import CABlock, _GPT2LMHeadModel
 from .configuration_prot2text import Prot2TextConfig
-import os
-import numpy as np
 from transformers.generation.configuration_utils import GenerationConfig
 from transformers.generation.logits_process import LogitsProcessorList
 from transformers.generation.stopping_criteria import StoppingCriteriaList
 
-from .pdb2graph import PDB2Graph, download_alphafold_structure
-from .graphs import *
-from .utils_dataset import *
-
-try:
-    from graphein.protein.config import ProteinGraphConfig, DSSPConfig
-    from graphein.protein.features.nodes.amino_acid import amino_acid_one_hot, meiler_embedding, expasy_protein_scale, hydrogen_bond_acceptor, hydrogen_bond_donor
-    from graphein.protein.features.nodes.dssp import  phi, psi, asa, rsa, secondary_structure
-    from graphein.protein.edges.distance import (add_peptide_bonds,
-                                                add_hydrogen_bond_interactions,
-                                                add_distance_threshold,
-                                                )
-except ImportError:
-    raise Exception('You need to install graphein from source in addition to DSSP to use this model please refer to https://github.com/a-r-j/graphein and https://ssbio.readthedocs.io/en/latest/instructions/dssp.html')
-
-try:
-    from torch_geometric.nn import RGCNConv, global_mean_pool
-except ImportError:  
-    raise Exception('You need to install torch geometric and its dependecies to use this model please refer to https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html')
-
-
-
-class EncoderRGCN(PreTrainedModel):
-    '''
-    This class implement the RGCN encoder to encode the protein structure
-    '''
-    def __init__(self, input_dim, hidden_dim=512, n_layers=6, emb_dim=512, dropout=0.2, num_relation=7, prot2text_version='1.0'):
-        super(EncoderRGCN, self).__init__(PretrainedConfig(name='RGCN'))
-        self.n_layers = n_layers
-        self.output_dim = emb_dim
-        self.prot2text_version = prot2text_version
-
-        self.fc0 = nn.Linear(input_dim, hidden_dim)
-        self.batchnorm_final = nn.BatchNorm1d(hidden_dim)
-        
-        self.batch_norms = nn.ModuleList()
-        self.batch_norms.append(nn.BatchNorm1d(hidden_dim))
-        lst = list()
-        
-        lst.append(RGCNConv(hidden_dim, hidden_dim, num_relations=num_relation))
-            
-        for i in range(n_layers-1):
-            lst.append(RGCNConv(hidden_dim,hidden_dim, num_relations=num_relation))
-
-        self.conv = nn.ModuleList(lst)
-      
-        self.fc1 = nn.Linear(hidden_dim, hidden_dim)
-        self.fc2 = nn.Linear(hidden_dim, self.output_dim)
-      
-        self.dropout = nn.Dropout(p=dropout)
-        self.relu = nn.LeakyReLU()
-        self.batchnorm = nn.BatchNorm1d(hidden_dim)
-        self.main_input_name = 'nothing'
-
-    def forward(self, x:Optional[torch.FloatTensor] = None, 
-                edge_index:Optional[torch.LongTensor] = None,
-                edge_type:Optional[torch.LongTensor] = None,
-                batch:Optional[torch.LongTensor] = None,
-                **kargs):
-        #construct pyg edge index shape (2, num_edges) from edge_list
-        x = self.relu(self.fc0(x))
-        
-        for i in range(self.n_layers):
-            x = self.conv[i](x, edge_index, edge_type)
-
-        out = global_mean_pool(x, batch)
-        out = self.relu(self.fc1(out))
-        out = self.relu(self.fc2(out))
-        
-        return out.unsqueeze(1)
 
 class Prot2TextModel(PreTrainedModel):
     config_class = Prot2TextConfig
@@ -92,10 +20,6 @@ class Prot2TextModel(PreTrainedModel):
         super().__init__(config)
 
         self.gpt_config = GPT2Config.from_dict(config.gpt_config)
-        
-        # if we are using RGCN to encode the protein's structure, define the RGCN encoder
-        if config.rgcn:
-            self.encoder = EncoderRGCN(input_dim=config.rgcn_input_dim, hidden_dim=self.gpt_config.n_embd, n_layers=config.rgcn_n_layers, emb_dim=self.gpt_config.n_embd, prot2text_version=self.config.prot2text_version)
 
         # define the GPT2 decoder
         self.decoder = _GPT2LMHeadModel(self.gpt_config)
@@ -164,10 +88,6 @@ class Prot2TextModel(PreTrainedModel):
         if decoder_input_ids is not None and len(decoder_input_ids.size()) == 3:
             decoder_input_ids = decoder_input_ids.squeeze(0) 
 
-        if x is not None and self.config.rgcn:
-            graph_emb = self.encoder(x, edge_index, edge_type, batch)
-            graph_mask = None
-            
         if self.config.esm:
             if self.config.prot2text_version=='1.0':
                 if encoder_input_ids.size()[1] != 1021:
@@ -175,38 +95,7 @@ class Prot2TextModel(PreTrainedModel):
             
             esm_emb = self.esm(input_ids=encoder_input_ids, attention_mask=attention_mask, return_dict=return_dict).last_hidden_state
             esm_emb = self.to_embedding(esm_emb)
-            if not self.config.cross_esm_graph and self.config.rgcn:
-                graph_emb = torch.cat((graph_emb, esm_emb), dim=1)
-                t_add = torch.ones((attention_mask.size(0), 1)).to(attention_mask.get_device())
-                attention_mask = torch.cat((t_add, attention_mask), dim=1) 
-            elif self.config.cross_esm_graph and self.config.rgcn:
-                if past_key_values_graph_esm is None:
-                    past_length = 0
-                    past_key_values_graph_esm = tuple([None] * len(self.h))
-                else:
-                    past_length = past_key_values_graph_esm[0][0].size(-2) 
-                output_shape = esm_emb.size()
-                
-                all_self_attentions = () if output_attentions else None
-                all_cross_attentions = () if output_attentions and self.gpt_config.add_cross_attention else None
-                all_hidden_states = () if output_hidden_states else None
-                for i, (block, layer_past) in enumerate(zip(self.h, past_key_values_graph_esm)):
-                    outputs = block(
-                        esm_emb,
-                        layer_past=layer_past,
-                        attention_mask=attention_mask,
-                        encoder_hidden_states=graph_emb,
-                        encoder_attention_mask=graph_mask,
-                        use_cache=use_cache,
-                        output_attentions=False,
-                    )
-                    esm_emb = outputs[0]
-
-                esm_emb = self.ln_f(esm_emb)
-                esm_emb = esm_emb.view(output_shape)  
-                graph_emb = esm_emb
-            else:
-                graph_emb = esm_emb
+            graph_emb = esm_emb
         else:
             attention_mask = None
         if self.config.prot2text_version=='1.0':
@@ -234,11 +123,7 @@ class Prot2TextModel(PreTrainedModel):
     
     @torch.no_grad()    
     def generate_protein_description(self,
-                                    protein_pdbID=None, 
                                     protein_sequence=None,
-                                    edge_index: Optional[torch.LongTensor] = None,
-                                    x: Optional[torch.FloatTensor] = None,
-                                    edge_type: Optional[torch.LongTensor] = None,
                                     tokenizer=None,
                                     device='cpu'
                                      ):
@@ -247,120 +132,24 @@ class Prot2TextModel(PreTrainedModel):
             raise ValueError(
                 "The model you are trying to use is based only on protein sequence, please provide an amino-acid protein_sequence"
             )
-        if self.config.rgcn and protein_pdbID==None and (x==None or edge_index==None or edge_type==None):
-            raise ValueError(
-                "The model you are trying to use is based on protein structure, please provide a AlphaFold ID (you must have to have internet connection using protein_pdbID, or provide the triplet inputs: x (node features), edge_index and edge_type"
-            )
         if self.config.esm:
             esmtokenizer = AutoTokenizer.from_pretrained(self.config.esm_model_name)
-        
-        if protein_pdbID==None and protein_sequence==None:
-            raise ValueError(
-                "you need to provide either a protein AlphaFold Id or an amino-acid sequence"
-            )
-            
-        if protein_pdbID!=None:
-            config = {"node_metadata_functions": [amino_acid_one_hot, 
-                                                expasy_protein_scale,
-                                                meiler_embedding,
-                                                hydrogen_bond_acceptor, hydrogen_bond_donor
-                                                ],
-                    "edge_construction_functions": [add_peptide_bonds,
-                                                    add_hydrogen_bond_interactions,
-                                                    partial(add_distance_threshold, long_interaction_threshold=3, threshold=10.),],
-                    "graph_metadata_functions":[asa,phi, psi, secondary_structure, rsa],
-                    "dssp_config": DSSPConfig()}
-            config = ProteinGraphConfig(**config)
-
-            PATH_TO_DATA = f"~/.tmp/pdb/pdb"
-            OUTPUT_FOLDER = f"~/.tmp/pdb/raw"
-            save_dir = f"~/.tmp/pdb/"
-            isExist = os.path.exists(PATH_TO_DATA)
-            if not isExist:
-                os.makedirs(PATH_TO_DATA)
-            isExist = os.path.exists(OUTPUT_FOLDER)
-            if not isExist:
-                os.makedirs(OUTPUT_FOLDER)
-            isExist = os.path.exists(save_dir+'processed')
-            if not isExist:
-                os.makedirs(save_dir+'processed')
-            
-            structure_filename = download_alphafold_structure(uniprot_id=protein_pdbID, out_dir=PATH_TO_DATA)
-            if structure_filename is None:
-                raise ValueError("Error! the ID does not exist in AlphaFoldDB or you do not have internet connection")
-            graph_filename = structure_filename.split('/')
-            graph_filename[-2] = 'raw'
-            graph_filename[-1] = graph_filename[-1].replace('.pdb', '.pt')
-            graph_filename = '/'.join(graph_filename)
-            process_filename = structure_filename.split('/')
-            process_filename[-2] = 'processed'
-            process_filename[-1] = process_filename[-1].replace('.pdb', '.pt')
-            process_filename = '/'.join(process_filename)    
-            try:            
-                gpdb = PDB2Graph(root = PATH_TO_DATA, output_folder = OUTPUT_FOLDER, config=config, n_processors=1).create_pyg_graph(structure_filename)
-                seq = esmtokenizer(gpdb.sequence, add_special_tokens=True, truncation=True, max_length=1021, padding='max_length',return_tensors="pt") #
-                torch.save(gpdb, graph_filename)
-                gpdb.edge_type = [np.array(gpdb.edge_type.transpose(0,1))]
-                gpdb.encoder_input_ids = seq['input_ids']
-                gpdb.attention_mask = seq['attention_mask']
-                torch.save(gpdb, process_filename)
-            except:
-                os.remove(structure_filename)
-                raise ValueError('creating graphs did not work, probably the pdb file of alphaFold is damaged')
-            
-            self.eval()
-            inputs = gpdb
-            inputs = inputs.to_dict()
-            
-            inputs['edge_type'] =  torch.cat([torch.tensor(inputs['edge_type'][i]) for i in range(len(inputs['edge_type']))], dim=0)
-            inputs['edge_type'] = torch.argmax(inputs['edge_type'], dim=1)
-            for key in ['num_nodes', 'node_id', 'name', 'sequence', 'distance_matrix', 'distance', 'coordinates']:
-                inputs.pop(key)
-            inputs['decoder_input_ids'] = inputs['encoder_input_ids'][:,0:1].clone()
-            inputs['decoder_input_ids'][:,0] = tokenizer.bos_token_id
-            inputs["decoder_attention_mask"] = torch.ones(inputs['decoder_input_ids'].shape[0], 1)
-            self.to(device)
-            inputs = {k: v.to(device=device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
-            encoder_state = dict()
-            encoder_state['hidden_states'] = self(**inputs, get_graph_emb=True, output_attentions=True)
-            encoder_state['attentions'] = inputs['attention_mask']
-            for key in ['edge_index', 'edge_type', 'x', 'encoder_input_ids']:
-                inputs.pop(key)
-            tok_ids = self.decoder.generate(input_ids=inputs['decoder_input_ids'], 
-                                            encoder_outputs=encoder_state, 
-                                            use_cache=True, 
-                                            output_attentions=False, 
-                                            output_scores=False, 
-                                            return_dict_in_generate=True, 
-                                            encoder_attention_mask=inputs['attention_mask'], 
-                                            length_penalty=1.0,
-                                            no_repeat_ngram_size=None,
-                                            early_stopping=False,
-                                            num_beams=1)
-
-            generated = tokenizer.batch_decode(tok_ids.get('sequences'), skip_special_tokens=True)
-
-            os.remove(structure_filename)
-            os.remove(graph_filename)
-            os.remove(process_filename)        
-                
-            return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
-            
-        else:
-            seq = esmtokenizer([protein_sequence], add_special_tokens=True, truncation=True, max_length=1021, padding='max_length', return_tensors="pt")
-            inputs={}
-            inputs['encoder_input_ids'] = seq['input_ids']
-            inputs['attention_mask'] = seq['attention_mask']
-            inputs['decoder_input_ids'] = inputs['encoder_input_ids'][:,0:1].clone()
-            inputs['decoder_input_ids'][:,0] = tokenizer.bos_token_id
             
-            self.to(device)
-            inputs = {k: v.to(device=device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
-            encoder_state = dict()
-            encoder_state['hidden_states'] = self(**inputs, get_graph_emb=True, output_attentions=True)
-            generated = tokenizer.batch_decode(self.decoder.generate(input_ids=inputs['decoder_input_ids'], encoder_outputs=encoder_state, use_cache=True), skip_special_tokens=True)
-            
-            return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
+        
+        seq = esmtokenizer([protein_sequence], add_special_tokens=True, truncation=True, max_length=1021, padding='max_length', return_tensors="pt")
+        inputs={}
+        inputs['encoder_input_ids'] = seq['input_ids']
+        inputs['attention_mask'] = seq['attention_mask']
+        inputs['decoder_input_ids'] = inputs['encoder_input_ids'][:,0:1].clone()
+        inputs['decoder_input_ids'][:,0] = tokenizer.bos_token_id
+        
+        self.to(device)
+        inputs = {k: v.to(device=device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
+        encoder_state = dict()
+        encoder_state['hidden_states'] = self(**inputs, get_graph_emb=True, output_attentions=True)
+        generated = tokenizer.batch_decode(self.decoder.generate(input_ids=inputs['decoder_input_ids'], encoder_outputs=encoder_state, use_cache=True), skip_special_tokens=True)
+        
+        return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
     
     @torch.no_grad()
     def generate(self,

utils.py

@@ -1,8 +1,7 @@
 import torch.nn as nn
 from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
 from typing import Optional, Tuple, Union, Any, Dict, List
-from transformers import Seq2SeqTrainer, GPT2LMHeadModel
-from torch.utils.data.distributed import DistributedSampler
+from transformers import GPT2LMHeadModel
 import torch
 from transformers.deepspeed import is_deepspeed_zero3_enabled
 from transformers.generation.logits_process import LogitsProcessorList
@@ -10,11 +9,6 @@ from transformers.generation.stopping_criteria import StoppingCriteriaList
 from transformers.generation.utils import GreedySearchOutput, GreedySearchEncoderDecoderOutput, BeamSearchOutput, BeamSearchEncoderDecoderOutput
 from transformers.generation.beam_search import BeamScorer
 
-try:
-    from torch_geometric.loader import DataLoader
-    from torch_geometric.data import Dataset
-except ImportError:  
-    raise Exception('You need to install torch geometric and its dependecies to use this model please refer to https://pytorch-geometric.readthedocs.io/en/latest/install/installation.html')
 
 class _GPT2LMHeadModel(GPT2LMHeadModel):
     def _init_(self, config):
@@ -593,153 +587,3 @@ class CABlock(nn.Module):
 
         return (hidden_states,)
     
-class Prot2TextTrainer(Seq2SeqTrainer):
-    '''
-    This function is an edited version of the Seq2SeqTrainer from HuggingFace's transformers 
-    '''
-    def get_eval_dataloader(self, eval_dataset: Optional[Dataset] = None) -> DataLoader:
-        if self.args.world_size > 1:
-            eval_sampler = DistributedSampler(self.eval_dataset, num_replicas=self.args.world_size, rank=self.args.process_index)
-        else:
-            eval_sampler = None
-        return DataLoader(
-            self.eval_dataset,
-            batch_size=self.args.eval_batch_size,
-            collate_fn=None,
-            num_workers=self.args.dataloader_num_workers,
-            pin_memory=self.args.dataloader_pin_memory,
-            sampler=eval_sampler,
-        )
-    def get_train_dataloader(self) -> DataLoader:
-        if self.args.world_size > 1:
-            train_sampler = DistributedSampler(self.train_dataset, num_replicas=self.args.world_size, rank=self.args.process_index)
-        else:
-            train_sampler = None
-        return DataLoader(
-            self.train_dataset,
-            batch_size=self.args.per_device_train_batch_size,
-            collate_fn=None,
-            num_workers=self.args.dataloader_num_workers,
-            pin_memory=self.args.dataloader_pin_memory,
-            sampler=train_sampler,
-        )
-    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
-        """
-        Prepare `inputs` before feeding them to the model, converting them to tensors if they are not already and
-        handling potential state.
-        """
-        inputs = self._prepare_input(inputs)
-        if len(inputs) == 0:
-            raise ValueError(
-                "The batch received was empty, your model won't be able to train on it. Double-check that your "
-                f"training dataset contains keys expected by the model: {','.join(self._signature_columns)}."
-            )
-        if self.args.past_index >= 0 and self._past is not None:
-            inputs["mems"] = self._past
-        
-        inputs = inputs.to_dict()
-        inputs['edge_type'] =  torch.cat([torch.tensor(inputs['edge_type'][i]) for i in range(len(inputs['edge_type']))], dim=0)
-        inputs['edge_type'] = torch.argmax(inputs['edge_type'], dim=1)
-        inputs = {k: v.to(device=self.args.device, non_blocking=True) if hasattr(v, 'to') else v for k, v in inputs.items()}
-        return inputs
-    
-    def prediction_step(
-        self,
-        model: nn.Module,
-        inputs: Dict[str, Union[torch.Tensor, Any]],
-        prediction_loss_only: bool,
-        ignore_keys: Optional[List[str]] = None,
-        ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
-        """
-        Perform an evaluation step on `model` using `inputs`.
-
-        Subclass and override to inject custom behavior.
-
-        Args:
-            model (`nn.Module`):
-                The model to evaluate.
-            inputs (`Dict[str, Union[torch.Tensor, Any]]`):
-                The inputs and targets of the model.
-
-                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
-                argument `labels`. Check your model's documentation for all accepted arguments.
-            prediction_loss_only (`bool`):
-                Whether or not to return the loss only.
-
-        Return:
-            Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and
-            labels (each being optional).
-        """
-
-        if not self.args.predict_with_generate or prediction_loss_only:
-            return super().prediction_step(
-                model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
-            )
-
-        has_labels = "labels" in inputs
-        inputs = self._prepare_inputs(inputs)
-
-        # XXX: adapt synced_gpus for fairscale as well
-        gen_kwargs = self._gen_kwargs.copy()
-        if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None:
-            gen_kwargs["max_length"] = self.model.config.max_length
-        gen_kwargs["num_beams"] = (
-            gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams
-        )
-        default_synced_gpus = True if is_deepspeed_zero3_enabled() else False
-        gen_kwargs["synced_gpus"] = (
-            gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus
-        )
-
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
-        if "global_attention_mask" in inputs:
-            gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None)
-
-        generation_inputs = None
-        gen_kwargs['x'] = inputs.get('x', None)
-        gen_kwargs['edge_index'] = inputs.get('edge_index', None)
-        gen_kwargs['edge_type'] = inputs.get('edge_type', None)
-        gen_kwargs['batch'] = inputs.get('batch', None)
-        gen_kwargs['encoder_input_ids'] = inputs.get('encoder_input_ids', None)
-        gen_kwargs['decoder_input_ids'] = inputs.get('decoder_input_ids', None)[:,0:1]
-        gen_kwargs["decoder_attention_mask"] = torch.ones(gen_kwargs['decoder_input_ids'].shape[0], 1).to(self.args.device)
-
-        generated_tokens = self.model.generate(
-            generation_inputs,
-            **gen_kwargs,
-        )
-        # in case the batch is shorter than max length, the output should be padded
-        if gen_kwargs.get("max_length") is not None and generated_tokens.shape[-1] < gen_kwargs["max_length"]:
-            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"])
-        elif gen_kwargs.get("max_new_tokens") is not None and generated_tokens.shape[-1] < (
-            gen_kwargs["max_new_tokens"] + 1
-        ):
-            generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_new_tokens"] + 1)
-
-        with torch.no_grad():
-            if has_labels:
-                with self.compute_loss_context_manager():
-                    outputs = model(**inputs)
-                if self.label_smoother is not None:
-                    loss = self.label_smoother(outputs, inputs["labels"]).mean().detach()
-                else:
-                    loss = (outputs["loss"] if isinstance(outputs, dict) else outputs[0]).mean().detach()
-            else:
-                loss = None
-
-        if self.args.prediction_loss_only:
-            return (loss, None, None)
-
-        if has_labels:
-            labels = inputs["labels"]
-            if gen_kwargs.get("max_length") is not None and labels.shape[-1] < gen_kwargs["max_length"]:
-                labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"])
-            elif gen_kwargs.get("max_new_tokens") is not None and labels.shape[-1] < (
-                gen_kwargs["max_new_tokens"] + 1
-            ):
-                labels = self._pad_tensors_to_max_len(labels, (gen_kwargs["max_new_tokens"] + 1))
-        else:
-            labels = None
-
-        return (loss, generated_tokens, labels)