



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):







							def __init__(   self):

														'''simple docstring'''

														_lowerCAmelCase     :							List[str]    =    []







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_init_end")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_train_begin")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_train_end")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_epoch_begin")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_epoch_end")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_step_begin")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_step_end")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_evaluate")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_predict")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_save")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_log")







							def    snake_case__      (   self,			__a,			__a,			__a,			**__a):

														'''simple docstring'''

														self.events.append("on_prediction_step")



@require_torch
class   UpperCAmelCase_    (       unittest.TestCase):







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    tempfile.mkdtemp()







							def    snake_case__      (   self):

														'''simple docstring'''

														shutil.rmtree(self.output_dir)







							def    snake_case__      (   self,			__a=0,			__a=0,			__a=64,			__a=64,			__a=None,			__a=False,			**__a):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[int]    =    RegressionDataset(length=__a)
														_lowerCAmelCase     :							Union[str, Any]    =    RegressionDataset(length=__a)
														_lowerCAmelCase     :							List[str]    =    RegressionModelConfig(a=__a,			b=__a)
														_lowerCAmelCase     :							str    =    RegressionPreTrainedModel(__a)

														_lowerCAmelCase     :							str    =    TrainingArguments(self.output_dir,			disable_tqdm=__a,			report_to=[],			**__a)
														return Trainer(
														    __a,			__a,			train_dataset=__a,			eval_dataset=__a,			callbacks=__a,			)







							def    snake_case__      (   self,			__a,			__a):

														'''simple docstring'''

														self.assertEqual(len(__a),			len(__a))

														# Order doesn't matter
														_lowerCAmelCase     :							Any    =    sorted(__a,			key=lambda __a: cb.__name__ if isinstance(__a,			__a) else cb.__class__.__name__)
														_lowerCAmelCase     :							str    =    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    snake_case__      (   self,			__a):

														'''simple docstring'''

														_lowerCAmelCase     :							Optional[Any]    =    ["on_init_end", "on_train_begin"]
														_lowerCAmelCase     :							str    =    0
														_lowerCAmelCase     :							List[Any]    =    len(trainer.get_eval_dataloader())
														_lowerCAmelCase     :							List[str]    =    ["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    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							int    =    self.get_trainer()
														_lowerCAmelCase     :							str    =    DEFAULT_CALLBACKS.copy() + [ProgressCallback]
														self.check_callbacks_equality(trainer.callback_handler.callbacks,			__a)

														# Callbacks passed at init are added to the default callbacks
														_lowerCAmelCase     :							List[Any]    =    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
														_lowerCAmelCase     :							Dict    =    self.get_trainer(disable_tqdm=__a)
														_lowerCAmelCase     :							List[Any]    =    DEFAULT_CALLBACKS.copy() + [PrinterCallback]
														self.check_callbacks_equality(trainer.callback_handler.callbacks,			__a)







							def    snake_case__      (   self):

														'''simple docstring'''

														_lowerCAmelCase     :							Dict    =    DEFAULT_CALLBACKS.copy() + [ProgressCallback]
														_lowerCAmelCase     :							Dict    =    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)

														_lowerCAmelCase     :							Tuple    =    self.get_trainer()
														_lowerCAmelCase     :							Dict    =    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
														_lowerCAmelCase     :							List[str]    =    self.get_trainer()
														_lowerCAmelCase     :							int    =    trainer.callback_handler.callbacks[0]
														trainer.remove_callback(__a)
														expected_callbacks.remove(__a)
														self.check_callbacks_equality(trainer.callback_handler.callbacks,			__a)

														_lowerCAmelCase     :							Optional[int]    =    self.get_trainer()
														_lowerCAmelCase     :							Optional[Any]    =    trainer.callback_handler.callbacks[0]
														_lowerCAmelCase     :							int    =    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    snake_case__      (   self):

														'''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=__a)

														_lowerCAmelCase     :							Tuple    =    self.get_trainer(callbacks=[MyTestTrainerCallback])
														trainer.train()
														_lowerCAmelCase     :							List[Any]    =    trainer.callback_handler.callbacks[-2].events
														self.assertEqual(__a,			self.get_expected_events(__a))

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

														_lowerCAmelCase     :							List[str]    =    self.get_trainer(callbacks=[MyTestTrainerCallback],			save_steps=5)
														trainer.train()
														_lowerCAmelCase     :							List[str]    =    trainer.callback_handler.callbacks[-2].events
														self.assertEqual(__a,			self.get_expected_events(__a))

														_lowerCAmelCase     :							Dict    =    self.get_trainer(callbacks=[MyTestTrainerCallback],			eval_steps=5,			evaluation_strategy="steps")
														trainer.train()
														_lowerCAmelCase     :							Tuple    =    trainer.callback_handler.callbacks[-2].events
														self.assertEqual(__a,			self.get_expected_events(__a))

														_lowerCAmelCase     :							str    =    self.get_trainer(callbacks=[MyTestTrainerCallback],			evaluation_strategy="epoch")
														trainer.train()
														_lowerCAmelCase     :							Any    =    trainer.callback_handler.callbacks[-2].events
														self.assertEqual(__a,			self.get_expected_events(__a))

														# A bit of everything
														_lowerCAmelCase     :							int    =    self.get_trainer(
														    callbacks=[MyTestTrainerCallback],			logging_steps=3,			save_steps=10,			eval_steps=5,			evaluation_strategy="steps",			)
														trainer.train()
														_lowerCAmelCase     :							Union[str, Any]    =    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:
																					_lowerCAmelCase     :							Optional[int]    =    self.get_trainer(
																					    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback],			)
																					assert str(__a) in warn_mock.call_args[0][0]




"""simple docstring"""
from ...utils import is_torch_available, is_transformers_available


if is_transformers_available() and is_torch_available():
			from .pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings, VQDiffusionPipeline





