






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


__magic_name__:					Dict							=     get_logger(__name__)


__magic_name__:					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 	snake_case__ :
	@add_start_docstrings(_lowerCamelCase			)
	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				List[str]:
							raise NotImplementedError(
							    F'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'			)





class 	snake_case__ :
	@add_start_docstrings(_lowerCamelCase			)
	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Dict:
							raise NotImplementedError(
							    F'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'			)





class 	snake_case__ (							a__				):
	@add_start_docstrings(_lowerCamelCase			)
	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       **lowerCAmelCase__			)							->				str:







							for processor in self:
													__magic_name__      :       int       						= inspect.signature(processor.__call__			).parameters
													if len(_lowerCamelCase			) > 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__      :       List[str]       						= processor(_lowerCamelCase      ,       _lowerCamelCase      ,       _lowerCamelCase      ,       **_lowerCamelCase			)
													else:
																			__magic_name__      :       Dict       						= processor(_lowerCamelCase      ,       _lowerCamelCase      ,       _lowerCamelCase			)
							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__			)							->				Dict:







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

							__magic_name__      :       int       						= temperature







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				List[Any]:
							__magic_name__      :       Optional[Any]       						= scores / self.temperature
							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__ = -float("""Inf"""			)      ,       lowerCAmelCase__ = 1			)							->				Any:







							if not isinstance(_lowerCamelCase      ,       _lowerCamelCase			) 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(_lowerCamelCase      ,       _lowerCamelCase			) 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__      :       int       						= top_p
							__magic_name__      :       Union[str, Any]       						= filter_value
							__magic_name__      :       Union[str, Any]       						= min_tokens_to_keep







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				List[Any]:
							__magic_name__      :       int       						= lax.top_k(_lowerCamelCase      ,       scores.shape[-1]			)

							__magic_name__      :       str       						= jnp.full_like(_lowerCamelCase      ,       self.filter_value			)
							__magic_name__      :       Dict       						= jax.nn.softmax(_lowerCamelCase      ,       axis=-1			).cumsum(axis=-1			)
							__magic_name__      :       Any       						= cumulative_probs < self.top_p

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

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

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

							return next_scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__ = -float("""Inf"""			)      ,       lowerCAmelCase__ = 1			)							->				Dict:







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

							__magic_name__      :       str       						= max(_lowerCamelCase      ,       _lowerCamelCase			)
							__magic_name__      :       Optional[int]       						= filter_value







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Any:
							__magic_name__      :       Dict       						= scores.shape
							__magic_name__      :       Optional[Any]       						= jnp.full(batch_size * vocab_size      ,       self.filter_value			)

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

							__magic_name__      :       Any       						= next_scores_flat.at[topk_indices_flat].set(_lowerCamelCase			)
							__magic_name__      :       Optional[Any]       						= next_scores_flat.reshape(_lowerCamelCase      ,       _lowerCamelCase			)
							return next_scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__			)							->				List[str]:
							__magic_name__      :       Any       						= bos_token_id







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

							__magic_name__      :       List[Any]       						= 1 - jnp.bool_(cur_len - 1			)

							__magic_name__      :       List[str]       						= jnp.where(_lowerCamelCase      ,       new_scores.at[:, self.bos_token_id].set(0			)      ,       _lowerCamelCase			)

							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Dict:
							__magic_name__      :       int       						= max_length
							__magic_name__      :       Any       						= eos_token_id







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

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

							__magic_name__      :       Optional[int]       						= jnp.where(_lowerCamelCase      ,       new_scores.at[:, self.eos_token_id].set(0			)      ,       _lowerCamelCase			)

							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				int:







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

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

							__magic_name__      :       List[str]       						= min_length
							__magic_name__      :       Optional[int]       						= eos_token_id







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				str:
							__magic_name__      :       int       						= 1 - jnp.clip(cur_len - self.min_length      ,       0      ,       1			)

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

							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				List[str]:
							__magic_name__      :       int       						= list(_lowerCamelCase			)
							__magic_name__      :       Tuple       						= begin_index







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Optional[Any]:
							__magic_name__      :       Any       						= 1 - jnp.bool_(cur_len - self.begin_index			)

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

							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__			)							->				Dict:
							__magic_name__      :       int       						= list(_lowerCamelCase			)







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Any:
							__magic_name__      :       Any       						= scores.at[..., self.suppress_tokens].set(-float("""inf"""			)			)

							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__			)							->				Optional[int]:
							__magic_name__      :       List[Any]       						= dict(_lowerCamelCase			)
							# 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__      :       int       						= 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__      :       List[str]       						= force_token_array.at[index].set(_lowerCamelCase			)
							__magic_name__      :       List[Any]       						= jnp.intaa(_lowerCamelCase			)







	def __call__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				Union[str, Any]:







							def _force_token(lowerCAmelCase__			):
													__magic_name__      :       List[str]       						= scores.shape[0]
													__magic_name__      :       List[Any]       						= self.force_token_array[generation_idx]

													__magic_name__      :       Optional[Any]       						= jnp.ones_like(_lowerCamelCase      ,       dtype=scores.dtype			) * -float("""inf"""			)
													__magic_name__      :       Union[str, Any]       						= jnp.zeros((batch_size, 1)      ,       dtype=scores.dtype			)
													__magic_name__      :       Union[str, Any]       						= lax.dynamic_update_slice(_lowerCamelCase      ,       _lowerCamelCase      ,       (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(_lowerCamelCase			)      ,       lambda: scores      ,       )      ,       )
							return scores





