




'''simple docstring'''


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


_snake_case      :      Optional[int]   							=		logging.get_logger(__name__)

_snake_case      :      Any   							=		{
    '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    (		_a    ):
							lowercase_							=				'mbart'
							lowercase_							=				['past_key_values']
							lowercase_							=				{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}







							def __init__(  self       :  Union[str, Any]							,			lowerCAmelCase_       :  Any=5_02_65							,			lowerCAmelCase_       :  Dict=10_24							,			lowerCAmelCase_       :  Tuple=12							,			lowerCAmelCase_       :  Optional[Any]=40_96							,			lowerCAmelCase_       :  List[Any]=16							,			lowerCAmelCase_       :  List[Any]=12							,			lowerCAmelCase_       :  Optional[Any]=40_96							,			lowerCAmelCase_       :  List[str]=16							,			lowerCAmelCase_       :  Dict=0.0							,			lowerCAmelCase_       :  List[str]=0.0							,			lowerCAmelCase_       :  Any=True							,			lowerCAmelCase_       :  Any=True							,			lowerCAmelCase_       :  str="gelu"							,			lowerCAmelCase_       :  Optional[Any]=10_24							,			lowerCAmelCase_       :  Optional[Any]=0.1							,			lowerCAmelCase_       :  List[str]=0.0							,			lowerCAmelCase_       :  Dict=0.0							,			lowerCAmelCase_       :  Dict=0.0_2							,			lowerCAmelCase_       :  int=0.0							,			lowerCAmelCase_       :  Optional[int]=False							,			lowerCAmelCase_       :  str=1							,			lowerCAmelCase_       :  str=0							,			lowerCAmelCase_       :  List[Any]=2							,			lowerCAmelCase_       :  Union[str, Any]=2							,			**lowerCAmelCase_       :  List[str]							,			)			->							Union[str, Any]:


														"""simple docstring"""


														_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 			A    (		_a    ):







							@property
							def    __lowerCAmelCase			(  self       :  List[Any]					)			->							Mapping[str, Mapping[int, str]]:


														"""simple docstring"""


														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    __lowerCAmelCase			(  self       :  str					)			->							Mapping[str, Mapping[int, str]]:


														"""simple docstring"""


														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    __lowerCAmelCase			(  self       :  Dict							,			lowerCAmelCase_       :  PreTrainedTokenizer							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  bool = False							,			lowerCAmelCase_       :  Optional[TensorType] = None							,			)			->							Mapping[str, Any]:


														"""simple docstring"""


														_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    __lowerCAmelCase			(  self       :  Union[str, Any]							,			lowerCAmelCase_       :  PreTrainedTokenizer							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  bool = False							,			lowerCAmelCase_       :  Optional[TensorType] = None							,			)			->							Mapping[str, Any]:


														"""simple docstring"""


														_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    __lowerCAmelCase			(  self       :  int							,			lowerCAmelCase_       :  PreTrainedTokenizer							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  bool = False							,			lowerCAmelCase_       :  Optional[TensorType] = None							,			)			->							Mapping[str, Any]:


														"""simple docstring"""


														_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    __lowerCAmelCase			(  self       :  List[str]							,			lowerCAmelCase_       :  PreTrainedTokenizer							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  int = -1							,			lowerCAmelCase_       :  bool = False							,			lowerCAmelCase_       :  Optional[TensorType] = None							,			)			->							Mapping[str, Any]:


														"""simple docstring"""


														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    __lowerCAmelCase			(  self       :  List[Any]							,			lowerCAmelCase_       :  List[Any]							,			lowerCAmelCase_       :  int							,			lowerCAmelCase_       :  Dict							,			lowerCAmelCase_       :  Union[str, Any]					)			->							Tuple:


														"""simple docstring"""


														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_					)







'''simple docstring'''


from dataclasses import dataclass
from typing import Callable, Optional

import datasets

@dataclass
class 			A    (		datasets.BuilderConfig    ):
							lowercase_							=				None
							lowercase_							=				None
							lowercase_							=				None







							def    __lowerCAmelCase			(  self       :  Dict					)			->							List[str]:


														"""simple docstring"""


														assert self.generator is not None, "generator must be specified"

														if self.gen_kwargs is None:
																					_a	     =				{}



class 			A    (		datasets.GeneratorBasedBuilder    ):
							lowercase_							=				GeneratorConfig







							def    __lowerCAmelCase			(  self       :  Any					)			->							List[str]:


														"""simple docstring"""


														return datasets.DatasetInfo(features=self.config.features					)







							def    __lowerCAmelCase			(  self       :  Dict							,			lowerCAmelCase_       :  Optional[int]					)			->							Optional[Any]:


														"""simple docstring"""


														return [datasets.SplitGenerator(name=datasets.Split.TRAIN							,			gen_kwargs=self.config.gen_kwargs					)]







							def    __lowerCAmelCase			(  self       :  str							,			**lowerCAmelCase_       :  int					)			->							List[str]:


														"""simple docstring"""


														for idx, ex in enumerate(self.config.generator(**lowerCAmelCase_					)					):
																					yield idx, ex

