

"""simple docstring"""




import os
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils import AddedToken
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
     from .tokenization_xlnet import XLNetTokenizer
else:
     SCREAMING_SNAKE_CASE_     :			Union[str, Any]	  =			None


SCREAMING_SNAKE_CASE_     :			int	  =			logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_     :			Tuple	  =			{'''vocab_file''': '''spiece.model''', '''tokenizer_file''': '''tokenizer.json'''}

SCREAMING_SNAKE_CASE_     :			Union[str, Any]	  =			{
    '''vocab_file''': {
        '''xlnet-base-cased''': '''https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model''',
        '''xlnet-large-cased''': '''https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model''',
    },
    '''tokenizer_file''': {
        '''xlnet-base-cased''': '''https://huggingface.co/xlnet-base-cased/resolve/main/tokenizer.json''',
        '''xlnet-large-cased''': '''https://huggingface.co/xlnet-large-cased/resolve/main/tokenizer.json''',
    },
}

SCREAMING_SNAKE_CASE_     :			int	  =			{
    '''xlnet-base-cased''': None,
    '''xlnet-large-cased''': None,
}

SCREAMING_SNAKE_CASE_     :			int	  =			'''▁'''

# Segments (not really needed)
SCREAMING_SNAKE_CASE_     :			Any	  =			0
SCREAMING_SNAKE_CASE_     :			str	  =			1
SCREAMING_SNAKE_CASE_     :			Dict	  =			2
SCREAMING_SNAKE_CASE_     :			Union[str, Any]	  =			3
SCREAMING_SNAKE_CASE_     :			Any	  =			4


class 		_A       (			lowercase__					):
      __a								=	VOCAB_FILES_NAMES
      __a								=	PRETRAINED_VOCAB_FILES_MAP
      __a								=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __a								=	'left'
      __a								=	XLNetTokenizer





      def __init__(  self    ,						SCREAMING_SNAKE_CASE__=None    ,						SCREAMING_SNAKE_CASE__=None    ,						SCREAMING_SNAKE_CASE__=False    ,						SCREAMING_SNAKE_CASE__=True    ,						SCREAMING_SNAKE_CASE__=False    ,						SCREAMING_SNAKE_CASE__="<s>"    ,						SCREAMING_SNAKE_CASE__="</s>"    ,						SCREAMING_SNAKE_CASE__="<unk>"    ,						SCREAMING_SNAKE_CASE__="<sep>"    ,						SCREAMING_SNAKE_CASE__="<pad>"    ,						SCREAMING_SNAKE_CASE__="<cls>"    ,						SCREAMING_SNAKE_CASE__="<mask>"    ,						SCREAMING_SNAKE_CASE__=["<eop>", "<eod>"]    ,						**SCREAMING_SNAKE_CASE__    ,						)   ->					Optional[int]:
         lowerCamelCase__						=			AddedToken(UpperCAmelCase__    ,						lstrip=UpperCAmelCase__    ,						rstrip=UpperCAmelCase__		) if isinstance(UpperCAmelCase__    ,						UpperCAmelCase__		) else mask_token

         super().__init__(
             vocab_file=UpperCAmelCase__    ,						tokenizer_file=UpperCAmelCase__    ,						do_lower_case=UpperCAmelCase__    ,						remove_space=UpperCAmelCase__    ,						keep_accents=UpperCAmelCase__    ,						bos_token=UpperCAmelCase__    ,						eos_token=UpperCAmelCase__    ,						unk_token=UpperCAmelCase__    ,						sep_token=UpperCAmelCase__    ,						pad_token=UpperCAmelCase__    ,						cls_token=UpperCAmelCase__    ,						mask_token=UpperCAmelCase__    ,						additional_special_tokens=UpperCAmelCase__    ,						**UpperCAmelCase__    ,						)

         lowerCamelCase__						=			3
         lowerCamelCase__						=			do_lower_case
         lowerCamelCase__						=			remove_space
         lowerCamelCase__						=			keep_accents
         lowerCamelCase__						=			vocab_file
         lowerCamelCase__						=			False if not self.vocab_file else True





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__ = None		)   ->					int:
         lowerCamelCase__						=			[self.sep_token_id]
         lowerCamelCase__						=			[self.cls_token_id]
         if token_ids_a is None:
            return token_ids_a + sep + cls
         return token_ids_a + sep + token_ids_a + sep + cls





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__ = None		)   ->					List[str]:
         lowerCamelCase__						=			[self.sep_token_id]
         lowerCamelCase__						=			[2]

         if token_ids_a is None:
            return len(token_ids_a + sep		) * [0] + cls_segment_id
         return len(token_ids_a + sep		) * [0] + len(token_ids_a + sep		) * [1] + cls_segment_id





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__ = None		)   ->					Union[str, Any]:
         if not self.can_save_slow_tokenizer:
            raise ValueError(
                "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
                "tokenizer."		)

         if not os.path.isdir(UpperCAmelCase__		):
            logger.error(f'Vocabulary path ({save_directory}) should be a directory'		)
            return
         lowerCamelCase__						=			os.path.join(
             UpperCAmelCase__    ,						(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]		)

         if os.path.abspath(self.vocab_file		) != os.path.abspath(UpperCAmelCase__		):
            copyfile(self.vocab_file    ,						UpperCAmelCase__		)

         return (out_vocab_file,)







