

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





  def 						A  (		self				:     Tuple							, _a				:     int							, _a				:     int			)     ->   List[str]:





       '''simple docstring'''






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





  def 						A  (		self				:     List[Any]			)     ->   Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =None
       _SCREAMING_SNAKE_CASE							   =20

       _SCREAMING_SNAKE_CASE							   =self._get_uniform_logits(batch_size=2							, length=_a			)

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

       # compute softmax
       _SCREAMING_SNAKE_CASE							   =jax.nn.softmax(_a							, 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(_a							, scores.copy()							, cur_len=_a			)							, axis=-1			)
       _SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     Dict			)     ->   Any:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =None
       _SCREAMING_SNAKE_CASE							   =10
       _SCREAMING_SNAKE_CASE							   =2

       # create ramp distribution
       _SCREAMING_SNAKE_CASE							   =np.broadcast_to(np.arange(_a			)[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(_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
       _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(_a			)[None, :]							, (batch_size, length)			).copy()
       _SCREAMING_SNAKE_CASE							   =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 						A  (		self				:     Any			)     ->   Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =None
       _SCREAMING_SNAKE_CASE							   =10
       _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(_a							, _a							, cur_len=_a			)			)

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

       # check edge cases with negative and extreme logits
       _SCREAMING_SNAKE_CASE							   =np.broadcast_to(np.arange(_a			)[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(_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 						A  (		self				:     Optional[Any]			)     ->   Tuple:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =20
       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =0

       _SCREAMING_SNAKE_CASE							   =FlaxMinLengthLogitsProcessor(min_length=10							, eos_token_id=_a			)

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





  def 						A  (		self				:     Dict			)     ->   Optional[Any]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =20
       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =0

       _SCREAMING_SNAKE_CASE							   =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_a			)

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





  def 						A  (		self				:     Any			)     ->   Optional[int]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =20
       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =0
       _SCREAMING_SNAKE_CASE							   =5

       _SCREAMING_SNAKE_CASE							   =FlaxForcedEOSTokenLogitsProcessor(max_length=_a							, eos_token_id=_a			)

       # 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=20			)
       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =self._get_uniform_logits(_a							, _a			)
       _SCREAMING_SNAKE_CASE							   =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
       _SCREAMING_SNAKE_CASE							   =3
       _SCREAMING_SNAKE_CASE							   =self._get_uniform_logits(_a							, _a			)
       _SCREAMING_SNAKE_CASE							   =logits_processor(_a							, _a							, cur_len=_a			)
       self.assertFalse(jnp.isinf(_a			).any()			)





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





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =10
       _SCREAMING_SNAKE_CASE							   =15
       _SCREAMING_SNAKE_CASE							   =2
       _SCREAMING_SNAKE_CASE							   =1
       _SCREAMING_SNAKE_CASE							   =15

       # dummy input_ids and scores
       _SCREAMING_SNAKE_CASE							   =ids_tensor((batch_size, sequence_length)							, _a			)
       _SCREAMING_SNAKE_CASE							   =input_ids.copy()

       _SCREAMING_SNAKE_CASE							   =self._get_uniform_logits(_a							, _a			)
       _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=10							, eos_token_id=_a			)
       _SCREAMING_SNAKE_CASE							   =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_a			)
       _SCREAMING_SNAKE_CASE							   =FlaxForcedEOSTokenLogitsProcessor(max_length=_a							, eos_token_id=_a			)

       _SCREAMING_SNAKE_CASE							   =10

       # no processor list
       _SCREAMING_SNAKE_CASE							   =temp_dist_warp(_a							, _a							, cur_len=_a			)
       _SCREAMING_SNAKE_CASE							   =top_k_warp(_a							, _a							, cur_len=_a			)
       _SCREAMING_SNAKE_CASE							   =top_p_warp(_a							, _a							, cur_len=_a			)
       _SCREAMING_SNAKE_CASE							   =min_dist_proc(_a							, _a							, cur_len=_a			)
       _SCREAMING_SNAKE_CASE							   =bos_dist_proc(_a							, _a							, cur_len=_a			)
       _SCREAMING_SNAKE_CASE							   =eos_dist_proc(_a							, _a							, cur_len=_a			)

       # 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(_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 						A  (		self				:     str			)     ->   List[str]:





       '''simple docstring'''






       _SCREAMING_SNAKE_CASE							   =4
       _SCREAMING_SNAKE_CASE							   =10
       _SCREAMING_SNAKE_CASE							   =15
       _SCREAMING_SNAKE_CASE							   =2
       _SCREAMING_SNAKE_CASE							   =1
       _SCREAMING_SNAKE_CASE							   =15

       # dummy input_ids and scores
       _SCREAMING_SNAKE_CASE							   =ids_tensor((batch_size, sequence_length)							, _a			)
       _SCREAMING_SNAKE_CASE							   =input_ids.copy()

       _SCREAMING_SNAKE_CASE							   =self._get_uniform_logits(_a							, _a			)
       _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=10							, eos_token_id=_a			)
       _SCREAMING_SNAKE_CASE							   =FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_a			)
       _SCREAMING_SNAKE_CASE							   =FlaxForcedEOSTokenLogitsProcessor(max_length=_a							, eos_token_id=_a			)

       _SCREAMING_SNAKE_CASE							   =10

       # no processor list
       def run_no_processor_list(_a				:     Dict							, _a				:     List[str]							, _a				:     List[Any]			):
            _SCREAMING_SNAKE_CASE							   =temp_dist_warp(_a							, _a							, cur_len=_a			)
            _SCREAMING_SNAKE_CASE							   =top_k_warp(_a							, _a							, cur_len=_a			)
            _SCREAMING_SNAKE_CASE							   =top_p_warp(_a							, _a							, cur_len=_a			)
            _SCREAMING_SNAKE_CASE							   =min_dist_proc(_a							, _a							, cur_len=_a			)
            _SCREAMING_SNAKE_CASE							   =bos_dist_proc(_a							, _a							, cur_len=_a			)
            _SCREAMING_SNAKE_CASE							   =eos_dist_proc(_a							, _a							, cur_len=_a			)
            return scores

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

       _SCREAMING_SNAKE_CASE							   =jax.jit(_a			)
       _SCREAMING_SNAKE_CASE							   =jax.jit(_a			)

       _SCREAMING_SNAKE_CASE							   =jitted_run_no_processor_list(_a							, _a							, _a			)
       _SCREAMING_SNAKE_CASE							   =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()			)



