

import shutil
import tempfile
import unittest
from unittest.mock import patch

from transformers import (
    DefaultFlowCallback,
    IntervalStrategy,
    PrinterCallback,
    ProgressCallback,
    Trainer,
    TrainerCallback,
    TrainingArguments,
    is_torch_available,
)
from transformers.testing_utils import require_torch


if is_torch_available():
		from transformers.trainer import DEFAULT_CALLBACKS

		from .test_trainer import RegressionDataset, RegressionModelConfig, RegressionPreTrainedModel




class 			UpperCAmelCase_      (   A       ):




							'''simple docstring'''



							def __init__(					self       :	int			)			->					int:
										SCREAMING_SNAKE_CASE           =					[]





							def    _UpperCAmelCase     (					self       :	Dict       ,     a       :	Union[str, Any]       ,     a       :	Union[str, Any]       ,     a       :	Optional[int]       ,     **a       :	int			)			->					List[str]:
										self.events.append("""on_init_end"""			)





							def    _UpperCAmelCase     (					self       :	List[str]       ,     a       :	List[str]       ,     a       :	Optional[int]       ,     a       :	Optional[Any]       ,     **a       :	Union[str, Any]			)			->					Union[str, Any]:
										self.events.append("""on_train_begin"""			)





							def    _UpperCAmelCase     (					self       :	str       ,     a       :	Dict       ,     a       :	str       ,     a       :	Union[str, Any]       ,     **a       :	int			)			->					Dict:
										self.events.append("""on_train_end"""			)





							def    _UpperCAmelCase     (					self       :	Any       ,     a       :	List[str]       ,     a       :	Any       ,     a       :	Tuple       ,     **a       :	Any			)			->					str:
										self.events.append("""on_epoch_begin"""			)





							def    _UpperCAmelCase     (					self       :	Optional[int]       ,     a       :	Any       ,     a       :	Dict       ,     a       :	str       ,     **a       :	List[str]			)			->					List[str]:
										self.events.append("""on_epoch_end"""			)





							def    _UpperCAmelCase     (					self       :	Dict       ,     a       :	List[Any]       ,     a       :	Tuple       ,     a       :	Union[str, Any]       ,     **a       :	int			)			->					Dict:
										self.events.append("""on_step_begin"""			)





							def    _UpperCAmelCase     (					self       :	Dict       ,     a       :	Dict       ,     a       :	Union[str, Any]       ,     a       :	Optional[Any]       ,     **a       :	str			)			->					Any:
										self.events.append("""on_step_end"""			)





							def    _UpperCAmelCase     (					self       :	List[Any]       ,     a       :	Optional[Any]       ,     a       :	Dict       ,     a       :	Dict       ,     **a       :	Any			)			->					Optional[Any]:
										self.events.append("""on_evaluate"""			)





							def    _UpperCAmelCase     (					self       :	int       ,     a       :	Any       ,     a       :	List[str]       ,     a       :	Optional[Any]       ,     **a       :	List[str]			)			->					Tuple:
										self.events.append("""on_predict"""			)





							def    _UpperCAmelCase     (					self       :	Tuple       ,     a       :	List[str]       ,     a       :	int       ,     a       :	str       ,     **a       :	Tuple			)			->					List[Any]:
										self.events.append("""on_save"""			)





							def    _UpperCAmelCase     (					self       :	int       ,     a       :	List[Any]       ,     a       :	str       ,     a       :	Optional[int]       ,     **a       :	Optional[int]			)			->					List[str]:
										self.events.append("""on_log"""			)







							def    _UpperCAmelCase     (					self       :	int       ,     a       :	Optional[Any]       ,     a       :	Tuple       ,     a       :	List[Any]       ,     **a       :	Optional[Any]			)			->					Optional[int]:
										self.events.append("""on_prediction_step"""			)





