'''simple docstring'''




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)

__lowercase        =    [
    '''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 							a__(   unittest.TestCase      ):


			'''simple docstring'''






			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None):




					"""simple docstring"""


					lowerCAmelCase     =    None
					lowerCAmelCase     =    os.path.abspath(os.path.join("""examples"""			,       """by_feature"""))
					lowerCAmelCase     =    os.path.abspath("""examples""")
					for item in os.listdir(__lowerCAmelCase):
							if item not in EXCLUDE_EXAMPLES:
									lowerCAmelCase     =    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     =    compare_against_test(
													    os.path.join(__lowerCAmelCase			,       __lowerCAmelCase)			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase)
													lowerCAmelCase     =    """\n""".join(__lowerCAmelCase)
													if special_strings is not None:
															for string in special_strings:
																	lowerCAmelCase     =    diff.replace(__lowerCAmelCase			,       """""")
													self.assertEqual(__lowerCAmelCase			,       """""")



			def 				a_			(  self):




					"""simple docstring"""


					self.one_complete_example("""complete_nlp_example.py"""			,       __lowerCAmelCase)
					self.one_complete_example("""complete_nlp_example.py"""			,       __lowerCAmelCase)







			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    os.path.abspath(os.path.join("""examples"""			,       """cv_example.py"""))
					lowerCAmelCase     =    [
					    """ """ * 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 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	Dict						=			False



			@classmethod
			def 				a_			(  cls):




					"""simple docstring"""


					super().setUpClass()
					lowerCAmelCase     =    tempfile.mkdtemp()
					lowerCAmelCase     =    os.path.join(cls._tmpdir			,       """default_config.yml""")

					write_basic_config(save_location=cls.configPath)
					lowerCAmelCase     =    ["""accelerate""", """launch""", """--config_file""", cls.configPath]



			@classmethod
			def 				a_			(  cls):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    f"\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps epoch\n        --output_dir {self.tmpdir}\n        ".split()
					run_command(self._launch_args + testargs)
					self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,       """epoch_0""")))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    f"\n        examples/by_feature/checkpointing.py\n        --checkpointing_steps 1\n        --output_dir {self.tmpdir}\n        ".split()
					lowerCAmelCase     =    run_command(self._launch_args + testargs)
					self.assertTrue(os.path.exists(os.path.join(self.tmpdir			,       """step_2""")))



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    f"\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir			,       'epoch_0')}\n        ".split()
					lowerCAmelCase     =    run_command(self._launch_args + testargs			,       return_stdout=__lowerCAmelCase)
					self.assertNotIn("""epoch 0:"""			,       __lowerCAmelCase)
					self.assertIn("""epoch 1:"""			,       __lowerCAmelCase)



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    f"\n        examples/by_feature/checkpointing.py\n        --resume_from_checkpoint {os.path.join(self.tmpdir			,       'step_2')}\n        ".split()
					lowerCAmelCase     =    run_command(self._launch_args + testargs			,       return_stdout=__lowerCAmelCase)
					if torch.cuda.is_available():
							lowerCAmelCase     =    torch.cuda.device_count()
					else:
							lowerCAmelCase     =    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 				a_			(  self):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ["""examples/by_feature/multi_process_metrics.py"""]
					run_command(self._launch_args + testargs)



			@require_trackers
			@mock.patch.dict(os.environ			,       {"""WANDB_MODE""": """offline"""})
			def 				a_			(  self):




					"""simple docstring"""


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



			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ["""examples/by_feature/gradient_accumulation.py"""]
					run_command(self._launch_args + testargs)







			def 				a_			(  self):




					"""simple docstring"""


					lowerCAmelCase     =    ["""examples/by_feature/local_sgd.py"""]
					run_command(self._launch_args + testargs)





'''simple docstring'''




from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
						import PIL


__lowercase        =    logging.get_logger(__name__)






def 							snake_case__				(     _A:   Dict ,  _A:   Any					)       ->      Union[str, Any]:

		'''simple docstring'''


		lowerCAmelCase     =    b.T
		lowerCAmelCase     =    np.sum(np.square(_A					) ,  axis=1					)
		lowerCAmelCase     =    np.sum(np.square(_A					) ,  axis=0					)
		lowerCAmelCase     =    np.matmul(_A ,  _A					)
		lowerCAmelCase     =    aa[:, None] - 2 * ab + ba[None, :]
		return d






def 							snake_case__				(     _A:   Dict ,  _A:   Dict					)       ->      Optional[Any]:

		'''simple docstring'''


		lowerCAmelCase     =    x.reshape(-1 ,  3					)
		lowerCAmelCase     =    squared_euclidean_distance(_A ,  _A					)
		return np.argmin(_A ,  axis=1					)






class 							a__(   lowerCAmelCase__      ):


			'''simple docstring'''






			UpperCAmelCase_     :	List[Any]						=			['''pixel_values''']



			def __init__(  self			,       __lowerCAmelCase = None			,       __lowerCAmelCase = True			,       __lowerCAmelCase = None			,       __lowerCAmelCase = PILImageResampling.BILINEAR			,       __lowerCAmelCase = True			,       __lowerCAmelCase = True			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					super().__init__(**__lowerCAmelCase)
					lowerCAmelCase     =    size if size is not None else {"""height""": 256, """width""": 256}
					lowerCAmelCase     =    get_size_dict(__lowerCAmelCase)
					lowerCAmelCase     =    np.array(__lowerCAmelCase) if clusters is not None else None
					lowerCAmelCase     =    do_resize
					lowerCAmelCase     =    size
					lowerCAmelCase     =    resample
					lowerCAmelCase     =    do_normalize
					lowerCAmelCase     =    do_color_quantize



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase			,       __lowerCAmelCase = PILImageResampling.BILINEAR			,       __lowerCAmelCase = None			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    get_size_dict(__lowerCAmelCase)
					if "height" not in size or "width" not in size:
							raise ValueError(f"Size dictionary must contain both height and width keys. Got {size.keys()}")
					return resize(
					    __lowerCAmelCase			,       size=(size["""height"""], size["""width"""])			,       resample=__lowerCAmelCase			,       data_format=__lowerCAmelCase			,       **__lowerCAmelCase)



			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = None			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    rescale(image=__lowerCAmelCase			,       scale=1 / 127.5			,       data_format=__lowerCAmelCase)
					lowerCAmelCase     =    image - 1
					return image







			def 				a_			(  self			,       __lowerCAmelCase			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = None			,       __lowerCAmelCase = ChannelDimension.FIRST			,       **__lowerCAmelCase			,       ):




					"""simple docstring"""


					lowerCAmelCase     =    do_resize if do_resize is not None else self.do_resize
					lowerCAmelCase     =    size if size is not None else self.size
					lowerCAmelCase     =    get_size_dict(__lowerCAmelCase)
					lowerCAmelCase     =    resample if resample is not None else self.resample
					lowerCAmelCase     =    do_normalize if do_normalize is not None else self.do_normalize
					lowerCAmelCase     =    do_color_quantize if do_color_quantize is not None else self.do_color_quantize
					lowerCAmelCase     =    clusters if clusters is not None else self.clusters
					lowerCAmelCase     =    np.array(__lowerCAmelCase)

					lowerCAmelCase     =    make_list_of_images(__lowerCAmelCase)

					if not valid_images(__lowerCAmelCase):
							raise ValueError(
							    """Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, """
							    """torch.Tensor, tf.Tensor or jax.ndarray.""")

					if do_resize and size is None or resample is None:
							raise ValueError("""Size and resample must be specified if do_resize is True.""")

					if do_color_quantize and clusters is None:
							raise ValueError("""Clusters must be specified if do_color_quantize is True.""")

					# All transformations expect numpy arrays.
					lowerCAmelCase     =    [to_numpy_array(__lowerCAmelCase) for image in images]

					if do_resize:
							lowerCAmelCase     =    [self.resize(image=__lowerCAmelCase			,       size=__lowerCAmelCase			,       resample=__lowerCAmelCase) for image in images]

					if do_normalize:
							lowerCAmelCase     =    [self.normalize(image=__lowerCAmelCase) for image in images]

					if do_color_quantize:
							lowerCAmelCase     =    [to_channel_dimension_format(__lowerCAmelCase			,       ChannelDimension.LAST) for image in images]
							# color quantize from (batch_size, height, width, 3) to (batch_size, height, width)
							lowerCAmelCase     =    np.array(__lowerCAmelCase)
							lowerCAmelCase     =    color_quantize(__lowerCAmelCase			,       __lowerCAmelCase).reshape(images.shape[:-1])

							# flatten to (batch_size, height*width)
							lowerCAmelCase     =    images.shape[0]
							lowerCAmelCase     =    images.reshape(__lowerCAmelCase			,       -1)

							# We need to convert back to a list of images to keep consistent behaviour across processors.
							lowerCAmelCase     =    list(__lowerCAmelCase)
					else:
							lowerCAmelCase     =    [to_channel_dimension_format(__lowerCAmelCase			,       __lowerCAmelCase) for image in images]

					lowerCAmelCase     =    {"""input_ids""": images}
					return BatchFeature(data=__lowerCAmelCase			,       tensor_type=__lowerCAmelCase)




