






import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / "utils"))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


UpperCamelCase        =				get_tests_dir("fixtures")
UpperCamelCase        =				get_tests_dir("fixtures/dummy_feature_extractor_config.json")
UpperCamelCase        =				get_tests_dir("fixtures/dummy-config.json")

class        lowerCAmelCase_				(      unittest.TestCase	):







				"""simple docstring"""




				def 				__a   ( self     :Optional[int]     ):
					UpperCamelCase__ :Optional[int]						=      0







				def 				__a   ( self     :str     ):
					UpperCamelCase__ :Optional[Any]						=      AutoFeatureExtractor.from_pretrained("""facebook/wav2vec2-base-960h"""     )
					self.assertIsInstance(lowerCamelCase__  ,     lowerCamelCase__     )







				def 				__a   ( self     :Dict     ):
					UpperCamelCase__ :Dict						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__     )
					self.assertIsInstance(lowerCamelCase__  ,     lowerCamelCase__     )







				def 				__a   ( self     :Optional[int]     ):
					with tempfile.TemporaryDirectory() as tmpdirname:
						UpperCamelCase__ :List[str]						=      WavaVecaConfig()

						# remove feature_extractor_type to make sure config.json alone is enough to load feature processor locally
						UpperCamelCase__ :Tuple						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__     ).to_dict()

						config_dict.pop("""feature_extractor_type"""     )
						UpperCamelCase__ :Union[str, Any]						=      WavaVecaFeatureExtractor(**lowerCamelCase__     )

						# save in new folder
						model_config.save_pretrained(lowerCamelCase__     )
						config.save_pretrained(lowerCamelCase__     )

						UpperCamelCase__ :Optional[Any]						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__     )

						# make sure private variable is not incorrectly saved
						UpperCamelCase__ :Tuple						=      json.loads(config.to_json_string()     )
						self.assertTrue("""_processor_class""" not in dict_as_saved     )

					self.assertIsInstance(lowerCamelCase__  ,     lowerCamelCase__     )







				def 				__a   ( self     :Union[str, Any]     ):
					UpperCamelCase__ :Dict						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__     )
					self.assertIsInstance(lowerCamelCase__  ,     lowerCamelCase__     )







				def 				__a   ( self     :Dict     ):
					with self.assertRaisesRegex(
					    lowerCamelCase__  ,     """bert-base is not a local folder and is not a valid model identifier"""     ):
						UpperCamelCase__ :Dict						=      AutoFeatureExtractor.from_pretrained("""bert-base"""     )







				def 				__a   ( self     :List[Any]     ):
					with self.assertRaisesRegex(
					    lowerCamelCase__  ,     r"""aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"""     ):
						UpperCamelCase__ :Optional[int]						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__  ,     revision="""aaaaaa"""     )







				def 				__a   ( self     :int     ):
					with self.assertRaisesRegex(
					    lowerCamelCase__  ,     """hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json."""  ,     ):
						UpperCamelCase__ :List[Any]						=      AutoFeatureExtractor.from_pretrained("""hf-internal-testing/config-no-model"""     )







				def 				__a   ( self     :Optional[int]     ):
					# If remote code is not set, we will time out when asking whether to load the model.
					with self.assertRaises(lowerCamelCase__     ):
						UpperCamelCase__ :List[Any]						=      AutoFeatureExtractor.from_pretrained(
						    """hf-internal-testing/test_dynamic_feature_extractor"""     )
					# If remote code is disabled, we can't load this config.
					with self.assertRaises(lowerCamelCase__     ):
						UpperCamelCase__ :Optional[Any]						=      AutoFeatureExtractor.from_pretrained(
						    """hf-internal-testing/test_dynamic_feature_extractor"""  ,     trust_remote_code=lowerCamelCase__     )

					UpperCamelCase__ :str						=      AutoFeatureExtractor.from_pretrained(
					    """hf-internal-testing/test_dynamic_feature_extractor"""  ,     trust_remote_code=lowerCamelCase__     )
					self.assertEqual(feature_extractor.__class__.__name__  ,     """NewFeatureExtractor"""     )

					# Test feature extractor can be reloaded.
					with tempfile.TemporaryDirectory() as tmp_dir:
						feature_extractor.save_pretrained(lowerCamelCase__     )
						UpperCamelCase__ :Any						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__  ,     trust_remote_code=lowerCamelCase__     )
					self.assertEqual(reloaded_feature_extractor.__class__.__name__  ,     """NewFeatureExtractor"""     )







				def 				__a   ( self     :Dict     ):
					try:
						AutoConfig.register("""custom"""  ,     lowerCamelCase__     )
						AutoFeatureExtractor.register(lowerCamelCase__  ,     lowerCamelCase__     )
						# Trying to register something existing in the Transformers library will raise an error
						with self.assertRaises(lowerCamelCase__     ):
							AutoFeatureExtractor.register(lowerCamelCase__  ,     lowerCamelCase__     )

						# Now that the config is registered, it can be used as any other config with the auto-API
						UpperCamelCase__ :Any						=      CustomFeatureExtractor.from_pretrained(lowerCamelCase__     )
						with tempfile.TemporaryDirectory() as tmp_dir:
							feature_extractor.save_pretrained(lowerCamelCase__     )
							UpperCamelCase__ :List[Any]						=      AutoFeatureExtractor.from_pretrained(lowerCamelCase__     )
							self.assertIsInstance(lowerCamelCase__  ,     lowerCamelCase__     )

					finally:
						if "custom" in CONFIG_MAPPING._extra_content:
							del CONFIG_MAPPING._extra_content["custom"]
						if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
							del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]







				def 				__a   ( self     :Optional[int]     ):

					class        lowerCAmelCase_				(      lowercase	):







									"""simple docstring"""




									_snake_case				:   Optional[int]        =			True

					try:
						AutoConfig.register("""custom"""  ,     lowerCamelCase__     )
						AutoFeatureExtractor.register(lowerCamelCase__  ,     lowerCamelCase__     )
						# If remote code is not set, the default is to use local
						UpperCamelCase__ :Optional[Any]						=      AutoFeatureExtractor.from_pretrained(
						    """hf-internal-testing/test_dynamic_feature_extractor"""     )
						self.assertEqual(feature_extractor.__class__.__name__  ,     """NewFeatureExtractor"""     )
						self.assertTrue(feature_extractor.is_local     )

						# If remote code is disabled, we load the local one.
						UpperCamelCase__ :str						=      AutoFeatureExtractor.from_pretrained(
						    """hf-internal-testing/test_dynamic_feature_extractor"""  ,     trust_remote_code=lowerCamelCase__     )
						self.assertEqual(feature_extractor.__class__.__name__  ,     """NewFeatureExtractor"""     )
						self.assertTrue(feature_extractor.is_local     )

						# If remote is enabled, we load from the Hub
						UpperCamelCase__ :Optional[int]						=      AutoFeatureExtractor.from_pretrained(
						    """hf-internal-testing/test_dynamic_feature_extractor"""  ,     trust_remote_code=lowerCamelCase__     )
						self.assertEqual(feature_extractor.__class__.__name__  ,     """NewFeatureExtractor"""     )
						self.assertTrue(not hasattr(lowerCamelCase__  ,     """is_local"""     )     )

					finally:
						if "custom" in CONFIG_MAPPING._extra_content:
							del CONFIG_MAPPING._extra_content["custom"]
						if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
							del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]




