Update generate.py

generate.py

@@ -25,6 +25,7 @@ def custom_generate(
     bos_token_id=None,
     pad_token_id=None,
     eos_token_id=None,
+    streamer=None,
     length_penalty=None,
     no_repeat_ngram_size=None,
     num_return_sequences=None,
@@ -51,6 +52,7 @@ def custom_generate(
             new_ids = self(
                 input_ids[~finished_generating],
                 attention_mask=attention_mask[~finished_generating] if attention_mask is not None else None,
+                streamer=streamer,
                 **kwargs
             )['logits']
 
@@ -100,6 +102,7 @@ def generate(
     early_stopping=None,
     num_beams=None,
     temperature=1.0,
+    streamer=None,
     top_k=None,
     top_p=None,
     repetition_penalty=None,
@@ -202,6 +205,7 @@ def generate(
         forced_eos_token_id=forced_eos_token_id,
         remove_invalid_values=remove_invalid_values,
         synced_gpus=synced_gpus,
+        streamer=streamer,
         **model_kwargs,
     )