
import importlib.util
import weakref
from functools import partialmethod

from .dependency_versions_check import dep_version_check
from .utils import is_accelerate_available, is_torch_available, logging


if is_torch_available():
   import torch

__lowerCamelCase							= logging.get_logger(__name__)


def     UpperCamelCase      (					):
       return importlib.util.find_spec("deepspeed"					) is not None


if is_accelerate_available() and is_deepspeed_available():
   from accelerate.utils.deepspeed import HfDeepSpeedConfig as DeepSpeedConfig
else:
   # Inherits from a dummy `object` if accelerate is not available, so that python succeeds to import this file.
   # Deepspeed glue code will never inherit this dummy object as it checks if accelerate is available.
   from builtins import object as DeepSpeedConfig






class 						UpperCAmelCase	(		A_				):



   def __init__(self	:      Tuple					,		snake_case__	:      List[str]				)	->			Any:


          '''simple docstring'''


          set_hf_deepspeed_config(self				)
          dep_version_check("accelerate"				)
          dep_version_check("deepspeed"				)
          super().__init__(snake_case__				)






class 						UpperCAmelCase	(		A_				):



   def __init__(self	:      Tuple					,		snake_case__	:      Optional[int]				)	->			str:


          '''simple docstring'''


          super().__init__(snake_case__				)
          snake_case :     Union[str, Any]           =   None
          snake_case :     Any           =   []



   def        _SCREAMING_SNAKE_CASE     (self	:      List[Any]				)	->			int:


          '''simple docstring'''


          if self._dtype is None:
                 raise ValueError("trainer_config_process() wasn't called yet to tell dtype"				)
          return self._dtype



   def        _SCREAMING_SNAKE_CASE     (self	:      Tuple					,		snake_case__	:      Optional[Any]				)	->			List[Any]:


          '''simple docstring'''


          snake_case :     Any           =   self.get_value(snake_case__				)
          if val is None:
                 return False
          else:
                 return val == "auto"



   def        _SCREAMING_SNAKE_CASE     (self	:      Optional[int]					,		snake_case__	:      Union[str, Any]					,		snake_case__	:      Any					,		snake_case__	:      List[Any]=None					,		snake_case__	:      Optional[Any]=True				)	->			Optional[Any]:


          '''simple docstring'''


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

          if config.get(snake_case__				) == "auto":
                 snake_case :     Tuple           =   hf_val
                 return

          if not must_match:
                 return

          snake_case :     Tuple           =   config.get(snake_case__				)
          if ds_val is not None and ds_val != hf_val:
                 self.mismatches.append(f"""- ds {ds_key_long}={ds_val} vs hf {hf_key}={hf_val}"""				)

   A__      :	List[str]		  =  partialmethod(A_ ,must_match=A_				)



   def        _SCREAMING_SNAKE_CASE     (self	:      Dict					,		snake_case__	:      Dict				)	->			List[Any]:


          '''simple docstring'''


          snake_case :     Tuple           =   args.world_size * args.per_device_train_batch_size * args.gradient_accumulation_steps
          self.fill_match(
              "train_micro_batch_size_per_gpu"					,		args.per_device_train_batch_size					,		"per_device_train_batch_size"				)
          self.fill_match("gradient_accumulation_steps"					,		args.gradient_accumulation_steps					,		"gradient_accumulation_steps"				)
          self.fill_match("train_batch_size"					,		snake_case__					,		"train_batch_size (calculated)"				)
          self.fill_match("gradient_clipping"					,		args.max_grad_norm					,		"max_grad_norm"				)

          self.fill_match("optimizer.params.lr"					,		args.learning_rate					,		"learning_rate"				)
          self.fill_match("optimizer.params.betas"					,		[args.adam_betaa, args.adam_betaa]					,		"adam_beta1+adam_beta2"				)
          self.fill_match("optimizer.params.eps"					,		args.adam_epsilon					,		"adam_epsilon"				)
          self.fill_match("optimizer.params.weight_decay"					,		args.weight_decay					,		"weight_decay"				)

          self.fill_only("scheduler.params.warmup_min_lr"					,		0				)  # not a trainer arg
          self.fill_match("scheduler.params.warmup_max_lr"					,		args.learning_rate					,		"learning_rate"				)
          # total_num_steps - will get set in trainer_config_finalize

          # fp16
          if args.fpaa or args.fpaa_full_eval:
                 snake_case :     str           =   "apex" if args.fpaa_backend == "apex" else "amp"
          else:
                 snake_case :     List[str]           =   None

          if args.save_on_each_node:
                 # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
                 snake_case :     List[str]           =   self.config.get("checkpoint"					,		{}				)
                 snake_case :     str           =   args.save_on_each_node

          # amp: similar to the pytorch native amp - it has a bunch of optional params but we won't set
          # any here unless the user did the work
          self.fill_match(
              "fp16.enabled"					,		((args.fpaa or args.fpaa_full_eval) and fpaa_backend == "amp")					,		"fp16|fp16_full_eval+fp16_backend(amp)"					,		)

          # apex: delegates amp work to apex (which needs to be available), but it cannot be used with any
          # ZeRO features
          self.fill_match("amp.enabled"					,		fpaa_backend == "apex"					,		"fp16+fp16_backend(apex)"				)
          self.fill_match("amp.opt_level"					,		args.fpaa_opt_level					,		"fp16_opt_level"				)

          self.fill_match("bf16.enabled"					,		(args.bfaa or args.bfaa_full_eval)					,		"bf16|bf16_full_eval"				)

          # deepspeed's default mode is fp16 unless there is a config that says differently
          if self.is_true("bf16.enabled"				):
                 snake_case :     List[Any]           =   torch.bfloataa
          elif self.is_false("fp16.enabled"				):
                 snake_case :     List[str]           =   torch.floataa
          else:
                 snake_case :     Optional[Any]           =   torch.floataa



   def        _SCREAMING_SNAKE_CASE     (self	:      List[str]					,		snake_case__	:      str					,		snake_case__	:      Any					,		snake_case__	:      Optional[int]				)	->			List[str]:


          '''simple docstring'''


          snake_case :     Any           =   [
              "zero_optimization.reduce_bucket_size",
              "zero_optimization.stage3_prefetch_bucket_size",
              "zero_optimization.stage3_param_persistence_threshold",
          ]
          snake_case :     int           =   [x for x in hidden_size_based_keys if self.is_auto(snake_case__				)]

          if len(snake_case__				) > 0:
                 if hasattr(model.config					,		"hidden_size"				):
                        snake_case :     Tuple           =   model.config.hidden_size
                 elif hasattr(model.config					,		"hidden_sizes"				):
                        # if there are many hidden sizes pick the largest one
                        snake_case :     Tuple           =   max(model.config.hidden_sizes				)
                 else:
                        raise ValueError(
                            "The model's config file has neither `hidden_size` nor `hidden_sizes` entry, "
                            "therefore it's not possible to automatically fill out the following `auto` entries "
                            f"""in the DeepSpeed config file: {hidden_size_auto_keys}. You can fix that by replacing """
                            "`auto` values for these keys with an integer value of your choice."				)

                 self.fill_only("zero_optimization.reduce_bucket_size"					,		hidden_size * hidden_size				)
                 if self.is_zeroa():
                        # automatically assign the optimal config values based on model config
                        self.fill_only("zero_optimization.stage3_prefetch_bucket_size"					,		0.9 * hidden_size * hidden_size				)
                        self.fill_only("zero_optimization.stage3_param_persistence_threshold"					,		10 * hidden_size				)

        # scheduler
          self.fill_match("scheduler.params.total_num_steps"					,		snake_case__					,		"num_training_steps (calculated)"				)
          self.fill_match("scheduler.params.warmup_num_steps"					,		args.get_warmup_steps(snake_case__				)					,		"warmup_steps"				)

          if len(self.mismatches				) > 0:
                 snake_case :     Dict           =   "\n".join(self.mismatches				)
                 raise ValueError(
                     "Please correct the following DeepSpeed config values that mismatch TrainingArguments"
                     f""" values:\n{mismatches}\nThe easiest method is to set these DeepSpeed config values to 'auto'."""				)



