
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      :				Any          = get_logger(__name__)


_UpperCAmelCase      :				Tuple          = 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        lowercase      :



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






class        lowercase      :



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






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		@add_start_docstrings(A_ )
		def __call__(				self       ,   A_       ,   A_       ,   A_       ,   **A_ )						->      jnp.ndarray:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				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:
				"""simple docstring"""
				UpperCamelCase						=  scores / self.temperature
				return scores






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_ = -float('Inf' )       ,   A_ = 1 )						->      Union[str, Any]:
				"""simple docstring"""
				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:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_ = -float('Inf' )       ,   A_ = 1 )						->      Any:
				"""simple docstring"""
				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:
				"""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(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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Optional[Any]:
				"""simple docstring"""
				UpperCamelCase						=  bos_token_id



		def __call__(				self       ,   A_       ,   A_       ,   A_ )						->      jnp.ndarray:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_ )						->      str:
				"""simple docstring"""
				UpperCamelCase						=  max_length
				UpperCamelCase						=  eos_token_id



		def __call__(				self       ,   A_       ,   A_       ,   A_ )						->      jnp.ndarray:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_ )						->      Tuple:
				"""simple docstring"""
				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:
				"""simple docstring"""
				# 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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_ )						->      Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase						=  list(A_ )
				UpperCamelCase						=  begin_index



		def __call__(				self       ,   A_       ,   A_       ,   A_ )						->      Any:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      Any:
				"""simple docstring"""
				UpperCamelCase						=  list(A_ )



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

				return scores






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_ )						->      List[Any]:
				"""simple docstring"""
				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:
				"""simple docstring"""
				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        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		def __init__(				self       ,   A_       ,   A_       ,   A_ )						->      int:
				"""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(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:
				"""simple docstring"""
				# 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


from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


_UpperCAmelCase      :				Tuple          = logging.get_logger(__name__)

_UpperCAmelCase      :				Union[str, Any]          = {
    "facebook/data2vec-text-base": "https://huggingface.co/data2vec/resolve/main/config.json",
}


class        lowercase      (	_SCREAMING_SNAKE_CASE     ):
		__lowercase : Dict			  =       "data2vec-text"



		def __init__(				self       ,   A_=30_522       ,   A_=768       ,   A_=12       ,   A_=12       ,   A_=3_072       ,   A_="gelu"       ,   A_=0.1       ,   A_=0.1       ,   A_=512       ,   A_=2       ,   A_=0.02       ,   A_=1e-12       ,   A_=1       ,   A_=0       ,   A_=2       ,   A_="absolute"       ,   A_=True       ,   A_=None       ,   **A_       ,   )						->      Any:
				"""simple docstring"""
				super().__init__(pad_token_id=A_       ,   bos_token_id=A_       ,   eos_token_id=A_       ,   **A_ )

				UpperCamelCase						=  vocab_size
				UpperCamelCase						=  hidden_size
				UpperCamelCase						=  num_hidden_layers
				UpperCamelCase						=  num_attention_heads
				UpperCamelCase						=  hidden_act
				UpperCamelCase						=  intermediate_size
				UpperCamelCase						=  hidden_dropout_prob
				UpperCamelCase						=  attention_probs_dropout_prob
				UpperCamelCase						=  max_position_embeddings
				UpperCamelCase						=  type_vocab_size
				UpperCamelCase						=  initializer_range
				UpperCamelCase						=  layer_norm_eps
				UpperCamelCase						=  position_embedding_type
				UpperCamelCase						=  use_cache
				UpperCamelCase						=  classifier_dropout






class        lowercase      (	_SCREAMING_SNAKE_CASE     ):



		@property
		def __UpperCamelCase				(				self )						->      Mapping[str, Mapping[int, str]]:
				"""simple docstring"""
				if self.task == "multiple-choice":
						UpperCamelCase						=  {0: 'batch', 1: 'choice', 2: 'sequence'}
				else:
						UpperCamelCase						=  {0: 'batch', 1: 'sequence'}
				return OrderedDict(
				    [
				        ('input_ids', dynamic_axis),
				        ('attention_mask', dynamic_axis),
				    ] )
