

"""simple docstring"""




from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed



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




   """simple docstring"""
   for param, grad_param in zip(model_a.parameters()	,    model_b.parameters()   ):
      if not param.requires_grad:
         continue
      if not did_step:
         # Grads should not be in sync
         assert (
             torch.allclose(param.grad	,    grad_param.grad   ) is False
         ), f'Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'
      else:
         # Grads should be in sync
         assert (
             torch.allclose(param.grad	,    grad_param.grad   ) is True
         ), f'Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'



def 	UpperCAmelCase__   (				A__	,    A__	,    A__	,    A__	,    A__=True   )      ->      Dict:




   """simple docstring"""
   model.train()
   lowerCamelCase__						=			model(A__   )
   lowerCamelCase__						=			F.mse_loss(A__	,    target.to(output.device   )   )
   if not do_backward:
      loss /= accelerator.gradient_accumulation_steps
      loss.backward()
   else:
      accelerator.backward(A__   )



def 	UpperCAmelCase__   (				A__	,    A__=False   )      ->      Optional[int]:




   """simple docstring"""
   set_seed(42   )
   lowerCamelCase__						=			RegressionModel()
   lowerCamelCase__						=			deepcopy(A__   )
   lowerCamelCase__						=			RegressionDataset(length=80   )
   lowerCamelCase__						=			DataLoader(A__	,    batch_size=16   )
   model.to(accelerator.device   )
   if sched:
      lowerCamelCase__						=			AdamW(params=model.parameters()	,    lr=1E-3   )
      lowerCamelCase__						=			AdamW(params=ddp_model.parameters()	,    lr=1E-3   )
      lowerCamelCase__						=			LambdaLR(A__	,    lr_lambda=lambda A__   : epoch**0.65   )
      lowerCamelCase__						=			LambdaLR(A__	,    lr_lambda=lambda A__   : epoch**0.65   )
   # Make a copy of `model`
   if sched:
      lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			accelerator.prepare(A__	,    A__	,    A__	,    A__   )
   else:
      lowerCamelCase__					,       lowerCamelCase__						=			accelerator.prepare(A__	,    A__   )
   if sched:
      return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
   return model, ddp_model, dataloader



def 	UpperCAmelCase__   (				A__   )      ->      Any:




   """simple docstring"""
   # Test when on a single CPU or GPU that the context manager does nothing
   lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			get_training_setup(A__   )
   # Use a single batch
   lowerCamelCase__					,       lowerCamelCase__						=			next(iter(A__   )   ).values()
   for iteration in range(3   ):
      # Gather the distributed inputs and targs for the base model
      lowerCamelCase__					,       lowerCamelCase__						=			accelerator.gather((ddp_input, ddp_target)   )
      lowerCamelCase__					,       lowerCamelCase__						=			input.to(accelerator.device   ), target.to(accelerator.device   )
      # Perform our initial ground truth step in non "DDP"
      step_model(A__	,    A__	,    A__	,    A__   )
      # Do "gradient accumulation" (noop)
      if iteration % 2 == 0:
         # Accumulate grads locally
         with accelerator.no_sync(A__   ):
            step_model(A__	,    A__	,    A__	,    A__   )
      else:
         # Sync grads
         step_model(A__	,    A__	,    A__	,    A__   )

      # Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
      check_model_parameters(A__	,    A__	,    A__	,    A__   )
      for param, ddp_param in zip(model.parameters()	,    ddp_model.parameters()   ):
         if not param.requires_grad:
            continue
         assert torch.allclose(
             param.grad	,    ddp_param.grad   ), f'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

      # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration   )
      lowerCamelCase__						=			ddp_input[torch.randperm(len(A__   )   )]



def 	UpperCAmelCase__   (				A__   )      ->      Optional[Any]:




   """simple docstring"""
   # Test on distributed setup that context manager behaves properly
   lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			get_training_setup(A__   )
   # Use a single batch
   lowerCamelCase__					,       lowerCamelCase__						=			next(iter(A__   )   ).values()
   for iteration in range(3   ):
      # Gather the distributed inputs and targs for the base model
      lowerCamelCase__					,       lowerCamelCase__						=			accelerator.gather((ddp_input, ddp_target)   )
      lowerCamelCase__					,       lowerCamelCase__						=			input.to(accelerator.device   ), target.to(accelerator.device   )
      # Perform our initial ground truth step in non "DDP"
      step_model(A__	,    A__	,    A__	,    A__   )
      # Do "gradient accumulation" (noop)
      if iteration % 2 == 0:
         # Accumulate grads locally
         with accelerator.no_sync(A__   ):
            step_model(A__	,    A__	,    A__	,    A__   )
      else:
         # Sync grads
         step_model(A__	,    A__	,    A__	,    A__   )

      # DDP model and model should only be in sync when not (iteration % 2 == 0)
      for param, ddp_param in zip(model.parameters()	,    ddp_model.parameters()   ):
         if not param.requires_grad:
            continue
         if iteration % 2 == 0:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad	,    ddp_param.grad   ) is False
            ), f'Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'
         else:
            # Grads should be in sync
            assert (
                torch.allclose(param.grad	,    ddp_param.grad   ) is True
            ), f'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration   )
      lowerCamelCase__						=			ddp_input[torch.randperm(len(A__   )   )]



