

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




  def __init__(							self   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=[]




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					List[str]:
        '''simple docstring'''






        self.events.append('on_init_end'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Tuple:
        '''simple docstring'''






        self.events.append('on_train_begin'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					int:
        '''simple docstring'''






        self.events.append('on_train_end'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Union[str, Any]:
        '''simple docstring'''






        self.events.append('on_epoch_begin'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					List[Any]:
        '''simple docstring'''






        self.events.append('on_epoch_end'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Union[str, Any]:
        '''simple docstring'''






        self.events.append('on_step_begin'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Tuple:
        '''simple docstring'''






        self.events.append('on_step_end'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					List[Any]:
        '''simple docstring'''






        self.events.append('on_evaluate'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Union[str, Any]:
        '''simple docstring'''






        self.events.append('on_predict'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					List[Any]:
        '''simple docstring'''






        self.events.append('on_save'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Any:
        '''simple docstring'''






        self.events.append('on_log'   )




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase				,  lowerCAmelCase				,  **lowerCAmelCase   )    ->					Optional[Any]:
        '''simple docstring'''






        self.events.append('on_prediction_step'   )





@require_torch
class 					__lowerCAmelCase	(   unittest.TestCase				):




  def     A__	(							self   )    ->					Union[str, Any]:
        '''simple docstring'''






        _lowercase      			=tempfile.mkdtemp()




  def     A__	(							self   )    ->					str:
        '''simple docstring'''






        shutil.rmtree(self.output_dir   )




  def     A__	(							self				,  lowerCAmelCase=0				,  lowerCAmelCase=0				,  lowerCAmelCase=64				,  lowerCAmelCase=64				,  lowerCAmelCase=None				,  lowerCAmelCase=False				,  **lowerCAmelCase   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=RegressionDataset(length=lowerCAmelCase   )
        _lowercase      			=RegressionDataset(length=lowerCAmelCase   )
        _lowercase      			=RegressionModelConfig(a=lowerCAmelCase				,  b=lowerCAmelCase   )
        _lowercase      			=RegressionPreTrainedModel(lowerCAmelCase   )

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




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					List[Any]:
        '''simple docstring'''






        self.assertEqual(len(lowerCAmelCase   )				,  len(lowerCAmelCase   )   )

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






        _lowercase      			=['on_init_end', 'on_train_begin']
        _lowercase      			=0
        _lowercase      			=len(trainer.get_eval_dataloader()   )
        _lowercase      			=['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     A__	(							self   )    ->					List[str]:
        '''simple docstring'''






        _lowercase      			=self.get_trainer()
        _lowercase      			=DEFAULT_CALLBACKS.copy() + [ProgressCallback]
        self.check_callbacks_equality(trainer.callback_handler.callbacks				,  lowerCAmelCase   )

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




  def     A__	(							self   )    ->					Union[str, Any]:
        '''simple docstring'''






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

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

        _lowercase      			=self.get_trainer()
        _lowercase      			=trainer.callback_handler.callbacks[0]
        _lowercase      			=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     A__	(							self   )    ->					Tuple:
        '''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   )

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

        # Independent log/save/eval
        _lowercase      			=self.get_trainer(callbacks=[MyTestTrainerCallback]				,  logging_steps=5   )
        trainer.train()
        _lowercase      			=trainer.callback_handler.callbacks[-2].events
        self.assertEqual(lowerCAmelCase				,  self.get_expected_events(lowerCAmelCase   )   )

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

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

        _lowercase      			=self.get_trainer(callbacks=[MyTestTrainerCallback]				,  evaluation_strategy='epoch'   )
        trainer.train()
        _lowercase      			=trainer.callback_handler.callbacks[-2].events
        self.assertEqual(lowerCAmelCase				,  self.get_expected_events(lowerCAmelCase   )   )

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





from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import BaseOutput, is_torch_available, is_transformers_available





@dataclass
class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE				):
  _a		     =    42
  _a		     =    42


if is_transformers_available() and is_torch_available():
  from .pipeline_semantic_stable_diffusion import SemanticStableDiffusionPipeline


