add stream generation following huggingface

deepjavalibrary · May 30, 2023 · cbbc99d · cbbc99d
1 parent d59fcb0
commit cbbc99d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/streaming_utils.py b/engines/python/setup/djl_python/streaming_utils.py
@@ -59,6 +59,9 @@ def _hf_model_stream_generator(model, tokenizer, inputs, **kwargs):
         if "engine" in kwargs.keys():
             engine = kwargs["engine"]
 
+        if engine and "transformers-neuronx" == engine:
+            model.reset_generation()
+
         if generic_model_class == "CausalLM":
             input_length = input_ids.shape[1]
             all_decoder_input_ids = tokenized_inputs["input_ids"]
@@ -93,9 +96,9 @@ def _hf_model_stream_generator(model, tokenizer, inputs, **kwargs):
                     "use_cache": True
                 }
                 if engine and "transformers-neuronx" == engine:
-                    model_inputs["return_dict"] = True
                     model_inputs = model.prepare_inputs_for_generation(
                         **model_inputs)
+                    model_inputs["return_dict"] = True
                 outputs = model.forward(**model_inputs)
 
             if generic_model_class == "Seq2SeqLM":

diff --git a/engines/python/setup/djl_python/transformers-neuronx.py b/engines/python/setup/djl_python/transformers-neuronx.py
@@ -189,7 +189,7 @@ def infer(self, inputs):
                     "transformers-neuronx")
                 model_kwargs["engine"] = "transformers-neuronx"
                 outputs.add_stream_content(
-                    stream_generator(self.model.model, self.tokenizer,
+                    stream_generator(self.model, self.tokenizer,
                                      input_text, **model_kwargs))
                 return outputs