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

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

class 							_UpperCamelCase					:






			'''simple docstring'''

			def __init__(							self			:							List[str]    ,						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
											UpperCamelCase_:				Optional[int]						=    deepcopy(snake_case_   )
							elif os.path.exists(snake_case_   ):
											with io.open(snake_case_    ,						"""r"""    ,						encoding="""utf-8"""   ) as f:
															UpperCamelCase_:				Union[str, Any]						=    json.load(snake_case_   )
							else:
											try:
															UpperCamelCase_:				Optional[int]						=    baseaa.urlsafe_baadecode(snake_case_   ).decode("""utf-8"""   )
															UpperCamelCase_:				Optional[int]						=    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}'''   )

							UpperCamelCase_:				Optional[int]						=    config

							self.set_stage_and_offload()

			def    lowerCAmelCase__						(							self			:							int   ):
							# 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.
							UpperCamelCase_:				int						=    self.get_value("""zero_optimization.stage"""    ,						-1   )

							# offload
							UpperCamelCase_:				Optional[Any]						=    False
							if self.is_zeroa() or self.is_zeroa():
											UpperCamelCase_:				Tuple						=    set(["""cpu""", """nvme"""]   )
											UpperCamelCase_:				Optional[int]						=    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:
															UpperCamelCase_:				Tuple						=    True

			def    lowerCAmelCase__						(							self			:							str    ,						snake_case_			:							Tuple   ):
							UpperCamelCase_:				Any						=    self.config

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

							return config, ds_key

			def    lowerCAmelCase__						(							self			:							List[str]    ,						snake_case_			:							Optional[Any]    ,						snake_case_			:							Optional[int]=None   ):
							UpperCamelCase_, UpperCamelCase_:				List[str]						=    self.find_config_node(snake_case_   )
							if config is None:
											return default
							return config.get(snake_case_    ,						snake_case_   )

			def    lowerCAmelCase__						(							self			:							Dict    ,						snake_case_			:							Union[str, Any]    ,						snake_case_			:							Optional[Any]=False   ):
							UpperCamelCase_:				Any						=    self.config

							# find the config node of interest if it exists
							UpperCamelCase_:				List[str]						=    ds_key_long.split("""."""   )
							for node in nodes:
											UpperCamelCase_:				Tuple						=    config
											UpperCamelCase_:				Optional[Any]						=    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    lowerCAmelCase__						(							self			:							str    ,						snake_case_			:							Dict   ):
							UpperCamelCase_:				Dict						=    self.get_value(snake_case_   )
							return False if value is None else bool(snake_case_   )

			def    lowerCAmelCase__						(							self			:							List[str]    ,						snake_case_			:							int   ):
							UpperCamelCase_:				Dict						=    self.get_value(snake_case_   )
							return False if value is None else not bool(snake_case_   )

			def    lowerCAmelCase__						(							self			:							List[Any]   ):
							return self._stage == 2

			def    lowerCAmelCase__						(							self			:							List[Any]   ):
							return self._stage == 3




			def    lowerCAmelCase__						(							self			:							Tuple   ):
							return self._offload




class 							_UpperCamelCase					:






			'''simple docstring'''

			def __init__(							self			:							Optional[Any]    ,						snake_case_			:							Union[str, Any]   ):
							UpperCamelCase_:				Optional[int]						=    engine




			def    lowerCAmelCase__						(							self			:							Optional[Any]    ,						snake_case_			:							List[str]    ,						**snake_case_			:							Dict   ):
							# 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 							_UpperCamelCase					(       _A	):






			'''simple docstring'''

			def __init__(							self			:							Dict    ,						snake_case_			:							Tuple   ):
							super().__init__(snake_case_    ,						device_placement=snake_case_    ,						scaler=snake_case_   )
							UpperCamelCase_:				Tuple						=    hasattr(self.optimizer    ,						"""overflow"""   )

			def    lowerCAmelCase__						(							self			:							int    ,						snake_case_			:							Dict=None   ):
							pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed

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




			@property
			def    lowerCAmelCase__						(							self			:							str   ):
							if self.__has_overflow__:
											return self.optimizer.overflow
							return False




class 							_UpperCamelCase					(       _A	):






			'''simple docstring'''

			def __init__(							self			:							Optional[int]    ,						snake_case_			:							List[Any]    ,						snake_case_			:							Union[str, Any]   ):
							super().__init__(snake_case_    ,						snake_case_   )




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




class 							_UpperCamelCase					:






			'''simple docstring'''

			def __init__(							self			:							Union[str, Any]    ,						snake_case_			:							List[Any]    ,						snake_case_			:							Tuple=0.001    ,						snake_case_			:							Tuple=0    ,						**snake_case_			:							Optional[Any]   ):
							UpperCamelCase_:				Tuple						=    params
							UpperCamelCase_:				Any						=    lr
							UpperCamelCase_:				Union[str, Any]						=    weight_decay
							UpperCamelCase_:				List[Any]						=    kwargs




class 							_UpperCamelCase					:






			'''simple docstring'''

			def __init__(							self			:							Optional[int]    ,						snake_case_			:							List[Any]    ,						snake_case_			:							int=None    ,						snake_case_			:							int=0    ,						**snake_case_			:							str   ):
							UpperCamelCase_:				Tuple						=    optimizer
							UpperCamelCase_:				Union[str, Any]						=    total_num_steps
							UpperCamelCase_:				Optional[int]						=    warmup_num_steps
							UpperCamelCase_:				str						=    kwargs








import collections
import os
import re
from pathlib import Path


lowerCamelCase_    :			Optional[Any]    =	"""src/transformers"""


# Matches is_xxx_available()
lowerCamelCase_    :			Union[str, Any]    =	re.compile(r"""is\_([a-z_]*)_available()""")
# Catches a one-line _import_struct = {xxx}
lowerCamelCase_    :			int    =	re.compile(r"""^_import_structure\s+=\s+\{([^\}]+)\}""")
# Catches a line with a key-values pattern: "bla": ["foo", "bar"]
lowerCamelCase_    :			Union[str, Any]    =	re.compile(r"""\s+\"\S*\":\s+\[([^\]]*)\]""")
# Catches a line if not is_foo_available
lowerCamelCase_    :			Any    =	re.compile(r"""^\s*if\s+not\s+is\_[a-z_]*\_available\(\)""")
# Catches a line _import_struct["bla"].append("foo")
lowerCamelCase_    :			Any    =	re.compile(r"""^\s*_import_structure\[\"\S*\"\]\.append\(\"(\S*)\"\)""")
# Catches a line _import_struct["bla"].extend(["foo", "bar"]) or _import_struct["bla"] = ["foo", "bar"]
lowerCamelCase_    :			List[Any]    =	re.compile(r"""^\s*_import_structure\[\S*\](?:\.extend\(|\s*=\s+)\[([^\]]*)\]""")
# Catches a line with an object between quotes and a comma:     "MyModel",
lowerCamelCase_    :			Any    =	re.compile(r"""^\s+\"([^\"]+)\",""")
# Catches a line with objects between brackets only:    ["foo", "bar"],
lowerCamelCase_    :			Tuple    =	re.compile(r"""^\s+\[([^\]]+)\]""")
# Catches a line with from foo import bar, bla, boo
lowerCamelCase_    :			Tuple    =	re.compile(r"""\s+from\s+\S*\s+import\s+([^\(\s].*)\n""")
# Catches a line with try:
lowerCamelCase_    :			Dict    =	re.compile(r"""^\s*try:""")
# Catches a line with else:
lowerCamelCase_    :			Union[str, Any]    =	re.compile(r"""^\s*else:""")







def 							A__       (    lowerCamelCase    )	->						List[Any]:
				if _re_test_backend.search(lowerCamelCase    ) is None:
								return None
				UpperCamelCase_:				Any						=    [b[0] for b in _re_backend.findall(lowerCamelCase    )]
				backends.sort()
				return "_and_".join(lowerCamelCase    )







def 							A__       (    lowerCamelCase    )	->						Union[str, Any]:
				with open(lowerCamelCase  ,     """r"""  ,     encoding="""utf-8"""  ,     newline="""\n"""    ) as f:
								UpperCamelCase_:				Dict						=    f.readlines()

				UpperCamelCase_:				Tuple						=    0
				while line_index < len(lowerCamelCase    ) and not lines[line_index].startswith("""_import_structure = {"""    ):
								line_index += 1

				# If this is a traditional init, just return.
				if line_index >= len(lowerCamelCase    ):
								return None

				# First grab the objects without a specific backend in _import_structure
				UpperCamelCase_:				Optional[int]						=    []
				while not lines[line_index].startswith("""if TYPE_CHECKING"""    ) and find_backend(lines[line_index]    ) is None:
								UpperCamelCase_:				Any						=    lines[line_index]
								# If we have everything on a single line, let's deal with it.
								if _re_one_line_import_struct.search(lowerCamelCase    ):
												UpperCamelCase_:				str						=    _re_one_line_import_struct.search(lowerCamelCase    ).groups()[0]
												UpperCamelCase_:				Tuple						=    re.findall(r"""\[([^\]]+)\]"""  ,     lowerCamelCase    )
												for imp in imports:
																objects.extend([obj[1:-1] for obj in imp.split(""", """    )]    )
												line_index += 1
												continue
								UpperCamelCase_:				Any						=    _re_import_struct_key_value.search(lowerCamelCase    )
								if single_line_import_search is not None:
												UpperCamelCase_:				Any						=    [obj[1:-1] for obj in single_line_import_search.groups()[0].split(""", """    ) if len(lowerCamelCase    ) > 0]
												objects.extend(lowerCamelCase    )
								elif line.startswith(""" """ * 8 + """\""""    ):
												objects.append(line[9:-3]    )
								line_index += 1

				UpperCamelCase_:				Optional[Any]						=    {"""none""": objects}
				# Let's continue with backend-specific objects in _import_structure
				while not lines[line_index].startswith("""if TYPE_CHECKING"""    ):
								# If the line is an if not is_backend_available, we grab all objects associated.
								UpperCamelCase_:				Any						=    find_backend(lines[line_index]    )
								# Check if the backend declaration is inside a try block:
								if _re_try.search(lines[line_index - 1]    ) is None:
												UpperCamelCase_:				Any						=    None

								if backend is not None:
												line_index += 1

												# Scroll until we hit the else block of try-except-else
												while _re_else.search(lines[line_index]    ) is None:
																line_index += 1

												line_index += 1

												UpperCamelCase_:				int						=    []
												# Until we unindent, add backend objects to the list
												while len(lines[line_index]    ) <= 1 or lines[line_index].startswith(""" """ * 4    ):
																UpperCamelCase_:				Tuple						=    lines[line_index]
																if _re_import_struct_add_one.search(lowerCamelCase    ) is not None:
																				objects.append(_re_import_struct_add_one.search(lowerCamelCase    ).groups()[0]    )
																elif _re_import_struct_add_many.search(lowerCamelCase    ) is not None:
																				UpperCamelCase_:				List[str]						=    _re_import_struct_add_many.search(lowerCamelCase    ).groups()[0].split(""", """    )
																				UpperCamelCase_:				str						=    [obj[1:-1] for obj in imports if len(lowerCamelCase    ) > 0]
																				objects.extend(lowerCamelCase    )
																elif _re_between_brackets.search(lowerCamelCase    ) is not None:
																				UpperCamelCase_:				Tuple						=    _re_between_brackets.search(lowerCamelCase    ).groups()[0].split(""", """    )
																				UpperCamelCase_:				List[Any]						=    [obj[1:-1] for obj in imports if len(lowerCamelCase    ) > 0]
																				objects.extend(lowerCamelCase    )
																elif _re_quote_object.search(lowerCamelCase    ) is not None:
																				objects.append(_re_quote_object.search(lowerCamelCase    ).groups()[0]    )
																elif line.startswith(""" """ * 8 + """\""""    ):
																				objects.append(line[9:-3]    )
																elif line.startswith(""" """ * 12 + """\""""    ):
																				objects.append(line[13:-3]    )
																line_index += 1

												UpperCamelCase_:				List[str]						=    objects
								else:
												line_index += 1

    # At this stage we are in the TYPE_CHECKING part, first grab the objects without a specific backend
				UpperCamelCase_:				List[str]						=    []
				while (
				    line_index < len(lowerCamelCase    )
				    and find_backend(lines[line_index]    ) is None
				    and not lines[line_index].startswith("""else"""    )
				):
								UpperCamelCase_:				List[str]						=    lines[line_index]
								UpperCamelCase_:				Union[str, Any]						=    _re_import.search(lowerCamelCase    )
								if single_line_import_search is not None:
												objects.extend(single_line_import_search.groups()[0].split(""", """    )    )
								elif line.startswith(""" """ * 8    ):
												objects.append(line[8:-2]    )
								line_index += 1

				UpperCamelCase_:				Any						=    {"""none""": objects}
				# Let's continue with backend-specific objects
				while line_index < len(lowerCamelCase    ):
								# If the line is an if is_backend_available, we grab all objects associated.
								UpperCamelCase_:				str						=    find_backend(lines[line_index]    )
								# Check if the backend declaration is inside a try block:
								if _re_try.search(lines[line_index - 1]    ) is None:
												UpperCamelCase_:				List[Any]						=    None

								if backend is not None:
												line_index += 1

												# Scroll until we hit the else block of try-except-else
												while _re_else.search(lines[line_index]    ) is None:
																line_index += 1

												line_index += 1

												UpperCamelCase_:				int						=    []
												# Until we unindent, add backend objects to the list
												while len(lines[line_index]    ) <= 1 or lines[line_index].startswith(""" """ * 8    ):
																UpperCamelCase_:				Tuple						=    lines[line_index]
																UpperCamelCase_:				Union[str, Any]						=    _re_import.search(lowerCamelCase    )
																if single_line_import_search is not None:
																				objects.extend(single_line_import_search.groups()[0].split(""", """    )    )
																elif line.startswith(""" """ * 12    ):
																				objects.append(line[12:-2]    )
																line_index += 1

												UpperCamelCase_:				Optional[Any]						=    objects
								else:
												line_index += 1

				return import_dict_objects, type_hint_objects







def 							A__       (    lowerCamelCase  ,     lowerCamelCase    )	->						List[Any]:

				def find_duplicates(lowerCamelCase    ):
								return [k for k, v in collections.Counter(lowerCamelCase    ).items() if v > 1]

				if list(import_dict_objects.keys()    ) != list(type_hint_objects.keys()    ):
								return ["Both sides of the init do not have the same backends!"]

				UpperCamelCase_:				Optional[int]						=    []
				for key in import_dict_objects.keys():
								UpperCamelCase_:				List[str]						=    find_duplicates(import_dict_objects[key]    )
								if duplicate_imports:
												errors.append(F'''Duplicate _import_structure definitions for: {duplicate_imports}'''    )
								UpperCamelCase_:				Tuple						=    find_duplicates(type_hint_objects[key]    )
								if duplicate_type_hints:
												errors.append(F'''Duplicate TYPE_CHECKING objects for: {duplicate_type_hints}'''    )

								if sorted(set(import_dict_objects[key]    )    ) != sorted(set(type_hint_objects[key]    )    ):
												UpperCamelCase_:				int						=    """base imports""" if key == """none""" else F'''{key} backend'''
												errors.append(F'''Differences for {name}:'''    )
												for a in type_hint_objects[key]:
																if a not in import_dict_objects[key]:
																				errors.append(F'''  {a} in TYPE_HINT but not in _import_structure.'''    )
												for a in import_dict_objects[key]:
																if a not in type_hint_objects[key]:
																				errors.append(F'''  {a} in _import_structure but not in TYPE_HINT.'''    )
				return errors







def 							A__       (    )	->						Tuple:
				UpperCamelCase_:				Optional[int]						=    []
				for root, _, files in os.walk(lowerCamelCase    ):
								if "__init__.py" in files:
												UpperCamelCase_:				Union[str, Any]						=    os.path.join(lowerCamelCase  ,     """__init__.py"""    )
												UpperCamelCase_:				Optional[int]						=    parse_init(lowerCamelCase    )
												if objects is not None:
																UpperCamelCase_:				Any						=    analyze_results(*lowerCamelCase    )
																if len(lowerCamelCase    ) > 0:
																				UpperCamelCase_:				Any						=    F'''Problem in {fname}, both halves do not define the same objects.\n{errors[0]}'''
																				failures.append("""\n""".join(lowerCamelCase    )    )
				if len(lowerCamelCase    ) > 0:
								raise ValueError("""\n\n""".join(lowerCamelCase    )    )







def 							A__       (    )	->						Any:
				UpperCamelCase_:				List[Any]						=    []
				for path, directories, files in os.walk(lowerCamelCase    ):
								for folder in directories:
												# Ignore private modules
												if folder.startswith("""_"""    ):
																directories.remove(lowerCamelCase    )
																continue
												# Ignore leftovers from branches (empty folders apart from pycache)
												if len(list((Path(lowerCamelCase    ) / folder).glob("""*.py"""    )    )    ) == 0:
																continue
												UpperCamelCase_:				str						=    str((Path(lowerCamelCase    ) / folder).relative_to(lowerCamelCase    )    )
												UpperCamelCase_:				Optional[Any]						=    short_path.replace(os.path.sep  ,     """."""    )
												submodules.append(lowerCamelCase    )
								for fname in files:
												if fname == "__init__.py":
																continue
												UpperCamelCase_:				Dict						=    str((Path(lowerCamelCase    ) / fname).relative_to(lowerCamelCase    )    )
												UpperCamelCase_:				Optional[int]						=    short_path.replace(""".py"""  ,     """"""    ).replace(os.path.sep  ,     """."""    )
												if len(submodule.split("""."""    )    ) == 1:
																submodules.append(lowerCamelCase    )
				return submodules


lowerCamelCase_    :			Optional[int]    =	[
    """convert_pytorch_checkpoint_to_tf2""",
    """modeling_flax_pytorch_utils""",
    """models.esm.openfold_utils""",
]







def 							A__       (    )	->						List[str]:
				# This is to make sure the transformers module imported is the one in the repo.
				from transformers.utils import direct_transformers_import

				UpperCamelCase_:				Optional[Any]						=    direct_transformers_import(lowerCamelCase    )

				UpperCamelCase_:				Tuple						=    set(transformers._import_structure.keys()    )
				# This contains all the base keys of the _import_structure object defined in the init, but if the user is missing
				# some optional dependencies, they may not have all of them. Thus we read the init to read all additions and
				# (potentiall re-) add them.
				with open(os.path.join(lowerCamelCase  ,     """__init__.py"""    )  ,     """r"""    ) as f:
								UpperCamelCase_:				List[Any]						=    f.read()
				import_structure_keys.update(set(re.findall(r"""import_structure\[\"([^\"]*)\"\]"""  ,     lowerCamelCase    )    )    )

				UpperCamelCase_:				Dict						=    [
				    module
				    for module in get_transformers_submodules()
				    if module not in IGNORE_SUBMODULES and module not in import_structure_keys
				]

				if len(lowerCamelCase    ) > 0:
								UpperCamelCase_:				Dict						=    """\n""".join(F'''- {module}''' for module in module_not_registered    )
								raise ValueError(
								    """The following submodules are not properly registed in the main init of Transformers:\n"""
								    F'''{list_of_modules}\n'''
								    """Make sure they appear somewhere in the keys of `_import_structure` with an empty list as value."""    )


if __name__ == "__main__":
							check_all_inits()
							check_submodules()






