




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)

__lowerCamelCase   :						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 							UpperCAmelCase__       ( self			:   Any       ,						lowerCamelCase_			:   str       ,						lowerCamelCase_			:   bool       ,						lowerCamelCase_			:   str = None       ,						lowerCamelCase_			:   list = None  )						->   List[str]:
								__magic_name__							:	Tuple    							=							None
								__magic_name__							:	Dict    							=							os.path.abspath(os.path.join('''examples'''       ,						'''by_feature'''  )  )
								__magic_name__							:	int    							=							os.path.abspath('''examples'''  )
								for item in os.listdir(lowerCamelCase_  ):
													if item not in EXCLUDE_EXAMPLES:
																		__magic_name__							:	Union[str, Any]    							=							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()'''       ,						):
																												__magic_name__							:	List[str]    							=							compare_against_test(
																												    os.path.join(lowerCamelCase_       ,						lowerCamelCase_  )       ,						lowerCamelCase_       ,						lowerCamelCase_       ,						lowerCamelCase_  )
																												__magic_name__							:	Dict    							=							'''\n'''.join(lowerCamelCase_  )
																												if special_strings is not None:
																																	for string in special_strings:
																																						__magic_name__							:	int    							=							diff.replace(lowerCamelCase_       ,						''''''  )
																												self.assertEqual(lowerCamelCase_       ,						''''''  )

			def 							UpperCAmelCase__       ( self			:   Any  )						->   str:
								self.one_complete_example('''complete_nlp_example.py'''       ,						lowerCamelCase_  )
								self.one_complete_example('''complete_nlp_example.py'''       ,						lowerCamelCase_  )


			def 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   str:
								__magic_name__							:	Any    							=							os.path.abspath(os.path.join('''examples'''       ,						'''cv_example.py'''  )  )
								__magic_name__							:	Dict    							=							[
								    ''' ''' * 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'''





			UpperCamelCase__								=False

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   Optional[int]  )						->   Dict:
								super().setUpClass()
								__magic_name__							:	Optional[int]    							=							tempfile.mkdtemp()
								__magic_name__							:	int    							=							os.path.join(cls._tmpdir       ,						'''default_config.yml'''  )

								write_basic_config(save_location=cls.configPath  )
								__magic_name__							:	List[Any]    							=							['''accelerate''', '''launch''', '''--config_file''', cls.configPath]

			@classmethod
			def 							UpperCAmelCase__       ( cls			:   Optional[Any]  )						->   Any:
								super().tearDownClass()
								shutil.rmtree(cls._tmpdir  )

			def 							UpperCAmelCase__       ( self			:   List[Any]  )						->   Any:
								__magic_name__							:	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 							UpperCAmelCase__       ( self			:   int  )						->   Optional[Any]:
								__magic_name__							:	Any    							=							F'''
        examples/by_feature/checkpointing.py
        --checkpointing_steps 1
        --output_dir {self.tmpdir}
        '''.split()
								__magic_name__							:	Optional[Any]    							=							run_command(self._launch_args + testargs  )
								self.assertTrue(os.path.exists(os.path.join(self.tmpdir       ,						'''step_2'''  )  )  )

			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   Any:
								__magic_name__							:	List[str]    							=							F'''
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir       ,						'epoch_0'  )}
        '''.split()
								__magic_name__							:	Union[str, Any]    							=							run_command(self._launch_args + testargs       ,						return_stdout=lowerCamelCase_  )
								self.assertNotIn('''epoch 0:'''       ,						lowerCamelCase_  )
								self.assertIn('''epoch 1:'''       ,						lowerCamelCase_  )

			def 							UpperCAmelCase__       ( self			:   List[str]  )						->   Optional[int]:
								__magic_name__							:	List[Any]    							=							F'''
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir       ,						'step_2'  )}
        '''.split()
								__magic_name__							:	Tuple    							=							run_command(self._launch_args + testargs       ,						return_stdout=lowerCamelCase_  )
								if torch.cuda.is_available():
													__magic_name__							:	Any    							=							torch.cuda.device_count()
								else:
													__magic_name__							:	Optional[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 							UpperCAmelCase__       ( self			:   Optional[int]  )						->   int:
								__magic_name__							:	List[str]    							=							'''
        examples/by_feature/cross_validation.py
        --num_folds 2
        '''.split()
								with mock.patch.dict(os.environ       ,						{'''TESTING_MOCKED_DATALOADERS''': '''0'''}  ):
													__magic_name__							:	int    							=							run_command(self._launch_args + testargs       ,						return_stdout=lowerCamelCase_  )
													__magic_name__							:	Any    							=							re.findall('''({.+})'''       ,						lowerCamelCase_  )
													__magic_name__							:	Union[str, Any]    							=							[r for r in results if '''accuracy''' in r][-1]
													__magic_name__							:	Optional[int]    							=							ast.literal_eval(lowerCamelCase_  )
													self.assertGreaterEqual(results['''accuracy''']       ,						0.7_5  )

			def 							UpperCAmelCase__       ( self			:   Union[str, Any]  )						->   Optional[Any]:
								__magic_name__							:	Dict    							=							['''examples/by_feature/multi_process_metrics.py''']
								run_command(self._launch_args + testargs  )

			@require_trackers
			@mock.patch.dict(os.environ       ,						{'''WANDB_MODE''': '''offline'''}  )
			def 							UpperCAmelCase__       ( self			:   Optional[Any]  )						->   List[Any]:
								with tempfile.TemporaryDirectory() as tmpdir:
													__magic_name__							:	int    							=							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 							UpperCAmelCase__       ( self			:   List[str]  )						->   Any:
								__magic_name__							:	List[str]    							=							['''examples/by_feature/gradient_accumulation.py''']
								run_command(self._launch_args + testargs  )


			def 							UpperCAmelCase__       ( self			:   Any  )						->   List[str]:
								__magic_name__							:	Tuple    							=							['''examples/by_feature/local_sgd.py''']
								run_command(self._launch_args + testargs  )









from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase   :						Optional[int]				=			logging.get_logger(__name__)

__lowerCamelCase   :						Tuple				=			{
    '''shi-labs/nat-mini-in1k-224''': '''https://huggingface.co/shi-labs/nat-mini-in1k-224/resolve/main/config.json''',
    # See all Nat models at https://huggingface.co/models?filter=nat
}


class    lowerCamelCase    (		_lowerCamelCase				,_lowerCamelCase						):






			'''simple docstring'''





			UpperCamelCase__								='''nat'''

			UpperCamelCase__								={
			    '''num_attention_heads''': '''num_heads''',
			    '''num_hidden_layers''': '''num_layers''',
			}

			def __init__( self			:   Union[str, Any]       ,						lowerCamelCase_			:   Dict=4       ,						lowerCamelCase_			:   Union[str, Any]=3       ,						lowerCamelCase_			:   List[str]=64       ,						lowerCamelCase_			:   Union[str, Any]=[3, 4, 6, 5]       ,						lowerCamelCase_			:   List[Any]=[2, 4, 8, 16]       ,						lowerCamelCase_			:   List[Any]=7       ,						lowerCamelCase_			:   Union[str, Any]=3.0       ,						lowerCamelCase_			:   int=True       ,						lowerCamelCase_			:   str=0.0       ,						lowerCamelCase_			:   int=0.0       ,						lowerCamelCase_			:   Dict=0.1       ,						lowerCamelCase_			:   Optional[int]="gelu"       ,						lowerCamelCase_			:   Tuple=0.0_2       ,						lowerCamelCase_			:   Tuple=1E-5       ,						lowerCamelCase_			:   List[Any]=0.0       ,						lowerCamelCase_			:   str=None       ,						lowerCamelCase_			:   Tuple=None       ,						**lowerCamelCase_			:   Any       ,						)						->   Any:
								super().__init__(**lowerCamelCase_  )

								__magic_name__							:	List[Any]    							=							patch_size
								__magic_name__							:	str    							=							num_channels
								__magic_name__							:	Union[str, Any]    							=							embed_dim
								__magic_name__							:	Dict    							=							depths
								__magic_name__							:	str    							=							len(lowerCamelCase_  )
								__magic_name__							:	List[str]    							=							num_heads
								__magic_name__							:	Optional[int]    							=							kernel_size
								__magic_name__							:	Any    							=							mlp_ratio
								__magic_name__							:	Any    							=							qkv_bias
								__magic_name__							:	Union[str, Any]    							=							hidden_dropout_prob
								__magic_name__							:	Any    							=							attention_probs_dropout_prob
								__magic_name__							:	List[str]    							=							drop_path_rate
								__magic_name__							:	str    							=							hidden_act
								__magic_name__							:	Tuple    							=							layer_norm_eps
								__magic_name__							:	Any    							=							initializer_range
								# we set the hidden_size attribute in order to make Nat work with VisionEncoderDecoderModel
								# this indicates the channel dimension after the last stage of the model
								__magic_name__							:	str    							=							int(embed_dim * 2 ** (len(lowerCamelCase_  ) - 1)  )
								__magic_name__							:	str    							=							layer_scale_init_value
								__magic_name__							:	Dict    							=							['''stem'''] + [F'''stage{idx}''' for idx in range(1       ,						len(lowerCamelCase_  ) + 1  )]
								__magic_name__							,    __magic_name__							:	int    							=							get_aligned_output_features_output_indices(
								    out_features=lowerCamelCase_       ,						out_indices=lowerCamelCase_       ,						stage_names=self.stage_names  )