def 	UpperCAmelCase__   (				A__=False	,    A__=False   )      ->      Any:




   """simple docstring"""
   lowerCamelCase__						=			Accelerator(
       split_batches=A__	,    dispatch_batches=A__	,    gradient_accumulation_steps=2   )
   # Test that context manager behaves properly
   lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			get_training_setup(A__   )
   for iteration, batch in enumerate(A__   ):
      lowerCamelCase__					,       lowerCamelCase__						=			batch.values()
      # Gather the distributed inputs and targs for the base model
      lowerCamelCase__					,       lowerCamelCase__						=			accelerator.gather((ddp_input, ddp_target)   )
      lowerCamelCase__					,       lowerCamelCase__						=			input.to(accelerator.device   ), target.to(accelerator.device   )
      # Perform our initial ground truth step in non "DDP"
      step_model(A__	,    A__	,    A__	,    A__	,    A__   )
      # Do "gradient accumulation" (noop)
      with accelerator.accumulate(A__   ):
         step_model(A__	,    A__	,    A__	,    A__   )

      # DDP model and model should only be in sync when not (iteration % 2 == 0)
      for param, ddp_param in zip(model.parameters()	,    ddp_model.parameters()   ):
         if not param.requires_grad:
            continue
         if ((iteration + 1) % 2 == 0) or (iteration == len(A__   ) - 1):
            # Grads should be in sync
            assert (
                torch.allclose(param.grad	,    ddp_param.grad   ) is True
            ), f'Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'
         else:
            # Grads should not be in sync
            assert (
                torch.allclose(param.grad	,    ddp_param.grad   ) is False
            ), f'Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration   )
      lowerCamelCase__						=			ddp_input[torch.randperm(len(A__   )   )]
   GradientState._reset_state()



def 	UpperCAmelCase__   (				A__=False	,    A__=False   )      ->      Optional[int]:




   """simple docstring"""
   lowerCamelCase__						=			Accelerator(
       split_batches=A__	,    dispatch_batches=A__	,    gradient_accumulation_steps=2   )
   # Test that context manager behaves properly
   lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			get_training_setup(A__	,    A__   )
   for iteration, batch in enumerate(A__   ):
      lowerCamelCase__					,       lowerCamelCase__						=			batch.values()
      # Gather the distributed inputs and targs for the base model
      lowerCamelCase__					,       lowerCamelCase__						=			accelerator.gather((ddp_input, ddp_target)   )
      lowerCamelCase__					,       lowerCamelCase__						=			input.to(accelerator.device   ), target.to(accelerator.device   )
      # Perform our initial ground truth step in non "DDP"
      model.train()
      ddp_model.train()
      step_model(A__	,    A__	,    A__	,    A__	,    A__   )
      opt.step()

      if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(A__   )):
         if split_batches:
            sched.step()
         else:
            for _ in range(accelerator.num_processes   ):
               sched.step()
      opt.zero_grad()
      # Perform gradient accumulation under wrapper
      with accelerator.accumulate(A__   ):
         step_model(A__	,    A__	,    A__	,    A__   )
         ddp_opt.step()
         ddp_sched.step()
         ddp_opt.zero_grad()

      # Learning rates should be the same
      assert (
          opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
      ), f'Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'
      lowerCamelCase__						=			(((iteration + 1) % 2) == 0) or ((iteration + 1) == len(A__   ))
      if accelerator.num_processes > 1:
         check_model_parameters(A__	,    A__	,    A__	,    A__   )
      # Shuffle ddp_input on each iteration
      torch.manual_seed(1337 + iteration   )
   GradientState._reset_state()



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




   """simple docstring"""
   lowerCamelCase__						=			Accelerator()

   lowerCamelCase__						=			RegressionDataset(length=80   )
   lowerCamelCase__						=			DataLoader(A__	,    batch_size=16   )
   lowerCamelCase__						=			RegressionDataset(length=96   )
   lowerCamelCase__						=			DataLoader(A__	,    batch_size=16   )
   lowerCamelCase__					,       lowerCamelCase__						=			accelerator.prepare(A__	,    A__   )
   assert accelerator.gradient_state.active_dataloader is None
   for iteration, _ in enumerate(A__   ):
      assert id(accelerator.gradient_state.active_dataloader   ) == id(A__   )
      if iteration < len(A__   ) - 1:
         assert not accelerator.gradient_state.end_of_dataloader
         if iteration == 1:
            for batch_num, _ in enumerate(A__   ):
               assert id(accelerator.gradient_state.active_dataloader   ) == id(A__   )
               if batch_num < len(A__   ) - 1:
                  assert not accelerator.gradient_state.end_of_dataloader
               else:
                  assert accelerator.gradient_state.end_of_dataloader
      else:
         assert accelerator.gradient_state.end_of_dataloader
   assert accelerator.gradient_state.active_dataloader is None



