


"""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						,     UpperCamelCase__						,     UpperCamelCase__						)     ->		Union[str, Any]:

							'''simple docstring'''




							snake_case   :				str				 =   jnp.ones((batch_size, length)						) / length
							return scores







			def 			lowerCamelCase		(     self						)     ->		Any:

							'''simple docstring'''




							snake_case   :				int				 =   None
							snake_case   :				List[str]				 =   20

							snake_case   :				Union[str, Any]				 =   self._get_uniform_logits(batch_size=2						,     length=snake_case_						)

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

							# compute softmax
							snake_case   :				Any				 =   jax.nn.softmax(snake_case_						,     axis=-1						)

							snake_case   :				str				 =   FlaxTemperatureLogitsWarper(temperature=0.5						)
							snake_case   :				Union[str, Any]				 =   FlaxTemperatureLogitsWarper(temperature=1.3						)

							snake_case   :				str				 =   jax.nn.softmax(temp_dist_warper_sharper(snake_case_						,     scores.copy()						,     cur_len=snake_case_						)						,     axis=-1						)
							snake_case   :				List[Any]				 =   jax.nn.softmax(temp_dist_warper_smoother(snake_case_						,     scores.copy()						,     cur_len=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						)     ->		List[Any]:

							'''simple docstring'''




							snake_case   :				Any				 =   None
							snake_case   :				int				 =   10
							snake_case   :				Union[str, Any]				 =   2

							# create ramp distribution
							snake_case   :				int				 =   np.broadcast_to(np.arange(snake_case_						)[None, :]						,     (batch_size, vocab_size)						).copy()
							snake_case   :				Union[str, Any]				 =   ramp_logits[1:, : vocab_size // 2] + vocab_size

							snake_case   :				Dict				 =   FlaxTopKLogitsWarper(3						)

							snake_case   :				Optional[Any]				 =   top_k_warp(snake_case_						,     snake_case_						,     cur_len=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
							snake_case   :				str				 =   5
							snake_case   :				int				 =   FlaxTopKLogitsWarper(top_k=1						,     filter_value=0.0						,     min_tokens_to_keep=3						)

							snake_case   :				str				 =   np.broadcast_to(np.arange(snake_case_						)[None, :]						,     (batch_size, length)						).copy()
							snake_case   :				Tuple				 =   top_k_warp_safety_check(snake_case_						,     snake_case_						,     cur_len=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						)     ->		Union[str, Any]:

							'''simple docstring'''




							snake_case   :				Dict				 =   None
							snake_case   :				str				 =   10
							snake_case   :				Tuple				 =   2

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

							snake_case   :				Union[str, Any]				 =   FlaxTopPLogitsWarper(0.8						)
							snake_case   :				Tuple				 =   np.exp(top_p_warp(snake_case_						,     snake_case_						,     cur_len=snake_case_						)						)

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

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

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

							# make sure at least 2 tokens are kept
							snake_case   :				Any				 =   FlaxTopPLogitsWarper(0.9						,     min_tokens_to_keep=2						,     filter_value=0.0						)
							snake_case   :				Any				 =   top_p_warp(snake_case_						,     snake_case_						,     cur_len=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						)     ->		Any:

							'''simple docstring'''




							snake_case   :				int				 =   20
							snake_case   :				Union[str, Any]				 =   4
							snake_case   :				str				 =   0

							snake_case   :				Optional[Any]				 =   FlaxMinLengthLogitsProcessor(min_length=10						,     eos_token_id=snake_case_						)

							# check that min length is applied at length 5
							snake_case   :				Union[str, Any]				 =   ids_tensor((batch_size, 20)						,     vocab_size=20						)
							snake_case   :				List[str]				 =   5
							snake_case   :				Any				 =   self._get_uniform_logits(snake_case_						,     snake_case_						)
							snake_case   :				Optional[int]				 =   min_dist_processor(snake_case_						,     snake_case_						,     cur_len=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
							snake_case   :				Optional[Any]				 =   self._get_uniform_logits(snake_case_						,     snake_case_						)
							snake_case   :				Any				 =   15
							snake_case   :				Any				 =   min_dist_processor(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
							self.assertFalse(jnp.isinf(snake_case_						).any()						)







			def 			lowerCamelCase		(     self						)     ->		Union[str, Any]:

							'''simple docstring'''




							snake_case   :				Dict				 =   20
							snake_case   :				Union[str, Any]				 =   4
							snake_case   :				Optional[int]				 =   0

							snake_case   :				int				 =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_						)

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







			def 			lowerCamelCase		(     self						)     ->		Union[str, Any]:

							'''simple docstring'''




							snake_case   :				Union[str, Any]				 =   20
							snake_case   :				str				 =   4
							snake_case   :				Optional[Any]				 =   0
							snake_case   :				int				 =   5

							snake_case   :				Union[str, Any]				 =   FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_						,     eos_token_id=snake_case_						)

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







			def 			lowerCamelCase		(     self						)     ->		Dict:

							'''simple docstring'''




							snake_case   :				Dict				 =   4
							snake_case   :				Optional[int]				 =   10
							snake_case   :				int				 =   15
							snake_case   :				Optional[Any]				 =   2
							snake_case   :				Union[str, Any]				 =   1
							snake_case   :				Dict				 =   15

							# dummy input_ids and scores
							snake_case   :				Tuple				 =   ids_tensor((batch_size, sequence_length)						,     snake_case_						)
							snake_case   :				Dict				 =   input_ids.copy()

							snake_case   :				List[Any]				 =   self._get_uniform_logits(snake_case_						,     snake_case_						)
							snake_case   :				Optional[Any]				 =   scores.copy()

							# instantiate all dist processors
							snake_case   :				Dict				 =   FlaxTemperatureLogitsWarper(temperature=0.5						)
							snake_case   :				Tuple				 =   FlaxTopKLogitsWarper(3						)
							snake_case   :				str				 =   FlaxTopPLogitsWarper(0.8						)

							# instantiate all logits processors
							snake_case   :				Optional[int]				 =   FlaxMinLengthLogitsProcessor(min_length=10						,     eos_token_id=snake_case_						)
							snake_case   :				Any				 =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_						)
							snake_case   :				Optional[Any]				 =   FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_						,     eos_token_id=snake_case_						)

							snake_case   :				List[str]				 =   10

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

							# with processor list
							snake_case   :				Dict				 =   FlaxLogitsProcessorList(
							    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]						)
							snake_case   :				Dict				 =   processor(snake_case_						,     snake_case_						,     cur_len=snake_case_						)

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

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







			def 			lowerCamelCase		(     self						)     ->		Any:

							'''simple docstring'''




							snake_case   :				List[Any]				 =   4
							snake_case   :				Optional[Any]				 =   10
							snake_case   :				Optional[Any]				 =   15
							snake_case   :				int				 =   2
							snake_case   :				Tuple				 =   1
							snake_case   :				int				 =   15

							# dummy input_ids and scores
							snake_case   :				Union[str, Any]				 =   ids_tensor((batch_size, sequence_length)						,     snake_case_						)
							snake_case   :				List[str]				 =   input_ids.copy()

							snake_case   :				int				 =   self._get_uniform_logits(snake_case_						,     snake_case_						)
							snake_case   :				Dict				 =   scores.copy()

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

							# instantiate all logits processors
							snake_case   :				Optional[int]				 =   FlaxMinLengthLogitsProcessor(min_length=10						,     eos_token_id=snake_case_						)
							snake_case   :				Optional[int]				 =   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case_						)
							snake_case   :				Dict				 =   FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case_						,     eos_token_id=snake_case_						)

							snake_case   :				int				 =   10

							# no processor list
							def run_no_processor_list(UpperCamelCase__						,     UpperCamelCase__						,     UpperCamelCase__						):
											snake_case   :				Any				 =   temp_dist_warp(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											snake_case   :				Optional[Any]				 =   top_k_warp(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											snake_case   :				Optional[int]				 =   top_p_warp(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											snake_case   :				Optional[int]				 =   min_dist_proc(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											snake_case   :				List[str]				 =   bos_dist_proc(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											snake_case   :				int				 =   eos_dist_proc(snake_case_						,     snake_case_						,     cur_len=snake_case_						)
											return scores

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

							snake_case   :				Dict				 =   jax.jit(snake_case_						)
							snake_case   :				Union[str, Any]				 =   jax.jit(snake_case_						)

							snake_case   :				List[str]				 =   jitted_run_no_processor_list(snake_case_						,     snake_case_						,     snake_case_						)
							snake_case   :				str				 =   jitted_run_processor_list(snake_case_						,     snake_case_						,     snake_case_						)

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

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









"""simple docstring"""



def   __lowerCAmelCase						(						lowercase  :  int			)     ->						int:


    """simple docstring"""





    if not isinstance(lowercase						,							lowercase			):
        raise ValueError("Input must be an integer"			)
    if input_num <= 0:
        raise ValueError("Input must be positive"			)
    return sum(
        divisor for divisor in range(1						,							input_num // 2 + 1			) if input_num % divisor == 0			)


if __name__ == "__main__":
  import doctest

  doctest.testmod()





