






'''simple docstring'''




from copy import deepcopy
from unittest.case import TestCase

import pytest

from datasets.arrow_dataset import Dataset
from datasets.features import Audio, ClassLabel, Features, Image, Sequence, Value
from datasets.info import DatasetInfo
from datasets.tasks import (
    AudioClassification,
    AutomaticSpeechRecognition,
    ImageClassification,
    LanguageModeling,
    QuestionAnsweringExtractive,
    Summarization,
    TextClassification,
    task_template_from_dict,
)
from datasets.utils.py_utils import asdict


lowerCamelCase_							      =      {
    '''id''': '''5733be284776f41900661182''',
    '''title''': '''University_of_Notre_Dame''',
    '''context''': '''Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.''',
    '''question''': '''To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?''',
    '''answers''': {'''text''': ['''Saint Bernadette Soubirous'''], '''answer_start''': [5_15]},
}

@pytest.mark.parametrize(
    "task_cls"     ,						[
        AudioClassification,
        AutomaticSpeechRecognition,
        ImageClassification,
        LanguageModeling,
        QuestionAnsweringExtractive,
        Summarization,
        TextClassification,
    ]     ,						)
def       __lowercase		(    __lowercase ) ->		Tuple:

      '''simple docstring'''


      _A      =			task_cls()
      _A      =			asdict(__lowercase )
      _A      =			task_template_from_dict(__lowercase )
      assert task == reloaded

