fix breaking tests

neuralmagic · Jul 27, 2023 · e599583 · e599583
1 parent 3987e1b
commit e599583
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/src/deepsparse/transformers/pipelines/text_generation.py b/src/deepsparse/transformers/pipelines/text_generation.py
@@ -393,7 +393,7 @@ def prompt_inference(
         if len(tokens) > self.prompt_processing_sequence_length:
             for engine_inputs in self.engine_inputs_for_prefill(tokens):
                 new_token, new_logits = self.multitoken_engine(engine_inputs)
-                num_tokens_processed = self.prompt_processing_sequence_length
+                num_tokens_processed += self.prompt_processing_sequence_length
                 prompt_logits.append(new_logits)
 
         if num_tokens_processed:
@@ -523,11 +523,16 @@ def engine_inputs_for_prefill(
                     # delay creation of the causal mask
                     continue
                 elif name == "positions":
-                    engine_input = (
-                        numpy.arange(self.prompt_processing_sequence_length)
-                        .reshape(1, -1)
-                        .astype(numpy.int64)
-                    )
+                    if self.prompt_processing_sequence_length == 1:
+                        # we need to treat `positions` as if we were in
+                        # the autoregressive mode
+                        engine_input = numpy.array([[idx]], dtype=numpy.int64)
+                    else:
+                        engine_input = (
+                            numpy.arange(self.prompt_processing_sequence_length)
+                            .reshape(1, -1)
+                            .astype(numpy.int64)
+                        )
 
                 engine_inputs.append(engine_input)
 

diff --git a/src/deepsparse/transformers/utils/decoder_kv_cache.py b/src/deepsparse/transformers/utils/decoder_kv_cache.py
@@ -204,7 +204,7 @@ def set_capacity(self, capacity: int):
             state = self._add_entries(state, indices=indices)
 
         else:
-            pass
+            return
 
         self._state = state