
import argparse
import os

import torch

from diffusers import (
    CMStochasticIterativeScheduler,
    ConsistencyModelPipeline,
    UNetaDModel,
)


lowerCamelCase__     :		Optional[Any]        = {
    """sample_size""": 3_2,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": 1_0_0_0,
    """block_out_channels""": [3_2, 6_4],
    """attention_head_dim""": 8,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

lowerCamelCase__     :		Union[str, Any]        = {
    """sample_size""": 6_4,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 3,
    """num_class_embeds""": 1_0_0_0,
    """block_out_channels""": [1_9_2, 1_9_2 * 2, 1_9_2 * 3, 1_9_2 * 4],
    """attention_head_dim""": 6_4,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """scale_shift""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

lowerCamelCase__     :		List[str]        = {
    """sample_size""": 2_5_6,
    """in_channels""": 3,
    """out_channels""": 3,
    """layers_per_block""": 2,
    """num_class_embeds""": None,
    """block_out_channels""": [2_5_6, 2_5_6, 2_5_6 * 2, 2_5_6 * 2, 2_5_6 * 4, 2_5_6 * 4],
    """attention_head_dim""": 6_4,
    """down_block_types""": [
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """ResnetDownsampleBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
        """AttnDownBlock2D""",
    ],
    """up_block_types""": [
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """AttnUpBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
        """ResnetUpsampleBlock2D""",
    ],
    """resnet_time_scale_shift""": """default""",
    """upsample_type""": """resnet""",
    """downsample_type""": """resnet""",
}

lowerCamelCase__     :		Any        = {
    """num_train_timesteps""": 4_0,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}

lowerCamelCase__     :		int        = {
    """num_train_timesteps""": 2_0_1,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}

lowerCamelCase__     :		List[Any]        = {
    """num_train_timesteps""": 1_5_1,
    """sigma_min""": 0.002,
    """sigma_max""": 80.0,
}

def 		UpperCamelCase						(   lowercase_ )		->						Tuple:


  '''simple docstring'''






  if isinstance(lowercase_			,			lowercase_ ):
    return v
  if v.lower() in ("yes", "true", "t", "y", "1"):
    return True
  elif v.lower() in ("no", "false", "f", "n", "0"):
    return False
  else:
    raise argparse.ArgumentTypeError("""boolean value expected""" )

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=False )		->						Any:


  '''simple docstring'''






  lowercase__	: List[str]     		=   checkpoint[F'{old_prefix}.in_layers.0.weight']
  lowercase__	: Tuple     		=   checkpoint[F'{old_prefix}.in_layers.0.bias']
  lowercase__	: Any     		=   checkpoint[F'{old_prefix}.in_layers.2.weight']
  lowercase__	: Union[str, Any]     		=   checkpoint[F'{old_prefix}.in_layers.2.bias']
  lowercase__	: Optional[int]     		=   checkpoint[F'{old_prefix}.emb_layers.1.weight']
  lowercase__	: List[Any]     		=   checkpoint[F'{old_prefix}.emb_layers.1.bias']
  lowercase__	: List[Any]     		=   checkpoint[F'{old_prefix}.out_layers.0.weight']
  lowercase__	: Optional[int]     		=   checkpoint[F'{old_prefix}.out_layers.0.bias']
  lowercase__	: Union[str, Any]     		=   checkpoint[F'{old_prefix}.out_layers.3.weight']
  lowercase__	: Tuple     		=   checkpoint[F'{old_prefix}.out_layers.3.bias']

  if has_skip:
    lowercase__	: Any     		=   checkpoint[F'{old_prefix}.skip_connection.weight']
    lowercase__	: Any     		=   checkpoint[F'{old_prefix}.skip_connection.bias']

  return new_checkpoint

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_=None )		->						List[str]:


  '''simple docstring'''






  lowercase__		,     lowercase__		,     lowercase__	: str     		=   checkpoint[F'{old_prefix}.qkv.weight'].chunk(3			,			dim=0 )
  lowercase__		,     lowercase__		,     lowercase__	: Any     		=   checkpoint[F'{old_prefix}.qkv.bias'].chunk(3			,			dim=0 )

  lowercase__	: Optional[Any]     		=   checkpoint[F'{old_prefix}.norm.weight']
  lowercase__	: Tuple     		=   checkpoint[F'{old_prefix}.norm.bias']

  lowercase__	: int     		=   weight_q.squeeze(-1 ).squeeze(-1 )
  lowercase__	: int     		=   bias_q.squeeze(-1 ).squeeze(-1 )
  lowercase__	: str     		=   weight_k.squeeze(-1 ).squeeze(-1 )
  lowercase__	: Dict     		=   bias_k.squeeze(-1 ).squeeze(-1 )
  lowercase__	: List[Any]     		=   weight_v.squeeze(-1 ).squeeze(-1 )
  lowercase__	: Dict     		=   bias_v.squeeze(-1 ).squeeze(-1 )

  lowercase__	: List[Any]     		=   (
      checkpoint[F'{old_prefix}.proj_out.weight'].squeeze(-1 ).squeeze(-1 )
  )
  lowercase__	: List[str]     		=   checkpoint[F'{old_prefix}.proj_out.bias'].squeeze(-1 ).squeeze(-1 )

  return new_checkpoint

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						Tuple:


  '''simple docstring'''






  lowercase__	: List[Any]     		=   torch.load(lowercase_			,			map_location="""cpu""" )
  lowercase__	: Optional[Any]     		=   {}

  lowercase__	: Tuple     		=   checkpoint["""time_embed.0.weight"""]
  lowercase__	: Any     		=   checkpoint["""time_embed.0.bias"""]
  lowercase__	: Dict     		=   checkpoint["""time_embed.2.weight"""]
  lowercase__	: str     		=   checkpoint["""time_embed.2.bias"""]

  if unet_config["num_class_embeds"] is not None:
    lowercase__	: int     		=   checkpoint["""label_emb.weight"""]

  lowercase__	: Optional[int]     		=   checkpoint["""input_blocks.0.0.weight"""]
  lowercase__	: Union[str, Any]     		=   checkpoint["""input_blocks.0.0.bias"""]

  lowercase__	: Union[str, Any]     		=   unet_config["""down_block_types"""]
  lowercase__	: str     		=   unet_config["""layers_per_block"""]
  lowercase__	: Optional[int]     		=   unet_config["""attention_head_dim"""]
  lowercase__	: Optional[Any]     		=   unet_config["""block_out_channels"""]
  lowercase__	: Tuple     		=   1
  lowercase__	: Optional[int]     		=   channels_list[0]

  for i, layer_type in enumerate(lowercase_ ):
    lowercase__	: Any     		=   channels_list[i]
    lowercase__	: List[Any]     		=   current_channels != prev_channels
    if layer_type == "ResnetDownsampleBlock2D":
      for j in range(lowercase_ ):
        lowercase__	: Any     		=   F'down_blocks.{i}.resnets.{j}'
        lowercase__	: int     		=   F'input_blocks.{current_layer}.0'
        lowercase__	: Tuple     		=   True if j == 0 and downsample_block_has_skip else False
        lowercase__	: Tuple     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			has_skip=lowercase_ )
        current_layer += 1

    elif layer_type == "AttnDownBlock2D":
      for j in range(lowercase_ ):
        lowercase__	: Optional[Any]     		=   F'down_blocks.{i}.resnets.{j}'
        lowercase__	: Union[str, Any]     		=   F'input_blocks.{current_layer}.0'
        lowercase__	: Any     		=   True if j == 0 and downsample_block_has_skip else False
        lowercase__	: Optional[int]     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			has_skip=lowercase_ )
        lowercase__	: Optional[int]     		=   F'down_blocks.{i}.attentions.{j}'
        lowercase__	: int     		=   F'input_blocks.{current_layer}.1'
        lowercase__	: str     		=   convert_attention(
            lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
        current_layer += 1

    if i != len(lowercase_ ) - 1:
      lowercase__	: Dict     		=   F'down_blocks.{i}.downsamplers.0'
      lowercase__	: Union[str, Any]     		=   F'input_blocks.{current_layer}.0'
      lowercase__	: List[str]     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
      current_layer += 1

    lowercase__	: Optional[Any]     		=   current_channels

  # hardcoded the mid-block for now
  lowercase__	: Any     		=   """mid_block.resnets.0"""
  lowercase__	: List[Any]     		=   """middle_block.0"""
  lowercase__	: Tuple     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
  lowercase__	: str     		=   """mid_block.attentions.0"""
  lowercase__	: Union[str, Any]     		=   """middle_block.1"""
  lowercase__	: Tuple     		=   convert_attention(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
  lowercase__	: Dict     		=   """mid_block.resnets.1"""
  lowercase__	: List[Any]     		=   """middle_block.2"""
  lowercase__	: int     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )

  lowercase__	: str     		=   0
  lowercase__	: Union[str, Any]     		=   unet_config["""up_block_types"""]

  for i, layer_type in enumerate(lowercase_ ):
    if layer_type == "ResnetUpsampleBlock2D":
      for j in range(layers_per_block + 1 ):
        lowercase__	: List[str]     		=   F'up_blocks.{i}.resnets.{j}'
        lowercase__	: Dict     		=   F'output_blocks.{current_layer}.0'
        lowercase__	: str     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			has_skip=lowercase_ )
        current_layer += 1

      if i != len(lowercase_ ) - 1:
        lowercase__	: List[Any]     		=   F'up_blocks.{i}.upsamplers.0'
        lowercase__	: Any     		=   F'output_blocks.{current_layer-1}.1'
        lowercase__	: List[str]     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
    elif layer_type == "AttnUpBlock2D":
      for j in range(layers_per_block + 1 ):
        lowercase__	: Dict     		=   F'up_blocks.{i}.resnets.{j}'
        lowercase__	: Any     		=   F'output_blocks.{current_layer}.0'
        lowercase__	: List[str]     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			has_skip=lowercase_ )
        lowercase__	: Tuple     		=   F'up_blocks.{i}.attentions.{j}'
        lowercase__	: Optional[Any]     		=   F'output_blocks.{current_layer}.1'
        lowercase__	: Any     		=   convert_attention(
            lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )
        current_layer += 1

      if i != len(lowercase_ ) - 1:
        lowercase__	: Any     		=   F'up_blocks.{i}.upsamplers.0'
        lowercase__	: Optional[int]     		=   F'output_blocks.{current_layer-1}.2'
        lowercase__	: int     		=   convert_resnet(lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )

  lowercase__	: str     		=   checkpoint["""out.0.weight"""]
  lowercase__	: Any     		=   checkpoint["""out.0.bias"""]
  lowercase__	: Any     		=   checkpoint["""out.2.weight"""]
  lowercase__	: Any     		=   checkpoint["""out.2.bias"""]

  return new_checkpoint


if __name__ == "__main__":
     lowerCamelCase__     :		Dict        = argparse.ArgumentParser()

     parser.add_argument("""--unet_path""", default=None, type=str, required=True, help="""Path to the unet.pt to convert.""")
     parser.add_argument(
         """--dump_path""", default=None, type=str, required=True, help="""Path to output the converted UNet model."""
     )
     parser.add_argument("""--class_cond""", default=True, type=str, help="""Whether the model is class-conditional.""")

     lowerCamelCase__     :		Tuple        = parser.parse_args()
     lowerCamelCase__     :		Dict        = strabool(args.class_cond)

     lowerCamelCase__     :		List[Any]        = os.path.basename(args.unet_path)
     print(f'''Checkpoint: {ckpt_name}''')

     # Get U-Net config
     if "imagenet64" in ckpt_name:
          lowerCamelCase__     :		Optional[Any]        = IMAGENET_64_UNET_CONFIG
     elif "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
          lowerCamelCase__     :		Optional[Any]        = LSUN_256_UNET_CONFIG
     elif "test" in ckpt_name:
          lowerCamelCase__     :		Dict        = TEST_UNET_CONFIG
     else:
          raise ValueError(f'''Checkpoint type {ckpt_name} is not currently supported.''')

     if not args.class_cond:
          lowerCamelCase__     :		Union[str, Any]        = None

     lowerCamelCase__     :		Optional[Any]        = con_pt_to_diffuser(args.unet_path, unet_config)

     lowerCamelCase__     :		List[Any]        = UNetaDModel(**unet_config)
     image_unet.load_state_dict(converted_unet_ckpt)

     # Get scheduler config
     if "cd" in ckpt_name or "test" in ckpt_name:
          lowerCamelCase__     :		Union[str, Any]        = CD_SCHEDULER_CONFIG
     elif "ct" in ckpt_name and "imagenet64" in ckpt_name:
          lowerCamelCase__     :		List[str]        = CT_IMAGENET_64_SCHEDULER_CONFIG
     elif "ct" in ckpt_name and "256" in ckpt_name and (("bedroom" in ckpt_name) or ("cat" in ckpt_name)):
          lowerCamelCase__     :		List[Any]        = CT_LSUN_256_SCHEDULER_CONFIG
     else:
          raise ValueError(f'''Checkpoint type {ckpt_name} is not currently supported.''')

     lowerCamelCase__     :		str        = CMStochasticIterativeScheduler(**scheduler_config)

     lowerCamelCase__     :		Optional[int]        = ConsistencyModelPipeline(unet=image_unet, scheduler=cm_scheduler)
     consistency_model.save_pretrained(args.dump_path)








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


 def __init__(       self):


   '''simple docstring'''




   lowercase__	: List[Any]     		=   []


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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



 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




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



@require_torch
class    _snake_case (						unittest.TestCase				):


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Dict     		=   tempfile.mkdtemp()


 def   lowercase__  (       self):


   '''simple docstring'''




   shutil.rmtree(self.output_dir)


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_=0       ,	SCREAMING_SNAKE_CASE_=0       ,	SCREAMING_SNAKE_CASE_=64       ,	SCREAMING_SNAKE_CASE_=64       ,	SCREAMING_SNAKE_CASE_=None       ,	SCREAMING_SNAKE_CASE_=False       ,	**SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: Any     		=   RegressionDataset(length=SCREAMING_SNAKE_CASE_)
   lowercase__	: Optional[int]     		=   RegressionDataset(length=SCREAMING_SNAKE_CASE_)
   lowercase__	: Dict     		=   RegressionModelConfig(a=SCREAMING_SNAKE_CASE_       ,	b=SCREAMING_SNAKE_CASE_)
   lowercase__	: Any     		=   RegressionPreTrainedModel(SCREAMING_SNAKE_CASE_)

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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   self.assertEqual(len(SCREAMING_SNAKE_CASE_)       ,	len(SCREAMING_SNAKE_CASE_))

   # Order doesn't matter
   lowercase__	: str     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: cb.__name__ if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else cb.__class__.__name__)
   lowercase__	: Tuple     		=   sorted(SCREAMING_SNAKE_CASE_       ,	key=lambda SCREAMING_SNAKE_CASE_: cb.__name__ if isinstance(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_) else cb.__class__.__name__)

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


 def   lowercase__  (       self       ,	SCREAMING_SNAKE_CASE_):


   '''simple docstring'''




   lowercase__	: int     		=   ["""on_init_end""", """on_train_begin"""]
   lowercase__	: Union[str, Any]     		=   0
   lowercase__	: Union[str, Any]     		=   len(trainer.get_eval_dataloader())
   lowercase__	: Dict     		=   ["""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(SCREAMING_SNAKE_CASE_):
       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   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: int     		=   self.get_trainer()
   lowercase__	: Union[str, Any]     		=   DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

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

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


 def   lowercase__  (       self):


   '''simple docstring'''




   lowercase__	: Any     		=   DEFAULT_CALLBACKS.copy() + [ProgressCallback]
   lowercase__	: Tuple     		=   self.get_trainer()

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

   lowercase__	: Optional[int]     		=   self.get_trainer()
   lowercase__	: List[Any]     		=   trainer.pop_callback(SCREAMING_SNAKE_CASE_)
   self.assertEqual(cb.__class__       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

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

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

   lowercase__	: str     		=   self.get_trainer()
   lowercase__	: Optional[Any]     		=   trainer.callback_handler.callbacks[0]
   lowercase__	: Union[str, Any]     		=   trainer.pop_callback(SCREAMING_SNAKE_CASE_)
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	SCREAMING_SNAKE_CASE_)
   self.check_callbacks_equality(trainer.callback_handler.callbacks       ,	SCREAMING_SNAKE_CASE_)

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



 def   lowercase__  (       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=SCREAMING_SNAKE_CASE_)

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

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

   lowercase__	: Optional[Any]     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	save_steps=5)
   trainer.train()
   lowercase__	: Dict     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

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

   lowercase__	: Tuple     		=   self.get_trainer(callbacks=[MyTestTrainerCallback]       ,	evaluation_strategy="""epoch""")
   trainer.train()
   lowercase__	: Optional[int]     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   # A bit of everything
   lowercase__	: Any     		=   self.get_trainer(
       callbacks=[MyTestTrainerCallback]       ,	logging_steps=3       ,	save_steps=10       ,	eval_steps=5       ,	evaluation_strategy="""steps"""       ,	)
   trainer.train()
   lowercase__	: str     		=   trainer.callback_handler.callbacks[-2].events
   self.assertEqual(SCREAMING_SNAKE_CASE_       ,	self.get_expected_events(SCREAMING_SNAKE_CASE_))

   # warning should be emitted for duplicated callbacks
   with patch("""transformers.trainer_callback.logger.warning""") as warn_mock:
     lowercase__	: Dict     		=   self.get_trainer(
         callbacks=[MyTestTrainerCallback, MyTestTrainerCallback]       ,	)
     assert str(SCREAMING_SNAKE_CASE_) in warn_mock.call_args[0][0]






