






'''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	(      lowerCAmelCase_  ):




       def __init__(					self      : Dict					):


            """simple docstring"""







            __lowerCamelCase						:       Any				   =				[]




       def 					a_					(					self      : Optional[Any]	,    A__      : Dict	,    A__      : Optional[int]	,    A__      : Tuple	,    **A__      : Tuple					):


            """simple docstring"""







            self.events.append("""on_init_end"""					)




       def 					a_					(					self      : str	,    A__      : List[Any]	,    A__      : Tuple	,    A__      : Tuple	,    **A__      : Optional[Any]					):


            """simple docstring"""







            self.events.append("""on_train_begin"""					)




       def 					a_					(					self      : Tuple	,    A__      : List[Any]	,    A__      : List[str]	,    A__      : List[str]	,    **A__      : Tuple					):


            """simple docstring"""







            self.events.append("""on_train_end"""					)




       def 					a_					(					self      : Union[str, Any]	,    A__      : str	,    A__      : str	,    A__      : Optional[Any]	,    **A__      : Dict					):


            """simple docstring"""







            self.events.append("""on_epoch_begin"""					)




       def 					a_					(					self      : List[Any]	,    A__      : List[str]	,    A__      : List[Any]	,    A__      : Dict	,    **A__      : List[str]					):


            """simple docstring"""







            self.events.append("""on_epoch_end"""					)




       def 					a_					(					self      : int	,    A__      : Optional[Any]	,    A__      : Any	,    A__      : Any	,    **A__      : int					):


            """simple docstring"""







            self.events.append("""on_step_begin"""					)




       def 					a_					(					self      : Any	,    A__      : int	,    A__      : Any	,    A__      : Tuple	,    **A__      : int					):


            """simple docstring"""







            self.events.append("""on_step_end"""					)




       def 					a_					(					self      : str	,    A__      : int	,    A__      : Optional[Any]	,    A__      : str	,    **A__      : Optional[Any]					):


            """simple docstring"""







            self.events.append("""on_evaluate"""					)




       def 					a_					(					self      : str	,    A__      : str	,    A__      : Union[str, Any]	,    A__      : Any	,    **A__      : int					):


            """simple docstring"""







            self.events.append("""on_predict"""					)




       def 					a_					(					self      : str	,    A__      : str	,    A__      : int	,    A__      : List[Any]	,    **A__      : Dict					):


            """simple docstring"""







            self.events.append("""on_save"""					)




       def 					a_					(					self      : List[str]	,    A__      : str	,    A__      : List[str]	,    A__      : Any	,    **A__      : Dict					):


            """simple docstring"""







            self.events.append("""on_log"""					)




       def 					a_					(					self      : int	,    A__      : List[Any]	,    A__      : List[str]	,    A__      : Any	,    **A__      : Union[str, Any]					):


            """simple docstring"""







            self.events.append("""on_prediction_step"""					)







@require_torch
class 	SCREAMING_SNAKE_CASE	(      unittest.TestCase  ):




       def 					a_					(					self      : Optional[int]					):


            """simple docstring"""







            __lowerCamelCase						:       List[Any]				   =				tempfile.mkdtemp()




       def 					a_					(					self      : Union[str, Any]					):


            """simple docstring"""







            shutil.rmtree(self.output_dir					)




       def 					a_					(					self      : Optional[Any]	,    A__      : Any=0	,    A__      : Any=0	,    A__      : Dict=64	,    A__      : Dict=64	,    A__      : List[str]=None	,    A__      : Union[str, Any]=False	,    **A__      : Tuple					):


            """simple docstring"""







            __lowerCamelCase						:       Optional[int]				   =				RegressionDataset(length=A__					)
            __lowerCamelCase						:       Optional[int]				   =				RegressionDataset(length=A__					)
            __lowerCamelCase						:       Any				   =				RegressionModelConfig(a=A__	,    b=A__					)
            __lowerCamelCase						:       Optional[Any]				   =				RegressionPreTrainedModel(A__					)

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




       def 					a_					(					self      : List[str]	,    A__      : Optional[int]	,    A__      : Dict					):


            """simple docstring"""







            self.assertEqual(len(A__					)	,    len(A__					)					)

            # Order doesn't matter
            __lowerCamelCase						:       List[Any]				   =				sorted(A__	,    key=lambda A__					: cb.__name__ if isinstance(A__	,    A__					) else cb.__class__.__name__					)
            __lowerCamelCase						:       int				   =				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 					a_					(					self      : Optional[int]	,    A__      : List[Any]					):


            """simple docstring"""







            __lowerCamelCase						:       Tuple				   =				["""on_init_end""", """on_train_begin"""]
            __lowerCamelCase						:       str				   =				0
            __lowerCamelCase						:       List[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(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 					a_					(					self      : Union[str, Any]					):


            """simple docstring"""







            __lowerCamelCase						:       int				   =				self.get_trainer()
            __lowerCamelCase						:       List[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						:       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						:       List[Any]				   =				self.get_trainer(disable_tqdm=A__					)
            __lowerCamelCase						:       int				   =				DEFAULT_CALLBACKS.copy() + [PrinterCallback]
            self.check_callbacks_equality(trainer.callback_handler.callbacks	,    A__					)




       def 					a_					(					self      : List[str]					):


            """simple docstring"""







            __lowerCamelCase						:       List[Any]				   =				DEFAULT_CALLBACKS.copy() + [ProgressCallback]
            __lowerCamelCase						:       Optional[int]				   =				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						:       List[Any]				   =				self.get_trainer()
            __lowerCamelCase						:       Any				   =				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						:       Optional[int]				   =				self.get_trainer()
            __lowerCamelCase						:       Union[str, Any]				   =				trainer.callback_handler.callbacks[0]
            trainer.remove_callback(A__					)
            expected_callbacks.remove(A__					)
            self.check_callbacks_equality(trainer.callback_handler.callbacks	,    A__					)

            __lowerCamelCase						:       int				   =				self.get_trainer()
            __lowerCamelCase						:       Dict				   =				trainer.callback_handler.callbacks[0]
            __lowerCamelCase						:       Dict				   =				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 					a_					(					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=A__					)

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

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

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

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

            __lowerCamelCase						:       Any				   =				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						:       Any				   =				self.get_trainer(
                callbacks=[MyTestTrainerCallback]	,    logging_steps=3	,    save_steps=10	,    eval_steps=5	,    evaluation_strategy="""steps"""	,    )
            trainer.train()
            __lowerCamelCase						:       Tuple				   =				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						:       Union[str, Any]				   =				self.get_trainer(
                     callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]	,    )
                 assert str(A__					) in warn_mock.call_args[0][0]











'''simple docstring'''

from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)
def     __lowercase  (_lowercase,						_lowercase     )   ->     List[str]:


     """simple docstring"""
     __lowerCamelCase						:       Union[str, Any]				   =				[]
     for part_id in partition_order:
          __lowerCamelCase						:       Union[str, Any]				   =				df.where(f"SPARK_PARTITION_ID() = {part_id}"     ).collect()
          for row_idx, row in enumerate(_lowercase     ):
               expected_row_ids_and_row_dicts.append((f"{part_id}_{row_idx}", row.asDict())     )
     return expected_row_ids_and_row_dicts
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     Optional[int]:


     """simple docstring"""
     __lowerCamelCase						:       int				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       List[str]				   =				spark.range(100     ).repartition(1     )
     __lowerCamelCase						:       str				   =				Spark(_lowercase     )
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16     )
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     Optional[Any]:


     """simple docstring"""
     __lowerCamelCase						:       Any				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       Dict				   =				spark.range(10     ).repartition(2     )
     __lowerCamelCase						:       List[Any]				   =				[1, 0]
     __lowerCamelCase						:       Any				   =				_generate_iterable_examples(_lowercase,						_lowercase     )  # Reverse the partitions.
     __lowerCamelCase						:       int				   =				_get_expected_row_ids_and_row_dicts_for_partition_order(_lowercase,						_lowercase     )

     for i, (row_id, row_dict) in enumerate(generate_fn()     ):
          __lowerCamelCase   ,      __lowerCamelCase						:       str				   =				expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     int:


     """simple docstring"""
     __lowerCamelCase						:       Dict				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       Union[str, Any]				   =				spark.range(10     ).repartition(1     )
     __lowerCamelCase						:       Dict				   =				SparkExamplesIterable(_lowercase     )
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(_lowercase     ):
          assert row_id == f"0_{i}"
          assert row_dict == {"id": i}
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     Tuple:


     """simple docstring"""
     __lowerCamelCase						:       int				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       Optional[Any]				   =				spark.range(30     ).repartition(3     )
     # Mock the generator so that shuffle reverses the partition indices.
     with patch("""numpy.random.Generator"""     ) as generator_mock:
          __lowerCamelCase						:       Dict				   =				lambda _lowercase     : x.reverse()
          __lowerCamelCase						:       Union[str, Any]				   =				_get_expected_row_ids_and_row_dicts_for_partition_order(_lowercase,						[2, 1, 0]     )

          __lowerCamelCase						:       Tuple				   =				SparkExamplesIterable(_lowercase     ).shuffle_data_sources(_lowercase     )
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(_lowercase     ):
               __lowerCamelCase   ,      __lowerCamelCase						:       Union[str, Any]				   =				expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     str:


     """simple docstring"""
     __lowerCamelCase						:       Union[str, Any]				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       List[str]				   =				spark.range(20     ).repartition(4     )

     # Partitions 0 and 2
     __lowerCamelCase						:       Union[str, Any]				   =				SparkExamplesIterable(_lowercase     ).shard_data_sources(worker_id=0,						num_workers=2     )
     assert shard_it_a.n_shards == 2
     __lowerCamelCase						:       Optional[int]				   =				_get_expected_row_ids_and_row_dicts_for_partition_order(_lowercase,						[0, 2]     )
     for i, (row_id, row_dict) in enumerate(_lowercase     ):
          __lowerCamelCase   ,      __lowerCamelCase						:       int				   =				expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     __lowerCamelCase						:       List[str]				   =				SparkExamplesIterable(_lowercase     ).shard_data_sources(worker_id=1,						num_workers=2     )
     assert shard_it_a.n_shards == 2
     __lowerCamelCase						:       int				   =				_get_expected_row_ids_and_row_dicts_for_partition_order(_lowercase,						[1, 3]     )
     for i, (row_id, row_dict) in enumerate(_lowercase     ):
          __lowerCamelCase   ,      __lowerCamelCase						:       int				   =				expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict
@require_not_windows
@require_dill_gt_0_3_2
def     __lowercase  ()   ->     int:


     """simple docstring"""
     __lowerCamelCase						:       Union[str, Any]				   =				pyspark.sql.SparkSession.builder.master("""local[*]"""     ).appName("""pyspark"""     ).getOrCreate()
     __lowerCamelCase						:       Dict				   =				spark.range(100     ).repartition(1     )
     __lowerCamelCase						:       Any				   =				Spark(_lowercase     )
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1     )
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 100



