




import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
		import jax
		import jax.numpy as jnp

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
		import torch




class    snake_case__							:
				"""simple docstring"""



				def __init__(				self    :					int ,				__lowerCamelCase    :					Optional[Any] ,				__lowerCamelCase    :					Optional[int]=14 ,				__lowerCamelCase    :					Any=7 ,				__lowerCamelCase    :					List[str]=True ,				__lowerCamelCase    :					str=True ,				__lowerCamelCase    :					Optional[Any]=False ,				__lowerCamelCase    :					int=True ,				__lowerCamelCase    :					str=99 ,				__lowerCamelCase    :					str=32 ,				__lowerCamelCase    :					Tuple=4 ,				__lowerCamelCase    :					Any=4 ,				__lowerCamelCase    :					int=[[["global", "local"], 2]] ,				__lowerCamelCase    :					str=37 ,				__lowerCamelCase    :					int="gelu" ,				__lowerCamelCase    :					Tuple=0.1 ,				__lowerCamelCase    :					int=0.1 ,				__lowerCamelCase    :					Union[str, Any]=5_12 ,				__lowerCamelCase    :					Any=7 ,				__lowerCamelCase    :					Any=0.02 ,				) ->					Optional[Any]:
							a   =      parent
							a   =      batch_size
							a   =      seq_length
							a   =      is_training
							a   =      use_input_mask
							a   =      use_token_type_ids
							a   =      use_labels
							a   =      vocab_size
							a   =      hidden_size
							a   =      num_hidden_layers
							a   =      num_attention_heads
							a   =      attention_types
							a   =      intermediate_size
							a   =      hidden_act
							a   =      hidden_dropout_prob
							a   =      attention_probs_dropout_prob
							a   =      max_position_embeddings
							a   =      window_size
							a   =      initializer_range
							a   =      None
							a   =      vocab_size - 1
							a   =      vocab_size - 1
							a   =      vocab_size - 1





				def 							__UpperCAmelCase							(				self    :					Any							) ->					Union[str, Any]:
							a   =      ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size							)

							a   =      None
							if self.use_input_mask:
										a   =      random_attention_mask([self.batch_size, self.seq_length]							)

							a   =      GPTNeoConfig(
							    vocab_size=self.vocab_size ,				hidden_size=self.hidden_size ,				num_layers=self.num_hidden_layers ,				num_heads=self.num_attention_heads ,				max_position_embeddings=self.max_position_embeddings ,				use_cache=__lowerCamelCase ,				bos_token_id=self.bos_token_id ,				eos_token_id=self.eos_token_id ,				pad_token_id=self.pad_token_id ,				window_size=self.window_size ,				attention_types=self.attention_types ,				)

							return (config, input_ids, input_mask)





				def 							__UpperCAmelCase							(				self    :					Dict							) ->					Union[str, Any]:
							a   =      self.prepare_config_and_inputs()
							a   ,		a   ,		a   =      config_and_inputs
							a   =      {"input_ids": input_ids, "attention_mask": attention_mask}
							return config, inputs_dict





				def 							__UpperCAmelCase							(				self    :					Union[str, Any] ,				__lowerCamelCase    :					Union[str, Any] ,				__lowerCamelCase    :					str ,				__lowerCamelCase    :					int ,				__lowerCamelCase    :					Any							) ->					str:
							a   =      20
							a   =      model_class_name(__lowerCamelCase							)

							a   =      model.init_cache(input_ids.shape[0] ,				__lowerCamelCase							)
							a   =      jnp.ones((input_ids.shape[0], max_decoder_length) ,				dtype="i4"							)

							a   =      jnp.broadcast_to(
							    jnp.arange(input_ids.shape[-1] - 1							)[None, :] ,				(input_ids.shape[0], input_ids.shape[-1] - 1)							)
							a   =      model(
							    input_ids[:, :-1] ,				attention_mask=__lowerCamelCase ,				past_key_values=__lowerCamelCase ,				position_ids=__lowerCamelCase ,				)

							a   =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,				dtype="i4"							)
							a   =      model(
							    input_ids[:, -1:] ,				attention_mask=__lowerCamelCase ,				past_key_values=outputs_cache.past_key_values ,				position_ids=__lowerCamelCase ,				)

							a   =      model(__lowerCamelCase							)

							a   =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
							self.parent.assertTrue(diff < 1e-3 ,				msg=f"""Max diff is {diff}"""							)





				def 							__UpperCAmelCase							(				self    :					List[str] ,				__lowerCamelCase    :					Optional[int] ,				__lowerCamelCase    :					int ,				__lowerCamelCase    :					Tuple ,				__lowerCamelCase    :					Tuple							) ->					List[str]:
							a   =      20
							a   =      model_class_name(__lowerCamelCase							)

							a   =      jnp.concatenate(
							    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])							)] ,				axis=-1 ,				)

							a   =      model.init_cache(input_ids.shape[0] ,				__lowerCamelCase							)
							a   =      jnp.broadcast_to(
							    jnp.arange(input_ids.shape[-1] - 1							)[None, :] ,				(input_ids.shape[0], input_ids.shape[-1] - 1)							)

							a   =      model(
							    input_ids[:, :-1] ,				attention_mask=__lowerCamelCase ,				past_key_values=__lowerCamelCase ,				position_ids=__lowerCamelCase ,				)
							a   =      jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]] ,				dtype="i4"							)
							a   =      model(
							    input_ids[:, -1:] ,				past_key_values=outputs_cache.past_key_values ,				attention_mask=__lowerCamelCase ,				position_ids=__lowerCamelCase ,				)

							a   =      model(__lowerCamelCase ,				attention_mask=__lowerCamelCase							)

							a   =      np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])							)							)
							self.parent.assertTrue(diff < 1e-3 ,				msg=f"""Max diff is {diff}"""							)




