






'''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_					= 1_6
a_					= 3_2





def  _a(  UpperCamelCase__				:						Accelerator,	UpperCamelCase__				:						int = 1_6							):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				List[Any]									=AutoTokenizer.from_pretrained('''bert-base-cased'''							)
      SCREAMING_SNAKE_CASE__     :				Dict									=load_dataset('''glue''',	'''mrpc'''							)

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

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

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

      # Instantiate dataloaders.
      SCREAMING_SNAKE_CASE__     :				List[Any]									=DataLoader(
          tokenized_datasets['''train'''],	shuffle=UpperCamelCase__,	collate_fn=UpperCamelCase__,	batch_size=UpperCamelCase__							)
      SCREAMING_SNAKE_CASE__     :				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

  a_					= mocked_dataloaders  # noqa: F811





def  _a(  UpperCamelCase__				:						List[Any],	UpperCamelCase__				:						List[Any]							):




      '''simple docstring'''






      if os.environ.get('''TESTING_MOCKED_DATALOADERS''',	UpperCamelCase__							) == "1":
            SCREAMING_SNAKE_CASE__     :				Optional[int]									=2
      # New Code #
      SCREAMING_SNAKE_CASE__     :				int									=int(args.gradient_accumulation_steps							)
      # Initialize accelerator
      SCREAMING_SNAKE_CASE__     :				Dict									=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
      SCREAMING_SNAKE_CASE__     :				Optional[int]									=config['''lr''']
      SCREAMING_SNAKE_CASE__     :				Dict									=int(config['''num_epochs''']							)
      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=int(config['''seed''']							)
      SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=int(config['''batch_size''']							)

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

      set_seed(UpperCamelCase__							)
      SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=get_dataloaders(UpperCamelCase__,	UpperCamelCase__							)
      # Instantiate the model (we build the model here so that the seed also control new weights initialization)
      SCREAMING_SNAKE_CASE__     :				List[Any]									=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).
      SCREAMING_SNAKE_CASE__     :				str									=model.to(accelerator.device							)

      # Instantiate optimizer
      SCREAMING_SNAKE_CASE__     :				Dict									=AdamW(params=model.parameters(),	lr=UpperCamelCase__							)

      # Instantiate scheduler
      SCREAMING_SNAKE_CASE__     :				Optional[Any]									=get_linear_schedule_with_warmup(
          optimizer=UpperCamelCase__,	num_warmup_steps=1_0_0,	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.
      SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__     :				Tuple									=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__							):
                        SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=model(**UpperCamelCase__							)
                        SCREAMING_SNAKE_CASE__     :				Union[str, 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():
                        SCREAMING_SNAKE_CASE__     :				Tuple									=model(**UpperCamelCase__							)
                  SCREAMING_SNAKE_CASE__     :				str									=outputs.logits.argmax(dim=-1							)
                  SCREAMING_SNAKE_CASE__		,	SCREAMING_SNAKE_CASE__     :				int									=accelerator.gather_for_metrics((predictions, batch['''labels'''])							)
                  metric.add_batch(
                      predictions=UpperCamelCase__,	references=UpperCamelCase__,	)

            SCREAMING_SNAKE_CASE__     :				List[str]									=metric.compute()
            # Use accelerator.print to print only on the main process.
            accelerator.print(f"epoch {epoch}:",	UpperCamelCase__							)





def  _a(  ):




      '''simple docstring'''






      SCREAMING_SNAKE_CASE__     :				List[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.'''							)
      SCREAMING_SNAKE_CASE__     :				List[Any]									=parser.parse_args()
      SCREAMING_SNAKE_CASE__     :				Dict									={'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
      training_function(UpperCamelCase__,	UpperCamelCase__							)


if __name__ == "__main__":
  main()







'''simple docstring'''
import unittest

from transformers import is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
  from PIL import Image
else:






  class      __SCREAMING_SNAKE_CASE      :




   @staticmethod
   def      __magic_name__	(      *__lowercase	:							int					,      **__lowercase	:							Optional[Any]      )	->						Optional[Any]:
         pass






@is_pipeline_test
@require_vision
class      __SCREAMING_SNAKE_CASE      (			unittest.TestCase				):




 @require_torch
 def      __magic_name__	(      self	:							Optional[int]      )	->						Tuple:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=pipeline(
           model='''hf-internal-testing/tiny-random-clip-zero-shot-image-classification'''					,      )
       SCREAMING_SNAKE_CASE__     :				List[str]									=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
       SCREAMING_SNAKE_CASE__     :				Any									=image_classifier(__lowercase					,      candidate_labels=['''a''', '''b''', '''c''']      )

       # The floating scores are so close, we enter floating error approximation and the order is not guaranteed across
       # python and torch versions.
       self.assertIn(
           nested_simplify(__lowercase      )					,      [
               [{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''b'''}, {'''score''': 0.333, '''label''': '''c'''}],
               [{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''c'''}, {'''score''': 0.333, '''label''': '''b'''}],
           ]					,      )

       SCREAMING_SNAKE_CASE__     :				int									=image_classifier([image] * 5					,      candidate_labels=['''A''', '''B''', '''C''']					,      batch_size=2      )
       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
           ]					,      )




 @require_tf
 def      __magic_name__	(      self	:							Union[str, Any]      )	->						Optional[int]:
       SCREAMING_SNAKE_CASE__     :				List[Any]									=pipeline(
           model='''hf-internal-testing/tiny-random-clip-zero-shot-image-classification'''					,      framework='''tf'''      )
       SCREAMING_SNAKE_CASE__     :				Dict									=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
       SCREAMING_SNAKE_CASE__     :				Tuple									=image_classifier(__lowercase					,      candidate_labels=['''a''', '''b''', '''c''']      )

       self.assertEqual(
           nested_simplify(__lowercase      )					,      [{'''score''': 0.333, '''label''': '''a'''}, {'''score''': 0.333, '''label''': '''b'''}, {'''score''': 0.333, '''label''': '''c'''}]					,      )

       SCREAMING_SNAKE_CASE__     :				List[Any]									=image_classifier([image] * 5					,      candidate_labels=['''A''', '''B''', '''C''']					,      batch_size=2      )
       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
               [
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
                   {'''score''': 0.333, '''label''': ANY(__lowercase      )},
               ],
           ]					,      )




 @slow
 @require_torch
 def      __magic_name__	(      self	:							List[Any]      )	->						Dict:
       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=pipeline(
           task='''zero-shot-image-classification'''					,      model='''openai/clip-vit-base-patch32'''					,      )
       # This is an image of 2 cats with remotes and no planes
       SCREAMING_SNAKE_CASE__     :				Any									=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
       SCREAMING_SNAKE_CASE__     :				List[str]									=image_classifier(__lowercase					,      candidate_labels=['''cat''', '''plane''', '''remote''']      )

       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               {'''score''': 0.511, '''label''': '''remote'''},
               {'''score''': 0.485, '''label''': '''cat'''},
               {'''score''': 0.004, '''label''': '''plane'''},
           ]					,      )

       SCREAMING_SNAKE_CASE__     :				Union[str, Any]									=image_classifier([image] * 5					,      candidate_labels=['''cat''', '''plane''', '''remote''']					,      batch_size=2      )
       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               [
                   {'''score''': 0.511, '''label''': '''remote'''},
                   {'''score''': 0.485, '''label''': '''cat'''},
                   {'''score''': 0.004, '''label''': '''plane'''},
               ],
           ]
           * 5					,      )




 @slow
 @require_tf
 def      __magic_name__	(      self	:							Union[str, Any]      )	->						int:
       SCREAMING_SNAKE_CASE__     :				str									=pipeline(
           task='''zero-shot-image-classification'''					,      model='''openai/clip-vit-base-patch32'''					,      framework='''tf'''      )
       # This is an image of 2 cats with remotes and no planes
       SCREAMING_SNAKE_CASE__     :				Optional[Any]									=Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''      )
       SCREAMING_SNAKE_CASE__     :				Any									=image_classifier(__lowercase					,      candidate_labels=['''cat''', '''plane''', '''remote''']      )
       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               {'''score''': 0.511, '''label''': '''remote'''},
               {'''score''': 0.485, '''label''': '''cat'''},
               {'''score''': 0.004, '''label''': '''plane'''},
           ]					,      )

       SCREAMING_SNAKE_CASE__     :				List[Any]									=image_classifier([image] * 5					,      candidate_labels=['''cat''', '''plane''', '''remote''']					,      batch_size=2      )
       self.assertEqual(
           nested_simplify(__lowercase      )					,      [
               [
                   {'''score''': 0.511, '''label''': '''remote'''},
                   {'''score''': 0.485, '''label''': '''cat'''},
                   {'''score''': 0.004, '''label''': '''plane'''},
               ],
           ]
           * 5					,      )