'''simple docstring'''






import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler

class 				a__:







    def __init__( self					:    List[Any]		,     __snake_case					:    str      ):
      if isinstance(__snake_case		,     __snake_case      ):
        # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
        # modified it, it will not be accepted here again, since `auto` values would have been overridden
        a		:   Union[str, Any]	  =	deepcopy(__snake_case      )
      elif os.path.exists(__snake_case      ):
        with io.open(__snake_case		,     'r'		,     encoding='utf-8'      ) as f:
          a		:   List[str]	  =	json.load(__snake_case      )
      else:
        try:
          a		:   str	  =	baseaa.urlsafe_baadecode(__snake_case      ).decode('utf-8'      )
          a		:   Union[str, Any]	  =	json.loads(__snake_case      )
        except (UnicodeDecodeError, AttributeError, ValueError):
          raise ValueError(
              F"""Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}"""      )

      a		:   Optional[int]	  =	config

      self.set_stage_and_offload()







    def      lowercase_			( self					:    Any      ):
      # zero stage - this is done as early as possible, before model is created, to allow
      # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
      # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
      a		:   List[str]	  =	self.get_value('zero_optimization.stage'		,     -1      )

      # offload
      a		:   Union[str, Any]	  =	False
      if self.is_zeroa() or self.is_zeroa():
        a		:   Any	  =	set(['cpu', 'nvme']      )
        a		:   Tuple	  =	set(
            [
                self.get_value('zero_optimization.offload_optimizer.device'      ),
                self.get_value('zero_optimization.offload_param.device'      ),
            ]      )
        if len(offload_devices & offload_devices_valid      ) > 0:
          a		:   Optional[int]	  =	True







    def      lowercase_			( self					:    str		,     __snake_case					:    Optional[Any]      ):
      a		:   Any	  =	self.config

      # find the config node of interest if it exists
      a		:   Dict	  =	ds_key_long.split('.'      )
      a		:   str	  =	nodes.pop()
      for node in nodes:
        a		:   Tuple	  =	config.get(__snake_case      )
        if config is None:
          return None, ds_key

      return config, ds_key







    def      lowercase_			( self					:    int		,     __snake_case					:    Optional[int]		,     __snake_case					:    List[str]=None      ):
      a  ,      a		:   int	  =	self.find_config_node(__snake_case      )
      if config is None:
        return default
      return config.get(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    Dict		,     __snake_case					:    List[str]		,     __snake_case					:    List[Any]=False      ):
      a		:   Union[str, Any]	  =	self.config

      # find the config node of interest if it exists
      a		:   Tuple	  =	ds_key_long.split('.'      )
      for node in nodes:
        a		:   Optional[int]	  =	config
        a		:   List[str]	  =	config.get(__snake_case      )
        if config is None:
          if must_exist:
            raise ValueError(F"""Can't find {ds_key_long} entry in the config: {self.config}"""      )
          else:
            return

        # if found remove it
      if parent_config is not None:
        parent_config.pop(__snake_case      )







    def      lowercase_			( self					:    List[str]		,     __snake_case					:    Dict      ):
      a		:   Optional[Any]	  =	self.get_value(__snake_case      )
      return False if value is None else bool(__snake_case      )







    def      lowercase_			( self					:    str		,     __snake_case					:    Optional[Any]      ):
      a		:   Optional[Any]	  =	self.get_value(__snake_case      )
      return False if value is None else not bool(__snake_case      )







    def      lowercase_			( self					:    Tuple      ):
      return self._stage == 2







    def      lowercase_			( self					:    str      ):
      return self._stage == 3







    def      lowercase_			( self					:    int      ):
      return self._offload

