






"""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      ,							lowerCAmelCase_      ,							lowerCAmelCase_=1_3      ,							lowerCAmelCase_=7      ,							lowerCAmelCase_=True      ,							lowerCAmelCase_=True      ,							lowerCAmelCase_=True      ,							lowerCAmelCase_=True      ,							lowerCAmelCase_=9_9      ,							lowerCAmelCase_=2_4      ,							lowerCAmelCase_=3_2      ,							lowerCAmelCase_=5      ,							lowerCAmelCase_=4      ,							lowerCAmelCase_=3_7      ,							lowerCAmelCase_="gelu"      ,							lowerCAmelCase_=0.1      ,							lowerCAmelCase_=0.1      ,							lowerCAmelCase_=5_1_2      ,							lowerCAmelCase_=1_6      ,							lowerCAmelCase_=2      ,							lowerCAmelCase_=0.02      ,							lowerCAmelCase_=4      ,							)		->       Union[str, Any]:
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						parent
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						batch_size
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						seq_length
										_SCREAMING_SNAKE_CASE      :		Tuple							=						is_training
										_SCREAMING_SNAKE_CASE      :		Tuple							=						use_attention_mask
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						use_token_type_ids
										_SCREAMING_SNAKE_CASE      :		Dict							=						use_labels
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						vocab_size
										_SCREAMING_SNAKE_CASE      :		int							=						hidden_size
										_SCREAMING_SNAKE_CASE      :		str							=						embedding_size
										_SCREAMING_SNAKE_CASE      :		Any							=						num_hidden_layers
										_SCREAMING_SNAKE_CASE      :		List[str]							=						num_attention_heads
										_SCREAMING_SNAKE_CASE      :		str							=						intermediate_size
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						hidden_act
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						hidden_dropout_prob
										_SCREAMING_SNAKE_CASE      :		List[str]							=						attention_probs_dropout_prob
										_SCREAMING_SNAKE_CASE      :		Tuple							=						max_position_embeddings
										_SCREAMING_SNAKE_CASE      :		Optional[int]							=						type_vocab_size
										_SCREAMING_SNAKE_CASE      :		int							=						type_sequence_label_size
										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						initializer_range
										_SCREAMING_SNAKE_CASE      :		Tuple							=						num_choices







						def  A			(						self       )		->       List[Any]:
										_SCREAMING_SNAKE_CASE      :		List[Any]							=						ids_tensor([self.batch_size, self.seq_length]      ,							self.vocab_size       )

										_SCREAMING_SNAKE_CASE      :		str							=						None
										if self.use_attention_mask:
														_SCREAMING_SNAKE_CASE      :		int							=						random_attention_mask([self.batch_size, self.seq_length]       )

										_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						None
										if self.use_token_type_ids:
														_SCREAMING_SNAKE_CASE      :		Tuple							=						ids_tensor([self.batch_size, self.seq_length]      ,							self.type_vocab_size       )

										_SCREAMING_SNAKE_CASE      :		List[str]							=						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  A			(						self       )		->       str:
										_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						self.prepare_config_and_inputs()
										_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						config_and_inputs
										_SCREAMING_SNAKE_CASE      :		Dict							=						{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
										return config, inputs_dict







@require_flax
class _lowerCAmelCase	(						__UpperCAmelCase      ,     unittest.TestCase					):
						SCREAMING_SNAKE_CASE_:				int   				=							True

						SCREAMING_SNAKE_CASE_:				Dict   				=							(
						    (
						        FlaxElectraModel,
						        FlaxElectraForCausalLM,
						        FlaxElectraForMaskedLM,
						        FlaxElectraForPreTraining,
						        FlaxElectraForTokenClassification,
						        FlaxElectraForQuestionAnswering,
						        FlaxElectraForMultipleChoice,
						        FlaxElectraForSequenceClassification,
						    )
						    if is_flax_available()
						    else ()
						)







						def  A			(						self       )		->       List[Any]:
										_SCREAMING_SNAKE_CASE      :		Tuple							=						FlaxElectraModelTester(self       )







						@slow
						def  A			(						self       )		->       str:
										for model_class_name in self.all_model_classes:
														if model_class_name == FlaxElectraForMaskedLM:
																		_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						model_class_name.from_pretrained('google/electra-small-generator'       )
														else:
																		_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						model_class_name.from_pretrained('google/electra-small-discriminator'       )
														_SCREAMING_SNAKE_CASE      :		Optional[int]							=						model(np.ones((1, 1)       )       )
														self.assertIsNotNone(lowerCAmelCase_       )









"""simple docstring"""
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Optional

import transformers
from transformers import (
    AutoConfig,
    AutoModelForQuestionAnswering,
    AutoTokenizer,
    DataCollatorWithPadding,
    HfArgumentParser,
    SquadDataset,
    Trainer,
    TrainingArguments,
)
from transformers import SquadDataTrainingArguments as DataTrainingArguments
from transformers.trainer_utils import is_main_process


lowerCAmelCase__							= logging.getLogger(__name__)


@dataclass
class _lowerCAmelCase	:
						SCREAMING_SNAKE_CASE_:				str   				=							field(
						    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}					)
						SCREAMING_SNAKE_CASE_:				Optional[str]   				=							field(
						    default=__UpperCAmelCase      ,     metadata={'help': 'Pretrained config name or path if not the same as model_name'}					)
						SCREAMING_SNAKE_CASE_:				Optional[str]   				=							field(
						    default=__UpperCAmelCase      ,     metadata={'help': 'Pretrained tokenizer name or path if not the same as model_name'}					)
						SCREAMING_SNAKE_CASE_:				bool   				=							field(default=__UpperCAmelCase      ,     metadata={'help': 'Set this flag to use fast tokenization.'}					)
						# If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
						# or just modify its tokenizer_config.json.
						SCREAMING_SNAKE_CASE_:				Optional[str]   				=							field(
						    default=__UpperCAmelCase      ,     metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}      ,     )







def 							lowercase__		(				):
				# See all possible arguments in src/transformers/training_args.py
				# or by passing the --help flag to this script.
				# We now keep distinct sets of args, for a cleaner separation of concerns.

				_SCREAMING_SNAKE_CASE      :		List[Any]							=						HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)

				if len(sys.argv					) == 2 and sys.argv[1].endswith('.json'					):
								# If we pass only one argument to the script and it's the path to a json file,
								# let's parse it to get our arguments.
								_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]					)					)
				else:
								_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE		,	_SCREAMING_SNAKE_CASE      :		List[Any]							=						parser.parse_args_into_dataclasses()

				if (
				    os.path.exists(training_args.output_dir					)
				    and os.listdir(training_args.output_dir					)
				    and training_args.do_train
				    and not training_args.overwrite_output_dir
				):
								raise ValueError(
								    f"""Output directory ({training_args.output_dir}) already exists and is not empty. Use"""
								    ' --overwrite_output_dir to overcome.'					)

				# Setup logging
				logging.basicConfig(
				    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',	datefmt='%m/%d/%Y %H:%M:%S',	level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,	)
				logger.warning(
				    'Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s',	training_args.local_rank,	training_args.device,	training_args.n_gpu,	bool(training_args.local_rank != -1					),	training_args.fpaa,	)
				# Set the verbosity to info of the Transformers logger (on main process only):
				if is_main_process(training_args.local_rank					):
								transformers.utils.logging.set_verbosity_info()
								transformers.utils.logging.enable_default_handler()
								transformers.utils.logging.enable_explicit_format()
				logger.info('Training/evaluation parameters %s',	lowerCamelCase					)

				# Prepare Question-Answering task
				# Load pretrained model and tokenizer
				#
				# Distributed training:
				# The .from_pretrained methods guarantee that only one local process can concurrently
				# download model & vocab.

				_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						AutoConfig.from_pretrained(
				    model_args.config_name if model_args.config_name else model_args.model_name_or_path,	cache_dir=model_args.cache_dir,	)
				_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						AutoTokenizer.from_pretrained(
				    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,	cache_dir=model_args.cache_dir,	use_fast=lowerCamelCase,	)
				_SCREAMING_SNAKE_CASE      :		Dict							=						AutoModelForQuestionAnswering.from_pretrained(
				    model_args.model_name_or_path,	from_tf=bool('.ckpt' in model_args.model_name_or_path					),	config=lowerCamelCase,	cache_dir=model_args.cache_dir,	)

				# Get datasets
				_SCREAMING_SNAKE_CASE      :		Optional[Any]							=						hasattr(model.config,	'lang2id'					)
				_SCREAMING_SNAKE_CASE      :		List[Any]							=						(
				    SquadDataset(
				        lowerCamelCase,	tokenizer=lowerCamelCase,	is_language_sensitive=lowerCamelCase,	cache_dir=model_args.cache_dir					)
				    if training_args.do_train
				    else None
				)
				_SCREAMING_SNAKE_CASE      :		Union[str, Any]							=						(
				    SquadDataset(
				        lowerCamelCase,	tokenizer=lowerCamelCase,	mode='dev',	is_language_sensitive=lowerCamelCase,	cache_dir=model_args.cache_dir,	)
				    if training_args.do_eval
				    else None
				)

				# Data collator
				_SCREAMING_SNAKE_CASE      :		int							=						DataCollatorWithPadding(lowerCamelCase,	pad_to_multiple_of=8					) if training_args.fpaa else None

				# Initialize our Trainer
				_SCREAMING_SNAKE_CASE      :		List[Any]							=						Trainer(
				    model=lowerCamelCase,	args=lowerCamelCase,	train_dataset=lowerCamelCase,	eval_dataset=lowerCamelCase,	data_collator=lowerCamelCase,	)

				# Training
				if training_args.do_train:
								trainer.train(
								    model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path					) else None					)
								trainer.save_model()
								# For convenience, we also re-save the tokenizer to the same directory,
								# so that you can share your model easily on huggingface.co/models =)
								if trainer.is_world_master():
												tokenizer.save_pretrained(training_args.output_dir					)





def 							lowercase__		(				lowerCamelCase					):
				# For xla_spawn (TPUs)
				main()


if __name__ == "__main__":
			main()

