
'''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      )					->							Any:






         """simple docstring"""






         super().__init__()
         A			:  List[str]   =		nn.Linear(3		,     4      )
         A			:  Tuple   =		nn.BatchNormad(4      )
         A			:  Optional[int]   =		nn.Linear(4		,     5      )





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE      )					->							List[str]:






         """simple docstring"""






         return self.lineara(self.batchnorm(self.lineara(SCREAMING_SNAKE_CASE      )      )      )





class 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     *SCREAMING_SNAKE_CASE		,     **SCREAMING_SNAKE_CASE      )					->							Any:






         """simple docstring"""






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





class 					A							(				__snake_case		):





     def 							__lowerCAmelCase				(				self		,     SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )					->							Union[str, Any]:






         """simple docstring"""






         return output + 1





class 					A							(				unittest.TestCase		):





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A			:  int   =		ModelForTest()
         A			:  List[Any]   =		ModelHook()

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





     def 							__lowerCAmelCase				(				self      )					->							str:






         """simple docstring"""






         A			:  Any   =		ModelForTest()
         A			:  Tuple   =		ModelHook()

         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     append=SCREAMING_SNAKE_CASE      )

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





     def 							__lowerCAmelCase				(				self      )					->							List[Any]:






         """simple docstring"""






         A			:  Tuple   =		ModelForTest()
         A			:  Tuple   =		torch.randn(2		,     3      )
         A			:  Optional[int]   =		test_model(x + 1      )
         A			:  str   =		test_model(x + 2      )

         A			:  Optional[Any]   =		PreForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
         A			:  Optional[Any]   =		test_model(SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     atol=1e-5      )      )

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

         # You need to use the sequential hook to chain two or more hooks
         A			:  int   =		SequentialHook(PreForwardHook()		,     PreForwardHook()      )
         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

         A			:  Dict   =		test_model(SCREAMING_SNAKE_CASE      )
         assert torch.allclose(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE		,     atol=1e-5      )





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






         """simple docstring"""






         A			:  List[str]   =		ModelForTest()
         A			:  List[str]   =		torch.randn(2		,     3      )
         A			:  Optional[Any]   =		test_model(SCREAMING_SNAKE_CASE      )

         A			:  List[str]   =		PostForwardHook()
         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )
         A			:  int   =		test_model(SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.allclose(SCREAMING_SNAKE_CASE		,     output + 1		,     atol=1e-5      )      )

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

         # You need to use the sequential hook to chain two or more hooks
         A			:  int   =		SequentialHook(PostForwardHook()		,     PostForwardHook()      )
         add_hook_to_module(SCREAMING_SNAKE_CASE		,     SCREAMING_SNAKE_CASE      )

         A			:  Tuple   =		test_model(SCREAMING_SNAKE_CASE      )
         assert torch.allclose(SCREAMING_SNAKE_CASE		,     output + 2		,     atol=1e-5      )





     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






         A			:  int   =		ModelForTest()
         A			:  Optional[int]   =		torch.randn(2		,     3      )
         A			:  Dict   =		test_model(SCREAMING_SNAKE_CASE      )

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

         A			:  Union[str, Any]   =		True
         A			:  Union[str, Any]   =		test_model(SCREAMING_SNAKE_CASE      )
         self.assertFalse(outputa.requires_grad      )





     @require_multi_gpu
     def 							__lowerCAmelCase				(				self      )					->							Any:






         """simple docstring"""






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





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






         """simple docstring"""






         A			:  Dict   =		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			:  Dict   =		{'''execution_device''': 0 if torch.cuda.is_available() else '''cpu''', '''offload''': True}

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

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

         A			:  List[Any]   =		torch.randn(2		,     3      )
         A			:  Dict   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

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

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

         # 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			:  str   =		torch.randn(2		,     3      )
         A			:  List[Any]   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

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






         """simple docstring"""






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

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

         A			:  int   =		torch.randn(2		,     3      )
         A			:  Optional[Any]   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

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

         # 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			:  Dict   =		torch.randn(2		,     3      )
         A			:  int   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

         # Removing hooks loads back the weights in the model.
         remove_hook_from_submodules(SCREAMING_SNAKE_CASE      )
         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 							__lowerCAmelCase				(				self      )					->							Optional[Any]:






         """simple docstring"""






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

         A			:  Optional[Any]   =		torch.randn(2		,     3      )
         A			:  Optional[Any]   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

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

         # 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			:  List[str]   =		torch.randn(2		,     3      )
         A			:  str   =		model(SCREAMING_SNAKE_CASE      )
         self.assertEqual(output.device		,     SCREAMING_SNAKE_CASE      )

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




# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
lowercase     :							Tuple       						=     '\nHuman: <<task>>\n\nAssistant: '


lowercase     :							Union[str, Any]       						=     'huggingface-tools/default-prompts'
lowercase     :							int       						=     {'chat': 'chat_prompt_template.txt', 'run': 'run_prompt_template.txt'}







def 						lowerCAmelCase_							(  snake_case__						,							snake_case__						,							snake_case__="run"		):




    '''simple docstring'''



    if prompt_or_repo_id is None:
        A			:  List[Any]   =		DEFAULT_PROMPTS_REPO

    # prompt is considered a repo ID when it does not contain any kind of space
    if re.search('''\\s'''						,							snake_case__		) is not None:
        return prompt_or_repo_id

    A			:  int   =		cached_file(
        snake_case__						,							PROMPT_FILES[mode]						,							repo_type='''dataset'''						,							user_agent={'''agent''': agent_name}		)
    with open(snake_case__						,							'''r'''						,							encoding='''utf-8'''		) as f:
        return f.read()




