



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


__magic_name__       =     logging.get_logger(__name__)

__magic_name__       =     {
    '''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 	a__       (       _snake_case       ):


					"""simple docstring"""



					A__  :				Dict					 =				'''mbart'''
					A__  :				Dict					 =				['''past_key_values''']
					A__  :				Union[str, Any]					 =				{'''num_attention_heads''': '''encoder_attention_heads''', '''hidden_size''': '''d_model'''}

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




class 	a__       (       _snake_case       ):


					"""simple docstring"""



					@property
					def 				__UpperCAmelCase		(		self     :Tuple    ):
										if self.task in ["default", "seq2seq-lm"]:
															lowercase						   =       OrderedDict(
															    [
															        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
															        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
															    ]    )

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

															if self.use_past:
																				self.fill_with_past_key_values_(lowercase__      ,  direction='inputs'    )
										elif self.task == "causal-lm":
															# TODO: figure this case out.
															lowercase						   =       OrderedDict(
															    [
															        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
															        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
															    ]    )
															if self.use_past:
																				lowercase   ,      lowercase						   =       self.num_layers
																				for i in range(lowercase__    ):
																									lowercase						   =       {0: 'batch', 2: 'past_sequence + sequence'}
																									lowercase						   =       {0: 'batch', 2: 'past_sequence + sequence'}
										else:
															lowercase						   =       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     :Optional[Any]    ):
										if self.task in ["default", "seq2seq-lm"]:
															lowercase						   =       super().outputs
										else:
															lowercase						   =       super(lowercase__      ,  self    ).outputs
															if self.use_past:
																				lowercase   ,      lowercase						   =       self.num_layers
																				for i in range(lowercase__    ):
																									lowercase						   =       {0: 'batch', 2: 'past_sequence + sequence'}
																									lowercase						   =       {0: 'batch', 2: 'past_sequence + sequence'}
										return common_outputs

					def 				__UpperCAmelCase		(		self     :int      ,  lowercase__     :PreTrainedTokenizer      ,  lowercase__     :int = -1      ,  lowercase__     :int = -1      ,  lowercase__     :bool = False      ,  lowercase__     :Optional[TensorType] = None      ,  ):
										lowercase						   =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    lowercase__      ,  lowercase__      ,  lowercase__      ,  lowercase__      ,  lowercase__    )

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

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

															lowercase						   =       torch.cat(
															    [common_inputs['decoder_attention_mask'], torch.ones(lowercase__      ,  lowercase__    )]      ,  dim=1    )

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

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

					def 				__UpperCAmelCase		(		self     :int      ,  lowercase__     :PreTrainedTokenizer      ,  lowercase__     :int = -1      ,  lowercase__     :int = -1      ,  lowercase__     :bool = False      ,  lowercase__     :Optional[TensorType] = None      ,  ):
										lowercase						   =       self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    lowercase__      ,  lowercase__      ,  lowercase__      ,  lowercase__      ,  lowercase__    )

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

															lowercase						   =       common_inputs['attention_mask'].dtype
															lowercase						   =       torch.cat(
															    [common_inputs['attention_mask'], torch.ones(lowercase__      ,  lowercase__      ,  dtype=lowercase__    )]      ,  dim=1    )
															lowercase						   =       [
															    (torch.zeros(lowercase__    ), torch.zeros(lowercase__    )) for _ in range(lowercase__    )
															]
										return common_inputs

					def 				__UpperCAmelCase		(		self     :List[Any]      ,  lowercase__     :PreTrainedTokenizer      ,  lowercase__     :int = -1      ,  lowercase__     :int = -1      ,  lowercase__     :bool = False      ,  lowercase__     :Optional[TensorType] = None      ,  ):
										# 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
										lowercase						   =       compute_effective_axis_dimension(
										    lowercase__      ,  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
										lowercase						   =       tokenizer.num_special_tokens_to_add(lowercase__    )
										lowercase						   =       compute_effective_axis_dimension(
										    lowercase__      ,  fixed_dimension=OnnxConfig.default_fixed_sequence      ,  num_token_to_add=lowercase__    )

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

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

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

										return common_inputs



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











from typing import List, Optional, Tuple, Union

from . import logging
from .import_utils import is_torch_available, is_torch_version


if is_torch_available():
		import torch

__magic_name__       =     logging.get_logger(__name__)  # pylint: disable=invalid-name

try:
		from torch._dynamo import allow_in_graph as maybe_allow_in_graph
except (ImportError, ModuleNotFoundError):





		def 		__snake_case						(    cls				):
							"""simple docstring"""






							return cls





def 		__snake_case						(    _UpperCAmelCase       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			_UpperCAmelCase = None       ,			):
					"""simple docstring"""






					lowercase						   =       device
					lowercase						   =       shape[0]

					lowercase						   =       layout or torch.strided
					lowercase						   =       device or torch.device('cpu'				)

					if generator is not None:
										lowercase						   =       generator.device.type if not isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				) else generator[0].device.type
										if gen_device_type != device.type and gen_device_type == "cpu":
															lowercase						   =       'cpu'
															if device != "mps":
																				logger.info(
																				    f"""The passed generator was created on 'cpu' even though a tensor on {device} was expected."""
																				    f""" Tensors will be created on 'cpu' and then moved to {device}. Note that one can probably"""
																				    f""" slighly speed up this function by passing a generator that was created on the {device} device."""				)
										elif gen_device_type != device.type and gen_device_type == "cuda":
															raise ValueError(f"""Cannot generate a {device} tensor from a generator of type {gen_device_type}."""				)

					if isinstance(_UpperCAmelCase       ,			_UpperCAmelCase				):
										lowercase						   =       (1,) + shape[1:]
										lowercase						   =       [
										    torch.randn(_UpperCAmelCase       ,			generator=generator[i]       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				)
										    for i in range(_UpperCAmelCase				)
										]
										lowercase						   =       torch.cat(_UpperCAmelCase       ,			dim=0				).to(_UpperCAmelCase				)
					else:
										lowercase						   =       torch.randn(_UpperCAmelCase       ,			generator=_UpperCAmelCase       ,			device=_UpperCAmelCase       ,			dtype=_UpperCAmelCase       ,			layout=_UpperCAmelCase				).to(_UpperCAmelCase				)

					return latents





def 		__snake_case						(    _UpperCAmelCase				):
					"""simple docstring"""






					if is_torch_version('<'       ,			'2.0.0'				) or not hasattr(_UpperCAmelCase       ,			'_dynamo'				):
										return False
					return isinstance(_UpperCAmelCase       ,			torch._dynamo.eval_frame.OptimizedModule				)






