






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   lowerCamelCase				(_snake_case      ):







		'''simple docstring'''



		def __init__(							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				Dict						=						[]



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[int]:
							self.events.append('on_init_end'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							str:
							self.events.append('on_train_begin'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[Any]:
							self.events.append('on_train_end'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[Any]:
							self.events.append('on_epoch_begin'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[int]:
							self.events.append('on_epoch_end'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							int:
							self.events.append('on_step_begin'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							List[str]:
							self.events.append('on_step_end'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Union[str, Any]:
							self.events.append('on_evaluate'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Dict:
							self.events.append('on_predict'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							Optional[Any]:
							self.events.append('on_save'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							List[str]:
							self.events.append('on_log'       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		**_UpperCamelCase       )     ->							List[Any]:
							self.events.append('on_prediction_step'       )





@require_torch
class   lowerCamelCase				(unittest.TestCase      ):







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Union[str, Any]:
							UpperCAmelCase_   :				Tuple						=						tempfile.mkdtemp()



		def  __UpperCAmelCase      (							self       )     ->							int:
							shutil.rmtree(self.output_dir       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase=0    ,		_UpperCamelCase=0    ,		_UpperCamelCase=6_4    ,		_UpperCamelCase=6_4    ,		_UpperCamelCase=None    ,		_UpperCamelCase=False    ,		**_UpperCamelCase       )     ->							int:
							# 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.
							UpperCAmelCase_   :				Any						=						RegressionDataset(length=_UpperCamelCase       )
							UpperCAmelCase_   :				int						=						RegressionDataset(length=_UpperCamelCase       )
							UpperCAmelCase_   :				Union[str, Any]						=						RegressionModelConfig(a=_UpperCamelCase    ,		b=_UpperCamelCase       )
							UpperCAmelCase_   :				Optional[int]						=						RegressionPreTrainedModel(_UpperCamelCase       )

							UpperCAmelCase_   :				Optional[Any]						=						TrainingArguments(self.output_dir    ,		disable_tqdm=_UpperCamelCase    ,		report_to=[]    ,		**_UpperCamelCase       )
							return Trainer(
							    _UpperCamelCase    ,		_UpperCamelCase    ,		train_dataset=_UpperCamelCase    ,		eval_dataset=_UpperCamelCase    ,		callbacks=_UpperCamelCase    ,		)



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase       )     ->							Optional[int]:
							self.assertEqual(len(_UpperCamelCase       )    ,		len(_UpperCamelCase       )       )

							# Order doesn't matter
							UpperCAmelCase_   :				Union[str, Any]						=						sorted(_UpperCamelCase    ,		key=lambda _UpperCamelCase       : cb.__name__ if isinstance(_UpperCamelCase    ,		_UpperCamelCase       ) else cb.__class__.__name__       )
							UpperCAmelCase_   :				List[Any]						=						sorted(_UpperCamelCase    ,		key=lambda _UpperCamelCase       : cb.__name__ if isinstance(_UpperCamelCase    ,		_UpperCamelCase       ) else cb.__class__.__name__       )

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



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase       )     ->							Optional[int]:
							UpperCAmelCase_   :				Optional[Any]						=						['on_init_end', 'on_train_begin']
							UpperCAmelCase_   :				str						=						0
							UpperCAmelCase_   :				List[Any]						=						len(trainer.get_eval_dataloader()       )
							UpperCAmelCase_   :				int						=						['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(_UpperCamelCase       ):
																	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[Any]:
							UpperCAmelCase_   :				List[str]						=						self.get_trainer()
							UpperCAmelCase_   :				int						=						DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks    ,		_UpperCamelCase       )

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

							# TrainingArguments.disable_tqdm controls if use ProgressCallback or PrinterCallback
							UpperCAmelCase_   :				Tuple						=						self.get_trainer(disable_tqdm=_UpperCamelCase       )
							UpperCAmelCase_   :				Any						=						DEFAULT_CALLBACKS.copy() + [PrinterCallback]
							self.check_callbacks_equality(trainer.callback_handler.callbacks    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							UpperCAmelCase_   :				str						=						DEFAULT_CALLBACKS.copy() + [ProgressCallback]
							UpperCAmelCase_   :				str						=						self.get_trainer()

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

							UpperCAmelCase_   :				str						=						self.get_trainer()
							UpperCAmelCase_   :				List[Any]						=						trainer.pop_callback(_UpperCamelCase       )
							self.assertEqual(cb.__class__    ,		_UpperCamelCase       )
							self.check_callbacks_equality(trainer.callback_handler.callbacks    ,		_UpperCamelCase       )

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

							# We can also add, pop, or remove by instance
							UpperCAmelCase_   :				List[str]						=						self.get_trainer()
							UpperCAmelCase_   :				Any						=						trainer.callback_handler.callbacks[0]
							trainer.remove_callback(_UpperCamelCase       )
							expected_callbacks.remove(_UpperCamelCase       )
							self.check_callbacks_equality(trainer.callback_handler.callbacks    ,		_UpperCamelCase       )

							UpperCAmelCase_   :				List[str]						=						self.get_trainer()
							UpperCAmelCase_   :				Optional[int]						=						trainer.callback_handler.callbacks[0]
							UpperCAmelCase_   :				str						=						trainer.pop_callback(_UpperCamelCase       )
							self.assertEqual(_UpperCamelCase    ,		_UpperCamelCase       )
							self.check_callbacks_equality(trainer.callback_handler.callbacks    ,		_UpperCamelCase       )

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



		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							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=_UpperCamelCase       )

							UpperCAmelCase_   :				Any						=						self.get_trainer(callbacks=[MyTestTrainerCallback]       )
							trainer.train()
							UpperCAmelCase_   :				Union[str, Any]						=						trainer.callback_handler.callbacks[-2].events
							self.assertEqual(_UpperCamelCase    ,		self.get_expected_events(_UpperCamelCase       )       )

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

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

							UpperCAmelCase_   :				Dict						=						self.get_trainer(callbacks=[MyTestTrainerCallback]    ,		eval_steps=5    ,		evaluation_strategy='steps'       )
							trainer.train()
							UpperCAmelCase_   :				Dict						=						trainer.callback_handler.callbacks[-2].events
							self.assertEqual(_UpperCamelCase    ,		self.get_expected_events(_UpperCamelCase       )       )

							UpperCAmelCase_   :				Union[str, Any]						=						self.get_trainer(callbacks=[MyTestTrainerCallback]    ,		evaluation_strategy='epoch'       )
							trainer.train()
							UpperCAmelCase_   :				Dict						=						trainer.callback_handler.callbacks[-2].events
							self.assertEqual(_UpperCamelCase    ,		self.get_expected_events(_UpperCamelCase       )       )

							# A bit of everything
							UpperCAmelCase_   :				Any						=						self.get_trainer(
							    callbacks=[MyTestTrainerCallback]    ,		logging_steps=3    ,		save_steps=1_0    ,		eval_steps=5    ,		evaluation_strategy='steps'    ,		)
							trainer.train()
							UpperCAmelCase_   :				List[Any]						=						trainer.callback_handler.callbacks[-2].events
							self.assertEqual(_UpperCamelCase    ,		self.get_expected_events(_UpperCamelCase       )       )

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














import os
import re
import shutil
import sys
import tempfile
import unittest

import black


__UpperCAmelCase     							=						os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, 'utils'))

import check_copies  # noqa: E402


# This is the reference code that will be used in the tests.
# If DDPMSchedulerOutput is changed in scheduling_ddpm.py, this code needs to be manually updated.
__UpperCAmelCase     							=						'    \"""\n    Output class for the scheduler\'s step function output.\n\n    Args:\n        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the\n            denoising loop.\n        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):\n            The predicted denoised sample (x_{0}) based on the model output from the current timestep.\n            `pred_original_sample` can be used to preview progress or for guidance.\n    \"""\n\n    prev_sample: torch.FloatTensor\n    pred_original_sample: Optional[torch.FloatTensor] = None\n'



class   lowerCamelCase				(unittest.TestCase      ):







		'''simple docstring'''



		def  __UpperCAmelCase      (							self       )     ->							Tuple:
							UpperCAmelCase_   :				str						=						tempfile.mkdtemp()
							os.makedirs(os.path.join(self.diffusers_dir    ,		'schedulers/'       )       )
							UpperCAmelCase_   :				int						=						self.diffusers_dir
							shutil.copy(
							    os.path.join(_UpperCamelCase    ,		'src/diffusers/schedulers/scheduling_ddpm.py'       )    ,		os.path.join(self.diffusers_dir    ,		'schedulers/scheduling_ddpm.py'       )    ,		)



		def  __UpperCAmelCase      (							self       )     ->							Optional[int]:
							UpperCAmelCase_   :				str						=						'src/diffusers'
							shutil.rmtree(self.diffusers_dir       )



		def  __UpperCAmelCase      (							self    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase    ,		_UpperCamelCase=None       )     ->							Optional[Any]:
							UpperCAmelCase_   :				Dict						=						comment + f"\nclass {class_name}(nn.Module):\n" + class_code
							if overwrite_result is not None:
												UpperCAmelCase_   :				Optional[Any]						=						comment + f"\nclass {class_name}(nn.Module):\n" + overwrite_result
							UpperCAmelCase_   :				Any						=						black.Mode(target_versions={black.TargetVersion.PYaa}    ,		line_length=1_1_9       )
							UpperCAmelCase_   :				Optional[Any]						=						black.format_str(_UpperCamelCase    ,		mode=_UpperCamelCase       )
							UpperCAmelCase_   :				Any						=						os.path.join(self.diffusers_dir    ,		'new_code.py'       )
							with open(_UpperCamelCase    ,		'w'    ,		newline='\n'       ) as f:
												f.write(_UpperCamelCase       )
							if overwrite_result is None:
												self.assertTrue(len(check_copies.is_copy_consistent(_UpperCamelCase       )       ) == 0       )
							else:
												check_copies.is_copy_consistent(f.name    ,		overwrite=_UpperCamelCase       )
												with open(_UpperCamelCase    ,		'r'       ) as f:
																	self.assertTrue(f.read()    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							List[str]:
							UpperCAmelCase_   :				List[Any]						=						check_copies.find_code_in_diffusers('schedulers.scheduling_ddpm.DDPMSchedulerOutput'       )
							self.assertEqual(_UpperCamelCase    ,		_UpperCamelCase       )



		def  __UpperCAmelCase      (							self       )     ->							str:
							# Base copy consistency
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput'    ,		'DDPMSchedulerOutput'    ,		REFERENCE_CODE + '\n'    ,		)

							# With no empty line at the end
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput'    ,		'DDPMSchedulerOutput'    ,		_UpperCamelCase    ,		)

							# Copy consistency with rename
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test'    ,		'TestSchedulerOutput'    ,		re.sub('DDPM'    ,		'Test'    ,		_UpperCamelCase       )    ,		)

							# Copy consistency with a really long name
							UpperCAmelCase_   :				Optional[int]						=						'TestClassWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReason'
							self.check_copy_consistency(
							    f"# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->{long_class_name}"    ,		f"{long_class_name}SchedulerOutput"    ,		re.sub('Bert'    ,		_UpperCamelCase    ,		_UpperCamelCase       )    ,		)

							# Copy consistency with overwrite
							self.check_copy_consistency(
							    '# Copied from diffusers.schedulers.scheduling_ddpm.DDPMSchedulerOutput with DDPM->Test'    ,		'TestSchedulerOutput'    ,		_UpperCamelCase    ,		overwrite_result=re.sub('DDPM'    ,		'Test'    ,		_UpperCamelCase       )    ,		)






