
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


UpperCamelCase_   	=	{
    "configuration_mega": ["MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP", "MegaConfig", "MegaOnnxConfig"],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					UpperCamelCase_   	=	[
					    "MEGA_PRETRAINED_MODEL_ARCHIVE_LIST",
					    "MegaForCausalLM",
					    "MegaForMaskedLM",
					    "MegaForMultipleChoice",
					    "MegaForQuestionAnswering",
					    "MegaForSequenceClassification",
					    "MegaForTokenClassification",
					    "MegaModel",
					    "MegaPreTrainedModel",
					]

if TYPE_CHECKING:
					from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mega import (
										    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MegaForCausalLM,
										    MegaForMaskedLM,
										    MegaForMultipleChoice,
										    MegaForQuestionAnswering,
										    MegaForSequenceClassification,
										    MegaForTokenClassification,
										    MegaModel,
										    MegaPreTrainedModel,
										)

else:
					import sys

					UpperCamelCase_   	=	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)








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







							def       UpperCAmelCase__					(						self  :  Tuple      ,				snake_case__  :  int      ,				snake_case__  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							jnp.ones((batch_size, length) ) / length
								return scores







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							20

								__lowerCAmelCase         =							self._get_uniform_logits(batch_size=2      ,				length=snake_case__ )

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

								# compute softmax
								__lowerCAmelCase         =							jax.nn.softmax(snake_case__      ,				axis=-1 )

								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=0.5 )
								__lowerCAmelCase         =							FlaxTemperatureLogitsWarper(temperature=1.3 )

								__lowerCAmelCase         =							jax.nn.softmax(temp_dist_warper_sharper(snake_case__      ,				scores.copy()      ,				cur_len=snake_case__ )      ,				axis=-1 )
								__lowerCAmelCase         =							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       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							2

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

								__lowerCAmelCase         =							FlaxTopKLogitsWarper(3 )

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

								__lowerCAmelCase         =							np.broadcast_to(np.arange(snake_case__ )[None, :]      ,				(batch_size, length) ).copy()
								__lowerCAmelCase         =							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       UpperCAmelCase__					(						self  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							None
								__lowerCAmelCase         =							10
								__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(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
								__lowerCAmelCase         =							np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]] )
								self.assertTrue(np.allclose(snake_case__      ,				snake_case__      ,				atol=1E-3 ) )

								# check edge cases with negative and extreme logits
								__lowerCAmelCase         =							np.broadcast_to(np.arange(snake_case__ )[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(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       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0

								__lowerCAmelCase         =							FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=snake_case__ )

								# check that min length is applied at length 5
								__lowerCAmelCase         =							ids_tensor((batch_size, 20)      ,				vocab_size=20 )
								__lowerCAmelCase         =							5
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							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
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							min_dist_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertFalse(jnp.isinf(snake_case__ ).any() )







							def       UpperCAmelCase__					(						self  :  int ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0

								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )

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







							def       UpperCAmelCase__					(						self  :  str ):
								"""simple docstring"""







								__lowerCAmelCase         =							20
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							0
								__lowerCAmelCase         =							5

								__lowerCAmelCase         =							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
								__lowerCAmelCase         =							ids_tensor((batch_size, 4)      ,				vocab_size=20 )
								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							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
								__lowerCAmelCase         =							3
								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							logits_processor(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								self.assertFalse(jnp.isinf(snake_case__ ).any() )







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							2
								__lowerCAmelCase         =							1
								__lowerCAmelCase         =							15

								# dummy input_ids and scores
								__lowerCAmelCase         =							ids_tensor((batch_size, sequence_length)      ,				snake_case__ )
								__lowerCAmelCase         =							input_ids.copy()

								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__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=10      ,				eos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,				eos_token_id=snake_case__ )

								__lowerCAmelCase         =							10

								# no processor list
								__lowerCAmelCase         =							temp_dist_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							top_k_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							min_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							bos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
								__lowerCAmelCase         =							eos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )

								# 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(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       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								__lowerCAmelCase         =							4
								__lowerCAmelCase         =							10
								__lowerCAmelCase         =							15
								__lowerCAmelCase         =							2
								__lowerCAmelCase         =							1
								__lowerCAmelCase         =							15

								# dummy input_ids and scores
								__lowerCAmelCase         =							ids_tensor((batch_size, sequence_length)      ,				snake_case__ )
								__lowerCAmelCase         =							input_ids.copy()

								__lowerCAmelCase         =							self._get_uniform_logits(snake_case__      ,				snake_case__ )
								__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=10      ,				eos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=snake_case__ )
								__lowerCAmelCase         =							FlaxForcedEOSTokenLogitsProcessor(max_length=snake_case__      ,				eos_token_id=snake_case__ )

								__lowerCAmelCase         =							10

								# no processor list
								def run_no_processor_list(snake_case__  :  int      ,				snake_case__  :  Any      ,				snake_case__  :  int ):
									__lowerCAmelCase         =							temp_dist_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							top_k_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							top_p_warp(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							min_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							bos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									__lowerCAmelCase         =							eos_dist_proc(snake_case__      ,				snake_case__      ,				cur_len=snake_case__ )
									return scores

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

								__lowerCAmelCase         =							jax.jit(snake_case__ )
								__lowerCAmelCase         =							jax.jit(snake_case__ )

								__lowerCAmelCase         =							jitted_run_no_processor_list(snake_case__      ,				snake_case__      ,				snake_case__ )
								__lowerCAmelCase         =							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() )






