




import unittest

from transformers import AlbertTokenizer, AlbertTokenizerFast
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCamelCase			:	str							 =     get_tests_dir("""fixtures/spiece.model""")


@require_sentencepiece
@require_tokenizers
class        UpperCamelCase_					(				UpperCAmelCase__	,      unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   AlbertTokenizer
    UpperCAmelCase__													=   AlbertTokenizerFast
    UpperCAmelCase__													=   True
    UpperCAmelCase__													=   True
    UpperCAmelCase__													=   True


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->Optional[int]:





           '''simple docstring'''

           super().setUp()

           # We have a SentencePiece fixture for testing
           A__   				=					AlbertTokenizer(UpperCAmelCase__)
           tokenizer.save_pretrained(self.tmpdirname)


    def  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			Union[str, Any])		->List[Any]:





           '''simple docstring'''

           A__   				=					'''this is a test'''
           A__   				=					'''this is a test'''
           return input_text, output_text


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int])		->str:





           '''simple docstring'''

           A__   				=					'''<pad>'''
           A__   				=					0

           self.assertEqual(self.get_tokenizer()._convert_token_to_id(UpperCAmelCase__)      ,				UpperCAmelCase__)
           self.assertEqual(self.get_tokenizer()._convert_id_to_token(UpperCAmelCase__)      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Union[str, Any])		->List[str]:





           '''simple docstring'''

           A__   				=					list(self.get_tokenizer().get_vocab().keys())

           self.assertEqual(vocab_keys[0]      ,				'''<pad>''')
           self.assertEqual(vocab_keys[1]      ,				'''<unk>''')
           self.assertEqual(vocab_keys[-1]      ,				'''▁eloquent''')
           self.assertEqual(len(UpperCAmelCase__)      ,				30_000)


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->List[Any]:





           '''simple docstring'''

           self.assertEqual(self.get_tokenizer().vocab_size      ,				30_000)


    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[Any])		->List[str]:





           '''simple docstring'''

           if not self.test_rust_tokenizer:
                  return

           A__   				=					self.get_tokenizer()
           A__   				=					self.get_rust_tokenizer()

           A__   				=					'''I was born in 92000, and this is falsé.'''

           A__   				=					tokenizer.tokenize(UpperCAmelCase__)
           A__   				=					rust_tokenizer.tokenize(UpperCAmelCase__)
           self.assertListEqual(UpperCAmelCase__      ,				UpperCAmelCase__)

           A__   				=					tokenizer.encode(UpperCAmelCase__      ,				add_special_tokens=UpperCAmelCase__)
           A__   				=					rust_tokenizer.encode(UpperCAmelCase__      ,				add_special_tokens=UpperCAmelCase__)
           self.assertListEqual(UpperCAmelCase__      ,				UpperCAmelCase__)

           A__   				=					self.get_rust_tokenizer()
           A__   				=					tokenizer.encode(UpperCAmelCase__)
           A__   				=					rust_tokenizer.encode(UpperCAmelCase__)
           self.assertListEqual(UpperCAmelCase__      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->Union[str, Any]:





           '''simple docstring'''

           A__   				=					AlbertTokenizer(UpperCAmelCase__      ,				keep_accents=UpperCAmelCase__)

           A__   				=					tokenizer.tokenize('''This is a test''')
           self.assertListEqual(UpperCAmelCase__      ,				['''▁this''', '''▁is''', '''▁a''', '''▁test'''])

           self.assertListEqual(tokenizer.convert_tokens_to_ids(UpperCAmelCase__)      ,				[48, 25, 21, 1_289])

           A__   				=					tokenizer.tokenize('''I was born in 92000, and this is falsé.''')
           self.assertListEqual(
               UpperCAmelCase__      ,				['''▁i''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''é''', '''.'''])
           A__   				=					tokenizer.convert_tokens_to_ids(UpperCAmelCase__)
           self.assertListEqual(UpperCAmelCase__      ,				[31, 23, 386, 19, 561, 3_050, 15, 17, 48, 25, 8_256, 18, 1, 9])

           A__   				=					tokenizer.convert_ids_to_tokens(UpperCAmelCase__)
           self.assertListEqual(
               UpperCAmelCase__      ,				['''▁i''', '''▁was''', '''▁born''', '''▁in''', '''▁9''', '''2000''', ''',''', '''▁and''', '''▁this''', '''▁is''', '''▁fal''', '''s''', '''<unk>''', '''.''']      ,				)


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





           '''simple docstring'''

           A__   				=					AlbertTokenizer(UpperCAmelCase__)

           A__   				=					tokenizer.encode('''sequence builders''')
           A__   				=					tokenizer.encode('''multi-sequence build''')

           A__   				=					tokenizer.build_inputs_with_special_tokens(UpperCAmelCase__)
           A__   				=					tokenizer.build_inputs_with_special_tokens(UpperCAmelCase__      ,				UpperCAmelCase__)

           assert encoded_sentence == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id]
           assert encoded_pair == [tokenizer.cls_token_id] + text + [tokenizer.sep_token_id] + text_a + [
               tokenizer.sep_token_id
           ]





    @slow
    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any])		->Optional[int]:





           '''simple docstring'''

           A__   				=					{'''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''input_ids''': [[2, 21_970, 13, 5, 6_092, 167, 28, 7_103, 2_153, 673, 8, 7_028, 12_051, 18, 17, 7_103, 2_153, 673, 8, 3_515, 18_684, 8, 4_461, 6, 1_927, 297, 8, 12_060, 2_607, 18, 13, 5, 4_461, 15, 10_538, 38, 8, 135, 15, 822, 58, 15, 993, 10_363, 15, 1_460, 8_005, 4_461, 15, 993, 255, 2_328, 9, 9, 9, 6, 26, 1_112, 816, 3_260, 13, 5, 103, 2_377, 6, 17, 1_112, 816, 2_782, 13, 5, 103, 10_641, 6, 29, 84, 2_512, 2_430, 782, 18_684, 2_761, 19, 808, 2_430, 2_556, 17, 855, 1_480, 9_477, 4_091, 128, 11_712, 15, 7_103, 2_153, 673, 17, 24_883, 9_990, 9, 3], [2, 11_502, 25, 1_006, 20, 782, 8, 11_809, 855, 1_732, 19_393, 18_667, 37, 367, 21_018, 69, 1_854, 34, 11_860, 19_124, 27, 156, 225, 17, 193, 4_141, 19, 65, 9_124, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [2, 14, 2_231, 886, 2_385, 17_659, 84, 14, 16_792, 1_952, 9, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], '''token_type_ids''': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
           # fmt: on

           self.tokenizer_integration_test_util(
               expected_encoding=UpperCAmelCase__      ,				model_name='''albert-base-v2'''      ,				revision='''6b6560eaf5ff2e250b00c50f380c5389a9c2d82e'''      ,				)








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  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			int      ,				UpperCAmelCase__   :			int)		->Dict:





           '''simple docstring'''

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


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





           '''simple docstring'''

           A__   				=					None
           A__   				=					20

           A__   				=					self._get_uniform_logits(batch_size=2      ,				length=UpperCAmelCase__)

           # 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(UpperCAmelCase__      ,				axis=-1)

           A__   				=					FlaxTemperatureLogitsWarper(temperature=0.5)
           A__   				=					FlaxTemperatureLogitsWarper(temperature=1.3)

           A__   				=					jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase__      ,				scores.copy()      ,				cur_len=UpperCAmelCase__)      ,				axis=-1)
           A__   				=					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)		->int:





           '''simple docstring'''

           A__   				=					None
           A__   				=					10
           A__   				=					2

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

           A__   				=					FlaxTopKLogitsWarper(3)

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

           A__   				=					np.broadcast_to(np.arange(UpperCAmelCase__)[None, :]      ,				(batch_size, length)).copy()
           A__   				=					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   :			List[Any])		->Union[str, Any]:





           '''simple docstring'''

           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(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__))

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

           # check edge cases with negative and extreme logits
           A__   				=					np.broadcast_to(np.arange(UpperCAmelCase__)[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(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   :			Tuple)		->Tuple:





           '''simple docstring'''

           A__   				=					20
           A__   				=					4
           A__   				=					0

           A__   				=					FlaxMinLengthLogitsProcessor(min_length=10      ,				eos_token_id=UpperCAmelCase__)

           # 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(UpperCAmelCase__      ,				UpperCAmelCase__)
           A__   				=					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
           A__   				=					self._get_uniform_logits(UpperCAmelCase__      ,				UpperCAmelCase__)
           A__   				=					15
           A__   				=					min_dist_processor(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           self.assertFalse(jnp.isinf(UpperCAmelCase__).any())


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





           '''simple docstring'''

           A__   				=					20
           A__   				=					4
           A__   				=					0

           A__   				=					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase__)

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


    def  SCREAMING_SNAKE_CASE							(    self   :			Tuple)		->str:





           '''simple docstring'''

           A__   				=					20
           A__   				=					4
           A__   				=					0
           A__   				=					5

           A__   				=					FlaxForcedEOSTokenLogitsProcessor(max_length=UpperCAmelCase__      ,				eos_token_id=UpperCAmelCase__)

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


    def  SCREAMING_SNAKE_CASE							(    self   :			List[str])		->Tuple:





           '''simple docstring'''

           A__   				=					4
           A__   				=					10
           A__   				=					15
           A__   				=					2
           A__   				=					1
           A__   				=					15

           # dummy input_ids and scores
           A__   				=					ids_tensor((batch_size, sequence_length)      ,				UpperCAmelCase__)
           A__   				=					input_ids.copy()

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

           A__   				=					10

           # no processor list
           A__   				=					temp_dist_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           A__   				=					top_k_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           A__   				=					top_p_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           A__   				=					min_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           A__   				=					bos_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
           A__   				=					eos_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)

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





           '''simple docstring'''

           A__   				=					4
           A__   				=					10
           A__   				=					15
           A__   				=					2
           A__   				=					1
           A__   				=					15

           # dummy input_ids and scores
           A__   				=					ids_tensor((batch_size, sequence_length)      ,				UpperCAmelCase__)
           A__   				=					input_ids.copy()

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

           A__   				=					10

           # no processor list
           def run_no_processor_list(UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Optional[int]      ,				UpperCAmelCase__   :			Union[str, Any]):
                  A__   				=					temp_dist_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  A__   				=					top_k_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  A__   				=					top_p_warp(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  A__   				=					min_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  A__   				=					bos_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  A__   				=					eos_dist_proc(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  return scores

           # with processor list
           def run_processor_list(UpperCAmelCase__   :			Any      ,				UpperCAmelCase__   :			Dict      ,				UpperCAmelCase__   :			Optional[Any]):
                  A__   				=					FlaxLogitsProcessorList(
                      [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
                  A__   				=					processor(UpperCAmelCase__      ,				UpperCAmelCase__      ,				cur_len=UpperCAmelCase__)
                  return scores

           A__   				=					jax.jit(UpperCAmelCase__)
           A__   				=					jax.jit(UpperCAmelCase__)

           A__   				=					jitted_run_no_processor_list(UpperCAmelCase__      ,				UpperCAmelCase__      ,				UpperCAmelCase__)
           A__   				=					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())


