

'''simple docstring'''

import unittest

import numpy as np

from transformers import ElectraConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
						from transformers.models.electra.modeling_flax_electra import (
						    FlaxElectraForCausalLM,
						    FlaxElectraForMaskedLM,
						    FlaxElectraForMultipleChoice,
						    FlaxElectraForPreTraining,
						    FlaxElectraForQuestionAnswering,
						    FlaxElectraForSequenceClassification,
						    FlaxElectraForTokenClassification,
						    FlaxElectraModel,
						)







class lowerCamelCase   (    unittest.TestCase				):

				def __init__(       self,						lowercase_,						lowercase_=13,						lowercase_=7,						lowercase_=True,						lowercase_=True,						lowercase_=True,						lowercase_=True,						lowercase_=99,						lowercase_=24,						lowercase_=32,						lowercase_=5,						lowercase_=4,						lowercase_=37,						lowercase_="gelu",						lowercase_=0.1,						lowercase_=0.1,						lowercase_=512,						lowercase_=16,						lowercase_=2,						lowercase_=0.02,						lowercase_=4,						)    ->							Union[str, Any]:
					snake_case							=					parent
					snake_case							=					batch_size
					snake_case							=					seq_length
					snake_case							=					is_training
					snake_case							=					use_attention_mask
					snake_case							=					use_token_type_ids
					snake_case							=					use_labels
					snake_case							=					vocab_size
					snake_case							=					hidden_size
					snake_case							=					embedding_size
					snake_case							=					num_hidden_layers
					snake_case							=					num_attention_heads
					snake_case							=					intermediate_size
					snake_case							=					hidden_act
					snake_case							=					hidden_dropout_prob
					snake_case							=					attention_probs_dropout_prob
					snake_case							=					max_position_embeddings
					snake_case							=					type_vocab_size
					snake_case							=					type_sequence_label_size
					snake_case							=					initializer_range
					snake_case							=					num_choices

				def 					_lowerCamelCase    (       self )    ->							int:
					snake_case							=					ids_tensor([self.batch_size, self.seq_length],						self.vocab_size )

					snake_case							=					None
					if self.use_attention_mask:
						snake_case							=					random_attention_mask([self.batch_size, self.seq_length] )

					snake_case							=					None
					if self.use_token_type_ids:
						snake_case							=					ids_tensor([self.batch_size, self.seq_length],						self.type_vocab_size )

					snake_case							=					ElectraConfig(
					    vocab_size=self.vocab_size,						hidden_size=self.hidden_size,						embedding_size=self.embedding_size,						num_hidden_layers=self.num_hidden_layers,						num_attention_heads=self.num_attention_heads,						intermediate_size=self.intermediate_size,						hidden_act=self.hidden_act,						hidden_dropout_prob=self.hidden_dropout_prob,						attention_probs_dropout_prob=self.attention_probs_dropout_prob,						max_position_embeddings=self.max_position_embeddings,						type_vocab_size=self.type_vocab_size,						initializer_range=self.initializer_range,						)

					return config, input_ids, token_type_ids, attention_mask






				def 					_lowerCamelCase    (       self )    ->							List[str]:
					snake_case							=					self.prepare_config_and_inputs()
					snake_case			,			snake_case			,			snake_case			,			snake_case							=					config_and_inputs
					snake_case							=					{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
					return config, inputs_dict







@require_flax
class lowerCamelCase   (    __lowerCAmelCase				,				unittest.TestCase				):
				snake_case_							     = True

				snake_case_							     = (
				    (
				        FlaxElectraModel,
				        FlaxElectraForCausalLM,
				        FlaxElectraForMaskedLM,
				        FlaxElectraForPreTraining,
				        FlaxElectraForTokenClassification,
				        FlaxElectraForQuestionAnswering,
				        FlaxElectraForMultipleChoice,
				        FlaxElectraForSequenceClassification,
				    )
				    if is_flax_available()
				    else ()
				)

				def 					_lowerCamelCase    (       self )    ->							str:
					snake_case							=					FlaxElectraModelTester(self )






				@slow
				def 					_lowerCamelCase    (       self )    ->							Any:
					for model_class_name in self.all_model_classes:
						if model_class_name == FlaxElectraForMaskedLM:
							snake_case							=					model_class_name.from_pretrained('google/electra-small-generator' )
						else:
							snake_case							=					model_class_name.from_pretrained('google/electra-small-discriminator' )
						snake_case							=					model(np.ones((1, 1) ) )
						self.assertIsNotNone(lowercase_ )











'''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


A					       =10_00_00_00_00_00
A					       =1_00_00

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





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


@dataclass
class      _a    :
       __a							:	int
       __a							:	int
       __a							:	int
       __a							:	int
       def        A (      self		:      List[str]							):



           '''simple docstring'''

           UpperCAmelCase     			=						np.random.default_rng(self.seed							)
           UpperCAmelCase     			=						rng.integers(low=self.low  ,							high=self.high  ,							size=self.size							).tolist()
       def __iter__(      self		:      Tuple							):



           '''simple docstring'''

           return iter(self._sampled_values							)







       def __len__(      self		:      Dict							):



           '''simple docstring'''

           return self.size






@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[0]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[-1]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[range(len(_a	) // 2       ,  len(_a	) // 2 + 1_0_2_4	)]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[RandIter(0       ,  len(_a	)       ,  1_0_2_4       ,  seed=4_2	)]





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

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


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




