


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

snake_case__  				=			logging.get_logger(__name__)
def 			__SCREAMING_SNAKE_CASE							(					)							->	List[Any]:
						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      snake_case__	(   _A	):
				def __init__(	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										set_hf_deepspeed_config(self   )
										dep_version_check('accelerate'   )
										dep_version_check('deepspeed'   )
										super().__init__(_UpperCAmelCase   )






class      snake_case__	(   _A	):
				def __init__(	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										super().__init__(_UpperCAmelCase   )
										__magic_name__					:				Dict				=   None
										__magic_name__					:				Optional[int]				=   []
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										if self._dtype is None:
																raise ValueError('trainer_config_process() wasn\'t called yet to tell dtype'   )
										return self._dtype
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Optional[Any]				=   self.get_value(_UpperCAmelCase   )
										if val is None:
																return False
										else:
																return val == "auto"
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase=None  ,   _UpperCAmelCase=True   ):




										"""simple docstring"""




										__magic_name__		,			__magic_name__					:				Tuple				=   self.find_config_node(_UpperCAmelCase   )
										if config is None:
																return

										if config.get(_UpperCAmelCase   ) == "auto":
																__magic_name__					:				Optional[Any]				=   hf_val
																return

										if not must_match:
																return

										__magic_name__					:				Optional[int]				=   config.get(_UpperCAmelCase   )
										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}"   )

				__lowerCAmelCase		=   partialmethod(_A    ,							must_match=_A	)
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Optional[Any]				=   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'  ,   _UpperCAmelCase  ,   '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:
																__magic_name__					:				List[str]				=   'apex' if args.fpaa_backend == 'apex' else 'amp'
										else:
																__magic_name__					:				Tuple				=   None

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





				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										__magic_name__					:				Optional[int]				=   [
										    'zero_optimization.reduce_bucket_size',
										    'zero_optimization.stage3_prefetch_bucket_size',
										    'zero_optimization.stage3_param_persistence_threshold',
										]
										__magic_name__					:				Tuple				=   [x for x in hidden_size_based_keys if self.is_auto(_UpperCAmelCase   )]

										if len(_UpperCAmelCase   ) > 0:
																if hasattr(model.config  ,   'hidden_size'   ):
																						__magic_name__					:				Optional[int]				=   model.config.hidden_size
																elif hasattr(model.config  ,   'hidden_sizes'   ):
																						# if there are many hidden sizes pick the largest one
																						__magic_name__					:				Optional[Any]				=   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'  ,   _UpperCAmelCase  ,   'num_training_steps (calculated)'   )
										self.fill_match('scheduler.params.warmup_num_steps'  ,   args.get_warmup_steps(_UpperCAmelCase   )  ,   'warmup_steps'   )

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






snake_case__  				=			None
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Tuple					)							->	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)
						__magic_name__					:				int				=   weakref.ref(lowercase__					)
def 			__SCREAMING_SNAKE_CASE							(					)							->	Any:
						# useful for unit tests to ensure the global state doesn't leak - call from `tearDown` method
						global _hf_deepspeed_config_weak_ref
						__magic_name__					:				List[Any]				=   None
def 			__SCREAMING_SNAKE_CASE							(					)							->	str:
						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 			__SCREAMING_SNAKE_CASE							(					)							->	Optional[Any]:
						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 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Union[str, Any]   ,    lowercase__	:Dict   ,    lowercase__	:Union[str, Any]   ,    lowercase__	:Optional[int]   ,    lowercase__	:str					)							->	List[Any]:
						from accelerate.utils import DummyOptim, DummyScheduler

						__magic_name__					:				Dict				=   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)

						__magic_name__					:				Any				=   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.'					)
												__magic_name__					:				Dict				=   DummyOptim(params=lowercase__					)
						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.
												__magic_name__					:				str				=   trainer.create_optimizer()
												# To use other optimizers requires voiding warranty with: `zero_allow_untested_optimizer`
												__magic_name__					:				List[Any]				=   True

						__magic_name__					:				Dict				=   None
						if "scheduler" in config:
												__magic_name__					:				Tuple				=   DummyScheduler(lowercase__					)
						else:
												if isinstance(lowercase__   ,    lowercase__					):
																		raise ValueError(
																		    'Found `optimizer` configured in the DeepSpeed config, but no `scheduler`. '
																		    'Please configure a scheduler in the DeepSpeed config.'					)
												__magic_name__					:				Optional[int]				=   trainer.create_scheduler(num_training_steps=lowercase__   ,    optimizer=lowercase__					)

						return optimizer, lr_scheduler
