



"""simple docstring"""







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_image_inputs


if is_torch_available():
 import torch

if is_vision_available():
 from PIL import Image

 from transformers import ViltImageProcessor

class       UpperCamelCase  (					unittest.TestCase				):





     """simple docstring"""






     def __init__(					self :							Union[str, Any]	,      _lowerCamelCase :							Union[str, Any]	,      _lowerCamelCase :							int=7	,      _lowerCamelCase :							List[str]=3	,      _lowerCamelCase :							Dict=1_8	,      _lowerCamelCase :							Optional[int]=3_0	,      _lowerCamelCase :							Optional[Any]=4_0_0	,      _lowerCamelCase :							int=True	,      _lowerCamelCase :							Dict=None	,      _lowerCamelCase :							List[str]=2	,      _lowerCamelCase :							Optional[Any]=True	,      _lowerCamelCase :							Optional[Any]=[0.5, 0.5, 0.5]	,      _lowerCamelCase :							str=[0.5, 0.5, 0.5]	,      ):
          A__				=   size if size is not None else {'''shortest_edge''': 3_0}
          A__				=   parent
          A__				=   batch_size
          A__				=   num_channels
          A__				=   image_size
          A__				=   min_resolution
          A__				=   max_resolution
          A__				=   do_resize
          A__				=   size
          A__				=   size_divisor
          A__				=   do_normalize
          A__				=   image_mean
          A__				=   image_std





     def 						A__  (					self :							Union[str, Any] ):
          return {
              "image_mean": self.image_mean,
              "image_std": self.image_std,
              "do_normalize": self.do_normalize,
              "do_resize": self.do_resize,
              "size": self.size,
              "size_divisor": self.size_divisor,
          }





     def 						A__  (					self :							Optional[int]	,      _lowerCamelCase :							Union[str, Any]	,      _lowerCamelCase :							Optional[Any]=False ):
          if not batched:
               A__				=   self.size['''shortest_edge''']
               A__				=   image_inputs[0]
               if isinstance(_lowerCamelCase	,      Image.Image ):
                    A__	,     A__				=   image.size
               else:
                    A__	,     A__				=   image.shape[1], image.shape[2]
               A__				=   size / min(_lowerCamelCase	,      _lowerCamelCase )
               if h < w:
                    A__	,     A__				=   size, scale * w
               else:
                    A__	,     A__				=   scale * h, size

               A__				=   int((1_3_3_3 / 8_0_0) * size )
               if max(_lowerCamelCase	,      _lowerCamelCase ) > max_size:
                    A__				=   max_size / max(_lowerCamelCase	,      _lowerCamelCase )
                    A__				=   newh * scale
                    A__				=   neww * scale

               A__	,     A__				=   int(newh + 0.5 ), int(neww + 0.5 )
               A__	,     A__				=   (
                   newh // self.size_divisor * self.size_divisor,
                   neww // self.size_divisor * self.size_divisor,
               )

          else:
               A__				=   []
               for image in image_inputs:
                    A__	,     A__				=   self.get_expected_values([image] )
                    expected_values.append((expected_height, expected_width) )
               A__				=   max(_lowerCamelCase	,      key=lambda _lowerCamelCase : item[0] )[0]
               A__				=   max(_lowerCamelCase	,      key=lambda _lowerCamelCase : item[1] )[1]

          return expected_height, expected_width


