


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 	snake_case_  (	self  ,       __A  ,       __A    ):
        __a					=  jnp.ones((batch_size, length)    ) / length
        return scores


     def 	snake_case_  (	self    ):
        __a					=  None
        __a					=  20

        __a					=  self._get_uniform_logits(batch_size=2  ,       length=__A    )

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

        # compute softmax
        __a					=  jax.nn.softmax(__A  ,       axis=-1    )

        __a					=  FlaxTemperatureLogitsWarper(temperature=0.5    )
        __a					=  FlaxTemperatureLogitsWarper(temperature=1.3    )

        __a					=  jax.nn.softmax(temp_dist_warper_sharper(__A  ,       scores.copy()  ,       cur_len=__A    )  ,       axis=-1    )
        __a					=  jax.nn.softmax(temp_dist_warper_smoother(__A  ,       scores.copy()  ,       cur_len=__A    )  ,       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    ):
        __a					=  None
        __a					=  10
        __a					=  2

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

        __a					=  FlaxTopKLogitsWarper(3    )

        __a					=  top_k_warp(__A  ,       __A  ,       cur_len=__A    )

        # 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
        __a					=  5
        __a					=  FlaxTopKLogitsWarper(top_k=1  ,       filter_value=0.0  ,       min_tokens_to_keep=3    )

        __a					=  np.broadcast_to(np.arange(__A    )[None, :]  ,       (batch_size, length)    ).copy()
        __a					=  top_k_warp_safety_check(__A  ,       __A  ,       cur_len=__A    )

        # 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    ):
        __a					=  None
        __a					=  10
        __a					=  2

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

        __a					=  FlaxTopPLogitsWarper(0.8    )
        __a					=  np.exp(top_p_warp(__A  ,       __A  ,       cur_len=__A    )    )

        # dist should be filtered to keep min num values so that sum is >= top_p
        # exp (-inf) => 0
        __a					=  np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]    )
        self.assertTrue(np.allclose(__A  ,       __A  ,       atol=1E-3    )    )

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

        # make ramp_logits more extreme
        __a					=  ramp_logits[1] * 100.0

        # make sure at least 2 tokens are kept
        __a					=  FlaxTopPLogitsWarper(0.9  ,       min_tokens_to_keep=2  ,       filter_value=0.0    )
        __a					=  top_p_warp(__A  ,       __A  ,       cur_len=__A    )

        # 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    ):
        __a					=  20
        __a					=  4
        __a					=  0

        __a					=  FlaxMinLengthLogitsProcessor(min_length=10  ,       eos_token_id=__A    )

        # check that min length is applied at length 5
        __a					=  ids_tensor((batch_size, 20)  ,       vocab_size=20    )
        __a					=  5
        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  min_dist_processor(__A  ,       __A  ,       cur_len=__A    )
        self.assertListEqual(scores_before_min_length[:, eos_token_id].tolist()  ,       4 * [-float("""inf"""    )]    )

        # check that min length is not applied anymore at length 15
        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  15
        __a					=  min_dist_processor(__A  ,       __A  ,       cur_len=__A    )
        self.assertFalse(jnp.isinf(__A    ).any()    )


     def 	snake_case_  (	self    ):
        __a					=  20
        __a					=  4
        __a					=  0

        __a					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A    )

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


     def 	snake_case_  (	self    ):
        __a					=  20
        __a					=  4
        __a					=  0
        __a					=  5

        __a					=  FlaxForcedEOSTokenLogitsProcessor(max_length=__A  ,       eos_token_id=__A    )

        # check that all scores are -inf except the eos_token_id when max_length is reached
        __a					=  ids_tensor((batch_size, 4)  ,       vocab_size=20    )
        __a					=  4
        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  logits_processor(__A  ,       __A  ,       cur_len=__A    )
        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
        __a					=  3
        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  logits_processor(__A  ,       __A  ,       cur_len=__A    )
        self.assertFalse(jnp.isinf(__A    ).any()    )


     def 	snake_case_  (	self    ):
        __a					=  4
        __a					=  10
        __a					=  15
        __a					=  2
        __a					=  1
        __a					=  15

        # dummy input_ids and scores
        __a					=  ids_tensor((batch_size, sequence_length)  ,       __A    )
        __a					=  input_ids.copy()

        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  scores.copy()

        # instantiate all dist processors
        __a					=  FlaxTemperatureLogitsWarper(temperature=0.5    )
        __a					=  FlaxTopKLogitsWarper(3    )
        __a					=  FlaxTopPLogitsWarper(0.8    )

        # instantiate all logits processors
        __a					=  FlaxMinLengthLogitsProcessor(min_length=10  ,       eos_token_id=__A    )
        __a					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A    )
        __a					=  FlaxForcedEOSTokenLogitsProcessor(max_length=__A  ,       eos_token_id=__A    )

        __a					=  10

        # no processor list
        __a					=  temp_dist_warp(__A  ,       __A  ,       cur_len=__A    )
        __a					=  top_k_warp(__A  ,       __A  ,       cur_len=__A    )
        __a					=  top_p_warp(__A  ,       __A  ,       cur_len=__A    )
        __a					=  min_dist_proc(__A  ,       __A  ,       cur_len=__A    )
        __a					=  bos_dist_proc(__A  ,       __A  ,       cur_len=__A    )
        __a					=  eos_dist_proc(__A  ,       __A  ,       cur_len=__A    )

        # with processor list
        __a					=  FlaxLogitsProcessorList(
            [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
        __a					=  processor(__A  ,       __A  ,       cur_len=__A    )

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

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


     def 	snake_case_  (	self    ):
        __a					=  4
        __a					=  10
        __a					=  15
        __a					=  2
        __a					=  1
        __a					=  15

        # dummy input_ids and scores
        __a					=  ids_tensor((batch_size, sequence_length)  ,       __A    )
        __a					=  input_ids.copy()

        __a					=  self._get_uniform_logits(__A  ,       __A    )
        __a					=  scores.copy()

        # instantiate all dist processors
        __a					=  FlaxTemperatureLogitsWarper(temperature=0.5    )
        __a					=  FlaxTopKLogitsWarper(3    )
        __a					=  FlaxTopPLogitsWarper(0.8    )

        # instantiate all logits processors
        __a					=  FlaxMinLengthLogitsProcessor(min_length=10  ,       eos_token_id=__A    )
        __a					=  FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__A    )
        __a					=  FlaxForcedEOSTokenLogitsProcessor(max_length=__A  ,       eos_token_id=__A    )

        __a					=  10

        # no processor list
        def run_no_processor_list(__A  ,       __A  ,       __A    ):
           __a					=  temp_dist_warp(__A  ,       __A  ,       cur_len=__A    )
           __a					=  top_k_warp(__A  ,       __A  ,       cur_len=__A    )
           __a					=  top_p_warp(__A  ,       __A  ,       cur_len=__A    )
           __a					=  min_dist_proc(__A  ,       __A  ,       cur_len=__A    )
           __a					=  bos_dist_proc(__A  ,       __A  ,       cur_len=__A    )
           __a					=  eos_dist_proc(__A  ,       __A  ,       cur_len=__A    )
           return scores

        # with processor list
        def run_processor_list(__A  ,       __A  ,       __A    ):
           __a					=  FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]    )
           __a					=  processor(__A  ,       __A  ,       cur_len=__A    )
           return scores

        __a					=  jax.jit(__A    )
        __a					=  jax.jit(__A    )

        __a					=  jitted_run_no_processor_list(__A  ,       __A  ,       __A    )
        __a					=  jitted_run_processor_list(__A  ,       __A  ,       __A    )

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

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










