


"""simple docstring"""
import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile, rmtree
from typing import Generator

import datasets.config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.filelock import logger as fl_logger
from datasets.utils.info_utils import VerificationMode
from datasets.utils.logging import ERROR, get_logger


__SCREAMING_SNAKE_CASE				:     str		=     get_logger(__name__)







def     _a					(   _SCREAMING_SNAKE_CASE					)     ->					Union[str, Any]:
							return TestCommand(
							    args.dataset					,			args.name					,			args.cache_dir					,			args.data_dir					,			args.all_configs					,			args.save_info or args.save_infos					,			args.ignore_verifications					,			args.force_redownload					,			args.clear_cache					,			)





class   __A  (snake_case__):





							'''simple docstring'''

							__lowercase:		int				   =     False  # to tell pytest it's not a test class



							@staticmethod
							def 		lowerCAmelCase							(					UpperCAmelCase_				:   ArgumentParser							)       ->int:



														"""simple docstring"""



														snake_case_     							=		parser.add_parser("""test"""		,							help="""Test dataset implementation."""							)
														test_parser.add_argument("""--name"""		,							type=UpperCAmelCase_		,							default=UpperCAmelCase_		,							help="""Dataset processing name"""							)
														test_parser.add_argument(
														    """--cache_dir"""		,							type=UpperCAmelCase_		,							default=UpperCAmelCase_		,							help="""Cache directory where the datasets are stored."""		,							)
														test_parser.add_argument(
														    """--data_dir"""		,							type=UpperCAmelCase_		,							default=UpperCAmelCase_		,							help="""Can be used to specify a manual directory to get the files from."""		,							)
														test_parser.add_argument("""--all_configs"""		,							action="""store_true"""		,							help="""Test all dataset configurations"""							)
														test_parser.add_argument(
														    """--save_info"""		,							action="""store_true"""		,							help="""Save the dataset infos in the dataset card (README.md)"""							)
														test_parser.add_argument(
														    """--ignore_verifications"""		,							action="""store_true"""		,							help="""Run the test without checksums and splits checks."""		,							)
														test_parser.add_argument("""--force_redownload"""		,							action="""store_true"""		,							help="""Force dataset redownload"""							)
														test_parser.add_argument(
														    """--clear_cache"""		,							action="""store_true"""		,							help="""Remove downloaded files and cached datasets after each config test"""		,							)
														# aliases
														test_parser.add_argument("""--save_infos"""		,							action="""store_true"""		,							help="""alias to save_info"""							)
														test_parser.add_argument("""dataset"""		,							type=UpperCAmelCase_		,							help="""Name of the dataset to download"""							)
														test_parser.set_defaults(func=_test_command_factory							)



							def __init__(					self				:   List[Any]		,							UpperCAmelCase_				:   str		,							UpperCAmelCase_				:   str		,							UpperCAmelCase_				:   str		,							UpperCAmelCase_				:   str		,							UpperCAmelCase_				:   bool		,							UpperCAmelCase_				:   bool		,							UpperCAmelCase_				:   bool		,							UpperCAmelCase_				:   bool		,							UpperCAmelCase_				:   bool		,							)       ->Optional[Any]:



														"""simple docstring"""



														snake_case_     							=		dataset
														snake_case_     							=		name
														snake_case_     							=		cache_dir
														snake_case_     							=		data_dir
														snake_case_     							=		all_configs
														snake_case_     							=		save_infos
														snake_case_     							=		ignore_verifications
														snake_case_     							=		force_redownload
														snake_case_     							=		clear_cache
														if clear_cache and not cache_dir:
																					print(
																					    """When --clear_cache is used, specifying a cache directory is mandatory.\n"""
																					    """The 'download' folder of the cache directory and the dataset builder cache will be deleted after each configuration test.\n"""
																					    """Please provide a --cache_dir that will be used to test the dataset script."""							)
																					exit(1							)
														if save_infos:
																					snake_case_     							=		True



							def 		lowerCAmelCase							(					self				:   Any							)       ->Any:



														"""simple docstring"""



														fl_logger().setLevel(UpperCAmelCase_							)
														if self._name is not None and self._all_configs:
																					print("""Both parameters `config` and `all_configs` can't be used at once."""							)
																					exit(1							)
														snake_case_							,    snake_case_     							=		self._dataset, self._name
														snake_case_     							=		dataset_module_factory(UpperCAmelCase_							)
														snake_case_     							=		import_main_class(module.module_path							)
														snake_case_     							=		len(builder_cls.BUILDER_CONFIGS							) if self._all_configs and builder_cls.BUILDER_CONFIGS else 1

														def get_builders() -> Generator[DatasetBuilder, None, None]:
																					if self._all_configs and builder_cls.BUILDER_CONFIGS:
																												for i, config in enumerate(builder_cls.BUILDER_CONFIGS							):
																																			if "config_name" in module.builder_kwargs:
																																										yield builder_cls(
																																										    cache_dir=self._cache_dir		,							data_dir=self._data_dir		,							**module.builder_kwargs		,							)
																																			else:
																																										yield builder_cls(
																																										    config_name=config.name		,							cache_dir=self._cache_dir		,							data_dir=self._data_dir		,							**module.builder_kwargs		,							)
																					else:
																												if "config_name" in module.builder_kwargs:
																																			yield builder_cls(cache_dir=self._cache_dir		,							data_dir=self._data_dir		,							**module.builder_kwargs							)
																												else:
																																			yield builder_cls(
																																			    config_name=UpperCAmelCase_		,							cache_dir=self._cache_dir		,							data_dir=self._data_dir		,							**module.builder_kwargs		,							)

														for j, builder in enumerate(get_builders()							):
																					print(F"""Testing builder '{builder.config.name}' ({j + 1}/{n_builders})"""							)
																					snake_case_     							=		os.path.exists(
																					    os.path.join(builder.get_imported_module_dir()		,							datasets.config.DATASETDICT_INFOS_FILENAME							)							)  # record checksums only if we need to update a (deprecated) dataset_infos.json
																					builder.download_and_prepare(
																					    download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
																					    if not self._force_redownload
																					    else DownloadMode.FORCE_REDOWNLOAD		,							verification_mode=VerificationMode.NO_CHECKS
																					    if self._ignore_verifications
																					    else VerificationMode.ALL_CHECKS		,							try_from_hf_gcs=UpperCAmelCase_		,							)
																					builder.as_dataset()
																					if self._save_infos:
																												builder._save_infos()

																					# If save_infos=True, the dataset card (README.md) is created next to the loaded module file.
																					# The dataset_infos are saved in the YAML part of the README.md

																					# Let's move it to the original directory of the dataset script, to allow the user to
																					# upload them on S3 at the same time afterwards.
																					if self._save_infos:
																												snake_case_     							=		os.path.join(builder_cls.get_imported_module_dir()		,							"""README.md"""							)
																												snake_case_     							=		Path(UpperCAmelCase_							).name + """.py"""
																												snake_case_     							=		os.path.join(UpperCAmelCase_		,							UpperCAmelCase_							)
																												if os.path.isfile(UpperCAmelCase_							):
																																			snake_case_     							=		os.path.dirname(UpperCAmelCase_							)
																												elif os.path.isfile(UpperCAmelCase_							):
																																			snake_case_     							=		path
																												elif os.path.isdir(UpperCAmelCase_							):  # for local directories containing only data files
																																			snake_case_     							=		path
																												else:  # in case of a remote dataset
																																			snake_case_     							=		None
																																			print(F"""Dataset card saved at {dataset_readme_path}"""							)

																												# Move dataset_info back to the user
																												if dataset_dir is not None:
																																			snake_case_     							=		os.path.join(UpperCAmelCase_		,							"""README.md"""							)
																																			copyfile(UpperCAmelCase_		,							UpperCAmelCase_							)
																																			print(F"""Dataset card saved at {user_dataset_readme_path}"""							)

            # If clear_cache=True, the download folder and the dataset builder cache directory are deleted
																					if self._clear_cache:
																												if os.path.isdir(builder._cache_dir							):
																																			logger.warning(F"""Clearing cache at {builder._cache_dir}"""							)
																																			rmtree(builder._cache_dir							)
																												snake_case_     							=		os.path.join(self._cache_dir		,							datasets.config.DOWNLOADED_DATASETS_DIR							)
																												if os.path.isdir(UpperCAmelCase_							):
																																			logger.warning(F"""Clearing cache at {download_dir}"""							)
																																			rmtree(UpperCAmelCase_							)

														print("""Test successful."""							)