@require_torch
@require_vision
class       UpperCamelCase  (					a      , unittest.TestCase				):





     """simple docstring"""






     _lowerCamelCase					: Union[str, Any]				   =ViltImageProcessor if is_vision_available() else None





     def 						A__  (					self :							Union[str, Any] ):
          A__				=   ViltImageProcessingTester(self )





     @property
     def 						A__  (					self :							str ):
          return self.image_processor_tester.prepare_image_processor_dict()





     def 						A__  (					self :							List[str] ):
          A__				=   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	,      '''size''' ) )
          self.assertTrue(hasattr(_lowerCamelCase	,      '''size_divisor''' ) )





     def 						A__  (					self :							str ):
          A__				=   self.image_processing_class.from_dict(self.image_processor_dict )
          self.assertEqual(image_processor.size	,      {'''shortest_edge''': 3_0} )

          A__				=   self.image_processing_class.from_dict(self.image_processor_dict	,      size=4_2 )
          self.assertEqual(image_processor.size	,      {'''shortest_edge''': 4_2} )





     def 						A__  (					self :							Tuple ):
          pass





     def 						A__  (					self :							Tuple ):
          # Initialize image_processing
          A__				=   self.image_processing_class(**self.image_processor_dict )
          # create random PIL images
          A__				=   prepare_image_inputs(self.image_processor_tester	,      equal_resolution=_lowerCamelCase )
          for image in image_inputs:
               self.assertIsInstance(_lowerCamelCase	,      Image.Image )

          # Test not batched input
          A__				=   image_processing(image_inputs[0]	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,      )

          # Test batched
          A__				=   image_processing(_lowerCamelCase	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase	,      batched=_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )	,      )





     def 						A__  (					self :							List[str] ):
          # Initialize image_processing
          A__				=   self.image_processing_class(**self.image_processor_dict )
          # create random numpy tensors
          A__				=   prepare_image_inputs(self.image_processor_tester	,      equal_resolution=_lowerCamelCase	,      numpify=_lowerCamelCase )
          for image in image_inputs:
               self.assertIsInstance(_lowerCamelCase	,      np.ndarray )

          # Test not batched input
          A__				=   image_processing(image_inputs[0]	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,      )

          # Test batched
          A__				=   image_processing(_lowerCamelCase	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase	,      batched=_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )	,      )





     def 						A__  (					self :							Tuple ):
          # Initialize image_processing
          A__				=   self.image_processing_class(**self.image_processor_dict )
          # create random PyTorch tensors
          A__				=   prepare_image_inputs(self.image_processor_tester	,      equal_resolution=_lowerCamelCase	,      torchify=_lowerCamelCase )
          for image in image_inputs:
               self.assertIsInstance(_lowerCamelCase	,      torch.Tensor )

          # Test not batched input
          A__				=   image_processing(image_inputs[0]	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (1, self.image_processor_tester.num_channels, expected_height, expected_width)	,      )

          # Test batched
          A__				=   image_processing(_lowerCamelCase	,      return_tensors='''pt''' ).pixel_values

          A__	,     A__				=   self.image_processor_tester.get_expected_values(_lowerCamelCase	,      batched=_lowerCamelCase )
          self.assertEqual(
              encoded_images.shape	,      (
                  self.image_processor_tester.batch_size,
                  self.image_processor_tester.num_channels,
                  expected_height,
                  expected_width,
              )	,      )








