'''simple docstring'''





import argparse
import os

import evaluate
import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType


########################################################################
# This is a fully working simple example to use Accelerate
# and perform gradient accumulation
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


SCREAMING_SNAKE_CASE_:	Union[str, Any]    =16
SCREAMING_SNAKE_CASE_:	List[str]    =32





def       lowerCAmelCase_							(   snake_case_			:					Accelerator   ,       snake_case_			:					int = 16     )       -> Union[str, Any]:


     '''simple docstring'''




     UpperCAmelCase_			       =  AutoTokenizer.from_pretrained("bert-base-cased"     )
     UpperCAmelCase_			       =  load_dataset("glue"   ,       "mrpc"     )

     def tokenize_function(snake_case_			:					List[Any]     ):
          # max_length=None => use the model max length (it's actually the default)
          UpperCAmelCase_			       =  tokenizer(examples["sentence1"]   ,       examples["sentence2"]   ,       truncation=snake_case_   ,       max_length=snake_case_     )
          return outputs

     # Apply the method we just defined to all the examples in all the splits of the dataset
     # starting with the main process first:
     with accelerator.main_process_first():
          UpperCAmelCase_			       =  datasets.map(
              snake_case_   ,       batched=snake_case_   ,       remove_columns=["idx", "sentence1", "sentence2"]   ,       )

     # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
     # transformers library
     UpperCAmelCase_			       =  tokenized_datasets.rename_column("label"   ,       "labels"     )

     def collate_fn(snake_case_			:					List[Any]     ):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          UpperCAmelCase_			       =  1_28 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               UpperCAmelCase_			       =  16
          elif accelerator.mixed_precision != "no":
               UpperCAmelCase_			       =  8
          else:
               UpperCAmelCase_			       =  None

          return tokenizer.pad(
              snake_case_   ,       padding="longest"   ,       max_length=snake_case_   ,       pad_to_multiple_of=snake_case_   ,       return_tensors="pt"   ,       )

     # Instantiate dataloaders.
     UpperCAmelCase_			       =  DataLoader(
         tokenized_datasets["train"]   ,       shuffle=snake_case_   ,       collate_fn=snake_case_   ,       batch_size=snake_case_     )
     UpperCAmelCase_			       =  DataLoader(
         tokenized_datasets["validation"]   ,       shuffle=snake_case_   ,       collate_fn=snake_case_   ,       batch_size=snake_case_     )

     return train_dataloader, eval_dataloader


# For testing only
if os.environ.get('TESTING_MOCKED_DATALOADERS', None) == "1":
    from accelerate.test_utils.training import mocked_dataloaders

    SCREAMING_SNAKE_CASE_:	Tuple    =mocked_dataloaders  # noqa: F811





