



"""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=_lowerCAmelCase )

									# 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(_lowerCAmelCase					,							axis=-1 )

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

									_lowerCamelCase   	=					jax.nn.softmax(temp_dist_warper_sharper(_lowerCAmelCase					,							scores.copy()					,							cur_len=_lowerCAmelCase )					,							axis=-1 )
									_lowerCamelCase   	=					jax.nn.softmax(temp_dist_warper_smoother(_lowerCAmelCase					,							scores.copy()					,							cur_len=_lowerCAmelCase )					,							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(_lowerCAmelCase )[None, :]					,							(batch_size, vocab_size) ).copy()
									_lowerCamelCase   	=					ramp_logits[1:, : vocab_size // 2] + vocab_size

									_lowerCamelCase   	=					FlaxTopKLogitsWarper(3 )

									_lowerCamelCase   	=					top_k_warp(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )

									# 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(_lowerCAmelCase )[None, :]					,							(batch_size, length) ).copy()
									_lowerCamelCase   	=					top_k_warp_safety_check(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )

									# 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(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase ) )

									# 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(_lowerCAmelCase					,							_lowerCAmelCase					,							atol=1e-3 ) )

									# check edge cases with negative and extreme logits
									_lowerCamelCase   	=					np.broadcast_to(np.arange(_lowerCAmelCase )[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(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )

									# 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=_lowerCAmelCase )

									# 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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					min_dist_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					1_5
									_lowerCamelCase   	=					min_dist_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									self.assertFalse(jnp.isinf(_lowerCAmelCase ).any() )

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

									_lowerCamelCase   	=					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase )

									# 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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					logits_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					logits_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									self.assertFalse(jnp.isinf(_lowerCAmelCase ).any() )

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

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

									# 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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					logits_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									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(_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					logits_processor(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )
									self.assertFalse(jnp.isinf(_lowerCAmelCase ).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)					,							_lowerCAmelCase )
									_lowerCamelCase   	=					input_ids.copy()

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

									_lowerCamelCase   	=					1_0

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

									# 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(_lowerCAmelCase					,							_lowerCAmelCase					,							cur_len=_lowerCAmelCase )

									# scores should be equal
									self.assertTrue(jnp.allclose(_lowerCAmelCase					,							_lowerCAmelCase					,							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)					,							_lowerCAmelCase )
									_lowerCamelCase   	=					input_ids.copy()

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

									_lowerCamelCase   	=					1_0

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

									_lowerCamelCase   	=					jax.jit(_lowerCAmelCase )
									_lowerCamelCase   	=					jax.jit(_lowerCAmelCase )

									_lowerCamelCase   	=					jitted_run_no_processor_list(_lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase )
									_lowerCamelCase   	=					jitted_run_processor_list(_lowerCAmelCase					,							_lowerCAmelCase					,							_lowerCAmelCase )

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

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






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


if is_sentencepiece_available():
   from ..ta.tokenization_ta import TaTokenizer
else:
   from ...utils.dummy_sentencepiece_objects import TaTokenizer

lowerCamelCase__					:  List[Any]      							=				TaTokenizer

if is_tokenizers_available():
   from ..ta.tokenization_ta_fast import TaTokenizerFast
else:
   from ...utils.dummy_tokenizers_objects import TaTokenizerFast

lowerCamelCase__					:  Union[str, Any]      							=				TaTokenizerFast

lowerCamelCase__					:  Dict      							=				{'configuration_mt5': ['MT5Config', 'MT5OnnxConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__					:  Tuple      							=				[
       'MT5EncoderModel',
       'MT5ForConditionalGeneration',
       'MT5ForQuestionAnswering',
       'MT5Model',
       'MT5PreTrainedModel',
       'MT5Stack',
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__					:  Tuple      							=				['TFMT5EncoderModel', 'TFMT5ForConditionalGeneration', 'TFMT5Model']

try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowerCamelCase__					:  Tuple      							=				['FlaxMT5EncoderModel', 'FlaxMT5ForConditionalGeneration', 'FlaxMT5Model']


if TYPE_CHECKING:
   from .configuration_mta import MTaConfig, MTaOnnxConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_mta import (
          MTaEncoderModel,
          MTaForConditionalGeneration,
          MTaForQuestionAnswering,
          MTaModel,
          MTaPreTrainedModel,
          MTaStack,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_mta import TFMTaEncoderModel, TFMTaForConditionalGeneration, TFMTaModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_mta import FlaxMTaEncoderModel, FlaxMTaForConditionalGeneration, FlaxMTaModel

else:
   import sys

   lowerCamelCase__					:  int      							=				_LazyModule(
       __name__,
       globals()['__file__'],
       _import_structure,
       extra_objects={'MT5Tokenizer': MTaTokenizer, 'MT5TokenizerFast': MTaTokenizerFast},
       module_spec=__spec__,
   )