



import ast
import os
import re
import shutil
import tempfile
import unittest
from unittest import mock

import torch

from accelerate.test_utils.examples import compare_against_test
from accelerate.test_utils.testing import TempDirTestCase, require_trackers, run_command, slow
from accelerate.utils import write_basic_config


# DataLoaders built from `test_samples/MRPC` for quick testing
# Should mock `{script_name}.get_dataloaders` via:
# @mock.patch("{script_name}.get_dataloaders", mocked_dataloaders)

__snake_case       :	int            =[
    'cross_validation.py',
    'gradient_accumulation.py',
    'local_sgd.py',
    'multi_process_metrics.py',
    'memory.py',
    'automatic_gradient_accumulation.py',
    'fsdp_with_peak_mem_tracking.py',
    'deepspeed_with_config_support.py',
    'megatron_lm_gpt_pretraining.py',
]



class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase = None	,__lowerCamelCase = None				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  None
						lowerCAmelCase__  :       List[str]				=  os.path.abspath(os.path.join('''examples'''	,'''by_feature'''				)				)
						lowerCAmelCase__  :       Tuple				=  os.path.abspath('''examples'''				)
						for item in os.listdir(__lowerCamelCase				):
										if item not in EXCLUDE_EXAMPLES:
														lowerCAmelCase__  :       Optional[int]				=  os.path.join(__lowerCamelCase	,__lowerCamelCase				)
														if os.path.isfile(__lowerCamelCase				) and ".py" in item_path:
																		with self.subTest(
																		    tested_script=__lowerCamelCase	,feature_script=__lowerCamelCase	,tested_section='''main()''' if parser_only else '''training_function()'''	,):
																						lowerCAmelCase__  :       Optional[int]				=  compare_against_test(
																						    os.path.join(__lowerCamelCase	,__lowerCamelCase				)	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)
																						lowerCAmelCase__  :       Tuple				=  '''\n'''.join(__lowerCamelCase				)
																						if special_strings is not None:
																										for string in special_strings:
																														lowerCAmelCase__  :       int				=  diff.replace(__lowerCamelCase	,''''''				)
																						self.assertEqual(__lowerCamelCase	,''''''				)



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

						"""simple docstring"""






						self.one_complete_example('''complete_nlp_example.py'''	,__lowerCamelCase				)
						self.one_complete_example('''complete_nlp_example.py'''	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  os.path.abspath(os.path.join('''examples'''	,'''cv_example.py'''				)				)
						lowerCAmelCase__  :       Tuple				=  [
						    ''' ''' * 16 + '''{\n\n''',
						    ''' ''' * 20 + '''"accuracy": eval_metric["accuracy"],\n\n''',
						    ''' ''' * 20 + '''"f1": eval_metric["f1"],\n\n''',
						    ''' ''' * 20 + '''"train_loss": total_loss.item() / len(train_dataloader),\n\n''',
						    ''' ''' * 20 + '''"epoch": epoch,\n\n''',
						    ''' ''' * 16 + '''},\n\n''',
						    ''' ''' * 16 + '''step=epoch,\n''',
						    ''' ''' * 12,
						    ''' ''' * 8 + '''for step, batch in enumerate(active_dataloader):\n''',
						]
						self.one_complete_example('''complete_cv_example.py'''	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)
						self.one_complete_example('''complete_cv_example.py'''	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)



@mock.patch.dict(os.environ   ,       {"""TESTING_MOCKED_DATALOADERS""": """1"""})
class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			=False



		@classmethod
		def 	lowerCAmelCase__   (cls				)  ->    Union[str, Any]:

						"""simple docstring"""






						super().setUpClass()
						lowerCAmelCase__  :       List[Any]				=  tempfile.mkdtemp()
						lowerCAmelCase__  :       str				=  os.path.join(cls._tmpdir	,'''default_config.yml'''				)

						write_basic_config(save_location=cls.configPath				)
						lowerCAmelCase__  :       Optional[int]				=  ['''accelerate''', '''launch''', '''--config_file''', cls.configPath]



		@classmethod
		def 	lowerCAmelCase__   (cls				)  ->    List[str]:

						"""simple docstring"""






						super().tearDownClass()
						shutil.rmtree(cls._tmpdir				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  f"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps epoch
        --output_dir {self.tmpdir}
        """.split()
						run_command(self._launch_args + testargs				)
						self.assertTrue(os.path.exists(os.path.join(self.tmpdir	,'''epoch_0'''				)				)				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  f"""
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        """.split()
						lowerCAmelCase__  :       str				=  run_command(self._launch_args + testargs				)
						self.assertTrue(os.path.exists(os.path.join(self.tmpdir	,'''step_2'''				)				)				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  f"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir	,'epoch_0'				)}
        """.split()
						lowerCAmelCase__  :       Optional[int]				=  run_command(self._launch_args + testargs	,return_stdout=__lowerCamelCase				)
						self.assertNotIn('''epoch 0:'''	,__lowerCamelCase				)
						self.assertIn('''epoch 1:'''	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  f"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir	,'step_2'				)}
        """.split()
						lowerCAmelCase__  :       Any				=  run_command(self._launch_args + testargs	,return_stdout=__lowerCamelCase				)
						if torch.cuda.is_available():
										lowerCAmelCase__  :       Union[str, Any]				=  torch.cuda.device_count()
						else:
										lowerCAmelCase__  :       int				=  1
						if num_processes > 1:
										self.assertNotIn('''epoch 0:'''	,__lowerCamelCase				)
										self.assertIn('''epoch 1:'''	,__lowerCamelCase				)
						else:
										self.assertIn('''epoch 0:'''	,__lowerCamelCase				)
										self.assertIn('''epoch 1:'''	,__lowerCamelCase				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  '''
        examples/by_feature/cross_validation.py
        --num_folds 2
        '''.split()
						with mock.patch.dict(os.environ	,{'''TESTING_MOCKED_DATALOADERS''': '''0'''}				):
										lowerCAmelCase__  :       Union[str, Any]				=  run_command(self._launch_args + testargs	,return_stdout=__lowerCamelCase				)
										lowerCAmelCase__  :       List[Any]				=  re.findall('''({.+})'''	,__lowerCamelCase				)
										lowerCAmelCase__  :       Union[str, Any]				=  [r for r in results if '''accuracy''' in r][-1]
										lowerCAmelCase__  :       List[str]				=  ast.literal_eval(__lowerCamelCase				)
										self.assertGreaterEqual(results['''accuracy''']	,0.75				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  ['''examples/by_feature/multi_process_metrics.py''']
						run_command(self._launch_args + testargs				)



		@require_trackers
		@mock.patch.dict(os.environ	,{'''WANDB_MODE''': '''offline'''}				)
		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						with tempfile.TemporaryDirectory() as tmpdir:
										lowerCAmelCase__  :       Any				=  f"""
            examples/by_feature/tracking.py
            --with_tracking
            --project_dir {tmpdir}
            """.split()
										run_command(self._launch_args + testargs				)
										self.assertTrue(os.path.exists(os.path.join(__lowerCamelCase	,'''tracking'''				)				)				)



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

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  ['''examples/by_feature/gradient_accumulation.py''']
						run_command(self._launch_args + testargs				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  ['''examples/by_feature/local_sgd.py''']
						run_command(self._launch_args + testargs				)






import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__snake_case       :	Any            =logging.get_logger(__name__)



class  lowerCamelCase__    (							lowerCamelCase__):




		'''simple docstring'''



		snake_case_			="""encoder-decoder"""
		snake_case_			=True



		def __init__(self	,**__lowerCamelCase				)  ->    Union[str, Any]:

						"""simple docstring"""






						super().__init__(**__lowerCamelCase				)
						assert (
						    "encoder" in kwargs and "decoder" in kwargs
						), "Config has to be initialized with encoder and decoder config"
						lowerCAmelCase__  :       Optional[Any]				=  kwargs.pop('''encoder'''				)
						lowerCAmelCase__  :       Any				=  encoder_config.pop('''model_type'''				)
						lowerCAmelCase__  :       str				=  kwargs.pop('''decoder'''				)
						lowerCAmelCase__  :       Tuple				=  decoder_config.pop('''model_type'''				)

						from ..auto.configuration_auto import AutoConfig

						lowerCAmelCase__  :       Tuple				=  AutoConfig.for_model(__lowerCamelCase	,**__lowerCamelCase				)
						lowerCAmelCase__  :       Any				=  AutoConfig.for_model(__lowerCamelCase	,**__lowerCamelCase				)
						lowerCAmelCase__  :       str				=  True



		@classmethod
		def 	lowerCAmelCase__   (cls	,__lowerCamelCase	,__lowerCamelCase	,**__lowerCamelCase				)  ->    PretrainedConfig:

						"""simple docstring"""






						logger.info('''Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config'''				)
						lowerCAmelCase__  :       Optional[int]				=  True
						lowerCAmelCase__  :       List[str]				=  True

						return cls(encoder=encoder_config.to_dict()	,decoder=decoder_config.to_dict()	,**__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  copy.deepcopy(self.__dict__				)
						lowerCAmelCase__  :       Optional[Any]				=  self.encoder.to_dict()
						lowerCAmelCase__  :       str				=  self.decoder.to_dict()
						lowerCAmelCase__  :       Optional[int]				=  self.__class__.model_type
						return output

