





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
#
########################################################################


__A     =    16
__A     =    32





def 					lowerCAmelCase_			(	__a     ,   __a = 16		)					->		Optional[Any]:







  """simple docstring"""





  lowerCamelCase__:				int        =AutoTokenizer.from_pretrained("bert-base-cased"		)
  lowerCamelCase__:				Any        =load_dataset("glue"     ,   "mrpc"		)

  def tokenize_function(__a		):
    # max_length=None => use the model max length (it's actually the default)
    lowerCamelCase__:				int        =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():
    lowerCamelCase__:				Tuple        =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
  lowerCamelCase__:				Optional[int]        =tokenized_datasets.rename_column("label"     ,   "labels"		)

  def collate_fn(__a		):
    # On TPU it's best to pad everything to the same length or training will be very slow.
    lowerCamelCase__:				Dict        =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":
      lowerCamelCase__:				str        =16
    elif accelerator.mixed_precision != "no":
      lowerCamelCase__:				str        =8
    else:
      lowerCamelCase__:				str        =None

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

  # Instantiate dataloaders.
  lowerCamelCase__:				Union[str, Any]        =DataLoader(
      tokenized_datasets["train"]     ,   shuffle=__a     ,   collate_fn=__a     ,   batch_size=__a		)
  lowerCamelCase__:				Optional[int]        =DataLoader(
      tokenized_datasets["validation"]     ,   shuffle=__a     ,   collate_fn=__a     ,   batch_size=__a		)

  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

     __A     =    mocked_dataloaders  # noqa: F811