def 	UpperCAmelCase__   (				)      ->      List[str]:




   """simple docstring"""
   lowerCamelCase__						=			Accelerator()
   lowerCamelCase__						=			accelerator.state
   if state.local_process_index == 0:
      print("**Test `accumulate` gradient accumulation with dataloader break**"   )
   test_dataloader_break()
   if state.distributed_type == DistributedType.NO:
      if state.local_process_index == 0:
         print("**Test NOOP `no_sync` context manager**"   )
      test_noop_sync(A__   )
   if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
      if state.local_process_index == 0:
         print("**Test Distributed `no_sync` context manager**"   )
      test_distributed_sync(A__   )
   if state.distributed_type == DistributedType.MULTI_GPU:
      for split_batch in [True, False]:
         for dispatch_batches in [True, False]:
            if state.local_process_index == 0:
               print(
                   "**Test `accumulate` gradient accumulation, "	,    f'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'	,    )
            test_gradient_accumulation(A__	,    A__   )

    # Currently will break on torch 2.0 +, need to investigate why
   if is_torch_version("<"	,    "2.0"   ) or state.distributed_type == DistributedType.NO:
      if state.local_process_index == 0:
         print(
             "**Test `accumulate` gradient accumulation with optimizer and scheduler, "	,    "`split_batches=False`, `dispatch_batches=False`**"	,    )
      test_gradient_accumulation_with_opt_and_scheduler()
      if state.distributed_type == DistributedType.MULTI_GPU:
         for split_batch in [True, False]:
            for dispatch_batches in [True, False]:
               if not split_batch and not dispatch_batches:
                  continue
               if state.local_process_index == 0:
                  print(
                      "**Test `accumulate` gradient accumulation with optimizer and scheduler, "	,    f'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'	,    )
               test_gradient_accumulation_with_opt_and_scheduler(A__	,    A__   )



def 	UpperCAmelCase__   (				A__   )      ->      str:




   """simple docstring"""
   # For xla_spawn (TPUs)
   main()


if __name__ == "__main__":
     main()







"""simple docstring"""




import os
import unittest

from transformers.models.phobert.tokenization_phobert import VOCAB_FILES_NAMES, PhobertTokenizer

from ...test_tokenization_common import TokenizerTesterMixin


class 		_A       (			__a				,       unittest.TestCase					):
      __a								=	PhobertTokenizer
      __a								=	False





      def 					_lowerCamelCase  (  self		)   ->					Any:
         super().setUp()

         # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
         lowerCamelCase__						=			["T@@", "i", "I", "R@@", "r", "e@@"]
         lowerCamelCase__						=			dict(zip(SCREAMING_SNAKE_CASE__    ,						range(len(SCREAMING_SNAKE_CASE__		)		)		)		)
         lowerCamelCase__						=			["#version: 0.2", "l à</w>"]
         lowerCamelCase__						=			{"unk_token": "<unk>"}

         lowerCamelCase__						=			os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["vocab_file"]		)
         lowerCamelCase__						=			os.path.join(self.tmpdirname    ,						VOCAB_FILES_NAMES["merges_file"]		)

         with open(self.vocab_file    ,						"w"    ,						encoding="utf-8"		) as fp:
            for token in vocab_tokens:
               fp.write(f'{token} {vocab_tokens[token]}\n'		)
         with open(self.merges_file    ,						"w"    ,						encoding="utf-8"		) as fp:
            fp.write("\n".join(SCREAMING_SNAKE_CASE__		)		)





      def 					_lowerCamelCase  (  self    ,						**SCREAMING_SNAKE_CASE__		)   ->					Tuple:
         kwargs.update(self.special_tokens_map		)
         return PhobertTokenizer.from_pretrained(self.tmpdirname    ,						**SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__		)   ->					Tuple:
         lowerCamelCase__						=			"Tôi là VinAI Research"
         lowerCamelCase__						=			"T<unk> i <unk> <unk> <unk> <unk> <unk> <unk> I Re<unk> e<unk> <unk> <unk> <unk>"
         return input_text, output_text





      def 					_lowerCamelCase  (  self		)   ->					Tuple:
         lowerCamelCase__						=			PhobertTokenizer(self.vocab_file    ,						self.merges_file    ,						**self.special_tokens_map		)
         lowerCamelCase__						=			"Tôi là VinAI Research"
         lowerCamelCase__						=			"T@@ ô@@ i l@@ à V@@ i@@ n@@ A@@ I R@@ e@@ s@@ e@@ a@@ r@@ c@@ h".split()
         lowerCamelCase__						=			tokenizer.tokenize(SCREAMING_SNAKE_CASE__		)
         print(SCREAMING_SNAKE_CASE__		)
         self.assertListEqual(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			tokens + [tokenizer.unk_token]

         lowerCamelCase__						=			[4, 3, 5, 3, 3, 3, 3, 3, 3, 6, 7, 9, 3, 9, 3, 3, 3, 3, 3]
         self.assertListEqual(tokenizer.convert_tokens_to_ids(SCREAMING_SNAKE_CASE__		)    ,						SCREAMING_SNAKE_CASE__		)




