




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

__SCREAMING_SNAKE_CASE			:      List[str]      			=logging.get_logger(__name__)

def 	lowerCAmelCase_	(     )    ->	List[str]:


      '''simple docstring'''




      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 	A	(						_lowerCAmelCase						):
     """simple docstring"""

     def __init__(			self     :      Union[str, Any]		,     A__     :      Optional[int]	) ->				Union[str, Any]:
           # set global weakref object
           set_hf_deepspeed_config(self	)
           dep_version_check('accelerate'	)
           dep_version_check('deepspeed'	)
           super().__init__(A__	)






class 	A	(						_lowerCAmelCase						):
     """simple docstring"""

     def __init__(			self     :      List[str]		,     A__     :      Union[str, Any]	) ->				Dict:
           super().__init__(A__	)
           lowerCamelCase__					     = None
           lowerCamelCase__					     = []




     def        __snake_case      (			self     :      List[str]	) ->				Optional[Any]:
           if self._dtype is None:
                 raise ValueError('trainer_config_process() wasn\'t called yet to tell dtype'	)
           return self._dtype




     def        __snake_case      (			self     :      Union[str, Any]		,     A__     :      Tuple	) ->				Optional[Any]:
           lowerCamelCase__					     = self.get_value(A__	)
           if val is None:
                 return False
           else:
                 return val == "auto"




     def        __snake_case      (			self     :      Union[str, Any]		,     A__     :      Tuple		,     A__     :      List[Any]		,     A__     :      List[Any]=None		,     A__     :      Union[str, Any]=True	) ->				Tuple:
           lowerCamelCase__     ,						lowerCamelCase__					     = self.find_config_node(A__	)
           if config is None:
                 return

           if config.get(A__	) == "auto":
                 lowerCamelCase__					     = hf_val
                 return

           if not must_match:
                 return

           lowerCamelCase__					     = config.get(A__	)
           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}"""	)

     _snake_case:		List[str]   =partialmethod(_lowerCAmelCase	,  must_match=_lowerCAmelCase						)




     def        __snake_case      (			self     :      List[str]		,     A__     :      List[str]	) ->				Union[str, Any]:
           lowerCamelCase__					     = 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'		,     A__		,     '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:
                 lowerCamelCase__					     = 'apex' if args.fpaa_backend == 'apex' else 'amp'
           else:
                 lowerCamelCase__					     = None

           if args.save_on_each_node:
                 # deepspeed uses shared storage by default. Let's override this setting if save_on_each_node == True
                 lowerCamelCase__					     = self.config.get('checkpoint'		,     {}	)
                 lowerCamelCase__					     = 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'	):
                 lowerCamelCase__					     = torch.bfloataa
           elif self.is_false('fp16.enabled'	):
                 lowerCamelCase__					     = torch.floataa
           else:
                 lowerCamelCase__					     = torch.floataa




     def        __snake_case      (			self     :      Dict		,     A__     :      Union[str, Any]		,     A__     :      Optional[int]		,     A__     :      List[Any]	) ->				str:
           lowerCamelCase__					     = [
               'zero_optimization.reduce_bucket_size',
               'zero_optimization.stage3_prefetch_bucket_size',
               'zero_optimization.stage3_param_persistence_threshold',
           ]
           lowerCamelCase__					     = [x for x in hidden_size_based_keys if self.is_auto(A__	)]

           if len(A__	) > 0:
                 if hasattr(model.config		,     'hidden_size'	):
                       lowerCamelCase__					     = model.config.hidden_size
                 elif hasattr(model.config		,     'hidden_sizes'	):
                       # if there are many hidden sizes pick the largest one
                       lowerCamelCase__					     = 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'		,     1_0 * hidden_size	)

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

           if len(self.mismatches	) > 0:
                 lowerCamelCase__					     = '\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'."""	)
__SCREAMING_SNAKE_CASE			:      List[str]      			=None

