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



	"""simple docstring"""




	def 		_lowercase    (			self    :     str			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     int			)			->	Dict:
					lowercase_	    =   jnp.ones((batch_size, length)			) / length
					return scores
	def 		_lowercase    (			self    :     List[Any]			)			->	List[Any]:
					lowercase_	    =   None
					lowercase_	    =   2_0

					lowercase_	    =   self._get_uniform_logits(batch_size=2			,				length=SCREAMING_SNAKE_CASE_			)

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

					# compute softmax
					lowercase_	    =   jax.nn.softmax(SCREAMING_SNAKE_CASE_			,				axis=-1			)

					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=0.5			)
					lowercase_	    =   FlaxTemperatureLogitsWarper(temperature=1.3			)

					lowercase_	    =   jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_SNAKE_CASE_			)			,				axis=-1			)
					lowercase_	    =   jax.nn.softmax(temp_dist_warper_smoother(SCREAMING_SNAKE_CASE_			,				scores.copy()			,				cur_len=SCREAMING_SNAKE_CASE_			)			,				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 		_lowercase    (			self    :     Optional[Any]			)			->	List[str]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

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

					lowercase_	    =   FlaxTopKLogitsWarper(3			)

					lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

					# 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
					lowercase_	    =   5
					lowercase_	    =   FlaxTopKLogitsWarper(top_k=1			,				filter_value=0.0			,				min_tokens_to_keep=3			)

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

					# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
					self.assertListEqual((scores == 0.0).sum(axis=-1			).tolist()			,				[2, 2]			)
	def 		_lowercase    (			self    :     List[Any]			)			->	Union[str, Any]:
					lowercase_	    =   None
					lowercase_	    =   1_0
					lowercase_	    =   2

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

					lowercase_	    =   FlaxTopPLogitsWarper(0.8			)
					lowercase_	    =   np.exp(top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)			)

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

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

					# make ramp_logits more extreme
					lowercase_	    =   ramp_logits[1] * 1_00.0

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

					# 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 		_lowercase    (			self    :     Tuple			)			->	Any:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

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

					# check that min length is not applied anymore at length 15
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   1_5
					lowercase_	    =   min_dist_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     List[str]			)			->	Dict:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0

					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the bos_token_id score
					lowercase_	    =   ids_tensor((batch_size, 1)			,				vocab_size=2_0			)
					lowercase_	    =   1
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					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
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Optional[int]:
					lowercase_	    =   2_0
					lowercase_	    =   4
					lowercase_	    =   0
					lowercase_	    =   5

					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					# check that all scores are -inf except the eos_token_id when max_length is reached
					lowercase_	    =   ids_tensor((batch_size, 4)			,				vocab_size=2_0			)
					lowercase_	    =   4
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					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
					lowercase_	    =   3
					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   logits_processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_			).any()			)
	def 		_lowercase    (			self    :     int			)			->	Tuple:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

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

					# instantiate all logits processors
					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   1_0

					# no processor list
					lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)

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

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

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




	def 		_lowercase    (			self    :     List[Any]			)			->	List[str]:
					lowercase_	    =   4
					lowercase_	    =   1_0
					lowercase_	    =   1_5
					lowercase_	    =   2
					lowercase_	    =   1
					lowercase_	    =   1_5

					# dummy input_ids and scores
					lowercase_	    =   ids_tensor((batch_size, sequence_length)			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   input_ids.copy()

					lowercase_	    =   self._get_uniform_logits(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   scores.copy()

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

					# instantiate all logits processors
					lowercase_	    =   FlaxMinLengthLogitsProcessor(min_length=1_0			,				eos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_			,				eos_token_id=SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   1_0

					# no processor list
					def run_no_processor_list(SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     Optional[int]			,				SCREAMING_SNAKE_CASE_    :     int			):
									lowercase_	    =   temp_dist_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_k_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   top_p_warp(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   min_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   bos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									lowercase_	    =   eos_dist_proc(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					# with processor list
					def run_processor_list(SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Dict			,				SCREAMING_SNAKE_CASE_    :     List[Any]			):
									lowercase_	    =   FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
									lowercase_	    =   processor(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				cur_len=SCREAMING_SNAKE_CASE_			)
									return scores

					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jax.jit(SCREAMING_SNAKE_CASE_			)

					lowercase_	    =   jitted_run_no_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   jitted_run_processor_list(SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			,				SCREAMING_SNAKE_CASE_			)

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

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












"""simple docstring"""



def 	SCREAMING_SNAKE_CASE__       (  snake_case       : int			,					snake_case       : int				)->		str:






     '''simple docstring'''



     if not isinstance(snake_case			,					snake_case				):
          raise ValueError("iterations must be defined as integers"				)
     if not isinstance(snake_case			,					snake_case				) or not number >= 1:
          raise ValueError(
              "starting number must be\n                         and integer and be more than 0"				)
     if not iterations >= 1:
          raise ValueError("Iterations must be done more than 0 times to play FizzBuzz"				)

     UpperCAmelCase__						:    str							=							""
     while number <= iterations:
          if number % 3 == 0:
               out += "Fizz"
          if number % 5 == 0:
               out += "Buzz"
          if 0 not in (number % 3, number % 5):
               out += str(snake_case				)

          # print(out)
          number += 1
          out += " "
     return out


if __name__ == "__main__":
 import doctest

 doctest.testmod()


