
import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


__a        = namedtuple('ExampleNamedTuple', 'a b c')
class 			lowercase__(    unittest.TestCase							):



	"""simple docstring"""




	def 						_lowercase  (	self  :      List[Any]	)	->					List[Any]:
					lowercase_					  = torch.randn(5	,							2	)
					lowercase_					  = torch.device('''cuda'''	) if torch.cuda.is_available() else torch.device('''cpu'''	)

					lowercase_					  = send_to_device(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta.cpu()	,							SCREAMING_SNAKE_CASE_	)	)

					lowercase_					  = send_to_device((tensor, [tensor, tensor], 1)	,							SCREAMING_SNAKE_CASE_	)
					self.assertIsInstance(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta[0].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertIsInstance(resulta[1]	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta[1][0].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertTrue(torch.equal(resulta[1][1].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertEqual(resulta[2]	,							1	)

					lowercase_					  = send_to_device({'''a''': tensor, '''b''': [tensor, tensor], '''c''': 1}	,							SCREAMING_SNAKE_CASE_	)
					self.assertIsInstance(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta['''a'''].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertIsInstance(resulta['''b''']	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta['''b'''][0].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertTrue(torch.equal(resulta['''b'''][1].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertEqual(resulta['''c''']	,							1	)

					lowercase_					  = send_to_device(ExampleNamedTuple(a=SCREAMING_SNAKE_CASE_	,							b=[tensor, tensor]	,							c=1	)	,							SCREAMING_SNAKE_CASE_	)
					self.assertIsInstance(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta.a.cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertIsInstance(resulta.b	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta.b[0].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertTrue(torch.equal(resulta.b[1].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertEqual(resulta.c	,							1	)

					lowercase_					  = send_to_device(UserDict({'''a''': tensor, '''b''': [tensor, tensor], '''c''': 1}	)	,							SCREAMING_SNAKE_CASE_	)
					self.assertIsInstance(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta['''a'''].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertIsInstance(resulta['''b''']	,							SCREAMING_SNAKE_CASE_	)
					self.assertTrue(torch.equal(resulta['''b'''][0].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertTrue(torch.equal(resulta['''b'''][1].cpu()	,							SCREAMING_SNAKE_CASE_	)	)
					self.assertEqual(resulta['''c''']	,							1	)
	def 						_lowercase  (	self  :      List[str]	)	->					int:



					with self.assertRaises(SCREAMING_SNAKE_CASE_	) as cm:
									lowercase_					  = recursively_apply(torch.tensor	,							(torch.tensor(1	), 1)	,							error_on_other_type=SCREAMING_SNAKE_CASE_	)
					self.assertEqual(
					    str(cm.exception	)	,							'''Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'''	,							)
	def 						_lowercase  (	self  :      Dict	)	->					Optional[int]:
					lowercase_					  = torch.tensor([1, 2, 3, 4, 5]	)
					self.assertEqual(listify(SCREAMING_SNAKE_CASE_	)	,							[1, 2, 3, 4, 5]	)

					lowercase_					  = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)
					self.assertEqual(listify(SCREAMING_SNAKE_CASE_	)	,							[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]	)

					lowercase_					  = torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)
					self.assertEqual(
					    listify(SCREAMING_SNAKE_CASE_	)	,							[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]	)
	def 						_lowercase  (	self  :      Optional[int]	)	->					Optional[Any]:



					with patch_environment(aa=1	,							BB=2	):
									self.assertEqual(os.environ.get('''AA'''	)	,							'''1'''	)
									self.assertEqual(os.environ.get('''BB'''	)	,							'''2'''	)

					self.assertNotIn('''AA'''	,							os.environ	)
					self.assertNotIn('''BB'''	,							os.environ	)
	def 						_lowercase  (	self  :      int	)	->					Union[str, Any]:
					lowercase_					  = RegressionModel()
					lowercase_					  = model.forward
					lowercase_					  = convert_outputs_to_fpaa(model.forward	)
					lowercase_					  = extract_model_from_parallel(SCREAMING_SNAKE_CASE_	,							keep_fpaa_wrapper=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = pickle.dumps(SCREAMING_SNAKE_CASE_	)
	@require_cuda
	def 						_lowercase  (	self  :      Union[str, Any]	)	->					List[str]:
					lowercase_					  = RegressionModel()
					lowercase_					  = model.forward
					lowercase_					  = torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
					lowercase_					  = convert_outputs_to_fpaa(model.forward	)
					lowercase_					  = extract_model_from_parallel(SCREAMING_SNAKE_CASE_	,							keep_fpaa_wrapper=SCREAMING_SNAKE_CASE_	)
					lowercase_					  = pickle.dumps(SCREAMING_SNAKE_CASE_	)
	@require_cuda
	@require_torch_min_version(version='''2.0'''	)
	def 						_lowercase  (	self  :      str	)	->					Dict:
					lowercase_					  = RegressionModel()
					lowercase_					  = model.forward
					lowercase_					  = torch.cuda.amp.autocast(dtype=torch.floataa	)(model.forward	)
					lowercase_					  = convert_outputs_to_fpaa(model.forward	)
					lowercase_					  = torch.compile(model.forward	,							backend='''inductor'''	)
					lowercase_					  = torch.randn(4	,							1_0	).cuda()
					lowercase_					  = model(SCREAMING_SNAKE_CASE_	)
	def 						_lowercase  (	self  :      Any	)	->					List[str]:
					lowercase_					  = RegressionModel()
					# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
					lowercase_					  = torch.nn.parallel.DataParallel(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = extract_model_from_parallel(SCREAMING_SNAKE_CASE_	)

					self.assertEqual(SCREAMING_SNAKE_CASE_	,							SCREAMING_SNAKE_CASE_	)
	@require_torch_min_version(version='''2.0'''	)
	def 						_lowercase  (	self  :      Optional[int]	)	->					Tuple:
					lowercase_					  = RegressionModel()
					lowercase_					  = torch.compile(SCREAMING_SNAKE_CASE_	)

					# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
					lowercase_					  = torch.nn.parallel.DataParallel(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = torch.compile(SCREAMING_SNAKE_CASE_	)
					lowercase_					  = extract_model_from_parallel(SCREAMING_SNAKE_CASE_	)

					self.assertEqual(compiled_model._orig_mod	,							compiled_model_unwrapped._orig_mod	)




	def 						_lowercase  (	self  :      Dict	)	->					int:
					self.assertEqual(find_device([1, '''a''', torch.tensor([1, 2, 3]	)]	)	,							torch.device('''cpu'''	)	)
					self.assertEqual(find_device({'''a''': 1, '''b''': torch.tensor([1, 2, 3]	)}	)	,							torch.device('''cpu'''	)	)
					self.assertIsNone(find_device([1, '''a''']	)	)






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


__a        = {
    'configuration_speech_to_text_2': ['SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP', 'Speech2Text2Config'],
    'processing_speech_to_text_2': ['Speech2Text2Processor'],
    'tokenization_speech_to_text_2': ['Speech2Text2Tokenizer'],
}


try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						__a        = [
						    'SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST',
						    'Speech2Text2ForCausalLM',
						    'Speech2Text2PreTrainedModel',
						]


if TYPE_CHECKING:
						from .configuration_speech_to_text_a import SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextaConfig
						from .processing_speech_to_text_a import SpeechaTextaProcessor
						from .tokenization_speech_to_text_a import SpeechaTextaTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_speech_to_text_a import (
												    SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST,
												    SpeechaTextaForCausalLM,
												    SpeechaTextaPreTrainedModel,
												)

else:
						import sys

						__a        = _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




