


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_  :	str							  =							logging.get_logger(__name__)

snake_case_  :	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   A__							(						UpperCamelCase__     ):
						UpperCAmelCase									=  "mbart"
						UpperCAmelCase									=  ["past_key_values"]
						UpperCAmelCase									=  {"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




						def __init__(				self    :					List[str] ,			_a    :					Tuple=5_0265 ,			_a    :					List[str]=1024 ,			_a    :					List[str]=12 ,			_a    :					Tuple=4096 ,			_a    :					Optional[int]=16 ,			_a    :					Union[str, Any]=12 ,			_a    :					Optional[Any]=4096 ,			_a    :					Optional[Any]=16 ,			_a    :					Union[str, Any]=0.0 ,			_a    :					List[Any]=0.0 ,			_a    :					Dict=True ,			_a    :					List[str]=True ,			_a    :					Optional[Any]="gelu" ,			_a    :					Optional[int]=1024 ,			_a    :					str=0.1 ,			_a    :					Tuple=0.0 ,			_a    :					Tuple=0.0 ,			_a    :					List[Any]=0.02 ,			_a    :					Dict=0.0 ,			_a    :					Any=False ,			_a    :					Union[str, Any]=1 ,			_a    :					Tuple=0 ,			_a    :					List[Any]=2 ,			_a    :					List[str]=2 ,			**_a    :					Union[str, Any] ,			)							->   Union[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =vocab_size
										_SCREAMING_SNAKE_CASE	     =max_position_embeddings
										_SCREAMING_SNAKE_CASE	     =d_model
										_SCREAMING_SNAKE_CASE	     =encoder_ffn_dim
										_SCREAMING_SNAKE_CASE	     =encoder_layers
										_SCREAMING_SNAKE_CASE	     =encoder_attention_heads
										_SCREAMING_SNAKE_CASE	     =decoder_ffn_dim
										_SCREAMING_SNAKE_CASE	     =decoder_layers
										_SCREAMING_SNAKE_CASE	     =decoder_attention_heads
										_SCREAMING_SNAKE_CASE	     =dropout
										_SCREAMING_SNAKE_CASE	     =attention_dropout
										_SCREAMING_SNAKE_CASE	     =activation_dropout
										_SCREAMING_SNAKE_CASE	     =activation_function
										_SCREAMING_SNAKE_CASE	     =init_std
										_SCREAMING_SNAKE_CASE	     =encoder_layerdrop
										_SCREAMING_SNAKE_CASE	     =decoder_layerdrop
										_SCREAMING_SNAKE_CASE	     =classifier_dropout
										_SCREAMING_SNAKE_CASE	     =use_cache
										_SCREAMING_SNAKE_CASE	     =encoder_layers
										_SCREAMING_SNAKE_CASE	     =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 ,			forced_eos_token_id=_a ,			**_a ,			)

class   A__							(						UpperCamelCase__     ):




						@property
						def 							__UpperCamelCase (				self    :					int     )							->   Mapping[str, Mapping[int, str]]:


										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
														_SCREAMING_SNAKE_CASE	     =OrderedDict(
														    [
														        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
														        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
														    ]     )

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

														if self.use_past:
																		self.fill_with_past_key_values_(_a ,			direction='''inputs'''     )
										elif self.task == "causal-lm":
														# TODO: figure this case out.
														_SCREAMING_SNAKE_CASE	     =OrderedDict(
														    [
														        ('''input_ids''', {0: '''batch''', 1: '''encoder_sequence'''}),
														        ('''attention_mask''', {0: '''batch''', 1: '''encoder_sequence'''}),
														    ]     )
														if self.use_past:
																		_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.num_layers
																		for i in range(_a     ):
																						_SCREAMING_SNAKE_CASE	     ={0: '''batch''', 2: '''past_sequence + sequence'''}
																						_SCREAMING_SNAKE_CASE	     ={0: '''batch''', 2: '''past_sequence + sequence'''}
										else:
														_SCREAMING_SNAKE_CASE	     =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    :					List[str]     )							->   Mapping[str, Mapping[int, str]]:


										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
														_SCREAMING_SNAKE_CASE	     =super().outputs
										else:
														_SCREAMING_SNAKE_CASE	     =super(_a ,			self     ).outputs
														if self.use_past:
																		_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.num_layers
																		for i in range(_a     ):
																						_SCREAMING_SNAKE_CASE	     ={0: '''batch''', 2: '''past_sequence + sequence'''}
																						_SCREAMING_SNAKE_CASE	     ={0: '''batch''', 2: '''past_sequence + sequence'''}
										return common_outputs




						def 							__UpperCamelCase (				self    :					Optional[Any] ,			_a    :					PreTrainedTokenizer ,			_a    :					int = -1 ,			_a    :					int = -1 ,			_a    :					bool = False ,			_a    :					Optional[TensorType] = None ,			)							->   Mapping[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    _a ,			_a ,			_a ,			_a ,			_a     )

										# Generate decoder inputs
										_SCREAMING_SNAKE_CASE	     =seq_length if not self.use_past else 1
										_SCREAMING_SNAKE_CASE	     =self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    _a ,			_a ,			_a ,			_a ,			_a     )
										_SCREAMING_SNAKE_CASE	     ={f"decoder_{name}": tensor for name, tensor in decoder_inputs.items()}
										_SCREAMING_SNAKE_CASE	     =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
														_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =common_inputs['''input_ids'''].shape
														_SCREAMING_SNAKE_CASE	     =common_inputs['''decoder_input_ids'''].shape[1]
														_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.num_attention_heads
														_SCREAMING_SNAKE_CASE	     =(
														    batch,
														    num_encoder_attention_heads,
														    encoder_seq_length,
														    self._config.hidden_size // num_encoder_attention_heads,
														)
														_SCREAMING_SNAKE_CASE	     =decoder_seq_length + 3
														_SCREAMING_SNAKE_CASE	     =(
														    batch,
														    num_decoder_attention_heads,
														    decoder_past_length,
														    self._config.hidden_size // num_decoder_attention_heads,
														)

														_SCREAMING_SNAKE_CASE	     =torch.cat(
														    [common_inputs['''decoder_attention_mask'''], torch.ones(_a ,			_a     )] ,			dim=1     )

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

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




						def 							__UpperCamelCase (				self    :					List[str] ,			_a    :					PreTrainedTokenizer ,			_a    :					int = -1 ,			_a    :					int = -1 ,			_a    :					bool = False ,			_a    :					Optional[TensorType] = None ,			)							->   Mapping[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =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
														_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =common_inputs['''input_ids'''].shape
														# Not using the same length for past_key_values
														_SCREAMING_SNAKE_CASE	     =seqlen + 2
														_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.num_layers
														_SCREAMING_SNAKE_CASE     , _SCREAMING_SNAKE_CASE	     =self.num_attention_heads
														_SCREAMING_SNAKE_CASE	     =(
														    batch,
														    num_encoder_attention_heads,
														    past_key_values_length,
														    self._config.hidden_size // num_encoder_attention_heads,
														)

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




						def 							__UpperCamelCase (				self    :					Optional[Any] ,			_a    :					PreTrainedTokenizer ,			_a    :					int = -1 ,			_a    :					int = -1 ,			_a    :					bool = False ,			_a    :					Optional[TensorType] = None ,			)							->   Mapping[str, Any]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE	     =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
										_SCREAMING_SNAKE_CASE	     =tokenizer.num_special_tokens_to_add(_a     )
										_SCREAMING_SNAKE_CASE	     =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
										_SCREAMING_SNAKE_CASE	     =[''' '''.join([tokenizer.unk_token]     ) * seq_length] * batch_size
										_SCREAMING_SNAKE_CASE	     =dict(tokenizer(_a ,			return_tensors=_a     )     )
										return common_inputs




						def 							__UpperCamelCase (				self    :					List[str] ,			_a    :					PreTrainedTokenizer ,			_a    :					int = -1 ,			_a    :					int = -1 ,			_a    :					bool = False ,			_a    :					Optional[TensorType] = None ,			)							->   Mapping[str, Any]:


										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
														_SCREAMING_SNAKE_CASE	     =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":
														_SCREAMING_SNAKE_CASE	     =self._generate_dummy_inputs_for_causal_lm(
														    _a ,			batch_size=_a ,			seq_length=_a ,			is_pair=_a ,			framework=_a     )
										else:
														_SCREAMING_SNAKE_CASE	     =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 							__UpperCamelCase (				self    :					Union[str, Any] ,			_a    :					Optional[Any] ,			_a    :					List[str] ,			_a    :					Tuple ,			_a    :					Any     )							->   Optional[Any]:


										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
														_SCREAMING_SNAKE_CASE	     =super()._flatten_past_key_values_(_a ,			_a ,			_a ,			_a     )
										else:
														_SCREAMING_SNAKE_CASE	     =super(_a ,			self     )._flatten_past_key_values_(
														    _a ,			_a ,			_a ,			_a     )

"""simple docstring"""




import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


snake_case    =		1_0_0_0_0_0_0_0_0_0_0_0
snake_case    =		1_0_0_0_0

snake_case   ,			snake_case    =		os.path.split(__file__)
snake_case    =		os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
					SCREAMING_SNAKE_CASE				=	pa.Table.from_pydict({'col': [0] * chunk_size}			)
					SCREAMING_SNAKE_CASE				=	pa.concat_tables([table] * (num_examples // chunk_size)			)
					return datasets.Dataset(SCREAMING_SNAKE_CASE_,     fingerprint='table_100B'			)




@dataclass
class      UpperCamelCase	:

				"""simple docstring"""



				UpperCAmelCase_				:      int
				UpperCAmelCase_				:      int
				UpperCAmelCase_				:      int
				UpperCAmelCase_				:      int


				def 							A					(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE				=	np.random.default_rng(self.seed			)
									SCREAMING_SNAKE_CASE				=	rng.integers(low=self.low		,	high=self.high		,	size=self.size			).tolist()


				def __iter__(				self			)    ->    Union[str, Any]:
									"""simple docstring"""







									return iter(self._sampled_values			)




				def __len__(				self			)    ->    Optional[Any]:
									"""simple docstring"""







									return self.size






@get_duration
def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_			):
					SCREAMING_SNAKE_CASE				=	dataset[0]

@get_duration
def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_			):
					SCREAMING_SNAKE_CASE				=	dataset[-1]

@get_duration
def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_			):
					SCREAMING_SNAKE_CASE				=	dataset[range(len(SCREAMING_SNAKE_CASE_			) // 2,     len(SCREAMING_SNAKE_CASE_			) // 2 + 1_0_2_4			)]

@get_duration
def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_			):
					SCREAMING_SNAKE_CASE				=	dataset[RandIter(0,     len(SCREAMING_SNAKE_CASE_			),     1_0_2_4,     seed=4_2			)]

def   UpperCamelCase_  (		):
					SCREAMING_SNAKE_CASE				=	{'num examples': SPEED_TEST_N_EXAMPLES}
					SCREAMING_SNAKE_CASE				=	(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print('generating dataset'			)
					SCREAMING_SNAKE_CASE				=	generate_100B_dataset(num_examples=SCREAMING_SNAKE_CASE_,     chunk_size=SCREAMING_SNAKE_CASE_			)
					print('Functions'			)
					for func in functions:
										print(func.__name__			)
										SCREAMING_SNAKE_CASE				=	func(SCREAMING_SNAKE_CASE_			)

					with open(SCREAMING_SNAKE_CASE_,     'wb'			) as f:
										f.write(json.dumps(SCREAMING_SNAKE_CASE_			).encode('utf-8'			)			)


if __name__ == "__main__":  # useful to run the profiler
	benchmark_table_100B()