def 	lowerCAmelCase_	(     lowerCamelCase:				Any  )    ->	Dict:


      '''simple docstring'''




      global _hf_deepspeed_config_weak_ref
      # will go away automatically when HfDeepSpeedConfig is destroyed (when TrainingArguments is destroyed)
      lowerCamelCase__					     = weakref.ref(lowerCamelCase  )

def 	lowerCAmelCase_	(     )    ->	Dict:


      '''simple docstring'''




      global _hf_deepspeed_config_weak_ref
      lowerCamelCase__					     = None

def 	lowerCAmelCase_	(     )    ->	Dict:


      '''simple docstring'''




      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 	lowerCAmelCase_	(     )    ->	Optional[int]:


      '''simple docstring'''




      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 	lowerCAmelCase_	(     lowerCamelCase:				Optional[int]   , lowerCamelCase:				Any   , lowerCamelCase:				Optional[Any]   , lowerCamelCase:				str   , lowerCamelCase:				Any  )    ->	str:


      '''simple docstring'''




      from accelerate.utils import DummyOptim, DummyScheduler

      lowerCamelCase__					     = 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)

      lowerCamelCase__					     = 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.'  )
            lowerCamelCase__					     = 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.
            lowerCamelCase__					     = trainer.create_optimizer()
            # To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
            lowerCamelCase__					     = True

      lowerCamelCase__					     = None
      if "scheduler" in config:
            lowerCamelCase__					     = 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.'  )
            lowerCamelCase__					     = trainer.create_scheduler(num_training_steps=lowerCamelCase   , optimizer=lowerCamelCase  )

      return optimizer, lr_scheduler

def 	lowerCAmelCase_	(     lowerCamelCase:				int   , lowerCamelCase:				str   , lowerCamelCase:				Dict=False  )    ->	Tuple:


      '''simple docstring'''




      from deepspeed.utils import logger as ds_logger

      lowerCamelCase__					     = trainer.model
      lowerCamelCase__					     = trainer.args

      lowerCamelCase__					     = 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'  )
            lowerCamelCase__     ,						lowerCamelCase__					     = None, None
            lowerCamelCase__					     = None
      else:
            lowerCamelCase__					     = None  # important for when deepspeed_init is used as re-init
            lowerCamelCase__					     = list(filter(lambda lowerCamelCase  : p.requires_grad   , model.parameters()  )  )
            lowerCamelCase__     ,						lowerCamelCase__					     = deepspeed_optim_sched(
                lowerCamelCase   , lowerCamelCase   , lowerCamelCase   , lowerCamelCase   , lowerCamelCase  )

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

      return optimizer, lr_scheduler


def 	lowerCAmelCase_	(     lowerCamelCase:				Dict   , lowerCamelCase:				Dict  )    ->	str:


      '''simple docstring'''




      import glob

      lowerCamelCase__					     = 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
            lowerCamelCase__     ,						lowerCamelCase__					     = 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 sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union

import pandas as pd
import pyarrow as pa

import datasets
import datasets.config
from datasets.features.features import require_storage_cast
from datasets.table import table_cast


if TYPE_CHECKING:
      import sqlitea

      import sqlalchemy


__SCREAMING_SNAKE_CASE			:      Union[str, Any]      			=datasets.utils.logging.get_logger(__name__)






