





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






		'''simple docstring'''






		def 						snake_case__				(							self,   snake_case__,   snake_case__  )       -> Tuple:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     jnp.ones((batch_size, length)  ) / length
									return scores






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






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     None
									lowercase_	:   Optional[int]  				=     20

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

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

									# compute softmax
									lowercase_	:   str  				=     jax.nn.softmax(snake_case__,   axis=-1  )

									lowercase_	:   List[str]  				=     FlaxTemperatureLogitsWarper(temperature=0.5  )
									lowercase_	:   Tuple  				=     FlaxTemperatureLogitsWarper(temperature=1.3  )

									lowercase_	:   List[Any]  				=     jax.nn.softmax(temp_dist_warper_sharper(snake_case__,   scores.copy(),   cur_len=snake_case__  ),   axis=-1  )
									lowercase_	:   Union[str, 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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Any  				=     None
									lowercase_	:   List[str]  				=     10
									lowercase_	:   int  				=     2

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

									lowercase_	:   int  				=     FlaxTopKLogitsWarper(3  )

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

									lowercase_	:   Union[str, Any]  				=     np.broadcast_to(np.arange(snake_case__  )[None, :],   (batch_size, length)  ).copy()
									lowercase_	:   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 						snake_case__				(							self  )       -> Optional[int]:






									"""simple docstring"""
									lowercase_	:   List[Any]  				=     None
									lowercase_	:   Dict  				=     10
									lowercase_	:   Optional[int]  				=     2

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

									lowercase_	:   Any  				=     FlaxTopPLogitsWarper(0.8  )
									lowercase_	:   int  				=     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
									lowercase_	:   List[Any]  				=     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
									lowercase_	:   Union[str, Any]  				=     np.broadcast_to(np.arange(snake_case__  )[None, :],   (batch_size, vocab_size)  ).copy() - (
									    vocab_size // 2
									)

									# make ramp_logits more extreme
									lowercase_	:   Optional[Any]  				=     ramp_logits[1] * 100.0

									# make sure at least 2 tokens are kept
									lowercase_	:   Optional[int]  				=     FlaxTopPLogitsWarper(0.9,   min_tokens_to_keep=2,   filter_value=0.0  )
									lowercase_	:   Optional[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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   List[str]  				=     20
									lowercase_	:   Any  				=     4
									lowercase_	:   List[Any]  				=     0

									lowercase_	:   int  				=     FlaxMinLengthLogitsProcessor(min_length=10,   eos_token_id=snake_case__  )

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






		def 						snake_case__				(							self  )       -> Any:






									"""simple docstring"""
									lowercase_	:   Any  				=     20
									lowercase_	:   Dict  				=     4
									lowercase_	:   List[str]  				=     0

									lowercase_	:   Any  				=     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__  )

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






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






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     20
									lowercase_	:   Dict  				=     4
									lowercase_	:   str  				=     0
									lowercase_	:   Dict  				=     5

									lowercase_	:   int  				=     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
									lowercase_	:   Tuple  				=     ids_tensor((batch_size, 4),   vocab_size=20  )
									lowercase_	:   List[str]  				=     4
									lowercase_	:   List[Any]  				=     self._get_uniform_logits(snake_case__,   snake_case__  )
									lowercase_	:   Tuple  				=     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
									lowercase_	:   List[str]  				=     3
									lowercase_	:   Dict  				=     self._get_uniform_logits(snake_case__,   snake_case__  )
									lowercase_	:   Union[str, Any]  				=     logits_processor(snake_case__,   snake_case__,   cur_len=snake_case__  )
									self.assertFalse(jnp.isinf(snake_case__  ).any()  )






		def 						snake_case__				(							self  )       -> List[Any]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     4
									lowercase_	:   Optional[int]  				=     10
									lowercase_	:   Optional[int]  				=     15
									lowercase_	:   List[Any]  				=     2
									lowercase_	:   Tuple  				=     1
									lowercase_	:   Tuple  				=     15

									# dummy input_ids and scores
									lowercase_	:   str  				=     ids_tensor((batch_size, sequence_length),   snake_case__  )
									lowercase_	:   Any  				=     input_ids.copy()

									lowercase_	:   List[Any]  				=     self._get_uniform_logits(snake_case__,   snake_case__  )
									lowercase_	:   List[Any]  				=     scores.copy()

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

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

									lowercase_	:   Optional[Any]  				=     10

									# no processor list
									lowercase_	:   Tuple  				=     temp_dist_warp(snake_case__,   snake_case__,   cur_len=snake_case__  )
									lowercase_	:   int  				=     top_k_warp(snake_case__,   snake_case__,   cur_len=snake_case__  )
									lowercase_	:   List[str]  				=     top_p_warp(snake_case__,   snake_case__,   cur_len=snake_case__  )
									lowercase_	:   int  				=     min_dist_proc(snake_case__,   snake_case__,   cur_len=snake_case__  )
									lowercase_	:   Tuple  				=     bos_dist_proc(snake_case__,   snake_case__,   cur_len=snake_case__  )
									lowercase_	:   Optional[int]  				=     eos_dist_proc(snake_case__,   snake_case__,   cur_len=snake_case__  )

									# with processor list
									lowercase_	:   List[Any]  				=     FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]  )
									lowercase_	:   List[str]  				=     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 						snake_case__				(							self  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Union[str, Any]  				=     4
									lowercase_	:   int  				=     10
									lowercase_	:   Dict  				=     15
									lowercase_	:   Optional[int]  				=     2
									lowercase_	:   List[Any]  				=     1
									lowercase_	:   Tuple  				=     15

									# dummy input_ids and scores
									lowercase_	:   Optional[int]  				=     ids_tensor((batch_size, sequence_length),   snake_case__  )
									lowercase_	:   int  				=     input_ids.copy()

									lowercase_	:   List[Any]  				=     self._get_uniform_logits(snake_case__,   snake_case__  )
									lowercase_	:   Any  				=     scores.copy()

									# instantiate all dist processors
									lowercase_	:   Optional[Any]  				=     FlaxTemperatureLogitsWarper(temperature=0.5  )
									lowercase_	:   Union[str, Any]  				=     FlaxTopKLogitsWarper(3  )
									lowercase_	:   Tuple  				=     FlaxTopPLogitsWarper(0.8  )

									# instantiate all logits processors
									lowercase_	:   Any  				=     FlaxMinLengthLogitsProcessor(min_length=10,   eos_token_id=snake_case__  )
									lowercase_	:   Union[str, Any]  				=     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__  )
									lowercase_	:   Dict  				=     FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__,   eos_token_id=snake_case__  )

									lowercase_	:   str  				=     10

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

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

									lowercase_	:   Dict  				=     jax.jit(snake_case__  )
									lowercase_	:   Dict  				=     jax.jit(snake_case__  )

									lowercase_	:   Optional[Any]  				=     jitted_run_no_processor_list(snake_case__,   snake_case__,   snake_case__  )
									lowercase_	:   int  				=     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()  )






