



import argparse
import os

# New Code #
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
from accelerate.utils import find_executable_batch_size


########################################################################
# This is a fully working simple example to use Accelerate,
# specifically showcasing how to ensure out-of-memory errors never
# interrupt training, and builds off the `nlp_example.py` script.
#
# 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)
#
# New additions from the base script can be found quickly by
# looking for the # New Code # tags
#
# 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
#
########################################################################


UpperCAmelCase_									=     1_6
UpperCAmelCase_									=     3_2
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Accelerator     ,				_SCREAMING_SNAKE_CASE   :	int = 1_6				)->Tuple:
      _lowerCAmelCase 			=  AutoTokenizer.from_pretrained('''bert-base-cased'''				)
      _lowerCAmelCase 			=  load_dataset('''glue'''     ,				'''mrpc'''				)

      def tokenize_function(_SCREAMING_SNAKE_CASE   :	List[Any]				):
            # max_length=None => use the model max length (it's actually the default)
            _lowerCAmelCase 			=  tokenizer(examples['''sentence1''']     ,				examples['''sentence2''']     ,				truncation=_SCREAMING_SNAKE_CASE     ,				max_length=_SCREAMING_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():
            _lowerCAmelCase 			=  datasets.map(
                _SCREAMING_SNAKE_CASE     ,				batched=_SCREAMING_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
      _lowerCAmelCase 			=  tokenized_datasets.rename_column('''label'''     ,				'''labels'''				)

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

            return tokenizer.pad(
                _SCREAMING_SNAKE_CASE     ,				padding='''longest'''     ,				max_length=_SCREAMING_SNAKE_CASE     ,				pad_to_multiple_of=_SCREAMING_SNAKE_CASE     ,				return_tensors='''pt'''     ,				)

      # Instantiate dataloaders.
      _lowerCAmelCase 			=  DataLoader(
          tokenized_datasets['''train''']     ,				shuffle=_SCREAMING_SNAKE_CASE     ,				collate_fn=_SCREAMING_SNAKE_CASE     ,				batch_size=_SCREAMING_SNAKE_CASE				)
      _lowerCAmelCase 			=  DataLoader(
          tokenized_datasets['''validation''']     ,				shuffle=_SCREAMING_SNAKE_CASE     ,				collate_fn=_SCREAMING_SNAKE_CASE     ,				batch_size=_SCREAMING_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

    UpperCAmelCase_									=     mocked_dataloaders  # noqa: F811
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Tuple     ,				_SCREAMING_SNAKE_CASE   :	Optional[int]				)->Optional[Any]:
      # For testing only
      if os.environ.get('''TESTING_MOCKED_DATALOADERS'''     ,				_SCREAMING_SNAKE_CASE				) == "1":
            _lowerCAmelCase 			=  2
      # Initialize accelerator
      _lowerCAmelCase 			=  Accelerator(cpu=args.cpu     ,				mixed_precision=args.mixed_precision				)
      # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
      _lowerCAmelCase 			=  config['''lr''']
      _lowerCAmelCase 			=  int(config['''num_epochs''']				)
      _lowerCAmelCase 			=  int(config['''seed''']				)
      _lowerCAmelCase 			=  int(config['''batch_size''']				)

      _lowerCAmelCase 			=  evaluate.load('''glue'''     ,				'''mrpc'''				)

      # New Code #
      # We now can define an inner training loop function. It should take a batch size as the only parameter,
      # and build the dataloaders in there.
      # It also gets our decorator
      @find_executable_batch_size(starting_batch_size=_SCREAMING_SNAKE_CASE				)
      def inner_training_loop(_SCREAMING_SNAKE_CASE   :	int				):
            # And now just move everything below under this function
            # We need to bring in the Accelerator object from earlier
            nonlocal accelerator
            # And reset all of its attributes that could hold onto any memory:
            accelerator.free_memory()

            # Then we can declare the model, optimizer, and everything else:
            set_seed(_SCREAMING_SNAKE_CASE				)

            # Instantiate the model (we build the model here so that the seed also control new weights initialization)
            _lowerCAmelCase 			=  AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''     ,				return_dict=_SCREAMING_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).
            _lowerCAmelCase 			=  model.to(accelerator.device				)

            # Instantiate optimizer
            _lowerCAmelCase 			=  AdamW(params=model.parameters()     ,				lr=_SCREAMING_SNAKE_CASE				)
            _lowerCAmelCase			,   _lowerCAmelCase 			=  get_dataloaders(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				)

            # Instantiate scheduler
            _lowerCAmelCase 			=  get_linear_schedule_with_warmup(
                optimizer=_SCREAMING_SNAKE_CASE     ,				num_warmup_steps=1_0_0     ,				num_training_steps=(len(_SCREAMING_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.
            _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase			,   _lowerCAmelCase 			=  accelerator.prepare(
                _SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				)

            # Now we train the model
            for epoch in range(_SCREAMING_SNAKE_CASE				):
                  model.train()
                  for step, batch in enumerate(_SCREAMING_SNAKE_CASE				):
                        # We could avoid this line since we set the accelerator with `device_placement=True`.
                        batch.to(accelerator.device				)
                        _lowerCAmelCase 			=  model(**_SCREAMING_SNAKE_CASE				)
                        _lowerCAmelCase 			=  outputs.loss
                        accelerator.backward(_SCREAMING_SNAKE_CASE				)
                        optimizer.step()
                        lr_scheduler.step()
                        optimizer.zero_grad()

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

                  _lowerCAmelCase 			=  metric.compute()
                  # Use accelerator.print to print only on the main process.
                  accelerator.print(f'''epoch {epoch}:'''     ,				_SCREAMING_SNAKE_CASE				)

    # New Code #
    # And call it at the end with no arguments
    # Note: You could also refactor this outside of your training loop function
      inner_training_loop()
def 							UpperCAmelCase__		(				)->int:
      _lowerCAmelCase 			=  argparse.ArgumentParser(description='''Simple example of training script.'''				)
      parser.add_argument(
          '''--mixed_precision'''     ,				type=_SCREAMING_SNAKE_CASE     ,				default=_SCREAMING_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.'''     ,				)
      parser.add_argument('''--cpu'''     ,				action='''store_true'''     ,				help='''If passed, will train on the CPU.'''				)
      _lowerCAmelCase 			=  parser.parse_args()
      _lowerCAmelCase 			=  {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
      training_function(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
    main()




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_video_inputs


if is_torch_available():
    import torch

if is_vision_available():
    from PIL import Image

    from transformers import VivitImageProcessor





class 		UpperCAmelCase						(						unittest.TestCase				):





   def __init__(	self       ,     _lowerCAmelCase       ,     _lowerCAmelCase=7       ,     _lowerCAmelCase=3       ,     _lowerCAmelCase=10       ,     _lowerCAmelCase=18       ,     _lowerCAmelCase=30       ,     _lowerCAmelCase=400       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=None       ,     _lowerCAmelCase=True       ,     _lowerCAmelCase=[0.5, 0.5, 0.5]       ,     _lowerCAmelCase=[0.5, 0.5, 0.5]       ,     _lowerCAmelCase=None       ,     ):
         _lowerCAmelCase 			=  size if size is not None else {'''shortest_edge''': 18}
         _lowerCAmelCase 			=  crop_size if crop_size is not None else {'''height''': 18, '''width''': 18}

         _lowerCAmelCase 			=  parent
         _lowerCAmelCase 			=  batch_size
         _lowerCAmelCase 			=  num_channels
         _lowerCAmelCase 			=  num_frames
         _lowerCAmelCase 			=  image_size
         _lowerCAmelCase 			=  min_resolution
         _lowerCAmelCase 			=  max_resolution
         _lowerCAmelCase 			=  do_resize
         _lowerCAmelCase 			=  size
         _lowerCAmelCase 			=  do_normalize
         _lowerCAmelCase 			=  image_mean
         _lowerCAmelCase 			=  image_std
         _lowerCAmelCase 			=  crop_size





   def 	__lowerCAmelCase							(	self			):
         return {
             "image_mean": self.image_mean,
             "image_std": self.image_std,
             "do_normalize": self.do_normalize,
             "do_resize": self.do_resize,
             "size": self.size,
             "crop_size": self.crop_size,
         }






@require_torch
@require_vision
class 		UpperCAmelCase						(						snake_case_       ,unittest.TestCase				):
   SCREAMING_SNAKE_CASE__							=     VivitImageProcessor if is_vision_available() else None





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  VivitImageProcessingTester(self			)





   @property
   def 	__lowerCAmelCase							(	self			):
         return self.image_processor_tester.prepare_image_processor_dict()





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.image_processing_class(**self.image_processor_dict			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''image_mean'''			)			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''image_std'''			)			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''do_normalize'''			)			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''do_resize'''			)			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''do_center_crop'''			)			)
         self.assertTrue(hasattr(_lowerCAmelCase       ,     '''size'''			)			)





   def 	__lowerCAmelCase							(	self			):
         _lowerCAmelCase 			=  self.image_processing_class.from_dict(self.image_processor_dict			)
         self.assertEqual(image_processor.size       ,     {'''shortest_edge''': 18}			)
         self.assertEqual(image_processor.crop_size       ,     {'''height''': 18, '''width''': 18}			)

         _lowerCAmelCase 			=  self.image_processing_class.from_dict(self.image_processor_dict       ,     size=42       ,     crop_size=84			)
         self.assertEqual(image_processor.size       ,     {'''shortest_edge''': 42}			)
         self.assertEqual(image_processor.crop_size       ,     {'''height''': 84, '''width''': 84}			)





   def 	__lowerCAmelCase							(	self			):
         # Initialize image_processing
         _lowerCAmelCase 			=  self.image_processing_class(**self.image_processor_dict			)
         # create random PIL videos
         _lowerCAmelCase 			=  prepare_video_inputs(self.image_processor_tester       ,     equal_resolution=_lowerCAmelCase			)
         for video in video_inputs:
               self.assertIsInstance(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertIsInstance(video[0]       ,     Image.Image			)

         # Test not batched input
         _lowerCAmelCase 			=  image_processing(video_inputs[0]       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 1,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )

         # Test batched
         _lowerCAmelCase 			=  image_processing(_lowerCAmelCase       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )





   def 	__lowerCAmelCase							(	self			):
         # Initialize image_processing
         _lowerCAmelCase 			=  self.image_processing_class(**self.image_processor_dict			)
         # create random numpy tensors
         _lowerCAmelCase 			=  prepare_video_inputs(self.image_processor_tester       ,     equal_resolution=_lowerCAmelCase       ,     numpify=_lowerCAmelCase			)
         for video in video_inputs:
               self.assertIsInstance(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertIsInstance(video[0]       ,     np.ndarray			)

         # Test not batched input
         _lowerCAmelCase 			=  image_processing(video_inputs[0]       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 1,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )

         # Test batched
         _lowerCAmelCase 			=  image_processing(_lowerCAmelCase       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )





   def 	__lowerCAmelCase							(	self			):
         # Initialize image_processing
         _lowerCAmelCase 			=  self.image_processing_class(**self.image_processor_dict			)
         # create random PyTorch tensors
         _lowerCAmelCase 			=  prepare_video_inputs(self.image_processor_tester       ,     equal_resolution=_lowerCAmelCase       ,     torchify=_lowerCAmelCase			)
         for video in video_inputs:
               self.assertIsInstance(_lowerCAmelCase       ,     _lowerCAmelCase			)
               self.assertIsInstance(video[0]       ,     torch.Tensor			)

         # Test not batched input
         _lowerCAmelCase 			=  image_processing(video_inputs[0]       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 1,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )

         # Test batched
         _lowerCAmelCase 			=  image_processing(_lowerCAmelCase       ,     return_tensors='''pt'''			).pixel_values
         self.assertEqual(
             encoded_videos.shape       ,     (
                 self.image_processor_tester.batch_size,
                 self.image_processor_tester.num_frames,
                 self.image_processor_tester.num_channels,
                 self.image_processor_tester.crop_size['''height'''],
                 self.image_processor_tester.crop_size['''width'''],
             )       ,     )