@require_flax
class    snake_case__							(_UpperCamelCase   ,	_UpperCamelCase   ,	unittest.TestCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							List[str]										= (FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
				SCREAMING_SNAKE_CASE_	:							Tuple										= (FlaxGPTNeoForCausalLM,) if is_flax_available() else ()





				def 							__UpperCAmelCase							(				self    :					List[Any]							) ->					Tuple:
							a   =      FlaxGPTNeoModelTester(self							)





				def 							__UpperCAmelCase							(				self    :					Union[str, Any]							) ->					Any:
							for model_class_name in self.all_model_classes:
										a   ,		a   ,		a   =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward(__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)





				def 							__UpperCAmelCase							(				self    :					List[str]							) ->					Optional[Any]:
							for model_class_name in self.all_model_classes:
										a   ,		a   ,		a   =      self.model_tester.prepare_config_and_inputs()
										self.model_tester.check_use_cache_forward_with_attn_mask(
										    __lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)





				@slow
				def 							__UpperCAmelCase							(				self    :					str							) ->					Union[str, Any]:
							a   =      GPTaTokenizer.from_pretrained("gpt2" ,				pad_token="<|endoftext|>" ,				padding_side="left"							)
							a   =      tokenizer(["Hello this is a long string", "Hey"] ,				return_tensors="np" ,				padding=__lowerCamelCase ,				truncation=__lowerCamelCase							)

							a   =      FlaxGPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M"							)
							a   =      False
							a   =      model.config.eos_token_id

							a   =      jax.jit(model.generate							)

							a   =      jit_generate(
							    inputs["input_ids"] ,				attention_mask=inputs["attention_mask"] ,				pad_token_id=tokenizer.pad_token_id							).sequences

							a   =      tokenizer.batch_decode(__lowerCamelCase ,				skip_special_tokens=__lowerCamelCase							)

							a   =      [
							    "Hello this is a long string of text.\n\nI'm trying to get the text of the",
							    "Hey, I'm a little late to the party. I'm going to",
							]

							self.assertListEqual(__lowerCamelCase ,				__lowerCamelCase							)





				@is_pt_flax_cross_test
				def 							__UpperCAmelCase							(				self    :					Dict							) ->					Optional[Any]:
							a   ,		a   =      self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__							):
													# prepare inputs
													a   =      self._prepare_for_class(__lowerCamelCase ,				__lowerCamelCase							)
													a   =      {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

													# load corresponding PyTorch class
													a   =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
													a   =      getattr(__lowerCamelCase ,				__lowerCamelCase							)

													a   ,		a   =      pt_inputs["input_ids"].shape
													a   =      np.random.randint(0 ,				seq_length - 1 ,				size=(batch_size,)							)
													for batch_idx, start_index in enumerate(__lowerCamelCase							):
																a   =      0
																a   =      1
																a   =      0
																a   =      1
													a   =      pt_model_class(__lowerCamelCase							).eval()
													a   =      model_class(__lowerCamelCase ,				dtype=jnp.floataa							)

													a   =      convert_pytorch_state_dict_to_flax(pt_model.state_dict() ,				__lowerCamelCase							)
													a   =      fx_state

													with torch.no_grad():
																a   =      pt_model(**__lowerCamelCase							).to_tuple()

													a   =      fx_model(**__lowerCamelCase							).to_tuple()
													self.assertEqual(len(__lowerCamelCase							) ,				len(__lowerCamelCase							) ,				"Output lengths differ between Flax and PyTorch"							)
													for fx_output, pt_output in zip(__lowerCamelCase ,				__lowerCamelCase							):
																self.assert_almost_equals(fx_output[:, -1] ,				pt_output[:, -1].numpy() ,				4e-2							)

													with tempfile.TemporaryDirectory() as tmpdirname:
																pt_model.save_pretrained(__lowerCamelCase							)
																a   =      model_class.from_pretrained(__lowerCamelCase ,				from_pt=__lowerCamelCase							)

													a   =      fx_model_loaded(**__lowerCamelCase							).to_tuple()
													self.assertEqual(
													    len(__lowerCamelCase							) ,				len(__lowerCamelCase							) ,				"Output lengths differ between Flax and PyTorch"							)
													for fx_output_loaded, pt_output in zip(__lowerCamelCase ,				__lowerCamelCase							):
																self.assert_almost_equals(fx_output_loaded[:, -1] ,				pt_output[:, -1].numpy() ,				4e-2							)





				@is_pt_flax_cross_test
				def 							__UpperCAmelCase							(				self    :					int							) ->					Any:
							a   ,		a   =      self.model_tester.prepare_config_and_inputs_for_common()
							for model_class in self.all_model_classes:
										with self.subTest(model_class.__name__							):
													# prepare inputs
													a   =      self._prepare_for_class(__lowerCamelCase ,				__lowerCamelCase							)
													a   =      {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

													# load corresponding PyTorch class
													a   =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
													a   =      getattr(__lowerCamelCase ,				__lowerCamelCase							)

													a   =      pt_model_class(__lowerCamelCase							).eval()
													a   =      model_class(__lowerCamelCase ,				dtype=jnp.floataa							)

													a   =      load_flax_weights_in_pytorch_model(__lowerCamelCase ,				fx_model.params							)
													a   ,		a   =      pt_inputs["input_ids"].shape
													a   =      np.random.randint(0 ,				seq_length - 1 ,				size=(batch_size,)							)
													for batch_idx, start_index in enumerate(__lowerCamelCase							):
																a   =      0
																a   =      1
																a   =      0
																a   =      1

													# make sure weights are tied in PyTorch
													pt_model.tie_weights()

													with torch.no_grad():
																a   =      pt_model(**__lowerCamelCase							).to_tuple()

													a   =      fx_model(**__lowerCamelCase							).to_tuple()
													self.assertEqual(len(__lowerCamelCase							) ,				len(__lowerCamelCase							) ,				"Output lengths differ between Flax and PyTorch"							)
													for fx_output, pt_output in zip(__lowerCamelCase ,				__lowerCamelCase							):
																self.assert_almost_equals(fx_output[:, -1] ,				pt_output[:, -1].numpy() ,				4e-2							)

													with tempfile.TemporaryDirectory() as tmpdirname:
																fx_model.save_pretrained(__lowerCamelCase							)
																a   =      pt_model_class.from_pretrained(__lowerCamelCase ,				from_flax=__lowerCamelCase							)

													with torch.no_grad():
																a   =      pt_model_loaded(**__lowerCamelCase							).to_tuple()

													self.assertEqual(
													    len(__lowerCamelCase							) ,				len(__lowerCamelCase							) ,				"Output lengths differ between Flax and PyTorch"							)
													for fx_output, pt_output in zip(__lowerCamelCase ,				__lowerCamelCase							):
																self.assert_almost_equals(fx_output[:, -1] ,				pt_output[:, -1].numpy() ,				4e-2							)





				@slow
				def 							__UpperCAmelCase							(				self    :					Tuple							) ->					str:
							for model_class_name in self.all_model_classes:
										a   =      model_class_name.from_pretrained("EleutherAI/gpt-neo-125M"							)
										a   =      model(np.ones((1, 1)							)							)
										self.assertIsNotNone(__lowerCamelCase							)












from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__lowerCAmelCase		:    Dict	      =		logging.get_logger(__name__)

__lowerCAmelCase		:    Optional[int]	      =		{
    'facebook/mbart-large-cc25': 'https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json',
    # See all MBART models at https://huggingface.co/models?filter=mbart
}




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_	:							int										= """mbart"""
				SCREAMING_SNAKE_CASE_	:							Optional[Any]										= ["""past_key_values"""]
				SCREAMING_SNAKE_CASE_	:							Any										= {"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}





				def __init__(				self    :					List[Any] ,				__lowerCamelCase    :					str=5_02_65 ,				__lowerCamelCase    :					List[Any]=10_24 ,				__lowerCamelCase    :					Optional[Any]=12 ,				__lowerCamelCase    :					List[Any]=40_96 ,				__lowerCamelCase    :					int=16 ,				__lowerCamelCase    :					List[Any]=12 ,				__lowerCamelCase    :					Tuple=40_96 ,				__lowerCamelCase    :					Tuple=16 ,				__lowerCamelCase    :					Optional[Any]=0.0 ,				__lowerCamelCase    :					int=0.0 ,				__lowerCamelCase    :					Any=True ,				__lowerCamelCase    :					List[Any]=True ,				__lowerCamelCase    :					str="gelu" ,				__lowerCamelCase    :					Tuple=10_24 ,				__lowerCamelCase    :					Any=0.1 ,				__lowerCamelCase    :					Dict=0.0 ,				__lowerCamelCase    :					Tuple=0.0 ,				__lowerCamelCase    :					Optional[int]=0.02 ,				__lowerCamelCase    :					Dict=0.0 ,				__lowerCamelCase    :					Union[str, Any]=False ,				__lowerCamelCase    :					Optional[int]=1 ,				__lowerCamelCase    :					Tuple=0 ,				__lowerCamelCase    :					List[Any]=2 ,				__lowerCamelCase    :					Dict=2 ,				**__lowerCamelCase    :					Optional[int] ,				) ->					Optional[int]:
							a   =      vocab_size
							a   =      max_position_embeddings
							a   =      d_model
							a   =      encoder_ffn_dim
							a   =      encoder_layers
							a   =      encoder_attention_heads
							a   =      decoder_ffn_dim
							a   =      decoder_layers
							a   =      decoder_attention_heads
							a   =      dropout
							a   =      attention_dropout
							a   =      activation_dropout
							a   =      activation_function
							a   =      init_std
							a   =      encoder_layerdrop
							a   =      decoder_layerdrop
							a   =      classifier_dropout
							a   =      use_cache
							a   =      encoder_layers
							a   =      scale_embedding  # scale factor will be sqrt(d_model) if True
							super().__init__(
							    pad_token_id=__lowerCamelCase ,				bos_token_id=__lowerCamelCase ,				eos_token_id=__lowerCamelCase ,				is_encoder_decoder=__lowerCamelCase ,				forced_eos_token_id=__lowerCamelCase ,				**__lowerCamelCase ,				)




class    snake_case__							(_UpperCamelCase       ):
				"""simple docstring"""



				@property
				def 							__UpperCAmelCase							(				self    :					Optional[Any]							) ->					Mapping[str, Mapping[int, str]]:
							if self.task in ["default", "seq2seq-lm"]:
										a   =      OrderedDict(
										    [
										        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
										        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
										    ]							)

										if self.use_past:
													a   =      {0: "batch"}
													a   =      {0: "batch", 1: "past_decoder_sequence + sequence"}
										else:
													a   =      {0: "batch", 1: "decoder_sequence"}
													a   =      {0: "batch", 1: "decoder_sequence"}

										if self.use_past:
													self.fill_with_past_key_values_(__lowerCamelCase ,				direction="inputs"							)
							elif self.task == "causal-lm":
										# TODO: figure this case out.
										a   =      OrderedDict(
										    [
										        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
										        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
										    ]							)
										if self.use_past:
													a   ,		a   =      self.num_layers
													for i in range(__lowerCamelCase							):
																a   =      {0: "batch", 2: "past_sequence + sequence"}
																a   =      {0: "batch", 2: "past_sequence + sequence"}
							else:
										a   =      OrderedDict(
										    [
										        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
										        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
										        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
										        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
										    ]							)

							return common_inputs





				@property
				def 							__UpperCAmelCase							(				self    :					Any							) ->					Mapping[str, Mapping[int, str]]:
							if self.task in ["default", "seq2seq-lm"]:
										a   =      super().outputs
							else:
										a   =      super(__lowerCamelCase ,				self							).outputs
										if self.use_past:
													a   ,		a   =      self.num_layers
													for i in range(__lowerCamelCase							):
																a   =      {0: "batch", 2: "past_sequence + sequence"}
																a   =      {0: "batch", 2: "past_sequence + sequence"}
							return common_outputs





				def 							__UpperCAmelCase							(				self    :					Union[str, Any] ,				__lowerCamelCase    :					PreTrainedTokenizer ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					bool = False ,				__lowerCamelCase    :					Optional[TensorType] = None ,				) ->					Mapping[str, Any]:
							a   =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)

							# Generate decoder inputs
							a   =      seq_length if not self.use_past else 1
							a   =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)
							a   =      {f"""decoder_{name}""": tensor for name, tensor in decoder_inputs.items()}
							a   =      dict(**__lowerCamelCase ,				**__lowerCamelCase							)

							if self.use_past:
										if not is_torch_available():
													raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."							)
										else:
													import torch
										a   ,		a   =      common_inputs["input_ids"].shape
										a   =      common_inputs["decoder_input_ids"].shape[1]
										a   ,		a   =      self.num_attention_heads
										a   =      (
										    batch,
										    num_encoder_attention_heads,
										    encoder_seq_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)
										a   =      decoder_seq_length + 3
										a   =      (
										    batch,
										    num_decoder_attention_heads,
										    decoder_past_length,
										    self._config.hidden_size // num_decoder_attention_heads,
										)

										a   =      torch.cat(
										    [common_inputs["decoder_attention_mask"], torch.ones(__lowerCamelCase ,				__lowerCamelCase							)] ,				dim=1							)

										a   =      []
										# If the number of encoder and decoder layers are present in the model configuration, both are considered
										a   ,		a   =      self.num_layers
										a   =      min(__lowerCamelCase ,				__lowerCamelCase							)
										a   =      max(__lowerCamelCase ,				__lowerCamelCase							) - min_num_layers
										a   =      "encoder" if num_encoder_layers > num_decoder_layers else "decoder"

										for _ in range(__lowerCamelCase							):
													common_inputs["past_key_values"].append(
													    (
													        torch.zeros(__lowerCamelCase							),
													        torch.zeros(__lowerCamelCase							),
													        torch.zeros(__lowerCamelCase							),
													        torch.zeros(__lowerCamelCase							),
													    )							)
										# TODO: test this.
										a   =      encoder_shape if remaining_side_name == "encoder" else decoder_shape
										for _ in range(__lowerCamelCase ,				__lowerCamelCase							):
													common_inputs["past_key_values"].append((torch.zeros(__lowerCamelCase							), torch.zeros(__lowerCamelCase							))							)
							return common_inputs





				def 							__UpperCAmelCase							(				self    :					str ,				__lowerCamelCase    :					PreTrainedTokenizer ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					bool = False ,				__lowerCamelCase    :					Optional[TensorType] = None ,				) ->					Mapping[str, Any]:
							a   =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
							    __lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)

							if self.use_past:
										if not is_torch_available():
													raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."							)
										else:
													import torch
										a   ,		a   =      common_inputs["input_ids"].shape
										# Not using the same length for past_key_values
										a   =      seqlen + 2
										a   ,		a   =      self.num_layers
										a   ,		a   =      self.num_attention_heads
										a   =      (
										    batch,
										    num_encoder_attention_heads,
										    past_key_values_length,
										    self._config.hidden_size // num_encoder_attention_heads,
										)

										a   =      common_inputs["attention_mask"].dtype
										a   =      torch.cat(
										    [common_inputs["attention_mask"], torch.ones(__lowerCamelCase ,				__lowerCamelCase ,				dtype=__lowerCamelCase							)] ,				dim=1							)
										a   =      [
										    (torch.zeros(__lowerCamelCase							), torch.zeros(__lowerCamelCase							)) for _ in range(__lowerCamelCase							)
										]
							return common_inputs





				def 							__UpperCAmelCase							(				self    :					Union[str, Any] ,				__lowerCamelCase    :					PreTrainedTokenizer ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					bool = False ,				__lowerCamelCase    :					Optional[TensorType] = None ,				) ->					Mapping[str, Any]:
							# Copied from OnnxConfig.generate_dummy_inputs
							# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
							# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
							a   =      compute_effective_axis_dimension(
							    __lowerCamelCase ,				fixed_dimension=OnnxConfig.default_fixed_batch ,				num_token_to_add=0							)

							# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
							a   =      tokenizer.num_special_tokens_to_add(__lowerCamelCase							)
							a   =      compute_effective_axis_dimension(
							    __lowerCamelCase ,				fixed_dimension=OnnxConfig.default_fixed_sequence ,				num_token_to_add=__lowerCamelCase							)

							# Generate dummy inputs according to compute batch and sequence
							a   =      [" ".join([tokenizer.unk_token]							) * seq_length] * batch_size
							a   =      dict(tokenizer(__lowerCamelCase ,				return_tensors=__lowerCamelCase							)							)
							return common_inputs





				def 							__UpperCAmelCase							(				self    :					str ,				__lowerCamelCase    :					PreTrainedTokenizer ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					int = -1 ,				__lowerCamelCase    :					bool = False ,				__lowerCamelCase    :					Optional[TensorType] = None ,				) ->					Mapping[str, Any]:
							if self.task in ["default", "seq2seq-lm"]:
										a   =      self._generate_dummy_inputs_for_default_and_seqaseq_lm(
										    __lowerCamelCase ,				batch_size=__lowerCamelCase ,				seq_length=__lowerCamelCase ,				is_pair=__lowerCamelCase ,				framework=__lowerCamelCase							)

							elif self.task == "causal-lm":
										a   =      self._generate_dummy_inputs_for_causal_lm(
										    __lowerCamelCase ,				batch_size=__lowerCamelCase ,				seq_length=__lowerCamelCase ,				is_pair=__lowerCamelCase ,				framework=__lowerCamelCase							)
							else:
										a   =      self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    __lowerCamelCase ,				batch_size=__lowerCamelCase ,				seq_length=__lowerCamelCase ,				is_pair=__lowerCamelCase ,				framework=__lowerCamelCase							)

							return common_inputs





				def 							__UpperCAmelCase							(				self    :					Optional[int] ,				__lowerCamelCase    :					Union[str, Any] ,				__lowerCamelCase    :					int ,				__lowerCamelCase    :					str ,				__lowerCamelCase    :					Optional[Any]							) ->					Dict:
							if self.task in ["default", "seq2seq-lm"]:
										a   =      super()._flatten_past_key_values_(__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)
							else:
										a   =      super(__lowerCamelCase ,				self							)._flatten_past_key_values_(
										    __lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase ,				__lowerCamelCase							)