__lowerCamelCase							= None


def     UpperCamelCase      (					__lowerCamelCase  :				int					):
       # this is a special weakref global object to allow us to get to Deepspeed config from APIs
       # that don't have an easy way to get to the Deepspeed config outside of the Trainer domain.
       global _hf_deepspeed_config_weak_ref
       # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
       snake_case :     Any           =   weakref.ref(__lowerCamelCase					)


def     UpperCamelCase      (					):
       # useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
       global _hf_deepspeed_config_weak_ref
       snake_case :     Any           =   None


def     UpperCamelCase      (					):
       if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
              return _hf_deepspeed_config_weak_ref().is_zeroa()
       else:
              return False


def     UpperCamelCase      (					):
       if _hf_deepspeed_config_weak_ref is not None and _hf_deepspeed_config_weak_ref() is not None:
              return _hf_deepspeed_config_weak_ref().config
       else:
              return None


def     UpperCamelCase      (					__lowerCamelCase  :				str       ,      __lowerCamelCase  :				str       ,      __lowerCamelCase  :				int       ,      __lowerCamelCase  :				Union[str, Any]       ,      __lowerCamelCase  :				int					):
       from accelerate.utils import DummyOptim, DummyScheduler

       snake_case :     Union[str, Any]           =   hf_deepspeed_config.config

       # Optimizer + Scheduler
       # Currently supported combos:
       # 1. DS scheduler + DS optimizer: Yes
       # 2. HF scheduler + HF optimizer: Yes
       # 3. DS scheduler + HF optimizer: Yes
       # 4. HF scheduler + DS optimizer: No
       #
       # Unless Offload is enabled in which case it's:
       # 1. DS scheduler + DS optimizer: Yes
       # 2. HF scheduler + HF optimizer: Mostly*
       # 3. DS scheduler + HF optimizer: Mostly*
       # 4. HF scheduler + DS optimizer: No
       #
       # Mostly*: All non-native DeepSpeed optimizers that have both CPU and GPU implementation should work (except LAMB)

       snake_case :     Optional[int]           =   None
       if "optimizer" in config:
              if args.adafactor:
                     raise ValueError(
                         "--adafactor was passed, but also found `optimizer` configured in the DeepSpeed config. "
                         "Only one optimizer can be configured."					)
              snake_case :     Union[str, Any]           =   DummyOptim(params=__lowerCamelCase					)
       else:
              if hf_deepspeed_config.is_offload():
                     logger.info(
                         "Detected ZeRO Offload and non-DeepSpeed optimizers: This combination should work as long as the"
                         " custom optimizer has both CPU and GPU implementation (except LAMB)"					)

              # ds supports Adam, OneBitAdam, and Lamb optimizers and can import other optimizers from torch.
              # But trainer uses AdamW by default.
              snake_case :     int           =   trainer.create_optimizer()
              # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
              snake_case :     Union[str, Any]           =   True

       snake_case :     List[Any]           =   None
       if "scheduler" in config:
              snake_case :     int           =   DummyScheduler(__lowerCamelCase					)
       else:
              if isinstance(__lowerCamelCase       ,      __lowerCamelCase					):
                     raise ValueError(
                         "Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. "
                         "Please configure a scheduler in the DeepSpeed config."					)
              snake_case :     Union[str, Any]           =   trainer.create_scheduler(num_training_steps=__lowerCamelCase       ,      optimizer=__lowerCamelCase					)

       return optimizer, lr_scheduler


