




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


_lowerCamelCase      :						Dict							=  get_logger(__name__)


_lowerCamelCase      :						Dict							=  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 							UpperCamelCase_       :

    '''simple docstring'''






    @add_start_docstrings(lowerCAmelCase_)
    def __call__(     self :      Optional[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      int)		->jnp.ndarray:





           '''simple docstring'''

           raise NotImplementedError(
               f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")


class 							UpperCamelCase_       :

    '''simple docstring'''






    @add_start_docstrings(lowerCAmelCase_)
    def __call__(     self :      Any   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      List[str])		->jnp.ndarray:





           '''simple docstring'''

           raise NotImplementedError(
               f"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    @add_start_docstrings(lowerCAmelCase_)
    def __call__(     self :      Any   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      str   ,					**UpperCAmelCase__ :      Union[str, Any])		->jnp.ndarray:





           '''simple docstring'''

           for processor in self:
                  A__ 					=       inspect.signature(processor.__call__).parameters
                  if len(lowerCAmelCase_) > 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.""")
                         A__ 					=       processor(lowerCAmelCase_   ,					lowerCAmelCase_   ,					lowerCAmelCase_   ,					**lowerCAmelCase_)
                  else:
                         A__ 					=       processor(lowerCAmelCase_   ,					lowerCAmelCase_   ,					lowerCAmelCase_)
           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      Optional[int]   ,					UpperCAmelCase__ :      List[str])		->Any:





           '''simple docstring'''

           if not isinstance(lowerCAmelCase_   ,					lowerCAmelCase_) or not (temperature > 0):
                  raise ValueError(f"""`temperature` has to be a strictly positive float, but is {temperature}""")

           A__ 					=       temperature





    def __call__(     self :      Any   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      int)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       scores / self.temperature
           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      int   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      List[str] = -float('''Inf''')   ,					UpperCAmelCase__ :      Optional[int] = 1)		->int:





           '''simple docstring'''

           if not isinstance(lowerCAmelCase_   ,					lowerCAmelCase_) 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(lowerCAmelCase_   ,					lowerCAmelCase_) 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}""")

           A__ 					=       top_p
           A__ 					=       filter_value
           A__ 					=       min_tokens_to_keep





    def __call__(     self :      Union[str, Any]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      int)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       lax.top_k(lowerCAmelCase_   ,					scores.shape[-1])

           A__ 					=       jnp.full_like(lowerCAmelCase_   ,					self.filter_value)
           A__ 					=       jax.nn.softmax(lowerCAmelCase_   ,					axis=-1).cumsum(axis=-1)
           A__ 					=       cumulative_probs < self.top_p

           # include the token that is higher than top_p as well
           A__ 					=       jnp.roll(lowerCAmelCase_   ,					1)
           score_mask |= score_mask.at[:, 0].set(lowerCAmelCase_)

           # min tokens to keep
           A__ 					=       score_mask.at[:, : self.min_tokens_to_keep].set(lowerCAmelCase_)

           A__ 					=       jnp.where(lowerCAmelCase_   ,					lowerCAmelCase_   ,					lowerCAmelCase_)
           A__ 					=       jax.lax.sort_key_val(lowerCAmelCase_   ,					lowerCAmelCase_)[-1]

           return next_scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      Optional[Any] = -float('''Inf''')   ,					UpperCAmelCase__ :      int = 1)		->Tuple:





           '''simple docstring'''

           if not isinstance(lowerCAmelCase_   ,					lowerCAmelCase_) or top_k <= 0:
                  raise ValueError(f"""`top_k` has to be a strictly positive integer, but is {top_k}""")

           A__ 					=       max(lowerCAmelCase_   ,					lowerCAmelCase_)
           A__ 					=       filter_value





    def __call__(     self :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      str)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       scores.shape
           A__ 					=       jnp.full(batch_size * vocab_size   ,					self.filter_value)

           A__ 					=       min(self.top_k   ,					scores.shape[-1])  # Safety check
           A__ 					=       lax.top_k(lowerCAmelCase_   ,					lowerCAmelCase_)
           A__ 					=       jnp.broadcast_to((jnp.arange(lowerCAmelCase_) * vocab_size)[:, None]   ,					(batch_size, topk)).flatten()
           A__ 					=       topk_scores.flatten()
           A__ 					=       topk_indices.flatten() + shift

           A__ 					=       next_scores_flat.at[topk_indices_flat].set(lowerCAmelCase_)
           A__ 					=       next_scores_flat.reshape(lowerCAmelCase_   ,					lowerCAmelCase_)
           return next_scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      List[Any]   ,					UpperCAmelCase__ :      List[Any])		->Dict:





           '''simple docstring'''

           A__ 					=       bos_token_id





    def __call__(     self :      Any   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      int)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       jnp.full(scores.shape   ,					-float('''inf'''))

           A__ 					=       1 - jnp.bool_(cur_len - 1)

           A__ 					=       jnp.where(lowerCAmelCase_   ,					new_scores.at[:, self.bos_token_id].set(0)   ,					lowerCAmelCase_)

           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      Optional[int]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Optional[int])		->List[Any]:





           '''simple docstring'''

           A__ 					=       max_length
           A__ 					=       eos_token_id





    def __call__(     self :      int   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      str)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       jnp.full(scores.shape   ,					-float('''inf'''))

           A__ 					=       1 - jnp.bool_(cur_len - self.max_length + 1)

           A__ 					=       jnp.where(lowerCAmelCase_   ,					new_scores.at[:, self.eos_token_id].set(0)   ,					lowerCAmelCase_)

           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      Optional[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Tuple)		->Optional[int]:





           '''simple docstring'''

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

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

           A__ 					=       min_length
           A__ 					=       eos_token_id





    def __call__(     self :      Any   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Optional[Any])		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       1 - jnp.clip(cur_len - self.min_length   ,					0   ,					1)

           A__ 					=       jnp.where(lowerCAmelCase_   ,					scores.at[:, self.eos_token_id].set(-float('''inf'''))   ,					lowerCAmelCase_)

           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      Union[str, Any]   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Any)		->List[str]:





           '''simple docstring'''

           A__ 					=       list(lowerCAmelCase_)
           A__ 					=       begin_index





    def __call__(     self :      Union[str, Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Optional[int])		->str:





           '''simple docstring'''

           A__ 					=       1 - jnp.bool_(cur_len - self.begin_index)

           A__ 					=       jnp.where(lowerCAmelCase_   ,					scores.at[:, self.begin_suppress_tokens].set(-float('''inf'''))   ,					lowerCAmelCase_)

           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      str   ,					UpperCAmelCase__ :      Any)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       list(lowerCAmelCase_)





    def __call__(     self :      Tuple   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Tuple)		->jnp.ndarray:





           '''simple docstring'''

           A__ 					=       scores.at[..., self.suppress_tokens].set(-float('''inf'''))

           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






    def __init__(     self :      int   ,					UpperCAmelCase__ :      str)		->Any:





           '''simple docstring'''

           A__ 					=       dict(lowerCAmelCase_)
           # 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.
           A__ 					=       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:
                         A__ 					=       force_token_array.at[index].set(lowerCAmelCase_)
           A__ 					=       jnp.intaa(lowerCAmelCase_)





    def __call__(     self :      Tuple   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      int)		->jnp.ndarray:





           '''simple docstring'''

           def _force_token(UpperCAmelCase__ :      List[Any]):
                  A__ 					=       scores.shape[0]
                  A__ 					=       self.force_token_array[generation_idx]

                  A__ 					=       jnp.ones_like(lowerCAmelCase_   ,					dtype=scores.dtype) * -float('''inf''')
                  A__ 					=       jnp.zeros((batch_size, 1)   ,					dtype=scores.dtype)
                  A__ 					=       lax.dynamic_update_slice(lowerCAmelCase_   ,					lowerCAmelCase_   ,					(0, current_token))
                  return new_scores

           A__ 					=       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(lowerCAmelCase_)   ,					lambda: scores   ,					)   ,					)
           return scores


