





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


A__							:  Any    =16
A__							:  Tuple    =32
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase = 16    ):





      """simple docstring"""


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

      def tokenize_function(lowerCAmelCase    ):
            # max_length=None => use the model max length (it's actually the default)
            _lowerCAmelCase					=				tokenizer(examples["""sentence1"""]				,      examples["""sentence2"""]				,      truncation=lowerCAmelCase				,      max_length=lowerCAmelCase    )
            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(
                lowerCAmelCase				,      batched=lowerCAmelCase				,      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(lowerCAmelCase    ):
            # On TPU it's best to pad everything to the same length or training will be very slow.
            _lowerCAmelCase					=				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":
                  _lowerCAmelCase					=				16
            elif accelerator.mixed_precision != "no":
                  _lowerCAmelCase					=				8
            else:
                  _lowerCAmelCase					=				None

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

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

      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__							:  int    =mocked_dataloaders  # noqa: F811
def   UpperCamelCase__ (						lowerCAmelCase				,      lowerCAmelCase    ):





      """simple docstring"""


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

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

      set_seed(lowerCAmelCase    )
      _lowerCAmelCase	,				_lowerCAmelCase					=				get_dataloaders(lowerCAmelCase				,      lowerCAmelCase    )
      # 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=lowerCAmelCase    )

      # 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=lowerCAmelCase    )

      # Instantiate scheduler
      _lowerCAmelCase					=				get_linear_schedule_with_warmup(
          optimizer=lowerCAmelCase				,      num_warmup_steps=1_00				,      num_training_steps=(len(lowerCAmelCase    ) * 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(
          lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase				,      lowerCAmelCase    )

      # Now we train the model
      for epoch in range(lowerCAmelCase    ):
            model.train()
            for step, batch in enumerate(lowerCAmelCase    ):
                  # 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(lowerCAmelCase    ):
                        _lowerCAmelCase					=				model(**lowerCAmelCase    )
                        _lowerCAmelCase					=				output.loss
                        accelerator.backward(lowerCAmelCase    )
                        optimizer.step()
                        lr_scheduler.step()
                        optimizer.zero_grad()

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

            _lowerCAmelCase					=				metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(f"epoch {epoch}:"				,      lowerCAmelCase    )
def   UpperCamelCase__ (						):





      """simple docstring"""


      _lowerCAmelCase					=				argparse.ArgumentParser(description="""Simple example of training script."""    )
      parser.add_argument(
          """--mixed_precision"""				,      type=lowerCAmelCase				,      default=lowerCAmelCase				,      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=lowerCAmelCase				,      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					=				parser.parse_args()
      _lowerCAmelCase					=				{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
      training_function(lowerCAmelCase				,      lowerCAmelCase    )


if __name__ == "__main__":
      main()













'''simple docstring'''






def   UpperCamelCase__ (						lowerCAmelCase    ):





      """simple docstring"""


      if not isinstance(lowerCAmelCase				,      lowerCAmelCase    ):
            raise ValueError("""Input must be an integer"""    )
      if input_num <= 0:
            raise ValueError("""Input must be positive"""    )
      return sum(
          divisor for divisor in range(1				,      input_num // 2 + 1    ) if input_num % divisor == 0    )


if __name__ == "__main__":
      import doctest

      doctest.testmod()






