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 	__snake_case	( nn.Module		):





      '''simple docstring'''






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


          """simple docstring"""


          super().__init__()
          _lowerCamelCase     :		Any 						=						nn.Linear(3    ,  4)
          _lowerCamelCase     :		List[Any] 						=						nn.BatchNormad(4)
          _lowerCamelCase     :		Any 						=						nn.Linear(4    ,  5)




      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any]    ,  _UpperCamelCase     :				Dict)				->Any:


          """simple docstring"""


          return self.lineara(self.batchnorm(self.lineara(_UpperCamelCase)))





class 	__snake_case	( lowercase__		):





      '''simple docstring'''






      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any]    ,  _UpperCamelCase     :				Union[str, Any]    ,  *_UpperCamelCase     :				Union[str, Any]    ,  **_UpperCamelCase     :				List[Any])				->Any:


          """simple docstring"""


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





class 	__snake_case	( lowercase__		):





      '''simple docstring'''






      def 	_SCREAMING_SNAKE_CASE    (					self     :				str    ,  _UpperCamelCase     :				Tuple    ,  _UpperCamelCase     :				Tuple)				->str:


          """simple docstring"""


          return output + 1





class 	__snake_case	( unittest.TestCase		):





      '''simple docstring'''






      def 	_SCREAMING_SNAKE_CASE    (					self     :				Any)				->str:


          """simple docstring"""


          _lowerCamelCase     :		List[str] 						=						ModelForTest()
          _lowerCamelCase     :		Dict 						=						ModelHook()

          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          self.assertEqual(test_model._hf_hook    ,  _UpperCamelCase)
          self.assertTrue(hasattr(_UpperCamelCase    ,  """_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(_UpperCamelCase)
          self.assertFalse(hasattr(_UpperCamelCase    ,  """_hf_hook"""))
          self.assertFalse(hasattr(_UpperCamelCase    ,  """_old_forward"""))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				str)				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		List[Any] 						=						ModelForTest()
          _lowerCamelCase     :		Optional[Any] 						=						ModelHook()

          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase    ,  append=_UpperCamelCase)

          self.assertEqual(isinstance(test_model._hf_hook    ,  _UpperCamelCase)    ,  _UpperCamelCase)
          self.assertEqual(len(test_model._hf_hook.hooks)    ,  2)
          self.assertTrue(hasattr(_UpperCamelCase    ,  """_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(_UpperCamelCase)
          self.assertFalse(hasattr(_UpperCamelCase    ,  """_hf_hook"""))
          self.assertFalse(hasattr(_UpperCamelCase    ,  """_old_forward"""))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				List[str])				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Any 						=						ModelForTest()
          _lowerCamelCase     :		Any 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Any 						=						test_model(x + 1)
          _lowerCamelCase     :		int 						=						test_model(x + 2)

          _lowerCamelCase     :		List[str] 						=						PreForwardHook()
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          _lowerCamelCase     :		Optional[int] 						=						test_model(_UpperCamelCase)
          self.assertTrue(torch.allclose(_UpperCamelCase    ,  _UpperCamelCase    ,  atol=1E-5))

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

          # You need to use the sequential hook to chain two or more hooks
          _lowerCamelCase     :		Optional[int] 						=						SequentialHook(PreForwardHook()    ,  PreForwardHook())
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)

          _lowerCamelCase     :		str 						=						test_model(_UpperCamelCase)
          assert torch.allclose(_UpperCamelCase    ,  _UpperCamelCase    ,  atol=1E-5)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Tuple:


          """simple docstring"""


          _lowerCamelCase     :		int 						=						ModelForTest()
          _lowerCamelCase     :		Tuple 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Tuple 						=						test_model(_UpperCamelCase)

          _lowerCamelCase     :		str 						=						PostForwardHook()
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          _lowerCamelCase     :		Union[str, Any] 						=						test_model(_UpperCamelCase)
          self.assertTrue(torch.allclose(_UpperCamelCase    ,  output + 1    ,  atol=1E-5))

          # Attaching a hook to a model when it already has one replaces, does not chain
          _lowerCamelCase     :		Tuple 						=						PostForwardHook()
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          _lowerCamelCase     :		Optional[Any] 						=						test_model(_UpperCamelCase)
          self.assertTrue(torch.allclose(_UpperCamelCase    ,  output + 1    ,  atol=1E-5))

          # You need to use the sequential hook to chain two or more hooks
          _lowerCamelCase     :		List[str] 						=						SequentialHook(PostForwardHook()    ,  PostForwardHook())
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)

          _lowerCamelCase     :		Tuple 						=						test_model(_UpperCamelCase)
          assert torch.allclose(_UpperCamelCase    ,  output + 2    ,  atol=1E-5)



      def 	_SCREAMING_SNAKE_CASE    (					self     :				int)				->Optional[int]:


          """simple docstring"""


          _lowerCamelCase     :		str 						=						ModelForTest()
          _lowerCamelCase     :		Tuple 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		int 						=						test_model(_UpperCamelCase)

          _lowerCamelCase     :		Dict 						=						PostForwardHook()
          add_hook_to_module(_UpperCamelCase    ,  _UpperCamelCase)
          _lowerCamelCase     :		int 						=						test_model(_UpperCamelCase)
          self.assertTrue(torch.allclose(_UpperCamelCase    ,  output + 1))
          self.assertTrue(outputa.requires_grad)

          _lowerCamelCase     :		Optional[Any] 						=						True
          _lowerCamelCase     :		Optional[Any] 						=						test_model(_UpperCamelCase)
          self.assertFalse(outputa.requires_grad)



      @require_multi_gpu
      def 	_SCREAMING_SNAKE_CASE    (					self     :				int)				->Dict:


          """simple docstring"""


          _lowerCamelCase     :		Any 						=						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
          _lowerCamelCase     :		Any 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Union[str, Any] 						=						model(_UpperCamelCase)
          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(_UpperCamelCase    ,  AlignDevicesHook(io_same_device=_UpperCamelCase))
          _lowerCamelCase     :		List[str] 						=						torch.randn(2    ,  3).to(0)
          _lowerCamelCase     :		Tuple 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  torch.device(0))



      def 	_SCREAMING_SNAKE_CASE    (					self     :				Tuple)				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Tuple 						=						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
          _lowerCamelCase     :		int 						=						{"execution_device": 0 if torch.cuda.is_available() else "cpu", "offload": True}

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

          # 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
          _lowerCamelCase     :		str 						=						torch.device(hook_kwargs["""execution_device"""])
          self.assertEqual(model.batchnorm.running_mean.device    ,  _UpperCamelCase)

          _lowerCamelCase     :		int 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		str 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # 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
          _lowerCamelCase     :		Tuple 						=						{
              "execution_device": 0 if torch.cuda.is_available() else "cpu",
              "offload": True,
              "offload_buffers": True,
          }

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

          # 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"""))

          _lowerCamelCase     :		int 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Tuple 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # 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 	_SCREAMING_SNAKE_CASE    (					self     :				Dict)				->Union[str, Any]:


          """simple docstring"""


          _lowerCamelCase     :		int 						=						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
          _lowerCamelCase     :		Any 						=						0 if torch.cuda.is_available() else "cpu"
          attach_align_device_hook(_UpperCamelCase    ,  execution_device=_UpperCamelCase    ,  offload=_UpperCamelCase)

          # 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
          _lowerCamelCase     :		Optional[int] 						=						torch.device(_UpperCamelCase)
          self.assertEqual(model.batchnorm.running_mean.device    ,  _UpperCamelCase)

          _lowerCamelCase     :		Any 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Dict 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(_UpperCamelCase)
          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(_UpperCamelCase    ,  execution_device=_UpperCamelCase    ,  offload=_UpperCamelCase    ,  offload_buffers=_UpperCamelCase)

          # 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"""))

          _lowerCamelCase     :		Tuple 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Any 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(_UpperCamelCase)
          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 	_SCREAMING_SNAKE_CASE    (					self     :				Tuple)				->List[str]:


          """simple docstring"""


          _lowerCamelCase     :		Tuple 						=						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
          _lowerCamelCase     :		Union[str, Any] 						=						0 if torch.cuda.is_available() else "cpu"
          attach_align_device_hook(
              _UpperCamelCase    ,  execution_device=_UpperCamelCase    ,  offload=_UpperCamelCase    ,  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
          _lowerCamelCase     :		str 						=						torch.device(_UpperCamelCase)
          self.assertEqual(model.batchnorm.running_mean.device    ,  _UpperCamelCase)

          _lowerCamelCase     :		Any 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		str 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(_UpperCamelCase)
          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(
              _UpperCamelCase    ,  execution_device=_UpperCamelCase    ,  offload=_UpperCamelCase    ,  weights_map=model.state_dict()    ,  offload_buffers=_UpperCamelCase    ,  )

          # 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"""))

          _lowerCamelCase     :		Optional[int] 						=						torch.randn(2    ,  3)
          _lowerCamelCase     :		Dict 						=						model(_UpperCamelCase)
          self.assertEqual(output.device    ,  _UpperCamelCase)

          # Removing hooks loads back the weights in the model.
          remove_hook_from_submodules(_UpperCamelCase)
          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 Optional

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader




class 	__snake_case	( __lowerCAmelCase		):





      '''simple docstring'''






      def __init__(					self     :				Dict    ,  _UpperCamelCase     :				NestedDataStructureLike[PathLike]    ,  _UpperCamelCase     :				Optional[NamedSplit] = None    ,  _UpperCamelCase     :				Optional[Features] = None    ,  _UpperCamelCase     :				str = None    ,  _UpperCamelCase     :				bool = False    ,  _UpperCamelCase     :				bool = False    ,  _UpperCamelCase     :				Optional[int] = None    ,  **_UpperCamelCase     :				Tuple    ,  )				->Union[str, Any]:


          """simple docstring"""


          super().__init__(
              _UpperCamelCase    ,  split=_UpperCamelCase    ,  features=_UpperCamelCase    ,  cache_dir=_UpperCamelCase    ,  keep_in_memory=_UpperCamelCase    ,  streaming=_UpperCamelCase    ,  num_proc=_UpperCamelCase    ,  **_UpperCamelCase    ,  )
          _lowerCamelCase     :		List[Any] 						=						path_or_paths if isinstance(_UpperCamelCase    ,  _UpperCamelCase) else {self.split: path_or_paths}
          _lowerCamelCase     :		Any 						=						Text(
              cache_dir=_UpperCamelCase    ,  data_files=_UpperCamelCase    ,  features=_UpperCamelCase    ,  **_UpperCamelCase    ,  )




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


          """simple docstring"""


          if self.streaming:
              _lowerCamelCase     :		Tuple 						=						self.builder.as_streaming_dataset(split=self.split)
          # Build regular (map-style) dataset
          else:
              _lowerCamelCase     :		List[Any] 						=						None
              _lowerCamelCase     :		Any 						=						None
              _lowerCamelCase     :		List[str] 						=						None
              _lowerCamelCase     :		Dict 						=						None

              self.builder.download_and_prepare(
                  download_config=_UpperCamelCase    ,  download_mode=_UpperCamelCase    ,  verification_mode=_UpperCamelCase    ,  base_path=_UpperCamelCase    ,  num_proc=self.num_proc    ,  )
              _lowerCamelCase     :		Optional[int] 						=						self.builder.as_dataset(
                  split=self.split    ,  verification_mode=_UpperCamelCase    ,  in_memory=self.keep_in_memory)
          return dataset