def     UpperCamelCase      (					__lowerCamelCase  :				List[Any]       ,      __lowerCamelCase  :				Union[str, Any]       ,      __lowerCamelCase  :				Tuple=False					):
       from deepspeed.utils import logger as ds_logger

       snake_case :     Optional[Any]           =   trainer.model
       snake_case :     List[str]           =   trainer.args

       snake_case :     Any           =   trainer.accelerator.state.deepspeed_plugin.hf_ds_config

       # resume config update - some bits like `model` and `num_training_steps` only become available during train
       hf_deepspeed_config.trainer_config_finalize(__lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase					)

       # set the Deepspeed log level consistent with the Trainer
       ds_logger.setLevel(args.get_process_log_level()					)

       if inference:
              # only Z3 makes sense for the inference
              if not hf_deepspeed_config.is_zeroa():
                     raise ValueError("ZeRO inference only makes sense with ZeRO Stage 3 - please adjust your config"					)

              # in case the training config is re-used for inference
              hf_deepspeed_config.del_config_sub_tree("optimizer"					)
              hf_deepspeed_config.del_config_sub_tree("lr_scheduler"					)
              snake_case       ,      snake_case :     List[Any]           =   None, None
              snake_case :     List[str]           =   None
       else:
              snake_case :     List[str]           =   None  # important for when deepspeed_init is used as re-init
              snake_case :     List[Any]           =   list(filter(lambda __lowerCamelCase					: p.requires_grad       ,      model.parameters()					)					)
              snake_case       ,      snake_case :     Any           =   deepspeed_optim_sched(
                  __lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase       ,      __lowerCamelCase					)

       # keep for quick debug:
       # from pprint import pprint; pprint(config)

       return optimizer, lr_scheduler




