


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 			SCREAMING_SNAKE_CASE__  (    unittest.TestCase      ):
      def 	SCREAMING_SNAKE_CASE							(	self    :	List[Any]					,							SCREAMING_SNAKE_CASE__    :	int					,							SCREAMING_SNAKE_CASE__    :	int					)		->      List[str]:
             a_ :		List[str]							=  jnp.ones((batch_size, length)					) / length
             return scores
      def 	SCREAMING_SNAKE_CASE							(	self    :	str					)		->      Tuple:
             a_ :		Dict							=  None
             a_ :		List[str]							=  2_0

             a_ :		List[str]							=  self._get_uniform_logits(batch_size=2					,							length=SCREAMING_SNAKE_CASE__					)

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

             # compute softmax
             a_ :		Union[str, Any]							=  jax.nn.softmax(SCREAMING_SNAKE_CASE__					,							axis=-1					)

             a_ :		List[Any]							=  FlaxTemperatureLogitsWarper(temperature=0.5					)
             a_ :		Tuple							=  FlaxTemperatureLogitsWarper(temperature=1.3					)

             a_ :		Tuple							=  jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE__					,							scores.copy()					,							cur_len=SCREAMING_SNAKE_CASE__					)					,							axis=-1					)
             a_ :		Optional[Any]							=  jax.nn.softmax(temp_dist_warper_smoother(SCREAMING_SNAKE_CASE__					,							scores.copy()					,							cur_len=SCREAMING_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 	SCREAMING_SNAKE_CASE							(	self    :	Any					)		->      Optional[int]:
             a_ :		Union[str, Any]							=  None
             a_ :		List[str]							=  1_0
             a_ :		Tuple							=  2

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

             a_ :		str							=  FlaxTopKLogitsWarper(3					)

             a_ :		List[Any]							=  top_k_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_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
             a_ :		Tuple							=  5
             a_ :		Union[str, Any]							=  FlaxTopKLogitsWarper(top_k=1					,							filter_value=0.0					,							min_tokens_to_keep=3					)

             a_ :		int							=  np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE__					)[None, :]					,							(batch_size, length)					).copy()
             a_ :		List[Any]							=  top_k_warp_safety_check(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_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 	SCREAMING_SNAKE_CASE							(	self    :	List[Any]					)		->      Optional[int]:
             a_ :		Dict							=  None
             a_ :		Optional[Any]							=  1_0
             a_ :		Optional[int]							=  2

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

             a_ :		Any							=  FlaxTopPLogitsWarper(0.8					)
             a_ :		Dict							=  np.exp(top_p_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)					)

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

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

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

             # make sure at least 2 tokens are kept
             a_ :		Dict							=  FlaxTopPLogitsWarper(0.9					,							min_tokens_to_keep=2					,							filter_value=0.0					)
             a_ :		Any							=  top_p_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_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 	SCREAMING_SNAKE_CASE							(	self    :	List[str]					)		->      Tuple:
             a_ :		Optional[int]							=  2_0
             a_ :		Tuple							=  4
             a_ :		str							=  0

             a_ :		Union[str, Any]							=  FlaxMinLengthLogitsProcessor(min_length=1_0					,							eos_token_id=SCREAMING_SNAKE_CASE__					)

             # check that min length is applied at length 5
             a_ :		List[Any]							=  ids_tensor((batch_size, 2_0)					,							vocab_size=2_0					)
             a_ :		Union[str, Any]							=  5
             a_ :		Union[str, Any]							=  self._get_uniform_logits(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
             a_ :		List[str]							=  min_dist_processor(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_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
             a_ :		str							=  self._get_uniform_logits(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
             a_ :		Dict							=  1_5
             a_ :		List[str]							=  min_dist_processor(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE__					).any()					)
      def 	SCREAMING_SNAKE_CASE							(	self    :	Union[str, Any]					)		->      Dict:
             a_ :		Any							=  2_0
             a_ :		int							=  4
             a_ :		Optional[int]							=  0

             a_ :		List[str]							=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__					)

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

             a_ :		Union[str, Any]							=  FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE__					,							eos_token_id=SCREAMING_SNAKE_CASE__					)

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

             # dummy input_ids and scores
             a_ :		List[str]							=  ids_tensor((batch_size, sequence_length)					,							SCREAMING_SNAKE_CASE__					)
             a_ :		Tuple							=  input_ids.copy()

             a_ :		Dict							=  self._get_uniform_logits(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
             a_ :		List[Any]							=  scores.copy()

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

             # instantiate all logits processors
             a_ :		Dict							=  FlaxMinLengthLogitsProcessor(min_length=1_0					,							eos_token_id=SCREAMING_SNAKE_CASE__					)
             a_ :		int							=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE__					)
             a_ :		Any							=  FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE__					,							eos_token_id=SCREAMING_SNAKE_CASE__					)

             a_ :		Optional[Any]							=  1_0

             # no processor list
             a_ :		List[str]							=  temp_dist_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             a_ :		str							=  top_k_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             a_ :		Union[str, Any]							=  top_p_warp(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             a_ :		Any							=  min_dist_proc(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             a_ :		Optional[Any]							=  bos_dist_proc(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)
             a_ :		Tuple							=  eos_dist_proc(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)

             # with processor list
             a_ :		Any							=  FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]					)
             a_ :		Optional[Any]							=  processor(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							cur_len=SCREAMING_SNAKE_CASE__					)

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

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

      def 	SCREAMING_SNAKE_CASE							(	self    :	Optional[Any]					)		->      Tuple:
             a_ :		Union[str, Any]							=  4
             a_ :		Tuple							=  1_0
             a_ :		Any							=  1_5
             a_ :		Optional[int]							=  2
             a_ :		List[Any]							=  1
             a_ :		List[Any]							=  1_5

             # dummy input_ids and scores
             a_ :		str							=  ids_tensor((batch_size, sequence_length)					,							SCREAMING_SNAKE_CASE__					)
             a_ :		Dict							=  input_ids.copy()

             a_ :		Tuple							=  self._get_uniform_logits(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
             a_ :		Any							=  scores.copy()

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

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

             a_ :		List[str]							=  1_0

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

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

             a_ :		Optional[int]							=  jax.jit(SCREAMING_SNAKE_CASE__					)
             a_ :		Optional[Any]							=  jax.jit(SCREAMING_SNAKE_CASE__					)

             a_ :		List[str]							=  jitted_run_no_processor_list(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)
             a_ :		List[Any]							=  jitted_run_processor_list(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					)

             # scores should be equal
             self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE__					,							SCREAMING_SNAKE_CASE__					,							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_				:			Dict   =	{'configuration_vit_msn': ['VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP', 'ViTMSNConfig']}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase_				:			int   =	[
       'VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST',
       'ViTMSNModel',
       'ViTMSNForImageClassification',
       'ViTMSNPreTrainedModel',
   ]

if TYPE_CHECKING:
   from .configuration_vit_msn import VIT_MSN_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTMSNConfig

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vit_msn import (
          VIT_MSN_PRETRAINED_MODEL_ARCHIVE_LIST,
          ViTMSNForImageClassification,
          ViTMSNModel,
          ViTMSNPreTrainedModel,
      )

else:
   import sys

   UpperCAmelCase_				:			Dict   =	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)