class 				a__:







    def __init__( self					:    str		,     __snake_case					:    Any      ):
      a		:   Any	  =	engine







    def      lowercase_			( self					:    Optional[Any]		,     __snake_case					:    int		,     **__snake_case					:    List[Any]      ):
      # runs backpropagation and handles mixed precision
      self.engine.backward(__snake_case		,     **__snake_case      )

      # Deepspeed's `engine.step` performs the following operations:
      # - gradient accumulation check
      # - gradient clipping
      # - optimizer step
      # - zero grad
      # - checking overflow
      # - lr_scheduler step (only if engine.lr_scheduler is not None)
      self.engine.step()
      # and this plugin overrides the above calls with no-ops when Accelerate runs under
      # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
      # training loop that works transparently under many training regimes.

class 				a__(   lowerCamelCase__      ):







    def __init__( self					:    Dict		,     __snake_case					:    Any      ):
      super().__init__(__snake_case		,     device_placement=__snake_case		,     scaler=__snake_case      )
      a		:   List[Any]	  =	hasattr(self.optimizer		,     'overflow'      )







    def      lowercase_			( self					:    Optional[Any]		,     __snake_case					:    str=None      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







    def      lowercase_			( self					:    int      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







    @property
    def      lowercase_			( self					:    Any      ):
      if self.__has_overflow__:
        return self.optimizer.overflow
      return False

class 				a__(   lowerCamelCase__      ):







    def __init__( self					:    str		,     __snake_case					:    List[Any]		,     __snake_case					:    str      ):
      super().__init__(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    Union[str, Any]      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed

class 				a__:







    def __init__( self					:    Optional[Any]		,     __snake_case					:    Union[str, Any]		,     __snake_case					:    int=0.001		,     __snake_case					:    Any=0		,     **__snake_case					:    List[Any]      ):
      a		:   Optional[Any]	  =	params
      a		:   Union[str, Any]	  =	lr
      a		:   Any	  =	weight_decay
      a		:   Any	  =	kwargs

class 				a__:







    def __init__( self					:    Any		,     __snake_case					:    int		,     __snake_case					:    int=None		,     __snake_case					:    Dict=0		,     **__snake_case					:    Optional[int]      ):
      a		:   Optional[int]	  =	optimizer
      a		:   List[str]	  =	total_num_steps
      a		:   Tuple	  =	warmup_num_steps
      a		:   Optional[Any]	  =	kwargs