




'''simple docstring'''

import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment


def __UpperCAmelCase	(   )      ->      Optional[int]:
			__snake_case			=		torch.nn.Linear(2		,				4						)
			__snake_case			=		torch.optim.AdamW(model.parameters()		,				lr=1.0						)
			__snake_case			=		torch.optim.lr_scheduler.OneCycleLR(_UpperCAmelCase		,				max_lr=0.01		,				steps_per_epoch=2		,				epochs=1						)
			__snake_case			=		DataLoader(TensorDataset(torch.tensor([1, 2, 3]						)						)						)
			__snake_case			=		DataLoader(TensorDataset(torch.tensor([4, 5, 6]						)						)						)

			return model, optimizer, scheduler, train_dl, valid_dl


def __UpperCAmelCase	(   _UpperCAmelCase		:   List[Any]						)      ->      List[Any]:
			return (model.weight.abs().sum() + model.bias.abs().sum()).item()


def __UpperCAmelCase	(   _UpperCAmelCase		:   int						)      ->      Union[str, Any]:
			__snake_case			=		torch.nn.Linear(*tuple(model.weight.T.shape						)						).state_dict()
			model.load_state_dict(_UpperCAmelCase						)






class 		SCREAMING_SNAKE_CASE__  (			_UpperCamelCase	):





			@require_cuda
			def 							A (      self						:    List[Any]							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						assert PartialState._shared_state["_cpu"] is False
						assert PartialState._shared_state["device"].type == "cuda"
						with self.assertRaises(a_							):
									__snake_case			=		Accelerator(cpu=a_							)





			def 							A (      self						:    int							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case			=		GradientState()
						assert state.num_steps == 1
						__snake_case			=		4
						assert state.num_steps == 4

						assert state.sync_gradients is True
						__snake_case			=		False
						assert state.sync_gradients is False
						GradientState._reset_state()





			def 							A (      self						:    int							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()

						(
						    (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     (
						    __snake_case
						)		,     
						)			=		accelerator.prepare(a_      ,		a_      ,		a_      ,		a_      ,		a_							)

						self.assertTrue(prepared_model in accelerator._models							)
						self.assertTrue(prepared_optimizer in accelerator._optimizers							)
						self.assertTrue(prepared_scheduler in accelerator._schedulers							)
						self.assertTrue(prepared_train_dl in accelerator._dataloaders							)
						self.assertTrue(prepared_valid_dl in accelerator._dataloaders							)





			def 							A (      self						:    int							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()
						accelerator.prepare(a_      ,		a_      ,		a_      ,		a_      ,		a_							)
						accelerator.free_memory()

						self.assertTrue(len(accelerator._models							) == 0							)
						self.assertTrue(len(accelerator._optimizers							) == 0							)
						self.assertTrue(len(accelerator._schedulers							) == 0							)
						self.assertTrue(len(accelerator._dataloaders							) == 0							)





			def 							A (      self						:    Tuple							):



						"""simple docstring"""




						PartialState._reset_state()

						# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
						def noop(*a_						:    Tuple      ,		**a_						:    Optional[int]							):
									pass

						with patch("torch.cuda.set_device"      ,		a_							), patch_environment(ACCELERATE_TORCH_DEVICE="cuda:64"							):
									__snake_case			=		Accelerator()
									self.assertEqual(str(accelerator.state.device							)      ,		"cuda:64"							)





			def 							A (      self						:    List[str]							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()
						accelerator.prepare(a_      ,		a_      ,		a_      ,		a_      ,		a_							)

						__snake_case			=		get_signature(a_							)

						with tempfile.TemporaryDirectory() as tmpdirname:
									accelerator.save_state(a_							)

									# make sure random weights don't match
									load_random_weights(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) > 1e-3							)

									# make sure loaded weights match
									accelerator.load_state(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) < 1e-3							)





			def 							A (      self						:    Dict							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()
						accelerator.prepare(a_      ,		a_      ,		a_      ,		a_      ,		a_							)

						__snake_case			=		get_signature(a_							)

						# saving hook
						def save_config(a_						:    int      ,		a_						:    str      ,		a_						:    List[str]							):
									__snake_case			=		{"class_name": models[0].__class__.__name__}

									with open(os.path.join(a_      ,		"data.json"							)      ,		"w"							) as f:
												json.dump(a_      ,		a_							)

        # loading hook
						def load_config(a_						:    int      ,		a_						:    int							):
									with open(os.path.join(a_      ,		"data.json"							)      ,		"r"							) as f:
												__snake_case			=		json.load(a_							)

									__snake_case			=		config["class_name"]

						__snake_case			=		accelerator.register_save_state_pre_hook(a_							)
						__snake_case			=		accelerator.register_load_state_pre_hook(a_							)

						with tempfile.TemporaryDirectory() as tmpdirname:
									accelerator.save_state(a_							)

									# make sure random weights don't match with hooks
									load_random_weights(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) > 1e-3							)

									# random class name to verify correct one is loaded
									__snake_case			=		"random"

									# make sure loaded weights match with hooks
									accelerator.load_state(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) < 1e-3							)

									# mode.class_name is loaded from config
									self.assertTrue(model.class_name == model.__class__.__name__							)

						# remove hooks
						save_hook.remove()
						load_hook.remove()

						with tempfile.TemporaryDirectory() as tmpdirname:
									accelerator.save_state(a_							)

									# make sure random weights don't match with hooks removed
									load_random_weights(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) > 1e-3							)

									# random class name to verify correct one is loaded
									__snake_case			=		"random"

									# make sure loaded weights match with hooks removed
									accelerator.load_state(a_							)
									self.assertTrue(abs(model_signature - get_signature(a_							)							) < 1e-3							)

									# mode.class_name is NOT loaded from config
									self.assertTrue(model.class_name != model.__class__.__name__							)





			def 							A (      self						:    int							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()
						__snake_case			=		None

						# This should work
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		accelerator.prepare(
						    a_      ,		a_      ,		a_      ,		a_      ,		a_      ,		a_							)
						self.assertTrue(dummy_obj is None							)





			def 							A (      self						:    List[str]							):



						"""simple docstring"""




						__snake_case			=		Accelerator()
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		create_components()
						__snake_case			=		[1, 2, 3]

						# This should work
						__snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case		,     __snake_case			=		accelerator.prepare(
						    a_      ,		a_      ,		a_      ,		a_      ,		a_      ,		a_							)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Dummy object should have `_is_accelerate_prepared` set to `True`"      ,		)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Model is missing `_is_accelerator_prepared` or is set to `False`"      ,		)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Optimizer is missing `_is_accelerator_prepared` or is set to `False`"      ,		)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Scheduler is missing `_is_accelerator_prepared` or is set to `False`"      ,		)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`"      ,		)
						self.assertEqual(
						    getattr(a_      ,		"_is_accelerate_prepared"      ,		a_							)      ,		a_      ,		"Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`"      ,		)





			@slow
			@require_bnb
			def 							A (      self						:    Union[str, Any]							):



						"""simple docstring"""




						from transformers import AutoModelForCausalLM

						__snake_case			=		AutoModelForCausalLM.from_pretrained(
						    "EleutherAI/gpt-neo-125m"      ,		load_in_abit=a_      ,		device_map={"": 0}      ,		)
						__snake_case			=		Accelerator()

						# This should work
						__snake_case			=		accelerator.prepare(a_							)





			@slow
			@require_bnb
			def 							A (      self						:    List[Any]							):



						"""simple docstring"""




						from transformers import AutoModelForCausalLM

						__snake_case			=		Accelerator()

						with init_empty_weights():
									__snake_case			=		AutoModelForCausalLM.from_pretrained(
									    "EleutherAI/gpt-neo-125m"      ,		)
									model.tie_weights()
									__snake_case			=		infer_auto_device_map(a_							)
									__snake_case			=		"cpu"

						__snake_case			=		AutoModelForCausalLM.from_pretrained(
						    "EleutherAI/gpt-neo-125m"      ,		device_map=a_      ,		load_in_abit=a_      ,		llm_inta_enable_fpaa_cpu_offload=a_							)

						# This should not work and get value error
						with self.assertRaises(a_							):
									__snake_case			=		accelerator.prepare(a_							)





			@slow
			@require_bnb
			@require_multi_gpu
			def 							A (      self						:    Optional[Any]							):



						"""simple docstring"""




						from transformers import AutoModelForCausalLM

						__snake_case			=		{"distributed_type": DistributedType.MULTI_GPU}

						with init_empty_weights():
									__snake_case			=		AutoModelForCausalLM.from_pretrained(
									    "EleutherAI/gpt-neo-125m"      ,		)
									model.tie_weights()
									__snake_case			=		infer_auto_device_map(a_							)
									__snake_case			=		1

						__snake_case			=		AutoModelForCausalLM.from_pretrained(
						    "EleutherAI/gpt-neo-125m"      ,		load_in_abit=a_      ,		device_map=a_      ,		)
						__snake_case			=		Accelerator()

						# This should not work and get value error
						with self.assertRaises(a_							):
									__snake_case			=		accelerator.prepare(a_							)

						PartialState._reset_state()





			@slow
			@require_bnb
			@require_multi_gpu
			def 							A (      self						:    Tuple							):



						"""simple docstring"""




						from transformers import AutoModelForCausalLM

						with init_empty_weights():
									__snake_case			=		AutoModelForCausalLM.from_pretrained(
									    "EleutherAI/gpt-neo-125m"      ,		)
									__snake_case			=		infer_auto_device_map(a_							)
									__snake_case			=		1

						__snake_case			=		AutoModelForCausalLM.from_pretrained(
						    "EleutherAI/gpt-neo-125m"      ,		load_in_abit=a_      ,		device_map=a_      ,		)
						__snake_case			=		Accelerator()

						# This should work
						__snake_case			=		accelerator.prepare(a_							)






			@require_cuda
			def 							A (      self						:    Optional[int]							):



						"""simple docstring"""




						__snake_case			=		torch.nn.Linear(10      ,		10							)
						__snake_case			=		torch.optim.SGD(model.parameters()      ,		lr=0.01							)
						__snake_case			=		Accelerator(cpu=a_							)
						__snake_case			=		accelerator.prepare(a_							)









"""simple docstring"""




import json
import os
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from requests.exceptions import HTTPError

from transformers.utils import (
    CONFIG_NAME,
    FLAX_WEIGHTS_NAME,
    TF2_WEIGHTS_NAME,
    TRANSFORMERS_CACHE,
    WEIGHTS_NAME,
    cached_file,
    get_file_from_repo,
    has_file,
)


__magic_name__						: Tuple    =				'hf-internal-testing/tiny-random-bert'
__magic_name__						: Dict    =				os.path.join(TRANSFORMERS_CACHE, 'models--hf-internal-testing--tiny-random-bert')
__magic_name__						: Optional[Any]    =				'9b8c223d42b2188cb49d29af482996f9d0f3e5a6'




class 			__snake_case   (unittest.TestCase  ):




							def   __a   (     self:    List[Any]   ):
										__lowerCamelCase          =      cached_file(A_     ,							A_   )
										# Should have downloaded the file in here
										self.assertTrue(os.path.isdir(A_   )   )
										# Cache should contain at least those three subfolders:
										for subfolder in ["blobs", "refs", "snapshots"]:
													self.assertTrue(os.path.isdir(os.path.join(A_     ,							A_   )   )   )
										with open(os.path.join(A_     ,							"""refs"""     ,							"""main"""   )   ) as f:
													__lowerCamelCase          =      f.read()
										self.assertEqual(A_     ,							os.path.join(A_     ,							"""snapshots"""     ,							A_     ,							A_   )   )
										self.assertTrue(os.path.isfile(A_   )   )

										# File is cached at the same place the second time.
										__lowerCamelCase          =      cached_file(A_     ,							A_   )
										self.assertEqual(A_     ,							A_   )

										# Using a specific revision to test the full commit hash.
										__lowerCamelCase          =      cached_file(A_     ,							A_     ,							revision="""9b8c223"""   )
										self.assertEqual(A_     ,							os.path.join(A_     ,							"""snapshots"""     ,							A_     ,							A_   )   )




							def   __a   (     self:    int   ):
										with self.assertRaisesRegex(A_     ,							"""is not a valid model identifier"""   ):
													__lowerCamelCase          =      cached_file("""tiny-random-bert"""     ,							A_   )

										with self.assertRaisesRegex(A_     ,							"""is not a valid git identifier"""   ):
													__lowerCamelCase          =      cached_file(A_     ,							A_     ,							revision="""aaaa"""   )

										with self.assertRaisesRegex(A_     ,							"""does not appear to have a file named"""   ):
													__lowerCamelCase          =      cached_file(A_     ,							"""conf"""   )




							def   __a   (     self:    Optional[int]   ):
										with self.assertRaisesRegex(A_     ,							"""does not appear to have a file named"""   ):
													__lowerCamelCase          =      cached_file(A_     ,							"""conf"""   )

										with open(os.path.join(A_     ,							"""refs"""     ,							"""main"""   )   ) as f:
													__lowerCamelCase          =      f.read()
										self.assertTrue(os.path.isfile(os.path.join(A_     ,							""".no_exist"""     ,							A_     ,							"""conf"""   )   )   )

										__lowerCamelCase          =      cached_file(A_     ,							"""conf"""     ,							_raise_exceptions_for_missing_entries=A_   )
										self.assertIsNone(A_   )

										__lowerCamelCase          =      cached_file(A_     ,							"""conf"""     ,							local_files_only=A_     ,							_raise_exceptions_for_missing_entries=A_   )
										self.assertIsNone(A_   )

										__lowerCamelCase          =      mock.Mock()
										__lowerCamelCase          =      5_00
										__lowerCamelCase          =      {}
										__lowerCamelCase          =      HTTPError
										__lowerCamelCase          =      {}

										# Under the mock environment we get a 500 error when trying to reach the tokenizer.
										with mock.patch("""requests.Session.request"""     ,							return_value=A_   ) as mock_head:
													__lowerCamelCase          =      cached_file(A_     ,							"""conf"""     ,							_raise_exceptions_for_connection_errors=A_   )
													self.assertIsNone(A_   )
													# This check we did call the fake head request
													mock_head.assert_called()




							def   __a   (     self:    str   ):
										self.assertTrue(has_file("""hf-internal-testing/tiny-bert-pt-only"""     ,							A_   )   )
										self.assertFalse(has_file("""hf-internal-testing/tiny-bert-pt-only"""     ,							A_   )   )
										self.assertFalse(has_file("""hf-internal-testing/tiny-bert-pt-only"""     ,							A_   )   )




							def   __a   (     self:    str   ):
										# `get_file_from_repo` returns None if the file does not exist
										self.assertIsNone(get_file_from_repo("""bert-base-cased"""     ,							"""ahah.txt"""   )   )

										# The function raises if the repository does not exist.
										with self.assertRaisesRegex(A_     ,							"""is not a valid model identifier"""   ):
													get_file_from_repo("""bert-base-case"""     ,							A_   )

										# The function raises if the revision does not exist.
										with self.assertRaisesRegex(A_     ,							"""is not a valid git identifier"""   ):
													get_file_from_repo("""bert-base-cased"""     ,							A_     ,							revision="""ahaha"""   )

										__lowerCamelCase          =      get_file_from_repo("""bert-base-cased"""     ,							A_   )
										# The name is the cached name which is not very easy to test, so instead we load the content.
										__lowerCamelCase          =      json.loads(open(A_     ,							"""r"""   ).read()   )
										self.assertEqual(config["""hidden_size"""]     ,							7_68   )




							def   __a   (     self:    Union[str, Any]   ):
										with tempfile.TemporaryDirectory() as tmp_dir:
													__lowerCamelCase          =      Path(A_   ) / """a.txt"""
													filename.touch()
													self.assertEqual(get_file_from_repo(A_     ,							"""a.txt"""   )     ,							str(A_   )   )

													self.assertIsNone(get_file_from_repo(A_     ,							"""b.txt"""   )   )




