


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 			_lowercase (							lowercase__):


       """simple docstring"""




       def __init__(   self							:			Dict       ):


             '''simple docstring'''



             lowerCamelCase__		:				Optional[int]    =  []





       def 						lowerCAmelCase     (   self							:			Union[str, Any]					,      __lowerCamelCase							:			Tuple					,      __lowerCamelCase							:			List[str]					,      __lowerCamelCase							:			str					,      **__lowerCamelCase							:			str       ):


             '''simple docstring'''



             self.events.append("on_init_end"       )





       def 						lowerCAmelCase     (   self							:			Union[str, Any]					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			Tuple					,      **__lowerCamelCase							:			List[str]       ):


             '''simple docstring'''



             self.events.append("on_train_begin"       )





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			Dict					,      __lowerCamelCase							:			Optional[Any]					,      __lowerCamelCase							:			str					,      **__lowerCamelCase							:			Any       ):


             '''simple docstring'''



             self.events.append("on_train_end"       )





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			List[str]					,      __lowerCamelCase							:			Tuple					,      __lowerCamelCase							:			Optional[int]					,      **__lowerCamelCase							:			str       ):


             '''simple docstring'''



             self.events.append("on_epoch_begin"       )





       def 						lowerCAmelCase     (   self							:			Optional[int]					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			Any					,      __lowerCamelCase							:			int					,      **__lowerCamelCase							:			Dict       ):


             '''simple docstring'''



             self.events.append("on_epoch_end"       )





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			Union[str, Any]					,      __lowerCamelCase							:			Optional[Any]					,      __lowerCamelCase							:			List[str]					,      **__lowerCamelCase							:			int       ):


             '''simple docstring'''



             self.events.append("on_step_begin"       )





       def 						lowerCAmelCase     (   self							:			Union[str, Any]					,      __lowerCamelCase							:			Optional[Any]					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			List[Any]					,      **__lowerCamelCase							:			Dict       ):


             '''simple docstring'''



             self.events.append("on_step_end"       )





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			str					,      __lowerCamelCase							:			List[Any]					,      __lowerCamelCase							:			Tuple					,      **__lowerCamelCase							:			Optional[int]       ):


             '''simple docstring'''



             self.events.append("on_evaluate"       )





       def 						lowerCAmelCase     (   self							:			Dict					,      __lowerCamelCase							:			str					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			Dict					,      **__lowerCamelCase							:			List[str]       ):


             '''simple docstring'''



             self.events.append("on_predict"       )





       def 						lowerCAmelCase     (   self							:			List[Any]					,      __lowerCamelCase							:			List[str]					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			List[str]					,      **__lowerCamelCase							:			Union[str, Any]       ):


             '''simple docstring'''



             self.events.append("on_save"       )





       def 						lowerCAmelCase     (   self							:			str					,      __lowerCamelCase							:			List[Any]					,      __lowerCamelCase							:			Union[str, Any]					,      __lowerCamelCase							:			List[str]					,      **__lowerCamelCase							:			List[str]       ):


             '''simple docstring'''



             self.events.append("on_log"       )





       def 						lowerCAmelCase     (   self							:			Any					,      __lowerCamelCase							:			int					,      __lowerCamelCase							:			Optional[int]					,      __lowerCamelCase							:			Tuple					,      **__lowerCamelCase							:			int       ):


             '''simple docstring'''



             self.events.append("on_prediction_step"       )







