


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






     def    lowerCamelCase__   (		self     :    Any    , UpperCamelCase_     :    int    , UpperCamelCase_     :    int ):
           lowerCAmelCase    :		Optional[int]				=      jnp.ones((batch_size, length) ) / length
           return scores






     def    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		List[Any]				=      None
           lowerCAmelCase    :		Optional[Any]				=      2_0

           lowerCAmelCase    :		Any				=      self._get_uniform_logits(batch_size=2    , length=UpperCamelCase_ )

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

           # compute softmax
           lowerCAmelCase    :		Dict				=      jax.nn.softmax(UpperCamelCase_    , axis=-1 )

           lowerCAmelCase    :		Optional[Any]				=      FlaxTemperatureLogitsWarper(temperature=0.5 )
           lowerCAmelCase    :		List[str]				=      FlaxTemperatureLogitsWarper(temperature=1.3 )

           lowerCAmelCase    :		Optional[Any]				=      jax.nn.softmax(temp_dist_warper_sharper(UpperCamelCase_    , scores.copy()    , cur_len=UpperCamelCase_ )    , axis=-1 )
           lowerCAmelCase    :		Optional[Any]				=      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    lowerCamelCase__   (		self     :    List[str] ):
           lowerCAmelCase    :		List[str]				=      None
           lowerCAmelCase    :		Optional[int]				=      1_0
           lowerCAmelCase    :		str				=      2

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

           lowerCAmelCase    :		Optional[int]				=      FlaxTopKLogitsWarper(3 )

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

           lowerCAmelCase    :		Dict				=      np.broadcast_to(np.arange(UpperCamelCase_ )[None, :]    , (batch_size, length) ).copy()
           lowerCAmelCase    :		Optional[int]				=      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    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		int				=      None
           lowerCAmelCase    :		Dict				=      1_0
           lowerCAmelCase    :		List[str]				=      2

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

           lowerCAmelCase    :		List[str]				=      FlaxTopPLogitsWarper(0.8 )
           lowerCAmelCase    :		Optional[Any]				=      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
           lowerCAmelCase    :		Dict				=      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
           lowerCAmelCase    :		Optional[Any]				=      np.broadcast_to(np.arange(UpperCamelCase_ )[None, :]    , (batch_size, vocab_size) ).copy() - (
               vocab_size // 2
           )

           # make ramp_logits more extreme
           lowerCAmelCase    :		Optional[int]				=      ramp_logits[1] * 100.0

           # make sure at least 2 tokens are kept
           lowerCAmelCase    :		Dict				=      FlaxTopPLogitsWarper(0.9    , min_tokens_to_keep=2    , filter_value=0.0 )
           lowerCAmelCase    :		str				=      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    lowerCamelCase__   (		self     :    Optional[Any] ):
           lowerCAmelCase    :		Any				=      2_0
           lowerCAmelCase    :		Optional[int]				=      4
           lowerCAmelCase    :		Dict				=      0

           lowerCAmelCase    :		Any				=      FlaxMinLengthLogitsProcessor(min_length=1_0    , eos_token_id=UpperCamelCase_ )

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






     def    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		str				=      2_0
           lowerCAmelCase    :		Union[str, Any]				=      4
           lowerCAmelCase    :		Optional[Any]				=      0

           lowerCAmelCase    :		int				=      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase_ )

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






     def    lowerCamelCase__   (		self     :    List[Any] ):
           lowerCAmelCase    :		Optional[Any]				=      2_0
           lowerCAmelCase    :		Dict				=      4
           lowerCAmelCase    :		Tuple				=      0
           lowerCAmelCase    :		Any				=      5

           lowerCAmelCase    :		Dict				=      FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase_    , eos_token_id=UpperCamelCase_ )

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






     def    lowerCamelCase__   (		self     :    str ):
           lowerCAmelCase    :		Union[str, Any]				=      4
           lowerCAmelCase    :		Tuple				=      1_0
           lowerCAmelCase    :		Union[str, Any]				=      1_5
           lowerCAmelCase    :		Union[str, Any]				=      2
           lowerCAmelCase    :		int				=      1
           lowerCAmelCase    :		Tuple				=      1_5

           # dummy input_ids and scores
           lowerCAmelCase    :		Union[str, Any]				=      ids_tensor((batch_size, sequence_length)    , UpperCamelCase_ )
           lowerCAmelCase    :		Optional[int]				=      input_ids.copy()

           lowerCAmelCase    :		Tuple				=      self._get_uniform_logits(UpperCamelCase_    , UpperCamelCase_ )
           lowerCAmelCase    :		List[str]				=      scores.copy()

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

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

           lowerCAmelCase    :		List[str]				=      1_0

           # no processor list
           lowerCAmelCase    :		Dict				=      temp_dist_warp(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
           lowerCAmelCase    :		Dict				=      top_k_warp(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
           lowerCAmelCase    :		Optional[Any]				=      top_p_warp(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
           lowerCAmelCase    :		Any				=      min_dist_proc(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
           lowerCAmelCase    :		int				=      bos_dist_proc(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
           lowerCAmelCase    :		List[Any]				=      eos_dist_proc(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )

           # with processor list
           lowerCAmelCase    :		Tuple				=      FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
           lowerCAmelCase    :		Optional[int]				=      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    lowerCamelCase__   (		self     :    Dict ):
           lowerCAmelCase    :		Dict				=      4
           lowerCAmelCase    :		str				=      1_0
           lowerCAmelCase    :		str				=      1_5
           lowerCAmelCase    :		Union[str, Any]				=      2
           lowerCAmelCase    :		List[Any]				=      1
           lowerCAmelCase    :		List[Any]				=      1_5

           # dummy input_ids and scores
           lowerCAmelCase    :		int				=      ids_tensor((batch_size, sequence_length)    , UpperCamelCase_ )
           lowerCAmelCase    :		Dict				=      input_ids.copy()

           lowerCAmelCase    :		str				=      self._get_uniform_logits(UpperCamelCase_    , UpperCamelCase_ )
           lowerCAmelCase    :		Any				=      scores.copy()

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

           # instantiate all logits processors
           lowerCAmelCase    :		str				=      FlaxMinLengthLogitsProcessor(min_length=1_0    , eos_token_id=UpperCamelCase_ )
           lowerCAmelCase    :		Tuple				=      FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCamelCase_ )
           lowerCAmelCase    :		List[str]				=      FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCamelCase_    , eos_token_id=UpperCamelCase_ )

           lowerCAmelCase    :		Optional[Any]				=      1_0

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

           # with processor list
           def run_processor_list(UpperCamelCase_     :    int    , UpperCamelCase_     :    Optional[int]    , UpperCamelCase_     :    Dict ):
                 lowerCAmelCase    :		List[Any]				=      FlaxLogitsProcessorList(
                     [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc] )
                 lowerCAmelCase    :		Optional[int]				=      processor(UpperCamelCase_    , UpperCamelCase_    , cur_len=UpperCamelCase_ )
                 return scores

           lowerCAmelCase    :		Any				=      jax.jit(UpperCamelCase_ )
           lowerCAmelCase    :		int				=      jax.jit(UpperCamelCase_ )

           lowerCAmelCase    :		str				=      jitted_run_no_processor_list(UpperCamelCase_    , UpperCamelCase_    , UpperCamelCase_ )
           lowerCAmelCase    :		int				=      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() )






"""simple docstring"""





from datetime import datetime

import matplotlib.pyplot as plt
import torch


def    _snake_case  (  _snake_case :	int	):
      for param in module.parameters():
            lowerCAmelCase    :		Optional[int]				=      False


def    _snake_case  (  ):
      lowerCAmelCase    :		List[str]				=      '''cuda''' if torch.cuda.is_available() else '''cpu'''
      if torch.backends.mps.is_available() and torch.backends.mps.is_built():
            lowerCAmelCase    :		Any				=      '''mps'''
      if device == "mps":
            print(
                '''WARNING: MPS currently doesn\'t seem to work, and messes up backpropagation without any visible torch'''
                ''' errors. I recommend using CUDA on a colab notebook or CPU instead if you\'re facing inexplicable issues'''
                ''' with generations.'''	)
      return device


def    _snake_case  (  _snake_case :	Dict	):
      lowerCAmelCase    :		Optional[int]				=      plt.imshow(_snake_case	)
      fig.axes.get_xaxis().set_visible(_snake_case	)
      fig.axes.get_yaxis().set_visible(_snake_case	)
      plt.show()






def    _snake_case  (  ):
      lowerCAmelCase    :		List[str]				=      datetime.now()
      lowerCAmelCase    :		Union[str, Any]				=      current_time.strftime('''%H:%M:%S'''	)
      return timestamp


