



"""simple docstring"""





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   ) ->						str:







								'''simple docstring'''

								super().__init__()
								snake_case_   : str						=			nn.Linear(3							,    4   )
								snake_case_   : Tuple						=			nn.BatchNormad(4   )
								snake_case_   : str						=			nn.Linear(4							,    5   )





		def    UpperCAmelCase__							(					self							,    _lowercase   ) ->						Any:







								'''simple docstring'''

								return self.lineara(self.batchnorm(self.lineara(_lowercase   )   )   )




class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self							,    _lowercase							,    *_lowercase							,    **_lowercase   ) ->						Optional[int]:







								'''simple docstring'''

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




class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self							,    _lowercase							,    _lowercase   ) ->						Any:







								'''simple docstring'''

								return output + 1




class    _lowerCAmelCase    (				unittest.TestCase		):





		"""simple docstring"""






		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								snake_case_   : Tuple						=			ModelForTest()
								snake_case_   : Union[str, Any]						=			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    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : str						=			ModelForTest()
								snake_case_   : List[Any]						=			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    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								snake_case_   : Union[str, Any]						=			ModelForTest()
								snake_case_   : List[Any]						=			torch.randn(2							,    3   )
								snake_case_   : List[Any]						=			test_model(x + 1   )
								snake_case_   : Tuple						=			test_model(x + 2   )

								snake_case_   : List[str]						=			PreForwardHook()
								add_hook_to_module(_lowercase							,    _lowercase   )
								snake_case_   : Dict						=			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
								snake_case_   : List[Any]						=			PreForwardHook()
								add_hook_to_module(_lowercase							,    _lowercase   )
								snake_case_   : Tuple						=			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
								snake_case_   : Any						=			SequentialHook(PreForwardHook()							,    PreForwardHook()   )
								add_hook_to_module(_lowercase							,    _lowercase   )

								snake_case_   : List[Any]						=			test_model(_lowercase   )
								assert torch.allclose(_lowercase							,    _lowercase							,    atol=1E-5   )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Tuple						=			ModelForTest()
								snake_case_   : Union[str, Any]						=			torch.randn(2							,    3   )
								snake_case_   : Optional[Any]						=			test_model(_lowercase   )

								snake_case_   : int						=			PostForwardHook()
								add_hook_to_module(_lowercase							,    _lowercase   )
								snake_case_   : Tuple						=			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
								snake_case_   : Dict						=			PostForwardHook()
								add_hook_to_module(_lowercase							,    _lowercase   )
								snake_case_   : Dict						=			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
								snake_case_   : List[str]						=			SequentialHook(PostForwardHook()							,    PostForwardHook()   )
								add_hook_to_module(_lowercase							,    _lowercase   )

								snake_case_   : str						=			test_model(_lowercase   )
								assert torch.allclose(_lowercase							,    output + 2							,    atol=1E-5   )





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : Tuple						=			ModelForTest()
								snake_case_   : Dict						=			torch.randn(2							,    3   )
								snake_case_   : Optional[int]						=			test_model(_lowercase   )

								snake_case_   : List[str]						=			PostForwardHook()
								add_hook_to_module(_lowercase							,    _lowercase   )
								snake_case_   : Tuple						=			test_model(_lowercase   )
								self.assertTrue(torch.allclose(_lowercase							,    output + 1   )   )
								self.assertTrue(outputa.requires_grad   )

								snake_case_   : Optional[int]						=			True
								snake_case_   : List[Any]						=			test_model(_lowercase   )
								self.assertFalse(outputa.requires_grad   )





		@require_multi_gpu
		def    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								snake_case_   : Optional[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
								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
								snake_case_   : Tuple						=			torch.randn(2							,    3   )
								snake_case_   : Optional[Any]						=			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   )   )
								snake_case_   : List[Any]						=			torch.randn(2							,    3   ).to(0   )
								snake_case_   : Tuple						=			model(_lowercase   )
								self.assertEqual(output.device							,    torch.device(0   )   )





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







								'''simple docstring'''

								snake_case_   : 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
								snake_case_   : Optional[Any]						=			{"""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
								snake_case_   : Union[str, Any]						=			torch.device(hook_kwargs["""execution_device"""]   )
								self.assertEqual(model.batchnorm.running_mean.device							,    _lowercase   )

								snake_case_   : Optional[int]						=			torch.randn(2							,    3   )
								snake_case_   : Dict						=			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
								snake_case_   : List[str]						=			{
								    """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"""   )   )

								snake_case_   : Dict						=			torch.randn(2							,    3   )
								snake_case_   : Dict						=			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    UpperCAmelCase__							(					self   ) ->						List[str]:







								'''simple docstring'''

								snake_case_   : Union[str, 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
								snake_case_   : Tuple						=			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
								snake_case_   : List[Any]						=			torch.device(_lowercase   )
								self.assertEqual(model.batchnorm.running_mean.device							,    _lowercase   )

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

								snake_case_   : List[Any]						=			torch.randn(2							,    3   )
								snake_case_   : int						=			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    UpperCAmelCase__							(					self   ) ->						int:







								'''simple docstring'''

								snake_case_   : Optional[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
								snake_case_   : Dict						=			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
								snake_case_   : List[Any]						=			torch.device(_lowercase   )
								self.assertEqual(model.batchnorm.running_mean.device							,    _lowercase   )

								snake_case_   : Any						=			torch.randn(2							,    3   )
								snake_case_   : int						=			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"""   )   )

								snake_case_   : Optional[int]						=			torch.randn(2							,    3   )
								snake_case_   : str						=			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"""   )   )






"""simple docstring"""





import unittest

from transformers import RoFormerTokenizer, RoFormerTokenizerFast
from transformers.testing_utils import require_rjieba, require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin



@require_rjieba
@require_tokenizers
class    _lowerCAmelCase    (				SCREAMING_SNAKE_CASE__  ,				unittest.TestCase		):





		"""simple docstring"""






		_lowerCamelCase			  =       RoFormerTokenizer
		_lowerCamelCase			  =       RoFormerTokenizerFast
		_lowerCamelCase			  =       True
		_lowerCamelCase			  =       True





		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								super().setUp()





		def    UpperCAmelCase__							(					self							,    **_lowercase   ) ->						str:







								'''simple docstring'''

								return self.tokenizer_class.from_pretrained("""junnyu/roformer_chinese_base"""							,    **_lowercase   )





		def    UpperCAmelCase__							(					self							,    **_lowercase   ) ->						Union[str, Any]:







								'''simple docstring'''

								return self.rust_tokenizer_class.from_pretrained("""junnyu/roformer_chinese_base"""							,    **_lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Dict:







								'''simple docstring'''

								snake_case_   : Tuple						=			"""永和服装饰品有限公司,今天天气非常好"""
								snake_case_   : int						=			"""永和 服装 饰品 有限公司 , 今 天 天 气 非常 好"""
								return input_text, output_text





		def    UpperCAmelCase__							(					self   ) ->						str:







								'''simple docstring'''

								snake_case_   : List[str]						=			self.get_tokenizer()
								snake_case_			,		snake_case_   : Optional[Any]						=			self.get_chinese_input_output_texts()
								snake_case_   : List[str]						=			tokenizer.tokenize(_lowercase   )

								self.assertListEqual(_lowercase							,    output_text.split()   )

								snake_case_   : str						=			tokens + [tokenizer.unk_token]
								snake_case_   : Tuple						=			[2_2_9_4_3, 2_1_3_3_2, 3_4_4_3_1, 4_5_9_0_4, 1_1_7, 3_0_6, 1_2_3_1, 1_2_3_1, 2_6_5_3, 3_3_9_9_4, 1_2_6_6, 1_0_0]
								self.assertListEqual(tokenizer.convert_tokens_to_ids(_lowercase   )							,    _lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Any:







								'''simple docstring'''

								snake_case_   : int						=			self.get_rust_tokenizer()
								snake_case_			,		snake_case_   : List[Any]						=			self.get_chinese_input_output_texts()
								snake_case_   : Union[str, Any]						=			tokenizer.tokenize(_lowercase   )
								self.assertListEqual(_lowercase							,    output_text.split()   )
								snake_case_   : Optional[int]						=			tokens + [tokenizer.unk_token]
								snake_case_   : Union[str, Any]						=			[2_2_9_4_3, 2_1_3_3_2, 3_4_4_3_1, 4_5_9_0_4, 1_1_7, 3_0_6, 1_2_3_1, 1_2_3_1, 2_6_5_3, 3_3_9_9_4, 1_2_6_6, 1_0_0]
								self.assertListEqual(tokenizer.convert_tokens_to_ids(_lowercase   )							,    _lowercase   )





		def    UpperCAmelCase__							(					self   ) ->						Tuple:







								'''simple docstring'''

								pass





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







								'''simple docstring'''

								pass





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







								'''simple docstring'''

								pass