def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Optional[int]   ,    lowercase__	:str   ,    lowercase__	:Any=False					)							->	Tuple:
						from deepspeed.utils import logger as ds_logger

						__magic_name__					:				Dict				=   trainer.model
						__magic_name__					:				Optional[Any]				=   trainer.args

						__magic_name__					:				Dict				=   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(lowercase__   ,    lowercase__   ,    lowercase__					)

						# 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'					)
												__magic_name__		,			__magic_name__					:				Union[str, Any]				=   None, None
												__magic_name__					:				str				=   None
						else:
												__magic_name__					:				int				=   None  # important for when deepspeed_init is used as re-init
												__magic_name__					:				Optional[int]				=   list(filter(lambda lowercase__					: p.requires_grad   ,    model.parameters()					)					)
												__magic_name__		,			__magic_name__					:				List[str]				=   deepspeed_optim_sched(
												    lowercase__   ,    lowercase__   ,    lowercase__   ,    lowercase__   ,    lowercase__					)

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

						return optimizer, lr_scheduler

def 			__SCREAMING_SNAKE_CASE							(					lowercase__	:Dict   ,    lowercase__	:Optional[int]					)							->	List[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

						__magic_name__					:				List[Any]				=   sorted(glob.glob(f"{checkpoint_path}/global_step*"					)					)

						if len(lowercase__					) > 0:
												logger.info(f"Attempting to resume from {checkpoint_path}"					)
												# this magically updates self.optimizer and self.lr_scheduler
												__magic_name__		,			__magic_name__					:				str				=   deepspeed_engine.load_checkpoint(
												    lowercase__   ,    load_optimizer_states=lowercase__   ,    load_lr_scheduler_states=lowercase__					)
												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


snake_case__  				=			datasets.utils.logging.get_logger(__name__)



@dataclass
class      snake_case__	(   datasets.BuilderConfig	):
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   True
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   None
				__lowerCAmelCase		=   10_000
				__lowerCAmelCase		=   None
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										if self.sql is None:
																raise ValueError('sql must be specified'   )
										if self.con is None:
																raise ValueError('con must be specified'   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase  ,   _UpperCAmelCase = None  ,   ):




										"""simple docstring"""




										__magic_name__					:				Optional[int]				=   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
										__magic_name__					:				Union[str, Any]				=   config_kwargs['sql']
										if not isinstance(_UpperCAmelCase  ,   _UpperCAmelCase   ):
																if datasets.config.SQLALCHEMY_AVAILABLE and "sqlalchemy" in sys.modules:
																						import sqlalchemy

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

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





				@property
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				str				=   {
										    '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      snake_case__	(   datasets.ArrowBasedBuilder	):
				__lowerCAmelCase		=   SqlConfig
				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										return datasets.DatasetInfo(features=self.config.features   )
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




										return [datasets.SplitGenerator(name=datasets.Split.TRAIN  ,   gen_kwargs={}   )]
				def 			UpperCamelCase_    (	self  ,   _UpperCAmelCase   ):




										"""simple docstring"""




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





				def 			UpperCamelCase_    (	self   ):




										"""simple docstring"""




										__magic_name__					:				int				=   self.config.chunksize
										__magic_name__					:				Tuple				=   pd.read_sql(
										    self.config.sql  ,   self.config.con  ,   chunksize=_UpperCAmelCase  ,   **self.config.pd_read_sql_kwargs   )
										__magic_name__					:				Dict				=   [sql_reader] if chunksize is None else sql_reader
										for chunk_idx, df in enumerate(_UpperCAmelCase   ):
																__magic_name__					:				Optional[Any]				=   pa.Table.from_pandas(_UpperCAmelCase   )
																yield chunk_idx, self._cast_table(_UpperCAmelCase   )