def 					lowerCAmelCase_			(	__a     ,   __a		)					->		int:







  """simple docstring"""





  if os.environ.get("TESTING_MOCKED_DATALOADERS"     ,   __a		) == "1":
    lowerCamelCase__:				Union[str, Any]        =2
  # New Code #
  lowerCamelCase__:				List[Any]        =int(args.gradient_accumulation_steps		)
  # Initialize accelerator
  lowerCamelCase__:				Optional[int]        =Accelerator(
      cpu=args.cpu     ,   mixed_precision=args.mixed_precision     ,   gradient_accumulation_steps=__a		)
  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
  lowerCamelCase__:				int        =config["lr"]
  lowerCamelCase__:				Tuple        =int(config["num_epochs"]		)
  lowerCamelCase__:				List[Any]        =int(config["seed"]		)
  lowerCamelCase__:				Tuple        =int(config["batch_size"]		)

  lowerCamelCase__:				List[Any]        =evaluate.load("glue"     ,   "mrpc"		)

  set_seed(__a		)
  lowerCamelCase__  , lowerCamelCase__:				Any        =get_dataloaders(__a     ,   __a		)
  # Instantiate the model (we build the model here so that the seed also control new weights initialization)
  lowerCamelCase__:				Dict        =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).
  lowerCamelCase__:				Any        =model.to(accelerator.device		)

  # Instantiate optimizer
  lowerCamelCase__:				List[Any]        =AdamW(params=model.parameters()     ,   lr=__a		)

  # Instantiate scheduler
  lowerCamelCase__:				str        =get_linear_schedule_with_warmup(
      optimizer=__a     ,   num_warmup_steps=100     ,   num_training_steps=(len(__a		) * 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__:				Optional[int]        =accelerator.prepare(
      __a     ,   __a     ,   __a     ,   __a     ,   __a		)

  # Now we train the model
  for epoch in range(__a		):
    model.train()
    for step, batch in enumerate(__a		):
      # 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(__a		):
        lowerCamelCase__:				List[str]        =model(**__a		)
        lowerCamelCase__:				Dict        =output.loss
        accelerator.backward(__a		)
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()

    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():
        lowerCamelCase__:				int        =model(**__a		)
      lowerCamelCase__:				int        =outputs.logits.argmax(dim=-1		)
      lowerCamelCase__  , lowerCamelCase__:				Optional[int]        =accelerator.gather_for_metrics((predictions, batch["labels"])		)
      metric.add_batch(
          predictions=__a     ,   references=__a     ,   )

    lowerCamelCase__:				List[Any]        =metric.compute()
    # Use accelerator.print to print only on the main process.
    accelerator.print(F"""epoch {epoch}:"""     ,   __a		)





def 					lowerCAmelCase_			(	)					->		List[Any]:







  """simple docstring"""





  lowerCamelCase__:				List[Any]        =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."     ,   )
  # New Code #
  parser.add_argument(
      "--gradient_accumulation_steps"     ,   type=__a     ,   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."		)
  lowerCamelCase__:				Tuple        =parser.parse_args()
  lowerCamelCase__:				Any        ={"lr": 2e-5, "num_epochs": 3, "seed": 42, "batch_size": 16}
  training_function(__a     ,   __a		)


if __name__ == "__main__":
     main()








from __future__ import annotations

import inspect
import unittest

from transformers import ViTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import tensorflow as tf

     from transformers import TFViTForImageClassification, TFViTModel


if is_vision_available():
     from PIL import Image

     from transformers import ViTImageProcessor



class        _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''

 def __init__(self    :       Optional[Any]	,			UpperCAmelCase_    :       Optional[int]	,			UpperCAmelCase_    :       List[Any]=13	,			UpperCAmelCase_    :       Dict=30	,			UpperCAmelCase_    :       str=2	,			UpperCAmelCase_    :       Union[str, Any]=3	,			UpperCAmelCase_    :       Union[str, Any]=True	,			UpperCAmelCase_    :       List[str]=True	,			UpperCAmelCase_    :       Dict=32	,			UpperCAmelCase_    :       List[Any]=2	,			UpperCAmelCase_    :       List[str]=4	,			UpperCAmelCase_    :       List[str]=37	,			UpperCAmelCase_    :       List[str]="gelu"	,			UpperCAmelCase_    :       List[Any]=0.1	,			UpperCAmelCase_    :       Optional[Any]=0.1	,			UpperCAmelCase_    :       str=10	,			UpperCAmelCase_    :       str=0.02	,			UpperCAmelCase_    :       Optional[int]=3	,			UpperCAmelCase_    :       Dict=None	,			)	->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__:				int        =parent
   lowerCamelCase__:				Any        =batch_size
   lowerCamelCase__:				List[str]        =image_size
   lowerCamelCase__:				Tuple        =patch_size
   lowerCamelCase__:				Tuple        =num_channels
   lowerCamelCase__:				Tuple        =is_training
   lowerCamelCase__:				Optional[int]        =use_labels
   lowerCamelCase__:				Dict        =hidden_size
   lowerCamelCase__:				Dict        =num_hidden_layers
   lowerCamelCase__:				Optional[Any]        =num_attention_heads
   lowerCamelCase__:				Union[str, Any]        =intermediate_size
   lowerCamelCase__:				List[str]        =hidden_act
   lowerCamelCase__:				Optional[Any]        =hidden_dropout_prob
   lowerCamelCase__:				str        =attention_probs_dropout_prob
   lowerCamelCase__:				str        =type_sequence_label_size
   lowerCamelCase__:				Optional[Any]        =initializer_range
   lowerCamelCase__:				Optional[int]        =scope

   # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
   lowerCamelCase__:				List[str]        =(image_size // patch_size) ** 2
   lowerCamelCase__:				Optional[Any]        =num_patches + 1

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->List[str]:



   '''simple docstring'''



   lowerCamelCase__:				Tuple        =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

   lowerCamelCase__:				List[str]        =None
   if self.use_labels:
     lowerCamelCase__:				Any        =ids_tensor([self.batch_size]	,			self.type_sequence_label_size)

   lowerCamelCase__:				List[str]        =self.get_config()

   return config, pixel_values, labels

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int])	->List[str]:



   '''simple docstring'''



   return ViTConfig(
       image_size=self.image_size	,			patch_size=self.patch_size	,			num_channels=self.num_channels	,			hidden_size=self.hidden_size	,			num_hidden_layers=self.num_hidden_layers	,			num_attention_heads=self.num_attention_heads	,			intermediate_size=self.intermediate_size	,			hidden_act=self.hidden_act	,			hidden_dropout_prob=self.hidden_dropout_prob	,			attention_probs_dropout_prob=self.attention_probs_dropout_prob	,			is_decoder=UpperCAmelCase_	,			initializer_range=self.initializer_range	,			)

 def 							SCREAMING_SNAKE_CASE_ (self    :       int	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Optional[Any]	,			UpperCAmelCase_    :       Tuple)	->Tuple:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =TFViTModel(config=UpperCAmelCase_)
   lowerCamelCase__:				Tuple        =model(UpperCAmelCase_	,			training=UpperCAmelCase_)
   self.parent.assertEqual(result.last_hidden_state.shape	,			(self.batch_size, self.seq_length, self.hidden_size))

   # Test with an image with different size than the one specified in config.
   lowerCamelCase__:				Dict        =self.image_size // 2
   lowerCamelCase__:				List[str]        =pixel_values[:, :, :image_size, :image_size]
   lowerCamelCase__:				List[str]        =model(UpperCAmelCase_	,			interpolate_pos_encoding=UpperCAmelCase_	,			training=UpperCAmelCase_)
   lowerCamelCase__:				List[str]        =(image_size // self.patch_size) ** 2 + 1
   self.parent.assertEqual(result.last_hidden_state.shape	,			(self.batch_size, seq_length, self.hidden_size))

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any]	,			UpperCAmelCase_    :       Tuple	,			UpperCAmelCase_    :       int	,			UpperCAmelCase_    :       Any)	->str:



   '''simple docstring'''



   lowerCamelCase__:				Any        =self.type_sequence_label_size
   lowerCamelCase__:				str        =TFViTForImageClassification(UpperCAmelCase_)
   lowerCamelCase__:				List[str]        =model(UpperCAmelCase_	,			labels=UpperCAmelCase_	,			training=UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.type_sequence_label_size))

   # Test with an image with different size than the one specified in config.
   lowerCamelCase__:				Union[str, Any]        =self.image_size // 2
   lowerCamelCase__:				Dict        =pixel_values[:, :, :image_size, :image_size]
   lowerCamelCase__:				Tuple        =model(UpperCAmelCase_	,			interpolate_pos_encoding=UpperCAmelCase_	,			training=UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.type_sequence_label_size))

   # test greyscale images
   lowerCamelCase__:				int        =1
   lowerCamelCase__:				List[str]        =TFViTForImageClassification(UpperCAmelCase_)
   lowerCamelCase__:				List[str]        =floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
   lowerCamelCase__:				Any        =model(UpperCAmelCase_)
   self.parent.assertEqual(result.logits.shape	,			(self.batch_size, self.type_sequence_label_size))







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



   '''simple docstring'''



   lowerCamelCase__:				Optional[Any]        =self.prepare_config_and_inputs()
   lowerCamelCase__  , lowerCamelCase__  , lowerCamelCase__:				Union[str, Any]        =config_and_inputs
   lowerCamelCase__:				Optional[Any]        ={"pixel_values": pixel_values}
   return config, inputs_dict



