import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


A__     =				get_logger(__name__)


A__     =				r'\n    Args:\n        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary.\n\n            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and\n            [`PreTrainedTokenizer.__call__`] for details.\n\n            [What are input IDs?](../glossary#input-ids)\n        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):\n            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam\n            search or log softmax for each vocabulary token when using beam search\n        kwargs (`Dict[str, Any]`, *optional*):\n            Additional logits processor specific kwargs.\n\n    Return:\n        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.\n\n'
class 	_lowerCAmelCase		:







							@add_start_docstrings(__snake_case		)
							def __call__(					self  :     Optional[Any]						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray		):
										raise NotImplementedError(
										    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."		)


class 	_lowerCAmelCase		:







							@add_start_docstrings(__snake_case		)
							def __call__(					self  :     Tuple						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray		):
										raise NotImplementedError(
										    F"{self.__class__} is an abstract class. Only classes inheriting this class can be called."		)


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							@add_start_docstrings(__snake_case		)
							def __call__(					self  :     Optional[int]						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int						,       **__snake_case  :     str		):
										for processor in self:
													lowerCamelCase    :Optional[Any]    	=  inspect.signature(processor.__call__		).parameters
													if len(__snake_case		) > 3:
																if not all(arg in kwargs for arg in list(function_args.keys()		)[2:]		):
																			raise ValueError(
																			    F"Make sure that all the required parameters: {list(function_args.keys()		)} for "
																			    F"{processor.__class__} are passed to the logits processor."		)
																lowerCamelCase    :str    	=  processor(__snake_case						,       __snake_case						,       __snake_case						,       **__snake_case		)
													else:
																lowerCamelCase    :Union[str, Any]    	=  processor(__snake_case						,       __snake_case						,       __snake_case		)
										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Any						,       __snake_case  :     float		):
										if not isinstance(__snake_case						,       __snake_case		) or not (temperature > 0):
													raise ValueError(F"`temperature` has to be a strictly positive float, but is {temperature}"		)

										lowerCamelCase    :int    	=  temperature







							def __call__(					self  :     Optional[int]						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :List[str]    	=  scores / self.temperature
										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Tuple						,       __snake_case  :     float						,       __snake_case  :     float = -float('''Inf'''		)						,       __snake_case  :     int = 1		):
										if not isinstance(__snake_case						,       __snake_case		) or (top_p < 0 or top_p > 1.0):
													raise ValueError(F"`top_p` has to be a float > 0 and < 1, but is {top_p}"		)
										if not isinstance(__snake_case						,       __snake_case		) or (min_tokens_to_keep < 1):
													raise ValueError(F"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}"		)

										lowerCamelCase    :str    	=  top_p
										lowerCamelCase    :Any    	=  filter_value
										lowerCamelCase    :Tuple    	=  min_tokens_to_keep







							def __call__(					self  :     str						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :List[Any]    	=  lax.top_k(__snake_case						,       scores.shape[-1]		)

										lowerCamelCase    :Tuple    	=  jnp.full_like(__snake_case						,       self.filter_value		)
										lowerCamelCase    :Dict    	=  jax.nn.softmax(__snake_case						,       axis=-1		).cumsum(axis=-1		)
										lowerCamelCase    :str    	=  cumulative_probs < self.top_p

										# include the token that is higher than top_p as well
										lowerCamelCase    :Union[str, Any]    	=  jnp.roll(__snake_case						,       1		)
										score_mask |= score_mask.at[:, 0].set(__snake_case		)

										# min tokens to keep
										lowerCamelCase    :Union[str, Any]    	=  score_mask.at[:, : self.min_tokens_to_keep].set(__snake_case		)

										lowerCamelCase    :Dict    	=  jnp.where(__snake_case						,       __snake_case						,       __snake_case		)
										lowerCamelCase    :Tuple    	=  jax.lax.sort_key_val(__snake_case						,       __snake_case		)[-1]

										return next_scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Dict						,       __snake_case  :     int						,       __snake_case  :     float = -float('''Inf'''		)						,       __snake_case  :     int = 1		):
										if not isinstance(__snake_case						,       __snake_case		) or top_k <= 0:
													raise ValueError(F"`top_k` has to be a strictly positive integer, but is {top_k}"		)

										lowerCamelCase    :Union[str, Any]    	=  max(__snake_case						,       __snake_case		)
										lowerCamelCase    :List[Any]    	=  filter_value







							def __call__(					self  :     Any						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :Dict    	=  scores.shape
										lowerCamelCase    :Any    	=  jnp.full(batch_size * vocab_size						,       self.filter_value		)

										lowerCamelCase    :Tuple    	=  min(self.top_k						,       scores.shape[-1]		)  # Safety check
										lowerCamelCase    :List[Any]    	=  lax.top_k(__snake_case						,       __snake_case		)
										lowerCamelCase    :int    	=  jnp.broadcast_to((jnp.arange(__snake_case		) * vocab_size)[:, None]						,       (batch_size, topk)		).flatten()
										lowerCamelCase    :Any    	=  topk_scores.flatten()
										lowerCamelCase    :str    	=  topk_indices.flatten() + shift

										lowerCamelCase    :Tuple    	=  next_scores_flat.at[topk_indices_flat].set(__snake_case		)
										lowerCamelCase    :Optional[int]    	=  next_scores_flat.reshape(__snake_case						,       __snake_case		)
										return next_scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Any						,       __snake_case  :     int		):
										lowerCamelCase    :Tuple    	=  bos_token_id







							def __call__(					self  :     Dict						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :Optional[int]    	=  jnp.full(scores.shape						,       -float('''inf'''		)		)

										lowerCamelCase    :int    	=  1 - jnp.bool_(cur_len - 1		)

										lowerCamelCase    :Any    	=  jnp.where(__snake_case						,       new_scores.at[:, self.bos_token_id].set(0		)						,       __snake_case		)

										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Union[str, Any]						,       __snake_case  :     int						,       __snake_case  :     int		):
										lowerCamelCase    :Optional[int]    	=  max_length
										lowerCamelCase    :Optional[Any]    	=  eos_token_id







							def __call__(					self  :     List[Any]						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :Dict    	=  jnp.full(scores.shape						,       -float('''inf'''		)		)

										lowerCamelCase    :Dict    	=  1 - jnp.bool_(cur_len - self.max_length + 1		)

										lowerCamelCase    :Tuple    	=  jnp.where(__snake_case						,       new_scores.at[:, self.eos_token_id].set(0		)						,       __snake_case		)

										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     int						,       __snake_case  :     int						,       __snake_case  :     int		):
										if not isinstance(__snake_case						,       __snake_case		) or min_length < 0:
													raise ValueError(F"`min_length` has to be a positive integer, but is {min_length}"		)

										if not isinstance(__snake_case						,       __snake_case		) or eos_token_id < 0:
													raise ValueError(F"`eos_token_id` has to be a positive integer, but is {eos_token_id}"		)

										lowerCamelCase    :str    	=  min_length
										lowerCamelCase    :Optional[Any]    	=  eos_token_id







							def __call__(					self  :     Optional[Any]						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										# create boolean flag to decide if min length penalty should be applied
										lowerCamelCase    :List[Any]    	=  1 - jnp.clip(cur_len - self.min_length						,       0						,       1		)

										lowerCamelCase    :Any    	=  jnp.where(__snake_case						,       scores.at[:, self.eos_token_id].set(-float('''inf'''		)		)						,       __snake_case		)

										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Dict						,       __snake_case  :     Any						,       __snake_case  :     Optional[int]		):
										lowerCamelCase    :Optional[Any]    	=  list(__snake_case		)
										lowerCamelCase    :Union[str, Any]    	=  begin_index







							def __call__(					self  :     Any						,       __snake_case  :     Dict						,       __snake_case  :     int						,       __snake_case  :     int		):
										lowerCamelCase    :Union[str, Any]    	=  1 - jnp.bool_(cur_len - self.begin_index		)

										lowerCamelCase    :int    	=  jnp.where(__snake_case						,       scores.at[:, self.begin_suppress_tokens].set(-float('''inf'''		)		)						,       __snake_case		)

										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Dict						,       __snake_case  :     list		):
										lowerCamelCase    :Any    	=  list(__snake_case		)







							def __call__(					self  :     Tuple						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										lowerCamelCase    :List[Any]    	=  scores.at[..., self.suppress_tokens].set(-float('''inf'''		)		)

										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     List[Any]						,       __snake_case  :     Dict		):
										lowerCamelCase    :Optional[int]    	=  dict(__snake_case		)
										# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
										# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
										# Indexes without forced tokens will have a negative value.
										lowerCamelCase    :Any    	=  jnp.ones((max(force_token_map.keys()		) + 1)						,       dtype=jnp.intaa		) * -1
										for index, token in force_token_map.items():
													if token is not None:
																lowerCamelCase    :Union[str, Any]    	=  force_token_array.at[index].set(__snake_case		)
										lowerCamelCase    :List[str]    	=  jnp.intaa(__snake_case		)







							def __call__(					self  :     int						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     jnp.ndarray						,       __snake_case  :     int		):
										def _force_token(__snake_case  :     Optional[Any]		):
													lowerCamelCase    :Dict    	=  scores.shape[0]
													lowerCamelCase    :Tuple    	=  self.force_token_array[generation_idx]

													lowerCamelCase    :int    	=  jnp.ones_like(__snake_case						,       dtype=scores.dtype		) * -float('''inf'''		)
													lowerCamelCase    :Union[str, Any]    	=  jnp.zeros((batch_size, 1)						,       dtype=scores.dtype		)
													lowerCamelCase    :Dict    	=  lax.dynamic_update_slice(__snake_case						,       __snake_case						,       (0, current_token)		)
													return new_scores

										lowerCamelCase    :int    	=  lax.cond(
										    cur_len >= self.force_token_array.shape[0]						,       lambda: scores						,       lambda: lax.cond(
										        self.force_token_array[cur_len] >= 0						,       lambda: _force_token(__snake_case		)						,       lambda: scores						,       )						,       )
										return scores