class 	snake_case__ (							a__				):
	def __init__(    self      ,       lowerCAmelCase__      ,       lowerCAmelCase__      ,       lowerCAmelCase__			)							->				int:
							__magic_name__      :       Tuple       						= generate_config.eos_token_id
							__magic_name__      :       Any       						= generate_config.no_timestamps_token_id
							__magic_name__      :       Union[str, Any]       						= generate_config.no_timestamps_token_id + 1

							__magic_name__      :       List[str]       						= decoder_input_length + 1

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







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

							def handle_pairs(lowerCAmelCase__      ,       lowerCAmelCase__			):
													__magic_name__      :       Optional[int]       						= jnp.where((cur_len - self.begin_index) >= 1      ,       _lowerCamelCase      ,       _lowerCamelCase			)
													__magic_name__      :       int       						= jnp.where(
													    input_ids_k[cur_len - 1] >= self.timestamp_begin      ,       True and last_was_timestamp      ,       _lowerCamelCase      ,       )

													__magic_name__      :       Union[str, Any]       						= jnp.where((cur_len - self.begin_index) < 2      ,       _lowerCamelCase      ,       _lowerCamelCase			)
													__magic_name__      :       Any       						= jnp.where(
													    input_ids_k[cur_len - 2] >= self.timestamp_begin      ,       _lowerCamelCase      ,       _lowerCamelCase      ,       )

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

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

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

							__magic_name__      :       str       						= self.timestamp_begin + self.max_initial_timestamp_index

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

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

							def handle_cumulative_probs(lowerCAmelCase__      ,       lowerCAmelCase__			):
													__magic_name__      :       Optional[Any]       						= jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]      ,       axis=-1			)
													__magic_name__      :       Union[str, Any]       						= 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"""			)			)      ,       _lowerCamelCase      ,       )

							__magic_name__      :       Any       						= jax.vmap(_lowerCamelCase			)(_lowerCamelCase      ,       _lowerCamelCase			)

							return scores








from __future__ import annotations

import math

import numpy as np
from numpy.linalg import norm

def    UpperCamelCase						(	_A,   _A			):






						"""simple docstring"""







						return math.sqrt(sum(pow(a - b,   2			) for a, b in zip(_A,   _A			)			)			)

def    UpperCamelCase						(	_A,   _A			):






						"""simple docstring"""







						if dataset.ndim != value_array.ndim:
												__magic_name__      :       str       						= (
												    """Wrong input data's dimensions... """
												    f'dataset : {dataset.ndim}, value_array : {value_array.ndim}'
												)
												raise ValueError(_A			)

						try:
												if dataset.shape[1] != value_array.shape[1]:
																		__magic_name__      :       Optional[Any]       						= (
																		    """Wrong input data's shape... """
																		    f'dataset : {dataset.shape[1]}, value_array : {value_array.shape[1]}'
																		)
																		raise ValueError(_A			)
						except IndexError:
												if dataset.ndim != value_array.ndim:
																		raise TypeError("""Wrong shape"""			)

						if dataset.dtype != value_array.dtype:
												__magic_name__      :       List[Any]       						= (
												    """Input data have different datatype... """
												    f'dataset : {dataset.dtype}, value_array : {value_array.dtype}'
												)
												raise TypeError(_A			)

						__magic_name__      :       Dict       						= []

						for value in value_array:
												__magic_name__      :       Tuple       						= euclidean(_A,   dataset[0]			)
												__magic_name__      :       Any       						= dataset[0].tolist()

												for dataset_value in dataset[1:]:
																		__magic_name__      :       Any       						= euclidean(_A,   _A			)

																		if dist > temp_dist:
																								__magic_name__      :       Dict       						= temp_dist
																								__magic_name__      :       Any       						= dataset_value.tolist()

												answer.append([vector, dist]			)

						return answer

def    UpperCamelCase						(	_A,   _A			):






						"""simple docstring"""







						return np.dot(_A,   _A			) / (norm(_A			) * norm(_A			))


if __name__ == "__main__":
		import doctest

		doctest.testmod()