@dataclass
class 	A	(						datasets.BuilderConfig						):
     """simple docstring"""

     _snake_case:		Union[str, "sqlalchemy.sql.Selectable"]   =None
     _snake_case:		Union[str, "sqlalchemy.engine.Connection", "sqlalchemy.engine.Engine", "sqlite3.Connection"]   =None
     _snake_case:		Optional[Union[str, List[str]]]   =None
     _snake_case:		bool   =True
     _snake_case:		Optional[Union[List, Tuple, Dict]]   =None
     _snake_case:		Optional[Union[List, Dict]]   =None
     _snake_case:		Optional[List[str]]   =None
     _snake_case:		Optional[int]   =10_000
     _snake_case:		Optional[datasets.Features]   =None




     def        __snake_case      (			self     :      Optional[Any]	) ->				Union[str, Any]:
           if self.sql is None:
                 raise ValueError('sql must be specified'	)
           if self.con is None:
                 raise ValueError('con must be specified'	)




     def        __snake_case      (			self     :      List[Any]		,     A__     :      dict		,     A__     :      Optional[datasets.Features] = None		,     ) ->				str:
           lowerCamelCase__					     = config_kwargs.copy()
           # We need to stringify the Selectable object to make its hash deterministic

           # The process of stringifying is explained here: http://docs.sqlalchemy.org/en/latest/faq/sqlexpressions.html
           lowerCamelCase__					     = config_kwargs['sql']
           if not isinstance(A__		,     A__	):
                 if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
                       import sqlalchemy

                       if isinstance(A__		,     sqlalchemy.sql.Selectable	):
                             lowerCamelCase__					     = sqlalchemy.create_engine(config_kwargs['con'].split('://'	)[0] + '://'	)
                             lowerCamelCase__					     = str(sql.compile(dialect=engine.dialect	)	)
                             lowerCamelCase__					     = sql_str
                       else:
                             raise TypeError(
                                 F"""Supported types for 'sql' are string and sqlalchemy.sql.Selectable but got {type(A__	)}: {sql}"""	)
                 else:
                       raise TypeError(
                           F"""Supported types for 'sql' are string and sqlalchemy.sql.Selectable but got {type(A__	)}: {sql}"""	)
           lowerCamelCase__					     = config_kwargs['con']
           if not isinstance(A__		,     A__	):
                 lowerCamelCase__					     = id(A__	)
                 logger.info(
                     F"""SQL connection 'con' of type {type(A__	)} couldn't be hashed properly. To enable hashing, specify 'con' as URI string instead."""	)

           return super().create_config_id(A__		,     custom_features=A__	)




     @property
     def        __snake_case      (			self     :      Any	) ->				Union[str, Any]:
           lowerCamelCase__					     = {
               'index_col': self.index_col,
               'columns': self.columns,
               'params': self.params,
               'coerce_float': self.coerce_float,
               'parse_dates': self.parse_dates,
           }
           return pd_read_sql_kwargs






class 	A	(						datasets.ArrowBasedBuilder						):
     """simple docstring"""

     _snake_case:		List[str]   =SqlConfig




     def        __snake_case      (			self     :      Tuple	) ->				Dict:
           return datasets.DatasetInfo(features=self.config.features	)




     def        __snake_case      (			self     :      List[Any]		,     A__     :      Any	) ->				List[Any]:
           return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,     gen_kwargs={}	)]




     def        __snake_case      (			self     :      Union[str, Any]		,     A__     :      pa.Table	) ->				pa.Table:
           if self.config.features is not None:
                 lowerCamelCase__					     = self.config.features.arrow_schema
                 if all(not require_storage_cast(A__	) for feature in self.config.features.values()	):
                       # cheaper cast
                       lowerCamelCase__					     = pa.Table.from_arrays([pa_table[field.name] for field in schema]		,     schema=A__	)
                 else:
                       # more expensive cast; allows str <-> int/float or str to Audio for example
                       lowerCamelCase__					     = table_cast(A__		,     A__	)
           return pa_table




     def        __snake_case      (			self     :      int	) ->				List[Any]:
           lowerCamelCase__					     = self.config.chunksize
           lowerCamelCase__					     = pd.read_sql(
               self.config.sql		,     self.config.con		,     chunksize=A__		,     **self.config.pd_read_sql_kwargs	)
           lowerCamelCase__					     = [sql_reader] if chunksize is None else sql_reader
           for chunk_idx, df in enumerate(A__	):
                 lowerCamelCase__					     = pa.Table.from_pandas(A__	)
                 yield chunk_idx, self._cast_table(A__	)