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







				"""simple docstring"""






				def 		_lowerCAmelCase		(   self				:				List[str]   ,     _A				:				int   ,     _A				:				Any):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Dict      	=     jnp.ones((batch_size, length)) / length
											return scores





				def 		_lowerCAmelCase		(   self				:				str):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     None
											_SCREAMING_SNAKE_CASE		:    str      	=     2_0

											_SCREAMING_SNAKE_CASE		:    Dict      	=     self._get_uniform_logits(batch_size=2   ,     length=_A)

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

											# compute softmax
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     jax.nn.softmax(_A   ,     axis=-1)

											_SCREAMING_SNAKE_CASE		:    str      	=     FlaxTemperatureLogitsWarper(temperature=0.5)
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     FlaxTemperatureLogitsWarper(temperature=1.3)

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



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     None
											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     1_0
											_SCREAMING_SNAKE_CASE		:    Dict      	=     2

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

											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     FlaxTopKLogitsWarper(3)

											_SCREAMING_SNAKE_CASE		:    Tuple      	=     top_k_warp(_A   ,     _A   ,     cur_len=_A)

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

											_SCREAMING_SNAKE_CASE		:    str      	=     np.broadcast_to(np.arange(_A)[None, :]   ,     (batch_size, length)).copy()
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     top_k_warp_safety_check(_A   ,     _A   ,     cur_len=_A)

											# 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				:				List[Any]):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    List[str]      	=     None
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     1_0
											_SCREAMING_SNAKE_CASE		:    int      	=     2

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

											_SCREAMING_SNAKE_CASE		:    List[str]      	=     FlaxTopPLogitsWarper(0.8)
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     np.exp(top_p_warp(_A   ,     _A   ,     cur_len=_A))

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

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

											# make ramp_logits more extreme
											_SCREAMING_SNAKE_CASE		:    Dict      	=     ramp_logits[1] * 1_0_0.0

											# make sure at least 2 tokens are kept
											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     FlaxTopPLogitsWarper(0.9   ,     min_tokens_to_keep=2   ,     filter_value=0.0)
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     top_p_warp(_A   ,     _A   ,     cur_len=_A)

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



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Dict      	=     2_0
											_SCREAMING_SNAKE_CASE		:    str      	=     4
											_SCREAMING_SNAKE_CASE		:    Dict      	=     0

											_SCREAMING_SNAKE_CASE		:    List[str]      	=     FlaxMinLengthLogitsProcessor(min_length=1_0   ,     eos_token_id=_A)

											# check that min length is applied at length 5
											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     ids_tensor((batch_size, 2_0)   ,     vocab_size=2_0)
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     5
											_SCREAMING_SNAKE_CASE		:    str      	=     self._get_uniform_logits(_A   ,     _A)
											_SCREAMING_SNAKE_CASE		:    List[str]      	=     min_dist_processor(_A   ,     _A   ,     cur_len=_A)
											self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()   ,     4 * [-float("""inf""")])

											# check that min length is not applied anymore at length 15
											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     self._get_uniform_logits(_A   ,     _A)
											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     1_5
											_SCREAMING_SNAKE_CASE		:    int      	=     min_dist_processor(_A   ,     _A   ,     cur_len=_A)
											self.assertFalse(jnp.isinf(_A).any())





				def 		_lowerCAmelCase		(   self				:				Any):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     2_0
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     4
											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     0

											_SCREAMING_SNAKE_CASE		:    str      	=     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A)

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





				def 		_lowerCAmelCase		(   self				:				Any):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Tuple      	=     2_0
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     4
											_SCREAMING_SNAKE_CASE		:    Dict      	=     0
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     5

											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     FlaxForcedEOSTokenLogitsProcessor(max_length=_A   ,     eos_token_id=_A)

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





				def 		_lowerCAmelCase		(   self				:				Tuple):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    Dict      	=     4
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     1_0
											_SCREAMING_SNAKE_CASE		:    str      	=     1_5
											_SCREAMING_SNAKE_CASE		:    Any      	=     2
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     1
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     1_5

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

											_SCREAMING_SNAKE_CASE		:    List[str]      	=     self._get_uniform_logits(_A   ,     _A)
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     scores.copy()

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

											# instantiate all logits processors
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     FlaxMinLengthLogitsProcessor(min_length=1_0   ,     eos_token_id=_A)
											_SCREAMING_SNAKE_CASE		:    str      	=     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A)
											_SCREAMING_SNAKE_CASE		:    Any      	=     FlaxForcedEOSTokenLogitsProcessor(max_length=_A   ,     eos_token_id=_A)

											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     1_0

											# no processor list
											_SCREAMING_SNAKE_CASE		:    str      	=     temp_dist_warp(_A   ,     _A   ,     cur_len=_A)
											_SCREAMING_SNAKE_CASE		:    str      	=     top_k_warp(_A   ,     _A   ,     cur_len=_A)
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     top_p_warp(_A   ,     _A   ,     cur_len=_A)
											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     min_dist_proc(_A   ,     _A   ,     cur_len=_A)
											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     bos_dist_proc(_A   ,     _A   ,     cur_len=_A)
											_SCREAMING_SNAKE_CASE		:    Optional[int]      	=     eos_dist_proc(_A   ,     _A   ,     cur_len=_A)

											# with processor list
											_SCREAMING_SNAKE_CASE		:    List[str]      	=     FlaxLogitsProcessorList(
											    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
											_SCREAMING_SNAKE_CASE		:    List[str]      	=     processor(_A   ,     _A   ,     cur_len=_A)

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

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





				def 		_lowerCAmelCase		(   self				:				Tuple):



											"""simple docstring"""







											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     4
											_SCREAMING_SNAKE_CASE		:    Any      	=     1_0
											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     1_5
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     2
											_SCREAMING_SNAKE_CASE		:    Any      	=     1
											_SCREAMING_SNAKE_CASE		:    Optional[Any]      	=     1_5

											# dummy input_ids and scores
											_SCREAMING_SNAKE_CASE		:    Tuple      	=     ids_tensor((batch_size, sequence_length)   ,     _A)
											_SCREAMING_SNAKE_CASE		:    str      	=     input_ids.copy()

											_SCREAMING_SNAKE_CASE		:    int      	=     self._get_uniform_logits(_A   ,     _A)
											_SCREAMING_SNAKE_CASE		:    str      	=     scores.copy()

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

											# instantiate all logits processors
											_SCREAMING_SNAKE_CASE		:    Any      	=     FlaxMinLengthLogitsProcessor(min_length=1_0   ,     eos_token_id=_A)
											_SCREAMING_SNAKE_CASE		:    List[Any]      	=     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_A)
											_SCREAMING_SNAKE_CASE		:    int      	=     FlaxForcedEOSTokenLogitsProcessor(max_length=_A   ,     eos_token_id=_A)

											_SCREAMING_SNAKE_CASE		:    Dict      	=     1_0

											# no processor list
											def run_no_processor_list(_A				:				Union[str, Any]   ,     _A				:				Any   ,     _A				:				Tuple):
																		_SCREAMING_SNAKE_CASE		:    str      	=     temp_dist_warp(_A   ,     _A   ,     cur_len=_A)
																		_SCREAMING_SNAKE_CASE		:    Dict      	=     top_k_warp(_A   ,     _A   ,     cur_len=_A)
																		_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     top_p_warp(_A   ,     _A   ,     cur_len=_A)
																		_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     min_dist_proc(_A   ,     _A   ,     cur_len=_A)
																		_SCREAMING_SNAKE_CASE		:    List[Any]      	=     bos_dist_proc(_A   ,     _A   ,     cur_len=_A)
																		_SCREAMING_SNAKE_CASE		:    Any      	=     eos_dist_proc(_A   ,     _A   ,     cur_len=_A)
																		return scores

											# with processor list
											def run_processor_list(_A				:				Tuple   ,     _A				:				Union[str, Any]   ,     _A				:				List[str]):
																		_SCREAMING_SNAKE_CASE		:    List[Any]      	=     FlaxLogitsProcessorList(
																		    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
																		_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     processor(_A   ,     _A   ,     cur_len=_A)
																		return scores

											_SCREAMING_SNAKE_CASE		:    Union[str, Any]      	=     jax.jit(_A)
											_SCREAMING_SNAKE_CASE		:    Any      	=     jax.jit(_A)

											_SCREAMING_SNAKE_CASE		:    Any      	=     jitted_run_no_processor_list(_A   ,     _A   ,     _A)
											_SCREAMING_SNAKE_CASE		:    Any      	=     jitted_run_processor_list(_A   ,     _A   ,     _A)

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

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




"""simple docstring"""


import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
						import optax
						from flax.training.common_utils import onehot

						from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
						from transformers.models.ta.modeling_flax_ta import shift_tokens_right




@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class  __lowerCamelCase  ( unittest.TestCase			):




				@slow
				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  FlaxMTaForConditionalGeneration.from_pretrained('''google/mt5-small'''   )
							lowerCamelCase_					=  AutoTokenizer.from_pretrained('''google/mt5-small'''   )

							lowerCamelCase_					=  tokenizer('''Hello there'''							,						return_tensors='''np'''   ).input_ids
							lowerCamelCase_					=  tokenizer('''Hi I am'''							,						return_tensors='''np'''   ).input_ids

							lowerCamelCase_					=  shift_tokens_right(UpperCAmelCase							,						model.config.pad_token_id							,						model.config.decoder_start_token_id   )

							lowerCamelCase_					=  model(UpperCAmelCase							,						decoder_input_ids=UpperCAmelCase   ).logits
							lowerCamelCase_					=  optax.softmax_cross_entropy(UpperCAmelCase							,						onehot(UpperCAmelCase							,						logits.shape[-1]   )   ).mean()

							lowerCamelCase_					=  -(labels.shape[-1] * loss.item())

							lowerCamelCase_					=  -8_4.9_1_2_7
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE   ) < 1e-4   )



