



"""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    ):





     '''simple docstring'''



     def 			snake_case__      (    self				,       lowerCamelCase__				,       lowerCamelCase__   ):
         _lowerCamelCase        =				jnp.ones((batch_size, length)   ) / length
         return scores

     def 			snake_case__      (    self   ):
         _lowerCamelCase        =				None
         _lowerCamelCase        =				2_0

         _lowerCamelCase        =				self._get_uniform_logits(batch_size=2				,       length=UpperCamelCase__   )

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

         # compute softmax
         _lowerCamelCase        =				jax.nn.softmax(UpperCamelCase__				,       axis=-1   )

         _lowerCamelCase        =				FlaxTemperatureLogitsWarper(temperature=0.5   )
         _lowerCamelCase        =				FlaxTemperatureLogitsWarper(temperature=1.3   )

         _lowerCamelCase        =				jax.nn.softmax(temp_dist_warper_sharper(UpperCamelCase__				,       scores.copy()				,       cur_len=UpperCamelCase__   )				,       axis=-1   )
         _lowerCamelCase        =				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 			snake_case__      (    self   ):
         _lowerCamelCase        =				None
         _lowerCamelCase        =				1_0
         _lowerCamelCase        =				2

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

         _lowerCamelCase        =				FlaxTopKLogitsWarper(3   )

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

         _lowerCamelCase        =				np.broadcast_to(np.arange(UpperCamelCase__   )[None, :]				,       (batch_size, length)   ).copy()
         _lowerCamelCase        =				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 			snake_case__      (    self   ):
         _lowerCamelCase        =				None
         _lowerCamelCase        =				1_0
         _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.1_5, 0.3, 0.3, 0.2_5]]   )   )

         _lowerCamelCase        =				FlaxTopPLogitsWarper(0.8   )
         _lowerCamelCase        =				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
         _lowerCamelCase        =				np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]   )
         self.assertTrue(np.allclose(UpperCamelCase__				,       UpperCamelCase__				,       atol=1e-3   )   )

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

         # make ramp_logits more extreme
         _lowerCamelCase        =				ramp_logits[1] * 1_0_0.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(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 			snake_case__      (    self   ):
         _lowerCamelCase        =				2_0
         _lowerCamelCase        =				4
         _lowerCamelCase        =				0

         _lowerCamelCase        =				FlaxMinLengthLogitsProcessor(min_length=1_0				,       eos_token_id=UpperCamelCase__   )

         # check that min length is applied at length 5
         _lowerCamelCase        =				ids_tensor((batch_size, 2_0)				,       vocab_size=2_0   )
         _lowerCamelCase        =				5
         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _lowerCamelCase        =				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
         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _lowerCamelCase        =				1_5
         _lowerCamelCase        =				min_dist_processor(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         self.assertFalse(jnp.isinf(UpperCamelCase__   ).any()   )

     def 			snake_case__      (    self   ):
         _lowerCamelCase        =				2_0
         _lowerCamelCase        =				4
         _lowerCamelCase        =				0

         _lowerCamelCase        =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__   )

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

     def 			snake_case__      (    self   ):
         _lowerCamelCase        =				2_0
         _lowerCamelCase        =				4
         _lowerCamelCase        =				0
         _lowerCamelCase        =				5

         _lowerCamelCase        =				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__				,       eos_token_id=UpperCamelCase__   )

         # check that all scores are -inf except the eos_token_id when max_length is reached
         _lowerCamelCase        =				ids_tensor((batch_size, 4)				,       vocab_size=2_0   )
         _lowerCamelCase        =				4
         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _lowerCamelCase        =				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
         _lowerCamelCase        =				3
         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _lowerCamelCase        =				logits_processor(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         self.assertFalse(jnp.isinf(UpperCamelCase__   ).any()   )

     def 			snake_case__      (    self   ):
         _lowerCamelCase        =				4
         _lowerCamelCase        =				1_0
         _lowerCamelCase        =				1_5
         _lowerCamelCase        =				2
         _lowerCamelCase        =				1
         _lowerCamelCase        =				1_5

         # dummy input_ids and scores
         _lowerCamelCase        =				ids_tensor((batch_size, sequence_length)				,       UpperCamelCase__   )
         _lowerCamelCase        =				input_ids.copy()

         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _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=1_0				,       eos_token_id=UpperCamelCase__   )
         _lowerCamelCase        =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__   )
         _lowerCamelCase        =				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__				,       eos_token_id=UpperCamelCase__   )

         _lowerCamelCase        =				1_0

         # no processor list
         _lowerCamelCase        =				temp_dist_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         _lowerCamelCase        =				top_k_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         _lowerCamelCase        =				top_p_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         _lowerCamelCase        =				min_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         _lowerCamelCase        =				bos_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
         _lowerCamelCase        =				eos_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )

         # 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(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 			snake_case__      (    self   ):
         _lowerCamelCase        =				4
         _lowerCamelCase        =				1_0
         _lowerCamelCase        =				1_5
         _lowerCamelCase        =				2
         _lowerCamelCase        =				1
         _lowerCamelCase        =				1_5

         # dummy input_ids and scores
         _lowerCamelCase        =				ids_tensor((batch_size, sequence_length)				,       UpperCamelCase__   )
         _lowerCamelCase        =				input_ids.copy()

         _lowerCamelCase        =				self._get_uniform_logits(UpperCamelCase__				,       UpperCamelCase__   )
         _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=1_0				,       eos_token_id=UpperCamelCase__   )
         _lowerCamelCase        =				FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase__   )
         _lowerCamelCase        =				FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase__				,       eos_token_id=UpperCamelCase__   )

         _lowerCamelCase        =				1_0

         # no processor list
         def run_no_processor_list(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
             _lowerCamelCase        =				temp_dist_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             _lowerCamelCase        =				top_k_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             _lowerCamelCase        =				top_p_warp(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             _lowerCamelCase        =				min_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             _lowerCamelCase        =				bos_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             _lowerCamelCase        =				eos_dist_proc(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             return scores

         # with processor list
         def run_processor_list(lowerCamelCase__				,       lowerCamelCase__				,       lowerCamelCase__   ):
             _lowerCamelCase        =				FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]   )
             _lowerCamelCase        =				processor(UpperCamelCase__				,       UpperCamelCase__				,       cur_len=UpperCamelCase__   )
             return scores

         _lowerCamelCase        =				jax.jit(UpperCamelCase__   )
         _lowerCamelCase        =				jax.jit(UpperCamelCase__   )

         _lowerCamelCase        =				jitted_run_no_processor_list(UpperCamelCase__				,       UpperCamelCase__				,       UpperCamelCase__   )
         _lowerCamelCase        =				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 numpy as np



def        lowerCAmelCase_(   lowercase_     :							np.array							)	->       np.array:
				return 1 / (1 + np.exp(-vector							))



def        lowerCAmelCase_(   lowercase_     :							np.array							)	->       np.array:
				return vector * sigmoid(1.7_0_2 * vector							)


if __name__ == "__main__":
			import doctest

			doctest.testmod()
