






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

	def __init__(							self				)->    str:







							'''simple docstring'''





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


	def 						UpperCAmelCase     (							self						,			__UpperCAmelCase				)->    Any:







							'''simple docstring'''





							return self.lineara(self.batchnorm(self.lineara(__UpperCAmelCase				)				)				)





class    lowercase__   (	lowercase_		):

	def 						UpperCAmelCase     (							self						,			__UpperCAmelCase						,			*__UpperCAmelCase						,			**__UpperCAmelCase				)->    int:







							'''simple docstring'''





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





class    lowercase__   (	lowercase_		):

	def 						UpperCAmelCase     (							self						,			__UpperCAmelCase						,			__UpperCAmelCase				)->    Dict:







							'''simple docstring'''





							return output + 1





class    lowercase__   (	unittest.TestCase		):

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







							'''simple docstring'''





							lowerCAmelCase__       					=  ModelForTest()
							lowerCAmelCase__       					=  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 						UpperCAmelCase     (							self				)->    Any:







							'''simple docstring'''





							lowerCAmelCase__       					=  ModelForTest()
							lowerCAmelCase__       					=  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 						UpperCAmelCase     (							self				)->    str:







							'''simple docstring'''





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

							lowerCAmelCase__       					=  PreForwardHook()
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)
							lowerCAmelCase__       					=  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__       					=  PreForwardHook()
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)
							lowerCAmelCase__       					=  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__       					=  SequentialHook(PreForwardHook()						,			PreForwardHook()				)
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)

							lowerCAmelCase__       					=  test_model(__UpperCAmelCase				)
							assert torch.allclose(__UpperCAmelCase						,			__UpperCAmelCase						,			atol=1E-5				)

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







							'''simple docstring'''





							lowerCAmelCase__       					=  ModelForTest()
							lowerCAmelCase__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  test_model(__UpperCAmelCase				)

							lowerCAmelCase__       					=  PostForwardHook()
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)
							lowerCAmelCase__       					=  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__       					=  PostForwardHook()
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)
							lowerCAmelCase__       					=  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__       					=  SequentialHook(PostForwardHook()						,			PostForwardHook()				)
							add_hook_to_module(__UpperCAmelCase						,			__UpperCAmelCase				)

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

	def 						UpperCAmelCase     (							self				)->    str:







							'''simple docstring'''





							lowerCAmelCase__       					=  ModelForTest()
							lowerCAmelCase__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  test_model(__UpperCAmelCase				)

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

							lowerCAmelCase__       					=  True
							lowerCAmelCase__       					=  test_model(__UpperCAmelCase				)
							self.assertFalse(outputa.requires_grad				)

	@require_multi_gpu
	def 						UpperCAmelCase     (							self				)->    int:







							'''simple docstring'''





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

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







							'''simple docstring'''





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

							lowerCAmelCase__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  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__       					=  {
							    "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__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  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 						UpperCAmelCase     (							self				)->    Optional[int]:







							'''simple docstring'''





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

							lowerCAmelCase__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  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__       					=  torch.randn(2						,			3				)
							lowerCAmelCase__       					=  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 						UpperCAmelCase     (							self				)->    Optional[int]:







							'''simple docstring'''





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

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









import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict





@pytest.mark.parametrize(
    'files'   ,							[
        ['full:README.md', 'dataset_infos.json'],
        ['empty:README.md', 'dataset_infos.json'],
        ['dataset_infos.json'],
        ['full:README.md'],
    ]   ,							)
def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    tmp_path_factory.mktemp('dset_infos_dir'      )
							if "full:README.md" in files:
														with open(dataset_infos_dir / 'README.md'   ,							'w'      ) as f:
																					f.write('---\ndataset_info:\n  dataset_size: 42\n---'      )
							if "empty:README.md" in files:
														with open(dataset_infos_dir / 'README.md'   ,							'w'      ) as f:
																					f.write(''      )
    # we want to support dataset_infos.json for backward compatibility
							if "dataset_infos.json" in files:
														with open(dataset_infos_dir / 'dataset_infos.json'   ,							'w'      ) as f:
																					f.write('{"default": {"dataset_size": 42}}'      )
							snake_case_							=    DatasetInfosDict.from_directory(UpperCamelCase__      )
							assert dataset_infos
							assert dataset_infos["default"].dataset_size == 42





@pytest.mark.parametrize(
    'dataset_info'   ,							[
        DatasetInfo(),
        DatasetInfo(
            description='foo'   ,							features=Features({'a': Value('int32'      )}      )   ,							builder_name='builder'   ,							config_name='config'   ,							version='1.0.0'   ,							splits=[{'name': 'train'}]   ,							download_size=42   ,							),
    ]   ,							)
def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    str(UpperCamelCase__      )
							dataset_info.write_to_directory(UpperCamelCase__      )
							snake_case_							=    DatasetInfo.from_directory(UpperCamelCase__      )
							assert dataset_info == reloaded
							assert os.path.exists(os.path.join(UpperCamelCase__   ,							'dataset_info.json'      )      )





def 					__lowerCamelCase	( ):







							'''simple docstring'''






							snake_case_							=    DatasetInfo(
							    description='foo'   ,							citation='bar'   ,							homepage='https://foo.bar'   ,							license='CC0'   ,							features=Features({'a': Value('int32'      )}      )   ,							post_processed={}   ,							supervised_keys=()   ,							task_templates=[]   ,							builder_name='builder'   ,							config_name='config'   ,							version='1.0.0'   ,							splits=[{'name': 'train', 'num_examples': 42}]   ,							download_checksums={}   ,							download_size=1337   ,							post_processing_size=442   ,							dataset_size=1234   ,							size_in_bytes=1337 + 442 + 1234   ,							)
							snake_case_							=    dataset_info._to_yaml_dict()
							assert sorted(UpperCamelCase__      ) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML      )
							for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
														assert key in dataset_info_yaml_dict
														assert isinstance(dataset_info_yaml_dict[key]   ,							(list, dict, int, str)      )
							snake_case_							=    yaml.safe_dump(UpperCamelCase__      )
							snake_case_							=    yaml.safe_load(UpperCamelCase__      )
							assert dataset_info_yaml_dict == reloaded