@require_torch
class 			UpperCAmelCase_      (   unittest.TestCase       ):




							'''simple docstring'''



							def    _UpperCAmelCase     (					self       :	Union[str, Any]			)			->					Union[str, Any]:
										SCREAMING_SNAKE_CASE           =					tempfile.mkdtemp()





							def    _UpperCAmelCase     (					self       :	str			)			->					Tuple:
										shutil.rmtree(self.output_dir			)





							def    _UpperCAmelCase     (					self       :	Optional[int]       ,     a       :	int=0       ,     a       :	Dict=0       ,     a       :	int=64       ,     a       :	Any=64       ,     a       :	Union[str, Any]=None       ,     a       :	Tuple=False       ,     **a       :	Union[str, Any]			)			->					str:
										# disable_tqdm in TrainingArguments has a flaky default since it depends on the level of logging. We make sure
										# its set to False since the tests later on depend on its value.
										SCREAMING_SNAKE_CASE           =					RegressionDataset(length=a			)
										SCREAMING_SNAKE_CASE           =					RegressionDataset(length=a			)
										SCREAMING_SNAKE_CASE           =					RegressionModelConfig(a=a       ,     b=a			)
										SCREAMING_SNAKE_CASE           =					RegressionPreTrainedModel(a			)

										SCREAMING_SNAKE_CASE           =					TrainingArguments(self.output_dir       ,     disable_tqdm=a       ,     report_to=[]       ,     **a			)
										return Trainer(
										    a       ,     a       ,     train_dataset=a       ,     eval_dataset=a       ,     callbacks=a       ,     )





							def    _UpperCAmelCase     (					self       :	List[Any]       ,     a       :	Union[str, Any]       ,     a       :	Dict			)			->					Optional[Any]:
										self.assertEqual(len(a			)       ,     len(a			)			)

										# Order doesn't matter
										SCREAMING_SNAKE_CASE           =					sorted(a       ,     key=lambda a			: cb.__name__ if isinstance(a       ,     a			) else cb.__class__.__name__			)
										SCREAMING_SNAKE_CASE           =					sorted(a       ,     key=lambda a			: cb.__name__ if isinstance(a       ,     a			) else cb.__class__.__name__			)

										for cba, cba in zip(a       ,     a			):
													if isinstance(a       ,     a			) and isinstance(a       ,     a			):
																self.assertEqual(a       ,     a			)
													elif isinstance(a       ,     a			) and not isinstance(a       ,     a			):
																self.assertEqual(a       ,     cba.__class__			)
													elif not isinstance(a       ,     a			) and isinstance(a       ,     a			):
																self.assertEqual(cba.__class__       ,     a			)
													else:
																self.assertEqual(a       ,     a			)





							def    _UpperCAmelCase     (					self       :	Tuple       ,     a       :	List[str]			)			->					Any:
										SCREAMING_SNAKE_CASE           =					["""on_init_end""", """on_train_begin"""]
										SCREAMING_SNAKE_CASE           =					0
										SCREAMING_SNAKE_CASE           =					len(trainer.get_eval_dataloader()			)
										SCREAMING_SNAKE_CASE           =					["""on_prediction_step"""] * len(trainer.get_eval_dataloader()			) + ["""on_log""", """on_evaluate"""]
										for _ in range(trainer.state.num_train_epochs			):
													expected_events.append("""on_epoch_begin"""			)
													for _ in range(a			):
																step += 1
																expected_events += ["on_step_begin", "on_step_end"]
																if step % trainer.args.logging_steps == 0:
																			expected_events.append("""on_log"""			)
																if trainer.args.evaluation_strategy == IntervalStrategy.STEPS and step % trainer.args.eval_steps == 0:
																			expected_events += evaluation_events.copy()
																if step % trainer.args.save_steps == 0:
																			expected_events.append("""on_save"""			)
													expected_events.append("""on_epoch_end"""			)
													if trainer.args.evaluation_strategy == IntervalStrategy.EPOCH:
																expected_events += evaluation_events.copy()
										expected_events += ["on_log", "on_train_end"]
										return expected_events





							def    _UpperCAmelCase     (					self       :	List[str]			)			->					Any:
										SCREAMING_SNAKE_CASE           =					self.get_trainer()
										SCREAMING_SNAKE_CASE           =					DEFAULT_CALLBACKS.copy() + [ProgressCallback]
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										# Callbacks passed at init are added to the default callbacks
										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]			)
										expected_callbacks.append(a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
										SCREAMING_SNAKE_CASE           =					self.get_trainer(disable_tqdm=a			)
										SCREAMING_SNAKE_CASE           =					DEFAULT_CALLBACKS.copy() + [PrinterCallback]
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)





							def    _UpperCAmelCase     (					self       :	List[Any]			)			->					Union[str, Any]:
										SCREAMING_SNAKE_CASE           =					DEFAULT_CALLBACKS.copy() + [ProgressCallback]
										SCREAMING_SNAKE_CASE           =					self.get_trainer()

										# We can add, pop, or remove by class name
										trainer.remove_callback(a			)
										expected_callbacks.remove(a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer()
										SCREAMING_SNAKE_CASE           =					trainer.pop_callback(a			)
										self.assertEqual(cb.__class__       ,     a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										trainer.add_callback(a			)
										expected_callbacks.insert(0       ,     a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										# We can also add, pop, or remove by instance
										SCREAMING_SNAKE_CASE           =					self.get_trainer()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[0]
										trainer.remove_callback(a			)
										expected_callbacks.remove(a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[0]
										SCREAMING_SNAKE_CASE           =					trainer.pop_callback(a			)
										self.assertEqual(a       ,     a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)

										trainer.add_callback(a			)
										expected_callbacks.insert(0       ,     a			)
										self.check_callbacks_equality(trainer.callback_handler.callbacks       ,     a			)







							def    _UpperCAmelCase     (					self       :	List[Any]			)			->					Union[str, Any]:
										import warnings

										# XXX: for now ignore scatter_gather warnings in this test since it's not relevant to what's being tested
										warnings.simplefilter(action="""ignore"""       ,     category=a			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]			)
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										# Independent log/save/eval
										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]       ,     logging_steps=5			)
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]       ,     save_steps=5			)
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]       ,     eval_steps=5       ,     evaluation_strategy="""steps"""			)
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										SCREAMING_SNAKE_CASE           =					self.get_trainer(callbacks=[MyTestTrainerCallback]       ,     evaluation_strategy="""epoch"""			)
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										# A bit of everything
										SCREAMING_SNAKE_CASE           =					self.get_trainer(
										    callbacks=[MyTestTrainerCallback]       ,     logging_steps=3       ,     save_steps=10       ,     eval_steps=5       ,     evaluation_strategy="""steps"""       ,     )
										trainer.train()
										SCREAMING_SNAKE_CASE           =					trainer.callback_handler.callbacks[-2].events
										self.assertEqual(a       ,     self.get_expected_events(a			)			)

										# warning should be emitted for duplicated callbacks
										with patch("""transformers.trainer_callback.logger.warning"""			) as warn_mock:
													SCREAMING_SNAKE_CASE           =					self.get_trainer(
													    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]       ,     )
													assert str(a			) in warn_mock.call_args[0][0]












