





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


lowercase__				:Optional[Any]				   =       1_6
lowercase__				:Dict				   =       3_2


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_ = 16  )						->Union[str, Any]:





  """simple docstring"""


  __UpperCAmelCase    :						Dict       =						AutoTokenizer.from_pretrained('''bert-base-cased'''  )
  __UpperCAmelCase    :						Any       =						load_dataset('''glue'''  ,			'''mrpc'''  )

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

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

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

  # Instantiate dataloaders.
  __UpperCAmelCase    :						Union[str, Any]       =						DataLoader(
      tokenized_datasets['''train''']  ,			shuffle=UpperCAmelCase_  ,			collate_fn=UpperCAmelCase_  ,			batch_size=UpperCAmelCase_  )
  __UpperCAmelCase    :						List[Any]       =						DataLoader(
      tokenized_datasets['''validation''']  ,			shuffle=UpperCAmelCase_  ,			collate_fn=UpperCAmelCase_  ,			batch_size=UpperCAmelCase_  )

  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

   lowercase__				:Tuple				   =       mocked_dataloaders  # noqa: F811


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->str:





  """simple docstring"""


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

  __UpperCAmelCase    :						Optional[int]       =						evaluate.load('''glue'''  ,			'''mrpc'''  )

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

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

  # Instantiate optimizer
  __UpperCAmelCase    :						int       =						AdamW(params=model.parameters()  ,			lr=UpperCAmelCase_  )

  # Instantiate scheduler
  __UpperCAmelCase    :						str       =						get_linear_schedule_with_warmup(
      optimizer=UpperCAmelCase_  ,			num_warmup_steps=1_00  ,			num_training_steps=(len(UpperCAmelCase_  ) * 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    :						Optional[Any]       =						accelerator.prepare(
      UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  ,			UpperCAmelCase_  )

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

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

    __UpperCAmelCase    :						List[Any]       =						metric.compute()
    # Use accelerator.print to print only on the main process.
    accelerator.print(f'''epoch {epoch}:'''  ,			UpperCAmelCase_  )


def       lowerCamelCase_       (   )						->List[str]:





  """simple docstring"""


  __UpperCAmelCase    :						Union[str, Any]       =						argparse.ArgumentParser(description='''Simple example of training script.'''  )
  parser.add_argument(
      '''--mixed_precision'''  ,			type=UpperCAmelCase_  ,			default=UpperCAmelCase_  ,			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=UpperCAmelCase_  ,			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    :						Tuple       =						parser.parse_args()
  __UpperCAmelCase    :						Tuple       =						{'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
  training_function(UpperCAmelCase_  ,			UpperCAmelCase_  )


if __name__ == "__main__":
   main()






"""simple docstring"""





from dataclasses import dataclass, field
from typing import Tuple

from ..utils import cached_property, is_tf_available, logging, requires_backends
from .benchmark_args_utils import BenchmarkArguments


if is_tf_available():
  import tensorflow as tf


lowercase__				:Optional[int]				   =       logging.get_logger(__name__)




@dataclass
class 							snake_case							(			__UpperCAmelCase ):

      '''simple docstring'''



      _A     :       Tuple         =       [
          'no_inference',
          'no_cuda',
          'no_tpu',
          'no_speed',
          'no_memory',
          'no_env_print',
          'no_multi_process',
      ]

      def __init__(							self       :      Dict						,				**__lowercase       :      Dict ):


        '''simple docstring'''




        for deprecated_arg in self.deprecated_args:
          if deprecated_arg in kwargs:
            __UpperCAmelCase    :						Optional[Any]       =						deprecated_arg[3:]
            __UpperCAmelCase    :						List[str]       =						not kwargs.pop(__lowercase )
            logger.warning(
                f'''{deprecated_arg} is depreciated. Please use --no-{positive_arg} or'''
                f''' {positive_arg}={kwargs[positive_arg]}''' )
        __UpperCAmelCase    :						int       =						kwargs.pop('''tpu_name'''						,				self.tpu_name )
        __UpperCAmelCase    :						Any       =						kwargs.pop('''device_idx'''						,				self.device_idx )
        __UpperCAmelCase    :						Tuple       =						kwargs.pop('''eager_mode'''						,				self.eager_mode )
        __UpperCAmelCase    :						Optional[int]       =						kwargs.pop('''use_xla'''						,				self.use_xla )
        super().__init__(**__lowercase )

      _A     :       str         =       field(
          default=__UpperCAmelCase      ,  metadata={'help': 'Name of TPU'}      ,  )
      _A     :       int         =       field(
          default=0      ,  metadata={'help': 'CPU / GPU device index. Defaults to 0.'}      ,  )
      _A     :       bool         =       field(default=__UpperCAmelCase      ,  metadata={'help': 'Benchmark models in eager model.'} )
      _A     :       bool         =       field(
          default=__UpperCAmelCase      ,  metadata={
              'help': 'Benchmark models using XLA JIT compilation. Note that `eager_model` has to be set to `False`.'
          }      ,  )

      @cached_property
      def       A_					(							self       :      int ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        __UpperCAmelCase    :						Any       =						None
        if self.tpu:
          try:
            if self.tpu_name:
              __UpperCAmelCase    :						List[str]       =						tf.distribute.cluster_resolver.TPUClusterResolver(self.tpu_name )
            else:
              __UpperCAmelCase    :						Any       =						tf.distribute.cluster_resolver.TPUClusterResolver()
          except ValueError:
            __UpperCAmelCase    :						int       =						None
        return tpu

      @cached_property
      def       A_					(							self       :      Any ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        if self.is_tpu:
          tf.config.experimental_connect_to_cluster(self._setup_tpu )
          tf.tpu.experimental.initialize_tpu_system(self._setup_tpu )

          __UpperCAmelCase    :						Any       =						tf.distribute.TPUStrategy(self._setup_tpu )
        else:
          # currently no multi gpu is allowed
          if self.is_gpu:
            # TODO: Currently only single GPU is supported
            tf.config.set_visible_devices(self.gpu_list[self.device_idx]						,				'''GPU''' )
            __UpperCAmelCase    :						Any       =						tf.distribute.OneDeviceStrategy(device=f'''/gpu:{self.device_idx}''' )
          else:
            tf.config.set_visible_devices([]						,				'''GPU''' )  # disable GPU
            __UpperCAmelCase    :						str       =						tf.distribute.OneDeviceStrategy(device=f'''/cpu:{self.device_idx}''' )

        return strategy

      @property
      def       A_					(							self       :      Optional[int] ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        return self._setup_tpu is not None

      @property
      def       A_					(							self       :      List[str] ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        return self._setup_strategy

      @property
      def       A_					(							self       :      Any ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        return tf.config.list_physical_devices('''GPU''' )

      @property
      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




        requires_backends(self						,				['''tf'''] )
        if self.cuda:
          return len(self.gpu_list )
        return 0







      @property
      def       A_					(							self       :      Optional[Any] ):


        '''simple docstring'''




        return self.n_gpu > 0