Update modeling_prot2text.py

modeling_prot2text.py

@@ -240,7 +240,15 @@ class Prot2TextModel(PreTrainedModel):
                                     x: Optional[torch.FloatTensor] = None,
                                     edge_type: Optional[torch.LongTensor] = None,
                                     tokenizer=None,
-                                    device='cpu'
+                                    device='cpu',
+                                     streamer=None, 
+                                     max_new_tokens=None, 
+                                     do_sample=None, 
+                                     top_p=None,
+                                     top_k=None, 
+                                     temperature=None, 
+                                     num_beams=1, 
+                                     repetition_penalty=None
                                      ):
         
         if self.config.esm and not self.config.rgcn and protein_sequence==None:
@@ -326,25 +334,45 @@ class Prot2TextModel(PreTrainedModel):
             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)
+            if streamer is None:
+                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)
+                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|>', '')
+                os.remove(structure_filename)
+                os.remove(graph_filename)
+                os.remove(process_filename)        
+                    
+                return generated[0].replace('<|stop_token|>', '').replace('<|graph_token|>', '')
+            else:
+                os.remove(structure_filename)
+                os.remove(graph_filename)
+                os.remove(process_filename) 
+                return self.decoder.generate(input_ids=inputs['decoder_input_ids'],
+                                         encoder_outputs=encoder_state, 
+                                         use_cache=True,
+                                        encoder_attention_mask=inputs['attention_mask'], 
+                                        length_penalty=1.0,
+                                        no_repeat_ngram_size=None,
+                                        early_stopping=False,
+                                         streamer=streamer, 
+                                         max_new_tokens=max_new_tokens, 
+                                         do_sample=do_sample, 
+                                         top_p=top_p,
+                                         top_k=top_k, 
+                                         temperature=temperature, 
+                                         num_beams=1, 
+                                         repetition_penalty=repetition_penalty)
             
         else:
             seq = esmtokenizer([protein_sequence], add_special_tokens=True, truncation=True, max_length=1021, padding='max_length', return_tensors="pt")