import tensorflow as tf

from ...tf_utils import shape_list
class     UpperCamelCase__					(     tf.keras.layers.Layer				):






		'''simple docstring'''






		def __init__(							self,   snake_case__,   snake_case__,   snake_case__,   snake_case__,   snake_case__=1,   snake_case__=False,   **snake_case__  )       -> Optional[int]:






									"""simple docstring"""
									super().__init__(**snake_case__  )

									lowercase_	:   Tuple  				=     vocab_size
									lowercase_	:   Union[str, Any]  				=     d_embed
									lowercase_	:   Optional[Any]  				=     d_proj

									lowercase_	:   Optional[Any]  				=     cutoffs + [vocab_size]
									lowercase_	:   Optional[int]  				=     [0] + self.cutoffs
									lowercase_	:   Union[str, Any]  				=     div_val

									lowercase_	:   Union[str, Any]  				=     self.cutoffs[0]
									lowercase_	:   List[str]  				=     len(self.cutoffs  ) - 1
									lowercase_	:   List[str]  				=     self.shortlist_size + self.n_clusters
									lowercase_	:   Tuple  				=     keep_order

									lowercase_	:   Dict  				=     []
									lowercase_	:   int  				=     []






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






									"""simple docstring"""
									if self.n_clusters > 0:
																lowercase_	:   List[str]  				=     self.add_weight(
																    shape=(self.n_clusters, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_weight"""  )
																lowercase_	:   List[Any]  				=     self.add_weight(
																    shape=(self.n_clusters,),   initializer="""zeros""",   trainable=snake_case__,   name="""cluster_bias"""  )

									if self.div_val == 1:
																for i in range(len(self.cutoffs  )  ):
																							if self.d_proj != self.d_embed:
																														lowercase_	:   Any  				=     self.add_weight(
																														    shape=(self.d_embed, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}""",   )
																														self.out_projs.append(snake_case__  )
																							else:
																														self.out_projs.append(snake_case__  )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size, self.d_embed),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   List[str]  				=     self.add_weight(
																							    shape=(self.vocab_size,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									else:
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[Any]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							lowercase_	:   Union[str, Any]  				=     self.d_embed // (self.div_val**i)

																							lowercase_	:   int  				=     self.add_weight(
																							    shape=(d_emb_i, self.d_proj),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_projs_._{i}"""  )
																							self.out_projs.append(snake_case__  )
																							lowercase_	:   Any  				=     self.add_weight(
																							    shape=(r_idx - l_idx, d_emb_i),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._weight""",   )
																							lowercase_	:   Tuple  				=     self.add_weight(
																							    shape=(r_idx - l_idx,),   initializer="""zeros""",   trainable=snake_case__,   name=f"""out_layers_._{i}_._bias""",   )
																							self.out_layers.append((weight, bias)  )
									super().build(snake_case__  )






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__,   snake_case__,   snake_case__=None  )       -> Optional[Any]:






									"""simple docstring"""
									lowercase_	:   Dict  				=     x
									if proj is not None:
																lowercase_	:   List[Any]  				=     tf.einsum("""ibd,ed->ibe""",   snake_case__,   snake_case__  )
									return tf.einsum("""ibd,nd->ibn""",   snake_case__,   snake_case__  ) + b






		@staticmethod
		def 						snake_case__				(							snake_case__,   snake_case__  )       -> List[str]:






									"""simple docstring"""
									lowercase_	:   Optional[int]  				=     shape_list(snake_case__  )
									lowercase_	:   Optional[Any]  				=     tf.range(lp_size[0],   dtype=target.dtype  )
									lowercase_	:   str  				=     tf.stack([r, target],   1  )
									return tf.gather_nd(snake_case__,   snake_case__  )






		def 						snake_case__				(							self,   snake_case__,   snake_case__,   snake_case__=True,   snake_case__=False  )       -> Dict:






									"""simple docstring"""
									lowercase_	:   Any  				=     0
									if self.n_clusters == 0:
																lowercase_	:   Union[str, Any]  				=     self._logit(snake_case__,   self.out_layers[0][0],   self.out_layers[0][1],   self.out_projs[0]  )
																if target is not None:
																							lowercase_	:   str  				=     tf.nn.sparse_softmax_cross_entropy_with_logits(labels=snake_case__,   logits=snake_case__  )
																lowercase_	:   Optional[int]  				=     tf.nn.log_softmax(snake_case__,   axis=-1  )
									else:
																lowercase_	:   int  				=     shape_list(snake_case__  )
																lowercase_	:   List[Any]  				=     []
																lowercase_	:   Tuple  				=     tf.zeros(hidden_sizes[:2]  )
																for i in range(len(self.cutoffs  )  ):
																							lowercase_				,		lowercase_	:   Optional[int]  				=     self.cutoff_ends[i], self.cutoff_ends[i + 1]
																							if target is not None:
																														lowercase_	:   Tuple  				=     (target >= l_idx) & (target < r_idx)
																														lowercase_	:   Union[str, Any]  				=     tf.where(snake_case__  )
																														lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  ) - l_idx

																							if self.div_val == 1:
																														lowercase_	:   Optional[Any]  				=     self.out_layers[0][0][l_idx:r_idx]
																														lowercase_	:   Union[str, Any]  				=     self.out_layers[0][1][l_idx:r_idx]
																							else:
																														lowercase_	:   Dict  				=     self.out_layers[i][0]
																														lowercase_	:   int  				=     self.out_layers[i][1]

																							if i == 0:
																														lowercase_	:   Optional[int]  				=     tf.concat([cur_W, self.cluster_weight],   0  )
																														lowercase_	:   Optional[Any]  				=     tf.concat([cur_b, self.cluster_bias],   0  )

																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[0]  )
																														lowercase_	:   List[str]  				=     tf.nn.log_softmax(snake_case__  )
																														out.append(head_logprob[..., : self.cutoffs[0]]  )
																														if target is not None:
																																					lowercase_	:   int  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																							else:
																														lowercase_	:   List[str]  				=     self._logit(snake_case__,   snake_case__,   snake_case__,   self.out_projs[i]  )
																														lowercase_	:   Dict  				=     tf.nn.log_softmax(snake_case__  )
																														lowercase_	:   Optional[int]  				=     self.cutoffs[0] + i - 1  # No probability for the head cluster
																														lowercase_	:   int  				=     head_logprob[..., cluster_prob_idx, None] + tail_logprob
																														out.append(snake_case__  )
																														if target is not None:
																																					lowercase_	:   Optional[Any]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Optional[int]  				=     tf.boolean_mask(snake_case__,   snake_case__  )
																																					lowercase_	:   Union[str, Any]  				=     self._gather_logprob(snake_case__,   snake_case__  )
																																					cur_logprob += cur_head_logprob[:, self.cutoff_ends[1] + i - 1]
																							if target is not None:
																														loss += tf.scatter_nd(snake_case__,   -cur_logprob,   shape_list(snake_case__  )  )
																lowercase_	:   List[str]  				=     tf.concat(snake_case__,   axis=-1  )

									if target is not None:
																if return_mean:
																							lowercase_	:   Tuple  				=     tf.reduce_mean(snake_case__  )
																# Add the training-time loss value to the layer using `self.add_loss()`.
																self.add_loss(snake_case__  )

																# Log the loss as a metric (we could log arbitrary metrics,
																# including different metrics for training and inference.
																self.add_metric(snake_case__,   name=self.name,   aggregation="""mean""" if return_mean else """"""  )

									return out