rebase

neuralmagic · Sep 21, 2023 · 3a18e5f · 3a18e5f
1 parent 592c81c
commit 3a18e5f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/src/deepsparse/transformers/pipelines/text_generation.py b/src/deepsparse/transformers/pipelines/text_generation.py
@@ -538,7 +538,9 @@ def process_engine_outputs(
         if streaming:
             return self._stream_engine_outputs(engine_outputs, prompts, kwargs)
 
-        generated_tokens, generated_logits, finished_reason, *debug = list(*engine_outputs)
+        generated_tokens, generated_logits, finished_reason, *debug = list(
+            *engine_outputs
+        )
         sequences = self.tokenizer.batch_decode(
             generated_tokens, skip_special_tokens=True
         )
@@ -587,7 +589,7 @@ def process_engine_outputs(
             )
             outputs.update(debug_params)
 
-        yield TextGenerationOutput(**outputs)
+        return TextGenerationOutput(**outputs)
 
     def engine_forward(
         self, engine_inputs: List[numpy.ndarray], context: Dict
@@ -699,7 +701,7 @@ def engine_forward(
                         [finished_reason[-1]],
                     )
 
-        if not self.streaming:
+        if not streaming:
             returns = (
                 numpy.array([generated_tokens]),
                 numpy.concatenate(generated_logits, axis=1),
@@ -953,12 +955,20 @@ def join_engine_outputs(
 
             if debug:
                 sessions = debug[0]
-                kv_cache_state = numpy.stack(session.cached_inputs for session in sessions)
+                kv_cache_state = numpy.stack(
+                    session.cached_inputs for session in sessions
+                )
                 num_processed_tokens = numpy.stack(
                     session.total_num_processed_tokens for session in sessions
                 )
 
-                yield [tokens, logits, finish_reason, kv_cache_state, num_processed_tokens]
+                yield [
+                    tokens,
+                    logits,
+                    finish_reason,
+                    kv_cache_state,
+                    num_processed_tokens,
+                ]
 
             yield [tokens, logits, finish_reason]
 

diff --git a/tests/deepsparse/transformers/pipelines/test_text_generation.py b/tests/deepsparse/transformers/pipelines/test_text_generation.py
@@ -432,8 +432,6 @@ def _test_output(
         run_cache_validation: bool = True,
     ):
         # extract numpy arrays from cached_inputs
-        kv_cache_array = list(cache_session.cached_inputs.values())
-
         (
             generated_logits,
             prompt_logits,