def       lowerCAmelCase_							(   snake_case_			:					Union[str, Any]   ,       snake_case_			:					Union[str, Any]     )       -> Optional[int]:


     '''simple docstring'''




     if os.environ.get("TESTING_MOCKED_DATALOADERS"   ,       snake_case_     ) == "1":
          UpperCAmelCase_			       =  2
     # New Code #
     UpperCAmelCase_			       =  int(args.gradient_accumulation_steps     )
     # Initialize accelerator
     UpperCAmelCase_			       =  Accelerator(
         cpu=args.cpu   ,       mixed_precision=args.mixed_precision   ,       gradient_accumulation_steps=snake_case_     )
     if accelerator.distributed_type == DistributedType.TPU and gradient_accumulation_steps > 1:
          raise NotImplementedError(
              "Gradient accumulation on TPUs is currently not supported. Pass `gradient_accumulation_steps=1`"     )
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     UpperCAmelCase_			       =  config["lr"]
     UpperCAmelCase_			       =  int(config["num_epochs"]     )
     UpperCAmelCase_			       =  int(config["seed"]     )
     UpperCAmelCase_			       =  int(config["batch_size"]     )

     UpperCAmelCase_			       =  evaluate.load("glue"   ,       "mrpc"     )

     set_seed(snake_case_     )
     UpperCAmelCase_       ,					UpperCAmelCase_			       =  get_dataloaders(snake_case_   ,       snake_case_     )
     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     UpperCAmelCase_			       =  AutoModelForSequenceClassification.from_pretrained("bert-base-cased"   ,       return_dict=snake_case_     )

     # We could avoid this line since the accelerator is set with `device_placement=True` (default value).
     # Note that if you are placing tensors on devices manually, this line absolutely needs to be before the optimizer
     # creation otherwise training will not work on TPU (`accelerate` will kindly throw an error to make us aware of that).
     UpperCAmelCase_			       =  model.to(accelerator.device     )

     # Instantiate optimizer
     UpperCAmelCase_			       =  AdamW(params=model.parameters()   ,       lr=snake_case_     )

     # Instantiate scheduler
     UpperCAmelCase_			       =  get_linear_schedule_with_warmup(
         optimizer=snake_case_   ,       num_warmup_steps=1_00   ,       num_training_steps=(len(snake_case_     ) * num_epochs)   ,       )

     # Prepare everything
     # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
     # prepare method.
     UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_       ,					UpperCAmelCase_			       =  accelerator.prepare(
         snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_   ,       snake_case_     )

     # Now we train the model
     for epoch in range(snake_case_     ):
          model.train()
          for step, batch in enumerate(snake_case_     ):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device     )
               # New code #
               # We use the new `accumulate` context manager to perform gradient accumulation
               # We also currently do not support TPUs nor advise it as bugs were found on the XLA side when running our tests.
               with accelerator.accumulate(snake_case_     ):
                    UpperCAmelCase_			       =  model(**snake_case_     )
                    UpperCAmelCase_			       =  output.loss
                    accelerator.backward(snake_case_     )
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

          model.eval()
          for step, batch in enumerate(snake_case_     ):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device     )
               with torch.no_grad():
                    UpperCAmelCase_			       =  model(**snake_case_     )
               UpperCAmelCase_			       =  outputs.logits.argmax(dim=-1     )
               UpperCAmelCase_       ,					UpperCAmelCase_			       =  accelerator.gather_for_metrics((predictions, batch["labels"])     )
               metric.add_batch(
                   predictions=snake_case_   ,       references=snake_case_   ,       )

          UpperCAmelCase_			       =  metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f"""epoch {epoch}:"""   ,       snake_case_     )





