



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


							def __init__(				self   :						Union[str, Any]			)			->					Optional[int]:
										super().__init__()
										snake_case:    List[Any]      							=nn.Linear(3      ,						4			)
										snake_case:    List[str]      							=nn.BatchNormad(4			)
										snake_case:    List[Any]      							=nn.Linear(4      ,						5			)




							def    UpperCamelCase     (				self   :						List[Any]      ,						a_   :						Any			)			->					int:
										return self.lineara(self.batchnorm(self.lineara(a_			)			)			)






class 						a_		(     snake_case    ):


							def    UpperCamelCase     (				self   :						Optional[Any]      ,						a_   :						str      ,						*a_   :						str      ,						**a_   :						List[str]			)			->					str:
										return (args[0] + 1,) + args[1:], kwargs






class 						a_		(     snake_case    ):


							def    UpperCamelCase     (				self   :						Dict      ,						a_   :						List[Any]      ,						a_   :						str			)			->					Any:
										return output + 1






class 						a_		(     unittest.TestCase    ):


							def    UpperCamelCase     (				self   :						int			)			->					Union[str, Any]:
										snake_case:    Tuple      							=ModelForTest()
										snake_case:    Optional[int]      							=ModelHook()

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


							def    UpperCamelCase     (				self   :						Optional[int]			)			->					Any:
										snake_case:    Optional[Any]      							=ModelForTest()
										snake_case:    Dict      							=ModelHook()

										add_hook_to_module(a_      ,						a_			)
										add_hook_to_module(a_      ,						a_      ,						append=a_			)

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


							def    UpperCamelCase     (				self   :						List[Any]			)			->					Any:
										snake_case:    Optional[Any]      							=ModelForTest()
										snake_case:    int      							=torch.randn(2      ,						3			)
										snake_case:    List[str]      							=test_model(x + 1			)
										snake_case:    Optional[int]      							=test_model(x + 2			)

										snake_case:    Optional[int]      							=PreForwardHook()
										add_hook_to_module(a_      ,						a_			)
										snake_case:    Tuple      							=test_model(a_			)
										self.assertTrue(torch.allclose(a_      ,						a_      ,						atol=1E-5			)			)

										# Attaching a hook to a model when it already has one replaces, does not chain
										snake_case:    Optional[int]      							=PreForwardHook()
										add_hook_to_module(a_      ,						a_			)
										snake_case:    Optional[int]      							=test_model(a_			)
										self.assertTrue(torch.allclose(a_      ,						a_      ,						atol=1E-5			)			)

										# You need to use the sequential hook to chain two or more hooks
										snake_case:    str      							=SequentialHook(PreForwardHook()      ,						PreForwardHook()			)
										add_hook_to_module(a_      ,						a_			)

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


							def    UpperCamelCase     (				self   :						str			)			->					Tuple:
										snake_case:    Any      							=ModelForTest()
										snake_case:    Any      							=torch.randn(2      ,						3			)
										snake_case:    List[Any]      							=test_model(a_			)

										snake_case:    List[Any]      							=PostForwardHook()
										add_hook_to_module(a_      ,						a_			)
										snake_case:    Any      							=test_model(a_			)
										self.assertTrue(torch.allclose(a_      ,						output + 1      ,						atol=1E-5			)			)

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

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

										snake_case:    Union[str, Any]      							=test_model(a_			)
										assert torch.allclose(a_      ,						output + 2      ,						atol=1E-5			)


							def    UpperCamelCase     (				self   :						List[Any]			)			->					Tuple:
										snake_case:    Union[str, Any]      							=ModelForTest()
										snake_case:    List[Any]      							=torch.randn(2      ,						3			)
										snake_case:    Union[str, Any]      							=test_model(a_			)

										snake_case:    Union[str, Any]      							=PostForwardHook()
										add_hook_to_module(a_      ,						a_			)
										snake_case:    List[Any]      							=test_model(a_			)
										self.assertTrue(torch.allclose(a_      ,						output + 1			)			)
										self.assertTrue(outputa.requires_grad			)

										snake_case:    List[str]      							=True
										snake_case:    Tuple      							=test_model(a_			)
										self.assertFalse(outputa.requires_grad			)


							@require_multi_gpu
							def    UpperCamelCase     (				self   :						int			)			->					List[Any]:
										snake_case:    Optional[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
										snake_case:    Union[str, Any]      							=torch.randn(2      ,						3			)
										snake_case:    List[Any]      							=model(a_			)
										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(a_      ,						AlignDevicesHook(io_same_device=a_			)			)
										snake_case:    List[Any]      							=torch.randn(2      ,						3			).to(0			)
										snake_case:    List[str]      							=model(a_			)
										self.assertEqual(output.device      ,						torch.device(0			)			)


							def    UpperCamelCase     (				self   :						List[Any]			)			->					Tuple:
										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:    List[str]      							={'execution_device': 0 if torch.cuda.is_available() else 'cpu', 'offload': True}

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

										# 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:    Dict      							=torch.device(hook_kwargs['execution_device']			)
										self.assertEqual(model.batchnorm.running_mean.device      ,						a_			)

										snake_case:    str      							=torch.randn(2      ,						3			)
										snake_case:    Optional[Any]      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

										# 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:    Union[str, Any]      							={
										    'execution_device': 0 if torch.cuda.is_available() else 'cpu',
										    'offload': True,
										    'offload_buffers': True,
										}

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

										# 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:    Tuple      							=torch.randn(2      ,						3			)
										snake_case:    Any      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

										# 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   :						int			)			->					Optional[Any]:
										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:    Dict      							=0 if torch.cuda.is_available() else 'cpu'
										attach_align_device_hook(a_      ,						execution_device=a_      ,						offload=a_			)

										# 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(a_			)
										self.assertEqual(model.batchnorm.running_mean.device      ,						a_			)

										snake_case:    List[Any]      							=torch.randn(2      ,						3			)
										snake_case:    List[str]      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

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

										# 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:    Tuple      							=torch.randn(2      ,						3			)
										snake_case:    Tuple      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a_			)
										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   :						Optional[int]			)			->					List[str]:
										snake_case:    str      							=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[int]      							=0 if torch.cuda.is_available() else 'cpu'
										attach_align_device_hook(
										    a_      ,						execution_device=a_      ,						offload=a_      ,						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:    Optional[Any]      							=torch.device(a_			)
										self.assertEqual(model.batchnorm.running_mean.device      ,						a_			)

										snake_case:    Optional[Any]      							=torch.randn(2      ,						3			)
										snake_case:    Union[str, Any]      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

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

										# 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:    Optional[Any]      							=model(a_			)
										self.assertEqual(output.device      ,						a_			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(a_			)
										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 json
import os
import unittest

from transformers import MgpstrTokenizer
from transformers.models.mgp_str.tokenization_mgp_str import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin






@require_tokenizers
class 						a_		(     snake_case    ,   unittest.TestCase    ):
							UpperCAmelCase			:		int        =				MgpstrTokenizer
							UpperCAmelCase			:		List[Any]        =				False
							UpperCAmelCase			:		Optional[int]        =				{}
							UpperCAmelCase			:		Any        =				False


							def    UpperCamelCase     (				self   :						Any			)			->					Dict:
										super().setUp()

										# fmt: off
										snake_case:    List[Any]      							=['[GO]', '[s]', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
										# fmt: on
										snake_case:    Union[str, Any]      							=dict(zip(a_      ,						range(len(a_			)			)			)			)

										snake_case:    List[Any]      							=os.path.join(self.tmpdirname      ,						VOCAB_FILES_NAMES['vocab_file']			)
										with open(self.vocab_file      ,						'w'      ,						encoding='utf-8'			) as fp:
													fp.write(json.dumps(a_			) + '\n'			)


							def    UpperCamelCase     (				self   :						Optional[Any]      ,						**a_   :						List[Any]			)			->					str:
										return MgpstrTokenizer.from_pretrained(self.tmpdirname      ,						**a_			)


							def    UpperCamelCase     (				self   :						int      ,						a_   :						Dict			)			->					int:
										snake_case:    Any      							='tester'
										snake_case:    List[str]      							='tester'
										return input_text, output_text


							@unittest.skip('MGP-STR always lower cases letters.'			)
							def    UpperCamelCase     (				self   :						Any			)			->					Dict:
										pass


							def    UpperCamelCase     (				self   :						Optional[int]			)			->					int:
										snake_case:    int      							=self.get_tokenizers(do_lower_case=a_			)
										for tokenizer in tokenizers:
													with self.subTest(F'''{tokenizer.__class__.__name__}'''			):
																snake_case:    Optional[Any]      							='[SPECIAL_TOKEN]'

																tokenizer.add_special_tokens({'cls_token': special_token}			)
																snake_case:    List[str]      							=tokenizer.encode([special_token]      ,						add_special_tokens=a_			)
																self.assertEqual(len(a_			)      ,						1			)

																snake_case:    Union[str, Any]      							=tokenizer.decode(a_      ,						skip_special_tokens=a_			)
																self.assertTrue(special_token not in decoded			)


							def    UpperCamelCase     (				self   :						int			)			->					str:
										snake_case:    int      							=self.get_tokenizers()
										for tokenizer in tokenizers:
													with self.subTest(F'''{tokenizer.__class__.__name__}'''			):
																snake_case							,  snake_case:    Dict      							=self.get_input_output_texts(a_			)

																snake_case:    Dict      							=tokenizer.tokenize(a_			)
																snake_case:    str      							=tokenizer.convert_tokens_to_ids(a_			)
																snake_case:    List[Any]      							=tokenizer.encode(a_      ,						add_special_tokens=a_			)
																self.assertListEqual(a_      ,						a_			)

																snake_case:    str      							=tokenizer.convert_ids_to_tokens(a_			)
																self.assertNotEqual(len(a_			)      ,						0			)
																snake_case:    Tuple      							=tokenizer.decode(a_			)
																self.assertIsInstance(a_      ,						a_			)

																self.assertEqual(text_a.replace(' '      ,						''			)      ,						a_			)


							@unittest.skip('MGP-STR tokenizer only handles one sequence.'			)
							def    UpperCamelCase     (				self   :						Dict			)			->					str:
										pass




							@unittest.skip('inputs cannot be pretokenized in MgpstrTokenizer'			)
							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					Optional[int]:
										pass




