Upload code

modeling_prot2text.py

@@ -123,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'
                                      ):
@@ -136,17 +132,8 @@ 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"
-            )
             
         
         seq = esmtokenizer([protein_sequence], add_special_tokens=True, truncation=True, max_length=1021, padding='max_length', return_tensors="pt")