





'''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 		lowerCamelCase	(   unittest.TestCase							):



	def 	_lowerCamelCase     ( self		,   a_		,   a_						):
			lowerCAmelCase  :     int       						= jnp.ones((batch_size, length)						) / length
			return scores



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Union[str, Any]       						= None
			lowerCAmelCase  :     Tuple       						= 20

			lowerCAmelCase  :     List[Any]       						= self._get_uniform_logits(batch_size=2		,   length=a_						)

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

			# compute softmax
			lowerCAmelCase  :     Tuple       						= jax.nn.softmax(a_		,   axis=-1						)

			lowerCAmelCase  :     int       						= FlaxTemperatureLogitsWarper(temperature=0.5						)
			lowerCAmelCase  :     Optional[int]       						= FlaxTemperatureLogitsWarper(temperature=1.3						)

			lowerCAmelCase  :     Dict       						= jax.nn.softmax(temp_dist_warper_sharper(a_		,   scores.copy()		,   cur_len=a_						)		,   axis=-1						)
			lowerCAmelCase  :     List[Any]       						= jax.nn.softmax(temp_dist_warper_smoother(a_		,   scores.copy()		,   cur_len=a_						)		,   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						):
			lowerCAmelCase  :     Tuple       						= None
			lowerCAmelCase  :     int       						= 10
			lowerCAmelCase  :     Optional[int]       						= 2

			# create ramp distribution
			lowerCAmelCase  :     List[str]       						= np.broadcast_to(np.arange(a_						)[None, :]		,   (batch_size, vocab_size)						).copy()
			lowerCAmelCase  :     Tuple       						= ramp_logits[1:, : vocab_size // 2] + vocab_size

			lowerCAmelCase  :     Optional[Any]       						= FlaxTopKLogitsWarper(3						)

			lowerCAmelCase  :     str       						= top_k_warp(a_		,   a_		,   cur_len=a_						)

			# 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
			lowerCAmelCase  :     Union[str, Any]       						= 5
			lowerCAmelCase  :     List[Any]       						= FlaxTopKLogitsWarper(top_k=1		,   filter_value=0.0		,   min_tokens_to_keep=3						)

			lowerCAmelCase  :     Tuple       						= np.broadcast_to(np.arange(a_						)[None, :]		,   (batch_size, length)						).copy()
			lowerCAmelCase  :     str       						= top_k_warp_safety_check(a_		,   a_		,   cur_len=a_						)

			# 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						):
			lowerCAmelCase  :     int       						= None
			lowerCAmelCase  :     Union[str, Any]       						= 10
			lowerCAmelCase  :     Optional[Any]       						= 2

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

			lowerCAmelCase  :     Union[str, Any]       						= FlaxTopPLogitsWarper(0.8						)
			lowerCAmelCase  :     Dict       						= np.exp(top_p_warp(a_		,   a_		,   cur_len=a_						)						)

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

			# check edge cases with negative and extreme logits
			lowerCAmelCase  :     Union[str, Any]       						= np.broadcast_to(np.arange(a_						)[None, :]		,   (batch_size, vocab_size)						).copy() - (
			    vocab_size // 2
			)

			# make ramp_logits more extreme
			lowerCAmelCase  :     int       						= ramp_logits[1] * 100.0

			# make sure at least 2 tokens are kept
			lowerCAmelCase  :     int       						= FlaxTopPLogitsWarper(0.9		,   min_tokens_to_keep=2		,   filter_value=0.0						)
			lowerCAmelCase  :     Tuple       						= top_p_warp(a_		,   a_		,   cur_len=a_						)

			# 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						):
			lowerCAmelCase  :     List[Any]       						= 20
			lowerCAmelCase  :     Dict       						= 4
			lowerCAmelCase  :     str       						= 0

			lowerCAmelCase  :     Any       						= FlaxMinLengthLogitsProcessor(min_length=10		,   eos_token_id=a_						)

			# check that min length is applied at length 5
			lowerCAmelCase  :     List[str]       						= ids_tensor((batch_size, 20)		,   vocab_size=20						)
			lowerCAmelCase  :     List[str]       						= 5
			lowerCAmelCase  :     List[Any]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     Dict       						= min_dist_processor(a_		,   a_		,   cur_len=a_						)
			self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()		,   4 * [-float("inf"						)]						)

			# check that min length is not applied anymore at length 15
			lowerCAmelCase  :     Optional[int]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     Optional[Any]       						= 15
			lowerCAmelCase  :     Optional[Any]       						= min_dist_processor(a_		,   a_		,   cur_len=a_						)
			self.assertFalse(jnp.isinf(a_						).any()						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[Any]       						= 20
			lowerCAmelCase  :     Optional[Any]       						= 4
			lowerCAmelCase  :     List[str]       						= 0

			lowerCAmelCase  :     Optional[int]       						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_						)

			# check that all scores are -inf except the bos_token_id score
			lowerCAmelCase  :     Optional[Any]       						= ids_tensor((batch_size, 1)		,   vocab_size=20						)
			lowerCAmelCase  :     Dict       						= 1
			lowerCAmelCase  :     Optional[Any]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     List[str]       						= logits_processor(a_		,   a_		,   cur_len=a_						)
			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
			lowerCAmelCase  :     List[Any]       						= 3
			lowerCAmelCase  :     List[Any]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     int       						= logits_processor(a_		,   a_		,   cur_len=a_						)
			self.assertFalse(jnp.isinf(a_						).any()						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     Optional[int]       						= 20
			lowerCAmelCase  :     Union[str, Any]       						= 4
			lowerCAmelCase  :     str       						= 0
			lowerCAmelCase  :     Tuple       						= 5

			lowerCAmelCase  :     List[str]       						= FlaxForcedEOSTokenLogitsProcessor(max_length=a_		,   eos_token_id=a_						)

			# check that all scores are -inf except the eos_token_id when max_length is reached
			lowerCAmelCase  :     str       						= ids_tensor((batch_size, 4)		,   vocab_size=20						)
			lowerCAmelCase  :     Tuple       						= 4
			lowerCAmelCase  :     Optional[Any]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     Optional[Any]       						= logits_processor(a_		,   a_		,   cur_len=a_						)
			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
			lowerCAmelCase  :     Optional[Any]       						= 3
			lowerCAmelCase  :     Any       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     Optional[int]       						= logits_processor(a_		,   a_		,   cur_len=a_						)
			self.assertFalse(jnp.isinf(a_						).any()						)



	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[Any]       						= 4
			lowerCAmelCase  :     Optional[int]       						= 10
			lowerCAmelCase  :     Dict       						= 15
			lowerCAmelCase  :     Optional[int]       						= 2
			lowerCAmelCase  :     Union[str, Any]       						= 1
			lowerCAmelCase  :     Optional[Any]       						= 15

			# dummy input_ids and scores
			lowerCAmelCase  :     Tuple       						= ids_tensor((batch_size, sequence_length)		,   a_						)
			lowerCAmelCase  :     Tuple       						= input_ids.copy()

			lowerCAmelCase  :     Any       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     Any       						= scores.copy()

			# instantiate all dist processors
			lowerCAmelCase  :     int       						= FlaxTemperatureLogitsWarper(temperature=0.5						)
			lowerCAmelCase  :     int       						= FlaxTopKLogitsWarper(3						)
			lowerCAmelCase  :     Optional[int]       						= FlaxTopPLogitsWarper(0.8						)

			# instantiate all logits processors
			lowerCAmelCase  :     str       						= FlaxMinLengthLogitsProcessor(min_length=10		,   eos_token_id=a_						)
			lowerCAmelCase  :     Dict       						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_						)
			lowerCAmelCase  :     int       						= FlaxForcedEOSTokenLogitsProcessor(max_length=a_		,   eos_token_id=a_						)

			lowerCAmelCase  :     Dict       						= 10

			# no processor list
			lowerCAmelCase  :     Union[str, Any]       						= temp_dist_warp(a_		,   a_		,   cur_len=a_						)
			lowerCAmelCase  :     Dict       						= top_k_warp(a_		,   a_		,   cur_len=a_						)
			lowerCAmelCase  :     Union[str, Any]       						= top_p_warp(a_		,   a_		,   cur_len=a_						)
			lowerCAmelCase  :     int       						= min_dist_proc(a_		,   a_		,   cur_len=a_						)
			lowerCAmelCase  :     Optional[int]       						= bos_dist_proc(a_		,   a_		,   cur_len=a_						)
			lowerCAmelCase  :     List[Any]       						= eos_dist_proc(a_		,   a_		,   cur_len=a_						)

			# with processor list
			lowerCAmelCase  :     List[Any]       						= FlaxLogitsProcessorList(
			    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
			lowerCAmelCase  :     Optional[int]       						= processor(a_		,   a_		,   cur_len=a_						)

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

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





	def 	_lowerCamelCase     ( self						):
			lowerCAmelCase  :     List[str]       						= 4
			lowerCAmelCase  :     Tuple       						= 10
			lowerCAmelCase  :     Optional[Any]       						= 15
			lowerCAmelCase  :     Dict       						= 2
			lowerCAmelCase  :     Optional[int]       						= 1
			lowerCAmelCase  :     List[str]       						= 15

			# dummy input_ids and scores
			lowerCAmelCase  :     Tuple       						= ids_tensor((batch_size, sequence_length)		,   a_						)
			lowerCAmelCase  :     int       						= input_ids.copy()

			lowerCAmelCase  :     Union[str, Any]       						= self._get_uniform_logits(a_		,   a_						)
			lowerCAmelCase  :     List[Any]       						= scores.copy()

			# instantiate all dist processors
			lowerCAmelCase  :     List[Any]       						= FlaxTemperatureLogitsWarper(temperature=0.5						)
			lowerCAmelCase  :     Optional[int]       						= FlaxTopKLogitsWarper(3						)
			lowerCAmelCase  :     Optional[Any]       						= FlaxTopPLogitsWarper(0.8						)

			# instantiate all logits processors
			lowerCAmelCase  :     Optional[Any]       						= FlaxMinLengthLogitsProcessor(min_length=10		,   eos_token_id=a_						)
			lowerCAmelCase  :     str       						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_						)
			lowerCAmelCase  :     Any       						= FlaxForcedEOSTokenLogitsProcessor(max_length=a_		,   eos_token_id=a_						)

			lowerCAmelCase  :     int       						= 10

			# no processor list
			def run_no_processor_list(a_		,   a_		,   a_						):
					lowerCAmelCase  :     Tuple       						= temp_dist_warp(a_		,   a_		,   cur_len=a_						)
					lowerCAmelCase  :     List[str]       						= top_k_warp(a_		,   a_		,   cur_len=a_						)
					lowerCAmelCase  :     str       						= top_p_warp(a_		,   a_		,   cur_len=a_						)
					lowerCAmelCase  :     Any       						= min_dist_proc(a_		,   a_		,   cur_len=a_						)
					lowerCAmelCase  :     Tuple       						= bos_dist_proc(a_		,   a_		,   cur_len=a_						)
					lowerCAmelCase  :     Any       						= eos_dist_proc(a_		,   a_		,   cur_len=a_						)
					return scores

			# with processor list
			def run_processor_list(a_		,   a_		,   a_						):
					lowerCAmelCase  :     Union[str, Any]       						= FlaxLogitsProcessorList(
					    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
					lowerCAmelCase  :     Any       						= processor(a_		,   a_		,   cur_len=a_						)
					return scores

			lowerCAmelCase  :     int       						= jax.jit(a_						)
			lowerCAmelCase  :     int       						= jax.jit(a_						)

			lowerCAmelCase  :     str       						= jitted_run_no_processor_list(a_		,   a_		,   a_						)
			lowerCAmelCase  :     Tuple       						= jitted_run_processor_list(a_		,   a_		,   a_						)

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

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










'''simple docstring'''






import random
from typing import Any






def   __A    (				a_    :	list						):
		for _ in range(len(a_						)						):
				lowerCAmelCase  :     List[Any]       						= random.randint(0   ,len(a_						) - 1						)
				lowerCAmelCase  :     Tuple       						= random.randint(0   ,len(a_						) - 1						)
				lowerCAmelCase						,		lowerCAmelCase  :     Optional[Any]       						= data[b], data[a]
		return data


if __name__ == "__main__":
					lowerCAmelCase									=   [0, 1, 2, 3, 4, 5, 6, 7]
					lowerCAmelCase									=   ["""python""", """says""", """hello""", """!"""]
					print("""Fisher-Yates Shuffle:""")
					print("""List""", integers, strings)
					print("""FY Shuffle""", fisher_yates_shuffle(integers), fisher_yates_shuffle(strings))



