
"""simple docstring"""

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


UpperCamelCase			:       Union[str, Any]									=				get_logger(__name__)


UpperCamelCase			:       List[str]									=				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(__UpperCAmelCase   )
					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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


class __lowerCAmelCase	:





					@add_start_docstrings(__UpperCAmelCase   )
					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					@add_start_docstrings(__UpperCAmelCase   )
					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase			,  **__UpperCAmelCase   ):

						'''simple docstring'''




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


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

						__UpperCamelCase       						= temperature







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= scores / self.temperature
						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase = -float('Inf'   )			,  __UpperCAmelCase = 1   ):

						'''simple docstring'''




						if not isinstance(__UpperCAmelCase			,  __UpperCAmelCase   ) 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(__UpperCAmelCase			,  __UpperCAmelCase   ) 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			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase						,			__UpperCamelCase       						= lax.top_k(__UpperCAmelCase			,  scores.shape[-1]   )

						__UpperCamelCase       						= jnp.full_like(__UpperCAmelCase			,  self.filter_value   )
						__UpperCamelCase       						= jax.nn.softmax(__UpperCAmelCase			,  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(__UpperCAmelCase			,  1   )
						score_mask |= score_mask.at[:, 0].set(__UpperCAmelCase   )

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

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

						return next_scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase = -float('Inf'   )			,  __UpperCAmelCase = 1   ):

						'''simple docstring'''




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

						__UpperCamelCase       						= max(__UpperCAmelCase			,  __UpperCAmelCase   )
						__UpperCamelCase       						= filter_value







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__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(__UpperCAmelCase			,  __UpperCAmelCase   )
						__UpperCamelCase       						= jnp.broadcast_to((jnp.arange(__UpperCAmelCase   ) * 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(__UpperCAmelCase   )
						__UpperCamelCase       						= next_scores_flat.reshape(__UpperCAmelCase			,  __UpperCAmelCase   )
						return next_scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= bos_token_id







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

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

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

						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= max_length
						__UpperCamelCase       						= eos_token_id







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

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

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

						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

						if not isinstance(__UpperCAmelCase			,  __UpperCAmelCase   ) 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			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

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

						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= list(__UpperCAmelCase   )
						__UpperCamelCase       						= begin_index







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

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

						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= list(__UpperCAmelCase   )







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





					def __init__(			self			,  __UpperCAmelCase   ):

						'''simple docstring'''




						__UpperCamelCase       						= dict(__UpperCAmelCase   )
						# 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(__UpperCAmelCase   )
						__UpperCamelCase       						= jnp.intaa(__UpperCAmelCase   )







					def __call__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




						def _force_token(__UpperCAmelCase   ):
							__UpperCamelCase       						= scores.shape[0]
							__UpperCamelCase       						= self.force_token_array[generation_idx]

							__UpperCamelCase       						= jnp.ones_like(__UpperCAmelCase			,  dtype=scores.dtype   ) * -float('inf'   )
							__UpperCamelCase       						= jnp.zeros((batch_size, 1)			,  dtype=scores.dtype   )
							__UpperCamelCase       						= lax.dynamic_update_slice(__UpperCAmelCase			,  __UpperCAmelCase			,  (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(__UpperCAmelCase   )			,  lambda: scores			,  )			,  )
						return scores


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





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

						'''simple docstring'''




						__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(__UpperCAmelCase			,  '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			,  __UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

						'''simple docstring'''




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

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

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

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

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

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

						__UpperCamelCase       						= self.timestamp_begin + self.max_initial_timestamp_index

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

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

						def handle_cumulative_probs(__UpperCAmelCase			,  __UpperCAmelCase   ):
							__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'   )   )			,  __UpperCAmelCase			,  )

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

						return scores





"""simple docstring"""

from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase			:       Dict									=				{
    "configuration_lxmert": ["LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "LxmertConfig"],
    "tokenization_lxmert": ["LxmertTokenizer"],
}

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase			:       List[str]									=				["LxmertTokenizerFast"]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase			:       Optional[Any]									=				[
          "LxmertEncoder",
          "LxmertForPreTraining",
          "LxmertForQuestionAnswering",
          "LxmertModel",
          "LxmertPreTrainedModel",
          "LxmertVisualFeatureEncoder",
          "LxmertXLayer",
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase			:       List[Any]									=				[
          "TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST",
          "TFLxmertForPreTraining",
          "TFLxmertMainLayer",
          "TFLxmertModel",
          "TFLxmertPreTrainedModel",
          "TFLxmertVisualFeatureEncoder",
      ]


if TYPE_CHECKING:
      from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
      from .tokenization_lxmert import LxmertTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_lxmert_fast import LxmertTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_lxmert import (
                LxmertEncoder,
                LxmertForPreTraining,
                LxmertForQuestionAnswering,
                LxmertModel,
                LxmertPreTrainedModel,
                LxmertVisualFeatureEncoder,
                LxmertXLayer,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_lxmert import (
                TF_LXMERT_PRETRAINED_MODEL_ARCHIVE_LIST,
                TFLxmertForPreTraining,
                TFLxmertMainLayer,
                TFLxmertModel,
                TFLxmertPreTrainedModel,
                TFLxmertVisualFeatureEncoder,
            )

else:
      import sys

      UpperCamelCase			:       Any									=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



