





import unittest

from transformers import MPNetConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          MPNetForMaskedLM,
          MPNetForMultipleChoice,
          MPNetForQuestionAnswering,
          MPNetForSequenceClassification,
          MPNetForTokenClassification,
          MPNetModel,
      )





class UpperCAmelCase_				:


  def __init__(							self ,__snake_case ,__snake_case=1_3 ,__snake_case=7 ,__snake_case=True ,__snake_case=True ,__snake_case=False ,__snake_case=True ,__snake_case=9_9 ,__snake_case=6_4 ,__snake_case=5 ,__snake_case=4 ,__snake_case=6_4 ,__snake_case="gelu" ,__snake_case=0.1 ,__snake_case=0.1 ,__snake_case=5_1_2 ,__snake_case=1_6 ,__snake_case=2 ,__snake_case=0.02 ,__snake_case=3 ,__snake_case=4 ,__snake_case=None ,):



       """simple docstring"""





       A_          =       parent
       A_          =       batch_size
       A_          =       seq_length
       A_          =       is_training
       A_          =       use_input_mask
       A_          =       use_token_type_ids
       A_          =       use_labels
       A_          =       vocab_size
       A_          =       hidden_size
       A_          =       num_hidden_layers
       A_          =       num_attention_heads
       A_          =       intermediate_size
       A_          =       hidden_act
       A_          =       hidden_dropout_prob
       A_          =       attention_probs_dropout_prob
       A_          =       max_position_embeddings
       A_          =       type_vocab_size
       A_          =       type_sequence_label_size
       A_          =       initializer_range
       A_          =       num_labels
       A_          =       num_choices
       A_          =       scope


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       return MPNetConfig.from_pretrained('''microsoft/mpnet-base''' )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size )

       A_          =       None
       if self.use_input_mask:
            A_          =       random_attention_mask([self.batch_size, self.seq_length] )

       A_          =       None
       A_          =       None
       A_          =       None
       if self.use_labels:
            A_          =       ids_tensor([self.batch_size] ,self.type_sequence_label_size )
            A_          =       ids_tensor([self.batch_size, self.seq_length] ,self.num_labels )
            A_          =       ids_tensor([self.batch_size] ,self.num_choices )

       A_          =       self.get_config()
       return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       return MPNetConfig(
           vocab_size=self.vocab_size ,hidden_size=self.hidden_size ,num_hidden_layers=self.num_hidden_layers ,num_attention_heads=self.num_attention_heads ,intermediate_size=self.intermediate_size ,hidden_act=self.hidden_act ,hidden_dropout_prob=self.hidden_dropout_prob ,attention_probs_dropout_prob=self.attention_probs_dropout_prob ,max_position_embeddings=self.max_position_embeddings ,initializer_range=self.initializer_range ,)


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       MPNetModel(config=__snake_case )
       model.to(__snake_case )
       model.eval()
       A_          =       model(__snake_case ,__snake_case )
       A_          =       model(__snake_case )
       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 ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       MPNetForQuestionAnswering(config=__snake_case )
       model.to(__snake_case )
       model.eval()
       A_          =       model(
           __snake_case ,attention_mask=__snake_case ,start_positions=__snake_case ,end_positions=__snake_case ,)
       self.parent.assertEqual(result.start_logits.shape ,(self.batch_size, self.seq_length) )
       self.parent.assertEqual(result.end_logits.shape ,(self.batch_size, self.seq_length) )


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       self.num_labels
       A_          =       MPNetForSequenceClassification(__snake_case )
       model.to(__snake_case )
       model.eval()
       A_          =       model(__snake_case ,attention_mask=__snake_case ,labels=__snake_case )
       self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_labels) )


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       self.num_choices
       A_          =       MPNetForMultipleChoice(config=__snake_case )
       model.to(__snake_case )
       model.eval()
       A_          =       input_ids.unsqueeze(1 ).expand(-1 ,self.num_choices ,-1 ).contiguous()
       A_          =       input_mask.unsqueeze(1 ).expand(-1 ,self.num_choices ,-1 ).contiguous()
       A_          =       model(
           __snake_case ,attention_mask=__snake_case ,labels=__snake_case ,)
       self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.num_choices) )


  def 						__UpperCAmelCase     (							self ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ,__snake_case ):



       """simple docstring"""





       A_          =       self.num_labels
       A_          =       MPNetForTokenClassification(config=__snake_case )
       model.to(__snake_case )
       model.eval()
       A_          =       model(__snake_case ,attention_mask=__snake_case ,labels=__snake_case )
       self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.num_labels) )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.prepare_config_and_inputs()
       ((A_)   ,    (A_)   ,    (A_)   ,    (A_)   ,    (A_)   ,    (A_))          =       config_and_inputs
       A_          =       {'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict





@require_torch
class UpperCAmelCase_				( _SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		unittest.TestCase ):
  __lowerCAmelCase   :						Optional[Any]							=   (
      (
          MPNetForMaskedLM,
          MPNetForMultipleChoice,
          MPNetForQuestionAnswering,
          MPNetForSequenceClassification,
          MPNetForTokenClassification,
          MPNetModel,
      )
      if is_torch_available()
      else ()
  )
  __lowerCAmelCase   :						Optional[int]							=   (
      {
          """feature-extraction""": MPNetModel,
          """fill-mask""": MPNetForMaskedLM,
          """question-answering""": MPNetForQuestionAnswering,
          """text-classification""": MPNetForSequenceClassification,
          """token-classification""": MPNetForTokenClassification,
          """zero-shot""": MPNetForSequenceClassification,
      }
      if is_torch_available()
      else {}
  )
  __lowerCAmelCase   :						Optional[Any]							=   False
  __lowerCAmelCase   :						Optional[int]							=   True


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       MPNetModelTester(self )
       A_          =       ConfigTester(self ,config_class=__snake_case ,hidden_size=3_7 )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       self.config_tester.run_common_tests()


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mpnet_model(*__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mpnet_for_sequence_classification(*__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mpnet_for_multiple_choice(*__snake_case )


  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mpnet_for_token_classification(*__snake_case )




  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_mpnet_for_question_answering(*__snake_case )





@require_torch
class UpperCAmelCase_				( unittest.TestCase ):


  @slow
  def 						__UpperCAmelCase     (							self ):



       """simple docstring"""





       A_          =       MPNetModel.from_pretrained('''microsoft/mpnet-base''' )
       A_          =       torch.tensor([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]] )
       A_          =       model(__snake_case )[0]
       A_          =       torch.Size((1, 1_1, 7_6_8) )
       self.assertEqual(output.shape ,__snake_case )
       A_          =       torch.tensor(
           [[[-0.0550, 0.1943, -0.0740], [-0.0562, 0.2211, -0.0579], [-0.0437, 0.3337, -0.0641]]] )
       # compare the actual values for a slice.
       self.assertTrue(torch.allclose(output[:, :3, :3] ,__snake_case ,atol=1E-4 ) )










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__      :     Any         =		16
a__      :     str         =		32

def      UpperCAmelCase_    (    _UpperCAmelCase							:Accelerator  ,						_UpperCAmelCase							:int = 16				)	->				int:
     '''simple docstring'''


     A_          =       AutoTokenizer.from_pretrained('''bert-base-cased'''				)
     A_          =       load_dataset('''glue'''  ,						'''mrpc'''				)

     def tokenize_function(_UpperCAmelCase							:Optional[Any]				):
          # max_length=None => use the model max length (it's actually the default)
          A_          =       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():
          A_          =       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
     A_          =       tokenized_datasets.rename_column('''label'''  ,						'''labels'''				)

     def collate_fn(_UpperCAmelCase							:List[Any]				):
          # On TPU it's best to pad everything to the same length or training will be very slow.
          A_          =       128 if accelerator.distributed_type == DistributedType.TPU else None
          # When using mixed precision we want round multiples of 8/16
          if accelerator.mixed_precision == "fp8":
               A_          =       16
          elif accelerator.mixed_precision != "no":
               A_          =       8
          else:
               A_          =       None

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

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

def      UpperCAmelCase_    (    _UpperCAmelCase							:List[str]  ,						_UpperCAmelCase							:Dict				)	->				Dict:
     '''simple docstring'''


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

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

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

     # Instantiate optimizer
     A_          =       AdamW(params=model.parameters()  ,						lr=_UpperCAmelCase				)

     # Instantiate scheduler
     A_          =       get_linear_schedule_with_warmup(
         optimizer=_UpperCAmelCase  ,						num_warmup_steps=100  ,						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.
     A_   ,    A_   ,    A_   ,    A_   ,    A_          =       accelerator.prepare(
         _UpperCAmelCase  ,						_UpperCAmelCase  ,						_UpperCAmelCase  ,						_UpperCAmelCase  ,						_UpperCAmelCase				)

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

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

          A_          =       metric.compute()
          # Use accelerator.print to print only on the main process.
          accelerator.print(f'epoch {epoch}:'  ,						_UpperCAmelCase				)

def      UpperCAmelCase_    (    )	->				Union[str, Any]:
     '''simple docstring'''


     A_          =       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(
         '''--local_sgd_steps'''  ,						type=_UpperCAmelCase  ,						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.'''				)
     A_          =       parser.parse_args()
     A_          =       {'''lr''': 2e-5, '''num_epochs''': 3, '''seed''': 42, '''batch_size''': 16}
     training_function(_UpperCAmelCase  ,						_UpperCAmelCase				)


if __name__ == "__main__":
      main()



