

"""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  lowerCamelCase    (unittest.TestCase      ):

			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]     ,					__UpperCAmelCase					:    int     ,					__UpperCAmelCase					:    int							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   jnp.ones((batch_size, length)							) / length
						return scores

			def        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   None
						SCREAMING_SNAKE_CASE__										=   2_0

						SCREAMING_SNAKE_CASE__										=   self._get_uniform_logits(batch_size=2     ,					length=__UpperCAmelCase							)

						# 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(__UpperCAmelCase     ,					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(__UpperCAmelCase     ,					scores.copy()     ,					cur_len=__UpperCAmelCase							)     ,					axis=-1							)
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    str							) ->     Union[str, Any]:
						SCREAMING_SNAKE_CASE__										=   None
						SCREAMING_SNAKE_CASE__										=   1_0
						SCREAMING_SNAKE_CASE__										=   2

						# create ramp distribution
						SCREAMING_SNAKE_CASE__										=   np.broadcast_to(np.arange(__UpperCAmelCase							)[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(__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
						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(__UpperCAmelCase							)[None, :]     ,					(batch_size, length)							).copy()
						SCREAMING_SNAKE_CASE__										=   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        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Union[str, Any]:
						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(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)							)

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

						# check edge cases with negative and extreme logits
						SCREAMING_SNAKE_CASE__										=   np.broadcast_to(np.arange(__UpperCAmelCase							)[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(__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        SCREAMING_SNAKE_CASE					(       self					:    Union[str, Any]							) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   2_0
						SCREAMING_SNAKE_CASE__										=   4
						SCREAMING_SNAKE_CASE__										=   0

						SCREAMING_SNAKE_CASE__										=   FlaxMinLengthLogitsProcessor(min_length=1_0     ,					eos_token_id=__UpperCAmelCase							)

						# 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(__UpperCAmelCase     ,					__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   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
						SCREAMING_SNAKE_CASE__										=   self._get_uniform_logits(__UpperCAmelCase     ,					__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   1_5
						SCREAMING_SNAKE_CASE__										=   min_dist_processor(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						self.assertFalse(jnp.isinf(__UpperCAmelCase							).any()							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     Dict:
						SCREAMING_SNAKE_CASE__										=   2_0
						SCREAMING_SNAKE_CASE__										=   4
						SCREAMING_SNAKE_CASE__										=   0

						SCREAMING_SNAKE_CASE__										=   FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__UpperCAmelCase							)

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

			def        SCREAMING_SNAKE_CASE					(       self					:    int							) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   2_0
						SCREAMING_SNAKE_CASE__										=   4
						SCREAMING_SNAKE_CASE__										=   0
						SCREAMING_SNAKE_CASE__										=   5

						SCREAMING_SNAKE_CASE__										=   FlaxForcedEOSTokenLogitsProcessor(max_length=__UpperCAmelCase     ,					eos_token_id=__UpperCAmelCase							)

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

			def        SCREAMING_SNAKE_CASE					(       self					:    Optional[Any]							) ->     Dict:
						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)     ,					__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   input_ids.copy()

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

						SCREAMING_SNAKE_CASE__										=   1_0

						# no processor list
						SCREAMING_SNAKE_CASE__										=   temp_dist_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   top_k_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   top_p_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   min_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   bos_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   eos_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)

						# 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(__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        SCREAMING_SNAKE_CASE					(       self					:    Optional[int]							) ->     str:
						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)     ,					__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   input_ids.copy()

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

						SCREAMING_SNAKE_CASE__										=   1_0

						# no processor list
						def run_no_processor_list(__UpperCAmelCase					:    str     ,					__UpperCAmelCase					:    Tuple     ,					__UpperCAmelCase					:    Any							):
									SCREAMING_SNAKE_CASE__										=   temp_dist_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   top_k_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   top_p_warp(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   min_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   bos_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									SCREAMING_SNAKE_CASE__										=   eos_dist_proc(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									return scores

						# with processor list
						def run_processor_list(__UpperCAmelCase					:    Dict     ,					__UpperCAmelCase					:    Optional[int]     ,					__UpperCAmelCase					:    List[Any]							):
									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(__UpperCAmelCase     ,					__UpperCAmelCase     ,					cur_len=__UpperCAmelCase							)
									return scores

						SCREAMING_SNAKE_CASE__										=   jax.jit(__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   jax.jit(__UpperCAmelCase							)

						SCREAMING_SNAKE_CASE__										=   jitted_run_no_processor_list(__UpperCAmelCase     ,					__UpperCAmelCase     ,					__UpperCAmelCase							)
						SCREAMING_SNAKE_CASE__										=   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"""




import logging
from dataclasses import dataclass, field
from typing import Optional

from seqaseq_trainer import arg_to_scheduler

from transformers import TrainingArguments


A_				:      Optional[Any]           =			logging.getLogger(__name__)





@dataclass
class  lowerCamelCase    (A__      ):
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=0.0      ,metadata={'help': 'The label smoothing epsilon to apply (if not zero).'}      )
			lowerCamelCase__ :		bool       		=				field(default=A__      ,metadata={'help': 'Whether to SortishSamler or not.'}      )
			lowerCamelCase__ :		bool       		=				field(
			    default=A__      ,metadata={'help': 'Whether to use generate to calculate generative metrics (ROUGE, BLEU).'}      )
			lowerCamelCase__ :		bool       		=				field(default=A__      ,metadata={'help': 'whether to use adafactor'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Encoder layer dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Decoder layer dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(default=A__      ,metadata={'help': 'Dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[float]       		=				field(
			    default=A__      ,metadata={'help': 'Attention dropout probability. Goes into model.config.'}      )
			lowerCamelCase__ :		Optional[str]       		=				field(
			    default='linear'      ,metadata={'help': f"Which lr scheduler to use. Selected in {sorted(arg_to_scheduler.keys()      )}"}      ,)


