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







	def  snake_case__     (     self      :Optional[int]   ,     lowercase      :int   ,     lowercase      :int			)		->      Optional[int]:







				"""simple docstring"""






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







	def  snake_case__     (     self      :Any			)		->      Tuple:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= 2_0

				SCREAMING_SNAKE_CASE     						= self._get_uniform_logits(batch_size=2   ,     length=lowercase			)

				# 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, 1_0].set((1 / length) - 0.4			)  # valley, 1st batch

				# compute softmax
				SCREAMING_SNAKE_CASE     						= jax.nn.softmax(lowercase   ,     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(lowercase   ,     scores.copy()   ,     cur_len=lowercase			)   ,     axis=-1			)
				SCREAMING_SNAKE_CASE     						= jax.nn.softmax(temp_dist_warper_smoother(lowercase   ,     scores.copy()   ,     cur_len=lowercase			)   ,     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  snake_case__     (     self      :Dict			)		->      str:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= 1_0
				SCREAMING_SNAKE_CASE     						= 2

				# create ramp distribution
				SCREAMING_SNAKE_CASE     						= np.broadcast_to(np.arange(lowercase			)[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(lowercase   ,     lowercase   ,     cur_len=lowercase			)

				# 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(lowercase			)[None, :]   ,     (batch_size, length)			).copy()
				SCREAMING_SNAKE_CASE     						= top_k_warp_safety_check(lowercase   ,     lowercase   ,     cur_len=lowercase			)

				# min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
				self.assertListEqual((scores == 0.0).sum(axis=-1			).tolist()   ,     [2, 2]			)







	def  snake_case__     (     self      :str			)		->      Optional[Any]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= None
				SCREAMING_SNAKE_CASE     						= 1_0
				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(lowercase   ,     lowercase   ,     cur_len=lowercase			)			)

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

				# check edge cases with negative and extreme logits
				SCREAMING_SNAKE_CASE     						= np.broadcast_to(np.arange(lowercase			)[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(lowercase   ,     lowercase   ,     cur_len=lowercase			)

				# 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  snake_case__     (     self      :Any			)		->      Optional[int]:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= 2_0
				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= 0

				SCREAMING_SNAKE_CASE     						= FlaxMinLengthLogitsProcessor(min_length=1_0   ,     eos_token_id=lowercase			)

				# check that min length is applied at length 5
				SCREAMING_SNAKE_CASE     						= ids_tensor((batch_size, 2_0)   ,     vocab_size=2_0			)
				SCREAMING_SNAKE_CASE     						= 5
				SCREAMING_SNAKE_CASE     						= self._get_uniform_logits(lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= min_dist_processor(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				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(lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= 1_5
				SCREAMING_SNAKE_CASE     						= min_dist_processor(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				self.assertFalse(jnp.isinf(lowercase			).any()			)







	def  snake_case__     (     self      :Union[str, Any]			)		->      Any:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= 2_0
				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= 0

				SCREAMING_SNAKE_CASE     						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase			)

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







	def  snake_case__     (     self      :Any			)		->      str:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= 2_0
				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= 0
				SCREAMING_SNAKE_CASE     						= 5

				SCREAMING_SNAKE_CASE     						= FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase   ,     eos_token_id=lowercase			)

				# 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=2_0			)
				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= self._get_uniform_logits(lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= logits_processor(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				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(lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= logits_processor(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				self.assertFalse(jnp.isinf(lowercase			).any()			)







	def  snake_case__     (     self      :int			)		->      Dict:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= 1_0
				SCREAMING_SNAKE_CASE     						= 1_5
				SCREAMING_SNAKE_CASE     						= 2
				SCREAMING_SNAKE_CASE     						= 1
				SCREAMING_SNAKE_CASE     						= 1_5

				# dummy input_ids and scores
				SCREAMING_SNAKE_CASE     						= ids_tensor((batch_size, sequence_length)   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= input_ids.copy()

				SCREAMING_SNAKE_CASE     						= self._get_uniform_logits(lowercase   ,     lowercase			)
				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=1_0   ,     eos_token_id=lowercase			)
				SCREAMING_SNAKE_CASE     						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase			)
				SCREAMING_SNAKE_CASE     						= FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase   ,     eos_token_id=lowercase			)

				SCREAMING_SNAKE_CASE     						= 1_0

				# no processor list
				SCREAMING_SNAKE_CASE     						= temp_dist_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				SCREAMING_SNAKE_CASE     						= top_k_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				SCREAMING_SNAKE_CASE     						= top_p_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				SCREAMING_SNAKE_CASE     						= min_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				SCREAMING_SNAKE_CASE     						= bos_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)
				SCREAMING_SNAKE_CASE     						= eos_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)

				# 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(lowercase   ,     lowercase   ,     cur_len=lowercase			)

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

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







	def  snake_case__     (     self      :Union[str, Any]			)		->      Any:







				"""simple docstring"""






				SCREAMING_SNAKE_CASE     						= 4
				SCREAMING_SNAKE_CASE     						= 1_0
				SCREAMING_SNAKE_CASE     						= 1_5
				SCREAMING_SNAKE_CASE     						= 2
				SCREAMING_SNAKE_CASE     						= 1
				SCREAMING_SNAKE_CASE     						= 1_5

				# dummy input_ids and scores
				SCREAMING_SNAKE_CASE     						= ids_tensor((batch_size, sequence_length)   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= input_ids.copy()

				SCREAMING_SNAKE_CASE     						= self._get_uniform_logits(lowercase   ,     lowercase			)
				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=1_0   ,     eos_token_id=lowercase			)
				SCREAMING_SNAKE_CASE     						= FlaxForcedBOSTokenLogitsProcessor(bos_token_id=lowercase			)
				SCREAMING_SNAKE_CASE     						= FlaxForcedEOSTokenLogitsProcessor(max_length=lowercase   ,     eos_token_id=lowercase			)

				SCREAMING_SNAKE_CASE     						= 1_0

				# no processor list
				def run_no_processor_list(lowercase      :int   ,     lowercase      :Tuple   ,     lowercase      :str			):
							SCREAMING_SNAKE_CASE     						= temp_dist_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							SCREAMING_SNAKE_CASE     						= top_k_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							SCREAMING_SNAKE_CASE     						= top_p_warp(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							SCREAMING_SNAKE_CASE     						= min_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							SCREAMING_SNAKE_CASE     						= bos_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							SCREAMING_SNAKE_CASE     						= eos_dist_proc(lowercase   ,     lowercase   ,     cur_len=lowercase			)
							return scores

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

				SCREAMING_SNAKE_CASE     						= jax.jit(lowercase			)
				SCREAMING_SNAKE_CASE     						= jax.jit(lowercase			)

				SCREAMING_SNAKE_CASE     						= jitted_run_no_processor_list(lowercase   ,     lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= jitted_run_processor_list(lowercase   ,     lowercase   ,     lowercase			)

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

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

import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin





class     lowerCamelCase      (       __lowerCamelCase					):
	UpperCamelCase_  :					Optional[Any]    		=	'MCTCTFeatureExtractor'
	UpperCamelCase_  :					List[Any]    		=	'AutoTokenizer'







	def __init__(     self      :Tuple   ,     lowercase      :List[str]   ,     lowercase      :Dict			)		->      Tuple:







				"""simple docstring"""






				super().__init__(lowercase   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= self.feature_extractor
				SCREAMING_SNAKE_CASE     						= False







	def __call__(     self      :Union[str, Any]   ,     *lowercase      :Union[str, Any]   ,     **lowercase      :str			)		->      int:







				"""simple docstring"""






				if self._in_target_context_manager:
							return self.current_processor(*lowercase   ,     **lowercase			)

				if "raw_speech" in kwargs:
							warnings.warn('''Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'''			)
							SCREAMING_SNAKE_CASE     						= kwargs.pop('''raw_speech'''			)
				else:
							SCREAMING_SNAKE_CASE     						= kwargs.pop('''audio'''   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= kwargs.pop('''sampling_rate'''   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= kwargs.pop('''text'''   ,     lowercase			)
				if len(lowercase			) > 0:
							SCREAMING_SNAKE_CASE     						= args[0]
							SCREAMING_SNAKE_CASE     						= args[1:]

				if audio is None and text is None:
							raise ValueError('''You need to specify either an `audio` or `text` input to process.'''			)

				if audio is not None:
							SCREAMING_SNAKE_CASE     						= self.feature_extractor(lowercase   ,     *lowercase   ,     sampling_rate=lowercase   ,     **lowercase			)
				if text is not None:
							SCREAMING_SNAKE_CASE     						= self.tokenizer(lowercase   ,     **lowercase			)

				if text is None:
							return inputs
				elif audio is None:
							return encodings
				else:
							SCREAMING_SNAKE_CASE     						= encodings['''input_ids''']
							return inputs







	def  snake_case__     (     self      :Dict   ,     *lowercase      :Union[str, Any]   ,     **lowercase      :List[str]			)		->      Union[str, Any]:







				"""simple docstring"""






				return self.tokenizer.batch_decode(*lowercase   ,     **lowercase			)







	def  snake_case__     (     self      :List[Any]   ,     *lowercase      :List[Any]   ,     **lowercase      :List[str]			)		->      int:







				"""simple docstring"""






				if self._in_target_context_manager:
							return self.current_processor.pad(*lowercase   ,     **lowercase			)

				SCREAMING_SNAKE_CASE     						= kwargs.pop('''input_features'''   ,     lowercase			)
				SCREAMING_SNAKE_CASE     						= kwargs.pop('''labels'''   ,     lowercase			)
				if len(lowercase			) > 0:
							SCREAMING_SNAKE_CASE     						= args[0]
							SCREAMING_SNAKE_CASE     						= args[1:]

				if input_features is not None:
							SCREAMING_SNAKE_CASE     						= self.feature_extractor.pad(lowercase   ,     *lowercase   ,     **lowercase			)
				if labels is not None:
							SCREAMING_SNAKE_CASE     						= self.tokenizer.pad(lowercase   ,     **lowercase			)

				if labels is None:
							return input_features
				elif input_features is None:
							return labels
				else:
							SCREAMING_SNAKE_CASE     						= labels['''input_ids''']
							return input_features







	def  snake_case__     (     self      :Dict   ,     *lowercase      :List[str]   ,     **lowercase      :Union[str, Any]			)		->      List[Any]:







				"""simple docstring"""






				return self.tokenizer.decode(*lowercase   ,     **lowercase			)







	@contextmanager
	def  snake_case__     (     self      :str			)		->      Union[str, Any]:







				"""simple docstring"""






				warnings.warn(
				    '''`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '''
				    '''labels by using the argument `text` of the regular `__call__` method (either in the same call as '''
				    '''your audio inputs, or in a separate call.'''			)
				SCREAMING_SNAKE_CASE     						= True
				SCREAMING_SNAKE_CASE     						= self.tokenizer
				yield
				SCREAMING_SNAKE_CASE     						= self.feature_extractor
				SCREAMING_SNAKE_CASE     						= False