




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 lowerCamelCase__       (      ):

     '''simple docstring'''


     UpperCAmelCase_      :							str       =      torch.nn.Linear(2			,						4				)
     UpperCAmelCase_      :							Tuple       =      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_      :							Optional[int]       =      DataLoader(TensorDataset(torch.tensor([1, 2, 3]				)				)				)
     UpperCAmelCase_      :							Tuple       =      DataLoader(TensorDataset(torch.tensor([4, 5, 6]				)				)				)

     return model, optimizer, scheduler, train_dl, valid_dl







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     return (model.weight.abs().sum() + model.bias.abs().sum()).item()







def lowerCamelCase__       (      _lowercase				):

     '''simple docstring'''


     UpperCAmelCase_      :							str       =      torch.nn.Linear(*tuple(model.weight.T.shape				)				).state_dict()
     model.load_state_dict(snake_case_				)



class   __a(					UpperCamelCase__       ):






    """simple docstring"""





    @require_cuda
    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							Optional[Any]       =      Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(__a  ):
              UpperCAmelCase_      :							Union[str, Any]       =      Accelerator(cpu=__a  )





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Optional[Any]       =      Accelerator()
         UpperCAmelCase_      :							int       =      GradientState()
         assert state.num_steps == 1
         UpperCAmelCase_      :							Dict       =      4
         assert state.num_steps == 4

         assert state.sync_gradients is True
         UpperCAmelCase_      :							List[Any]       =      False
         assert state.sync_gradients is False
         GradientState._reset_state()





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							List[Any]       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Any       =      create_components()

         (
             (
             UpperCAmelCase_			
         ),   (
             UpperCAmelCase_			
         ),   (
             UpperCAmelCase_			
         ),   (
             UpperCAmelCase_			
         ),   (
             UpperCAmelCase_			
         ),   			
         )      :							List[str]       =      accelerator.prepare(__a       ,__a       ,__a       ,__a       ,__a  )

         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 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							Any       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Dict       =      create_components()
         accelerator.prepare(__a       ,__a       ,__a       ,__a       ,__a  )
         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 							a__				(				self  )			->		Tuple:
         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'''       ,__a  ), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64'''  ):
              UpperCAmelCase_      :							List[str]       =      Accelerator()
              self.assertEqual(str(accelerator.state.device  )       ,'''cuda:64'''  )





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							Optional[Any]       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Union[str, Any]       =      create_components()
         accelerator.prepare(__a       ,__a       ,__a       ,__a       ,__a  )

         UpperCAmelCase_      :							Optional[Any]       =      get_signature(__a  )

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

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

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





    def 							a__				(				self  )			->		List[Any]:
         UpperCAmelCase_      :							Tuple       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							str       =      create_components()
         accelerator.prepare(__a       ,__a       ,__a       ,__a       ,__a  )

         UpperCAmelCase_      :							List[Any]       =      get_signature(__a  )

         # saving hook
         def save_config(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ):
              UpperCAmelCase_      :							Union[str, Any]       =      {'''class_name''': models[0].__class__.__name__}

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

        # loading hook
         def load_config(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ):
              with open(os.path.join(__a       ,'''data.json'''  )       ,'''r'''  ) as f:
                   UpperCAmelCase_      :							List[Any]       =      json.load(__a  )

              UpperCAmelCase_      :							List[Any]       =      config['''class_name''']

         UpperCAmelCase_      :							Dict       =      accelerator.register_save_state_pre_hook(__a  )
         UpperCAmelCase_      :							List[str]       =      accelerator.register_load_state_pre_hook(__a  )

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

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

              # random class name to verify correct one is loaded
              UpperCAmelCase_      :							int       =      '''random'''

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

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

              # random class name to verify correct one is loaded
              UpperCAmelCase_      :							str       =      '''random'''

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

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





    def 							a__				(				self  )			->		List[str]:
         UpperCAmelCase_      :							List[Any]       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Any       =      create_components()
         UpperCAmelCase_      :							List[str]       =      None

         # This should work
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							List[str]       =      accelerator.prepare(
             __a       ,__a       ,__a       ,__a       ,__a       ,__a  )
         self.assertTrue(dummy_obj is None  )





    def 							a__				(				self  )			->		Any:
         UpperCAmelCase_      :							Optional[int]       =      Accelerator()
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							Any       =      create_components()
         UpperCAmelCase_      :							Optional[int]       =      [1, 2, 3]

         # This should work
         UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_,   UpperCAmelCase_      :							List[str]       =      accelerator.prepare(
             __a       ,__a       ,__a       ,__a       ,__a       ,__a  )
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Dummy object should have `_is_accelerate_prepared` set to `True`'''       ,)
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Model is missing `_is_accelerator_prepared` or is set to `False`'''       ,)
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Optimizer is missing `_is_accelerator_prepared` or is set to `False`'''       ,)
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Scheduler is missing `_is_accelerator_prepared` or is set to `False`'''       ,)
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''       ,)
         self.assertEqual(
             getattr(__a       ,'''_is_accelerate_prepared'''       ,__a  )       ,__a       ,'''Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'''       ,)





    @slow
    @require_bnb
    def 							a__				(				self  )			->		str:
         from transformers import AutoModelForCausalLM

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

         # This should work
         UpperCAmelCase_      :							List[Any]       =      accelerator.prepare(__a  )





    @slow
    @require_bnb
    def 							a__				(				self  )			->		Optional[int]:
         from transformers import AutoModelForCausalLM

         UpperCAmelCase_      :							List[Any]       =      Accelerator()

         with init_empty_weights():
              UpperCAmelCase_      :							Optional[int]       =      AutoModelForCausalLM.from_pretrained(
                  '''EleutherAI/gpt-neo-125m'''       ,)
              model.tie_weights()
              UpperCAmelCase_      :							Any       =      infer_auto_device_map(__a  )
              UpperCAmelCase_      :							Any       =      '''cpu'''

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

         # This should not work and get value error
         with self.assertRaises(__a  ):
              UpperCAmelCase_      :							Dict       =      accelerator.prepare(__a  )





    @slow
    @require_bnb
    @require_multi_gpu
    def 							a__				(				self  )			->		str:
         from transformers import AutoModelForCausalLM

         UpperCAmelCase_      :							int       =      {'''distributed_type''': DistributedType.MULTI_GPU}

         with init_empty_weights():
              UpperCAmelCase_      :							Any       =      AutoModelForCausalLM.from_pretrained(
                  '''EleutherAI/gpt-neo-125m'''       ,)
              model.tie_weights()
              UpperCAmelCase_      :							Tuple       =      infer_auto_device_map(__a  )
              UpperCAmelCase_      :							Tuple       =      1

         UpperCAmelCase_      :							List[str]       =      AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''       ,load_in_abit=__a       ,device_map=__a       ,)
         UpperCAmelCase_      :							List[Any]       =      Accelerator()

         # This should not work and get value error
         with self.assertRaises(__a  ):
              UpperCAmelCase_      :							Dict       =      accelerator.prepare(__a  )

         PartialState._reset_state()





    @slow
    @require_bnb
    @require_multi_gpu
    def 							a__				(				self  )			->		Union[str, Any]:
         from transformers import AutoModelForCausalLM

         with init_empty_weights():
              UpperCAmelCase_      :							Optional[Any]       =      AutoModelForCausalLM.from_pretrained(
                  '''EleutherAI/gpt-neo-125m'''       ,)
              UpperCAmelCase_      :							Optional[int]       =      infer_auto_device_map(__a  )
              UpperCAmelCase_      :							List[str]       =      1

         UpperCAmelCase_      :							List[str]       =      AutoModelForCausalLM.from_pretrained(
             '''EleutherAI/gpt-neo-125m'''       ,load_in_abit=__a       ,device_map=__a       ,)
         UpperCAmelCase_      :							Dict       =      Accelerator()

         # This should work
         UpperCAmelCase_      :							str       =      accelerator.prepare(__a  )





    @require_cuda
    def 							a__				(				self  )			->		Optional[Any]:
         UpperCAmelCase_      :							List[Any]       =      torch.nn.Linear(10       ,10  )
         UpperCAmelCase_      :							Optional[int]       =      torch.optim.SGD(model.parameters()       ,lr=0.01  )
         UpperCAmelCase_      :							Tuple       =      Accelerator(cpu=__a  )
         UpperCAmelCase_      :							List[str]       =      accelerator.prepare(__a  )











import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s', datefmt='%m/%d/%Y %H:%M:%S', level=logging.INFO
)
__a          =			logging.getLogger(__name__)

if __name__ == "__main__":
     __a          =			argparse.ArgumentParser(
         description='Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)'
     )
     parser.add_argument(
         '--data_file', type=str, default='data/dump.bert-base-uncased.pickle', help='The binarized dataset.'
     )
     parser.add_argument(
         '--token_counts_dump', type=str, default='data/token_counts.bert-base-uncased.pickle', help='The dump file.'
     )
     parser.add_argument('--vocab_size', default=30_522, type=int)
     __a          =			parser.parse_args()

     logger.info(F"""Loading data from {args.data_file}""")
     with open(args.data_file, 'rb') as fp:
          __a          =			pickle.load(fp)

     logger.info('Counting occurrences for MLM.')
     __a          =			Counter()
     for tk_ids in data:
          counter.update(tk_ids)
     __a          =			[0] * args.vocab_size
     for k, v in counter.items():
          __a          =			v

     logger.info(F"""Dump to {args.token_counts_dump}""")
     with open(args.token_counts_dump, 'wb') as handle:
          pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)