



from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCAmelCase						  =     logging.get_logger(__name__)

__UpperCAmelCase						  =     {
    '''mnaylor/mega-base-wikitext''': '''https://huggingface.co/mnaylor/mega-base-wikitext/resolve/main/config.json''',
}


class 	lowerCamelCase__ (		__UpperCAmelCase   ):
						_lowerCAmelCase								= """mega"""


						def __init__( self						:  Union[str, Any]      ,       _a						:  Optional[Any]=3_0_5_2_2      ,       _a						:  List[Any]=1_2_8      ,       _a						:  Union[str, Any]=4      ,       _a						:  str=2_5_6      ,       _a						:  int=1_6      ,       _a						:  Dict=True      ,       _a						:  Any=6_4      ,       _a						:  Tuple=False      ,       _a						:  Any=-1      ,       _a						:  List[str]=None      ,       _a						:  str=True      ,       _a						:  int="scalenorm"      ,       _a						:  int=True      ,       _a						:  Optional[Any]="silu"      ,       _a						:  List[str]="softmax"      ,       _a						:  Union[str, Any]=0.1      ,       _a						:  Union[str, Any]=0.1      ,       _a						:  int=0.1      ,       _a						:  Optional[Any]=False      ,       _a						:  Dict=True      ,       _a						:  List[Any]=2_5_6      ,       _a						:  Dict=True      ,       _a						:  int=0.1      ,       _a						:  Dict=2_0_4_8      ,       _a						:  Tuple=False      ,       _a						:  Union[str, Any]=2      ,       _a						:  List[str]=0.0_2      ,       _a						:  Optional[Any]=0.2      ,       _a						:  List[Any]=0.0_2      ,       _a						:  Union[str, Any]=1.0      ,       _a						:  Union[str, Any]=1      ,       _a						:  int=0      ,       _a						:  Dict=2      ,       _a						:  List[Any]="rotary"      ,       _a						:  str=None      ,       _a						:  List[str]=True      ,       _a						:  str=True      ,       **_a						:  Optional[int]      ,       ):
							super().__init__(pad_token_id=lowerCamelCase__      ,       bos_token_id=lowerCamelCase__      ,       eos_token_id=lowerCamelCase__      ,       **lowerCamelCase__			)

							a__:       str											=vocab_size
							a__:       str											=hidden_size
							a__:       Tuple											=num_hidden_layers
							a__:       Union[str, Any]											=activation
							a__:       List[Any]											=attention_activation
							a__:       Optional[Any]											=intermediate_size
							a__:       Optional[int]											=ema_projection_size
							a__:       str											=bidirectional
							a__:       Optional[Any]											=shared_representation_size
							a__:       List[str]											=use_chunking
							a__:       Any											=chunk_size
							a__:       Any											=truncation
							a__:       int											=normalize_before_mega
							a__:       List[Any]											=normalization_type
							a__:       Dict											=norm_affine
							a__:       Optional[Any]											=dropout_prob
							a__:       List[str]											=hidden_dropout_prob
							a__:       Dict											=attention_probs_dropout_prob
							a__:       Union[str, Any]											=use_feature_dropout
							a__:       Optional[int]											=use_normalized_ffn
							a__:       Optional[Any]											=nffn_hidden_size
							a__:       Optional[int]											=normalize_before_ffn
							a__:       Any											=nffn_activation_dropout_prob
							a__:       int											=max_positions
							a__:       int											=add_token_type_embeddings
							a__:       Any											=type_vocab_size
							a__:       int											=initializer_range
							a__:       int											=ema_delta_alpha_range
							a__:       Any											=ema_beta_range
							a__:       Union[str, Any]											=ema_gamma_omega_range
							a__:       Union[str, Any]											=relative_positional_bias
							a__:       Optional[Any]											=use_cache
							a__:       Optional[Any]											=classifier_dropout
							a__:       Tuple											=add_lm_hidden_dense_layer
							a__:       int											=1  # not used but required by Hugging Face


class 	lowerCamelCase__ (		__UpperCAmelCase   ):


						@property
						def   _lowerCamelCase			( self						:  Any			):



							if self.task == "multiple-choice":
								a__:       Optional[int]											={0: "batch", 1: "choice", 2: "sequence"}
							else:
								a__:       Optional[int]											={0: "batch", 1: "sequence"}
							return OrderedDict(
							    [
							        ("input_ids", dynamic_axis),
							        ("attention_mask", dynamic_axis),
							    ]			)










import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


__UpperCAmelCase						  =     10_00_00_00_00_00
__UpperCAmelCase						  =     1_00_00

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





def 						__lowerCamelCase       (      __magic_name__	:  int     , __magic_name__	:  int						):
							a__:       Any											=pa.Table.from_pydict({"col": [0] * chunk_size}						)
							a__:       Optional[Any]											=pa.concat_tables([table] * (num_examples // chunk_size)						)
							return datasets.Dataset(__magic_name__     , fingerprint="table_100B"						)


@dataclass
class 	lowerCamelCase__ :
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42
						_lowerCAmelCase								= 42


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Tuple											=np.random.default_rng(self.seed			)
													a__:       Tuple											=rng.integers(low=self.low      ,       high=self.high      ,       size=self.size			).tolist()


						def __iter__( self						:  Any			):
													return iter(self._sampled_values			)







						def __len__( self						:  Any			):
													return self.size





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Union[str, Any]											=dataset[0]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[Any]											=dataset[-1]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       List[str]											=dataset[range(len(__magic_name__						) // 2     , len(__magic_name__						) // 2 + 1_024						)]





@get_duration
def 						__lowerCamelCase       (      __magic_name__	:  datasets.Dataset						):
							a__:       Optional[int]											=dataset[RandIter(0     , len(__magic_name__						)     , 1_024     , seed=42						)]





def 						__lowerCamelCase       (      ):
							a__:       Dict											={"num examples": SPEED_TEST_N_EXAMPLES}
							a__:       Union[str, Any]											=(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
							print("generating dataset"						)
							a__:       Dict											=generate_100B_dataset(num_examples=__magic_name__     , chunk_size=__magic_name__						)
							print("Functions"						)
							for func in functions:
														print(func.__name__						)
														a__:       Dict											=func(__magic_name__						)

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


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





