

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






					'''simple docstring'''



					def    _UpperCAmelCase    (				self				,      lowerCamelCase__				,      lowerCamelCase__						):
						UpperCAmelCase__:	Dict									= jnp.ones((batch_size, length)						) / length
						return scores





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Dict									= None
						UpperCAmelCase__:	int									= 2_0

						UpperCAmelCase__:	Optional[Any]									= self._get_uniform_logits(batch_size=2				,      length=lowerCamelCase__						)

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

						# compute softmax
						UpperCAmelCase__:	Optional[Any]									= jax.nn.softmax(lowerCamelCase__				,      axis=-1						)

						UpperCAmelCase__:	Union[str, Any]									= FlaxTemperatureLogitsWarper(temperature=0.5						)
						UpperCAmelCase__:	Any									= FlaxTemperatureLogitsWarper(temperature=1.3						)

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

						# create ramp distribution
						UpperCAmelCase__:	Union[str, Any]									= np.broadcast_to(np.arange(lowerCamelCase__						)[None, :]				,      (batch_size, vocab_size)						).copy()
						UpperCAmelCase__:	int									= ramp_logits[1:, : vocab_size // 2] + vocab_size

						UpperCAmelCase__:	str									= FlaxTopKLogitsWarper(3						)

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

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

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

						UpperCAmelCase__:	Tuple									= FlaxTopPLogitsWarper(0.8						)
						UpperCAmelCase__:	Tuple									= 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
						UpperCAmelCase__:	Optional[int]									= 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
						UpperCAmelCase__:	Union[str, Any]									= np.broadcast_to(np.arange(lowerCamelCase__						)[None, :]				,      (batch_size, vocab_size)						).copy() - (
						    vocab_size // 2
						)

						# make ramp_logits more extreme
						UpperCAmelCase__:	Dict									= ramp_logits[1] * 100.0

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

						UpperCAmelCase__:	Union[str, Any]									= FlaxMinLengthLogitsProcessor(min_length=1_0				,      eos_token_id=lowerCamelCase__						)

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





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Union[str, Any]									= 2_0
						UpperCAmelCase__:	List[Any]									= 4
						UpperCAmelCase__:	Optional[int]									= 0

						UpperCAmelCase__:	Tuple									= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase__						)

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





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	Union[str, Any]									= 2_0
						UpperCAmelCase__:	Any									= 4
						UpperCAmelCase__:	str									= 0
						UpperCAmelCase__:	Optional[int]									= 5

						UpperCAmelCase__:	Any									= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase__				,      eos_token_id=lowerCamelCase__						)

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





					def    _UpperCAmelCase    (				self						):
						UpperCAmelCase__:	int									= 4
						UpperCAmelCase__:	Tuple									= 1_0
						UpperCAmelCase__:	List[Any]									= 1_5
						UpperCAmelCase__:	str									= 2
						UpperCAmelCase__:	Optional[Any]									= 1
						UpperCAmelCase__:	List[Any]									= 1_5

						# dummy input_ids and scores
						UpperCAmelCase__:	Any									= ids_tensor((batch_size, sequence_length)				,      lowerCamelCase__						)
						UpperCAmelCase__:	List[str]									= input_ids.copy()

						UpperCAmelCase__:	Optional[int]									= self._get_uniform_logits(lowerCamelCase__				,      lowerCamelCase__						)
						UpperCAmelCase__:	List[str]									= scores.copy()

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

						# instantiate all logits processors
						UpperCAmelCase__:	int									= FlaxMinLengthLogitsProcessor(min_length=1_0				,      eos_token_id=lowerCamelCase__						)
						UpperCAmelCase__:	List[Any]									= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowerCamelCase__						)
						UpperCAmelCase__:	Optional[Any]									= FlaxForcedEOSTokenLogitsProcessor(max_length=lowerCamelCase__				,      eos_token_id=lowerCamelCase__						)

						UpperCAmelCase__:	Any									= 1_0

						# no processor list
						UpperCAmelCase__:	Union[str, Any]									= temp_dist_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
						UpperCAmelCase__:	List[str]									= top_k_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
						UpperCAmelCase__:	List[Any]									= top_p_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
						UpperCAmelCase__:	int									= min_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
						UpperCAmelCase__:	Optional[Any]									= bos_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
						UpperCAmelCase__:	Any									= eos_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)

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

						# dummy input_ids and scores
						UpperCAmelCase__:	List[str]									= ids_tensor((batch_size, sequence_length)				,      lowerCamelCase__						)
						UpperCAmelCase__:	Union[str, Any]									= input_ids.copy()

						UpperCAmelCase__:	Any									= self._get_uniform_logits(lowerCamelCase__				,      lowerCamelCase__						)
						UpperCAmelCase__:	Tuple									= scores.copy()

						# instantiate all dist processors
						UpperCAmelCase__:	Any									= FlaxTemperatureLogitsWarper(temperature=0.5						)
						UpperCAmelCase__:	List[str]									= FlaxTopKLogitsWarper(3						)
						UpperCAmelCase__:	Dict									= FlaxTopPLogitsWarper(0.8						)

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

						UpperCAmelCase__:	Optional[int]									= 1_0

						# no processor list
						def run_no_processor_list(lowerCamelCase__				,      lowerCamelCase__				,      lowerCamelCase__						):
							UpperCAmelCase__:	str									= temp_dist_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							UpperCAmelCase__:	Tuple									= top_k_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							UpperCAmelCase__:	Union[str, Any]									= top_p_warp(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							UpperCAmelCase__:	Dict									= min_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							UpperCAmelCase__:	Dict									= bos_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							UpperCAmelCase__:	Optional[Any]									= eos_dist_proc(lowerCamelCase__				,      lowerCamelCase__				,      cur_len=lowerCamelCase__						)
							return scores

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

						UpperCAmelCase__:	Union[str, Any]									= jax.jit(lowerCamelCase__						)
						UpperCAmelCase__:	int									= jax.jit(lowerCamelCase__						)

						UpperCAmelCase__:	List[Any]									= jitted_run_no_processor_list(lowerCamelCase__				,      lowerCamelCase__				,      lowerCamelCase__						)
						UpperCAmelCase__:	List[str]									= 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()						)


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_lowerCAmelCase  :      Union[str, Any]			 ={
    """configuration_clap""": [
        """CLAP_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """ClapAudioConfig""",
        """ClapConfig""",
        """ClapTextConfig""",
    ],
    """processing_clap""": ["""ClapProcessor"""],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					_lowerCAmelCase  :      Union[str, Any]			 =[
					    """CLAP_PRETRAINED_MODEL_ARCHIVE_LIST""",
					    """ClapModel""",
					    """ClapPreTrainedModel""",
					    """ClapTextModel""",
					    """ClapTextModelWithProjection""",
					    """ClapAudioModel""",
					    """ClapAudioModelWithProjection""",
					]
					_lowerCAmelCase  :      List[str]			 =["""ClapFeatureExtractor"""]

if TYPE_CHECKING:
					from .configuration_clap import (
					    CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
					    ClapAudioConfig,
					    ClapConfig,
					    ClapTextConfig,
					)
					from .processing_clap import ClapProcessor

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .feature_extraction_clap import ClapFeatureExtractor
										from .modeling_clap import (
										    CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
										    ClapAudioModel,
										    ClapAudioModelWithProjection,
										    ClapModel,
										    ClapPreTrainedModel,
										    ClapTextModel,
										    ClapTextModelWithProjection,
										)


else:
					import sys

					_lowerCAmelCase  :      Optional[int]			 =_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)