def       lowerCAmelCase_							(   )       -> Optional[Any]:


     '''simple docstring'''




     UpperCAmelCase_			       =  argparse.ArgumentParser(description="Simple example of training script."     )
     parser.add_argument(
         "--mixed_precision"   ,       type=snake_case_   ,       default=snake_case_   ,       choices=["no", "fp16", "bf16", "fp8"]   ,       help="Whether to use mixed precision. Choose"
         "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
         "and an Nvidia Ampere GPU."   ,       )
     # New Code #
     parser.add_argument(
         "--gradient_accumulation_steps"   ,       type=snake_case_   ,       default=1   ,       help="The number of minibatches to be ran before gradients are accumulated."   ,       )
     parser.add_argument("--cpu"   ,       action="store_true"   ,       help="If passed, will train on the CPU."     )
     UpperCAmelCase_			       =  parser.parse_args()
     UpperCAmelCase_			       =  {"lr": 2E-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
     training_function(snake_case_   ,       snake_case_     )


if __name__ == "__main__":
    main()












import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import TransformeraDModel, VQDiffusionPipeline, VQDiffusionScheduler, VQModel
from diffusers.pipelines.vq_diffusion.pipeline_vq_diffusion import LearnedClassifierFreeSamplingEmbeddings
from diffusers.utils import load_numpy, slow, torch_device
from diffusers.utils.testing_utils import require_torch_gpu


UpperCAmelCase__							  =						False






class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Dict    )       ->str:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      @property
      def  __lowerCAmelCase   ( self :		List[str]    )       ->Any:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Optional[Any]    )       ->Dict:




       """simple docstring"""


       return 12







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Any:




       """simple docstring"""


       return 32







      @property
      def  __lowerCAmelCase   ( self :		int    )       ->List[str]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = VQModel(
           block_out_channels=[32, 64]					,   in_channels=3					,   out_channels=3					,   down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,   up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,   latent_channels=3					,   num_vq_embeddings=self.num_embed					,   vq_embed_dim=3					,   )
       return model







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->str:




       """simple docstring"""


       a						    = CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''    )
       return tokenizer







      @property
      def  __lowerCAmelCase   ( self :		str    )       ->Union[str, Any]:




       """simple docstring"""


       torch.manual_seed(0    )
       a						    = CLIPTextConfig(
           bos_token_id=0					,   eos_token_id=2					,   hidden_size=self.text_embedder_hidden_size					,   intermediate_size=37					,   layer_norm_eps=1e-0_5					,   num_attention_heads=4					,   num_hidden_layers=5					,   pad_token_id=1					,   vocab_size=1_000					,   )
       return CLIPTextModel(__UpperCAmelCase    )







      @property
      def  __lowerCAmelCase   ( self :		Any    )       ->Tuple:




       """simple docstring"""


       torch.manual_seed(0    )

       a						    = 12
       a						    = 12

       a						    = {
           '''attention_bias''': True,
           '''cross_attention_dim''': 32,
           '''attention_head_dim''': height * width,
           '''num_attention_heads''': 1,
           '''num_vector_embeds''': self.num_embed,
           '''num_embeds_ada_norm''': self.num_embeds_ada_norm,
           '''norm_num_groups''': 32,
           '''sample_size''': width,
           '''activation_fn''': '''geglu-approximate''',
       }

       a						    = TransformeraDModel(**__UpperCAmelCase    )
       return model







      def  __lowerCAmelCase   ( self :		Dict    )       ->Any:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(learnable=__UpperCAmelCase    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6551, 0.6168, 0.5008, 0.5676, 0.5659, 0.4295, 0.6073, 0.5599, 0.4992]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 1e-2
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2







      def  __lowerCAmelCase   ( self :		Tuple    )       ->str:




       """simple docstring"""


       a						    = '''cpu'''

       a						    = self.dummy_vqvae
       a						    = self.dummy_text_encoder
       a						    = self.dummy_tokenizer
       a						    = self.dummy_transformer
       a						    = VQDiffusionScheduler(self.num_embed    )
       a						    = LearnedClassifierFreeSamplingEmbeddings(
           learnable=__UpperCAmelCase					,   hidden_size=self.text_embedder_hidden_size					,   length=tokenizer.model_max_length    )

       a						    = VQDiffusionPipeline(
           vqvae=__UpperCAmelCase					,   text_encoder=__UpperCAmelCase					,   tokenizer=__UpperCAmelCase					,   transformer=__UpperCAmelCase					,   scheduler=__UpperCAmelCase					,   learned_classifier_free_sampling_embeddings=__UpperCAmelCase					,   )
       a						    = pipe.to(__UpperCAmelCase    )
       pipe.set_progress_bar_config(disable=__UpperCAmelCase    )

       a						    = '''teddy bear playing in the pool'''

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe([prompt]					,   generator=__UpperCAmelCase					,   num_inference_steps=2					,   output_type='''np'''    )
       a						    = output.images

       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipe(
           [prompt]					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   return_dict=__UpperCAmelCase					,   num_inference_steps=2    )[0]

       a						    = image[0, -3:, -3:, -1]
       a						    = image_from_tuple[0, -3:, -3:, -1]

       assert image.shape == (1, 24, 24, 3)

       a						    = np.array([0.6693, 0.6075, 0.4959, 0.5701, 0.5583, 0.4333, 0.6171, 0.5684, 0.4988]    )

       assert np.abs(image_slice.flatten() - expected_slice    ).max() < 2.0
       assert np.abs(image_from_tuple_slice.flatten() - expected_slice    ).max() < 1e-2






@slow
@require_torch_gpu
class 						lowercase_   (      unittest.TestCase  ):






      '''simple docstring'''




      def  __lowerCAmelCase   ( self :		Union[str, Any]    )       ->Union[str, Any]:




       """simple docstring"""


       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()







      def  __lowerCAmelCase   ( self :		Dict    )       ->int:




       """simple docstring"""


       a						    = load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/vq_diffusion/teddy_bear_pool_classifier_free_sampling.npy'''    )

       a						    = VQDiffusionPipeline.from_pretrained('''microsoft/vq-diffusion-ithq'''    )
       a						    = pipeline.to(__UpperCAmelCase    )
       pipeline.set_progress_bar_config(disable=__UpperCAmelCase    )

       # requires GPU generator for gumbel softmax
       # don't use GPU generator in tests though
       a						    = torch.Generator(device=__UpperCAmelCase    ).manual_seed(0    )
       a						    = pipeline(
           '''teddy bear playing in the pool'''					,   num_images_per_prompt=1					,   generator=__UpperCAmelCase					,   output_type='''np'''					,   )

       a						    = output.images[0]

       assert image.shape == (256, 256, 3)
       assert np.abs(expected_image - image    ).max() < 2.0


