




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







   def lowerCamelCase						(      self		:						List[str]					,						_snake_case		:						int					,						_snake_case		:						int):

        """simple docstring"""


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







   def lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        UpperCAmelCase_          =							None
        UpperCAmelCase_          =							20

        UpperCAmelCase_          =							self._get_uniform_logits(batch_size=2					,						length=_snake_case)

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

        # compute softmax
        UpperCAmelCase_          =							jax.nn.softmax(_snake_case					,						axis=-1)

        UpperCAmelCase_          =							FlaxTemperatureLogitsWarper(temperature=0.5)
        UpperCAmelCase_          =							FlaxTemperatureLogitsWarper(temperature=1.3)

        UpperCAmelCase_          =							jax.nn.softmax(temp_dist_warper_sharper(_snake_case					,						scores.copy()					,						cur_len=_snake_case)					,						axis=-1)
        UpperCAmelCase_          =							jax.nn.softmax(temp_dist_warper_smoother(_snake_case					,						scores.copy()					,						cur_len=_snake_case)					,						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 lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							None
        UpperCAmelCase_          =							10
        UpperCAmelCase_          =							2

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

        UpperCAmelCase_          =							FlaxTopKLogitsWarper(3)

        UpperCAmelCase_          =							top_k_warp(_snake_case					,						_snake_case					,						cur_len=_snake_case)

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

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

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







   def lowerCamelCase						(      self		:						Dict):

        """simple docstring"""


        UpperCAmelCase_          =							None
        UpperCAmelCase_          =							10
        UpperCAmelCase_          =							2

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

        UpperCAmelCase_          =							FlaxTopPLogitsWarper(0.8)
        UpperCAmelCase_          =							np.exp(top_p_warp(_snake_case					,						_snake_case					,						cur_len=_snake_case))

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

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

        # make ramp_logits more extreme
        UpperCAmelCase_          =							ramp_logits[1] * 1_0_0.0

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

        # 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 lowerCamelCase						(      self		:						Union[str, Any]):

        """simple docstring"""


        UpperCAmelCase_          =							20
        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							0

        UpperCAmelCase_          =							FlaxMinLengthLogitsProcessor(min_length=10					,						eos_token_id=_snake_case)

        # check that min length is applied at length 5
        UpperCAmelCase_          =							ids_tensor((batch_size, 20)					,						vocab_size=20)
        UpperCAmelCase_          =							5
        UpperCAmelCase_          =							self._get_uniform_logits(_snake_case					,						_snake_case)
        UpperCAmelCase_          =							min_dist_processor(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        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_          =							self._get_uniform_logits(_snake_case					,						_snake_case)
        UpperCAmelCase_          =							15
        UpperCAmelCase_          =							min_dist_processor(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        self.assertFalse(jnp.isinf(_snake_case).any())







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							20
        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							0

        UpperCAmelCase_          =							FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_snake_case)

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







   def lowerCamelCase						(      self		:						List[str]):

        """simple docstring"""


        UpperCAmelCase_          =							20
        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							0
        UpperCAmelCase_          =							5

        UpperCAmelCase_          =							FlaxForcedEOSTokenLogitsProcessor(max_length=_snake_case					,						eos_token_id=_snake_case)

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







   def lowerCamelCase						(      self		:						List[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							10
        UpperCAmelCase_          =							15
        UpperCAmelCase_          =							2
        UpperCAmelCase_          =							1
        UpperCAmelCase_          =							15

        # dummy input_ids and scores
        UpperCAmelCase_          =							ids_tensor((batch_size, sequence_length)					,						_snake_case)
        UpperCAmelCase_          =							input_ids.copy()

        UpperCAmelCase_          =							self._get_uniform_logits(_snake_case					,						_snake_case)
        UpperCAmelCase_          =							scores.copy()

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

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

        UpperCAmelCase_          =							10

        # no processor list
        UpperCAmelCase_          =							temp_dist_warp(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        UpperCAmelCase_          =							top_k_warp(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        UpperCAmelCase_          =							top_p_warp(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        UpperCAmelCase_          =							min_dist_proc(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        UpperCAmelCase_          =							bos_dist_proc(_snake_case					,						_snake_case					,						cur_len=_snake_case)
        UpperCAmelCase_          =							eos_dist_proc(_snake_case					,						_snake_case					,						cur_len=_snake_case)

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

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

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







   def lowerCamelCase						(      self		:						Optional[Any]):

        """simple docstring"""


        UpperCAmelCase_          =							4
        UpperCAmelCase_          =							10
        UpperCAmelCase_          =							15
        UpperCAmelCase_          =							2
        UpperCAmelCase_          =							1
        UpperCAmelCase_          =							15

        # dummy input_ids and scores
        UpperCAmelCase_          =							ids_tensor((batch_size, sequence_length)					,						_snake_case)
        UpperCAmelCase_          =							input_ids.copy()

        UpperCAmelCase_          =							self._get_uniform_logits(_snake_case					,						_snake_case)
        UpperCAmelCase_          =							scores.copy()

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

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

        UpperCAmelCase_          =							10

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

        # with processor list
        def run_processor_list(_snake_case		:						Optional[Any]					,						_snake_case		:						Optional[Any]					,						_snake_case		:						Union[str, Any]):
             UpperCAmelCase_          =							FlaxLogitsProcessorList(
                 [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
             UpperCAmelCase_          =							processor(_snake_case					,						_snake_case					,						cur_len=_snake_case)
             return scores

        UpperCAmelCase_          =							jax.jit(_snake_case)
        UpperCAmelCase_          =							jax.jit(_snake_case)

        UpperCAmelCase_          =							jitted_run_no_processor_list(_snake_case					,						_snake_case					,						_snake_case)
        UpperCAmelCase_          =							jitted_run_processor_list(_snake_case					,						_snake_case					,						_snake_case)

        # scores should be equal
        self.assertTrue(jnp.allclose(_snake_case					,						_snake_case					,						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


UpperCAmelCase__				   =		{
    "configuration_nllb_moe": [
        "NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "NllbMoeConfig",
    ]
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     UpperCAmelCase__				   =		[
         "NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST",
         "NllbMoeForConditionalGeneration",
         "NllbMoeModel",
         "NllbMoePreTrainedModel",
         "NllbMoeTop2Router",
         "NllbMoeSparseMLP",
     ]


if TYPE_CHECKING:
     from .configuration_nllb_moe import (
         NLLB_MOE_PRETRAINED_CONFIG_ARCHIVE_MAP,
         NllbMoeConfig,
     )

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_nllb_moe import (
              NLLB_MOE_PRETRAINED_MODEL_ARCHIVE_LIST,
              NllbMoeForConditionalGeneration,
              NllbMoeModel,
              NllbMoePreTrainedModel,
              NllbMoeSparseMLP,
              NllbMoeTopaRouter,
          )


else:
     import sys

     UpperCAmelCase__				   =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