class 							UpperCamelCase_       (					__UpperCAmelCase	):

    '''simple docstring'''






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





           '''simple docstring'''

           A__ 					=       generate_config.eos_token_id
           A__ 					=       generate_config.no_timestamps_token_id
           A__ 					=       generate_config.no_timestamps_token_id + 1

           A__ 					=       decoder_input_length + 1

           if generate_config.is_multilingual:
                  # room for language token and task token
                  self.begin_index += 2
           if hasattr(lowerCAmelCase_   ,					'''max_initial_timestamp_index'''):
                  A__ 					=       generate_config.max_initial_timestamp_index
           else:
                  A__ 					=       model_config.vocab_size
           if self.max_initial_timestamp_index is None:
                  A__ 					=       model_config.vocab_size





    def __call__(     self :      List[Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      Dict)		->Dict:





           '''simple docstring'''

           A__ 					=       scores.at[:, self.no_timestamps_token_id].set(-float('''inf'''))

           def handle_pairs(UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Optional[Any]):
                  A__ 					=       jnp.where((cur_len - self.begin_index) >= 1   ,					lowerCAmelCase_   ,					lowerCAmelCase_)
                  A__ 					=       jnp.where(
                      input_ids_k[cur_len - 1] >= self.timestamp_begin   ,					True and last_was_timestamp   ,					lowerCAmelCase_   ,					)

                  A__ 					=       jnp.where((cur_len - self.begin_index) < 2   ,					lowerCAmelCase_   ,					lowerCAmelCase_)
                  A__ 					=       jnp.where(
                      input_ids_k[cur_len - 2] >= self.timestamp_begin   ,					lowerCAmelCase_   ,					lowerCAmelCase_   ,					)

                  return jnp.where(
                      lowerCAmelCase_   ,					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'''))   ,					)   ,					lowerCAmelCase_   ,					)

           A__ 					=       jax.vmap(lowerCAmelCase_)(lowerCAmelCase_   ,					lowerCAmelCase_)

           A__ 					=       jnp.where(cur_len == self.begin_index   ,					lowerCAmelCase_   ,					lowerCAmelCase_)
           A__ 					=       jnp.where(
               self.max_initial_timestamp_index is not None   ,					True and apply_max_initial_timestamp   ,					lowerCAmelCase_   ,					)

           A__ 					=       self.timestamp_begin + self.max_initial_timestamp_index

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

           # if sum of probability over timestamps is above any other token, sample timestamp
           A__ 					=       jax.nn.log_softmax(lowerCAmelCase_   ,					axis=-1)

           def handle_cumulative_probs(UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      List[Any]):
                  A__ 					=       jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]   ,					axis=-1)
                  A__ 					=       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'''))   ,					lowerCAmelCase_   ,					)

           A__ 					=       jax.vmap(lowerCAmelCase_)(lowerCAmelCase_   ,					lowerCAmelCase_)

           return scores










