






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       _a   ( self  ,      _lowerCamelCase  ,      _lowerCamelCase							):
        UpperCamelCase_: Optional[Any]   				=    jnp.ones((batch_size, length)							) / length
        return scores





       def       _a   ( self							):
        UpperCamelCase_: int   				=    None
        UpperCamelCase_: Optional[int]   				=    2_0

        UpperCamelCase_: Any   				=    self._get_uniform_logits(batch_size=2  ,      length=__lowerCAmelCase							)

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

        # compute softmax
        UpperCamelCase_: List[str]   				=    jax.nn.softmax(__lowerCAmelCase  ,      axis=-1							)

        UpperCamelCase_: Tuple   				=    FlaxTemperatureLogitsWarper(temperature=0.5							)
        UpperCamelCase_: List[Any]   				=    FlaxTemperatureLogitsWarper(temperature=1.3							)

        UpperCamelCase_: Optional[int]   				=    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       _a   ( self							):
        UpperCamelCase_: List[Any]   				=    None
        UpperCamelCase_: str   				=    1_0
        UpperCamelCase_: str   				=    2

        # create ramp distribution
        UpperCamelCase_: Optional[int]   				=    np.broadcast_to(np.arange(__lowerCAmelCase							)[None, :]  ,      (batch_size, vocab_size)							).copy()
        UpperCamelCase_: Tuple   				=    ramp_logits[1:, : vocab_size // 2] + vocab_size

        UpperCamelCase_: Tuple   				=    FlaxTopKLogitsWarper(3							)

        UpperCamelCase_: Optional[int]   				=    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_: Tuple   				=    5
        UpperCamelCase_: Optional[int]   				=    FlaxTopKLogitsWarper(top_k=1  ,      filter_value=0.0  ,      min_tokens_to_keep=3							)

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

        # create distribution and take log (inverse to Softmax as taken in TopPLogitsWarper)
        UpperCamelCase_: Tuple   				=    np.log(np.array([[0.3, 0.1, 0.1, 0.5], [0.1_5, 0.3, 0.3, 0.2_5]]							)							)

        UpperCamelCase_: int   				=    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_: Tuple   				=    np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.2_5]]							)
        self.assertTrue(np.allclose(__lowerCAmelCase  ,      __lowerCAmelCase  ,      atol=1e-3							)							)

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

        # make ramp_logits more extreme
        UpperCamelCase_: int   				=    ramp_logits[1] * 1_0_0.0

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

        UpperCamelCase_: List[str]   				=    FlaxMinLengthLogitsProcessor(min_length=1_0  ,      eos_token_id=__lowerCAmelCase							)

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





       def       _a   ( self							):
        UpperCamelCase_: Tuple   				=    2_0
        UpperCamelCase_: int   				=    4
        UpperCamelCase_: Optional[Any]   				=    0

        UpperCamelCase_: List[Any]   				=    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCAmelCase							)

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





       def       _a   ( self							):
        UpperCamelCase_: str   				=    2_0
        UpperCamelCase_: List[Any]   				=    4
        UpperCamelCase_: int   				=    0
        UpperCamelCase_: List[Any]   				=    5

        UpperCamelCase_: List[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_: str   				=    ids_tensor((batch_size, 4)  ,      vocab_size=2_0							)
        UpperCamelCase_: Any   				=    4
        UpperCamelCase_: List[str]   				=    self._get_uniform_logits(__lowerCAmelCase  ,      __lowerCAmelCase							)
        UpperCamelCase_: Optional[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_: Optional[Any]   				=    3
        UpperCamelCase_: Optional[int]   				=    self._get_uniform_logits(__lowerCAmelCase  ,      __lowerCAmelCase							)
        UpperCamelCase_: str   				=    logits_processor(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        self.assertFalse(jnp.isinf(__lowerCAmelCase							).any()							)





       def       _a   ( self							):
        UpperCamelCase_: List[Any]   				=    4
        UpperCamelCase_: int   				=    1_0
        UpperCamelCase_: Any   				=    1_5
        UpperCamelCase_: Optional[int]   				=    2
        UpperCamelCase_: List[Any]   				=    1
        UpperCamelCase_: str   				=    1_5

        # dummy input_ids and scores
        UpperCamelCase_: str   				=    ids_tensor((batch_size, sequence_length)  ,      __lowerCAmelCase							)
        UpperCamelCase_: str   				=    input_ids.copy()

        UpperCamelCase_: List[Any]   				=    self._get_uniform_logits(__lowerCAmelCase  ,      __lowerCAmelCase							)
        UpperCamelCase_: Tuple   				=    scores.copy()

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

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

        UpperCamelCase_: List[Any]   				=    1_0

        # no processor list
        UpperCamelCase_: Tuple   				=    temp_dist_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        UpperCamelCase_: List[str]   				=    top_k_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        UpperCamelCase_: Optional[Any]   				=    top_p_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        UpperCamelCase_: List[str]   				=    min_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        UpperCamelCase_: Dict   				=    bos_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
        UpperCamelCase_: Any   				=    eos_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)

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

        # dummy input_ids and scores
        UpperCamelCase_: List[Any]   				=    ids_tensor((batch_size, sequence_length)  ,      __lowerCAmelCase							)
        UpperCamelCase_: Tuple   				=    input_ids.copy()

        UpperCamelCase_: List[Any]   				=    self._get_uniform_logits(__lowerCAmelCase  ,      __lowerCAmelCase							)
        UpperCamelCase_: str   				=    scores.copy()

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

        # instantiate all logits processors
        UpperCamelCase_: str   				=    FlaxMinLengthLogitsProcessor(min_length=1_0  ,      eos_token_id=__lowerCAmelCase							)
        UpperCamelCase_: str   				=    FlaxForcedBOSTokenLogitsProcessor(bos_token_id=__lowerCAmelCase							)
        UpperCamelCase_: List[str]   				=    FlaxForcedEOSTokenLogitsProcessor(max_length=__lowerCAmelCase  ,      eos_token_id=__lowerCAmelCase							)

        UpperCamelCase_: Any   				=    1_0

        # no processor list
        def run_no_processor_list(_lowerCamelCase  ,      _lowerCamelCase  ,      _lowerCamelCase							):
         UpperCamelCase_: str   				=    temp_dist_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         UpperCamelCase_: List[Any]   				=    top_k_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         UpperCamelCase_: int   				=    top_p_warp(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         UpperCamelCase_: Union[str, Any]   				=    min_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         UpperCamelCase_: str   				=    bos_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         UpperCamelCase_: Any   				=    eos_dist_proc(__lowerCAmelCase  ,      __lowerCAmelCase  ,      cur_len=__lowerCAmelCase							)
         return scores

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

        UpperCamelCase_: Dict   				=    jax.jit(__lowerCAmelCase							)
        UpperCamelCase_: Union[str, Any]   				=    jax.jit(__lowerCAmelCase							)

        UpperCamelCase_: Optional[Any]   				=    jitted_run_no_processor_list(__lowerCAmelCase  ,      __lowerCAmelCase  ,      __lowerCAmelCase							)
        UpperCamelCase_: List[Any]   				=    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 abc import ABC, abstractmethod
from argparse import ArgumentParser




class     _lowerCAmelCase(	UpperCAmelCase_ ):





       """simple docstring"""







       @staticmethod
       @abstractmethod
       def       _a   ( _lowerCamelCase							):
         raise NotImplementedError()





       @abstractmethod
       def       _a   ( self							):
         raise NotImplementedError()