





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	:					Optional[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  A__       (	unittest.TestCase      ):





      def 	__UpperCamelCase(			self	,							A_	,							A_	,							A_ = None	,							A_ = None       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : int	     =       os.path.abspath(os.path.join("examples"	,							"cv_example.py"       )       )
       UpperCamelCase     : Optional[int]	     =       [
           " " * 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"	,							A_	,							A_	,							A_       )
       self.one_complete_example("complete_cv_example.py"	,							A_	,							A_	,							A_       )



@mock.patch.dict(os.environ	,       {'TESTING_MOCKED_DATALOADERS': '1'}      )
class  A__       (	__snake_case      ):
      _UpperCAmelCase     :Optional[Any]						   =       False





      @classmethod
      def 	__UpperCamelCase(			cls       ):
       '''simple docstring'''




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

       write_basic_config(save_location=cls.configPath       )
       UpperCamelCase     : Optional[Any]	     =       ["accelerate", "launch", "--config_file", cls.configPath]





      @classmethod
      def 	__UpperCamelCase(			cls       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Tuple	     =       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       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Tuple	     =       F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir	,							"epoch_0"       )}
        """.split()
       UpperCamelCase     : Dict	     =       run_command(self._launch_args + testargs	,							return_stdout=A_       )
       self.assertNotIn("epoch 0:"	,							A_       )
       self.assertIn("epoch 1:"	,							A_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Any	     =       F"""
        examples/by_feature/checkpointing.py
        --resume_from_checkpoint {os.path.join(self.tmpdir	,							"step_2"       )}
        """.split()
       UpperCamelCase     : Tuple	     =       run_command(self._launch_args + testargs	,							return_stdout=A_       )
       if torch.cuda.is_available():
        UpperCamelCase     : Tuple	     =       torch.cuda.device_count()
       else:
        UpperCamelCase     : Optional[Any]	     =       1
       if num_processes > 1:
        self.assertNotIn("epoch 0:"	,							A_       )
        self.assertIn("epoch 1:"	,							A_       )
       else:
        self.assertIn("epoch 0:"	,							A_       )
        self.assertIn("epoch 1:"	,							A_       )





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : int	     =       ["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       ):
       '''simple docstring'''




       with tempfile.TemporaryDirectory() as tmpdir:
        UpperCamelCase     : 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(A_	,							"tracking"       )       )       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[int]	     =       ["examples/by_feature/local_sgd.py"]
       run_command(self._launch_args + testargs       )







__lowerCamelCase	:					Tuple			  =					"""ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"""



def 		A_    (					_lowerCAmelCase					)						->			bytes:
 # Make sure the supplied data is a bytes-like object
 if not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  UpperCamelCase     : str	     =       F"""a bytes-like object is required, not '{data.__class__.__name__}'"""
  raise TypeError(_lowerCAmelCase					)

 UpperCamelCase     : Optional[int]	     =       "".join(bin(_lowerCAmelCase					)[2:].zfill(8					) for byte in data					)

 UpperCamelCase     : str	     =       len(_lowerCAmelCase					) % 6 != 0

 if padding_needed:
  # The padding that will be added later
  UpperCamelCase     : Optional[Any]	     =       b"=" * ((6 - len(_lowerCAmelCase					) % 6) // 2)

  # Append binary_stream with arbitrary binary digits (0's by default) to make its
  # length a multiple of 6.
  binary_stream += "0" * (6 - len(_lowerCAmelCase					) % 6)
 else:
  UpperCamelCase     : List[Any]	     =       b""

 # Encode every 6 binary digits to their corresponding Base64 character
 return (
     "".join(
         B64_CHARSET[int(binary_stream[index : index + 6]							,       2					)]
         for index in range(0							,       len(_lowerCAmelCase					)							,       6					)					).encode()
     + padding
 )



def 		A_    (					_lowerCAmelCase					)						->			bytes:
 # Make sure encoded_data is either a string or a bytes-like object
 if not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					) and not isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  UpperCamelCase     : List[Any]	     =       (
      "argument should be a bytes-like object or ASCII string, "
      F"""not '{encoded_data.__class__.__name__}'"""
  )
  raise TypeError(_lowerCAmelCase					)

 # In case encoded_data is a bytes-like object, make sure it contains only
 # ASCII characters so we convert it to a string object
 if isinstance(_lowerCAmelCase							,       _lowerCAmelCase					):
  try:
   UpperCamelCase     : Any	     =       encoded_data.decode("utf-8"					)
  except UnicodeDecodeError:
   raise ValueError("base64 encoded data should only contain ASCII characters"					)

 UpperCamelCase     : Union[str, Any]	     =       encoded_data.count("="					)

 # Check if the encoded string contains non base64 characters
 if padding:
  assert all(
      char in B64_CHARSET for char in encoded_data[:-padding]					), "Invalid base64 character(s) found."
 else:
  assert all(
      char in B64_CHARSET for char in encoded_data					), "Invalid base64 character(s) found."

 # Check the padding
 assert len(_lowerCAmelCase					) % 4 == 0 and padding < 3, "Incorrect padding"

 if padding:
  # Remove padding if there is one
  UpperCamelCase     : List[str]	     =       encoded_data[:-padding]

  UpperCamelCase     : List[Any]	     =       "".join(
      bin(B64_CHARSET.index(_lowerCAmelCase					)					)[2:].zfill(6					) for char in encoded_data					)[: -padding * 2]
 else:
  UpperCamelCase     : List[Any]	     =       "".join(
      bin(B64_CHARSET.index(_lowerCAmelCase					)					)[2:].zfill(6					) for char in encoded_data					)

 UpperCamelCase     : Any	     =       [
     int(binary_stream[index : index + 8]							,       2					)
     for index in range(0							,       len(_lowerCAmelCase					)							,       8					)
 ]

 return bytes(_lowerCAmelCase					)


if __name__ == "__main__":
    import doctest

    doctest.testmod()