"""simple docstring"""
import platform
from argparse import ArgumentParser

import huggingface_hub

from .. import __version__ as version
from ..utils import is_accelerate_available, is_torch_available, is_transformers_available, is_xformers_available
from . import BaseDiffusersCLICommand





def 				lowercase__		(					lowerCamelCase       ):
				return EnvironmentCommand()






class _lowerCAmelCase						(    __UpperCAmelCase       ):







						@staticmethod
						def 			A      (    lowerCAmelCase_							)		->							List[Any]:
										_SCREAMING_SNAKE_CASE   :		int       		=    parser.add_parser('env'							)
										download_parser.set_defaults(func=lowerCAmelCase_							)







						def 			A      (    self							)		->							Any:
										_SCREAMING_SNAKE_CASE   :		int       		=    huggingface_hub.__version__

										_SCREAMING_SNAKE_CASE   :		Optional[Any]       		=    'not installed'
										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    'NA'
										if is_torch_available():
														import torch

														_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    torch.__version__
														_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    torch.cuda.is_available()

										_SCREAMING_SNAKE_CASE   :		Union[str, Any]       		=    'not installed'
										if is_transformers_available():
														import transformers

														_SCREAMING_SNAKE_CASE   :		Tuple       		=    transformers.__version__

										_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    'not installed'
										if is_accelerate_available():
														import accelerate

														_SCREAMING_SNAKE_CASE   :		str       		=    accelerate.__version__

										_SCREAMING_SNAKE_CASE   :		str       		=    'not installed'
										if is_xformers_available():
														import xformers

														_SCREAMING_SNAKE_CASE   :		Optional[int]       		=    xformers.__version__

										_SCREAMING_SNAKE_CASE   :		Any       		=    {
										    '`diffusers` version': version,
										    'Platform': platform.platform(),
										    'Python version': platform.python_version(),
										    'PyTorch version (GPU?)': F"""{pt_version} ({pt_cuda_available})""",
										    'Huggingface_hub version': hub_version,
										    'Transformers version': transformers_version,
										    'Accelerate version': accelerate_version,
										    'xFormers version': xformers_version,
										    'Using GPU in script?': '<fill in>',
										    'Using distributed or parallel set-up in script?': '<fill in>',
										}

										print('\nCopy-and-paste the text below in your GitHub issue and FILL OUT the two last points.\n'							)
										print(self.format_dict(lowerCAmelCase_							)							)

										return info







						@staticmethod
						def 			A      (    lowerCAmelCase_							)		->							int:
										return "\n".join([F"""- {prop}: {val}""" for prop, val in d.items()]							) + "\n"

