



'''simple docstring'''





from scipy.stats import pearsonr

import datasets


SCREAMING_SNAKE_CASE__       =      "\nPearson correlation coefficient and p-value for testing non-correlation.\nThe Pearson correlation coefficient measures the linear relationship between two datasets. The calculation of the p-value relies on the assumption that each dataset is normally distributed. Like other correlation coefficients, this one varies between -1 and +1 with 0 implying no correlation. Correlations of -1 or +1 imply an exact linear relationship. Positive correlations imply that as x increases, so does y. Negative correlations imply that as x increases, y decreases.\nThe p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets.\n"


SCREAMING_SNAKE_CASE__       =      "\nArgs:\n    predictions (`list` of `int`): Predicted class labels, as returned by a model.\n    references (`list` of `int`): Ground truth labels.\n    return_pvalue (`boolean`): If `True`, returns the p-value, along with the correlation coefficient. If `False`, returns only the correlation coefficient. Defaults to `False`.\n\nReturns:\n    pearsonr (`float`): Pearson correlation coefficient. Minimum possible value is -1. Maximum possible value is 1. Values of 1 and -1 indicate exact linear positive and negative relationships, respectively. A value of 0 implies no correlation.\n    p-value (`float`): P-value, which roughly indicates the probability of an The p-value roughly indicates the probability of an uncorrelated system producing datasets that have a Pearson correlation at least as extreme as the one computed from these datasets. Minimum possible value is 0. Maximum possible value is 1. Higher values indicate higher probabilities.\n\nExamples:\n\n    Example 1-A simple example using only predictions and references.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5])\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n\n    Example 2-The same as Example 1, but that also returns the `p-value`.\n        >>> pearsonr_metric = datasets.load_metric(\"pearsonr\")\n        >>> results = pearsonr_metric.compute(predictions=[10, 9, 2.5, 6, 4], references=[1, 2, 3, 4, 5], return_pvalue=True)\n        >>> print(sorted(list(results.keys())))\n        ['p-value', 'pearsonr']\n        >>> print(round(results['pearsonr'], 2))\n        -0.74\n        >>> print(round(results['p-value'], 2))\n        0.15\n"


SCREAMING_SNAKE_CASE__       =      "\n@article{2020SciPy-NMeth,\nauthor  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and\n      Haberland, Matt and Reddy, Tyler and Cournapeau, David and\n      Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and\n      Bright, Jonathan and {van der Walt}, St{\'e}fan J. and\n      Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and\n      Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and\n      Kern, Robert and Larson, Eric and Carey, C J and\n      Polat, Ilhan and Feng, Yu and Moore, Eric W. and\n      {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and\n      Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and\n      Harris, Charles R. and Archibald, Anne M. and\n      Ribeiro, Antonio H. and Pedregosa, Fabian and\n      {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},\ntitle   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific\n      Computing in Python}},\njournal = {Nature Methods},\nyear    = {2020},\nvolume  = {17},\npages   = {261--272},\nadsurl  = {https://rdcu.be/b08Wh},\ndoi = {10.1038/s41592-019-0686-2},\n}\n"


@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,     _KWARGS_DESCRIPTION      )
class      __lowerCAmelCase     (		datasets.Metric      ):
					"""simple docstring"""






					def     _a							(     self					:       List[str]					):






								"""simple docstring"""


								return datasets.MetricInfo(
								    description=_DESCRIPTION				,					citation=_CITATION				,					inputs_description=_KWARGS_DESCRIPTION				,					features=datasets.Features(
								        {
								            'predictions': datasets.Value('float'					),
								            'references': datasets.Value('float'					),
								        }					)				,					reference_urls=['https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.pearsonr.html']				,					)



					def     _a							(     self					:       Optional[Any]				,					_snake_case					:       Optional[Any]				,					_snake_case					:       Any				,					_snake_case					:       Dict=False					):






								"""simple docstring"""


								if return_pvalue:
											A__       =    pearsonr(__UpperCamelCase				,					__UpperCamelCase					)
											return {"pearsonr": results[0], "p-value": results[1]}
								else:
											return {"pearsonr": float(pearsonr(__UpperCamelCase				,					__UpperCamelCase					)[0]					)}









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


