



"""simple docstring"""






import argparse
import importlib
from pathlib import Path


# Test all the extensions added in the setup
__snake_case				:  str =       [
    'kernels/rwkv/wkv_cuda.cu',
    'kernels/rwkv/wkv_op.cpp',
    'kernels/deformable_detr/ms_deform_attn.h',
    'kernels/deformable_detr/cuda/ms_deform_im2col_cuda.cuh',
    'models/graphormer/algos_graphormer.pyx',
]


def      _lowercase							(    __snake_case      )		->  int:
			# Test all the extensions added in the setup
			for file in FILES_TO_FIND:
						if not (transformers_path / file).exists():
									return False
			return True


if __name__ == "__main__":
	__snake_case				:  List[str] =       argparse.ArgumentParser()
	parser.add_argument('--check_lib', action='store_true', help='Whether to check the build or the actual package.')
	__snake_case				:  Any =       parser.parse_args()
	if args.check_lib:
		__snake_case				:  List[Any] =       importlib.import_module('transformers')
		__snake_case				:  int =       Path(transformers_module.__file__).parent
	else:
		__snake_case				:  Optional[Any] =       Path.cwd() / 'build/lib/transformers'
	if not test_custom_files_are_present(transformers_path):
		raise ValueError('The built release does not contain the custom files. Fix this before going further!')




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


__snake_case				:  List[str] =       get_logger(__name__)


__snake_case				:  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      A__  :




	'''simple docstring'''





	@add_start_docstrings(_SCREAMING_SNAKE_CASE)
	def __call__(	self:						Tuple			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray)							->			jnp.ndarray:

				"""simple docstring"""

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






class      A__  :




	'''simple docstring'''





	@add_start_docstrings(_SCREAMING_SNAKE_CASE)
	def __call__(	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray)							->			jnp.ndarray:

				"""simple docstring"""

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






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





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

				"""simple docstring"""

				for processor in self:
							__lowerCAmelCase :    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.""")
										__lowerCAmelCase :    int				=			processor(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     **_SCREAMING_SNAKE_CASE)
							else:
										__lowerCAmelCase :    Union[str, Any]				=			processor(_SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
				return scores






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





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

				"""simple docstring"""

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

				__lowerCAmelCase :    List[str]				=			temperature







	def __call__(	self:						Dict			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						int)							->			jnp.ndarray:

				"""simple docstring"""

				__lowerCAmelCase :    Union[str, Any]				=			scores / self.temperature
				return scores






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





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

				"""simple docstring"""

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

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







	def __call__(	self:						Union[str, Any]			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						int)							->			jnp.ndarray:

				"""simple docstring"""

				__lowerCAmelCase		,      __lowerCAmelCase :    Optional[Any]				=			lax.top_k(_SCREAMING_SNAKE_CASE			,     scores.shape[-1])

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

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

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

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

				return next_scores






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





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

				"""simple docstring"""

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

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







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

				"""simple docstring"""

				__lowerCAmelCase		,      __lowerCAmelCase :    Any				=			scores.shape
				__lowerCAmelCase :    Union[str, Any]				=			jnp.full(batch_size * vocab_size			,     self.filter_value)

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

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






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	def __init__(	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						int)							->			Optional[int]:

				"""simple docstring"""

				__lowerCAmelCase :    Tuple				=			bos_token_id







	def __call__(	self:						Optional[int]			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						int)							->			jnp.ndarray:

				"""simple docstring"""

				__lowerCAmelCase :    int				=			jnp.full(scores.shape			,     -float("inf"))

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

				__lowerCAmelCase :    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				):




	'''simple docstring'''





	def __init__(	self:						List[Any]			,     _SCREAMING_SNAKE_CASE:						int			,     _SCREAMING_SNAKE_CASE:						int)							->			int:

				"""simple docstring"""

				__lowerCAmelCase :    List[str]				=			max_length
				__lowerCAmelCase :    int				=			eos_token_id







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

				"""simple docstring"""

				__lowerCAmelCase :    Optional[int]				=			jnp.full(scores.shape			,     -float("inf"))

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

				__lowerCAmelCase :    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				):




	'''simple docstring'''





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

				"""simple docstring"""

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

				__lowerCAmelCase :    Any				=			min_length
				__lowerCAmelCase :    Tuple				=			eos_token_id







	def __call__(	self:						str			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						int)							->			jnp.ndarray:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			1 - jnp.clip(cur_len - self.min_length			,     0			,     1)

				__lowerCAmelCase :    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				):




	'''simple docstring'''





	def __init__(	self:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						str)							->			List[Any]:

				"""simple docstring"""

				__lowerCAmelCase :    int				=			list(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Optional[Any]				=			begin_index







	def __call__(	self:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						List[str]			,     _SCREAMING_SNAKE_CASE:						str			,     _SCREAMING_SNAKE_CASE:						int)							->			Dict:

				"""simple docstring"""

				__lowerCAmelCase :    Any				=			1 - jnp.bool_(cur_len - self.begin_index)

				__lowerCAmelCase :    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				):




	'''simple docstring'''





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

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			list(_SCREAMING_SNAKE_CASE)







	def __call__(	self:						int			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						jnp.ndarray			,     _SCREAMING_SNAKE_CASE:						int)							->			jnp.ndarray:

				"""simple docstring"""

				__lowerCAmelCase :    Optional[Any]				=			scores.at[..., self.suppress_tokens].set(-float("inf"))

				return scores






class      A__  (    __SCREAMING_SNAKE_CASE				):




	'''simple docstring'''





	def __init__(	self:						List[str]			,     _SCREAMING_SNAKE_CASE:						int)							->			str:

				"""simple docstring"""

				__lowerCAmelCase :    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.
				__lowerCAmelCase :    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:
										__lowerCAmelCase :    int				=			force_token_array.at[index].set(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase :    Tuple				=			jnp.intaa(_SCREAMING_SNAKE_CASE)







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

				"""simple docstring"""

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

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

				__lowerCAmelCase :    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				):




	'''simple docstring'''





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

				"""simple docstring"""

				__lowerCAmelCase :    List[str]				=			generate_config.eos_token_id
				__lowerCAmelCase :    Optional[int]				=			generate_config.no_timestamps_token_id
				__lowerCAmelCase :    Dict				=			generate_config.no_timestamps_token_id + 1

				__lowerCAmelCase :    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"):
							__lowerCAmelCase :    Optional[int]				=			generate_config.max_initial_timestamp_index
				else:
							__lowerCAmelCase :    Optional[int]				=			model_config.vocab_size
				if self.max_initial_timestamp_index is None:
							__lowerCAmelCase :    int				=			model_config.vocab_size







	def __call__(	self:						Tuple			,     _SCREAMING_SNAKE_CASE:						Union[str, Any]			,     _SCREAMING_SNAKE_CASE:						Optional[Any]			,     _SCREAMING_SNAKE_CASE:						int)							->			Optional[int]:

				"""simple docstring"""

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

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

							__lowerCAmelCase :    Tuple				=			jnp.where((cur_len - self.begin_index) < 2			,     _SCREAMING_SNAKE_CASE			,     _SCREAMING_SNAKE_CASE)
							__lowerCAmelCase :    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			,     )

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

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

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

				__lowerCAmelCase :    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
				__lowerCAmelCase :    List[str]				=			jax.nn.log_softmax(_SCREAMING_SNAKE_CASE			,     axis=-1)

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

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

				return scores