




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       lowerCAmelCase_  (							self  :		Optional[int]				,				_lowerCAmelCase  :		int				,				_lowerCAmelCase  :		int						):
       SCREAMING_SNAKE_CASE_							=			jnp.ones((batch_size, length)						) / length
       return scores



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			None
       SCREAMING_SNAKE_CASE_							=			20

       SCREAMING_SNAKE_CASE_							=			self._get_uniform_logits(batch_size=2				,				length=_lowerCAmelCase						)

       # 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(_lowerCAmelCase				,				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(_lowerCAmelCase				,				scores.copy()				,				cur_len=_lowerCAmelCase						)				,				axis=-1						)
       SCREAMING_SNAKE_CASE_							=			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       lowerCAmelCase_  (							self  :		int						):
       SCREAMING_SNAKE_CASE_							=			None
       SCREAMING_SNAKE_CASE_							=			10
       SCREAMING_SNAKE_CASE_							=			2

       # create ramp distribution
       SCREAMING_SNAKE_CASE_							=			np.broadcast_to(np.arange(_lowerCAmelCase						)[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(_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
       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(_lowerCAmelCase						)[None, :]				,				(batch_size, length)						).copy()
       SCREAMING_SNAKE_CASE_							=			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       lowerCAmelCase_  (							self  :		List[str]						):
       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(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)						)

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

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

       # make ramp_logits more extreme
       SCREAMING_SNAKE_CASE_							=			ramp_logits[1] * 100.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(_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       lowerCAmelCase_  (							self  :		Any						):
       SCREAMING_SNAKE_CASE_							=			20
       SCREAMING_SNAKE_CASE_							=			4
       SCREAMING_SNAKE_CASE_							=			0

       SCREAMING_SNAKE_CASE_							=			FlaxMinLengthLogitsProcessor(min_length=10				,				eos_token_id=_lowerCAmelCase						)

       # 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(_lowerCAmelCase				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			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
       SCREAMING_SNAKE_CASE_							=			self._get_uniform_logits(_lowerCAmelCase				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			15
       SCREAMING_SNAKE_CASE_							=			min_dist_processor(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
       self.assertFalse(jnp.isinf(_lowerCAmelCase						).any()						)



 def       lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_							=			20
       SCREAMING_SNAKE_CASE_							=			4
       SCREAMING_SNAKE_CASE_							=			0

       SCREAMING_SNAKE_CASE_							=			FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_lowerCAmelCase						)

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



 def       lowerCAmelCase_  (							self  :		str						):
       SCREAMING_SNAKE_CASE_							=			20
       SCREAMING_SNAKE_CASE_							=			4
       SCREAMING_SNAKE_CASE_							=			0
       SCREAMING_SNAKE_CASE_							=			5

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

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



 def       lowerCAmelCase_  (							self  :		List[Any]						):
       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)				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			input_ids.copy()

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

       SCREAMING_SNAKE_CASE_							=			10

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

       # 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(_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       lowerCAmelCase_  (							self  :		Tuple						):
       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)				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			input_ids.copy()

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

       SCREAMING_SNAKE_CASE_							=			10

       # no processor list
       def run_no_processor_list(_lowerCAmelCase  :		List[str]				,				_lowerCAmelCase  :		str				,				_lowerCAmelCase  :		List[str]						):
             SCREAMING_SNAKE_CASE_							=			temp_dist_warp(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			top_k_warp(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			top_p_warp(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			min_dist_proc(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			bos_dist_proc(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             SCREAMING_SNAKE_CASE_							=			eos_dist_proc(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             return scores

       # with processor list
       def run_processor_list(_lowerCAmelCase  :		Dict				,				_lowerCAmelCase  :		Union[str, Any]				,				_lowerCAmelCase  :		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(_lowerCAmelCase				,				_lowerCAmelCase				,				cur_len=_lowerCAmelCase						)
             return scores

       SCREAMING_SNAKE_CASE_							=			jax.jit(_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			jax.jit(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_							=			jitted_run_no_processor_list(_lowerCAmelCase				,				_lowerCAmelCase				,				_lowerCAmelCase						)
       SCREAMING_SNAKE_CASE_							=			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()						)





import warnings

from ...utils import logging
from .image_processing_imagegpt import ImageGPTImageProcessor


lowerCamelCase__ :   List[str]					=							logging.get_logger(__name__)





class        lowerCamelCase_							(      _SCREAMING_SNAKE_CASE					):


 '''simple docstring'''







 def __init__(							self  :		str				,				*_lowerCAmelCase  :		Optional[Any]				,				**_lowerCAmelCase  :		Optional[Any]						):
       warnings.warn(
           'The class ImageGPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers.'
           ' Please use ImageGPTImageProcessor instead.'				,				_lowerCAmelCase				,				)
       super().__init__(*_lowerCAmelCase				,				**_lowerCAmelCase						)