'''simple docstring'''
from datasets.utils.patching import _PatchedModuleObj, patch_submodule

from . import _test_patching


def 	_lowerCAmelCase (						)				->      int:
     """simple docstring"""




     import os as original_os
     from os import path as original_path
     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_mock__'
     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          # Every way to access os.path.join must be patched, and the rest must stay untouched

          # check os.path.join
          assert isinstance(_test_patching.os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.os.path							,			_PatchedModuleObj    )
          assert _test_patching.os.path.join is mock

          # check path.join
          assert isinstance(_test_patching.path							,			_PatchedModuleObj    )
          assert _test_patching.path.join is mock

          # check join
          assert _test_patching.join is mock

          # check that the other attributes are untouched
          assert _test_patching.os.rename is original_rename
          assert _test_patching.path.dirname is original_dirname
          assert _test_patching.os.path.dirname is original_dirname

          # Even renamed modules or objects must be patched

          # check renamed_os.path.join
          assert isinstance(_test_patching.renamed_os							,			_PatchedModuleObj    )
          assert isinstance(_test_patching.renamed_os.path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_os.path.join is mock

          # check renamed_path.join
          assert isinstance(_test_patching.renamed_path							,			_PatchedModuleObj    )
          assert _test_patching.renamed_path.join is mock

          # check renamed_join
          assert _test_patching.renamed_join is mock

          # check that the other attributes are untouched
          assert _test_patching.renamed_os.rename is original_rename
          assert _test_patching.renamed_path.dirname is original_dirname
          assert _test_patching.renamed_os.path.dirname is original_dirname

     # check that everthing is back to normal when the patch is over

     assert _test_patching.os is original_os
     assert _test_patching.path is original_path
     assert _test_patching.join is original_join

     assert _test_patching.renamed_os is original_os
     assert _test_patching.renamed_path is original_path
     assert _test_patching.renamed_join is original_join


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     assert _test_patching.open is open

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_builtin_mock__'
     # _test_patching has "open" in its globals
     assert _test_patching.open is open
     with patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    ):
          assert _test_patching.open is mock

     # check that everthing is back to normal when the patch is over

     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_mock__'
     with patch_submodule(_test_patching							,			'pandas.read_csv'							,			_UpperCamelCase    ):
          pass


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_missing_builtin_mock__'
     # _test_patching doesn't have "len" in its globals
     assert getattr(_test_patching							,			'len'							,			_UpperCamelCase    ) is None
     with patch_submodule(_test_patching							,			'len'							,			_UpperCamelCase    ):
          assert _test_patching.len is mock
     assert _test_patching.len is len


def 	_lowerCAmelCase (						)				->      str:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_start_and_stop_mock__'
     _SCREAMING_SNAKE_CASE							   =patch_submodule(_test_patching							,			'open'							,			_UpperCamelCase    )
     assert _test_patching.open is open
     patch.start()
     assert _test_patching.open is mock
     patch.stop()
     assert _test_patching.open is open


def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     from os import rename as original_rename
     from os.path import dirname as original_dirname
     from os.path import join as original_join

     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_join__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_dirname__'
     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_successive_rename__'
     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename

     with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

    # try another order
     with patch_submodule(_test_patching							,			'os.rename'							,			_UpperCamelCase    ):
          with patch_submodule(_test_patching							,			'os.path.join'							,			_UpperCamelCase    ):
               with patch_submodule(_test_patching							,			'os.path.dirname'							,			_UpperCamelCase    ):
                    assert _test_patching.os.path.join is mock_join
                    assert _test_patching.os.path.dirname is mock_dirname
                    assert _test_patching.os.rename is mock_rename

     assert _test_patching.os.path.join is original_join
     assert _test_patching.os.path.dirname is original_dirname
     assert _test_patching.os.rename is original_rename







def 	_lowerCAmelCase (						)				->      Optional[Any]:
     """simple docstring"""




     _SCREAMING_SNAKE_CASE							   ='__test_patch_submodule_doesnt_exist_mock__'
     with patch_submodule(_test_patching							,			'__module_that_doesn_exist__.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
     with patch_submodule(_test_patching							,			'os.__attribute_that_doesn_exist__'							,			_UpperCamelCase    ):
          pass