def     UpperCamelCase      (					__lowerCamelCase  :				str       ,      __lowerCamelCase  :				str					):
       # it's possible that the user is trying to resume from model_path, which doesn't necessarily
       # contain a deepspeed checkpoint. e.g. examples just check if the dir exists and assume it's
       # a resume from a checkpoint and not just a local pretrained weight. So we check here if the
       # path contains what looks like a deepspeed checkpoint
       import glob

       snake_case :     int           =   sorted(glob.glob(f"""{checkpoint_path}/global_step*"""					)					)

       if len(__lowerCamelCase					) > 0:
              logger.info(f"""Attempting to resume from {checkpoint_path}"""					)
              # this magically updates self.optimizer and self.lr_scheduler
              snake_case       ,      snake_case :     Dict           =   deepspeed_engine.load_checkpoint(
                  __lowerCamelCase       ,      load_optimizer_states=__lowerCamelCase       ,      load_lr_scheduler_states=__lowerCamelCase					)
              if load_path is None:
                     raise ValueError(f"""[deepspeed] failed to resume from checkpoint {checkpoint_path}"""					)
       else:
              raise ValueError(f"""Can't find a valid checkpoint at {checkpoint_path}"""					)







import platform
from argparse import ArgumentParser

import huggingface_hub
import pandas
import pyarrow

from datasets import __version__ as version
from datasets.commands import BaseDatasetsCLICommand


def     UpperCamelCase      (					__lowerCamelCase  :				int					):
       return EnvironmentCommand()






class 						UpperCAmelCase	(		A_				):



   @staticmethod
   def        _SCREAMING_SNAKE_CASE     (snake_case__	:      ArgumentParser				)	->			Any:


          '''simple docstring'''


          snake_case :     Optional[Any]           =   parser.add_parser("env"					,		help="Print relevant system environment info."				)
          download_parser.set_defaults(func=snake_case__				)



   def        _SCREAMING_SNAKE_CASE     (self	:      Dict				)	->			Optional[Any]:


          '''simple docstring'''


          snake_case :     List[str]           =   {
              "`datasets` version": version,
              "Platform": platform.platform(),
              "Python version": platform.python_version(),
              "Huggingface_hub version": huggingface_hub.__version__,
              "PyArrow version": pyarrow.__version__,
              "Pandas version": pandas.__version__,
          }

          print("\nCopy-and-paste the text below in your GitHub issue.\n"				)
          print(self.format_dict(snake_case__				)				)

          return info



   @staticmethod
   def        _SCREAMING_SNAKE_CASE     (snake_case__	:      Tuple				)	->			Any:


          '''simple docstring'''


          return "\n".join([f"""- {prop}: {val}""" for prop, val in d.items()]				) + "\n"