"""simple docstring"""
import importlib
import importlib.metadata
import os
import warnings
from distutils.util import strtobool
from functools import lru_cache

import torch
from packaging import version
from packaging.version import parse

from .environment import parse_flag_from_env
from .versions import compare_versions, is_torch_version


try:
				import torch_xla.core.xla_model as xm  # noqa: F401

				__SCREAMING_SNAKE_CASE				:     List[Any]		=     True
except ImportError:
				__SCREAMING_SNAKE_CASE				:     Dict		=     False


# Cache this result has it's a C FFI call which can be pretty time-consuming
__SCREAMING_SNAKE_CASE				:     Optional[Any]		=     torch.distributed.is_available()







def     _a					(   _SCREAMING_SNAKE_CASE					)     ->					int:
							# Check we're not importing a "pkg_name" directory somewhere but the actual library by trying to grab the version
							snake_case_     							=		importlib.util.find_spec(_SCREAMING_SNAKE_CASE					) is not None
							if package_exists:
														try:
																					snake_case_     							=		importlib.metadata.metadata(_SCREAMING_SNAKE_CASE					)
																					return True
														except importlib.metadata.PackageNotFoundError:
																					return False







def     _a					(   )     ->					bool:
							return _torch_distributed_available







def     _a					(   )     ->					Union[str, Any]:
							try:
														pass
							except ImportError:
														print(
														    """Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not"""
														    """ detected. If you see \"ValueError: Invalid backend: 'ccl'\" error, please install Intel(R) oneCCL"""
														    """ Bindings for PyTorch*."""					)
							return (
							    importlib.util.find_spec("""torch_ccl"""					) is not None
							    or importlib.util.find_spec("""oneccl_bindings_for_pytorch"""					) is not None
							)







