test beam search

Author: eaidova

optimum/intel/openvino/modeling_decoder.py

@@ -559,9 +559,9 @@ def _expand_outputs_for_generation(self, indicies, logits: torch.Tensor, past_ke
         if indicies.shape[0] != 1:
             logits = logits[indicies]
             if past_key_values and not self.stateful:
-                if (self.config.model_type not in MULTI_QUERY_ATTN_MODELS
-                    or (self.config.model_type == "falcon"
-                    and self.config.new_decoder_architecture)):
+                if self.config.model_type not in MULTI_QUERY_ATTN_MODELS or (
+                    self.config.model_type == "falcon" and self.config.new_decoder_architecture
+                ):
                     past_key_values = tuple(
                         tuple(
                             past_state[indicies]

tests/openvino/test_modeling.py

@@ -787,8 +787,8 @@ def test_default_filling_attention_mask_and_position_ids(self):
         gc.collect()
 
     @parameterized.expand(SUPPORTED_ARCHITECTURES)
-    @pytest.mark.run_slow
-    @slow
+    # @pytest.mark.run_slow
+    # @slow
     def test_beam_search(self, model_arch):
         model_kwargs = {}
         model_id = MODEL_NAMES[model_arch]