Fix gpt_bigcode input generator for transformers 4.54 (#2336)

* Fix gpt_bigcode input generator for transformers 4.54

* style

---------

Co-authored-by: IlyasMoutawwakil <moutawwakil.ilyas.tsi@gmail.com>

Author: echarlaix

optimum/utils/input_generators.py

@@ -1128,27 +1128,41 @@ def __init__(
         self.multi_query = normalized_config.multi_query
 
     def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        if self.multi_query:
-            past_key_value_shape = (
-                self.batch_size,
-                self.sequence_length,
-                self.hidden_size // self.num_attention_heads * 2,
-            )
-            return [
-                self.random_float_tensor(past_key_value_shape, framework=framework, dtype=float_dtype)
-                for _ in range(self.num_layers)
+        if is_transformers_version("<", "4.54"):
+            if self.multi_query:
+                shape = (
+                    self.batch_size,
+                    self.sequence_length,
+                    self.hidden_size // self.num_attention_heads * 2,
+                )
+            else:
+                shape = (
+                    self.batch_size,
+                    self.num_attention_heads,
+                    self.sequence_length,
+                    self.hidden_size // self.num_attention_heads * 2,
+                )
+            pkv = [
+                self.random_float_tensor(shape, framework=framework, dtype=float_dtype) for _ in range(self.num_layers)
             ]
+
         else:
             shape = (
                 self.batch_size,
-                self.num_attention_heads,
+                self.num_attention_heads if not self.multi_query else 1,
                 self.sequence_length,
-                self.hidden_size // self.num_attention_heads * 2,
+                self.hidden_size // self.num_attention_heads,
             )
-            return [
-                self.random_float_tensor(shape, framework=framework, dtype=float_dtype) for _ in range(self.num_layers)
+            pkv = [
+                (
+                    self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
+                    self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
+                )
+                for _ in range(self.num_layers)
             ]
 
+        return pkv
+
 
 class BloomDummyPastKeyValuesGenerator(DummyPastKeyValuesGenerator):
     def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
@@ -1278,30 +1292,24 @@ def __init__(
             random_sequence_length_range=random_sequence_length_range,
             **kwargs,
         )
-        self.num_kv_heads = self.num_kv_heads = (
+        self.num_kv_heads = (
             normalized_config.num_kv_heads
             if (normalized_config.new_decoder_architecture or not normalized_config.multi_query)
             else 1
         )
         self.head_dim = self.hidden_size // self.num_attention_heads
 
     def generate(self, input_name: str, framework: str = "pt", int_dtype: str = "int64", float_dtype: str = "fp32"):
-        past_key_shape = (
-            self.batch_size,
-            self.num_kv_heads,
-            self.sequence_length,
-            self.head_dim,
-        )
-        past_value_shape = (
+        shape = (
             self.batch_size,
             self.num_kv_heads,
             self.sequence_length,
             self.head_dim,
         )
         return [
             (
-                self.random_float_tensor(past_key_shape, framework=framework, dtype=float_dtype),
-                self.random_float_tensor(past_value_shape, framework=framework, dtype=float_dtype),
+                self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
+                self.random_float_tensor(shape, framework=framework, dtype=float_dtype),
             )
             for _ in range(self.num_layers)
         ]