

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





      def 					_lowerCamelCase  (  self    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)   ->					Union[str, Any]:
         lowerCamelCase__						=			jnp.ones((batch_size, length)		) / length
         return scores





      def 					_lowerCamelCase  (  self		)   ->					str:
         lowerCamelCase__						=			None
         lowerCamelCase__						=			20

         lowerCamelCase__						=			self._get_uniform_logits(batch_size=2    ,						length=SCREAMING_SNAKE_CASE__		)

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

         # compute softmax
         lowerCamelCase__						=			jax.nn.softmax(SCREAMING_SNAKE_CASE__    ,						axis=-1		)

         lowerCamelCase__						=			FlaxTemperatureLogitsWarper(temperature=0.5		)
         lowerCamelCase__						=			FlaxTemperatureLogitsWarper(temperature=1.3		)

         lowerCamelCase__						=			jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE__    ,						scores.copy()    ,						cur_len=SCREAMING_SNAKE_CASE__		)    ,						axis=-1		)
         lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         lowerCamelCase__						=			None
         lowerCamelCase__						=			10
         lowerCamelCase__						=			2

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

         lowerCamelCase__						=			FlaxTopKLogitsWarper(3		)

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

         lowerCamelCase__						=			np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE__		)[None, :]    ,						(batch_size, length)		).copy()
         lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					str:
         lowerCamelCase__						=			None
         lowerCamelCase__						=			10
         lowerCamelCase__						=			2

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

         lowerCamelCase__						=			FlaxTopPLogitsWarper(0.8		)
         lowerCamelCase__						=			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
         lowerCamelCase__						=			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
         lowerCamelCase__						=			np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE__		)[None, :]    ,						(batch_size, vocab_size)		).copy() - (
             vocab_size // 2
         )

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

         # make sure at least 2 tokens are kept
         lowerCamelCase__						=			FlaxTopPLogitsWarper(0.9    ,						min_tokens_to_keep=2    ,						filter_value=0.0		)
         lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					Tuple:
         lowerCamelCase__						=			20
         lowerCamelCase__						=			4
         lowerCamelCase__						=			0

         lowerCamelCase__						=			FlaxMinLengthLogitsProcessor(min_length=10    ,						eos_token_id=SCREAMING_SNAKE_CASE__		)

         # check that min length is applied at length 5
         lowerCamelCase__						=			ids_tensor((batch_size, 20)    ,						vocab_size=20		)
         lowerCamelCase__						=			5
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			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
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			15
         lowerCamelCase__						=			min_dist_processor(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__		).any()		)





      def 					_lowerCamelCase  (  self		)   ->					List[Any]:
         lowerCamelCase__						=			20
         lowerCamelCase__						=			4
         lowerCamelCase__						=			0

         lowerCamelCase__						=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__		)

         # check that all scores are -inf except the bos_token_id score
         lowerCamelCase__						=			ids_tensor((batch_size, 1)    ,						vocab_size=20		)
         lowerCamelCase__						=			1
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			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
         lowerCamelCase__						=			3
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			logits_processor(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__		).any()		)





      def 					_lowerCamelCase  (  self		)   ->					List[str]:
         lowerCamelCase__						=			20
         lowerCamelCase__						=			4
         lowerCamelCase__						=			0
         lowerCamelCase__						=			5

         lowerCamelCase__						=			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
         lowerCamelCase__						=			ids_tensor((batch_size, 4)    ,						vocab_size=20		)
         lowerCamelCase__						=			4
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			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
         lowerCamelCase__						=			3
         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			logits_processor(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__		).any()		)





      def 					_lowerCamelCase  (  self		)   ->					Dict:
         lowerCamelCase__						=			4
         lowerCamelCase__						=			10
         lowerCamelCase__						=			15
         lowerCamelCase__						=			2
         lowerCamelCase__						=			1
         lowerCamelCase__						=			15

         # dummy input_ids and scores
         lowerCamelCase__						=			ids_tensor((batch_size, sequence_length)    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			input_ids.copy()

         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			scores.copy()

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

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

         lowerCamelCase__						=			10

         # no processor list
         lowerCamelCase__						=			temp_dist_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			top_k_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			top_p_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			min_dist_proc(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			bos_dist_proc(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			eos_dist_proc(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)

         # with processor list
         lowerCamelCase__						=			FlaxLogitsProcessorList(
             [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
         lowerCamelCase__						=			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 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         lowerCamelCase__						=			4
         lowerCamelCase__						=			10
         lowerCamelCase__						=			15
         lowerCamelCase__						=			2
         lowerCamelCase__						=			1
         lowerCamelCase__						=			15

         # dummy input_ids and scores
         lowerCamelCase__						=			ids_tensor((batch_size, sequence_length)    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			input_ids.copy()

         lowerCamelCase__						=			self._get_uniform_logits(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			scores.copy()

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

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

         lowerCamelCase__						=			10

         # no processor list
         def run_no_processor_list(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		):
            lowerCamelCase__						=			temp_dist_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			top_k_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			top_p_warp(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			min_dist_proc(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			bos_dist_proc(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            lowerCamelCase__						=			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__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		):
            lowerCamelCase__						=			FlaxLogitsProcessorList(
                [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]		)
            lowerCamelCase__						=			processor(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						cur_len=SCREAMING_SNAKE_CASE__		)
            return scores

         lowerCamelCase__						=			jax.jit(SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			jax.jit(SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			jitted_run_no_processor_list(SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__    ,						SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			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 	UpperCAmelCase__   (				A__   )      ->      str:




   """simple docstring"""
   return " ".join(input_str.split()[::-1]   )


if __name__ == "__main__":
     import doctest

     doctest.testmod()




