


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

   return model, optimizer, scheduler, train_dl, valid_dl
def   _UpperCamelCase				(						lowerCAmelCase_   )						->Optional[Any]:
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()
def   _UpperCamelCase				(						lowerCAmelCase_   )						->Dict:
   UpperCAmelCase							=     torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
   model.load_state_dict(lowerCAmelCase_   )







class   __lowercase				(  __snake_case  ):





      @require_cuda
      def        _lowercase				(		self						:					str       )    ->   Optional[Any]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(__lowerCamelCase       ):
            UpperCAmelCase							=     Accelerator(cpu=__lowerCamelCase       )





      def        _lowercase				(		self						:					Optional[Any]       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase							=     GradientState()
         assert state.num_steps == 1
         UpperCAmelCase							=     4
         assert state.num_steps == 4

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





      def        _lowercase				(		self						:					Optional[int]       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     create_components()

         (
             (
             UpperCAmelCase
         )				,  (
             UpperCAmelCase
         )				,  (
             UpperCAmelCase
         )				,  (
             UpperCAmelCase
         )				,  (
             UpperCAmelCase
         )				,  
         )							=     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        _lowercase				(		self						:					Tuple       )    ->   Optional[int]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     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        _lowercase				(		self						:					Tuple       )    ->   Dict:






         """simple docstring"""







         PartialState._reset_state()

         # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
         def noop(*__lowerCamelCase						:					Union[str, Any]					,       **__lowerCamelCase						:					List[Any]       ):
            pass

         with patch("""torch.cuda.set_device"""					,       __lowerCamelCase       ), patch_environment(ACCELERATE_TORCH_DEVICE="""cuda:64"""       ):
            UpperCAmelCase							=     Accelerator()
            self.assertEqual(str(accelerator.state.device       )					,       """cuda:64"""       )





      def        _lowercase				(		self						:					Union[str, Any]       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     create_components()
         accelerator.prepare(__lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase       )

         UpperCAmelCase							=     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        _lowercase				(		self						:					Optional[int]       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     create_components()
         accelerator.prepare(__lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase       )

         UpperCAmelCase							=     get_signature(__lowerCamelCase       )

         # saving hook
         def save_config(__lowerCamelCase						:					str					,       __lowerCamelCase						:					Union[str, Any]					,       __lowerCamelCase						:					List[str]       ):
            UpperCAmelCase							=     {"""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						:					Dict					,       __lowerCamelCase						:					List[str]       ):
            with open(os.path.join(__lowerCamelCase					,       """data.json"""       )					,       """r"""       ) as f:
               UpperCAmelCase							=     json.load(__lowerCamelCase       )

            UpperCAmelCase							=     config["""class_name"""]

         UpperCAmelCase							=     accelerator.register_save_state_pre_hook(__lowerCamelCase       )
         UpperCAmelCase							=     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
            UpperCAmelCase							=     """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
            UpperCAmelCase							=     """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        _lowercase				(		self						:					Optional[int]       )    ->   Union[str, Any]:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     create_components()
         UpperCAmelCase							=     None

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





      def        _lowercase				(		self						:					List[Any]       )    ->   Dict:






         """simple docstring"""







         UpperCAmelCase							=     Accelerator()
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     create_components()
         UpperCAmelCase							=     [1, 2, 3]

         # This should work
         UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase				,  UpperCAmelCase							=     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        _lowercase				(		self						:					Union[str, Any]       )    ->   List[str]:






         """simple docstring"""







         from transformers import AutoModelForCausalLM

         UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""					,       load_in_abit=__lowerCamelCase					,       device_map={"""""": 0}					,       )
         UpperCAmelCase							=     Accelerator()

         # This should work
         UpperCAmelCase							=     accelerator.prepare(__lowerCamelCase       )





      @slow
      @require_bnb
      def        _lowercase				(		self						:					Dict       )    ->   List[str]:






         """simple docstring"""







         from transformers import AutoModelForCausalLM

         UpperCAmelCase							=     Accelerator()

         with init_empty_weights():
            UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""					,       )
            model.tie_weights()
            UpperCAmelCase							=     infer_auto_device_map(__lowerCamelCase       )
            UpperCAmelCase							=     """cpu"""

         UpperCAmelCase							=     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       ):
            UpperCAmelCase							=     accelerator.prepare(__lowerCamelCase       )





      @slow
      @require_bnb
      @require_multi_gpu
      def        _lowercase				(		self						:					Dict       )    ->   Tuple:






         """simple docstring"""







         from transformers import AutoModelForCausalLM

         UpperCAmelCase							=     {"""distributed_type""": DistributedType.MULTI_GPU}

         with init_empty_weights():
            UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""					,       )
            model.tie_weights()
            UpperCAmelCase							=     infer_auto_device_map(__lowerCamelCase       )
            UpperCAmelCase							=     1

         UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""					,       load_in_abit=__lowerCamelCase					,       device_map=__lowerCamelCase					,       )
         UpperCAmelCase							=     Accelerator()

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

         PartialState._reset_state()





      @slow
      @require_bnb
      @require_multi_gpu
      def        _lowercase				(		self						:					str       )    ->   str:






         """simple docstring"""







         from transformers import AutoModelForCausalLM

         with init_empty_weights():
            UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
                """EleutherAI/gpt-neo-125m"""					,       )
            UpperCAmelCase							=     infer_auto_device_map(__lowerCamelCase       )
            UpperCAmelCase							=     1

         UpperCAmelCase							=     AutoModelForCausalLM.from_pretrained(
             """EleutherAI/gpt-neo-125m"""					,       load_in_abit=__lowerCamelCase					,       device_map=__lowerCamelCase					,       )
         UpperCAmelCase							=     Accelerator()

         # This should work
         UpperCAmelCase							=     accelerator.prepare(__lowerCamelCase       )





      @require_cuda
      def        _lowercase				(		self						:					int       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     torch.nn.Linear(1_0					,       1_0       )
         UpperCAmelCase							=     torch.optim.SGD(model.parameters()					,       lr=0.01       )
         UpperCAmelCase							=     Accelerator(cpu=__lowerCamelCase       )
         UpperCAmelCase							=     accelerator.prepare(__lowerCamelCase       )







import multiprocessing
import time

from arguments import PretokenizationArguments
from datasets import load_dataset

from transformers import AutoTokenizer, HfArgumentParser
def   _UpperCamelCase				(						lowerCAmelCase_   )						->int:
   UpperCAmelCase							=     {}
   UpperCAmelCase							=     tokenizer(example["""content"""]  ,						truncation=lowerCAmelCase_   )["""input_ids"""]
   UpperCAmelCase							=     len(example["""content"""]   ) / len(output["""input_ids"""]   )
   return output


__a     =			HfArgumentParser(PretokenizationArguments)
__a     =			parser.parse_args()
if args.num_workers is None:
 __a     =			multiprocessing.cpu_count()
__a     =			AutoTokenizer.from_pretrained(args.tokenizer_dir)

__a     =			time.time()
__a     =			load_dataset(args.dataset_name, split="""train""")
print(F"""Dataset loaded in {time.time()-t_start:.2f}s""")

__a     =			time.time()
__a     =			ds.map(
    tokenize,
    num_proc=args.num_workers,
    remove_columns=[
        """repo_name""",
        """path""",
        """copies""",
        """size""",
        """content""",
        """license""",
        """hash""",
        """line_mean""",
        """line_max""",
        """alpha_frac""",
        """autogenerated""",
    ],
)
print(F"""Dataset tokenized in {time.time()-t_start:.2f}s""")

__a     =			time.time()
ds.push_to_hub(args.tokenized_data_repo)







print(F"""Data pushed to the hub in {time.time()-t_start:.2f}s""")