"""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 	UpperCAmelCase__   (				)      ->      List[Any]:




   """simple docstring"""
   lowerCamelCase__						=			torch.nn.Linear(2	,    4   )
   lowerCamelCase__						=			torch.optim.AdamW(model.parameters()	,    lr=1.0   )
   lowerCamelCase__						=			torch.optim.lr_scheduler.OneCycleLR(A__	,    max_lr=0.01	,    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 	UpperCAmelCase__   (				A__   )      ->      Any:




   """simple docstring"""
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()



def 	UpperCAmelCase__   (				A__   )      ->      Any:




   """simple docstring"""
   lowerCamelCase__						=			torch.nn.Linear(*tuple(model.weight.T.shape   )   ).state_dict()
   model.load_state_dict(A__   )


class 		_A       (			__a					):





      @require_cuda
      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			Accelerator()
         assert PartialState._shared_state["_cpu"] is False
         assert PartialState._shared_state["device"].type == "cuda"
         with self.assertRaises(SCREAMING_SNAKE_CASE__		):
            lowerCamelCase__						=			Accelerator(cpu=SCREAMING_SNAKE_CASE__		)





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         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 					_lowerCamelCase  (  self		)   ->					str:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			create_components()

         (
             (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       (
             lowerCamelCase__
         )					,       
         )						=			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 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					Dict:
         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"		):
            lowerCamelCase__						=			Accelerator()
            self.assertEqual(str(accelerator.state.device		)    ,						"cuda:64"		)





      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			create_components()
         accelerator.prepare(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					Any:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			create_components()
         accelerator.prepare(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			get_signature(SCREAMING_SNAKE_CASE__		)

         # saving hook
         def save_config(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		):
            lowerCamelCase__						=			{"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:
               lowerCamelCase__						=			json.load(SCREAMING_SNAKE_CASE__		)

            lowerCamelCase__						=			config["class_name"]

         lowerCamelCase__						=			accelerator.register_save_state_pre_hook(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			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
            lowerCamelCase__						=			"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
            lowerCamelCase__						=			"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 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			create_components()
         lowerCamelCase__						=			None

         # This should work
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					List[str]:
         lowerCamelCase__						=			Accelerator()
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			create_components()
         lowerCamelCase__						=			[1, 2, 3]

         # This should work
         lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__					,       lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					str:
         from transformers import AutoModelForCausalLM

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

         # This should work
         lowerCamelCase__						=			accelerator.prepare(SCREAMING_SNAKE_CASE__		)





      @slow
      @require_bnb
      def 					_lowerCamelCase  (  self		)   ->					Tuple:
         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(SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			"cpu"

         lowerCamelCase__						=			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__		):
            lowerCamelCase__						=			accelerator.prepare(SCREAMING_SNAKE_CASE__		)





      @slow
      @require_bnb
      @require_multi_gpu
      def 					_lowerCamelCase  (  self		)   ->					Dict:
         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(SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			1

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

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

         PartialState._reset_state()





      @slow
      @require_bnb
      @require_multi_gpu
      def 					_lowerCamelCase  (  self		)   ->					List[Any]:
         from transformers import AutoModelForCausalLM

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

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

         # This should work
         lowerCamelCase__						=			accelerator.prepare(SCREAMING_SNAKE_CASE__		)





      @require_cuda
      def 					_lowerCamelCase  (  self		)   ->					List[str]:
         lowerCamelCase__						=			torch.nn.Linear(10    ,						10		)
         lowerCamelCase__						=			torch.optim.SGD(model.parameters()    ,						lr=0.01		)
         lowerCamelCase__						=			Accelerator(cpu=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			accelerator.prepare(SCREAMING_SNAKE_CASE__		)