from ...configuration_utils import PretrainedConfig
from ...utils import logging


__A						:							Any							=  logging.get_logger(__name__)

__A						:							int							=  {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}




class 			UpperCAmelCase_      (   A       ):




							'''simple docstring'''



							a__						  =     '''megatron-bert'''





							def __init__(					self       :	str       ,     a       :	Union[str, Any]=29_056       ,     a       :	Optional[int]=1_024       ,     a       :	str=24       ,     a       :	str=16       ,     a       :	str=4_096       ,     a       :	Any="gelu"       ,     a       :	Any=0.1       ,     a       :	Optional[int]=0.1       ,     a       :	Optional[int]=512       ,     a       :	Optional[Any]=2       ,     a       :	List[str]=0.02       ,     a       :	List[Any]=1E-1_2       ,     a       :	Any=0       ,     a       :	List[Any]="absolute"       ,     a       :	List[str]=True       ,     **a       :	Optional[Any]       ,     )			->					List[str]:
										super().__init__(pad_token_id=a       ,     **a			)

										SCREAMING_SNAKE_CASE           =					vocab_size
										SCREAMING_SNAKE_CASE           =					hidden_size
										SCREAMING_SNAKE_CASE           =					num_hidden_layers
										SCREAMING_SNAKE_CASE           =					num_attention_heads
										SCREAMING_SNAKE_CASE           =					hidden_act
										SCREAMING_SNAKE_CASE           =					intermediate_size
										SCREAMING_SNAKE_CASE           =					hidden_dropout_prob
										SCREAMING_SNAKE_CASE           =					attention_probs_dropout_prob
										SCREAMING_SNAKE_CASE           =					max_position_embeddings
										SCREAMING_SNAKE_CASE           =					type_vocab_size
										SCREAMING_SNAKE_CASE           =					initializer_range
										SCREAMING_SNAKE_CASE           =					layer_norm_eps
										SCREAMING_SNAKE_CASE           =					position_embedding_type
										SCREAMING_SNAKE_CASE           =					use_cache