def     _a					(   )     ->					Optional[Any]:
							return importlib.metadata.version("""oneccl_bind_pt"""					)







def     _a					(   )     ->					Optional[Any]:
							return _is_package_available("""transformer_engine"""					)







@lru_cache
def     _a					(   _SCREAMING_SNAKE_CASE=True					)     ->					int:
							# Due to bugs on the amp series GPUs, we disable torch-xla on them
							if torch.cuda.is_available():
														return False
							if _tpu_available and check_device:
														try:
																					# Will raise a RuntimeError if no XLA configuration is found
																					snake_case_     							=		xm.xla_device()
																					return True
														except RuntimeError:
																					return False
							return _tpu_available







def     _a					(   )     ->					List[str]:
							return _is_package_available("""deepspeed"""					)







def     _a					(   _SCREAMING_SNAKE_CASE=False					)     ->					Optional[int]:
							if is_tpu_available():
														return not ignore_tpu
							if torch.cuda.is_available():
														return torch.cuda.is_bfaa_supported()
							if is_npu_available():
														return False
							return True







def     _a					(   )     ->					Any:
							snake_case_     							=		_is_package_available("""bitsandbytes"""					)
							if package_exists:
														snake_case_     							=		version.parse(importlib.metadata.version("""bitsandbytes"""					)					)
														return compare_versions(_SCREAMING_SNAKE_CASE					,			""">="""					,			"""0.39.0"""					)
							return False







def     _a					(   )     ->					List[Any]:
							snake_case_     							=		_is_package_available("""bitsandbytes"""					)
							if package_exists:
														snake_case_     							=		version.parse(importlib.metadata.version("""bitsandbytes"""					)					)
														return compare_versions(_SCREAMING_SNAKE_CASE					,			""">="""					,			"""0.37.2"""					)
							return False







def     _a					(   )     ->					Union[str, Any]:
							return _is_package_available("""bitsandbytes"""					)







def     _a					(   )     ->					Optional[int]:
							if strtobool(os.environ.get("""ACCELERATE_USE_MEGATRON_LM"""					,			"""False"""					)					) == 1:
														snake_case_     							=		importlib.util.find_spec("""megatron"""					) is not None
														if package_exists:
																					try:
																												snake_case_     							=		parse(importlib.metadata.version("""megatron-lm"""					)					)
																												return compare_versions(_SCREAMING_SNAKE_CASE					,			""">="""					,			"""2.2.0"""					)
																					except Exception as e:
																												warnings.warn(f"""Parse Megatron version failed. Exception:{e}"""					)
																												return False







