
'''simple docstring'''







import unittest

import numpy as np

from transformers import is_flax_available
from transformers.testing_utils import require_flax

from ..test_modeling_flax_common import ids_tensor


if is_flax_available():
					import jax
					import jax.numpy as jnp

					from transformers.generation import (
					    FlaxForcedBOSTokenLogitsProcessor,
					    FlaxForcedEOSTokenLogitsProcessor,
					    FlaxLogitsProcessorList,
					    FlaxMinLengthLogitsProcessor,
					    FlaxTemperatureLogitsWarper,
					    FlaxTopKLogitsWarper,
					    FlaxTopPLogitsWarper,
					)


@require_flax
class 		_UpperCAmelCase   (      unittest.TestCase		):


							def 				lowerCamelCase       ( self					,     lowerCAmelCase_					,     lowerCAmelCase_						):





								"""simple docstring"""






								_snake_case     		=	jnp.ones((batch_size, length)						) / length
								return scores


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	None
								_snake_case     		=	20

								_snake_case     		=	self._get_uniform_logits(batch_size=2					,     length=UpperCAmelCase__						)

								# tweak scores to not be uniform anymore
								_snake_case     		=	scores.at[1, 5].set((1 / length) + 0.1						)  # peak, 1st batch
								_snake_case     		=	scores.at[1, 10].set((1 / length) - 0.4						)  # valley, 1st batch

								# compute softmax
								_snake_case     		=	jax.nn.softmax(UpperCAmelCase__					,     axis=-1						)

								_snake_case     		=	FlaxTemperatureLogitsWarper(temperature=0.5						)
								_snake_case     		=	FlaxTemperatureLogitsWarper(temperature=1.3						)

								_snake_case     		=	jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__					,     scores.copy()					,     cur_len=UpperCAmelCase__						)					,     axis=-1						)
								_snake_case     		=	jax.nn.softmax(temp_dist_warper_smoother(UpperCAmelCase__					,     scores.copy()					,     cur_len=UpperCAmelCase__						)					,     axis=-1						)

								# uniform distribution stays uniform
								self.assertTrue(jnp.allclose(probs[0, :]					,     warped_prob_sharp[0, :]					,     atol=1E-3						)						)
								self.assertTrue(jnp.allclose(probs[0, :]					,     warped_prob_smooth[0, :]					,     atol=1E-3						)						)

								# sharp peaks get higher, valleys get lower
								self.assertLess(probs[1, :].max()					,     warped_prob_sharp[1, :].max()						)
								self.assertGreater(probs[1, :].min()					,     warped_prob_sharp[1, :].min()						)

								# smooth peaks get lower, valleys get higher
								self.assertGreater(probs[1, :].max()					,     warped_prob_smooth[1, :].max()						)
								self.assertLess(probs[1, :].min()					,     warped_prob_smooth[1, :].min()						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	None
								_snake_case     		=	10
								_snake_case     		=	2

								# create ramp distribution
								_snake_case     		=	np.broadcast_to(np.arange(UpperCAmelCase__						)[None, :]					,     (batch_size, vocab_size)						).copy()
								_snake_case     		=	ramp_logits[1:, : vocab_size // 2] + vocab_size

								_snake_case     		=	FlaxTopKLogitsWarper(3						)

								_snake_case     		=	top_k_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)

								# check that correct tokens are filtered
								self.assertListEqual(jnp.isinf(scores[0]						).tolist()					,     7 * [True] + 3 * [False]						)
								self.assertListEqual(jnp.isinf(scores[1]						).tolist()					,     2 * [True] + 3 * [False] + 5 * [True]						)

								# check special case
								_snake_case     		=	5
								_snake_case     		=	FlaxTopKLogitsWarper(top_k=1					,     filter_value=0.0					,     min_tokens_to_keep=3						)

								_snake_case     		=	np.broadcast_to(np.arange(UpperCAmelCase__						)[None, :]					,     (batch_size, length)						).copy()
								_snake_case     		=	top_k_warp_safety_check(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)

								# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
								self.assertListEqual((scores == 0.0).sum(axis=-1						).tolist()					,     [2, 2]						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	None
								_snake_case     		=	10
								_snake_case     		=	2

								# create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
								_snake_case     		=	np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.15, 0.3, 0.3, 0.25]]						)						)

								_snake_case     		=	FlaxTopPLogitsWarper(0.8						)
								_snake_case     		=	np.exp(top_p_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)						)

								# dist should be filtered to keep min num values so that sum is >= top_p
								# exp (-inf) => 0
								_snake_case     		=	np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]						)
								self.assertTrue(np.allclose(UpperCAmelCase__					,     UpperCAmelCase__					,     atol=1E-3						)						)

								# check edge cases with negative and extreme logits
								_snake_case     		=	np.broadcast_to(np.arange(UpperCAmelCase__						)[None, :]					,     (batch_size, vocab_size)						).copy() - (
								    vocab_size // 2
								)

								# make ramp_logits more extreme
								_snake_case     		=	ramp_logits[1] * 100.0

								# make sure at least 2 tokens are kept
								_snake_case     		=	FlaxTopPLogitsWarper(0.9					,     min_tokens_to_keep=2					,     filter_value=0.0						)
								_snake_case     		=	top_p_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)

								# first batch should keep three tokens, second batch would keep only 1, but due to `min_tokens_to_keep=2` keeps 2.
								self.assertListEqual((filtered_dist != 0.0).sum(axis=-1						).tolist()					,     [3, 2]						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	20
								_snake_case     		=	4
								_snake_case     		=	0

								_snake_case     		=	FlaxMinLengthLogitsProcessor(min_length=10					,     eos_token_id=UpperCAmelCase__						)

								# check that min length is applied at length 5
								_snake_case     		=	ids_tensor((batch_size, 20)					,     vocab_size=20						)
								_snake_case     		=	5
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	min_dist_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()					,     4 * [-float('inf'						)]						)

								# check that min length is not applied anymore at length 15
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	15
								_snake_case     		=	min_dist_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertFalse(jnp.isinf(UpperCAmelCase__						).any()						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	20
								_snake_case     		=	4
								_snake_case     		=	0

								_snake_case     		=	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__						)

								# check that all scores are -inf except the bos_token_id score
								_snake_case     		=	ids_tensor((batch_size, 1)					,     vocab_size=20						)
								_snake_case     		=	1
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	logits_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertTrue(jnp.isneginf(scores[:, bos_token_id + 1 :]						).all()						)
								self.assertListEqual(scores[:, bos_token_id].tolist()					,     4 * [0]						)  # score for bos_token_id shold be zero

								# check that bos_token_id is not forced if current length is greater than 1
								_snake_case     		=	3
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	logits_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertFalse(jnp.isinf(UpperCAmelCase__						).any()						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	20
								_snake_case     		=	4
								_snake_case     		=	0
								_snake_case     		=	5

								_snake_case     		=	FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__					,     eos_token_id=UpperCAmelCase__						)

								# check that all scores are -inf except the eos_token_id when max_length is reached
								_snake_case     		=	ids_tensor((batch_size, 4)					,     vocab_size=20						)
								_snake_case     		=	4
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	logits_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertTrue(jnp.isneginf(scores[:, eos_token_id + 1 :]						).all()						)
								self.assertListEqual(scores[:, eos_token_id].tolist()					,     4 * [0]						)  # score for eos_token_id should be zero

								# check that eos_token_id is not forced if max_length is not reached
								_snake_case     		=	3
								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	logits_processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								self.assertFalse(jnp.isinf(UpperCAmelCase__						).any()						)


							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	4
								_snake_case     		=	10
								_snake_case     		=	15
								_snake_case     		=	2
								_snake_case     		=	1
								_snake_case     		=	15

								# dummy input_ids and scores
								_snake_case     		=	ids_tensor((batch_size, sequence_length)					,     UpperCAmelCase__						)
								_snake_case     		=	input_ids.copy()

								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	scores.copy()

								# instantiate all dist processors
								_snake_case     		=	FlaxTemperatureLogitsWarper(temperature=0.5						)
								_snake_case     		=	FlaxTopKLogitsWarper(3						)
								_snake_case     		=	FlaxTopPLogitsWarper(0.8						)

								# instantiate all logits processors
								_snake_case     		=	FlaxMinLengthLogitsProcessor(min_length=10					,     eos_token_id=UpperCAmelCase__						)
								_snake_case     		=	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__						)
								_snake_case     		=	FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__					,     eos_token_id=UpperCAmelCase__						)

								_snake_case     		=	10

								# no processor list
								_snake_case     		=	temp_dist_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								_snake_case     		=	top_k_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								_snake_case     		=	top_p_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								_snake_case     		=	min_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								_snake_case     		=	bos_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
								_snake_case     		=	eos_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)

								# with processor list
								_snake_case     		=	FlaxLogitsProcessorList(
								    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
								_snake_case     		=	processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)

								# scores should be equal
								self.assertTrue(jnp.allclose(UpperCAmelCase__					,     UpperCAmelCase__					,     atol=1E-3						)						)

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()					,     input_ids_comp.tolist()						)




							def 				lowerCamelCase       ( self						):





								"""simple docstring"""






								_snake_case     		=	4
								_snake_case     		=	10
								_snake_case     		=	15
								_snake_case     		=	2
								_snake_case     		=	1
								_snake_case     		=	15

								# dummy input_ids and scores
								_snake_case     		=	ids_tensor((batch_size, sequence_length)					,     UpperCAmelCase__						)
								_snake_case     		=	input_ids.copy()

								_snake_case     		=	self._get_uniform_logits(UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	scores.copy()

								# instantiate all dist processors
								_snake_case     		=	FlaxTemperatureLogitsWarper(temperature=0.5						)
								_snake_case     		=	FlaxTopKLogitsWarper(3						)
								_snake_case     		=	FlaxTopPLogitsWarper(0.8						)

								# instantiate all logits processors
								_snake_case     		=	FlaxMinLengthLogitsProcessor(min_length=10					,     eos_token_id=UpperCAmelCase__						)
								_snake_case     		=	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__						)
								_snake_case     		=	FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__					,     eos_token_id=UpperCAmelCase__						)

								_snake_case     		=	10

								# no processor list
								def run_no_processor_list(lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):
									_snake_case     		=	temp_dist_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									_snake_case     		=	top_k_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									_snake_case     		=	top_p_warp(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									_snake_case     		=	min_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									_snake_case     		=	bos_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									_snake_case     		=	eos_dist_proc(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									return scores

								# with processor list
								def run_processor_list(lowerCAmelCase_					,     lowerCAmelCase_					,     lowerCAmelCase_						):
									_snake_case     		=	FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
									_snake_case     		=	processor(UpperCAmelCase__					,     UpperCAmelCase__					,     cur_len=UpperCAmelCase__						)
									return scores

								_snake_case     		=	jax.jit(UpperCAmelCase__						)
								_snake_case     		=	jax.jit(UpperCAmelCase__						)

								_snake_case     		=	jitted_run_no_processor_list(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)
								_snake_case     		=	jitted_run_processor_list(UpperCAmelCase__					,     UpperCAmelCase__					,     UpperCAmelCase__						)

								# scores should be equal
								self.assertTrue(jnp.allclose(UpperCAmelCase__					,     UpperCAmelCase__					,     atol=1E-3						)						)

								# input_ids should never be changed
								self.assertListEqual(input_ids.tolist()					,     input_ids_comp.tolist()						)



'''simple docstring'''







import math







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			float:
	if (
	    not isinstance(__A    ,       (int, float)			)
	    or power_factor < -1
	    or power_factor > 1
	):
		raise ValueError('power_factor must be a valid float value between -1 and 1.'			)
	return apparent_power * power_factor







def        SCREAMING_SNAKE_CASE__			(     __A    ,       __A			) ->			float:
	if (
	    not isinstance(__A    ,       (int, float)			)
	    or power_factor < -1
	    or power_factor > 1
	):
		raise ValueError('power_factor must be a valid float value between -1 and 1.'			)
	return apparent_power * math.sqrt(1 - power_factor**2			)


if __name__ == "__main__":
					import doctest

					doctest.testmod()

