import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

from transformers import AutoConfig, TFGPTaLMHeadModel, is_keras_nlp_available, is_tf_available
from transformers.models.gpta.tokenization_gpta import GPTaTokenizer
from transformers.testing_utils import require_keras_nlp, require_tf, slow


if is_tf_available():
    import tensorflow as tf

if is_keras_nlp_available():
    from transformers.models.gpta import TFGPTaTokenizer


__lowercase										= ['''gpt2''']
__lowercase										= '''gpt2'''

if is_tf_available():


    class      lowerCamelCase_					(       tf.Module			):






        '''simple docstring'''






        def __init__(     self			,			__lowercase)	->     List[str]:
            super().__init__()
            __UpperCamelCase					:Any        =		tokenizer
            __UpperCamelCase					:Tuple        =		AutoConfig.from_pretrained(__lowercase)
            __UpperCamelCase					:Optional[Any]        =		TFGPTaLMHeadModel.from_config(__lowercase)







        @tf.function(input_signature=(tf.TensorSpec((None,)			,			tf.string			,			name='''text'''),))
        def       UpperCamelCase__       (     self			,			__lowercase)	->     Union[str, Any]:
            __UpperCamelCase					:Tuple        =		self.tokenizer(__lowercase)
            __UpperCamelCase					:int        =		tokenized['''input_ids'''].to_tensor()

            __UpperCamelCase					:int        =		tf.cast(input_ids_dense > 0			,			tf.intaa)
            # input_mask = tf.reshape(input_mask, [-1, MAX_SEQ_LEN])

            __UpperCamelCase					:int        =		self.model(input_ids=__lowercase			,			attention_mask=__lowercase)['''logits''']

            return outputs