class 	_lowerCAmelCase		(  lowerCamelCase__					):







							def __init__(					self  :     Any						,       __snake_case  :     int						,       __snake_case  :     Optional[Any]						,       __snake_case  :     Union[str, Any]		):
										lowerCamelCase    :Union[str, Any]    	=  generate_config.eos_token_id
										lowerCamelCase    :Dict    	=  generate_config.no_timestamps_token_id
										lowerCamelCase    :List[str]    	=  generate_config.no_timestamps_token_id + 1

										lowerCamelCase    :Any    	=  decoder_input_length + 1

										if generate_config.is_multilingual:
													# room for language token and task token
													self.begin_index += 2
										if hasattr(__snake_case						,       '''max_initial_timestamp_index'''		):
													lowerCamelCase    :List[str]    	=  generate_config.max_initial_timestamp_index
										else:
													lowerCamelCase    :int    	=  model_config.vocab_size
										if self.max_initial_timestamp_index is None:
													lowerCamelCase    :Any    	=  model_config.vocab_size







							def __call__(					self  :     str						,       __snake_case  :     Optional[Any]						,       __snake_case  :     List[Any]						,       __snake_case  :     Union[str, Any]		):
										# suppress <|notimestamps|> which is handled by without_timestamps
										lowerCamelCase    :Union[str, Any]    	=  scores.at[:, self.no_timestamps_token_id].set(-float('''inf'''		)		)

										def handle_pairs(__snake_case  :     str						,       __snake_case  :     Any		):
													lowerCamelCase    :Tuple    	=  jnp.where((cur_len - self.begin_index) >= 1						,       __snake_case						,       __snake_case		)
													lowerCamelCase    :List[str]    	=  jnp.where(
													    input_ids_k[cur_len - 1] >= self.timestamp_begin						,       True and last_was_timestamp						,       __snake_case						,       )

													lowerCamelCase    :str    	=  jnp.where((cur_len - self.begin_index) < 2						,       __snake_case						,       __snake_case		)
													lowerCamelCase    :Dict    	=  jnp.where(
													    input_ids_k[cur_len - 2] >= self.timestamp_begin						,       __snake_case						,       __snake_case						,       )

													return jnp.where(
													    __snake_case						,       jnp.where(
													        penultimate_was_timestamp > 0						,       scores_k.at[self.timestamp_begin :].set(-float('''inf'''		)		)						,       scores_k.at[: self.eos_token_id].set(-float('''inf'''		)		)						,       )						,       __snake_case						,       )

										lowerCamelCase    :List[str]    	=  jax.vmap(__snake_case		)(__snake_case						,       __snake_case		)

										lowerCamelCase    :Any    	=  jnp.where(cur_len == self.begin_index						,       __snake_case						,       __snake_case		)
										lowerCamelCase    :List[Any]    	=  jnp.where(
										    self.max_initial_timestamp_index is not None						,       True and apply_max_initial_timestamp						,       __snake_case						,       )

										lowerCamelCase    :Dict    	=  self.timestamp_begin + self.max_initial_timestamp_index

										lowerCamelCase    :Dict    	=  jnp.where(
										    __snake_case						,       scores.at[:, last_allowed + 1 :].set(-float('''inf'''		)		)						,       __snake_case						,       )

										# if sum of probability over timestamps is above any other token, sample timestamp
										lowerCamelCase    :Union[str, Any]    	=  jax.nn.log_softmax(__snake_case						,       axis=-1		)

										def handle_cumulative_probs(__snake_case  :     Any						,       __snake_case  :     Union[str, Any]		):
													lowerCamelCase    :Any    	=  jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]						,       axis=-1		)
													lowerCamelCase    :Optional[int]    	=  jnp.max(logprobs_k[: self.timestamp_begin]		)
													return jnp.where(
													    timestamp_logprob > max_text_token_logprob						,       scores_k.at[: self.timestamp_begin].set(-float('''inf'''		)		)						,       __snake_case						,       )

										lowerCamelCase    :str    	=  jax.vmap(__snake_case		)(__snake_case						,       __snake_case		)

										return scores








