Update modeling_prot2text.py

modeling_prot2text.py

@@ -123,9 +123,17 @@ class Prot2TextModel(PreTrainedModel):
     
     @torch.no_grad()    
     def generate_protein_description(self,
-                                    protein_sequence=None,
-                                    tokenizer=None,
-                                    device='cpu'
+                                     protein_sequence=None,
+                                     tokenizer=None,
+                                     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:
@@ -147,9 +155,23 @@ class Prot2TextModel(PreTrainedModel):
         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|>', '')
+        if streamer is None:
+            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|>', '')
+        else:
+            return self.decoder.generate(input_ids=inputs['decoder_input_ids'],
+                                         encoder_outputs=encoder_state, 
+                                         use_cache=True,
+                                         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)
+                                         
+                                         
     
     @torch.no_grad()
     def generate(self,