


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    UpperCAmelCase__	:



    """simple docstring"""






    @add_start_docstrings(A_   )
    def __call__(      self ,       A_ ,       A_   )  ->	jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'   )






class    UpperCAmelCase__	:



    """simple docstring"""






    @add_start_docstrings(A_   )
    def __call__(      self ,       A_ ,       A_   )  ->	jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'   )






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    @add_start_docstrings(A_   )
    def __call__(      self ,       A_ ,       A_ ,       A_ ,       **A_   )  ->	jnp.ndarray:
         for processor in self:
              __UpperCamelCase   							=inspect.signature(processor.__call__   ).parameters
              if len(A_   ) > 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.'   )
                   __UpperCamelCase   							=processor(A_ ,       A_ ,       A_ ,       **A_   )
              else:
                   __UpperCamelCase   							=processor(A_ ,       A_ ,       A_   )
         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_   )  ->	List[str]:
         if not isinstance(A_ ,       A_   ) or not (temperature > 0):
              raise ValueError(f'`temperature` has to be a strictly positive float, but is {temperature}'   )

         __UpperCamelCase   							=temperature






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase   							=scores / self.temperature
         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ = -float('Inf'   ) ,       A_ = 1   )  ->	Dict:
         if not isinstance(A_ ,       A_   ) 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(A_ ,       A_   ) 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}'   )

         __UpperCamelCase   							=top_p
         __UpperCamelCase   							=filter_value
         __UpperCamelCase   							=min_tokens_to_keep






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase							,				__UpperCamelCase   							=lax.top_k(A_ ,       scores.shape[-1]   )

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

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

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

         __UpperCamelCase   							=jnp.where(A_ ,       A_ ,       A_   )
         __UpperCamelCase   							=jax.lax.sort_key_val(A_ ,       A_   )[-1]

         return next_scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






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

         __UpperCamelCase   							=max(A_ ,       A_   )
         __UpperCamelCase   							=filter_value






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase							,				__UpperCamelCase   							=scores.shape
         __UpperCamelCase   							=jnp.full(batch_size * vocab_size ,       self.filter_value   )

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

         __UpperCamelCase   							=next_scores_flat.at[topk_indices_flat].set(A_   )
         __UpperCamelCase   							=next_scores_flat.reshape(A_ ,       A_   )
         return next_scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_   )  ->	int:
         __UpperCamelCase   							=bos_token_id






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase   							=jnp.full(scores.shape ,       -float('inf'   )   )

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

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

         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_   )  ->	Optional[int]:
         __UpperCamelCase   							=max_length
         __UpperCamelCase   							=eos_token_id






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase   							=jnp.full(scores.shape ,       -float('inf'   )   )

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

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

         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_   )  ->	Union[str, Any]:
         if not isinstance(A_ ,       A_   ) or min_length < 0:
              raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'   )

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

         __UpperCamelCase   							=min_length
         __UpperCamelCase   							=eos_token_id






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         # create boolean flag to decide if min length penalty should be applied
         __UpperCamelCase   							=1 - jnp.clip(cur_len - self.min_length ,       0 ,       1   )

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

         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_   )  ->	Tuple:
         __UpperCamelCase   							=list(A_   )
         __UpperCamelCase   							=begin_index






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	Union[str, Any]:
         __UpperCamelCase   							=1 - jnp.bool_(cur_len - self.begin_index   )

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

         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_   )  ->	Optional[Any]:
         __UpperCamelCase   							=list(A_   )






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         __UpperCamelCase   							=scores.at[..., self.suppress_tokens].set(-float('inf'   )   )

         return scores






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_   )  ->	str:
         __UpperCamelCase   							=dict(A_   )
         # 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.
         __UpperCamelCase   							=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:
                   __UpperCamelCase   							=force_token_array.at[index].set(A_   )
         __UpperCamelCase   							=jnp.intaa(A_   )






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	jnp.ndarray:
         def _force_token(A_   ):
              __UpperCamelCase   							=scores.shape[0]
              __UpperCamelCase   							=self.force_token_array[generation_idx]

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

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






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    def __init__(      self ,       A_ ,       A_ ,       A_   )  ->	int:
         __UpperCamelCase   							=generate_config.eos_token_id
         __UpperCamelCase   							=generate_config.no_timestamps_token_id
         __UpperCamelCase   							=generate_config.no_timestamps_token_id + 1

         __UpperCamelCase   							=decoder_input_length + 1

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






    def __call__(      self ,       A_ ,       A_ ,       A_   )  ->	int:
         # suppress <|notimestamps|> which is handled by without_timestamps
         __UpperCamelCase   							=scores.at[:, self.no_timestamps_token_id].set(-float('inf'   )   )

         def handle_pairs(A_ ,       A_   ):
              __UpperCamelCase   							=jnp.where((cur_len - self.begin_index) >= 1 ,       A_ ,       A_   )
              __UpperCamelCase   							=jnp.where(
                  input_ids_k[cur_len - 1] >= self.timestamp_begin ,       True and last_was_timestamp ,       A_ ,       )

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

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

         __UpperCamelCase   							=jax.vmap(A_   )(A_ ,       A_   )

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

         __UpperCamelCase   							=self.timestamp_begin + self.max_initial_timestamp_index

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

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

         def handle_cumulative_probs(A_ ,       A_   ):
              __UpperCamelCase   							=jax.nn.logsumexp(logprobs_k[self.timestamp_begin :] ,       axis=-1   )
              __UpperCamelCase   							=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'   )   ) ,       A_ ,       )

         __UpperCamelCase   							=jax.vmap(A_   )(A_ ,       A_   )

         return scores









import sys

from .dependency_versions_table import deps
from .utils.versions import require_version, require_version_core


# define which module versions we always want to check at run time
# (usually the ones defined in `install_requires` in setup.py)
#
# order specific notes:
# - tqdm must be checked before tokenizers

_A			=						'python tqdm regex requests packaging filelock numpy tokenizers'.split()
if sys.version_info < (3, 7):
   pkgs_to_check_at_runtime.append('dataclasses')
if sys.version_info < (3, 8):
   pkgs_to_check_at_runtime.append('importlib_metadata')

for pkg in pkgs_to_check_at_runtime:
   if pkg in deps:
      if pkg == "tokenizers":
         # must be loaded here, or else tqdm check may fail
         from .utils import is_tokenizers_available

         if not is_tokenizers_available():
            continue  # not required, check version only if installed

      require_version_core(deps[pkg])
   else:
      raise ValueError(f"""can't find {pkg} in {deps.keys()}, check dependency_versions_table.py""")

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Union[str, Any]      ,					SCREAMING_SNAKE_CASE__      : Optional[int]=None    ):
     require_version(deps[pkg]      ,					SCREAMING_SNAKE_CASE__    )