@require_tf
class        _SCREAMING_SNAKE_CASE						(     __SCREAMING_SNAKE_CASE   ,							__SCREAMING_SNAKE_CASE   ,							unittest.TestCase   ):







 '''simple docstring'''
 lowercase_		=      (TFViTModel, TFViTForImageClassification) if is_tf_available() else ()
 lowercase_		=      (
     {"feature-extraction": TFViTModel, "image-classification": TFViTForImageClassification}
     if is_tf_available()
     else {}
 )

 lowercase_		=      False
 lowercase_		=      False
 lowercase_		=      False

 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str])	->Dict:



   '''simple docstring'''



   lowerCamelCase__:				int        =TFViTModelTester(self)
   lowerCamelCase__:				List[str]        =ConfigTester(self	,			config_class=UpperCAmelCase_	,			has_text_modality=UpperCAmelCase_	,			hidden_size=37)

 def 							SCREAMING_SNAKE_CASE_ (self    :       str)	->Tuple:



   '''simple docstring'''



   self.config_tester.run_common_tests()

 @unittest.skip(reason="ViT does not use inputs_embeds")
 def 							SCREAMING_SNAKE_CASE_ (self    :       List[str])	->Any:



   '''simple docstring'''



   pass

 @unittest.skip(reason="ViT does not use inputs_embeds")
 def 							SCREAMING_SNAKE_CASE_ (self    :       int)	->str:



   '''simple docstring'''



   pass

 def 							SCREAMING_SNAKE_CASE_ (self    :       Tuple)	->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				int        =self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowerCamelCase__:				int        =model_class(UpperCAmelCase_)
     self.assertIsInstance(model.get_input_embeddings()	,			(tf.keras.layers.Layer))
     lowerCamelCase__:				Tuple        =model.get_output_embeddings()
     self.assertTrue(x is None or isinstance(UpperCAmelCase_	,			tf.keras.layers.Layer))

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



   '''simple docstring'''



   lowerCamelCase__  , lowerCamelCase__:				Optional[Any]        =self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowerCamelCase__:				List[Any]        =model_class(UpperCAmelCase_)
     lowerCamelCase__:				Tuple        =inspect.signature(model.call)
     # signature.parameters is an OrderedDict => so arg_names order is deterministic
     lowerCamelCase__:				str        =[*signature.parameters.keys()]

     lowerCamelCase__:				str        =["pixel_values"]
     self.assertListEqual(arg_names[:1]	,			UpperCAmelCase_)

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



   '''simple docstring'''



   lowerCamelCase__:				Optional[int]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[Any])	->Any:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase_)







 @slow
 def 							SCREAMING_SNAKE_CASE_ (self    :       Dict)	->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:				Any        =TFViTModel.from_pretrained("google/vit-base-patch16-224")
   self.assertIsNotNone(UpperCAmelCase_)





