


'''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     a   (    )						->			List[str]:



					'''simple docstring'''






					UpperCamelCase__     :List[str]							      =				torch.nn.Linear(2				,      4     )
					UpperCamelCase__     :int							      =				torch.optim.AdamW(model.parameters()				,      lr=1.0     )
					UpperCamelCase__     :Union[str, Any]							      =				torch.optim.lr_scheduler.OneCycleLR(__a				,      max_lr=0.0_1				,      steps_per_epoch=2				,      epochs=1     )
					UpperCamelCase__     :Optional[Any]							      =				DataLoader(TensorDataset(torch.tensor([1, 2, 3]     )     )     )
					UpperCamelCase__     :int							      =				DataLoader(TensorDataset(torch.tensor([4, 5, 6]     )     )     )

					return model, optimizer, scheduler, train_dl, valid_dl


def     a   (    __a     )						->			Union[str, Any]:



					'''simple docstring'''






					return (model.weight.abs().sum() + model.bias.abs().sum()).item()


def     a   (    __a     )						->			Any:



					'''simple docstring'''






					UpperCamelCase__     :List[str]							      =				torch.nn.Linear(*tuple(model.weight.T.shape     )     ).state_dict()
					model.load_state_dict(__a     )







class      lowercase			(							A__      ):

		"""simple docstring"""

		@require_cuda
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :str							      =				Accelerator()
							assert PartialState._shared_state["_cpu"] is False
							assert PartialState._shared_state["device"].type == "cuda"
							with self.assertRaises(UpperCamelCase_							):
												UpperCamelCase__     :Any							      =				Accelerator(cpu=UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :str							      =				Accelerator()
							UpperCamelCase__     :int							      =				GradientState()
							assert state.num_steps == 1
							UpperCamelCase__     :Tuple							      =				4
							assert state.num_steps == 4

							assert state.sync_gradients is True
							UpperCamelCase__     :Any							      =				False
							assert state.sync_gradients is False
							GradientState._reset_state()




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Union[str, Any]							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :str							      =				create_components()

							(
							    (
							    UpperCamelCase__
							)      ,       (
							    UpperCamelCase__
							)      ,       (
							    UpperCamelCase__
							)      ,       (
							    UpperCamelCase__
							)      ,       (
							    UpperCamelCase__
							)      ,       
							)     :str							      =				accelerator.prepare(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)

							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 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Any							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :Tuple							      =				create_components()
							accelerator.prepare(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)
							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 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							PartialState._reset_state()

							# Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
							def noop(*UpperCamelCase_					,     **UpperCamelCase_							):
												pass

							with patch('''torch.cuda.set_device'''					,     UpperCamelCase_							), patch_environment(ACCELERATE_TORCH_DEVICE='''cuda:64'''							):
												UpperCamelCase__     :Optional[int]							      =				Accelerator()
												self.assertEqual(str(accelerator.state.device							)					,     '''cuda:64'''							)




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :List[Any]							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :List[str]							      =				create_components()
							accelerator.prepare(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)

							UpperCamelCase__     :Union[str, Any]							      =				get_signature(UpperCamelCase_							)

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

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

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




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :str							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :Union[str, Any]							      =				create_components()
							accelerator.prepare(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)

							UpperCamelCase__     :List[Any]							      =				get_signature(UpperCamelCase_							)

							# saving hook
							def save_config(UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							):
												UpperCamelCase__     :List[Any]							      =				{'''class_name''': models[0].__class__.__name__}

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

        # loading hook
							def load_config(UpperCamelCase_					,     UpperCamelCase_							):
												with open(os.path.join(UpperCamelCase_					,     '''data.json'''							)					,     '''r'''							) as f:
																	UpperCamelCase__     :List[str]							      =				json.load(UpperCamelCase_							)

												UpperCamelCase__     :int							      =				config['''class_name''']

							UpperCamelCase__     :List[str]							      =				accelerator.register_save_state_pre_hook(UpperCamelCase_							)
							UpperCamelCase__     :Optional[Any]							      =				accelerator.register_load_state_pre_hook(UpperCamelCase_							)

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

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

												# random class name to verify correct one is loaded
												UpperCamelCase__     :str							      =				'''random'''

												# make sure loaded weights match with hooks
												accelerator.load_state(UpperCamelCase_							)
												self.assertTrue(abs(model_signature - get_signature(UpperCamelCase_							)							) < 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(UpperCamelCase_							)

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

												# random class name to verify correct one is loaded
												UpperCamelCase__     :Optional[Any]							      =				'''random'''

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

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




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Dict							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :Dict							      =				create_components()
							UpperCamelCase__     :str							      =				None

							# This should work
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :List[str]							      =				accelerator.prepare(
							    UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_					,     UpperCamelCase_							)
							self.assertTrue(dummy_obj is None							)




		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Tuple							      =				Accelerator()
							UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__      ,       UpperCamelCase__     :Tuple							      =				create_components()
							UpperCamelCase__     :List[str]							      =				[1, 2, 3]

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




		@slow
		@require_bnb
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							from transformers import AutoModelForCausalLM

							UpperCamelCase__     :Optional[int]							      =				AutoModelForCausalLM.from_pretrained(
							    '''EleutherAI/gpt-neo-125m'''					,     load_in_abit=UpperCamelCase_					,     device_map={'''''': 0}					,     )
							UpperCamelCase__     :Tuple							      =				Accelerator()

							# This should work
							UpperCamelCase__     :Optional[int]							      =				accelerator.prepare(UpperCamelCase_							)




		@slow
		@require_bnb
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							from transformers import AutoModelForCausalLM

							UpperCamelCase__     :Optional[int]							      =				Accelerator()

							with init_empty_weights():
												UpperCamelCase__     :Tuple							      =				AutoModelForCausalLM.from_pretrained(
												    '''EleutherAI/gpt-neo-125m'''					,     )
												model.tie_weights()
												UpperCamelCase__     :List[Any]							      =				infer_auto_device_map(UpperCamelCase_							)
												UpperCamelCase__     :str							      =				'''cpu'''

							UpperCamelCase__     :Optional[int]							      =				AutoModelForCausalLM.from_pretrained(
							    '''EleutherAI/gpt-neo-125m'''					,     device_map=UpperCamelCase_					,     load_in_abit=UpperCamelCase_					,     llm_inta_enable_fpaa_cpu_offload=UpperCamelCase_							)

							# This should not work and get value error
							with self.assertRaises(UpperCamelCase_							):
												UpperCamelCase__     :Optional[int]							      =				accelerator.prepare(UpperCamelCase_							)




		@slow
		@require_bnb
		@require_multi_gpu
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							from transformers import AutoModelForCausalLM

							UpperCamelCase__     :List[Any]							      =				{'''distributed_type''': DistributedType.MULTI_GPU}

							with init_empty_weights():
												UpperCamelCase__     :List[Any]							      =				AutoModelForCausalLM.from_pretrained(
												    '''EleutherAI/gpt-neo-125m'''					,     )
												model.tie_weights()
												UpperCamelCase__     :Tuple							      =				infer_auto_device_map(UpperCamelCase_							)
												UpperCamelCase__     :Optional[Any]							      =				1

							UpperCamelCase__     :int							      =				AutoModelForCausalLM.from_pretrained(
							    '''EleutherAI/gpt-neo-125m'''					,     load_in_abit=UpperCamelCase_					,     device_map=UpperCamelCase_					,     )
							UpperCamelCase__     :int							      =				Accelerator()

							# This should not work and get value error
							with self.assertRaises(UpperCamelCase_							):
												UpperCamelCase__     :str							      =				accelerator.prepare(UpperCamelCase_							)

							PartialState._reset_state()




		@slow
		@require_bnb
		@require_multi_gpu
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							from transformers import AutoModelForCausalLM

							with init_empty_weights():
												UpperCamelCase__     :Tuple							      =				AutoModelForCausalLM.from_pretrained(
												    '''EleutherAI/gpt-neo-125m'''					,     )
												UpperCamelCase__     :str							      =				infer_auto_device_map(UpperCamelCase_							)
												UpperCamelCase__     :Optional[int]							      =				1

							UpperCamelCase__     :Union[str, Any]							      =				AutoModelForCausalLM.from_pretrained(
							    '''EleutherAI/gpt-neo-125m'''					,     load_in_abit=UpperCamelCase_					,     device_map=UpperCamelCase_					,     )
							UpperCamelCase__     :Any							      =				Accelerator()

							# This should work
							UpperCamelCase__     :Tuple							      =				accelerator.prepare(UpperCamelCase_							)




		@require_cuda
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							UpperCamelCase__     :Tuple							      =				torch.nn.Linear(10					,     10							)
							UpperCamelCase__     :List[Any]							      =				torch.optim.SGD(model.parameters()					,     lr=0.01							)
							UpperCamelCase__     :Optional[int]							      =				Accelerator(cpu=UpperCamelCase_							)
							UpperCamelCase__     :int							      =				accelerator.prepare(UpperCamelCase_							)



'''simple docstring'''
import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
class      lowercase			(							A__      ):

		"""simple docstring"""

		_a									=       'Speech2TextFeatureExtractor'
		_a									=       'Speech2TextTokenizer'




		def __init__(	self					,     UpperCamelCase_					,     UpperCamelCase_							):

							'''simple docstring'''


							super().__init__(UpperCamelCase_					,     UpperCamelCase_							)
							UpperCamelCase__     :List[str]							      =				self.feature_extractor
							UpperCamelCase__     :Optional[Any]							      =				False




		def __call__(	self					,     *UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							if self._in_target_context_manager:
												return self.current_processor(*UpperCamelCase_					,     **UpperCamelCase_							)

							if "raw_speech" in kwargs:
												warnings.warn('''Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'''							)
												UpperCamelCase__     :Optional[Any]							      =				kwargs.pop('''raw_speech'''							)
							else:
												UpperCamelCase__     :List[Any]							      =				kwargs.pop('''audio'''					,     UpperCamelCase_							)
							UpperCamelCase__     :List[Any]							      =				kwargs.pop('''sampling_rate'''					,     UpperCamelCase_							)
							UpperCamelCase__     :Dict							      =				kwargs.pop('''text'''					,     UpperCamelCase_							)
							if len(UpperCamelCase_							) > 0:
												UpperCamelCase__     :Tuple							      =				args[0]
												UpperCamelCase__     :Optional[int]							      =				args[1:]

							if audio is None and text is None:
												raise ValueError('''You need to specify either an `audio` or `text` input to process.'''							)

							if audio is not None:
												UpperCamelCase__     :Any							      =				self.feature_extractor(UpperCamelCase_					,     *UpperCamelCase_					,     sampling_rate=UpperCamelCase_					,     **UpperCamelCase_							)
							if text is not None:
												UpperCamelCase__     :Dict							      =				self.tokenizer(UpperCamelCase_					,     **UpperCamelCase_							)

							if text is None:
												return inputs
							elif audio is None:
												return encodings
							else:
												UpperCamelCase__     :Dict							      =				encodings['''input_ids''']
												return inputs




		def 					lowerCAmelCase__       (	self					,     *UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return self.tokenizer.batch_decode(*UpperCamelCase_					,     **UpperCamelCase_							)




		def 					lowerCAmelCase__       (	self					,     *UpperCamelCase_					,     **UpperCamelCase_							):

							'''simple docstring'''


							return self.tokenizer.decode(*UpperCamelCase_					,     **UpperCamelCase_							)




		@contextmanager
		def 					lowerCAmelCase__       (	self							):

							'''simple docstring'''


							warnings.warn(
							    '''`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '''
							    '''labels by using the argument `text` of the regular `__call__` method (either in the same call as '''
							    '''your audio inputs, or in a separate call.'''							)
							UpperCamelCase__     :Optional[Any]							      =				True
							UpperCamelCase__     :Any							      =				self.tokenizer
							yield
							UpperCamelCase__     :Any							      =				self.feature_extractor
							UpperCamelCase__     :Any							      =				False