




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

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







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__ ):
							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
										SCREAMING_SNAKE_CASE_   : Any					  =      deepcopy(snake_case__ )
							elif os.path.exists(snake_case__ ):
										with io.open(snake_case__      ,'r'      ,encoding='utf-8' ) as f:
													SCREAMING_SNAKE_CASE_   : Dict					  =      json.load(snake_case__ )
							else:
										try:
													SCREAMING_SNAKE_CASE_   : Dict					  =      baseaa.urlsafe_baadecode(snake_case__ ).decode('utf-8' )
													SCREAMING_SNAKE_CASE_   : Optional[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}' )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      config

							self.set_stage_and_offload()






				def 		snake_case	(			self ):
							# 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.
							SCREAMING_SNAKE_CASE_   : List[str]					  =      self.get_value('zero_optimization.stage'      ,-1 )

							# offload
							SCREAMING_SNAKE_CASE_   : int					  =      False
							if self.is_zeroa() or self.is_zeroa():
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      set(['cpu', 'nvme'] )
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      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:
													SCREAMING_SNAKE_CASE_   : int					  =      True






				def 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      self.config

							# find the config node of interest if it exists
							SCREAMING_SNAKE_CASE_   : Dict					  =      ds_key_long.split('.' )
							SCREAMING_SNAKE_CASE_   : Dict					  =      nodes.pop()
							for node in nodes:
										SCREAMING_SNAKE_CASE_   : Optional[int]					  =      config.get(snake_case__ )
										if config is None:
													return None, ds_key

							return config, ds_key






				def 		snake_case	(			self      ,snake_case__      ,snake_case__=None ):
							SCREAMING_SNAKE_CASE_  ,    SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      self.find_config_node(snake_case__ )
							if config is None:
										return default
							return config.get(snake_case__      ,snake_case__ )






				def 		snake_case	(			self      ,snake_case__      ,snake_case__=False ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.config

							# find the config node of interest if it exists
							SCREAMING_SNAKE_CASE_   : int					  =      ds_key_long.split('.' )
							for node in nodes:
										SCREAMING_SNAKE_CASE_   : Any					  =      config
										SCREAMING_SNAKE_CASE_   : 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 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.get_value(snake_case__ )
							return False if value is None else bool(snake_case__ )






				def 		snake_case	(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : int					  =      self.get_value(snake_case__ )
							return False if value is None else not bool(snake_case__ )






				def 		snake_case	(			self ):
							return self._stage == 2






				def 		snake_case	(			self ):
							return self._stage == 3






				def 		snake_case	(			self ):
							return self._offload







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      engine






				def 		snake_case	(			self      ,snake_case__      ,**snake_case__ ):
							# 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 				lowerCAmelCase_       (					lowerCamelCase_							):






				def __init__(			self      ,snake_case__ ):
							super().__init__(snake_case__      ,device_placement=snake_case__      ,scaler=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Dict					  =      hasattr(self.optimizer      ,'overflow' )






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






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






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







class 				lowerCAmelCase_       (					lowerCamelCase_							):






				def __init__(			self      ,snake_case__      ,snake_case__ ):
							super().__init__(snake_case__      ,snake_case__ )






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







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__      ,snake_case__=0.001      ,snake_case__=0      ,**snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      params
							SCREAMING_SNAKE_CASE_   : Dict					  =      lr
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      weight_decay
							SCREAMING_SNAKE_CASE_   : Dict					  =      kwargs







class 				lowerCAmelCase_       :






				def __init__(			self      ,snake_case__      ,snake_case__=None      ,snake_case__=0      ,**snake_case__ ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      optimizer
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      total_num_steps
							SCREAMING_SNAKE_CASE_   : Dict					  =      warmup_num_steps
							SCREAMING_SNAKE_CASE_   : str					  =      kwargs










import json
import os
from pathlib import Path

import pytest

from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadManager
from datasets.utils.file_utils import hash_url_to_filename


__A							:			List[Any]         =			"http://www.mocksite.com/file1.txt"
__A							:			List[Any]         =			"\"text\": [\"foo\", \"foo\"]"
__A							:			Dict         =			"6d8ce9aa78a471c7477201efbeabd3bb01ac2e7d100a6dc024ba1608361f90a8"







class      lowerCamelCase:
     '''simple docstring'''


     __magic_name__   =			200
     __magic_name__   =			{'Content-Length': '100'}
     __magic_name__   =			{}




     def        lowerCAmelCase__		(    self				,     **snake_case_							):
         return [bytes(snake_case_				,     'utf-8'							)]





def 				__lowerCAmelCase(       *_SCREAMING_SNAKE_CASE							,     **_SCREAMING_SNAKE_CASE				)							-> Optional[Any]:




    """simple docstring"""




    return MockResponse()





@pytest.mark.parametrize('urls_type'							,     [str, list, dict]				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> str:




    """simple docstring"""




    import requests

    monkeypatch.setattr(_SCREAMING_SNAKE_CASE							,     'request'							,     _SCREAMING_SNAKE_CASE				)

    _A     =							URL
    if issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							url
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							[url]
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							{'train': url}
    _A     =							'dummy'
    _A     =							'downloads'
    _A     =							tmp_path
    _A     =							DownloadConfig(
        cache_dir=os.path.join(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							,     use_etag=_SCREAMING_SNAKE_CASE							,     )
    _A     =							DownloadManager(dataset_name=_SCREAMING_SNAKE_CASE							,     download_config=_SCREAMING_SNAKE_CASE				)
    _A     =							dl_manager.download(_SCREAMING_SNAKE_CASE				)
    _A     =							urls
    for downloaded_paths in [downloaded_paths]:
        if isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            _A     =							[downloaded_paths]
            _A     =							[urls]
        elif isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert "train" in downloaded_paths.keys()
            _A     =							downloaded_paths.values()
            _A     =							urls.values()
        assert downloaded_paths
        for downloaded_path, input_url in zip(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert downloaded_path == dl_manager.downloaded_paths[input_url]
            _A     =							Path(_SCREAMING_SNAKE_CASE				)
            _A     =							downloaded_path.parts
            assert parts[-1] == HASH
            assert parts[-2] == cache_subdir
            assert downloaded_path.exists()
            _A     =							downloaded_path.read_text()
            assert content == CONTENT
            _A     =							downloaded_path.with_suffix('.json'				)
            assert metadata_downloaded_path.exists()
            _A     =							json.loads(metadata_downloaded_path.read_text()				)
            assert metadata_content == {"url": URL, "etag": None}





@pytest.mark.parametrize('paths_type'							,     [str, list, dict]				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Optional[Any]:




    """simple docstring"""




    _A     =							str(_SCREAMING_SNAKE_CASE				)
    if issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							filename
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							[filename]
    elif issubclass(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
        _A     =							{'train': filename}
    _A     =							'dummy'
    _A     =							xz_file.parent
    _A     =							'extracted'
    _A     =							DownloadConfig(
        cache_dir=_SCREAMING_SNAKE_CASE							,     use_etag=_SCREAMING_SNAKE_CASE							,     )
    _A     =							DownloadManager(dataset_name=_SCREAMING_SNAKE_CASE							,     download_config=_SCREAMING_SNAKE_CASE				)
    _A     =							dl_manager.extract(_SCREAMING_SNAKE_CASE				)
    _A     =							paths
    for extracted_paths in [extracted_paths]:
        if isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            _A     =							[extracted_paths]
            _A     =							[paths]
        elif isinstance(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert "train" in extracted_paths.keys()
            _A     =							extracted_paths.values()
            _A     =							paths.values()
        assert extracted_paths
        for extracted_path, input_path in zip(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				):
            assert extracted_path == dl_manager.extracted_paths[input_path]
            _A     =							Path(_SCREAMING_SNAKE_CASE				)
            _A     =							extracted_path.parts
            assert parts[-1] == hash_url_to_filename(_SCREAMING_SNAKE_CASE							,     etag=_SCREAMING_SNAKE_CASE				)
            assert parts[-2] == extracted_subdir
            assert extracted_path.exists()
            _A     =							extracted_path.read_text()
            _A     =							text_file.read_text()
            assert extracted_file_content == expected_file_content





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Dict:




    """simple docstring"""




    assert path.endswith('.jsonl'				)
    for num_items, line in enumerate(_SCREAMING_SNAKE_CASE							,     start=1				):
        _A     =							json.loads(line.decode('utf-8'				)				)
        assert item.keys() == {"col_1", "col_2", "col_3"}
    assert num_items == 4





@pytest.mark.parametrize('archive_jsonl'							,     ['tar_jsonl_path', 'zip_jsonl_path']				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Optional[int]:




    """simple docstring"""




    _A     =							request.getfixturevalue(_SCREAMING_SNAKE_CASE				)
    _A     =							DownloadManager()
    for num_jsonl, (path, file) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        _test_jsonl(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
    assert num_jsonl == 2





@pytest.mark.parametrize('archive_nested_jsonl'							,     ['tar_nested_jsonl_path', 'zip_nested_jsonl_path']				)
def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)							-> Any:




    """simple docstring"""




    _A     =							request.getfixturevalue(_SCREAMING_SNAKE_CASE				)
    _A     =							DownloadManager()
    for num_tar, (path, file) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        for num_jsonl, (subpath, subfile) in enumerate(dl_manager.iter_archive(_SCREAMING_SNAKE_CASE				)							,     start=1				):
            _test_jsonl(_SCREAMING_SNAKE_CASE							,     _SCREAMING_SNAKE_CASE				)
    assert num_tar == 1
    assert num_jsonl == 2





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE				)							-> int:




    """simple docstring"""




    _A     =							DownloadManager()
    for num_file, file in enumerate(dl_manager.iter_files(_SCREAMING_SNAKE_CASE				)							,     start=1				):
        assert os.path.basename(_SCREAMING_SNAKE_CASE				) == ("test.txt" if num_file == 1 else "train.txt")
    assert num_file == 2