def 					lowerCAmelCase_			(	)					->		Tuple:







  """simple docstring"""





  lowerCamelCase__:				Any        =Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
  return image






@require_tf
@require_vision
class        _SCREAMING_SNAKE_CASE						(     unittest.TestCase   ):







 '''simple docstring'''

 @cached_property
 def 							SCREAMING_SNAKE_CASE_ (self    :       str)	->List[Any]:



   '''simple docstring'''



   return ViTImageProcessor.from_pretrained("google/vit-base-patch16-224") if is_vision_available() else None







 @slow
 def 							SCREAMING_SNAKE_CASE_ (self    :       Optional[int])	->Tuple:



   '''simple docstring'''



   lowerCamelCase__:				List[str]        =TFViTForImageClassification.from_pretrained("google/vit-base-patch16-224")

   lowerCamelCase__:				List[str]        =self.default_image_processor
   lowerCamelCase__:				Union[str, Any]        =prepare_img()
   lowerCamelCase__:				List[Any]        =image_processor(images=UpperCAmelCase_	,			return_tensors="tf")

   # forward pass
   lowerCamelCase__:				Dict        =model(**UpperCAmelCase_)

   # verify the logits
   lowerCamelCase__:				List[Any]        =tf.TensorShape((1, 1_000))
   self.assertEqual(outputs.logits.shape	,			UpperCAmelCase_)

   lowerCamelCase__:				List[Any]        =tf.constant([-0.2744, 0.8215, -0.0836])

   tf.debugging.assert_near(outputs.logits[0, :3]	,			UpperCAmelCase_	,			atol=1E-4)