def     _a					(   )     ->					str:
							return _is_package_available("""safetensors"""					)







def     _a					(   )     ->					str:
							return _is_package_available("""transformers"""					)







def     _a					(   )     ->					Union[str, Any]:
							return _is_package_available("""datasets"""					)







def     _a					(   )     ->					List[str]:
							return _is_package_available("""aim"""					)







def     _a					(   )     ->					List[str]:
							return _is_package_available("""tensorboard"""					) or _is_package_available("""tensorboardX"""					)







def     _a					(   )     ->					int:
							return _is_package_available("""wandb"""					)







def     _a					(   )     ->					List[Any]:
							return _is_package_available("""comet_ml"""					)







def     _a					(   )     ->					Union[str, Any]:
							return _is_package_available("""boto3"""					)







def     _a					(   )     ->					Optional[Any]:
							if _is_package_available("""rich"""					):
														if "ACCELERATE_DISABLE_RICH" in os.environ:
																					warnings.warn(
																					    """`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`."""					)
																					return not parse_flag_from_env("""ACCELERATE_DISABLE_RICH"""					,			_SCREAMING_SNAKE_CASE					)
														return parse_flag_from_env("""ACCELERATE_ENABLE_RICH"""					,			_SCREAMING_SNAKE_CASE					)
							return False







def     _a					(   )     ->					List[Any]:
							return _is_package_available("""sagemaker"""					)







def     _a					(   )     ->					int:
							return _is_package_available("""tqdm"""					)







def     _a					(   )     ->					List[Any]:
							return _is_package_available("""mlflow"""					)







def     _a					(   )     ->					int:
							return is_torch_version(""">="""					,			"""1.12"""					) and torch.backends.mps.is_available() and torch.backends.mps.is_built()







def     _a					(   )     ->					Any:
							def get_major_and_minor_from_version(_SCREAMING_SNAKE_CASE					):
														return str(version.parse(_SCREAMING_SNAKE_CASE					).major					) + "." + str(version.parse(_SCREAMING_SNAKE_CASE					).minor					)

							snake_case_     							=		importlib.metadata.version("""torch"""					)
							if importlib.util.find_spec("""intel_extension_for_pytorch"""					) is None:
														return False
							snake_case_     							=		"""N/A"""
							try:
														snake_case_     							=		importlib.metadata.version("""intel_extension_for_pytorch"""					)
							except importlib.metadata.PackageNotFoundError:
														return False
							snake_case_     							=		get_major_and_minor_from_version(_torch_version					)
							snake_case_     							=		get_major_and_minor_from_version(_ipex_version					)
							if torch_major_and_minor != ipex_major_and_minor:
														warnings.warn(
														    f"""Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,"""
														    f""" but PyTorch {_torch_version} is found. Please switch to the matching version and run again."""					)
														return False
							return True







@lru_cache
def     _a					(   _SCREAMING_SNAKE_CASE=False					)     ->					List[str]:
							if importlib.util.find_spec("""torch"""					) is None or importlib.util.find_spec("""torch_npu"""					) is None:
														return False

							import torch
							import torch_npu  # noqa: F401

							if check_device:
														try:
																					# Will raise a RuntimeError if no NPU is found
																					snake_case_     							=		torch.npu.device_count()
																					return torch.npu.is_available()
														except RuntimeError:
																					return False
							return hasattr(_SCREAMING_SNAKE_CASE					,			"""npu"""					) and torch.npu.is_available()







@lru_cache
def     _a					(   _SCREAMING_SNAKE_CASE=False					)     ->					int:
							if not parse_flag_from_env("""ACCELERATE_USE_XPU"""					,			default=_SCREAMING_SNAKE_CASE					):
														return False
							if is_ipex_available():
														import torch

														if is_torch_version("""<="""					,			"""1.12"""					):
																					return False
							else:
														return False

							import intel_extension_for_pytorch  # noqa: F401

							if check_device:
														try:
																					# Will raise a RuntimeError if no XPU  is found
																					snake_case_     							=		torch.xpu.device_count()
																					return torch.xpu.is_available()
														except RuntimeError:
																					return False
							return hasattr(_SCREAMING_SNAKE_CASE					,			"""xpu"""					) and torch.xpu.is_available()