SCREAMING_SNAKE_CASE__       =      get_logger(__name__)


SCREAMING_SNAKE_CASE__       =      r'''
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

'''


class      __lowerCAmelCase     :
					"""simple docstring"""






					@add_start_docstrings(_snake_case					)
					def __call__(     self					:       Optional[int]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray					):






								"""simple docstring"""


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







class      __lowerCAmelCase     :
					"""simple docstring"""






					@add_start_docstrings(_snake_case					)
					def __call__(     self					:       List[Any]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray					):






								"""simple docstring"""


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







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					@add_start_docstrings(_snake_case					)
					def __call__(     self					:       Any				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int				,					**_snake_case					:       Optional[int]					):






								"""simple docstring"""


								for processor in self:
											A__       =    inspect.signature(processor.__call__					).parameters
											if len(_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.'''					)
														A__       =    processor(_snake_case				,					_snake_case				,					_snake_case				,					**_snake_case					)
											else:
														A__       =    processor(_snake_case				,					_snake_case				,					_snake_case					)
								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Any				,					_snake_case					:       float					):






								"""simple docstring"""


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

								A__       =    temperature



					def __call__(     self					:       str				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    scores / self.temperature
								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Optional[Any]				,					_snake_case					:       float				,					_snake_case					:       float = -float('Inf'					)				,					_snake_case					:       int = 1					):






								"""simple docstring"""


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

								A__       =    top_p
								A__       =    filter_value
								A__       =    min_tokens_to_keep



					def __call__(     self					:       str				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__     , A__       =    lax.top_k(_snake_case				,					scores.shape[-1]					)

								A__       =    jnp.full_like(_snake_case				,					self.filter_value					)
								A__       =    jax.nn.softmax(_snake_case				,					axis=-1					).cumsum(axis=-1					)
								A__       =    cumulative_probs < self.top_p

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

								# min tokens to keep
								A__       =    score_mask.at[:, : self.min_tokens_to_keep].set(_snake_case					)

								A__       =    jnp.where(_snake_case				,					_snake_case				,					_snake_case					)
								A__       =    jax.lax.sort_key_val(_snake_case				,					_snake_case					)[-1]

								return next_scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Union[str, Any]				,					_snake_case					:       int				,					_snake_case					:       float = -float('Inf'					)				,					_snake_case					:       int = 1					):






								"""simple docstring"""


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

								A__       =    max(_snake_case				,					_snake_case					)
								A__       =    filter_value



					def __call__(     self					:       Optional[Any]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


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

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

								A__       =    next_scores_flat.at[topk_indices_flat].set(_snake_case					)
								A__       =    next_scores_flat.reshape(_snake_case				,					_snake_case					)
								return next_scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Any				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    bos_token_id



					def __call__(     self					:       Optional[int]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    jnp.full(scores.shape				,					-float('inf'					)					)

								A__       =    1 - jnp.bool_(cur_len - 1					)

								A__       =    jnp.where(_snake_case				,					new_scores.at[:, self.bos_token_id].set(0					)				,					_snake_case					)

								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Any				,					_snake_case					:       int				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    max_length
								A__       =    eos_token_id



					def __call__(     self					:       List[Any]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    jnp.full(scores.shape				,					-float('inf'					)					)

								A__       =    1 - jnp.bool_(cur_len - self.max_length + 1					)

								A__       =    jnp.where(_snake_case				,					new_scores.at[:, self.eos_token_id].set(0					)				,					_snake_case					)

								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Dict				,					_snake_case					:       int				,					_snake_case					:       int					):






								"""simple docstring"""


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

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

								A__       =    min_length
								A__       =    eos_token_id



					def __call__(     self					:       int				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    1 - jnp.clip(cur_len - self.min_length				,					0				,					1					)

								A__       =    jnp.where(_snake_case				,					scores.at[:, self.eos_token_id].set(-float('inf'					)					)				,					_snake_case					)

								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       int				,					_snake_case					:       Tuple				,					_snake_case					:       Union[str, Any]					):






								"""simple docstring"""


								A__       =    list(_snake_case					)
								A__       =    begin_index



					def __call__(     self					:       Union[str, Any]				,					_snake_case					:       Optional[Any]				,					_snake_case					:       str				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    1 - jnp.bool_(cur_len - self.begin_index					)

								A__       =    jnp.where(_snake_case				,					scores.at[:, self.begin_suppress_tokens].set(-float('inf'					)					)				,					_snake_case					)

								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       int				,					_snake_case					:       list					):






								"""simple docstring"""


								A__       =    list(_snake_case					)



					def __call__(     self					:       List[Any]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								A__       =    scores.at[..., self.suppress_tokens].set(-float('inf'					)					)

								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       List[str]				,					_snake_case					:       Optional[Any]					):






								"""simple docstring"""


								A__       =    dict(_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.
								A__       =    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:
														A__       =    force_token_array.at[index].set(_snake_case					)
								A__       =    jnp.intaa(_snake_case					)



					def __call__(     self					:       List[Any]				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       jnp.ndarray				,					_snake_case					:       int					):






								"""simple docstring"""


								def _force_token(_snake_case					:       Dict					):
											A__       =    scores.shape[0]
											A__       =    self.force_token_array[generation_idx]

											A__       =    jnp.ones_like(_snake_case				,					dtype=scores.dtype					) * -float('inf'					)
											A__       =    jnp.zeros((batch_size, 1)				,					dtype=scores.dtype					)
											A__       =    lax.dynamic_update_slice(_snake_case				,					_snake_case				,					(0, current_token)					)
											return new_scores

								A__       =    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(_snake_case					)				,					lambda: scores				,					)				,					)
								return scores







class      __lowerCAmelCase     (		UpperCAmelCase_      ):
					"""simple docstring"""






					def __init__(     self					:       Union[str, Any]				,					_snake_case					:       Union[str, Any]				,					_snake_case					:       Dict				,					_snake_case					:       List[Any]					):






								"""simple docstring"""


								A__       =    generate_config.eos_token_id
								A__       =    generate_config.no_timestamps_token_id
								A__       =    generate_config.no_timestamps_token_id + 1

								A__       =    decoder_input_length + 1

								if generate_config.is_multilingual:
											# room for language token and task token
											self.begin_index += 2
								if hasattr(_snake_case				,					'max_initial_timestamp_index'					):
											A__       =    generate_config.max_initial_timestamp_index
								else:
											A__       =    model_config.vocab_size
								if self.max_initial_timestamp_index is None:
											A__       =    model_config.vocab_size



					def __call__(     self					:       Tuple				,					_snake_case					:       List[Any]				,					_snake_case					:       Dict				,					_snake_case					:       Dict					):






								"""simple docstring"""


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

								def handle_pairs(_snake_case					:       Dict				,					_snake_case					:       str					):
											A__       =    jnp.where((cur_len - self.begin_index) >= 1				,					_snake_case				,					_snake_case					)
											A__       =    jnp.where(
											    input_ids_k[cur_len - 1] >= self.timestamp_begin				,					True and last_was_timestamp				,					_snake_case				,					)

											A__       =    jnp.where((cur_len - self.begin_index) < 2				,					_snake_case				,					_snake_case					)
											A__       =    jnp.where(
											    input_ids_k[cur_len - 2] >= self.timestamp_begin				,					_snake_case				,					_snake_case				,					)

											return jnp.where(
											    _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'					)					)				,					)				,					_snake_case				,					)

								A__       =    jax.vmap(_snake_case					)(_snake_case				,					_snake_case					)

								A__       =    jnp.where(cur_len == self.begin_index				,					_snake_case				,					_snake_case					)
								A__       =    jnp.where(
								    self.max_initial_timestamp_index is not None				,					True and apply_max_initial_timestamp				,					_snake_case				,					)

								A__       =    self.timestamp_begin + self.max_initial_timestamp_index

								A__       =    jnp.where(
								    _snake_case				,					scores.at[:, last_allowed + 1 :].set(-float('inf'					)					)				,					_snake_case				,					)

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

								def handle_cumulative_probs(_snake_case					:       List[Any]				,					_snake_case					:       Union[str, Any]					):
											A__       =    jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]				,					axis=-1					)
											A__       =    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'					)					)				,					_snake_case				,					)

								A__       =    jax.vmap(_snake_case					)(_snake_case				,					_snake_case					)

								return scores