'''simple docstring'''






import baseaa
import io
import json
import os
from copy import deepcopy

from ..optimizer import AcceleratedOptimizer
from ..scheduler import AcceleratedScheduler

class 				a__:







    def __init__( self					:    List[Any]		,     __snake_case					:    str      ):
      if isinstance(__snake_case		,     __snake_case      ):
        # Don't modify user's data should they want to reuse it (e.g. in tests), because once we
        # modified it, it will not be accepted here again, since `auto` values would have been overridden
        a		:   Union[str, Any]	  =	deepcopy(__snake_case      )
      elif os.path.exists(__snake_case      ):
        with io.open(__snake_case		,     'r'		,     encoding='utf-8'      ) as f:
          a		:   List[str]	  =	json.load(__snake_case      )
      else:
        try:
          a		:   str	  =	baseaa.urlsafe_baadecode(__snake_case      ).decode('utf-8'      )
          a		:   Union[str, Any]	  =	json.loads(__snake_case      )
        except (UnicodeDecodeError, AttributeError, ValueError):
          raise ValueError(
              F"""Expected a string path to an existing deepspeed config, or a dictionary, or a base64 encoded string. Received: {config_file_or_dict}"""      )

      a		:   Optional[int]	  =	config

      self.set_stage_and_offload()







    def      lowercase_			( self					:    Any      ):
      # zero stage - this is done as early as possible, before model is created, to allow
      # ``is_deepspeed_zero3_enabled`` query and getting to the early deepspeed config object
      # during ``zero.Init()`` which needs to know the dtype, and some other hparams.
      a		:   List[str]	  =	self.get_value('zero_optimization.stage'		,     -1      )

      # offload
      a		:   Union[str, Any]	  =	False
      if self.is_zeroa() or self.is_zeroa():
        a		:   Any	  =	set(['cpu', 'nvme']      )
        a		:   Tuple	  =	set(
            [
                self.get_value('zero_optimization.offload_optimizer.device'      ),
                self.get_value('zero_optimization.offload_param.device'      ),
            ]      )
        if len(offload_devices & offload_devices_valid      ) > 0:
          a		:   Optional[int]	  =	True







    def      lowercase_			( self					:    str		,     __snake_case					:    Optional[Any]      ):
      a		:   Any	  =	self.config

      # find the config node of interest if it exists
      a		:   Dict	  =	ds_key_long.split('.'      )
      a		:   str	  =	nodes.pop()
      for node in nodes:
        a		:   Tuple	  =	config.get(__snake_case      )
        if config is None:
          return None, ds_key

      return config, ds_key







    def      lowercase_			( self					:    int		,     __snake_case					:    Optional[int]		,     __snake_case					:    List[str]=None      ):
      a  ,      a		:   int	  =	self.find_config_node(__snake_case      )
      if config is None:
        return default
      return config.get(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    Dict		,     __snake_case					:    List[str]		,     __snake_case					:    List[Any]=False      ):
      a		:   Union[str, Any]	  =	self.config

      # find the config node of interest if it exists
      a		:   Tuple	  =	ds_key_long.split('.'      )
      for node in nodes:
        a		:   Optional[int]	  =	config
        a		:   List[str]	  =	config.get(__snake_case      )
        if config is None:
          if must_exist:
            raise ValueError(F"""Can't find {ds_key_long} entry in the config: {self.config}"""      )
          else:
            return

        # if found remove it
      if parent_config is not None:
        parent_config.pop(__snake_case      )







    def      lowercase_			( self					:    List[str]		,     __snake_case					:    Dict      ):
      a		:   Optional[Any]	  =	self.get_value(__snake_case      )
      return False if value is None else bool(__snake_case      )







    def      lowercase_			( self					:    str		,     __snake_case					:    Optional[Any]      ):
      a		:   Optional[Any]	  =	self.get_value(__snake_case      )
      return False if value is None else not bool(__snake_case      )







    def      lowercase_			( self					:    Tuple      ):
      return self._stage == 2







    def      lowercase_			( self					:    str      ):
      return self._stage == 3







    def      lowercase_			( self					:    int      ):
      return self._offload

