simplify after PR review round

Author: dbogunowicz

src/deepsparse/v2/text_generation/generate_new_token.py

@@ -36,9 +36,9 @@ def can_operate(self, inp: NLEngineOutputs):
             return True
         return False
 
-    def run(self, *args, inference_state: InferenceState, **kwargs):
-        logits = args[0].engine_outputs if args else kwargs.get("logits")
-        kv_cache = args[0].kv_cache if args else kwargs.get("kv_cache")
+    def run(self, inp: NLEngineOutputs, inference_state: InferenceState, **kwargs):
+        logits = inp.engine_outputs
+        kv_cache = inp.kv_cache
 
         token_generator = inference_state.current_state.get("token_generator")
         token = token_generator.generate(logits=logits[0, -1, :])

src/deepsparse/v2/text_generation/prep_for_generation.py

@@ -20,6 +20,7 @@
 from deepsparse.transformers.utils.helpers import set_generated_length
 from deepsparse.v2.operators import Operator
 from deepsparse.v2.text_generation import TokenGeneratorOperator
+from deepsparse.v2.text_generation.nl_engine_operator import NLEngineOutputs
 from deepsparse.v2.utils import InferenceState
 
 
@@ -41,10 +42,11 @@ def can_operate(self, inp: Any):
         kv_cache = inp.get("kv_cache")
         tokens = inp.get("tokens")
 
-        # If the number of prompt tokens is greater than what we've processed,
-        # don't start generation. Should be equal when started as all prompt logits
-        # should be accounted for and we should have updated the kv_cache for the single
-        # token engine.
+        # If the number of prompt tokens is greater
+        # than what we've processed, don't start generation.
+        # Should be equal when started as all prompt logits
+        # should be accounted for, and we should have updated
+        # the kv_cache for the single token engine.
         if len(tokens) == kv_cache.total_num_processed_tokens:
             return True
         return False
@@ -90,10 +92,13 @@ def run(
             "finished_reason": [],
             "token_generator": token_generator,
         }
+
         output = {
-            "logits": prompt_logits,
             "tokens": token_generator.tokens,
             "kv_cache": kv_cache,
             "in_generation": True,
         }
+        if kv_cache is None:
+            output = NLEngineOutputs(**output, engine_outputs=prompt_logits)
+
         return output, state_update

tests/deepsparse/v2/unit/text_generation/test_token_generation.py

@@ -93,9 +93,7 @@ def test_generate_new_token(
         in_generation=True,
     )
     outputs, state = generate_new_token.run(
-        logits=inp.engine_outputs,
-        kv_cache=inp.kv_cache,
-        inference_state=mock_inference_state,
+        inp=inp, inference_state=mock_inference_state
     )
     # The new_token generated/returned by ths operator should match the last token in
     # token_generator