






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   ):







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  int					):







									'''simple docstring'''




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







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  20

									_SCREAMING_SNAKE_CASE     				=  self._get_uniform_logits(batch_size=2				,     length=UpperCAmelCase_					)

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

									# compute softmax
									_SCREAMING_SNAKE_CASE     				=  jax.nn.softmax(UpperCAmelCase_				,     axis=-1					)

									_SCREAMING_SNAKE_CASE     				=  FlaxTemperatureLogitsWarper(temperature=0.5					)
									_SCREAMING_SNAKE_CASE     				=  FlaxTemperatureLogitsWarper(temperature=1.3					)

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







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  10
									_SCREAMING_SNAKE_CASE     				=  2

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

									_SCREAMING_SNAKE_CASE     				=  FlaxTopKLogitsWarper(3					)

									_SCREAMING_SNAKE_CASE     				=  top_k_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)

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

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

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







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  10
									_SCREAMING_SNAKE_CASE     				=  2

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

									_SCREAMING_SNAKE_CASE     				=  FlaxTopPLogitsWarper(0.8					)
									_SCREAMING_SNAKE_CASE     				=  np.exp(top_p_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)					)

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

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

									# make ramp_logits more extreme
									_SCREAMING_SNAKE_CASE     				=  ramp_logits[1] * 1_00.0

									# make sure at least 2 tokens are kept
									_SCREAMING_SNAKE_CASE     				=  FlaxTopPLogitsWarper(0.9				,     min_tokens_to_keep=2				,     filter_value=0.0					)
									_SCREAMING_SNAKE_CASE     				=  top_p_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)

									# 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'''




									_SCREAMING_SNAKE_CASE     				=  20
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  0

									_SCREAMING_SNAKE_CASE     				=  FlaxMinLengthLogitsProcessor(min_length=10				,     eos_token_id=UpperCAmelCase_					)

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







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  20
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  0

									_SCREAMING_SNAKE_CASE     				=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_					)

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







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  20
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  0
									_SCREAMING_SNAKE_CASE     				=  5

									_SCREAMING_SNAKE_CASE     				=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,     eos_token_id=UpperCAmelCase_					)

									# check that all scores are -inf except the eos_token_id when max_length is reached
									_SCREAMING_SNAKE_CASE     				=  ids_tensor((batch_size, 4)				,     vocab_size=20					)
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  self._get_uniform_logits(UpperCAmelCase_				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  logits_processor(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									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     				=  3
									_SCREAMING_SNAKE_CASE     				=  self._get_uniform_logits(UpperCAmelCase_				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  logits_processor(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									self.assertFalse(jnp.isinf(UpperCAmelCase_					).any()					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  10
									_SCREAMING_SNAKE_CASE     				=  15
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  1
									_SCREAMING_SNAKE_CASE     				=  15

									# dummy input_ids and scores
									_SCREAMING_SNAKE_CASE     				=  ids_tensor((batch_size, sequence_length)				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  input_ids.copy()

									_SCREAMING_SNAKE_CASE     				=  self._get_uniform_logits(UpperCAmelCase_				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  scores.copy()

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

									# instantiate all logits processors
									_SCREAMING_SNAKE_CASE     				=  FlaxMinLengthLogitsProcessor(min_length=10				,     eos_token_id=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,     eos_token_id=UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  10

									# no processor list
									_SCREAMING_SNAKE_CASE     				=  temp_dist_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  top_k_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  top_p_warp(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  min_dist_proc(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  bos_dist_proc(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  eos_dist_proc(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)

									# with processor list
									_SCREAMING_SNAKE_CASE     				=  FlaxLogitsProcessorList(
									    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]					)
									_SCREAMING_SNAKE_CASE     				=  processor(UpperCAmelCase_				,     UpperCAmelCase_				,     cur_len=UpperCAmelCase_					)

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

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







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  10
									_SCREAMING_SNAKE_CASE     				=  15
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  1
									_SCREAMING_SNAKE_CASE     				=  15

									# dummy input_ids and scores
									_SCREAMING_SNAKE_CASE     				=  ids_tensor((batch_size, sequence_length)				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  input_ids.copy()

									_SCREAMING_SNAKE_CASE     				=  self._get_uniform_logits(UpperCAmelCase_				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  scores.copy()

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

									# instantiate all logits processors
									_SCREAMING_SNAKE_CASE     				=  FlaxMinLengthLogitsProcessor(min_length=10				,     eos_token_id=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase_				,     eos_token_id=UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  10

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

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

									_SCREAMING_SNAKE_CASE     				=  jax.jit(UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  jax.jit(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  jitted_run_no_processor_list(UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  jitted_run_processor_list(UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_					)

									# scores should be equal
									self.assertTrue(jnp.allclose(UpperCAmelCase_				,     UpperCAmelCase_				,     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_torch_available


UpperCamelCase			    =     {'''configuration_mmbt''': ['''MMBTConfig''']}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCamelCase			    =     ['''MMBTForClassification''', '''MMBTModel''', '''ModalEmbeddings''']


if TYPE_CHECKING:
		from .configuration_mmbt import MMBTConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings

else:
		import sys

		UpperCamelCase			    =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