@require_tf
@require_keras_nlp
class      lowerCamelCase_					(       unittest.TestCase			):






    '''simple docstring'''






    def       UpperCamelCase__       (     self)	->     Tuple:
        super().setUp()

        __UpperCamelCase					:Tuple        =		[GPTaTokenizer.from_pretrained(__lowercase) for checkpoint in (TOKENIZER_CHECKPOINTS)]
        __UpperCamelCase					:List[str]        =		[TFGPTaTokenizer.from_pretrained(__lowercase) for checkpoint in TOKENIZER_CHECKPOINTS]
        assert len(self.tokenizers) == len(self.tf_tokenizers)

        __UpperCamelCase					:List[Any]        =		[
            '''This is a straightforward English test sentence.''',
            '''This one has some weird characters\rto\nsee\r\nif  those\u00E9break things.''',
            '''Now we\'re going to add some Chinese: 一 二 三 一二三''',
            '''And some much more rare Chinese: 齉 堃 齉堃''',
            '''Je vais aussi écrire en français pour tester les accents''',
            '''Classical Irish also has some unusual characters, so in they go: Gaelaċ, ꝼ''',
        ]
        __UpperCamelCase					:Optional[Any]        =		list(zip(self.test_sentences			,			self.test_sentences[::-1]))







    def       UpperCamelCase__       (     self)	->     Tuple:
        for tokenizer, tf_tokenizer in zip(self.tokenizers			,			self.tf_tokenizers):
            for test_inputs in self.test_sentences:
                __UpperCamelCase					:Any        =		tokenizer([test_inputs]			,			return_tensors='''tf''')
                __UpperCamelCase					:str        =		tf_tokenizer([test_inputs])

                for key in python_outputs.keys():
                    # convert them to numpy to avoid messing with ragged tensors
                    __UpperCamelCase					:Optional[Any]        =		python_outputs[key].numpy()
                    __UpperCamelCase					:Optional[int]        =		tf_outputs[key].numpy()

                    self.assertTrue(tf.reduce_all(python_outputs_values.shape == tf_outputs_values.shape))
                    self.assertTrue(tf.reduce_all(tf.cast(__lowercase			,			tf.intaa) == tf_outputs_values))







    @slow
    def       UpperCamelCase__       (     self)	->     str:
        for tf_tokenizer in self.tf_tokenizers:
            __UpperCamelCase					:Dict        =		tf.function(__lowercase)
            for test_inputs in self.test_sentences:
                __UpperCamelCase					:List[str]        =		tf.constant(__lowercase)
                __UpperCamelCase					:Optional[int]        =		compiled_tokenizer(__lowercase)
                __UpperCamelCase					:str        =		tf_tokenizer(__lowercase)

                for key in eager_outputs.keys():
                    self.assertTrue(tf.reduce_all(eager_outputs[key] == compiled_outputs[key]))







    @slow
    def       UpperCamelCase__       (     self)	->     List[Any]:
        for tf_tokenizer in self.tf_tokenizers:
            __UpperCamelCase					:Tuple        =		ModelToSave(tokenizer=__lowercase)
            __UpperCamelCase					:List[Any]        =		tf.convert_to_tensor([self.test_sentences[0]])
            __UpperCamelCase					:int        =		model.serving(__lowercase)  # Build model with some sample inputs
            with TemporaryDirectory() as tempdir:
                __UpperCamelCase					:int        =		Path(__lowercase) / '''saved.model'''
                tf.saved_model.save(__lowercase			,			__lowercase			,			signatures={'''serving_default''': model.serving})
                __UpperCamelCase					:Union[str, Any]        =		tf.saved_model.load(__lowercase)
            __UpperCamelCase					:Union[str, Any]        =		loaded_model.signatures['''serving_default'''](__lowercase)['''output_0''']
            # We may see small differences because the loaded model is compiled, so we need an epsilon for the test
            self.assertTrue(tf.reduce_all(out == loaded_output))







    @slow
    def       UpperCamelCase__       (     self)	->     str:
        for tf_tokenizer in self.tf_tokenizers:
            __UpperCamelCase					:str        =		tf.convert_to_tensor([self.test_sentences[0]])
            __UpperCamelCase					:Dict        =		tf_tokenizer(__lowercase)  # Build model with some sample inputs

            __UpperCamelCase					:Dict        =		tf_tokenizer.get_config()
            __UpperCamelCase					:Tuple        =		TFGPTaTokenizer.from_config(__lowercase)
            __UpperCamelCase					:str        =		model_from_config(__lowercase)

            for key in from_config_output.keys():
                self.assertTrue(tf.reduce_all(from_config_output[key] == out[key]))







    @slow
    def       UpperCamelCase__       (     self)	->     int:
        for tf_tokenizer in self.tf_tokenizers:
            # for the test to run
            __UpperCamelCase					:Union[str, Any]        =		123_123

            for max_length in [3, 5, 1_024]:
                __UpperCamelCase					:Optional[int]        =		tf.convert_to_tensor([self.test_sentences[0]])
                __UpperCamelCase					:Tuple        =		tf_tokenizer(__lowercase			,			max_length=__lowercase)

                __UpperCamelCase					:Any        =		out['''input_ids'''].numpy().shape[1]

                assert out_length == max_length






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

    return model, optimizer, scheduler, train_dl, valid_dl



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







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



def lowerCamelCase  (       SCREAMING_SNAKE_CASE     ):
    '''simple docstring'''







    __UpperCamelCase					:Union[str, Any]        =		torch.nn.Linear(*tuple(model.weight.T.shape     )     ).state_dict()
    model.load_state_dict(SCREAMING_SNAKE_CASE     )


