



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     :							Union[str, Any]											=				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  __magic_name__   :






		@add_start_docstrings(a_      )
		def __call__(						self      :     Union[str, Any]				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     List[Any]      )				->					Tuple:
								'''simple docstring'''







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





class  __magic_name__   :






		@add_start_docstrings(a_      )
		def __call__(						self      :     List[Any]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Dict      )				->					Tuple:
								'''simple docstring'''







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





class  __magic_name__   (							__lowerCamelCase						):






		@add_start_docstrings(a_      )
		def __call__(						self      :     Any				,		UpperCamelCase__      :     Tuple				,		UpperCamelCase__      :     Tuple				,		UpperCamelCase__      :     Optional[Any]				,		**UpperCamelCase__      :     List[str]      )				->					Optional[int]:
								'''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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Optional[int]				,		UpperCamelCase__      :     str      )				->					List[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      :     str				,		UpperCamelCase__      :     Optional[Any]				,		UpperCamelCase__      :     List[Any]				,		UpperCamelCase__      :     Tuple      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     scores / self.temperature
								return scores





class  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Any				,		UpperCamelCase__      :     Optional[Any]				,		UpperCamelCase__      :     Any = -float("Inf"      )				,		UpperCamelCase__      :     Tuple = 1      )				->					List[str]:
								'''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      :     List[Any]				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     List[Any]				,		UpperCamelCase__      :     List[str]      )				->					int:
								'''simple docstring'''







								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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Dict				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     List[str] = -float("Inf"      )				,		UpperCamelCase__      :     Any = 1      )				->					int:
								'''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      :     Any				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Union[str, Any]      )				->					Tuple:
								'''simple docstring'''







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

								UpperCAmelCase         =     min(self.top_k				,		scores.shape[-1]      )  # Safety check
								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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Optional[int]				,		UpperCamelCase__      :     str      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     bos_token_id







		def __call__(						self      :     Tuple				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Dict      )				->					Tuple:
								'''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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     str				,		UpperCamelCase__      :     Tuple				,		UpperCamelCase__      :     Tuple      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     eos_token_id







		def __call__(						self      :     Union[str, Any]				,		UpperCamelCase__      :     List[str]				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Tuple      )				->					Tuple:
								'''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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Optional[int]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Optional[int]:
								'''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      :     str				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     List[str]				,		UpperCamelCase__      :     Dict      )				->					str:
								'''simple docstring'''







								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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     List[str]				,		UpperCamelCase__      :     Tuple				,		UpperCamelCase__      :     List[Any]      )				->					Dict:
								'''simple docstring'''







								UpperCAmelCase         =     list(a_      )
								UpperCAmelCase         =     begin_index







		def __call__(						self      :     List[str]				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Optional[Any]      )				->					int:
								'''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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     List[Any]				,		UpperCamelCase__      :     Tuple      )				->					int:
								'''simple docstring'''







								UpperCAmelCase         =     list(a_      )







		def __call__(						self      :     Tuple				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     List[str]      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     scores.at[..., self.suppress_tokens].set(-float("inf"      )      )

								return scores





class  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Tuple				,		UpperCamelCase__      :     Dict      )				->					Dict:
								'''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      :     Optional[Any]				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Optional[int]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Any:
								'''simple docstring'''







								def _force_token(UpperCamelCase__      :     int      ):
														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  __magic_name__   (							__lowerCamelCase						):






		def __init__(						self      :     Any				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Union[str, Any]				,		UpperCamelCase__      :     Union[str, Any]      )				->					Optional[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      :     Optional[int]				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     Dict				,		UpperCamelCase__      :     Tuple      )				->					str:
								'''simple docstring'''







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

								def handle_pairs(UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Dict      ):
														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(UpperCamelCase__      :     List[Any]				,		UpperCamelCase__      :     Tuple      ):
														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 typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


__lowerCamelCase     :							Optional[int]  							=				{
    "configuration_roc_bert": ["ROC_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "RoCBertConfig"],
    "tokenization_roc_bert": ["RoCBertTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		pass

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCamelCase     :							List[Any]  							=				[
		    "ROC_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "RoCBertForCausalLM",
		    "RoCBertForMaskedLM",
		    "RoCBertForMultipleChoice",
		    "RoCBertForPreTraining",
		    "RoCBertForQuestionAnswering",
		    "RoCBertForSequenceClassification",
		    "RoCBertForTokenClassification",
		    "RoCBertLayer",
		    "RoCBertModel",
		    "RoCBertPreTrainedModel",
		    "load_tf_weights_in_roc_bert",
		]

if TYPE_CHECKING:
		from .configuration_roc_bert import ROC_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RoCBertConfig
		from .tokenization_roc_bert import RoCBertTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				raise OptionalDependencyNotAvailable()

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_roc_bert import (
				    ROC_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
				    RoCBertForCausalLM,
				    RoCBertForMaskedLM,
				    RoCBertForMultipleChoice,
				    RoCBertForPreTraining,
				    RoCBertForQuestionAnswering,
				    RoCBertForSequenceClassification,
				    RoCBertForTokenClassification,
				    RoCBertLayer,
				    RoCBertModel,
				    RoCBertPreTrainedModel,
				    load_tf_weights_in_roc_bert,
				)


else:
		import sys

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



