

import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu


class  A						(							nn.Module ):







    """simple docstring"""





    def __init__(      self :       List[str]		)->    Union[str, Any]:





        '''simple docstring'''







        super().__init__()
        A__          =   nn.Linear(3,4		)
        A__          =   nn.BatchNormad(4		)
        A__          =   nn.Linear(4,5		)


    def 						snake_case__     (      self :       Optional[int],lowercase_ :       List[Any]		)->    Any:





        '''simple docstring'''







        return self.lineara(self.batchnorm(self.lineara(lowercase_		)		)		)




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    def 						snake_case__     (      self :       List[Any],lowercase_ :       List[str],*lowercase_ :       Optional[int],**lowercase_ :       Tuple		)->    Any:





        '''simple docstring'''







        return (args[0] + 1,) + args[1:], kwargs




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    def 						snake_case__     (      self :       List[str],lowercase_ :       int,lowercase_ :       Optional[Any]		)->    List[str]:





        '''simple docstring'''







        return output + 1




class  A						(							unittest.TestCase ):







    """simple docstring"""





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





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   ModelHook()

        add_hook_to_module(lowercase_,lowercase_		)
        self.assertEqual(test_model._hf_hook,lowercase_		)
        self.assertTrue(hasattr(lowercase_,'_old_forward'		)		)

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__,'forward'		)
        self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		),['x']		)

        remove_hook_from_module(lowercase_		)
        self.assertFalse(hasattr(lowercase_,'_hf_hook'		)		)
        self.assertFalse(hasattr(lowercase_,'_old_forward'		)		)


    def 						snake_case__     (      self :       Optional[Any]		)->    str:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   ModelHook()

        add_hook_to_module(lowercase_,lowercase_		)
        add_hook_to_module(lowercase_,lowercase_,append=lowercase_		)

        self.assertEqual(isinstance(test_model._hf_hook,lowercase_		),lowercase_		)
        self.assertEqual(len(test_model._hf_hook.hooks		),2		)
        self.assertTrue(hasattr(lowercase_,'_old_forward'		)		)

        # Check adding the hook did not change the name or the signature
        self.assertEqual(test_model.forward.__name__,'forward'		)
        self.assertListEqual(list(inspect.signature(test_model.forward		).parameters		),['x']		)

        remove_hook_from_module(lowercase_		)
        self.assertFalse(hasattr(lowercase_,'_hf_hook'		)		)
        self.assertFalse(hasattr(lowercase_,'_old_forward'		)		)


    def 						snake_case__     (      self :       str		)->    Any:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   torch.randn(2,3		)
        A__          =   test_model(x + 1		)
        A__          =   test_model(x + 2		)

        A__          =   PreForwardHook()
        add_hook_to_module(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,lowercase_,atol=1E-5		)		)

        # Attaching a hook to a model when it already has one replaces, does not chain
        A__          =   PreForwardHook()
        add_hook_to_module(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,lowercase_,atol=1E-5		)		)

        # You need to use the sequential hook to chain two or more hooks
        A__          =   SequentialHook(PreForwardHook(),PreForwardHook()		)
        add_hook_to_module(lowercase_,lowercase_		)

        A__          =   test_model(lowercase_		)
        assert torch.allclose(lowercase_,lowercase_,atol=1E-5		)


    def 						snake_case__     (      self :       Optional[Any]		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   torch.randn(2,3		)
        A__          =   test_model(lowercase_		)

        A__          =   PostForwardHook()
        add_hook_to_module(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,output + 1,atol=1E-5		)		)

        # Attaching a hook to a model when it already has one replaces, does not chain
        A__          =   PostForwardHook()
        add_hook_to_module(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,output + 1,atol=1E-5		)		)

        # You need to use the sequential hook to chain two or more hooks
        A__          =   SequentialHook(PostForwardHook(),PostForwardHook()		)
        add_hook_to_module(lowercase_,lowercase_		)

        A__          =   test_model(lowercase_		)
        assert torch.allclose(lowercase_,output + 2,atol=1E-5		)


    def 						snake_case__     (      self :       Dict		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   ModelForTest()
        A__          =   torch.randn(2,3		)
        A__          =   test_model(lowercase_		)

        A__          =   PostForwardHook()
        add_hook_to_module(lowercase_,lowercase_		)
        A__          =   test_model(lowercase_		)
        self.assertTrue(torch.allclose(lowercase_,output + 1		)		)
        self.assertTrue(outputa.requires_grad		)

        A__          =   True
        A__          =   test_model(lowercase_		)
        self.assertFalse(outputa.requires_grad		)


    @require_multi_gpu
    def 						snake_case__     (      self :       Optional[int]		)->    Any:





        '''simple docstring'''







        A__          =   ModelForTest()
        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # This will move each submodule on different devices
        add_hook_to_module(model.lineara,AlignDevicesHook(execution_device=0		)		)
        add_hook_to_module(model.batchnorm,AlignDevicesHook(execution_device=0		)		)
        add_hook_to_module(model.lineara,AlignDevicesHook(execution_device=1		)		)

        self.assertEqual(model.lineara.weight.device,torch.device(0		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device(0		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device(0		)		)
        self.assertEqual(model.lineara.weight.device,torch.device(1		)		)

        # We can still make a forward pass. The input does not need to be on any particular device
        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,torch.device(1		)		)

        # We can add a general hook to put back output on same device as input.
        add_hook_to_module(lowercase_,AlignDevicesHook(io_same_device=lowercase_		)		)
        A__          =   torch.randn(2,3		).to(0		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,torch.device(0		)		)


    def 						snake_case__     (      self :       Dict		)->    Optional[Any]:





        '''simple docstring'''







        A__          =   ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # This will move each submodule on different devices
        A__          =   {'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

        add_hook_to_module(model.lineara,AlignDevicesHook(**lowercase_		)		)
        add_hook_to_module(model.batchnorm,AlignDevicesHook(**lowercase_		)		)
        add_hook_to_module(model.lineara,AlignDevicesHook(**lowercase_		)		)

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(hook_kwargs['execution_device']		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara		)
        remove_hook_from_module(model.batchnorm		)
        remove_hook_from_module(model.lineara		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        A__          =   {
            'execution_device': 0 if torch.cuda.is_available() else 'cpu',
            'offload': True,
            'offload_buffers': True,
        }

        add_hook_to_module(model.lineara,AlignDevicesHook(**lowercase_		)		)
        add_hook_to_module(model.batchnorm,AlignDevicesHook(**lowercase_		)		)
        add_hook_to_module(model.lineara,AlignDevicesHook(**lowercase_		)		)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_module(model.lineara		)
        remove_hook_from_module(model.batchnorm		)
        remove_hook_from_module(model.lineara		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)


    def 						snake_case__     (      self :       Tuple		)->    Any:





        '''simple docstring'''







        A__          =   ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # This will move each submodule on different devices
        A__          =   0 if torch.cuda.is_available() else 'cpu'
        attach_align_device_hook(lowercase_,execution_device=lowercase_,offload=lowercase_		)

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(lowercase_		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        attach_align_device_hook(lowercase_,execution_device=lowercase_,offload=lowercase_,offload_buffers=lowercase_		)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)


    def 						snake_case__     (      self :       List[Any]		)->    Union[str, Any]:





        '''simple docstring'''







        A__          =   ModelForTest()

        # Everything is on CPU
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # This will move each submodule on different devices
        A__          =   0 if torch.cuda.is_available() else 'cpu'
        attach_align_device_hook(
            lowercase_,execution_device=lowercase_,offload=lowercase_,weights_map=model.state_dict()		)

        # Parameters have been offloaded, so on the meta device
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        # Buffers are not included in the offload by default, so are on the execution device
        A__          =   torch.device(lowercase_		)
        self.assertEqual(model.batchnorm.running_mean.device,lowercase_		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)

        # Now test with buffers included in the offload
        attach_align_device_hook(
            lowercase_,execution_device=lowercase_,offload=lowercase_,weights_map=model.state_dict(),offload_buffers=lowercase_,)

        # Parameters have been offloaded, so on the meta device, buffers included
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('meta'		)		)
        self.assertEqual(model.batchnorm.running_mean.device,torch.device('meta'		)		)

        A__          =   torch.randn(2,3		)
        A__          =   model(lowercase_		)
        self.assertEqual(output.device,lowercase_		)

        # Removing hooks loads back the weights in the model.
        remove_hook_from_submodules(lowercase_		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.batchnorm.weight.device,torch.device('cpu'		)		)
        self.assertEqual(model.lineara.weight.device,torch.device('cpu'		)		)




from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxSeqaSeqConfigWithPast
from ...utils import logging


lowercase_       		=  logging.get_logger(__name__)

lowercase_       		=  {
    "t5-small": "https://huggingface.co/t5-small/resolve/main/config.json",
    "t5-base": "https://huggingface.co/t5-base/resolve/main/config.json",
    "t5-large": "https://huggingface.co/t5-large/resolve/main/config.json",
    "t5-3b": "https://huggingface.co/t5-3b/resolve/main/config.json",
    "t5-11b": "https://huggingface.co/t5-11b/resolve/main/config.json",
}


class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    lowerCamelCase			=							't5'
    lowerCamelCase			=							['past_key_values']
    lowerCamelCase			=							{'hidden_size': 'd_model', 'num_attention_heads': 'num_heads', 'num_hidden_layers': 'num_layers'}


    def __init__(      self :       Union[str, Any],lowercase_ :       int=3_2_1_2_8,lowercase_ :       int=5_1_2,lowercase_ :       List[str]=6_4,lowercase_ :       Tuple=2_0_4_8,lowercase_ :       Any=6,lowercase_ :       List[str]=None,lowercase_ :       Union[str, Any]=8,lowercase_ :       int=3_2,lowercase_ :       Dict=1_2_8,lowercase_ :       Optional[int]=0.1,lowercase_ :       List[str]=1E-6,lowercase_ :       Tuple=1.0,lowercase_ :       Any="relu",lowercase_ :       Union[str, Any]=True,lowercase_ :       Optional[Any]=True,lowercase_ :       int=0,lowercase_ :       str=1,**lowercase_ :       str,)->    Any:





        '''simple docstring'''







        A__          =   vocab_size
        A__          =   d_model
        A__          =   d_kv
        A__          =   d_ff
        A__          =   num_layers
        A__          =   (
            num_decoder_layers if num_decoder_layers is not None else self.num_layers
        )  # default = symmetry
        A__          =   num_heads
        A__          =   relative_attention_num_buckets
        A__          =   relative_attention_max_distance
        A__          =   dropout_rate
        A__          =   layer_norm_epsilon
        A__          =   initializer_factor
        A__          =   feed_forward_proj
        A__          =   use_cache

        A__          =   self.feed_forward_proj.split('-'		)
        A__          =   act_info[-1]
        A__          =   act_info[0] == 'gated'

        if len(lowercase_		) > 1 and act_info[0] != "gated" or len(lowercase_		) > 2:
            raise ValueError(
                F'`feed_forward_proj`: {feed_forward_proj} is not a valid activation function of the dense layer.'
                'Please make sure `feed_forward_proj` is of the format `gated-{ACT_FN}` or `{ACT_FN}`, e.g. '
                '\'gated-gelu\' or \'relu\''		)

        # for backwards compatibility
        if feed_forward_proj == "gated-gelu":
            A__          =   'gelu_new'

        super().__init__(
            pad_token_id=lowercase_,eos_token_id=lowercase_,is_encoder_decoder=lowercase_,**lowercase_,)




class  A						(							_UpperCAmelCase ):







    """simple docstring"""





    @property
    def 						snake_case__     (      self :       Tuple		)->    Mapping[str, Mapping[int, str]]:





        '''simple docstring'''







        A__          =   {
            'input_ids': {0: 'batch', 1: 'encoder_sequence'},
            'attention_mask': {0: 'batch', 1: 'encoder_sequence'},
        }
        if self.use_past:
            A__          =   'past_encoder_sequence + sequence'
            A__          =   {0: 'batch'}
            A__          =   {0: 'batch', 1: 'past_decoder_sequence + sequence'}
        else:
            A__          =   {0: 'batch', 1: 'decoder_sequence'}
            A__          =   {0: 'batch', 1: 'decoder_sequence'}

        if self.use_past:
            self.fill_with_past_key_values_(lowercase_,direction='inputs'		)

        return common_inputs


    @property
    def 						snake_case__     (      self :       Any		)->    int:





        '''simple docstring'''







        return 1_3