class      _UpperCAmelCase    :




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			LanguageModeling(text_column="input_text"	)
         assert {"input_text": "text"} == task.column_mapping

   def 		lowerCAmelCase				(	self				:			Optional[int]	):





         '''simple docstring'''





         _A      =			Features({"text": Value("string"	)}	)
         _A      =			{"text_column": "input_text"}
         _A      =			LanguageModeling.from_dict(__UpperCAmelCase	)
         assert "language-modeling" == task.task
         assert input_schema == task.input_schema





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			sorted(["pos", "neg"]	)

   def 		lowerCAmelCase				(	self				:			Optional[Any]	):





         '''simple docstring'''





         _A      =			TextClassification(text_column="input_text"   ,   label_column="input_label"	)
         self.assertDictEqual({"input_text": "text", "input_label": "labels"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			Features({"text": Value("string"	)}	)
         # Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
         _A      =			Features({"labels": ClassLabel}	)
         _A      =			{"text_column": "input_text", "label_column": "input_labels"}
         _A      =			TextClassification.from_dict(__UpperCAmelCase	)
         self.assertEqual("text-classification"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)

   def 		lowerCAmelCase				(	self				:			Tuple	):





         '''simple docstring'''





         _A      =			TextClassification(text_column="input_text"   ,   label_column="input_label"	)
         self.assertEqual(task.label_schema["labels"]   ,   __UpperCAmelCase	)
         _A      =			task.align_with_features(Features({"input_label": ClassLabel(names=self.labels	)}	)	)
         self.assertEqual(task.label_schema["labels"]   ,   ClassLabel(names=self.labels	)	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			QuestionAnsweringExtractive(
             context_column="input_context"   ,   question_column="input_question"   ,   answers_column="input_answers"	)
         self.assertDictEqual(
             {"input_context": "context", "input_question": "question", "input_answers": "answers"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			Any	):





         '''simple docstring'''





         _A      =			Features({"question": Value("string"	), "context": Value("string"	)}	)
         _A      =			Features(
             {
                 "answers": Sequence(
                     {
                         "text": Value("string"	),
                         "answer_start": Value("int32"	),
                     }	)
             }	)
         _A      =			{
             "context_column": "input_input_context",
             "question_column": "input_question",
             "answers_column": "input_answers",
         }
         _A      =			QuestionAnsweringExtractive.from_dict(__UpperCAmelCase	)
         self.assertEqual("question-answering-extractive"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			Summarization(text_column="input_text"   ,   summary_column="input_summary"	)
         self.assertDictEqual({"input_text": "text", "input_summary": "summary"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			Features({"text": Value("string"	)}	)
         _A      =			Features({"summary": Value("string"	)}	)
         _A      =			{"text_column": "input_text", "summary_column": "input_summary"}
         _A      =			Summarization.from_dict(__UpperCAmelCase	)
         self.assertEqual("summarization"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			AutomaticSpeechRecognition(audio_column="input_audio"   ,   transcription_column="input_transcription"	)
         self.assertDictEqual({"input_audio": "audio", "input_transcription": "transcription"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			Tuple	):





         '''simple docstring'''





         _A      =			Features({"audio": Audio()}	)
         _A      =			Features({"transcription": Value("string"	)}	)
         _A      =			{
             "audio_column": "input_audio",
             "transcription_column": "input_transcription",
         }
         _A      =			AutomaticSpeechRecognition.from_dict(__UpperCAmelCase	)
         self.assertEqual("automatic-speech-recognition"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			List[Any]	):





         '''simple docstring'''





         _A      =			sorted(["pos", "neg"]	)

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			AudioClassification(audio_column="input_audio"   ,   label_column="input_label"	)
         self.assertDictEqual({"input_audio": "audio", "input_label": "labels"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			List[Any]	):





         '''simple docstring'''





         _A      =			Features({"audio": Audio()}	)
         _A      =			Features({"labels": ClassLabel}	)
         _A      =			{
             "audio_column": "input_image",
             "label_column": "input_label",
         }
         _A      =			AudioClassification.from_dict(__UpperCAmelCase	)
         self.assertEqual("audio-classification"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)

   def 		lowerCAmelCase				(	self				:			str	):





         '''simple docstring'''





         _A      =			AudioClassification(audio_column="input_audio"   ,   label_column="input_label"	)
         self.assertEqual(task.label_schema["labels"]   ,   __UpperCAmelCase	)
         _A      =			task.align_with_features(Features({"input_label": ClassLabel(names=self.labels	)}	)	)
         self.assertEqual(task.label_schema["labels"]   ,   ClassLabel(names=self.labels	)	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			List[str]	):





         '''simple docstring'''





         _A      =			sorted(["pos", "neg"]	)

   def 		lowerCAmelCase				(	self				:			Union[str, Any]	):





         '''simple docstring'''





         _A      =			ImageClassification(image_column="input_image"   ,   label_column="input_label"	)
         self.assertDictEqual({"input_image": "image", "input_label": "labels"}   ,   task.column_mapping	)

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			Features({"image": Image()}	)
         _A      =			Features({"labels": ClassLabel}	)
         _A      =			{
             "image_column": "input_image",
             "label_column": "input_label",
         }
         _A      =			ImageClassification.from_dict(__UpperCAmelCase	)
         self.assertEqual("image-classification"   ,   task.task	)
         self.assertEqual(__UpperCAmelCase   ,   task.input_schema	)
         self.assertEqual(__UpperCAmelCase   ,   task.label_schema	)

   def 		lowerCAmelCase				(	self				:			Union[str, Any]	):





         '''simple docstring'''





         _A      =			ImageClassification(image_column="input_image"   ,   label_column="input_label"	)
         self.assertEqual(task.label_schema["labels"]   ,   __UpperCAmelCase	)
         _A      =			task.align_with_features(Features({"input_label": ClassLabel(names=self.labels	)}	)	)
         self.assertEqual(task.label_schema["labels"]   ,   ClassLabel(names=self.labels	)	)





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def 		lowerCAmelCase				(	self				:			Dict	):





         '''simple docstring'''





         _A      =			DatasetInfo(task_templates=QuestionAnsweringExtractive()	)
         _A      =			Dataset.from_dict({k: [v] for k, v in SAMPLE_QUESTION_ANSWERING_EXTRACTIVE.items()}   ,   info=__UpperCAmelCase	)
         assert isinstance(dataset.info.task_templates   ,   __UpperCAmelCase	)
         assert len(dataset.info.task_templates	) == 1

         def keep_task(__UpperCAmelCase				:			Optional[Any]	):
               return x

         def dont_keep_task(__UpperCAmelCase				:			Tuple	):
               _A      =			deepcopy(__UpperCAmelCase	)
               _A      =			0
               return out

         _A      =			dataset.map(__UpperCAmelCase	)
         assert mapped_dataset.info.task_templates == dataset.info.task_templates
         # reload from cache
         _A      =			dataset.map(__UpperCAmelCase	)
         assert mapped_dataset.info.task_templates == dataset.info.task_templates

         _A      =			dataset.map(__UpperCAmelCase	)
         assert mapped_dataset.info.task_templates == []
         # reload from cache
         _A      =			dataset.map(__UpperCAmelCase	)
         assert mapped_dataset.info.task_templates == []

   def 		lowerCAmelCase				(	self				:			Optional[Any]	):





         '''simple docstring'''





         _A      =			Features({"text": Value("string"	), "label": ClassLabel(names=("pos", "neg")	)}	)
         _A      =			TextClassification(text_column="text"   ,   label_column="label"	)
         _A      =			DatasetInfo(features=__UpperCAmelCase   ,   task_templates=__UpperCAmelCase	)
         _A      =			Dataset.from_dict({"text": ["A sentence."], "label": ["pos"]}   ,   info=__UpperCAmelCase	)

         def process(__UpperCAmelCase				:			Any	):
               return example

         _A      =			dataset.remove_columns("label"	)
         _A      =			modified_dataset.map(__UpperCAmelCase	)
         assert mapped_dataset.info.task_templates == []













'''simple docstring'''




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       __lowercase		(    ) ->		Optional[Any]:

      '''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      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def __init__(	self				:			Dict   ,   __UpperCAmelCase				:			Tuple	):





         '''simple docstring'''





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





class      _UpperCAmelCase    (  snake_case_		):




   """simple docstring"""

   def __init__(	self				:			List[str]   ,   __UpperCAmelCase				:			List[str]	):





         '''simple docstring'''





         super().__init__(__UpperCAmelCase	)
         _A      =			None
         _A      =			[]

   def 		lowerCAmelCase				(	self				:			Optional[int]	):





         '''simple docstring'''





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

   def 		lowerCAmelCase				(	self				:			Dict   ,   __UpperCAmelCase				:			Dict	):





         '''simple docstring'''





         _A      =			self.get_value(__UpperCAmelCase	)
         if val is None:
               return False
         else:
               return val == "auto"

   def 		lowerCAmelCase				(	self				:			List[Any]   ,   __UpperCAmelCase				:			int   ,   __UpperCAmelCase				:			Optional[int]   ,   __UpperCAmelCase				:			Any=None   ,   __UpperCAmelCase				:			Union[str, Any]=True	):





         '''simple docstring'''





         _A  ,			_A      =			self.find_config_node(__UpperCAmelCase	)
         if config is None:
               return

         if config.get(__UpperCAmelCase	) == "auto":
               _A      =			hf_val
               return

         if not must_match:
               return

         _A      =			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}'''	)

   snake_case					       =		partialmethod(snake_case_							,					must_match=snake_case_		)

   def 		lowerCAmelCase				(	self				:			str   ,   __UpperCAmelCase				:			Any	):





         '''simple docstring'''





         _A      =			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:
               _A      =			"apex" if args.fpaa_backend == "apex" else "amp"
         else:
               _A      =			None

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

   def 		lowerCAmelCase				(	self				:			Optional[int]   ,   __UpperCAmelCase				:			Dict   ,   __UpperCAmelCase				:			str   ,   __UpperCAmelCase				:			Tuple	):





         '''simple docstring'''





         _A      =			[
             "zero_optimization.reduce_bucket_size",
             "zero_optimization.stage3_prefetch_bucket_size",
             "zero_optimization.stage3_param_persistence_threshold",
         ]
         _A      =			[x for x in hidden_size_based_keys if self.is_auto(__UpperCAmelCase	)]

         if len(__UpperCAmelCase	) > 0:
               if hasattr(model.config   ,   "hidden_size"	):
                     _A      =			model.config.hidden_size
               elif hasattr(model.config   ,   "hidden_sizes"	):
                     # if there are many hidden sizes pick the largest one
                     _A      =			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:
               _A      =			"\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       __lowercase		(    __lowercase ) ->		Optional[Any]:

      '''simple docstring'''


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

def       __lowercase		(    ) ->		str:

      '''simple docstring'''


      global _hf_deepspeed_config_weak_ref
      _A      =			None

def       __lowercase		(    ) ->		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().is_zeroa()
      else:
            return False

def       __lowercase		(    ) ->		str:

      '''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       __lowercase		(    __lowercase     ,						__lowercase     ,						__lowercase     ,						__lowercase     ,						__lowercase ) ->		int:

      '''simple docstring'''


      from accelerate.utils import DummyOptim, DummyScheduler

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

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

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

      return optimizer, lr_scheduler

def       __lowercase		(    __lowercase     ,						__lowercase     ,						__lowercase=False ) ->		Tuple:

      '''simple docstring'''


      from deepspeed.utils import logger as ds_logger

      _A      =			trainer.model
      _A      =			trainer.args

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

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

      return optimizer, lr_scheduler


def       __lowercase		(    __lowercase     ,						__lowercase ) ->		Union[str, Any]:

      '''simple docstring'''


      import glob

      _A      =			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
            _A  ,			_A      =			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}''' )