"""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
#
# 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)
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# 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
#
########################################################################


__snake_case	:				Tuple  					=			16
__snake_case	:				Any  					=			32



def   a_	(			__a		,							__a      ):
     # Initialize accelerator
     if args.with_tracking:
          A__				=   Accelerator(
              cpu=args.cpu		,							mixed_precision=args.mixed_precision		,							log_with='''all'''		,							project_dir=args.project_dir      )
     else:
          A__				=   Accelerator(cpu=args.cpu		,							mixed_precision=args.mixed_precision      )

     if hasattr(args.checkpointing_steps		,							'''isdigit'''      ):
          if args.checkpointing_steps == "epoch":
               A__				=   args.checkpointing_steps
          elif args.checkpointing_steps.isdigit():
               A__				=   int(args.checkpointing_steps      )
          else:
               raise ValueError(
                   f'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''      )
     else:
          A__				=   None
     # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
     A__				=   config['''lr''']
     A__				=   int(config['''num_epochs''']      )
     A__				=   int(config['''seed''']      )
     A__				=   int(config['''batch_size''']      )

     # We need to initialize the trackers we use, and also store our configuration
     if args.with_tracking:
          A__				=   os.path.split(__a      )[-1].split('''.'''      )[0]
          accelerator.init_trackers(__a		,							__a      )

     A__				=   AutoTokenizer.from_pretrained('''bert-base-cased'''      )
     A__				=   load_dataset('''glue'''		,							'''mrpc'''      )
     A__				=   evaluate.load('''glue'''		,							'''mrpc'''      )

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

     # If the batch size is too big we use gradient accumulation
     A__				=   1
     if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
          A__				=   batch_size // MAX_GPU_BATCH_SIZE
          A__				=   MAX_GPU_BATCH_SIZE

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

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

     # Instantiate dataloaders.
     A__				=   DataLoader(
         tokenized_datasets['''train''']		,							shuffle=__a		,							collate_fn=__a		,							batch_size=__a      )
     A__				=   DataLoader(
         tokenized_datasets['''validation''']		,							shuffle=__a		,							collate_fn=__a		,							batch_size=__a      )

     set_seed(__a      )

     # Instantiate the model (we build the model here so that the seed also control new weights initialization)
     A__				=   AutoModelForSequenceClassification.from_pretrained('''bert-base-cased'''		,							return_dict=__a      )

     # 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).
     A__				=   model.to(accelerator.device      )

     # Instantiate optimizer
     A__				=   AdamW(params=model.parameters()		,							lr=__a      )

     # Instantiate scheduler
     A__				=   get_linear_schedule_with_warmup(
         optimizer=__a		,							num_warmup_steps=100		,							num_training_steps=(len(__a      ) * num_epochs) // gradient_accumulation_steps		,							)

     # 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.
     A__	,     A__	,     A__	,     A__	,     A__				=   accelerator.prepare(
         __a		,							__a		,							__a		,							__a		,							__a      )

     # We need to keep track of how many total steps we have iterated over
     A__				=   0
     # We also need to keep track of the stating epoch so files are named properly
     A__				=   0

     # Potentially load in the weights and states from a previous save
     if args.resume_from_checkpoint:
          if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
               accelerator.print(f'''Resumed from checkpoint: {args.resume_from_checkpoint}'''      )
               accelerator.load_state(args.resume_from_checkpoint      )
               A__				=   os.path.basename(args.resume_from_checkpoint      )
          else:
               # Get the most recent checkpoint
               A__				=   [f.name for f in os.scandir(os.getcwd()      ) if f.is_dir()]
               dirs.sort(key=os.path.getctime      )
               A__				=   dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
          # Extract `epoch_{i}` or `step_{i}`
          A__				=   os.path.splitext(__a      )[0]

          if "epoch" in training_difference:
               A__				=   int(training_difference.replace('''epoch_'''		,							''''''      )      ) + 1
               A__				=   None
          else:
               A__				=   int(training_difference.replace('''step_'''		,							''''''      )      )
               A__				=   resume_step // len(__a      )
               resume_step -= starting_epoch * len(__a      )

    # Now we train the model
     for epoch in range(__a		,							__a      ):
          model.train()
          if args.with_tracking:
               A__				=   0
          if args.resume_from_checkpoint and epoch == starting_epoch and resume_step is not None:
               # We need to skip steps until we reach the resumed step
               A__				=   accelerator.skip_first_batches(__a		,							__a      )
               overall_step += resume_step
          else:
               # After the first iteration though, we need to go back to the original dataloader
               A__				=   train_dataloader
          for step, batch in enumerate(__a      ):
               # We could avoid this line since we set the accelerator with `device_placement=True`.
               batch.to(accelerator.device      )
               A__				=   model(**__a      )
               A__				=   outputs.loss
               A__				=   loss / gradient_accumulation_steps
               # We keep track of the loss at each epoch
               if args.with_tracking:
                    total_loss += loss.detach().float()
               accelerator.backward(__a      )
               if step % gradient_accumulation_steps == 0:
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()

               overall_step += 1

               if isinstance(__a		,							__a      ):
                    A__				=   f'''step_{overall_step}'''
                    if overall_step % checkpointing_steps == 0:
                         if args.output_dir is not None:
                              A__				=   os.path.join(args.output_dir		,							__a      )
                         accelerator.save_state(__a      )

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

          A__				=   metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f'''epoch {epoch}:'''		,							__a      )
          if args.with_tracking:
               accelerator.log(
                   {
                       '''accuracy''': eval_metric['''accuracy'''],
                       '''f1''': eval_metric['''f1'''],
                       '''train_loss''': total_loss.item() / len(__a      ),
                       '''epoch''': epoch,
                   }		,							step=__a		,							)

          if checkpointing_steps == "epoch":
               A__				=   f'''epoch_{epoch}'''
               if args.output_dir is not None:
                    A__				=   os.path.join(args.output_dir		,							__a      )
               accelerator.save_state(__a      )

     if args.with_tracking:
          accelerator.end_training()



def   a_	(			):
     A__				=   argparse.ArgumentParser(description='''Simple example of training script.'''      )
     parser.add_argument(
         '''--mixed_precision'''		,							type=__a		,							default=__a		,							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.'''      )
     parser.add_argument(
         '''--checkpointing_steps'''		,							type=__a		,							default=__a		,							help='''Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch.'''		,							)
     parser.add_argument(
         '''--resume_from_checkpoint'''		,							type=__a		,							default=__a		,							help='''If the training should continue from a checkpoint folder.'''		,							)
     parser.add_argument(
         '''--with_tracking'''		,							action='''store_true'''		,							help='''Whether to load in all available experiment trackers from the environment and use them for logging.'''		,							)
     parser.add_argument(
         '''--output_dir'''		,							type=__a		,							default='''.'''		,							help='''Optional save directory where all checkpoint folders will be stored. Default is the current working directory.'''		,							)
     parser.add_argument(
         '''--project_dir'''		,							type=__a		,							default='''logs'''		,							help='''Location on where to store experiment tracking logs` and relevent project information'''		,							)
     A__				=   parser.parse_args()
     A__				=   {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
     training_function(__a		,							__a      )


if __name__ == "__main__":
 main()



