



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






			"""simple docstring"""





			def __init__(				self	:       int   )						->					List[str]:
									"""simple docstring"""







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







			def 		snake_case__			(				self	:       int			,     lowercase	:       Dict   )						->					Optional[int]:
									"""simple docstring"""







									return self.lineara(self.batchnorm(self.lineara(lowercase   )   )   )





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       List[Any]			,     lowercase	:       List[Any]			,     *lowercase	:       Tuple			,     **lowercase	:       List[str]   )						->					Optional[Any]:
									"""simple docstring"""







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





class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Optional[Any]			,     lowercase	:       Dict			,     lowercase	:       Dict   )						->					Dict:
									"""simple docstring"""







									return output + 1





class 			_lowerCAmelCase   ( unittest.TestCase			):






			"""simple docstring"""





			def 		snake_case__			(				self	:       Any   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				ModelForTest()
									__lowercase					=				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	:       str   )						->					Any:
									"""simple docstring"""







									__lowercase					=				ModelForTest()
									__lowercase					=				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	:       Union[str, Any]   )						->					str:
									"""simple docstring"""







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

									__lowercase					=				PreForwardHook()
									add_hook_to_module(lowercase			,     lowercase   )
									__lowercase					=				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
									__lowercase					=				PreForwardHook()
									add_hook_to_module(lowercase			,     lowercase   )
									__lowercase					=				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
									__lowercase					=				SequentialHook(PreForwardHook()			,     PreForwardHook()   )
									add_hook_to_module(lowercase			,     lowercase   )

									__lowercase					=				test_model(lowercase   )
									assert torch.allclose(lowercase			,     lowercase			,     atol=1E-5   )





			def 		snake_case__			(				self	:       Optional[Any]   )						->					Union[str, Any]:
									"""simple docstring"""







									__lowercase					=				ModelForTest()
									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				test_model(lowercase   )

									__lowercase					=				PostForwardHook()
									add_hook_to_module(lowercase			,     lowercase   )
									__lowercase					=				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
									__lowercase					=				PostForwardHook()
									add_hook_to_module(lowercase			,     lowercase   )
									__lowercase					=				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
									__lowercase					=				SequentialHook(PostForwardHook()			,     PostForwardHook()   )
									add_hook_to_module(lowercase			,     lowercase   )

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





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







									__lowercase					=				ModelForTest()
									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				test_model(lowercase   )

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

									__lowercase					=				True
									__lowercase					=				test_model(lowercase   )
									self.assertFalse(outputa.requires_grad   )





			@require_multi_gpu
			def 		snake_case__			(				self	:       Dict   )						->					Optional[int]:
									"""simple docstring"""







									__lowercase					=				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
									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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   )   )
									__lowercase					=				torch.randn(2			,     3   ).to(0   )
									__lowercase					=				model(lowercase   )
									self.assertEqual(output.device			,     torch.device(0   )   )





			def 		snake_case__			(				self	:       List[Any]   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				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
									__lowercase					=				{"""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
									__lowercase					=				torch.device(hook_kwargs["""execution_device"""]   )
									self.assertEqual(model.batchnorm.running_mean.device			,     lowercase   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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
									__lowercase					=				{
									    """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"""   )   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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   )						->					Optional[Any]:
									"""simple docstring"""







									__lowercase					=				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
									__lowercase					=				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
									__lowercase					=				torch.device(lowercase   )
									self.assertEqual(model.batchnorm.running_mean.device			,     lowercase   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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"""   )   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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]   )						->					List[str]:
									"""simple docstring"""







									__lowercase					=				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
									__lowercase					=				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
									__lowercase					=				torch.device(lowercase   )
									self.assertEqual(model.batchnorm.running_mean.device			,     lowercase   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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"""   )   )

									__lowercase					=				torch.randn(2			,     3   )
									__lowercase					=				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 dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Features, Value
from .base import TaskTemplate
@dataclass(frozen=_UpperCAmelCase			)
class 			_lowerCAmelCase   ( _UpperCAmelCase			):






			"""simple docstring"""





			lowercase__      :  str      =			field(default="""summarization"""    ,					metadata={"""include_in_asdict_even_if_is_default""": True}			)
			lowercase__      :  ClassVar[Features]      =			Features({"""text""": Value("""string"""			)}			)
			lowercase__      :  ClassVar[Features]      =			Features({"""summary""": Value("""string"""			)}			)
			lowercase__      :  str      =			"text"
			lowercase__      :  str      =			"summary"





			@property
			def 		snake_case__			(				self	:       List[Any]   )						->					Dict[str, str]:
									"""simple docstring"""







									return {self.text_column: "text", self.summary_column: "summary"}





