


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_ :int  							=			16
a_ :List[str]  							=			32




def    a   (							A__      ,    A__ = 1_6							)	-> Any:




      '''simple docstring'''

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

      def tokenize_function(A__							):
            # max_length=None => use the model max length (it's actually the default)
            SCREAMING_SNAKE_CASE__       :  Union[str, Any]						     =		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():
            SCREAMING_SNAKE_CASE__       :  Any						     =		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
      SCREAMING_SNAKE_CASE__       :  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.
            SCREAMING_SNAKE_CASE__       :  Dict						     =		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__       :  Optional[Any]						     =		1_6
            elif accelerator.mixed_precision != "no":
                  SCREAMING_SNAKE_CASE__       :  List[Any]						     =		8
            else:
                  SCREAMING_SNAKE_CASE__       :  Tuple						     =		None

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

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




def    a   (							A__      ,    A__							)	-> Any:




      '''simple docstring'''

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

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

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

      # Instantiate optimizer
      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		AdamW(params=model.parameters()      ,    lr=__A							)

      # Instantiate scheduler
      SCREAMING_SNAKE_CASE__       :  str						     =		get_linear_schedule_with_warmup(
          optimizer=__A      ,    num_warmup_steps=1_0_0      ,    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.
      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		accelerator.prepare(
          __A      ,    __A      ,    __A      ,    __A      ,    __A							)

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

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

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




def    a   (							)	-> Tuple:




      '''simple docstring'''

      SCREAMING_SNAKE_CASE__       :  Union[str, 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(
          '''--local_sgd_steps'''      ,    type=__A      ,    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.'''							)
      SCREAMING_SNAKE_CASE__       :  Optional[Any]						     =		parser.parse_args()
      SCREAMING_SNAKE_CASE__       :  str						     =		{"lr": 2e-5, "num_epochs": 3, "seed": 4_2, "batch_size": 1_6}
      training_function(__A      ,    __A							)


if __name__ == "__main__":
    main()














from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
			import tensorflow as tf

			from transformers import TFBlipTextModel
			from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST




class 							snake_case_       :

						"""simple docstring"""



						def __init__(self:					Optional[Any] ,					__UpperCAmelCase:					List[str] ,					__UpperCAmelCase:					int=12 ,					__UpperCAmelCase:					Dict=7 ,					__UpperCAmelCase:					str=True ,					__UpperCAmelCase:					Tuple=True ,					__UpperCAmelCase:					Tuple=True ,					__UpperCAmelCase:					Optional[int]=99 ,					__UpperCAmelCase:					str=32 ,					__UpperCAmelCase:					Any=32 ,					__UpperCAmelCase:					int=2 ,					__UpperCAmelCase:					List[str]=4 ,					__UpperCAmelCase:					Tuple=37 ,					__UpperCAmelCase:					Optional[int]=0.1 ,					__UpperCAmelCase:					List[Any]=0.1 ,					__UpperCAmelCase:					Dict=512 ,					__UpperCAmelCase:					Union[str, Any]=0.02 ,					__UpperCAmelCase:					List[Any]=0 ,					__UpperCAmelCase:					str=None ,					)    ->     Tuple:




											'''simple docstring'''



											__a					:						Union[str, Any]      =					parent
											__a					:						Dict      =					batch_size
											__a					:						Optional[Any]      =					seq_length
											__a					:						int      =					is_training
											__a					:						Union[str, Any]      =					use_input_mask
											__a					:						str      =					use_labels
											__a					:						Tuple      =					vocab_size
											__a					:						Dict      =					hidden_size
											__a					:						Optional[int]      =					projection_dim
											__a					:						Tuple      =					num_hidden_layers
											__a					:						Dict      =					num_attention_heads
											__a					:						Any      =					intermediate_size
											__a					:						Optional[int]      =					dropout
											__a					:						Dict      =					attention_dropout
											__a					:						List[Any]      =					max_position_embeddings
											__a					:						Any      =					initializer_range
											__a					:						Tuple      =					scope
											__a					:						Dict      =					bos_token_id



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




											'''simple docstring'''



											__a					:						Optional[Any]      =					ids_tensor([self.batch_size, self.seq_length] ,					self.vocab_size   )

											__a					:						Any      =					None
											if self.use_input_mask:
																__a					:						Tuple      =					random_attention_mask([self.batch_size, self.seq_length]   )

											if input_mask is not None:
																__a					:						Tuple      =					input_mask.numpy()
																__a      ,			__a					:						int      =					input_mask.shape
																__a					:						Dict      =					np.random.randint(1 ,					seq_length - 1 ,					size=(batch_size,)   )
																for batch_idx, start_index in enumerate(__UpperCAmelCase   ):
																					__a					:						Any      =					1
																					__a					:						int      =					0

											__a					:						Union[str, Any]      =					self.get_config()

											return config, input_ids, tf.convert_to_tensor(__UpperCAmelCase   )



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




											'''simple docstring'''



											return BlipTextConfig(
											    vocab_size=self.vocab_size ,					hidden_size=self.hidden_size ,					projection_dim=self.projection_dim ,					num_hidden_layers=self.num_hidden_layers ,					num_attention_heads=self.num_attention_heads ,					intermediate_size=self.intermediate_size ,					dropout=self.dropout ,					attention_dropout=self.attention_dropout ,					max_position_embeddings=self.max_position_embeddings ,					initializer_range=self.initializer_range ,					bos_token_id=self.bos_token_id ,					)



						def        UpperCAmelCase__					(self:					List[str] ,					__UpperCAmelCase:					Any ,					__UpperCAmelCase:					Dict ,					__UpperCAmelCase:					str   )    ->     List[str]:




											'''simple docstring'''



											__a					:						int      =					TFBlipTextModel(config=__UpperCAmelCase   )
											__a					:						List[str]      =					model(__UpperCAmelCase ,					attention_mask=__UpperCAmelCase ,					training=__UpperCAmelCase   )
											__a					:						Any      =					model(__UpperCAmelCase ,					training=__UpperCAmelCase   )
											self.parent.assertEqual(result.last_hidden_state.shape ,					(self.batch_size, self.seq_length, self.hidden_size)   )
											self.parent.assertEqual(result.pooler_output.shape ,					(self.batch_size, self.hidden_size)   )



						def        UpperCAmelCase__					(self:					List[str]   )    ->     str:




											'''simple docstring'''



											__a					:						Tuple      =					self.prepare_config_and_inputs()
											__a      ,			__a      ,			__a					:						Dict      =					config_and_inputs
											__a					:						Union[str, Any]      =					{"input_ids": input_ids, "attention_mask": input_mask}
											return config, inputs_dict




@require_tf
class 							snake_case_       (    __UpperCamelCase		,							unittest.TestCase      ):

						"""simple docstring"""



						snake_case__        =      (TFBlipTextModel,) if is_tf_available() else ()
						snake_case__        =      False
						snake_case__        =      False
						snake_case__        =      False



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




											'''simple docstring'''



											__a					:						str      =					BlipTextModelTester(self   )
											__a					:						str      =					ConfigTester(self ,					config_class=__UpperCAmelCase ,					hidden_size=37   )



						def        UpperCAmelCase__					(self:					str   )    ->     str:




											'''simple docstring'''



											self.config_tester.run_common_tests()



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




											'''simple docstring'''



											__a					:						int      =					self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					Tuple   )    ->     Dict:




											'''simple docstring'''



											pass



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




											'''simple docstring'''



											pass



						@unittest.skip(reason="Blip does not use inputs_embeds"   )
						def        UpperCAmelCase__					(self:					Any   )    ->     Dict:




											'''simple docstring'''



											pass



						@unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"   )
						def        UpperCAmelCase__					(self:					List[Any]   )    ->     List[str]:




											'''simple docstring'''



											pass



						@unittest.skip(reason="BlipTextModel has no base class and is not available in MODEL_MAPPING"   )
						def        UpperCAmelCase__					(self:					Union[str, Any]   )    ->     int:




											'''simple docstring'''



											pass



						@slow
						def        UpperCAmelCase__					(self:					Optional[Any]   )    ->     Tuple:




											'''simple docstring'''



											for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__a					:						List[str]      =					TFBlipTextModel.from_pretrained(__UpperCAmelCase   )
																self.assertIsNotNone(__UpperCAmelCase   )



						def        UpperCAmelCase__					(self:					Dict ,					__UpperCAmelCase:					Optional[Any]=True   )    ->     Any:




											'''simple docstring'''



											super().test_pt_tf_model_equivalence(allow_missing_keys=__UpperCAmelCase   )





