




'''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
from accelerate.local_sgd import LocalSGD


########################################################################
# This is a fully working simple example to use Accelerate
# with LocalSGD, which is a method to synchronize model
# parameters every K batches. It is different, but complementary
# to 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 						snake_case_							(_a	:							Accelerator    ,   _a	:							int = 1_6     ):
    UpperCAmelCase			       =		AutoTokenizer.from_pretrained('''bert-base-cased'''     )
    UpperCAmelCase			       =		load_dataset('''glue'''    ,   '''mrpc'''     )

    def tokenize_function(_a	:							str     ):
        # max_length=None => use the model max length (it's actually the default)
        UpperCAmelCase			       =		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():
        UpperCAmelCase			       =		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
    UpperCAmelCase			       =		tokenized_datasets.rename_column('''label'''    ,   '''labels'''     )

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

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

    # Instantiate dataloaders.
    UpperCAmelCase			       =		DataLoader(
        tokenized_datasets['''train''']    ,   shuffle=__lowerCamelCase    ,   collate_fn=__lowerCamelCase    ,   batch_size=__lowerCamelCase     )
    UpperCAmelCase			       =		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     =mocked_dataloaders  # noqa: F811





def 						snake_case_							(_a	:							Optional[int]    ,   _a	:							Union[str, Any]     ):
    # For testing only
    if os.environ.get('''TESTING_MOCKED_DATALOADERS'''    ,   __lowerCamelCase     ) == "1":
        UpperCAmelCase			       =		2
    # New Code #
    UpperCAmelCase			       =		int(args.gradient_accumulation_steps     )
    UpperCAmelCase			       =		int(args.local_sgd_steps     )
    # Initialize accelerator
    UpperCAmelCase			       =		Accelerator(
        cpu=args.cpu    ,   mixed_precision=args.mixed_precision    ,   gradient_accumulation_steps=__lowerCamelCase     )
    if accelerator.distributed_type not in [DistributedType.NO, DistributedType.MULTI_CPU, DistributedType.MULTI_GPU]:
        raise NotImplementedError('''LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)'''     )
    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
    UpperCAmelCase			       =		config['''lr''']
    UpperCAmelCase			       =		int(config['''num_epochs''']     )
    UpperCAmelCase			       =		int(config['''seed''']     )
    UpperCAmelCase			       =		int(config['''batch_size''']     )

    UpperCAmelCase			       =		evaluate.load('''glue'''    ,   '''mrpc'''     )

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

    # Instantiate optimizer
    UpperCAmelCase			       =		AdamW(params=model.parameters()    ,   lr=__lowerCamelCase     )

    # Instantiate scheduler
    UpperCAmelCase			       =		get_linear_schedule_with_warmup(
        optimizer=__lowerCamelCase    ,   num_warmup_steps=1_0_0    ,   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.
    UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase       ,     UpperCAmelCase			       =		accelerator.prepare(
        __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase     )

    # Now we train the model
    for epoch in range(__lowerCamelCase     ):
        model.train()
        with LocalSGD(
            accelerator=__lowerCamelCase    ,   model=__lowerCamelCase    ,   local_sgd_steps=__lowerCamelCase    ,   enabled=local_sgd_steps is not None     ) as local_sgd:
            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     ):
                    UpperCAmelCase			       =		model(**__lowerCamelCase     )
                    UpperCAmelCase			       =		output.loss
                    accelerator.backward(__lowerCamelCase     )
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                    # LocalSGD-specific line
                    local_sgd.step()

        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():
                UpperCAmelCase			       =		model(**__lowerCamelCase     )
            UpperCAmelCase			       =		outputs.logits.argmax(dim=-1     )
            UpperCAmelCase       ,     UpperCAmelCase			       =		accelerator.gather_for_metrics((predictions, batch['''labels'''])     )
            metric.add_batch(
                predictions=__lowerCamelCase    ,   references=__lowerCamelCase    ,   )

        UpperCAmelCase			       =		metric.compute()
        # Use accelerator.print to print only on the main process.
        accelerator.print(F"epoch {epoch}:"    ,   __lowerCamelCase     )





def 						snake_case_							():
    UpperCAmelCase			       =		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(
        '''--local_sgd_steps'''    ,   type=__lowerCamelCase    ,   default=8    ,   help='''Number of local SGD steps or None to disable local SGD'''     )
    parser.add_argument('''--cpu'''    ,   action='''store_true'''    ,   help='''If passed, will train on the CPU.'''     )
    UpperCAmelCase			       =		parser.parse_args()
    UpperCAmelCase			       =		{'''lr''': 2E-5, '''num_epochs''': 3, '''seed''': 4_2, '''batch_size''': 1_6}
    training_function(__lowerCamelCase    ,   __lowerCamelCase     )


if __name__ == "__main__":
   main()






"""simple docstring"""



import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
 from PIL import Image

 from transformers import CLIPImageProcessor, CLIPProcessor

@require_vision
class       lowerCAmelCase__				(						unittest.TestCase	):







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       tempfile.mkdtemp()

       # fmt: off
       _snake_case					     =       ['''l''', '''o''', '''w''', '''e''', '''r''', '''s''', '''t''', '''i''', '''d''', '''n''', '''lo''', '''l</w>''', '''w</w>''', '''r</w>''', '''t</w>''', '''low</w>''', '''er</w>''', '''lowest</w>''', '''newer</w>''', '''wider''', '''<unk>''', '''<|startoftext|>''', '''<|endoftext|>''']
       # fmt: on
       _snake_case					     =       dict(zip(_lowerCamelCase    ,      range(len(_lowerCamelCase			)			)			)			)
       _snake_case					     =       ['''#version: 0.2''', '''l o''', '''lo w</w>''', '''e r</w>''', '''''']
       _snake_case					     =       {'''unk_token''': '''<unk>'''}

       _snake_case					     =       os.path.join(self.tmpdirname    ,      VOCAB_FILES_NAMES['''vocab_file''']			)
       _snake_case					     =       os.path.join(self.tmpdirname    ,      VOCAB_FILES_NAMES['''merges_file''']			)
       with open(self.vocab_file    ,      '''w'''    ,      encoding='''utf-8'''			) as fp:
         fp.write(json.dumps(_lowerCamelCase			) + '''\n'''			)
       with open(self.merges_file    ,      '''w'''    ,      encoding='''utf-8'''			) as fp:
         fp.write('''\n'''.join(_lowerCamelCase			)			)

       _snake_case					     =       {
           '''do_resize''': True,
           '''size''': 20,
           '''do_center_crop''': True,
           '''crop_size''': 18,
           '''do_normalize''': True,
           '''image_mean''': [0.4_8_1_4_5_4_6_6, 0.4_5_7_8_2_7_5, 0.4_0_8_2_1_0_7_3],
           '''image_std''': [0.2_6_8_6_2_9_5_4, 0.2_6_1_3_0_2_5_8, 0.2_7_5_7_7_7_1_1],
       }
       _snake_case					     =       os.path.join(self.tmpdirname    ,      _lowerCamelCase			)
       with open(self.image_processor_file    ,      '''w'''    ,      encoding='''utf-8'''			) as fp:
         json.dump(_lowerCamelCase    ,      _lowerCamelCase			)







     def      lowercase	(							self			:			Tuple    ,      **_lowerCamelCase			:			Any			):
       return CLIPTokenizer.from_pretrained(self.tmpdirname    ,      **_lowerCamelCase			)







     def      lowercase	(							self			:			str    ,      **_lowerCamelCase			:			Any			):
       return CLIPTokenizerFast.from_pretrained(self.tmpdirname    ,      **_lowerCamelCase			)







     def      lowercase	(							self			:			int    ,      **_lowerCamelCase			:			Optional[int]			):
       return CLIPImageProcessor.from_pretrained(self.tmpdirname    ,      **_lowerCamelCase			)







     def      lowercase	(							self			:			Union[str, Any]			):
       shutil.rmtree(self.tmpdirname			)







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       [np.random.randint(255    ,      size=(3, 30, 400)    ,      dtype=np.uinta			)]

       _snake_case					     =       [Image.fromarray(np.moveaxis(_lowerCamelCase    ,      0    ,      -1			)			) for x in image_inputs]

       return image_inputs







     def      lowercase	(							self			:			Optional[Any]			):
       _snake_case					     =       self.get_tokenizer()
       _snake_case					     =       self.get_rust_tokenizer()
       _snake_case					     =       self.get_image_processor()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)
       processor_slow.save_pretrained(self.tmpdirname			)
       _snake_case					     =       CLIPProcessor.from_pretrained(self.tmpdirname    ,      use_fast=_lowerCamelCase			)

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)
       processor_fast.save_pretrained(self.tmpdirname			)
       _snake_case					     =       CLIPProcessor.from_pretrained(self.tmpdirname			)

       self.assertEqual(processor_slow.tokenizer.get_vocab()    ,      tokenizer_slow.get_vocab()			)
       self.assertEqual(processor_fast.tokenizer.get_vocab()    ,      tokenizer_fast.get_vocab()			)
       self.assertEqual(tokenizer_slow.get_vocab()    ,      tokenizer_fast.get_vocab()			)
       self.assertIsInstance(processor_slow.tokenizer    ,      _lowerCamelCase			)
       self.assertIsInstance(processor_fast.tokenizer    ,      _lowerCamelCase			)

       self.assertEqual(processor_slow.image_processor.to_json_string()    ,      image_processor.to_json_string()			)
       self.assertEqual(processor_fast.image_processor.to_json_string()    ,      image_processor.to_json_string()			)
       self.assertIsInstance(processor_slow.image_processor    ,      _lowerCamelCase			)
       self.assertIsInstance(processor_fast.image_processor    ,      _lowerCamelCase			)







     def      lowercase	(							self			:			List[Any]			):
       _snake_case					     =       CLIPProcessor(tokenizer=self.get_tokenizer()    ,      image_processor=self.get_image_processor()			)
       processor.save_pretrained(self.tmpdirname			)

       _snake_case					     =       self.get_tokenizer(bos_token='''(BOS)'''    ,      eos_token='''(EOS)'''			)
       _snake_case					     =       self.get_image_processor(do_normalize=_lowerCamelCase    ,      padding_value=1.0			)

       _snake_case					     =       CLIPProcessor.from_pretrained(
           self.tmpdirname    ,      bos_token='''(BOS)'''    ,      eos_token='''(EOS)'''    ,      do_normalize=_lowerCamelCase    ,      padding_value=1.0			)

       self.assertEqual(processor.tokenizer.get_vocab()    ,      tokenizer_add_kwargs.get_vocab()			)
       self.assertIsInstance(processor.tokenizer    ,      _lowerCamelCase			)

       self.assertEqual(processor.image_processor.to_json_string()    ,      image_processor_add_kwargs.to_json_string()			)
       self.assertIsInstance(processor.image_processor    ,      _lowerCamelCase			)







     def      lowercase	(							self			:			int			):
       _snake_case					     =       self.get_image_processor()
       _snake_case					     =       self.get_tokenizer()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)

       _snake_case					     =       self.prepare_image_inputs()

       _snake_case					     =       image_processor(_lowerCamelCase    ,      return_tensors='''np'''			)
       _snake_case					     =       processor(images=_lowerCamelCase    ,      return_tensors='''np'''			)

       for key in input_image_proc.keys():
         self.assertAlmostEqual(input_image_proc[key].sum()    ,      input_processor[key].sum()    ,      delta=1e-2			)







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       self.get_image_processor()
       _snake_case					     =       self.get_tokenizer()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)

       _snake_case					     =       '''lower newer'''

       _snake_case					     =       processor(text=_lowerCamelCase			)

       _snake_case					     =       tokenizer(_lowerCamelCase			)

       for key in encoded_tok.keys():
         self.assertListEqual(encoded_tok[key]    ,      encoded_processor[key]			)







     def      lowercase	(							self			:			Any			):
       _snake_case					     =       self.get_image_processor()
       _snake_case					     =       self.get_tokenizer()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)

       _snake_case					     =       '''lower newer'''
       _snake_case					     =       self.prepare_image_inputs()

       _snake_case					     =       processor(text=_lowerCamelCase    ,      images=_lowerCamelCase			)

       self.assertListEqual(list(inputs.keys()			)    ,      ['''input_ids''', '''attention_mask''', '''pixel_values''']			)

       # test if it raises when no input is passed
       with pytest.raises(_lowerCamelCase			):
         processor()







     def      lowercase	(							self			:			List[str]			):
       _snake_case					     =       self.get_image_processor()
       _snake_case					     =       self.get_tokenizer()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)

       _snake_case					     =       [[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

       _snake_case					     =       processor.batch_decode(_lowerCamelCase			)
       _snake_case					     =       tokenizer.batch_decode(_lowerCamelCase			)

       self.assertListEqual(_lowerCamelCase    ,      _lowerCamelCase			)







     def      lowercase	(							self			:			List[Any]			):
       _snake_case					     =       self.get_image_processor()
       _snake_case					     =       self.get_tokenizer()

       _snake_case					     =       CLIPProcessor(tokenizer=_lowerCamelCase    ,      image_processor=_lowerCamelCase			)

       _snake_case					     =       '''lower newer'''
       _snake_case					     =       self.prepare_image_inputs()

       _snake_case					     =       processor(text=_lowerCamelCase    ,      images=_lowerCamelCase			)

       self.assertListEqual(list(inputs.keys()			)    ,      processor.model_input_names			)



