
"""simple docstring"""





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 						SCREAMING_SNAKE_CASE_ (						__a ):




		"""simple docstring"""





		def __init__(     self):
				__SCREAMING_SNAKE_CASE            = []




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_init_end""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_train_begin""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_train_end""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_epoch_begin""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_epoch_end""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_step_begin""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_step_end""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_evaluate""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_predict""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_save""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_log""")




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__							,  lowerCAmelCase__							,  **lowerCAmelCase__):
				self.events.append("""on_prediction_step""")



@require_torch
class 						SCREAMING_SNAKE_CASE_ (						unittest.TestCase ):




		"""simple docstring"""





		def 	snake_case_     (     self):
				__SCREAMING_SNAKE_CASE            = tempfile.mkdtemp()




		def 	snake_case_     (     self):
				shutil.rmtree(self.output_dir)




		def 	snake_case_     (     self							,  lowerCAmelCase__=0							,  lowerCAmelCase__=0							,  lowerCAmelCase__=6_4							,  lowerCAmelCase__=6_4							,  lowerCAmelCase__=None							,  lowerCAmelCase__=False							,  **lowerCAmelCase__):
				# 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=lowerCAmelCase__)
				__SCREAMING_SNAKE_CASE            = RegressionDataset(length=lowerCAmelCase__)
				__SCREAMING_SNAKE_CASE            = RegressionModelConfig(a=lowerCAmelCase__							,  b=lowerCAmelCase__)
				__SCREAMING_SNAKE_CASE            = RegressionPreTrainedModel(lowerCAmelCase__)

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




		def 	snake_case_     (     self							,  lowerCAmelCase__							,  lowerCAmelCase__):
				self.assertEqual(len(lowerCAmelCase__)							,  len(lowerCAmelCase__))

				# Order doesn't matter
				__SCREAMING_SNAKE_CASE            = sorted(lowerCAmelCase__							,  key=lambda lowerCAmelCase__: cb.__name__ if isinstance(lowerCAmelCase__							,  lowerCAmelCase__) else cb.__class__.__name__)
				__SCREAMING_SNAKE_CASE            = 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 	snake_case_     (     self							,  lowerCAmelCase__):
				__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(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 	snake_case_     (     self):
				__SCREAMING_SNAKE_CASE            = self.get_trainer()
				__SCREAMING_SNAKE_CASE            = DEFAULT_CALLBACKS.copy() + [ProgressCallback]
				self.check_callbacks_equality(trainer.callback_handler.callbacks							,  lowerCAmelCase__)

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




		def 	snake_case_     (     self):
				__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(lowerCAmelCase__)
				expected_callbacks.remove(lowerCAmelCase__)
				self.check_callbacks_equality(trainer.callback_handler.callbacks							,  lowerCAmelCase__)

				__SCREAMING_SNAKE_CASE            = self.get_trainer()
				__SCREAMING_SNAKE_CASE            = 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
				__SCREAMING_SNAKE_CASE            = self.get_trainer()
				__SCREAMING_SNAKE_CASE            = trainer.callback_handler.callbacks[0]
				trainer.remove_callback(lowerCAmelCase__)
				expected_callbacks.remove(lowerCAmelCase__)
				self.check_callbacks_equality(trainer.callback_handler.callbacks							,  lowerCAmelCase__)

				__SCREAMING_SNAKE_CASE            = self.get_trainer()
				__SCREAMING_SNAKE_CASE            = trainer.callback_handler.callbacks[0]
				__SCREAMING_SNAKE_CASE            = 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 	snake_case_     (     self):
				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__)

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

				# 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(lowerCAmelCase__							,  self.get_expected_events(lowerCAmelCase__))

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

				__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(lowerCAmelCase__							,  self.get_expected_events(lowerCAmelCase__))

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

				# A bit of everything
				__SCREAMING_SNAKE_CASE            = self.get_trainer(
				    callbacks=[MyTestTrainerCallback]							,  logging_steps=3							,  save_steps=1_0							,  eval_steps=5							,  evaluation_strategy="""steps"""							,  )
				trainer.train()
				__SCREAMING_SNAKE_CASE            = 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:
						__SCREAMING_SNAKE_CASE            = self.get_trainer(
						    callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]							,  )
						assert str(lowerCAmelCase__) in warn_mock.call_args[0][0]






"""simple docstring"""





import pytest

from datasets.utils.sharding import _distribute_shards, _number_of_shards_in_gen_kwargs, _split_gen_kwargs



@pytest.mark.parametrize(
    """kwargs, expected"""  ,				[
        ({"""num_shards""": 0, """max_num_jobs""": 1}, []),
        ({"""num_shards""": 10, """max_num_jobs""": 1}, [range(10				)]),
        ({"""num_shards""": 10, """max_num_jobs""": 10}, [range(UpperCamelCase_  ,				i + 1				) for i in range(10				)]),
        ({"""num_shards""": 1, """max_num_jobs""": 10}, [range(1				)]),
        ({"""num_shards""": 10, """max_num_jobs""": 3}, [range(0  ,				4				), range(4  ,				7				), range(7  ,				10				)]),
        ({"""num_shards""": 3, """max_num_jobs""": 10}, [range(0  ,				1				), range(1  ,				2				), range(2  ,				3				)]),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = _distribute_shards(**UpperCamelCase_				)
		assert out == expected



@pytest.mark.parametrize(
    """gen_kwargs, max_num_jobs, expected"""  ,				[
        ({"""foo""": 0}, 10, [{"""foo""": 0}]),
        ({"""shards""": [0, 1, 2, 3]}, 1, [{"""shards""": [0, 1, 2, 3]}]),
        ({"""shards""": [0, 1, 2, 3]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}, {"""shards""": [2]}, {"""shards""": [3]}]),
        ({"""shards""": [0, 1]}, 4, [{"""shards""": [0]}, {"""shards""": [1]}]),
        ({"""shards""": [0, 1, 2, 3]}, 2, [{"""shards""": [0, 1]}, {"""shards""": [2, 3]}]),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_  ,				UpperCamelCase_				):
		__SCREAMING_SNAKE_CASE            = _split_gen_kwargs(UpperCamelCase_  ,				UpperCamelCase_				)
		assert out == expected



@pytest.mark.parametrize(
    """gen_kwargs, expected"""  ,				[
        ({"""foo""": 0}, 1),
        ({"""shards""": [0]}, 1),
        ({"""shards""": [0, 1, 2, 3]}, 4),
        ({"""shards""": [0, 1, 2, 3], """foo""": 0}, 4),
        ({"""shards""": [0, 1, 2, 3], """other""": (0, 1)}, 4),
        ({"""shards""": [0, 1, 2, 3], """shards2""": [0, 1]}, RuntimeError),
    ]  ,				)
def 						_lowerCAmelCase				(				UpperCamelCase_  ,				UpperCamelCase_				):
		if expected is RuntimeError:
				with pytest.raises(UpperCamelCase_				):
						_number_of_shards_in_gen_kwargs(UpperCamelCase_				)
		else:
				__SCREAMING_SNAKE_CASE            = _number_of_shards_in_gen_kwargs(UpperCamelCase_				)
				assert out == expected




