






"""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


__lowerCamelCase							=  get_logger(__name__)


__lowerCamelCase							=  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       __A		:

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



class       __A		:

					@add_start_docstrings(_SCREAMING_SNAKE_CASE					)
					def __call__(			self  :      List[str]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray					)			->							jnp.ndarray:
											raise NotImplementedError(
											    F'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'					)



class       __A		(   __SCREAMING_SNAKE_CASE			):

					@add_start_docstrings(_SCREAMING_SNAKE_CASE					)
					def __call__(			self  :      Optional[int]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int     ,		**__snake_case  :      str					)			->							jnp.ndarray:

											for processor in self:
																	__magic_name__:	int							=       inspect.signature(processor.__call__					).parameters
																	if len(_SCREAMING_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.'					)
																							__magic_name__:	int							=       processor(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		**_SCREAMING_SNAKE_CASE					)
																	else:
																							__magic_name__:	Union[str, Any]							=       processor(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)
											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      int     ,		__snake_case  :      float					)			->							str:

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

											__magic_name__:	List[str]							=       temperature






					def __call__(			self  :      Dict     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Union[str, Any]							=       scores / self.temperature
											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      Dict     ,		__snake_case  :      float     ,		__snake_case  :      float = -float("""Inf"""					)     ,		__snake_case  :      int = 1					)			->							Optional[int]:

											if not isinstance(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_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(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_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}'					)

											__magic_name__:	Union[str, Any]							=       top_p
											__magic_name__:	Optional[int]							=       filter_value
											__magic_name__:	str							=       min_tokens_to_keep






					def __call__(			self  :      Union[str, Any]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Optional[Any]							=       lax.top_k(_SCREAMING_SNAKE_CASE     ,		scores.shape[-1]					)

											__magic_name__:	List[str]							=       jnp.full_like(_SCREAMING_SNAKE_CASE     ,		self.filter_value					)
											__magic_name__:	Any							=       jax.nn.softmax(_SCREAMING_SNAKE_CASE     ,		axis=-1					).cumsum(axis=-1					)
											__magic_name__:	Optional[Any]							=       cumulative_probs < self.top_p

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

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

											__magic_name__:	Optional[Any]							=       jnp.where(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)
											__magic_name__:	Dict							=       jax.lax.sort_key_val(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)[-1]

											return next_scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      Tuple     ,		__snake_case  :      int     ,		__snake_case  :      float = -float("""Inf"""					)     ,		__snake_case  :      int = 1					)			->							Optional[int]:

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

											__magic_name__:	Any							=       max(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)
											__magic_name__:	Optional[int]							=       filter_value






					def __call__(			self  :      List[Any]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Any							=       scores.shape
											__magic_name__:	Union[str, Any]							=       jnp.full(batch_size * vocab_size     ,		self.filter_value					)

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

											__magic_name__:	str							=       next_scores_flat.at[topk_indices_flat].set(_SCREAMING_SNAKE_CASE					)
											__magic_name__:	Union[str, Any]							=       next_scores_flat.reshape(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)
											return next_scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      List[str]     ,		__snake_case  :      int					)			->							Optional[int]:
											__magic_name__:	Tuple							=       bos_token_id






					def __call__(			self  :      Optional[int]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	int							=       jnp.full(scores.shape     ,		-float("""inf"""					)					)

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

											__magic_name__:	Tuple							=       jnp.where(_SCREAMING_SNAKE_CASE     ,		new_scores.at[:, self.bos_token_id].set(0					)     ,		_SCREAMING_SNAKE_CASE					)

											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      List[Any]     ,		__snake_case  :      int     ,		__snake_case  :      int					)			->							int:
											__magic_name__:	List[str]							=       max_length
											__magic_name__:	int							=       eos_token_id






					def __call__(			self  :      List[Any]     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Optional[int]							=       jnp.full(scores.shape     ,		-float("""inf"""					)					)

											__magic_name__:	Any							=       1 - jnp.bool_(cur_len - self.max_length + 1					)

											__magic_name__:	Union[str, Any]							=       jnp.where(_SCREAMING_SNAKE_CASE     ,		new_scores.at[:, self.eos_token_id].set(0					)     ,		_SCREAMING_SNAKE_CASE					)

											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      Optional[Any]     ,		__snake_case  :      int     ,		__snake_case  :      int					)			->							Optional[Any]:

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

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

											__magic_name__:	Any							=       min_length
											__magic_name__:	Tuple							=       eos_token_id






					def __call__(			self  :      str     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Optional[Any]							=       1 - jnp.clip(cur_len - self.min_length     ,		0     ,		1					)

											__magic_name__:	Union[str, Any]							=       jnp.where(_SCREAMING_SNAKE_CASE     ,		scores.at[:, self.eos_token_id].set(-float("""inf"""					)					)     ,		_SCREAMING_SNAKE_CASE					)

											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      Optional[Any]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      str					)			->							List[Any]:
											__magic_name__:	int							=       list(_SCREAMING_SNAKE_CASE					)
											__magic_name__:	Optional[Any]							=       begin_index






					def __call__(			self  :      Optional[Any]     ,		__snake_case  :      List[str]     ,		__snake_case  :      str     ,		__snake_case  :      int					)			->							Dict:
											__magic_name__:	Any							=       1 - jnp.bool_(cur_len - self.begin_index					)

											__magic_name__:	Optional[Any]							=       jnp.where(_SCREAMING_SNAKE_CASE     ,		scores.at[:, self.begin_suppress_tokens].set(-float("""inf"""					)					)     ,		_SCREAMING_SNAKE_CASE					)

											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      List[Any]     ,		__snake_case  :      list					)			->							Optional[int]:
											__magic_name__:	Optional[Any]							=       list(_SCREAMING_SNAKE_CASE					)






					def __call__(			self  :      int     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      jnp.ndarray     ,		__snake_case  :      int					)			->							jnp.ndarray:
											__magic_name__:	Optional[Any]							=       scores.at[..., self.suppress_tokens].set(-float("""inf"""					)					)

											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      List[str]     ,		__snake_case  :      int					)			->							str:
											__magic_name__:	Optional[Any]							=       dict(_SCREAMING_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.
											__magic_name__:	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:
																							__magic_name__:	int							=       force_token_array.at[index].set(_SCREAMING_SNAKE_CASE					)
											__magic_name__:	Tuple							=       jnp.intaa(_SCREAMING_SNAKE_CASE					)






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

											def _force_token(__snake_case  :      Optional[Any]					):
																	__magic_name__:	int							=       scores.shape[0]
																	__magic_name__:	Optional[int]							=       self.force_token_array[generation_idx]

																	__magic_name__:	int							=       jnp.ones_like(_SCREAMING_SNAKE_CASE     ,		dtype=scores.dtype					) * -float("""inf"""					)
																	__magic_name__:	Optional[int]							=       jnp.zeros((batch_size, 1)     ,		dtype=scores.dtype					)
																	__magic_name__:	Dict							=       lax.dynamic_update_slice(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		(0, current_token)					)
																	return new_scores

											__magic_name__:	Union[str, Any]							=       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(_SCREAMING_SNAKE_CASE					)     ,		lambda: scores     ,		)     ,		)
											return scores



class       __A		(   __SCREAMING_SNAKE_CASE			):

					def __init__(			self  :      Optional[Any]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Tuple     ,		__snake_case  :      Optional[Any]					)			->							int:
											__magic_name__:	List[str]							=       generate_config.eos_token_id
											__magic_name__:	Optional[int]							=       generate_config.no_timestamps_token_id
											__magic_name__:	Dict							=       generate_config.no_timestamps_token_id + 1

											__magic_name__:	Optional[Any]							=       decoder_input_length + 1

											if generate_config.is_multilingual:
																	# room for language token and task token
																	self.begin_index += 2
											if hasattr(_SCREAMING_SNAKE_CASE     ,		"""max_initial_timestamp_index"""					):
																	__magic_name__:	Optional[int]							=       generate_config.max_initial_timestamp_index
											else:
																	__magic_name__:	Optional[int]							=       model_config.vocab_size
											if self.max_initial_timestamp_index is None:
																	__magic_name__:	int							=       model_config.vocab_size






					def __call__(			self  :      Tuple     ,		__snake_case  :      Union[str, Any]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      int					)			->							Optional[int]:
											__magic_name__:	Optional[Any]							=       scores.at[:, self.no_timestamps_token_id].set(-float("""inf"""					)					)

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

																	__magic_name__:	Tuple							=       jnp.where((cur_len - self.begin_index) < 2     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)
																	__magic_name__:	int							=       jnp.where(
																	    input_ids_k[cur_len - 2] >= self.timestamp_begin     ,		_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE     ,		)

																	return jnp.where(
																	    _SCREAMING_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"""					)					)     ,		)     ,		_SCREAMING_SNAKE_CASE     ,		)

											__magic_name__:	int							=       jax.vmap(_SCREAMING_SNAKE_CASE					)(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)

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

											__magic_name__:	List[Any]							=       self.timestamp_begin + self.max_initial_timestamp_index

											__magic_name__:	Tuple							=       jnp.where(
											    _SCREAMING_SNAKE_CASE     ,		scores.at[:, last_allowed + 1 :].set(-float("""inf"""					)					)     ,		_SCREAMING_SNAKE_CASE     ,		)

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

											def handle_cumulative_probs(__snake_case  :      Any     ,		__snake_case  :      Any					):
																	__magic_name__:	Tuple							=       jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]     ,		axis=-1					)
																	__magic_name__:	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"""					)					)     ,		_SCREAMING_SNAKE_CASE     ,		)

											__magic_name__:	Dict							=       jax.vmap(_SCREAMING_SNAKE_CASE					)(_SCREAMING_SNAKE_CASE     ,		_SCREAMING_SNAKE_CASE					)

											return scores








"""simple docstring"""






from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


__snake_case				:  Tuple =       {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__snake_case				:  Optional[Any] =       [
	    'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'AutoformerForPrediction',
	    'AutoformerModel',
	    'AutoformerPreTrainedModel',
	]


if TYPE_CHECKING:
	from .configuration_autoformer import (
	    AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    AutoformerConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_autoformer import (
		    AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    AutoformerForPrediction,
		    AutoformerModel,
		    AutoformerPreTrainedModel,
		)

else:
	import sys

	__snake_case				:  Optional[int] =       _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)