'''simple docstring'''

import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from huggingface_hub import HfFolder, Repository, create_repo, delete_repo
from requests.exceptions import HTTPError

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoProcessor,
    AutoTokenizer,
    BertTokenizer,
    ProcessorMixin,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
)
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test
from transformers.tokenization_utils import TOKENIZER_CONFIG_FILE
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_tokenizers_available


sys.path.append(str(Path(__file__).parent.parent.parent.parent / '''utils'''))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402
from test_module.custom_processing import CustomProcessor  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


A_     :						List[Any]    							=get_tests_dir('''fixtures/dummy_feature_extractor_config.json''')
A_     :						str    							=get_tests_dir('''fixtures/vocab.json''')
A_     :						int    							=get_tests_dir('''fixtures''')


class     __UpperCAmelCase       ( unittest.TestCase		):
  __A :				Any											= ['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]', 'bla', 'blou']





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      0





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h''' )
     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_			=      WavaVecaConfig()
        lowerCAmelCase_			=      AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h''' )

        # save in new folder
        model_config.save_pretrained(_lowerCamelCase )
        processor.save_pretrained(_lowerCamelCase )

        lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )

     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     with tempfile.TemporaryDirectory() as tmpdirname:
        # copy relevant files
        copyfile(_lowerCamelCase    ,     os.path.join(_lowerCamelCase    ,     _lowerCamelCase ) )
        copyfile(_lowerCamelCase    ,     os.path.join(_lowerCamelCase    ,     '''vocab.json''' ) )

        lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )

     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_			=      WavaVecaFeatureExtractor()
        lowerCAmelCase_			=      AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h''' )

        lowerCAmelCase_			=      WavaVecaProcessor(_lowerCamelCase    ,     _lowerCamelCase )

        # save in new folder
        processor.save_pretrained(_lowerCamelCase )

        # drop `processor_class` in tokenizer
        with open(os.path.join(_lowerCamelCase    ,     _lowerCamelCase )    ,     '''r''' ) as f:
           lowerCAmelCase_			=      json.load(_lowerCamelCase )
           config_dict.pop('''processor_class''' )

        with open(os.path.join(_lowerCamelCase    ,     _lowerCamelCase )    ,     '''w''' ) as f:
           f.write(json.dumps(_lowerCamelCase ) )

        lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )

     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_			=      WavaVecaFeatureExtractor()
        lowerCAmelCase_			=      AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h''' )

        lowerCAmelCase_			=      WavaVecaProcessor(_lowerCamelCase    ,     _lowerCamelCase )

        # save in new folder
        processor.save_pretrained(_lowerCamelCase )

        # drop `processor_class` in feature extractor
        with open(os.path.join(_lowerCamelCase    ,     _lowerCamelCase )    ,     '''r''' ) as f:
           lowerCAmelCase_			=      json.load(_lowerCamelCase )
           config_dict.pop('''processor_class''' )

        with open(os.path.join(_lowerCamelCase    ,     _lowerCamelCase )    ,     '''w''' ) as f:
           f.write(json.dumps(_lowerCamelCase ) )

        lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )

     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     with tempfile.TemporaryDirectory() as tmpdirname:
        lowerCAmelCase_			=      WavaVecaConfig(processor_class='''Wav2Vec2Processor''' )
        model_config.save_pretrained(_lowerCamelCase )
        # copy relevant files
        copyfile(_lowerCamelCase    ,     os.path.join(_lowerCamelCase    ,     '''vocab.json''' ) )
        # create emtpy sample processor
        with open(os.path.join(_lowerCamelCase    ,     _lowerCamelCase )    ,     '''w''' ) as f:
           f.write('''{}''' )

        lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )

     self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     # If remote code is not set, we will time out when asking whether to load the model.
     with self.assertRaises(_lowerCamelCase ):
        lowerCAmelCase_			=      AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor''' )
     # If remote code is disabled, we can't load this config.
     with self.assertRaises(_lowerCamelCase ):
        lowerCAmelCase_			=      AutoProcessor.from_pretrained(
            '''hf-internal-testing/test_dynamic_processor'''    ,     trust_remote_code=_lowerCamelCase )

     lowerCAmelCase_			=      AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''    ,     trust_remote_code=_lowerCamelCase )
     self.assertTrue(processor.special_attribute_present )
     self.assertEqual(processor.__class__.__name__    ,     '''NewProcessor''' )

     lowerCAmelCase_			=      processor.feature_extractor
     self.assertTrue(feature_extractor.special_attribute_present )
     self.assertEqual(feature_extractor.__class__.__name__    ,     '''NewFeatureExtractor''' )

     lowerCAmelCase_			=      processor.tokenizer
     self.assertTrue(tokenizer.special_attribute_present )
     if is_tokenizers_available():
        self.assertEqual(tokenizer.__class__.__name__    ,     '''NewTokenizerFast''' )

        # Test we can also load the slow version
        lowerCAmelCase_			=      AutoProcessor.from_pretrained(
            '''hf-internal-testing/test_dynamic_processor'''    ,     trust_remote_code=_lowerCamelCase    ,     use_fast=_lowerCamelCase )
        lowerCAmelCase_			=      new_processor.tokenizer
        self.assertTrue(new_tokenizer.special_attribute_present )
        self.assertEqual(new_tokenizer.__class__.__name__    ,     '''NewTokenizer''' )
     else:
        self.assertEqual(tokenizer.__class__.__name__    ,     '''NewTokenizer''' )





  def 						UpperCAmelCase_	(	self ):
     try:
        AutoConfig.register('''custom'''    ,     _lowerCamelCase )
        AutoFeatureExtractor.register(_lowerCamelCase    ,     _lowerCamelCase )
        AutoTokenizer.register(_lowerCamelCase    ,     slow_tokenizer_class=_lowerCamelCase )
        AutoProcessor.register(_lowerCamelCase    ,     _lowerCamelCase )
        # Trying to register something existing in the Transformers library will raise an error
        with self.assertRaises(_lowerCamelCase ):
           AutoProcessor.register(_lowerCamelCase    ,     _lowerCamelCase )

        # Now that the config is registered, it can be used as any other config with the auto-API
        lowerCAmelCase_			=      CustomFeatureExtractor.from_pretrained(_lowerCamelCase )

        with tempfile.TemporaryDirectory() as tmp_dir:
           lowerCAmelCase_			=      os.path.join(_lowerCamelCase    ,     '''vocab.txt''' )
           with open(_lowerCamelCase    ,     '''w'''    ,     encoding='''utf-8''' ) as vocab_writer:
              vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens] ) )
           lowerCAmelCase_			=      CustomTokenizer(_lowerCamelCase )

        lowerCAmelCase_			=      CustomProcessor(_lowerCamelCase    ,     _lowerCamelCase )

        with tempfile.TemporaryDirectory() as tmp_dir:
           processor.save_pretrained(_lowerCamelCase )
           lowerCAmelCase_			=      AutoProcessor.from_pretrained(_lowerCamelCase )
           self.assertIsInstance(_lowerCamelCase    ,     _lowerCamelCase )

     finally:
        if "custom" in CONFIG_MAPPING._extra_content:
           del CONFIG_MAPPING._extra_content["custom"]
        if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
           del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
        if CustomConfig in TOKENIZER_MAPPING._extra_content:
           del TOKENIZER_MAPPING._extra_content[CustomConfig]
        if CustomConfig in PROCESSOR_MAPPING._extra_content:
           del PROCESSOR_MAPPING._extra_content[CustomConfig]





  def 						UpperCAmelCase_	(	self ):


     class     __UpperCAmelCase       ( __a		):
       __A :				Tuple											= False






     class     __UpperCAmelCase       ( __a		):
       __A :				List[Any]											= False






     class     __UpperCAmelCase       ( __a		):
       __A :				int											= 'AutoFeatureExtractor'
       __A :				Optional[int]											= 'AutoTokenizer'
       __A :				List[str]											= False

     try:
        AutoConfig.register('''custom'''    ,     _lowerCamelCase )
        AutoFeatureExtractor.register(_lowerCamelCase    ,     _lowerCamelCase )
        AutoTokenizer.register(_lowerCamelCase    ,     slow_tokenizer_class=_lowerCamelCase )
        AutoProcessor.register(_lowerCamelCase    ,     _lowerCamelCase )
        # If remote code is not set, the default is to use local classes.
        lowerCAmelCase_			=      AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor''' )
        self.assertEqual(processor.__class__.__name__    ,     '''NewProcessor''' )
        self.assertFalse(processor.special_attribute_present )
        self.assertFalse(processor.feature_extractor.special_attribute_present )
        self.assertFalse(processor.tokenizer.special_attribute_present )

        # If remote code is disabled, we load the local ones.
        lowerCAmelCase_			=      AutoProcessor.from_pretrained(
            '''hf-internal-testing/test_dynamic_processor'''    ,     trust_remote_code=_lowerCamelCase )
        self.assertEqual(processor.__class__.__name__    ,     '''NewProcessor''' )
        self.assertFalse(processor.special_attribute_present )
        self.assertFalse(processor.feature_extractor.special_attribute_present )
        self.assertFalse(processor.tokenizer.special_attribute_present )

        # If remote is enabled, we load from the Hub.
        lowerCAmelCase_			=      AutoProcessor.from_pretrained(
            '''hf-internal-testing/test_dynamic_processor'''    ,     trust_remote_code=_lowerCamelCase )
        self.assertEqual(processor.__class__.__name__    ,     '''NewProcessor''' )
        self.assertTrue(processor.special_attribute_present )
        self.assertTrue(processor.feature_extractor.special_attribute_present )
        self.assertTrue(processor.tokenizer.special_attribute_present )

     finally:
        if "custom" in CONFIG_MAPPING._extra_content:
           del CONFIG_MAPPING._extra_content["custom"]
        if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
           del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
        if CustomConfig in TOKENIZER_MAPPING._extra_content:
           del TOKENIZER_MAPPING._extra_content[CustomConfig]
        if CustomConfig in PROCESSOR_MAPPING._extra_content:
           del PROCESSOR_MAPPING._extra_content[CustomConfig]





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-bert''' )
     self.assertEqual(processor.__class__.__name__    ,     '''BertTokenizerFast''' )







  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-convnext''' )
     self.assertEqual(processor.__class__.__name__    ,     '''ConvNextImageProcessor''' )