from string import ascii_uppercase

UpperCAmelCase_							=				{char: i for i, char in enumerate(ascii_uppercase)}
UpperCAmelCase_							=				dict(enumerate(ascii_uppercase))





def  __magic_name__  (				lowercase				,       lowercase				)				-> str:

							"""simple docstring"""



							lowercase_	:   Dict  				=     len(lowercase				)
							lowercase_	:   Tuple  				=     0
							while True:
														if x == i:
																					lowercase_	:   Dict  				=     0
														if len(lowercase				) == len(lowercase				):
																					break
														key += key[i]
														i += 1
							return key





def  __magic_name__  (				lowercase				,       lowercase				)				-> str:

							"""simple docstring"""



							lowercase_	:   Any  				=     """"""
							lowercase_	:   Optional[Any]  				=     0
							for letter in message:
														if letter == " ":
																					cipher_text += " "
														else:
																					lowercase_	:   Any  				=     (dicta[letter] - dicta[key_new[i]]) % 26
																					i += 1
																					cipher_text += dicta[x]
							return cipher_text





def  __magic_name__  (				lowercase				,       lowercase				)				-> str:

							"""simple docstring"""



							lowercase_	:   Optional[Any]  				=     """"""
							lowercase_	:   Optional[Any]  				=     0
							for letter in cipher_text:
														if letter == " ":
																					or_txt += " "
														else:
																					lowercase_	:   Optional[int]  				=     (dicta[letter] + dicta[key_new[i]] + 26) % 26
																					i += 1
																					or_txt += dicta[x]
							return or_txt





def  __magic_name__  (				)				-> None:

							"""simple docstring"""



							lowercase_	:   Union[str, Any]  				=     """THE GERMAN ATTACK"""
							lowercase_	:   str  				=     """SECRET"""
							lowercase_	:   int  				=     generate_key(lowercase				,       lowercase				)
							lowercase_	:   List[str]  				=     cipher_text(lowercase				,       lowercase				)
							print(f"""Encrypted Text = {s}"""				)
							print(f"""Original Text = {original_text(lowercase				,       lowercase				)}"""				)


if __name__ == "__main__":
							import doctest

							doctest.testmod()
							main()