class 				a__:







    def __init__( self					:    str		,     __snake_case					:    Any      ):
      a		:   Any	  =	engine







    def      lowercase_			( self					:    Optional[Any]		,     __snake_case					:    int		,     **__snake_case					:    List[Any]      ):
      # runs backpropagation and handles mixed precision
      self.engine.backward(__snake_case		,     **__snake_case      )

      # Deepspeed's `engine.step` performs the following operations:
      # - gradient accumulation check
      # - gradient clipping
      # - optimizer step
      # - zero grad
      # - checking overflow
      # - lr_scheduler step (only if engine.lr_scheduler is not None)
      self.engine.step()
      # and this plugin overrides the above calls with no-ops when Accelerate runs under
      # Deepspeed, but allows normal functionality for non-Deepspeed cases thus enabling a simple
      # training loop that works transparently under many training regimes.

class 				a__(   lowerCamelCase__      ):







    def __init__( self					:    Dict		,     __snake_case					:    Any      ):
      super().__init__(__snake_case		,     device_placement=__snake_case		,     scaler=__snake_case      )
      a		:   List[Any]	  =	hasattr(self.optimizer		,     'overflow'      )







    def      lowercase_			( self					:    Optional[Any]		,     __snake_case					:    str=None      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







    def      lowercase_			( self					:    int      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed







    @property
    def      lowercase_			( self					:    Any      ):
      if self.__has_overflow__:
        return self.optimizer.overflow
      return False

class 				a__(   lowerCamelCase__      ):







    def __init__( self					:    str		,     __snake_case					:    List[Any]		,     __snake_case					:    str      ):
      super().__init__(__snake_case		,     __snake_case      )







    def      lowercase_			( self					:    Union[str, Any]      ):
      pass  # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed

class 				a__:







    def __init__( self					:    Optional[Any]		,     __snake_case					:    Union[str, Any]		,     __snake_case					:    int=0.001		,     __snake_case					:    Any=0		,     **__snake_case					:    List[Any]      ):
      a		:   Optional[Any]	  =	params
      a		:   Union[str, Any]	  =	lr
      a		:   Any	  =	weight_decay
      a		:   Any	  =	kwargs

class 				a__:







    def __init__( self					:    Any		,     __snake_case					:    int		,     __snake_case					:    int=None		,     __snake_case					:    Dict=0		,     **__snake_case					:    Optional[int]      ):
      a		:   Optional[int]	  =	optimizer
      a		:   List[str]	  =	total_num_steps
      a		:   Tuple	  =	warmup_num_steps
      a		:   Optional[Any]	  =	kwargs