@is_staging_test
class     __UpperCAmelCase       ( unittest.TestCase		):
  __A :				Optional[int]											= ['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]', 'bla', 'blou']





  @classmethod
  def 						UpperCAmelCase_	(	cls ):
     lowerCAmelCase_			=      TOKEN
     HfFolder.save_token(_lowerCamelCase )





  @classmethod
  def 						UpperCAmelCase_	(	cls ):
     try:
        delete_repo(token=cls._token    ,     repo_id='''test-processor''' )
     except HTTPError:
        pass

     try:
        delete_repo(token=cls._token    ,     repo_id='''valid_org/test-processor-org''' )
     except HTTPError:
        pass

     try:
        delete_repo(token=cls._token    ,     repo_id='''test-dynamic-processor''' )
     except HTTPError:
        pass





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      WavaVecaProcessor.from_pretrained(_lowerCamelCase )
     with tempfile.TemporaryDirectory() as tmp_dir:
        processor.save_pretrained(
            os.path.join(_lowerCamelCase    ,     '''test-processor''' )    ,     push_to_hub=_lowerCamelCase    ,     use_auth_token=self._token )

        lowerCAmelCase_			=      WavaVecaProcessor.from_pretrained(F'''{USER}/test-processor''' )
        for k, v in processor.feature_extractor.__dict__.items():
           self.assertEqual(_lowerCamelCase    ,     getattr(new_processor.feature_extractor    ,     _lowerCamelCase ) )
        self.assertDictEqual(new_processor.tokenizer.get_vocab()    ,     processor.tokenizer.get_vocab() )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      WavaVecaProcessor.from_pretrained(_lowerCamelCase )

     with tempfile.TemporaryDirectory() as tmp_dir:
        processor.save_pretrained(
            os.path.join(_lowerCamelCase    ,     '''test-processor-org''' )    ,     push_to_hub=_lowerCamelCase    ,     use_auth_token=self._token    ,     organization='''valid_org'''    ,     )

        lowerCAmelCase_			=      WavaVecaProcessor.from_pretrained('''valid_org/test-processor-org''' )
        for k, v in processor.feature_extractor.__dict__.items():
           self.assertEqual(_lowerCamelCase    ,     getattr(new_processor.feature_extractor    ,     _lowerCamelCase ) )
        self.assertDictEqual(new_processor.tokenizer.get_vocab()    ,     processor.tokenizer.get_vocab() )







  def 						UpperCAmelCase_	(	self ):
     CustomFeatureExtractor.register_for_auto_class()
     CustomTokenizer.register_for_auto_class()
     CustomProcessor.register_for_auto_class()

     lowerCAmelCase_			=      CustomFeatureExtractor.from_pretrained(_lowerCamelCase )

     with tempfile.TemporaryDirectory() as tmp_dir:
        lowerCAmelCase_			=      os.path.join(_lowerCamelCase    ,     '''vocab.txt''' )
        with open(_lowerCamelCase    ,     '''w'''    ,     encoding='''utf-8''' ) as vocab_writer:
           vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens] ) )
        lowerCAmelCase_			=      CustomTokenizer(_lowerCamelCase )

     lowerCAmelCase_			=      CustomProcessor(_lowerCamelCase    ,     _lowerCamelCase )

     with tempfile.TemporaryDirectory() as tmp_dir:
        create_repo(F'''{USER}/test-dynamic-processor'''    ,     token=self._token )
        lowerCAmelCase_			=      Repository(_lowerCamelCase    ,     clone_from=F'''{USER}/test-dynamic-processor'''    ,     token=self._token )
        processor.save_pretrained(_lowerCamelCase )

        # This has added the proper auto_map field to the feature extractor config
        self.assertDictEqual(
            processor.feature_extractor.auto_map    ,     {
                '''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor''',
                '''AutoProcessor''': '''custom_processing.CustomProcessor''',
            }    ,     )

        # This has added the proper auto_map field to the tokenizer config
        with open(os.path.join(_lowerCamelCase    ,     '''tokenizer_config.json''' ) ) as f:
           lowerCAmelCase_			=      json.load(_lowerCamelCase )
        self.assertDictEqual(
            tokenizer_config['''auto_map''']    ,     {
                '''AutoTokenizer''': ['''custom_tokenization.CustomTokenizer''', None],
                '''AutoProcessor''': '''custom_processing.CustomProcessor''',
            }    ,     )

        # The code has been copied from fixtures
        self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase    ,     '''custom_feature_extraction.py''' ) ) )
        self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase    ,     '''custom_tokenization.py''' ) ) )
        self.assertTrue(os.path.isfile(os.path.join(_lowerCamelCase    ,     '''custom_processing.py''' ) ) )

        repo.push_to_hub()

     lowerCAmelCase_			=      AutoProcessor.from_pretrained(F'''{USER}/test-dynamic-processor'''    ,     trust_remote_code=_lowerCamelCase )
     # Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
     self.assertEqual(new_processor.__class__.__name__    ,     '''CustomProcessor''' )