@require_torch
class 			_lowercase (							unittest.TestCase):


       """simple docstring"""




       def 						lowerCAmelCase     (   self							:			int       ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  tempfile.mkdtemp()





       def 						lowerCAmelCase     (   self							:			List[str]       ):


             '''simple docstring'''



             shutil.rmtree(self.output_dir       )





       def 						lowerCAmelCase     (   self							:			Tuple					,      __lowerCamelCase							:			Optional[int]=0					,      __lowerCamelCase							:			Tuple=0					,      __lowerCamelCase							:			Optional[Any]=64					,      __lowerCamelCase							:			Dict=64					,      __lowerCamelCase							:			Optional[Any]=None					,      __lowerCamelCase							:			Dict=False					,      **__lowerCamelCase							:			Optional[int]       ):


             '''simple docstring'''



             lowerCamelCase__		:				str    =  RegressionDataset(length=__lowerCamelCase       )
             lowerCamelCase__		:				int    =  RegressionDataset(length=__lowerCamelCase       )
             lowerCamelCase__		:				List[Any]    =  RegressionModelConfig(a=__lowerCamelCase					,      b=__lowerCamelCase       )
             lowerCamelCase__		:				Union[str, Any]    =  RegressionPreTrainedModel(__lowerCamelCase       )

             lowerCamelCase__		:				Union[str, Any]    =  TrainingArguments(self.output_dir					,      disable_tqdm=__lowerCamelCase					,      report_to=[]					,      **__lowerCamelCase       )
             return Trainer(
                 __lowerCamelCase					,      __lowerCamelCase					,      train_dataset=__lowerCamelCase					,      eval_dataset=__lowerCamelCase					,      callbacks=__lowerCamelCase					,      )





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			Union[str, Any]					,      __lowerCamelCase							:			Tuple       ):


             '''simple docstring'''



             self.assertEqual(len(__lowerCamelCase       )					,      len(__lowerCamelCase       )       )

             # Order doesn't matter
             lowerCamelCase__		:				Optional[int]    =  sorted(__lowerCamelCase					,      key=lambda __lowerCamelCase       : cb.__name__ if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else cb.__class__.__name__       )
             lowerCamelCase__		:				Optional[int]    =  sorted(__lowerCamelCase					,      key=lambda __lowerCamelCase       : cb.__name__ if isinstance(__lowerCamelCase					,      __lowerCamelCase       ) else cb.__class__.__name__       )

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





       def 						lowerCAmelCase     (   self							:			Optional[Any]					,      __lowerCamelCase							:			Dict       ):


             '''simple docstring'''



             lowerCamelCase__		:				int    =  ["on_init_end", "on_train_begin"]
             lowerCamelCase__		:				List[Any]    =  0
             lowerCamelCase__		:				Optional[Any]    =  len(trainer.get_eval_dataloader()       )
             lowerCamelCase__		:				List[Any]    =  ["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(__lowerCamelCase       ):
                         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 						lowerCAmelCase     (   self							:			Optional[int]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Any    =  self.get_trainer()
             lowerCamelCase__		:				Any    =  DEFAULT_CALLBACKS.copy() + [ProgressCallback]
             self.check_callbacks_equality(trainer.callback_handler.callbacks					,      __lowerCamelCase       )

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

             # TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
             lowerCamelCase__		:				Optional[Any]    =  self.get_trainer(disable_tqdm=__lowerCamelCase       )
             lowerCamelCase__		:				str    =  DEFAULT_CALLBACKS.copy() + [PrinterCallback]
             self.check_callbacks_equality(trainer.callback_handler.callbacks					,      __lowerCamelCase       )





       def 						lowerCAmelCase     (   self							:			List[Any]       ):


             '''simple docstring'''



             lowerCamelCase__		:				Tuple    =  DEFAULT_CALLBACKS.copy() + [ProgressCallback]
             lowerCamelCase__		:				Optional[Any]    =  self.get_trainer()

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

             lowerCamelCase__		:				Dict    =  self.get_trainer()
             lowerCamelCase__		:				Optional[int]    =  trainer.pop_callback(__lowerCamelCase       )
             self.assertEqual(cb.__class__					,      __lowerCamelCase       )
             self.check_callbacks_equality(trainer.callback_handler.callbacks					,      __lowerCamelCase       )

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

             # We can also add, pop, or remove by instance
             lowerCamelCase__		:				Tuple    =  self.get_trainer()
             lowerCamelCase__		:				int    =  trainer.callback_handler.callbacks[0]
             trainer.remove_callback(__lowerCamelCase       )
             expected_callbacks.remove(__lowerCamelCase       )
             self.check_callbacks_equality(trainer.callback_handler.callbacks					,      __lowerCamelCase       )

             lowerCamelCase__		:				Dict    =  self.get_trainer()
             lowerCamelCase__		:				Union[str, Any]    =  trainer.callback_handler.callbacks[0]
             lowerCamelCase__		:				Any    =  trainer.pop_callback(__lowerCamelCase       )
             self.assertEqual(__lowerCamelCase					,      __lowerCamelCase       )
             self.check_callbacks_equality(trainer.callback_handler.callbacks					,      __lowerCamelCase       )

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





       def 						lowerCAmelCase     (   self							:			str       ):


             '''simple docstring'''



             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=__lowerCamelCase       )

             lowerCamelCase__		:				Tuple    =  self.get_trainer(callbacks=[MyTestTrainerCallback]       )
             trainer.train()
             lowerCamelCase__		:				str    =  trainer.callback_handler.callbacks[-2].events
             self.assertEqual(__lowerCamelCase					,      self.get_expected_events(__lowerCamelCase       )       )

             # Independent log/save/eval
             lowerCamelCase__		:				Tuple    =  self.get_trainer(callbacks=[MyTestTrainerCallback]					,      logging_steps=5       )
             trainer.train()
             lowerCamelCase__		:				int    =  trainer.callback_handler.callbacks[-2].events
             self.assertEqual(__lowerCamelCase					,      self.get_expected_events(__lowerCamelCase       )       )

             lowerCamelCase__		:				Any    =  self.get_trainer(callbacks=[MyTestTrainerCallback]					,      save_steps=5       )
             trainer.train()
             lowerCamelCase__		:				List[Any]    =  trainer.callback_handler.callbacks[-2].events
             self.assertEqual(__lowerCamelCase					,      self.get_expected_events(__lowerCamelCase       )       )

             lowerCamelCase__		:				List[Any]    =  self.get_trainer(callbacks=[MyTestTrainerCallback]					,      eval_steps=5					,      evaluation_strategy="steps"       )
             trainer.train()
             lowerCamelCase__		:				List[str]    =  trainer.callback_handler.callbacks[-2].events
             self.assertEqual(__lowerCamelCase					,      self.get_expected_events(__lowerCamelCase       )       )

             lowerCamelCase__		:				List[str]    =  self.get_trainer(callbacks=[MyTestTrainerCallback]					,      evaluation_strategy="epoch"       )
             trainer.train()
             lowerCamelCase__		:				int    =  trainer.callback_handler.callbacks[-2].events
             self.assertEqual(__lowerCamelCase					,      self.get_expected_events(__lowerCamelCase       )       )

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

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






from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
    from .deepspeed import (
        DeepSpeedEngineWrapper,
        DeepSpeedOptimizerWrapper,
        DeepSpeedSchedulerWrapper,
        DummyOptim,
        DummyScheduler,
        HfDeepSpeedConfig,
    )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers


