

"""simple docstring"""






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import MBartConfig, is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_flax_available():
						import os

						# The slow tests are often failing with OOM error on GPU
						# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
						# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
						lowerCamelCase__			:					Dict							  =					'''platform'''

						import jax
						import jax.numpy as jnp

						from transformers import AutoTokenizer
						from transformers.models.mbart.modeling_flax_mbart import (
						    FlaxMBartForConditionalGeneration,
						    FlaxMBartForQuestionAnswering,
						    FlaxMBartForSequenceClassification,
						    FlaxMBartModel,
						    shift_tokens_right,
						)
def  UpperCamelCase   (     _lowerCAmelCase       :  Tuple,					_lowerCAmelCase       :  Any,					_lowerCAmelCase       :  Optional[int]=None,					_lowerCAmelCase       :  str=None,					_lowerCAmelCase       :  Tuple=None,					_lowerCAmelCase       :  str=None,					_lowerCAmelCase       :  List[Any]=None,					_lowerCAmelCase       :  Tuple=None,					)			->     Dict:
	if attention_mask is None:
		_UpperCAmelCase      :							Optional[int]       =   np.where(input_ids != config.pad_token_id,					1,					0	)
	if decoder_attention_mask is None:
		_UpperCAmelCase      :							str       =   np.where(decoder_input_ids != config.pad_token_id,					1,					0	)
	if head_mask is None:
		_UpperCAmelCase      :							int       =   np.ones((config.encoder_layers, config.encoder_attention_heads)	)
	if decoder_head_mask is None:
		_UpperCAmelCase      :							Optional[int]       =   np.ones((config.decoder_layers, config.decoder_attention_heads)	)
	if cross_attn_head_mask is None:
		_UpperCAmelCase      :							Union[str, Any]       =   np.ones((config.decoder_layers, config.decoder_attention_heads)	)
	return {
	    "input_ids": input_ids,
	    "decoder_input_ids": decoder_input_ids,
	    "attention_mask": attention_mask,
	    "decoder_attention_mask": decoder_attention_mask,
	}






class       _UpperCAmelCase     (						unittest.TestCase):
			def __init__(						self ,						_A ,						_A=13 ,						_A=7 ,						_A=True ,						_A=False ,						_A=99 ,						_A=16 ,						_A=2 ,						_A=4 ,						_A=4 ,						_A="gelu" ,						_A=0.1 ,						_A=0.1 ,						_A=32 ,						_A=2 ,						_A=1 ,						_A=0 ,						_A=2 ,						_A=0.02 ,						)      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[Any]       =   parent
				_UpperCAmelCase      :							List[str]       =   batch_size
				_UpperCAmelCase      :							Optional[int]       =   seq_length
				_UpperCAmelCase      :							Union[str, Any]       =   is_training
				_UpperCAmelCase      :							Optional[Any]       =   use_labels
				_UpperCAmelCase      :							Optional[Any]       =   vocab_size
				_UpperCAmelCase      :							Tuple       =   hidden_size
				_UpperCAmelCase      :							List[str]       =   num_hidden_layers
				_UpperCAmelCase      :							Dict       =   num_attention_heads
				_UpperCAmelCase      :							int       =   intermediate_size
				_UpperCAmelCase      :							Dict       =   hidden_act
				_UpperCAmelCase      :							str       =   hidden_dropout_prob
				_UpperCAmelCase      :							int       =   attention_probs_dropout_prob
				_UpperCAmelCase      :							Tuple       =   max_position_embeddings
				_UpperCAmelCase      :							str       =   eos_token_id
				_UpperCAmelCase      :							Tuple       =   pad_token_id
				_UpperCAmelCase      :							List[str]       =   bos_token_id
				_UpperCAmelCase      :							Optional[Any]       =   decoder_start_token_id
				_UpperCAmelCase      :							str       =   initializer_range
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				_UpperCAmelCase      :							Tuple       =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1] ,						self.vocab_size       ) ,						3 ,						self.vocab_size       )
				_UpperCAmelCase      :							List[Any]       =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1) ,						dtype=np.intaa       )) ,						-1       )

				_UpperCAmelCase      :							str       =   shift_tokens_right(_A ,						1       )

				_UpperCAmelCase      :							Optional[Any]       =   MBartConfig(
				    vocab_size=self.vocab_size ,						d_model=self.hidden_size ,						encoder_layers=self.num_hidden_layers ,						decoder_layers=self.num_hidden_layers ,						encoder_attention_heads=self.num_attention_heads ,						decoder_attention_heads=self.num_attention_heads ,						encoder_ffn_dim=self.intermediate_size ,						decoder_ffn_dim=self.intermediate_size ,						dropout=self.hidden_dropout_prob ,						attention_dropout=self.attention_probs_dropout_prob ,						max_position_embeddings=self.max_position_embeddings ,						eos_token_id=self.eos_token_id ,						bos_token_id=self.bos_token_id ,						pad_token_id=self.pad_token_id ,						decoder_start_token_id=self.decoder_start_token_id ,						initializer_range=self.initializer_range ,						use_cache=_A ,						)
				_UpperCAmelCase      :							List[Any]       =   prepare_mbart_inputs_dict(_A ,						_A ,						_A       )
				return config, inputs_dict
			def 							__snake_case	(						self       )      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[int]       =   self.prepare_config_and_inputs()
				return config, inputs_dict
			def 							__snake_case	(						self ,						_A ,						_A ,						_A       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   20
				_UpperCAmelCase      :							Any       =   model_class_name(_A       )

				_UpperCAmelCase      :							Tuple       =   model.encode(inputs_dict["""input_ids"""]       )

				_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   (
				    inputs_dict["""decoder_input_ids"""],
				    inputs_dict["""decoder_attention_mask"""],
				)

				_UpperCAmelCase      :							List[Any]       =   model.init_cache(decoder_input_ids.shape[0] ,						_A ,						_A       )
				_UpperCAmelCase      :							Tuple       =   jnp.ones((decoder_input_ids.shape[0], max_decoder_length) ,						dtype="""i4"""       )

				_UpperCAmelCase      :							Optional[int]       =   jnp.broadcast_to(
				    jnp.arange(decoder_input_ids.shape[-1] - 1       )[None, :] ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,						)
				_UpperCAmelCase      :							Tuple       =   model.decode(
				    decoder_input_ids[:, :-1] ,						_A ,						decoder_attention_mask=_A ,						past_key_values=_A ,						decoder_position_ids=_A ,						)

				_UpperCAmelCase      :							List[str]       =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]] ,						dtype="""i4"""       )
				_UpperCAmelCase      :							Optional[Any]       =   model.decode(
				    decoder_input_ids[:, -1:] ,						_A ,						decoder_attention_mask=_A ,						past_key_values=outputs_cache.past_key_values ,						decoder_position_ids=_A ,						)

				_UpperCAmelCase      :							Optional[int]       =   model.decode(_A ,						_A       )

				_UpperCAmelCase      :							Optional[int]       =   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 							__snake_case	(						self ,						_A ,						_A ,						_A       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[int]       =   20
				_UpperCAmelCase      :							Tuple       =   model_class_name(_A       )

				_UpperCAmelCase      :							Optional[Any]       =   model.encode(inputs_dict["""input_ids"""]       )

				_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[int]       =   (
				    inputs_dict["""decoder_input_ids"""],
				    inputs_dict["""decoder_attention_mask"""],
				)

				_UpperCAmelCase      :							Union[str, Any]       =   jnp.concatenate(
				    [
				        decoder_attention_mask,
				        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])       ),
				    ] ,						axis=-1 ,						)

				_UpperCAmelCase      :							Union[str, Any]       =   model.init_cache(decoder_input_ids.shape[0] ,						_A ,						_A       )
				_UpperCAmelCase      :							Optional[Any]       =   jnp.broadcast_to(
				    jnp.arange(decoder_input_ids.shape[-1] - 1       )[None, :] ,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1) ,						)

				_UpperCAmelCase      :							int       =   model.decode(
				    decoder_input_ids[:, :-1] ,						_A ,						decoder_attention_mask=_A ,						past_key_values=_A ,						decoder_position_ids=_A ,						)
				_UpperCAmelCase      :							Tuple       =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]] ,						dtype="""i4"""       )
				_UpperCAmelCase      :							Tuple       =   model.decode(
				    decoder_input_ids[:, -1:] ,						_A ,						past_key_values=outputs_cache.past_key_values ,						decoder_attention_mask=_A ,						decoder_position_ids=_A ,						)

				_UpperCAmelCase      :							int       =   model.decode(_A ,						_A ,						decoder_attention_mask=_A       )

				_UpperCAmelCase      :							Optional[int]       =   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       _UpperCAmelCase     (						unittest.TestCase):
			__a	:    str			   =							9_9
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Optional[Any]       =   np.array(
				    [
				        [71, 82, 18, 33, 46, 91, 2],
				        [68, 34, 26, 58, 30, 82, 2],
				        [5, 97, 17, 39, 94, 40, 2],
				        [76, 83, 94, 25, 70, 78, 2],
				        [87, 59, 41, 35, 48, 66, 2],
				        [55, 13, 16, 58, 5, 2, 1],  # note padding
				        [64, 27, 31, 51, 12, 75, 2],
				        [52, 64, 86, 17, 83, 39, 2],
				        [48, 61, 9, 24, 71, 82, 2],
				        [26, 1, 60, 48, 22, 13, 2],
				        [21, 5, 62, 28, 14, 76, 2],
				        [45, 98, 37, 86, 59, 48, 2],
				        [70, 70, 50, 9, 28, 0, 2],
				    ] ,						dtype=np.intaa ,						)

				_UpperCAmelCase      :							Optional[int]       =   input_ids.shape[0]
				_UpperCAmelCase      :							Optional[Any]       =   MBartConfig(
				    vocab_size=self.vocab_size ,						d_model=24 ,						encoder_layers=2 ,						decoder_layers=2 ,						encoder_attention_heads=2 ,						decoder_attention_heads=2 ,						encoder_ffn_dim=32 ,						decoder_ffn_dim=32 ,						max_position_embeddings=48 ,						eos_token_id=2 ,						pad_token_id=1 ,						bos_token_id=0 ,						)
				return config, input_ids, batch_size
			def 							__snake_case	(						self       )      ->      str:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   self._get_config_and_data()
				_UpperCAmelCase      :							Any       =   FlaxMBartForSequenceClassification(_A       )
				_UpperCAmelCase      :							Tuple       =   model(input_ids=_A ,						decoder_input_ids=_A       )
				_UpperCAmelCase      :							str       =   (batch_size, config.num_labels)
				self.assertEqual(outputs["""logits"""].shape ,						_A       )
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase  ,				_UpperCAmelCase      :							Union[str, Any]       =   self._get_config_and_data()
				_UpperCAmelCase      :							Tuple       =   FlaxMBartForQuestionAnswering(_A       )
				_UpperCAmelCase      :							str       =   model(input_ids=_A       )

				self.assertEqual(outputs["""start_logits"""].shape ,						input_ids.shape       )
				self.assertEqual(outputs["""end_logits"""].shape ,						input_ids.shape       )
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self._get_config_and_data()
				_UpperCAmelCase      :							Optional[Any]       =   FlaxMBartForConditionalGeneration(_A       )
				_UpperCAmelCase      :							Union[str, Any]       =   lm_model(input_ids=_A       )
				_UpperCAmelCase      :							Tuple       =   (batch_size, input_ids.shape[1], config.vocab_size)
				self.assertEqual(outputs["""logits"""].shape ,						_A       )
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							Union[str, Any]       =   MBartConfig(
				    vocab_size=self.vocab_size ,						d_model=14 ,						encoder_layers=2 ,						decoder_layers=2 ,						encoder_attention_heads=2 ,						decoder_attention_heads=2 ,						encoder_ffn_dim=8 ,						decoder_ffn_dim=8 ,						max_position_embeddings=48 ,						)
				_UpperCAmelCase      :							Dict       =   FlaxMBartForConditionalGeneration(_A       )
				_UpperCAmelCase      :							List[Any]       =   np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]] ,						dtype=np.intaa       )
				_UpperCAmelCase      :							str       =   np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]] ,						dtype=np.intaa       )
				_UpperCAmelCase      :							Union[str, Any]       =   lm_model(input_ids=_A ,						decoder_input_ids=_A       )
				_UpperCAmelCase      :							Optional[int]       =   (*summary.shape, config.vocab_size)
				self.assertEqual(outputs["""logits"""].shape ,						_A       )
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Tuple       =   np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]] ,						dtype=np.intaa       )
				_UpperCAmelCase      :							int       =   shift_tokens_right(_A ,						1       )
				_UpperCAmelCase      :							List[Any]       =   np.equal(_A ,						1       ).astype(np.floataa       ).sum()
				_UpperCAmelCase      :							int       =   np.equal(_A ,						1       ).astype(np.floataa       ).sum()
				self.assertEqual(shifted.shape ,						input_ids.shape       )
				self.assertEqual(_A ,						n_pad_before - 1       )
				self.assertTrue(np.equal(shifted[:, 0] ,						2       ).all()       )






