



import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment
def 						UpperCamelCase      (  )      ->	Optional[Any]:
   UpperCamelCase       :    Union[str, Any]		      =      torch.nn.Linear(2			,	4      )
   UpperCamelCase       :    Optional[int]		      =      torch.optim.AdamW(model.parameters()			,	lr=1.0      )
   UpperCamelCase       :    str		      =      torch.optim.lr_scheduler.OneCycleLR(snake_case__			,	max_lr=0.01			,	steps_per_epoch=2			,	epochs=1      )
   UpperCamelCase       :    Dict		      =      DataLoader(TensorDataset(torch.tensor([1, 2, 3]      )      )      )
   UpperCamelCase       :    Optional[Any]		      =      DataLoader(TensorDataset(torch.tensor([4, 5, 6]      )      )      )

   return model, optimizer, scheduler, train_dl, valid_dl
def 						UpperCamelCase      (  snake_case__			:       List[str]      )      ->	Any:
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def 						UpperCamelCase      (  snake_case__			:       List[str]      )      ->	Union[str, Any]:
   UpperCamelCase       :    Tuple		      =      torch.nn.Linear(*tuple(model.weight.T.shape      )      ).state_dict()
   model.load_state_dict(snake_case__      )



class 							lowerCAmelCase_						(    a__				):
       @require_cuda
       def    snake_case_			(     self			)	->						Optional[Any]:
          UpperCamelCase       :    Optional[Any]		      =      Accelerator()
          assert PartialState._shared_state["_cpu"] is False
          assert PartialState._shared_state["device"].type == "cuda"
          with self.assertRaises(SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    List[Any]		      =      Accelerator(cpu=SCREAMING_SNAKE_CASE_			)
       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase       :    Optional[int]		      =      Accelerator()
          UpperCamelCase       :    Any		      =      GradientState()
          assert state.num_steps == 1
          UpperCamelCase       :    Union[str, Any]		      =      4
          assert state.num_steps == 4

          assert state.sync_gradients is True
          UpperCamelCase       :    Dict		      =      False
          assert state.sync_gradients is False
          GradientState._reset_state()
       def    snake_case_			(     self			)	->						Dict:
          UpperCamelCase       :    Optional[Any]		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Optional[int]		      =      create_components()

          (
              (
              UpperCamelCase
          )      ,	(
              UpperCamelCase
          )      ,	(
              UpperCamelCase
          )      ,	(
              UpperCamelCase
          )      ,	(
              UpperCamelCase
          )      ,	
          )       :    int		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)

          self.assertTrue(prepared_model in accelerator._models			)
          self.assertTrue(prepared_optimizer in accelerator._optimizers			)
          self.assertTrue(prepared_scheduler in accelerator._schedulers			)
          self.assertTrue(prepared_train_dl in accelerator._dataloaders			)
          self.assertTrue(prepared_valid_dl in accelerator._dataloaders			)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    Any		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Optional[int]		      =      create_components()
          accelerator.prepare(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
          accelerator.free_memory()

          self.assertTrue(len(accelerator._models			) == 0			)
          self.assertTrue(len(accelerator._optimizers			) == 0			)
          self.assertTrue(len(accelerator._schedulers			) == 0			)
          self.assertTrue(len(accelerator._dataloaders			) == 0			)
       def    snake_case_			(     self			)	->						List[str]:
          PartialState._reset_state()

          # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
          def noop(*SCREAMING_SNAKE_CASE_,     **SCREAMING_SNAKE_CASE_			):
             pass

          with patch('torch.cuda.set_device',     SCREAMING_SNAKE_CASE_			), patch_environment(ACCELERATE_TORCH_DEVICE='cuda:64'			):
             UpperCamelCase       :    int		      =      Accelerator()
             self.assertEqual(str(accelerator.state.device			),     'cuda:64'			)
       def    snake_case_			(     self			)	->						Optional[Any]:
          UpperCamelCase       :    int		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Tuple		      =      create_components()
          accelerator.prepare(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    Dict		      =      get_signature(SCREAMING_SNAKE_CASE_			)

          with tempfile.TemporaryDirectory() as tmpdirname:
             accelerator.save_state(SCREAMING_SNAKE_CASE_			)

             # make sure random weights don't match
             load_random_weights(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) > 1e-3			)

             # make sure loaded weights match
             accelerator.load_state(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) < 1e-3			)
       def    snake_case_			(     self			)	->						Any:
          UpperCamelCase       :    int		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    str		      =      create_components()
          accelerator.prepare(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)

          UpperCamelCase       :    int		      =      get_signature(SCREAMING_SNAKE_CASE_			)

          # saving hook
          def save_config(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Tuple		      =      {'class_name': models[0].__class__.__name__}

             with open(os.path.join(SCREAMING_SNAKE_CASE_,     'data.json'			),     'w'			) as f:
                json.dump(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)

        # loading hook
          def load_config(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			):
             with open(os.path.join(SCREAMING_SNAKE_CASE_,     'data.json'			),     'r'			) as f:
                UpperCamelCase       :    str		      =      json.load(SCREAMING_SNAKE_CASE_			)

             UpperCamelCase       :    Optional[Any]		      =      config['class_name']

          UpperCamelCase       :    Union[str, Any]		      =      accelerator.register_save_state_pre_hook(SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    List[str]		      =      accelerator.register_load_state_pre_hook(SCREAMING_SNAKE_CASE_			)

          with tempfile.TemporaryDirectory() as tmpdirname:
             accelerator.save_state(SCREAMING_SNAKE_CASE_			)

             # make sure random weights don't match with hooks
             load_random_weights(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) > 1e-3			)

             # random class name to verify correct one is loaded
             UpperCamelCase       :    List[Any]		      =      'random'

             # make sure loaded weights match with hooks
             accelerator.load_state(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) < 1e-3			)

             # mode.class_name is loaded from config
             self.assertTrue(model.class_name == model.__class__.__name__			)

          # remove hooks
          save_hook.remove()
          load_hook.remove()

          with tempfile.TemporaryDirectory() as tmpdirname:
             accelerator.save_state(SCREAMING_SNAKE_CASE_			)

             # make sure random weights don't match with hooks removed
             load_random_weights(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) > 1e-3			)

             # random class name to verify correct one is loaded
             UpperCamelCase       :    List[str]		      =      'random'

             # make sure loaded weights match with hooks removed
             accelerator.load_state(SCREAMING_SNAKE_CASE_			)
             self.assertTrue(abs(model_signature - get_signature(SCREAMING_SNAKE_CASE_			)			) < 1e-3			)

             # mode.class_name is NOT loaded from config
             self.assertTrue(model.class_name != model.__class__.__name__			)
       def    snake_case_			(     self			)	->						Optional[int]:
          UpperCamelCase       :    str		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    str		      =      create_components()
          UpperCamelCase       :    Optional[int]		      =      None

          # This should work
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Tuple		      =      accelerator.prepare(
              SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
          self.assertTrue(dummy_obj is None			)
       def    snake_case_			(     self			)	->						str:
          UpperCamelCase       :    Optional[int]		      =      Accelerator()
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Dict		      =      create_components()
          UpperCamelCase       :    Union[str, Any]		      =      [1, 2, 3]

          # This should work
          UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase      ,	UpperCamelCase       :    Optional[Any]		      =      accelerator.prepare(
              SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Dummy object should have `_is_accelerate_prepared` set to `True`',     )
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Model is missing `_is_accelerator_prepared` or is set to `False`',     )
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Optimizer is missing `_is_accelerator_prepared` or is set to `False`',     )
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Scheduler is missing `_is_accelerator_prepared` or is set to `False`',     )
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`',     )
          self.assertEqual(
              getattr(SCREAMING_SNAKE_CASE_,     '_is_accelerate_prepared',     SCREAMING_SNAKE_CASE_			),     SCREAMING_SNAKE_CASE_,     'Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`',     )
       @slow
       @require_bnb
       def    snake_case_			(     self			)	->						List[Any]:
          from transformers import AutoModelForCausalLM

          UpperCamelCase       :    str		      =      AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m',     load_in_abit=SCREAMING_SNAKE_CASE_,     device_map={'': 0},     )
          UpperCamelCase       :    Optional[int]		      =      Accelerator()

          # This should work
          UpperCamelCase       :    Tuple		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_			)
       @slow
       @require_bnb
       def    snake_case_			(     self			)	->						Any:
          from transformers import AutoModelForCausalLM

          UpperCamelCase       :    Dict		      =      Accelerator()

          with init_empty_weights():
             UpperCamelCase       :    int		      =      AutoModelForCausalLM.from_pretrained(
                 'EleutherAI/gpt-neo-125m',     )
             model.tie_weights()
             UpperCamelCase       :    List[str]		      =      infer_auto_device_map(SCREAMING_SNAKE_CASE_			)
             UpperCamelCase       :    int		      =      'cpu'

          UpperCamelCase       :    str		      =      AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m',     device_map=SCREAMING_SNAKE_CASE_,     load_in_abit=SCREAMING_SNAKE_CASE_,     llm_inta_enable_fpaa_cpu_offload=SCREAMING_SNAKE_CASE_			)

          # This should not work and get value error
          with self.assertRaises(SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Union[str, Any]		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_			)
       @slow
       @require_bnb
       @require_multi_gpu
       def    snake_case_			(     self			)	->						Union[str, Any]:
          from transformers import AutoModelForCausalLM

          UpperCamelCase       :    Tuple		      =      {'distributed_type': DistributedType.MULTI_GPU}

          with init_empty_weights():
             UpperCamelCase       :    Union[str, Any]		      =      AutoModelForCausalLM.from_pretrained(
                 'EleutherAI/gpt-neo-125m',     )
             model.tie_weights()
             UpperCamelCase       :    Dict		      =      infer_auto_device_map(SCREAMING_SNAKE_CASE_			)
             UpperCamelCase       :    Dict		      =      1

          UpperCamelCase       :    str		      =      AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m',     load_in_abit=SCREAMING_SNAKE_CASE_,     device_map=SCREAMING_SNAKE_CASE_,     )
          UpperCamelCase       :    Optional[int]		      =      Accelerator()

          # This should not work and get value error
          with self.assertRaises(SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    str		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_			)

          PartialState._reset_state()
       @slow
       @require_bnb
       @require_multi_gpu
       def    snake_case_			(     self			)	->						int:
          from transformers import AutoModelForCausalLM

          with init_empty_weights():
             UpperCamelCase       :    Tuple		      =      AutoModelForCausalLM.from_pretrained(
                 'EleutherAI/gpt-neo-125m',     )
             UpperCamelCase       :    str		      =      infer_auto_device_map(SCREAMING_SNAKE_CASE_			)
             UpperCamelCase       :    Optional[int]		      =      1

          UpperCamelCase       :    Tuple		      =      AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m',     load_in_abit=SCREAMING_SNAKE_CASE_,     device_map=SCREAMING_SNAKE_CASE_,     )
          UpperCamelCase       :    Optional[Any]		      =      Accelerator()

          # This should work
          UpperCamelCase       :    List[Any]		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_			)






       @require_cuda
       def    snake_case_			(     self			)	->						int:
          UpperCamelCase       :    Tuple		      =      torch.nn.Linear(10,     10			)
          UpperCamelCase       :    Any		      =      torch.optim.SGD(model.parameters(),     lr=0.01			)
          UpperCamelCase       :    Union[str, Any]		      =      Accelerator(cpu=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Union[str, Any]		      =      accelerator.prepare(SCREAMING_SNAKE_CASE_			)








'''simple docstring'''





import os
import unittest

from huggingface_hub.utils import are_progress_bars_disabled

import transformers.models.bart.tokenization_bart
from transformers import logging
from transformers.testing_utils import CaptureLogger, mockenv, mockenv_context
from transformers.utils.logging import disable_progress_bar, enable_progress_bar







class        UpperCAmelCase_			(			unittest.TestCase      ):
    """simple docstring"""







    def lowerCamelCase			(       self   :	Optional[int]		):
          snake_case__    :       int       		= logging.get_logger()

          # the current default level is logging.WARNING
          snake_case__    :       str       		= logging.get_verbosity()

          logging.set_verbosity_error()
          self.assertEqual(logger.getEffectiveLevel()					,	logging.get_verbosity()		)

          logging.set_verbosity_warning()
          self.assertEqual(logger.getEffectiveLevel()					,	logging.get_verbosity()		)

          logging.set_verbosity_info()
          self.assertEqual(logger.getEffectiveLevel()					,	logging.get_verbosity()		)

          logging.set_verbosity_debug()
          self.assertEqual(logger.getEffectiveLevel()					,	logging.get_verbosity()		)

          # restore to the original level
          logging.set_verbosity(snake_case_		)




    def lowerCamelCase			(       self   :	Union[str, Any]		):
          snake_case__    :       List[Any]       		= logging.get_verbosity()

          snake_case__    :       Any       		= logging.get_logger("""transformers.models.bart.tokenization_bart"""		)
          snake_case__    :       int       		= """Testing 1, 2, 3"""

          # should be able to log warnings (if default settings weren't overridden by `pytest --log-level-all`)
          if level_origin <= logging.WARNING:
                with CaptureLogger(snake_case_		) as cl:
                      logger.warning(snake_case_		)
                self.assertEqual(cl.out					,	msg + """\n"""		)

          # this is setting the level for all of `transformers.*` loggers
          logging.set_verbosity_error()

          # should not be able to log warnings
          with CaptureLogger(snake_case_		) as cl:
                logger.warning(snake_case_		)
          self.assertEqual(cl.out					,	""""""		)

          # should be able to log warnings again
          logging.set_verbosity_warning()
          with CaptureLogger(snake_case_		) as cl:
                logger.warning(snake_case_		)
          self.assertEqual(cl.out					,	msg + """\n"""		)

          # restore to the original level
          logging.set_verbosity(snake_case_		)




    @mockenv(TRANSFORMERS_VERBOSITY="""error"""		)
    def lowerCamelCase			(       self   :	Any		):
          # reset for the env var to take effect, next time some logger call is made
          transformers.utils.logging._reset_library_root_logger()
          # this action activates the env var
          snake_case__    :       Union[str, Any]       		= logging.get_logger("""transformers.models.bart.tokenization_bart"""		)

          snake_case__    :       int       		= os.getenv("""TRANSFORMERS_VERBOSITY"""					,	snake_case_		)
          snake_case__    :       List[Any]       		= logging.log_levels[env_level_str]

          snake_case__    :       str       		= logging.get_verbosity()
          self.assertEqual(
              snake_case_					,	snake_case_					,	f"TRANSFORMERS_VERBOSITY={env_level_str}/{env_level}, but internal verbosity is {current_level}"					,	)

          # restore to the original level
          snake_case__    :       List[Any]       		= """"""
          transformers.utils.logging._reset_library_root_logger()




    @mockenv(TRANSFORMERS_VERBOSITY="""super-error"""		)
    def lowerCamelCase			(       self   :	int		):
          # reset for the env var to take effect, next time some logger call is made
          transformers.utils.logging._reset_library_root_logger()
          snake_case__    :       List[Any]       		= logging.logging.getLogger()
          with CaptureLogger(snake_case_		) as cl:
                # this action activates the env var
                logging.get_logger("""transformers.models.bart.tokenization_bart"""		)
          self.assertIn("""Unknown option TRANSFORMERS_VERBOSITY=super-error"""					,	cl.out		)

          # no need to restore as nothing was changed




    def lowerCamelCase			(       self   :	int		):
          # testing `logger.warning_advice()`
          transformers.utils.logging._reset_library_root_logger()

          snake_case__    :       Optional[int]       		= logging.get_logger("""transformers.models.bart.tokenization_bart"""		)
          snake_case__    :       Tuple       		= """Testing 1, 2, 3"""

          with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS="""1"""		):
                # nothing should be logged as env var disables this method
                with CaptureLogger(snake_case_		) as cl:
                      logger.warning_advice(snake_case_		)
                self.assertEqual(cl.out					,	""""""		)

          with mockenv_context(TRANSFORMERS_NO_ADVISORY_WARNINGS=""""""		):
                # should log normally as TRANSFORMERS_NO_ADVISORY_WARNINGS is unset
                with CaptureLogger(snake_case_		) as cl:
                      logger.warning_advice(snake_case_		)
                self.assertEqual(cl.out					,	msg + """\n"""		)





def 		__snake_case(							)						->					List[str]:
      disable_progress_bar()
      assert are_progress_bars_disabled()

      enable_progress_bar()
      assert not are_progress_bars_disabled()