class      lowerCamelCase_					(       UpperCAmelCase_			):






    '''simple docstring'''






    @require_cuda
    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Dict        =		Accelerator()
        assert PartialState._shared_state["_cpu"] is False
        assert PartialState._shared_state["device"].type == "cuda"
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:Any        =		Accelerator(cpu=__lowercase)







    def       UpperCamelCase__       (     self)	->     Any:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase					:List[Any]        =		GradientState()
        assert state.num_steps == 1
        __UpperCamelCase					:Any        =		4
        assert state.num_steps == 4

        assert state.sync_gradients is True
        __UpperCamelCase					:int        =		False
        assert state.sync_gradients is False
        GradientState._reset_state()







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Tuple        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Optional[Any]        =		create_components()

        (
            (
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	(
            __UpperCamelCase
        )  ,	
        )					:int        =		accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)

        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)	->     Optional[Any]:
        __UpperCamelCase					:str        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Optional[Any]        =		create_components()
        accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)
        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)	->     Union[str, Any]:
        PartialState._reset_state()

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

        with patch('''torch.cuda.set_device'''			,			__lowercase), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64'''):
            __UpperCamelCase					:Optional[Any]        =		Accelerator()
            self.assertEqual(str(accelerator.state.device)			,			'''cuda:64''')







    def       UpperCamelCase__       (     self)	->     List[Any]:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:List[Any]        =		create_components()
        accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)

        __UpperCamelCase					:Tuple        =		get_signature(__lowercase)

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

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

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







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:List[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Any        =		create_components()
        accelerator.prepare(__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)

        __UpperCamelCase					:Any        =		get_signature(__lowercase)

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

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

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

            __UpperCamelCase					:Dict        =		config['''class_name''']

        __UpperCamelCase					:Union[str, Any]        =		accelerator.register_save_state_pre_hook(__lowercase)
        __UpperCamelCase					:Any        =		accelerator.register_load_state_pre_hook(__lowercase)

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

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

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

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

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

            # random class name to verify correct one is loaded
            __UpperCamelCase					:Dict        =		'''random'''

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

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







    def       UpperCamelCase__       (     self)	->     Dict:
        __UpperCamelCase					:Optional[Any]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Union[str, Any]        =		create_components()
        __UpperCamelCase					:Optional[Any]        =		None

        # This should work
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:List[Any]        =		accelerator.prepare(
            __lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase			,			__lowercase)
        self.assertTrue(dummy_obj is None)







    def       UpperCamelCase__       (     self)	->     List[str]:
        __UpperCamelCase					:List[str]        =		Accelerator()
        __UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase  ,	__UpperCamelCase					:Any        =		create_components()
        __UpperCamelCase					:Dict        =		[1, 2, 3]

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







    @slow
    @require_bnb
    def       UpperCamelCase__       (     self)	->     int:
        from transformers import AutoModelForCausalLM

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

        # This should work
        __UpperCamelCase					:int        =		accelerator.prepare(__lowercase)







    @slow
    @require_bnb
    def       UpperCamelCase__       (     self)	->     List[str]:
        from transformers import AutoModelForCausalLM

        __UpperCamelCase					:str        =		Accelerator()

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

        __UpperCamelCase					:List[Any]        =		AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''			,			device_map=__lowercase			,			load_in_abit=__lowercase			,			llm_inta_enable_fpaa_cpu_offload=__lowercase)

        # This should not work and get value error
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:Union[str, Any]        =		accelerator.prepare(__lowercase)







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

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

        with init_empty_weights():
            __UpperCamelCase					:Tuple        =		AutoModelForCausalLM.from_pretrained(
                '''EleutherAI/gpt-neo-125m'''			,			)
            model.tie_weights()
            __UpperCamelCase					:int        =		infer_auto_device_map(__lowercase)
            __UpperCamelCase					:List[Any]        =		1

        __UpperCamelCase					:int        =		AutoModelForCausalLM.from_pretrained(
            '''EleutherAI/gpt-neo-125m'''			,			load_in_abit=__lowercase			,			device_map=__lowercase			,			)
        __UpperCamelCase					:Dict        =		Accelerator()

        # This should not work and get value error
        with self.assertRaises(__lowercase):
            __UpperCamelCase					:Any        =		accelerator.prepare(__lowercase)

        PartialState._reset_state()







    @slow
    @require_bnb
    @require_multi_gpu
    def       UpperCamelCase__       (     self)	->     Dict:
        from transformers import AutoModelForCausalLM

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

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

        # This should work
        __UpperCamelCase					:int        =		accelerator.prepare(__lowercase)







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