@require_flax
class       _UpperCAmelCase     (						__a      ,       unittest.TestCase      ,       __a):
			__a	:    Any			   =							True
			__a	:    Optional[Any]			   =							(
			    (
			        FlaxMBartModel,
			        FlaxMBartForConditionalGeneration,
			        FlaxMBartForSequenceClassification,
			        FlaxMBartForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)
			__a	:    List[Any]			   =							(FlaxMBartForConditionalGeneration,) if is_flax_available() else ()
			def 							__snake_case	(						self       )      ->      List[Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Union[str, Any]       =   FlaxMBartModelTester(self       )
			def 							__snake_case	(						self       )      ->      int:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							Dict       =   self.model_tester.prepare_config_and_inputs()
				for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward(_A ,						_A ,						_A       )
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   self.model_tester.prepare_config_and_inputs()
				for model_class in self.all_model_classes:
					self.model_tester.check_use_cache_forward_with_attn_mask(_A ,						_A ,						_A       )
			def 							__snake_case	(						self       )      ->      Optional[Any]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__       ):
						_UpperCAmelCase      :							int       =   self._prepare_for_class(_A ,						_A       )
						_UpperCAmelCase      :							Optional[Any]       =   model_class(_A       )

						@jax.jit
						def encode_jitted(_A ,						_A=None ,						**_A       ):
							return model.encode(input_ids=_A ,						attention_mask=_A       )

						with self.subTest("""JIT Enabled"""       ):
							_UpperCAmelCase      :							Optional[Any]       =   encode_jitted(**_A       ).to_tuple()

						with self.subTest("""JIT Disabled"""       ):
							with jax.disable_jit():
								_UpperCAmelCase      :							Optional[Any]       =   encode_jitted(**_A       ).to_tuple()

						self.assertEqual(len(_A       ) ,						len(_A       )       )
						for jitted_output, output in zip(_A ,						_A       ):
							self.assertEqual(jitted_output.shape ,						output.shape       )
			def 							__snake_case	(						self       )      ->      List[str]:






				'''simple docstring'''







				_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					with self.subTest(model_class.__name__       ):
						_UpperCAmelCase      :							Union[str, Any]       =   model_class(_A       )
						_UpperCAmelCase      :							str       =   model.encode(inputs_dict["""input_ids"""] ,						inputs_dict["""attention_mask"""]       )

						_UpperCAmelCase      :							Tuple       =   {
						    """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
						    """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
						    """encoder_outputs""": encoder_outputs,
						}

						@jax.jit
						def decode_jitted(_A ,						_A ,						_A       ):
							return model.decode(
							    decoder_input_ids=_A ,						decoder_attention_mask=_A ,						encoder_outputs=_A ,						)

						with self.subTest("""JIT Enabled"""       ):
							_UpperCAmelCase      :							Any       =   decode_jitted(**_A       ).to_tuple()

						with self.subTest("""JIT Disabled"""       ):
							with jax.disable_jit():
								_UpperCAmelCase      :							Optional[Any]       =   decode_jitted(**_A       ).to_tuple()

						self.assertEqual(len(_A       ) ,						len(_A       )       )
						for jitted_output, output in zip(_A ,						_A       ):
							self.assertEqual(jitted_output.shape ,						output.shape       )
			@slow
			def 							__snake_case	(						self       )      ->      str:






				'''simple docstring'''







				for model_class_name in self.all_model_classes:
					_UpperCAmelCase      :							List[Any]       =   model_class_name.from_pretrained("""facebook/mbart-large-cc25""" ,						from_pt=_A       )
					# FlaxMBartForSequenceClassification expects eos token in input_ids
					_UpperCAmelCase      :							List[Any]       =   np.ones((1, 1)       ) * model.config.eos_token_id
					_UpperCAmelCase      :							List[str]       =   model(_A       )
					self.assertIsNotNone(_A       )






@require_flax
@require_sentencepiece
@require_tokenizers
class       _UpperCAmelCase     (						unittest.TestCase):
			__a	:    Any			   =							[
			    """ UN Chief Says There Is No Military Solution in Syria""",
			]
			__a	:    Union[str, Any]			   =							[
			    """Şeful ONU declară că nu există o soluţie militară în Siria""",
			]
			__a	:    Any			   =							"""facebook/mbart-large-en-ro"""
			@cached_property
			def 							__snake_case	(						self       )      ->      Dict:






				'''simple docstring'''







				return AutoTokenizer.from_pretrained(self.model_name       )
			@cached_property
			def 							__snake_case	(						self       )      ->      Tuple:






				'''simple docstring'''







				_UpperCAmelCase      :							List[str]       =   FlaxMBartForConditionalGeneration.from_pretrained(self.model_name ,						from_pt=_A       )
				return model
			def 							__snake_case	(						self ,						**_A       )      ->      Optional[int]:






				'''simple docstring'''







				_UpperCAmelCase      :							Tuple       =   self.translate_src_text(**_A       )
				self.assertListEqual(self.expected_text ,						_A       )
			def 							__snake_case	(						self ,						**_A       )      ->      Dict:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   self.tokenizer(self.src_text ,						**_A ,						return_tensors="""np"""       )
				_UpperCAmelCase      :							List[str]       =   self.model.generate(
				    model_inputs.input_ids ,						attention_mask=model_inputs.attention_mask ,						decoder_start_token_id=self.tokenizer.lang_code_to_id["""ro_RO"""] ,						early_stopping=_A ,						num_beams=2 ,						).sequences
				_UpperCAmelCase      :							int       =   self.tokenizer.batch_decode(_A ,						skip_special_tokens=_A       )
				return generated_words
			@slow
			def 							__snake_case	(						self       )      ->      Optional[int]:






				'''simple docstring'''







				self._assert_generated_batch_equal_expected()








"""simple docstring"""






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

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


lowerCamelCase__			:					Dict							  =					logging.get_logger(__name__)

lowerCamelCase__			:					Optional[Any]							  =					{
    '''facebook/blenderbot-3B''': '''https://huggingface.co/facebook/blenderbot-3B/resolve/main/config.json''',
    # See all Blenderbot models at https://huggingface.co/models?filter=blenderbot
}






class       _UpperCAmelCase     (						__a):
			__a	:    Optional[int]			   =							"""blenderbot"""
			__a	:    Tuple			   =							["""past_key_values"""]
			__a	:    int			   =							{"""num_attention_heads""": """encoder_attention_heads""", """hidden_size""": """d_model"""}
			def __init__(						self ,						_A=80_08 ,						_A=1_28 ,						_A=2 ,						_A=1_02_40 ,						_A=32 ,						_A=24 ,						_A=1_02_40 ,						_A=32 ,						_A=0.0 ,						_A=0.0 ,						_A=True ,						_A=True ,						_A="gelu" ,						_A=25_60 ,						_A=0.1 ,						_A=0.0 ,						_A=0.0 ,						_A=0.02 ,						_A=1 ,						_A=False ,						_A=0 ,						_A=1 ,						_A=2 ,						_A=3 ,						_A=2 ,						**_A ,						)      ->      Union[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Any       =   vocab_size
				_UpperCAmelCase      :							Dict       =   max_position_embeddings
				_UpperCAmelCase      :							Optional[int]       =   d_model
				_UpperCAmelCase      :							Optional[Any]       =   encoder_ffn_dim
				_UpperCAmelCase      :							Optional[Any]       =   encoder_layers
				_UpperCAmelCase      :							List[Any]       =   encoder_attention_heads
				_UpperCAmelCase      :							str       =   decoder_ffn_dim
				_UpperCAmelCase      :							int       =   decoder_layers
				_UpperCAmelCase      :							List[str]       =   decoder_attention_heads
				_UpperCAmelCase      :							List[Any]       =   dropout
				_UpperCAmelCase      :							List[str]       =   attention_dropout
				_UpperCAmelCase      :							int       =   activation_dropout
				_UpperCAmelCase      :							Union[str, Any]       =   activation_function
				_UpperCAmelCase      :							List[str]       =   init_std
				_UpperCAmelCase      :							Dict       =   encoder_layerdrop
				_UpperCAmelCase      :							Union[str, Any]       =   decoder_layerdrop
				_UpperCAmelCase      :							str       =   use_cache
				_UpperCAmelCase      :							Optional[int]       =   encoder_layers
				_UpperCAmelCase      :							Union[str, Any]       =   scale_embedding  # scale factor will be sqrt(d_model) if True

				super().__init__(
				    pad_token_id=_A ,						bos_token_id=_A ,						eos_token_id=_A ,						is_encoder_decoder=_A ,						decoder_start_token_id=_A ,						encoder_no_repeat_ngram_size=_A ,						forced_eos_token_id=_A ,						**_A ,						)






class       _UpperCAmelCase     (						__a):
			@property
			def 							__snake_case	(						self       )      ->      Mapping[str, Mapping[int, str]]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Tuple       =   OrderedDict(
					    [
					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
					    ]       )
					if self.use_past:
						_UpperCAmelCase      :							List[str]       =   {0: """batch"""}
						_UpperCAmelCase      :							List[str]       =   {0: """batch""", 1: """past_decoder_sequence + sequence"""}
					else:
						_UpperCAmelCase      :							Optional[int]       =   {0: """batch""", 1: """decoder_sequence"""}
						_UpperCAmelCase      :							int       =   {0: """batch""", 1: """decoder_sequence"""}
					if self.use_past:
						self.fill_with_past_key_values_(_A ,						direction="""inputs"""       )
				elif self.task == "causal-lm":
					_UpperCAmelCase      :							List[str]       =   OrderedDict(
					    [
					        ("""input_ids""", {0: """batch""", 1: """encoder_sequence"""}),
					        ("""attention_mask""", {0: """batch""", 1: """encoder_sequence"""}),
					    ]       )
					if self.use_past:
						_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.num_layers
						for i in range(_A       ):
							_UpperCAmelCase      :							Any       =   {0: """batch""", 2: """past_sequence + sequence"""}
							_UpperCAmelCase      :							Dict       =   {0: """batch""", 2: """past_sequence + sequence"""}
				else:
					_UpperCAmelCase      :							List[Any]       =   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
			# Copied from transformers.models.bart.configuration_bart.BartOnnxConfig.outputs
			def 							__snake_case	(						self       )      ->      Mapping[str, Mapping[int, str]]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Optional[Any]       =   super().outputs
				else:
					_UpperCAmelCase      :							Optional[int]       =   super(_A ,						self       ).outputs
					if self.use_past:
						_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.num_layers
						for i in range(_A       ):
							_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: """past_sequence + sequence"""}
							_UpperCAmelCase      :							Optional[Any]       =   {0: """batch""", 2: """past_sequence + sequence"""}
				return common_outputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Dict       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )
				# Generate decoder inputs
				_UpperCAmelCase      :							Any       =   seq_length if not self.use_past else 1
				_UpperCAmelCase      :							List[str]       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )
				_UpperCAmelCase      :							str       =   {f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
				_UpperCAmelCase      :							Dict       =   dict(**_A ,						**_A       )

				if self.use_past:
					if not is_torch_available():
						raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
					else:
						import torch
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   common_inputs["""input_ids"""].shape
					_UpperCAmelCase      :							List[Any]       =   common_inputs["""decoder_input_ids"""].shape[1]
					_UpperCAmelCase  ,				_UpperCAmelCase      :							List[str]       =   self.num_attention_heads
					_UpperCAmelCase      :							List[str]       =   (
					    batch,
					    num_encoder_attention_heads,
					    encoder_seq_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					_UpperCAmelCase      :							List[str]       =   decoder_seq_length
					_UpperCAmelCase      :							Tuple       =   (
					    batch,
					    num_decoder_attention_heads,
					    decoder_past_length,
					    self._config.hidden_size // num_decoder_attention_heads,
					)
					_UpperCAmelCase      :							Tuple       =   torch.cat(
					    [common_inputs["""decoder_attention_mask"""], torch.ones(_A ,						_A       )] ,						dim=1       )
					_UpperCAmelCase      :							Tuple       =   []
					_UpperCAmelCase  ,				_UpperCAmelCase      :							int       =   self.num_layers

					for _ in range(_A       ):
						common_inputs["past_key_values"].append(
						    (
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						        torch.zeros(_A       ),
						    )       )
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							Union[str, Any]       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
				    _A ,						_A ,						_A ,						_A ,						_A       )

				if self.use_past:
					if not is_torch_available():
						raise ValueError("""Cannot generate dummy past_keys inputs without PyTorch installed."""       )
					else:
						import torch
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   common_inputs["""input_ids"""].shape
					_UpperCAmelCase      :							Union[str, Any]       =   seqlen
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Optional[Any]       =   self.num_layers
					_UpperCAmelCase  ,				_UpperCAmelCase      :							Tuple       =   self.num_attention_heads
					_UpperCAmelCase      :							Optional[Any]       =   (
					    batch,
					    num_encoder_attention_heads,
					    past_key_values_length,
					    self._config.hidden_size // num_encoder_attention_heads,
					)
					_UpperCAmelCase      :							str       =   common_inputs["""attention_mask"""].dtype
					_UpperCAmelCase      :							Tuple       =   torch.cat(
					    [common_inputs["""attention_mask"""], torch.ones(_A ,						_A ,						dtype=_A       )] ,						dim=1       )
					_UpperCAmelCase      :							str       =   [
					    (torch.zeros(_A       ), torch.zeros(_A       )) for _ in range(_A       )
					]
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				_UpperCAmelCase      :							str       =   compute_effective_axis_dimension(
				    _A ,						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
				_UpperCAmelCase      :							List[str]       =   tokenizer.num_special_tokens_to_add(_A       )
				_UpperCAmelCase      :							Tuple       =   compute_effective_axis_dimension(
				    _A ,						fixed_dimension=OnnxConfig.default_fixed_sequence ,						num_token_to_add=_A       )

				# Generate dummy inputs according to compute batch and sequence
				_UpperCAmelCase      :							str       =   [""" """.join([tokenizer.unk_token]       ) * seq_length] * batch_size
				_UpperCAmelCase      :							str       =   dict(tokenizer(_A ,						return_tensors=_A       )       )
				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A = -1 ,						_A = -1 ,						_A = False ,						_A = None ,						)      ->      Mapping[str, Any]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Any       =   self._generate_dummy_inputs_for_default_and_seqaseq_lm(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )

				elif self.task == "causal-lm":
					_UpperCAmelCase      :							int       =   self._generate_dummy_inputs_for_causal_lm(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
					    _A ,						batch_size=_A ,						seq_length=_A ,						is_pair=_A ,						framework=_A       )

				return common_inputs
			def 							__snake_case	(						self ,						_A ,						_A ,						_A ,						_A       )      ->      Optional[int]:






				'''simple docstring'''







				if self.task in ["default", "seq2seq-lm"]:
					_UpperCAmelCase      :							Union[str, Any]       =   super()._flatten_past_key_values_(_A ,						_A ,						_A ,						_A       )
				else:
					_UpperCAmelCase      :							Tuple       =   super(_A ,						self       )._flatten_past_key_values_(
					    _A ,						_A ,						_A ,						_A       )
			def 							__snake_case	(						self ,						_A ,						_A       )      ->      Optional[Any]:






				'''simple docstring'''







				if direction not in ["inputs", "outputs"]:
					raise ValueError(f'''direction must either be "inputs" or "outputs", but {direction} was given'''       )

				_UpperCAmelCase      :							Any       =   """past_key_values""" if direction == """inputs""" else """present"""
				_UpperCAmelCase  ,				_UpperCAmelCase      :							Any       =   self.num_layers

				_UpperCAmelCase      :							List[str]       =   """past_encoder_sequence"""
				_UpperCAmelCase      :							int       =   """past_decoder_sequence""" if direction == """inputs""" else """past_decoder_sequence + sequence"""

				for i in range(_A       ):
					_UpperCAmelCase      :							Union[str, Any]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Optional[int]       =   {0: """batch""", 2: decoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}
					_UpperCAmelCase      :							Tuple       =   {0: """batch""", 2: encoder_sequence}