'''simple docstring'''

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 snake_case_      (					)			->					str:
   lowerCAmelCase_			=      torch.nn.Linear(2	,   4)
   lowerCAmelCase_			=      torch.optim.AdamW(model.parameters()	,   lr=1.0)
   lowerCAmelCase_			=      torch.optim.lr_scheduler.OneCycleLR(__snake_case	,   max_lr=0.0_1	,   steps_per_epoch=2	,   epochs=1)
   lowerCAmelCase_			=      DataLoader(TensorDataset(torch.tensor([1, 2, 3])))
   lowerCAmelCase_			=      DataLoader(TensorDataset(torch.tensor([4, 5, 6])))

   return model, optimizer, scheduler, train_dl, valid_dl



def snake_case_      (					__snake_case  :							Any)			->					List[Any]:
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def snake_case_      (					__snake_case  :							Optional[int])			->					Any:
   lowerCAmelCase_			=      torch.nn.Linear(*tuple(model.weight.T.shape)).state_dict()
   model.load_state_dict(__snake_case)






class     __UpperCAmelCase       ( __a		):





  @require_cuda
  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     assert PartialState._shared_state["_cpu"] is False
     assert PartialState._shared_state["device"].type == "cuda"
     with self.assertRaises(_lowerCamelCase ):
        lowerCAmelCase_			=      Accelerator(cpu=_lowerCamelCase )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_			=      GradientState()
     assert state.num_steps == 1
     lowerCAmelCase_			=      4
     assert state.num_steps == 4

     assert state.sync_gradients is True
     lowerCAmelCase_			=      False
     assert state.sync_gradients is False
     GradientState._reset_state()





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()

     (
         (
         lowerCAmelCase_
     )	,(
         lowerCAmelCase_
     )	,(
         lowerCAmelCase_
     )	,(
         lowerCAmelCase_
     )	,(
         lowerCAmelCase_
     )	,
     )			=      accelerator.prepare(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )

     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 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()
     accelerator.prepare(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )
     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 						UpperCAmelCase_	(	self ):
     PartialState._reset_state()

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

     with patch('''torch.cuda.set_device'''    ,     _lowerCamelCase ), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64''' ):
        lowerCAmelCase_			=      Accelerator()
        self.assertEqual(str(accelerator.state.device )    ,     '''cuda:64''' )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()
     accelerator.prepare(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )

     lowerCAmelCase_			=      get_signature(_lowerCamelCase )

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

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

        # make sure loaded weights match
        accelerator.load_state(_lowerCamelCase )
        self.assertTrue(abs(model_signature - get_signature(_lowerCamelCase ) ) < 1E-3 )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()
     accelerator.prepare(_lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )

     lowerCAmelCase_			=      get_signature(_lowerCamelCase )

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

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

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

        lowerCAmelCase_			=      config['''class_name''']

     lowerCAmelCase_			=      accelerator.register_save_state_pre_hook(_lowerCamelCase )
     lowerCAmelCase_			=      accelerator.register_load_state_pre_hook(_lowerCamelCase )

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

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

        # random class name to verify correct one is loaded
        lowerCAmelCase_			=      '''random'''

        # make sure loaded weights match with hooks
        accelerator.load_state(_lowerCamelCase )
        self.assertTrue(abs(model_signature - get_signature(_lowerCamelCase ) ) < 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(_lowerCamelCase )

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

        # random class name to verify correct one is loaded
        lowerCAmelCase_			=      '''random'''

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

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





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()
     lowerCAmelCase_			=      None

     # This should work
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      accelerator.prepare(
         _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )
     self.assertTrue(dummy_obj is None )





  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      Accelerator()
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      create_components()
     lowerCAmelCase_			=      [1, 2, 3]

     # This should work
     lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_	,lowerCAmelCase_			=      accelerator.prepare(
         _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase    ,     _lowerCamelCase )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Dummy object should have `_is_accelerate_prepared` set to `True`'''    ,     )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Model is missing `_is_accelerator_prepared` or is set to `False`'''    ,     )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Optimizer is missing `_is_accelerator_prepared` or is set to `False`'''    ,     )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Scheduler is missing `_is_accelerator_prepared` or is set to `False`'''    ,     )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''    ,     )
     self.assertEqual(
         getattr(_lowerCamelCase    ,     '''_is_accelerate_prepared'''    ,     _lowerCamelCase )    ,     _lowerCamelCase    ,     '''Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''    ,     )





  @slow
  @require_bnb
  def 						UpperCAmelCase_	(	self ):
     from transformers import AutoModelForCausalLM

     lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
         '''EleutherAI/gpt-neo-125m'''    ,     load_in_abit=_lowerCamelCase    ,     device_map={'''''': 0}    ,     )
     lowerCAmelCase_			=      Accelerator()

     # This should work
     lowerCAmelCase_			=      accelerator.prepare(_lowerCamelCase )





  @slow
  @require_bnb
  def 						UpperCAmelCase_	(	self ):
     from transformers import AutoModelForCausalLM

     lowerCAmelCase_			=      Accelerator()

     with init_empty_weights():
        lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''    ,     )
        model.tie_weights()
        lowerCAmelCase_			=      infer_auto_device_map(_lowerCamelCase )
        lowerCAmelCase_			=      '''cpu'''

     lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
         '''EleutherAI/gpt-neo-125m'''    ,     device_map=_lowerCamelCase    ,     load_in_abit=_lowerCamelCase    ,     llm_inta_enable_fpaa_cpu_offload=_lowerCamelCase )

     # This should not work and get value error
     with self.assertRaises(_lowerCamelCase ):
        lowerCAmelCase_			=      accelerator.prepare(_lowerCamelCase )





  @slow
  @require_bnb
  @require_multi_gpu
  def 						UpperCAmelCase_	(	self ):
     from transformers import AutoModelForCausalLM

     lowerCAmelCase_			=      {'''distributed_type''': DistributedType.MULTI_GPU}

     with init_empty_weights():
        lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''    ,     )
        model.tie_weights()
        lowerCAmelCase_			=      infer_auto_device_map(_lowerCamelCase )
        lowerCAmelCase_			=      1

     lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
         '''EleutherAI/gpt-neo-125m'''    ,     load_in_abit=_lowerCamelCase    ,     device_map=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      Accelerator()

     # This should not work and get value error
     with self.assertRaises(_lowerCamelCase ):
        lowerCAmelCase_			=      accelerator.prepare(_lowerCamelCase )

     PartialState._reset_state()





  @slow
  @require_bnb
  @require_multi_gpu
  def 						UpperCAmelCase_	(	self ):
     from transformers import AutoModelForCausalLM

     with init_empty_weights():
        lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''    ,     )
        lowerCAmelCase_			=      infer_auto_device_map(_lowerCamelCase )
        lowerCAmelCase_			=      1

     lowerCAmelCase_			=      AutoModelForCausalLM.from_pretrained(
         '''EleutherAI/gpt-neo-125m'''    ,     load_in_abit=_lowerCamelCase    ,     device_map=_lowerCamelCase    ,     )
     lowerCAmelCase_			=      Accelerator()

     # This should work
     lowerCAmelCase_			=      accelerator.prepare(_lowerCamelCase )







  @require_cuda
  def 						UpperCAmelCase_	(	self ):
     lowerCAmelCase_			=      torch.nn.Linear(10    ,     10 )
     lowerCAmelCase_			=      torch.optim.SGD(model.parameters()    ,     lr=0.01 )
     lowerCAmelCase_			=      Accelerator(cpu=_lowerCamelCase )
     lowerCAmelCase_			=      accelerator.prepare(_lowerCamelCase )