def 					__lowerCamelCase	( ):







							'''simple docstring'''






							snake_case_							=    DatasetInfo()
							snake_case_							=    dataset_info._to_yaml_dict()
							assert dataset_info_yaml_dict == {}





@pytest.mark.parametrize(
    'dataset_infos_dict'   ,							[
        DatasetInfosDict(),
        DatasetInfosDict({'default': DatasetInfo()}      ),
        DatasetInfosDict({'my_config_name': DatasetInfo()}      ),
        DatasetInfosDict(
            {
                'default': DatasetInfo(
                    description='foo'   ,							features=Features({'a': Value('int32'      )}      )   ,							builder_name='builder'   ,							config_name='config'   ,							version='1.0.0'   ,							splits=[{'name': 'train'}]   ,							download_size=42   ,							)
            }      ),
        DatasetInfosDict(
            {
                'v1': DatasetInfo(dataset_size=42      ),
                'v2': DatasetInfo(dataset_size=1337      ),
            }      ),
    ]   ,							)
def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    str(UpperCamelCase__      )
							dataset_infos_dict.write_to_directory(UpperCamelCase__      )
							snake_case_							=    DatasetInfosDict.from_directory(UpperCamelCase__      )

							# the config_name of the dataset_infos_dict take over the attribute
							for config_name, dataset_info in dataset_infos_dict.items():
														snake_case_							=    config_name
														# the yaml representation doesn't include fields like description or citation
														# so we just test that we can recover what we can from the yaml
														snake_case_							=    DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()      )
							assert dataset_infos_dict == reloaded

							if dataset_infos_dict:
														assert os.path.exists(os.path.join(UpperCamelCase__   ,							'README.md'      )      )




