
"""simple docstring"""







import json
import os
import unittest

from transformers import DebertaTokenizer, DebertaTokenizerFast
from transformers.models.deberta.tokenization_deberta import VOCAB_FILES_NAMES
from transformers.testing_utils import slow

from ...test_tokenization_common import TokenizerTesterMixin



class 		UpperCAmelCase_						(lowerCamelCase_	,			unittest.TestCase    ):
 """simple docstring"""




 UpperCamelCase_    :  Any					 =     DebertaTokenizer
 UpperCamelCase_    :  List[str]					 =     True
 UpperCamelCase_    :  Dict					 =     DebertaTokenizerFast




 def 	a   (			self							:			Optional[int]			)->       Union[str, Any]:



      """simple docstring"""

      super().setUp()

      # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
      UpperCAmelCase_    :     Optional[Any]            =	[
          """l""",
          """o""",
          """w""",
          """e""",
          """r""",
          """s""",
          """t""",
          """i""",
          """d""",
          """n""",
          """\u0120""",
          """\u0120l""",
          """\u0120n""",
          """\u0120lo""",
          """\u0120low""",
          """er""",
          """\u0120lowest""",
          """\u0120newer""",
          """\u0120wider""",
          """[UNK]""",
      ]
      UpperCAmelCase_    :     List[str]            =	dict(zip(a_  , range(len(a_			)			)			)			)
      UpperCAmelCase_    :     Union[str, Any]            =	["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
      UpperCAmelCase_    :     Tuple            =	{"""unk_token""": """[UNK]"""}

      UpperCAmelCase_    :     str            =	os.path.join(self.tmpdirname  , VOCAB_FILES_NAMES["""vocab_file"""]			)
      UpperCAmelCase_    :     Optional[int]            =	os.path.join(self.tmpdirname  , VOCAB_FILES_NAMES["""merges_file"""]			)
      with open(self.vocab_file  , """w"""  , encoding="""utf-8"""			) as fp:
           fp.write(json.dumps(a_			) + """\n"""			)
      with open(self.merges_file  , """w"""  , encoding="""utf-8"""			) as fp:
           fp.write("""\n""".join(a_			)			)




 def 	a   (			self							:			Dict  , **a_							:			int			)->       Any:



      """simple docstring"""

      kwargs.update(self.special_tokens_map			)
      return self.tokenizer_class.from_pretrained(self.tmpdirname  , **a_			)




 def 	a   (			self							:			Dict  , a_							:			Optional[int]			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	"""lower newer"""
      UpperCAmelCase_    :     Optional[int]            =	"""lower newer"""
      return input_text, output_text




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



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	self.get_tokenizer()
      UpperCAmelCase_    :     str            =	"""lower newer"""
      UpperCAmelCase_    :     List[Any]            =	["""l""", """o""", """w""", """er""", """\u0120""", """n""", """e""", """w""", """er"""]
      UpperCAmelCase_    :     Union[str, Any]            =	tokenizer.tokenize(a_			)
      self.assertListEqual(a_  , a_			)

      UpperCAmelCase_    :     Dict            =	tokens + [tokenizer.unk_token]
      UpperCAmelCase_    :     Union[str, Any]            =	[0, 1, 2, 15, 10, 9, 3, 2, 15, 19]
      self.assertListEqual(tokenizer.convert_tokens_to_ids(a_			)  , a_			)




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



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	self.get_tokenizer()
      UpperCAmelCase_    :     Optional[int]            =	tokenizer("""Hello"""  , """World"""			)
      UpperCAmelCase_    :     Optional[Any]            =	[0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
      self.assertListEqual(tokd["""token_type_ids"""]  , a_			)




 @slow
 def 	a   (			self							:			int			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     str            =	self.tokenizer_class.from_pretrained("""microsoft/deberta-base"""			)

      UpperCAmelCase_    :     Optional[Any]            =	tokenizer.encode("""sequence builders"""  , add_special_tokens=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	tokenizer.encode("""multi-sequence build"""  , add_special_tokens=a_			)

      UpperCAmelCase_    :     Any            =	tokenizer.encode(
          """sequence builders"""  , add_special_tokens=a_  , add_prefix_space=a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	tokenizer.encode(
          """sequence builders"""  , """multi-sequence build"""  , add_special_tokens=a_  , add_prefix_space=a_			)

      UpperCAmelCase_    :     List[Any]            =	tokenizer.build_inputs_with_special_tokens(a_			)
      UpperCAmelCase_    :     Dict            =	tokenizer.build_inputs_with_special_tokens(a_  , a_			)

      assert encoded_sentence == encoded_text_from_decode
      assert encoded_pair == encoded_pair_from_decode






 @slow
 def 	a   (			self							:			List[Any]			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[Any]            =	[self.tokenizer_class]
      if self.test_rust_tokenizer:
           tokenizer_classes.append(self.rust_tokenizer_class			)

      for tokenizer_class in tokenizer_classes:
           UpperCAmelCase_    :     Tuple            =	tokenizer_class.from_pretrained("""microsoft/deberta-base"""			)

           UpperCAmelCase_    :     List[Any]            =	[
               """ALBERT: A Lite BERT for Self-supervised Learning of Language Representations""",
               """ALBERT incorporates two parameter reduction techniques""",
               """The first one is a factorized embedding parameterization. By decomposing the large vocabulary"""
               """ embedding matrix into two small matrices, we separate the size of the hidden layers from the size of"""
               """ vocabulary embedding.""",
           ]

           UpperCAmelCase_    :     List[str]            =	tokenizer(a_  , padding=a_			)
           UpperCAmelCase_    :     Optional[Any]            =	[tokenizer.decode(a_  , skip_special_tokens=a_			) for seq in encoding["""input_ids"""]]

           # fmt: off
           UpperCAmelCase_    :     Any            =	{
               """input_ids""": [
                   [1, 21_18, 1_11_26, 5_65, 35, 83, 2_51_91, 1_63, 1_88_54, 13, 1_21_56, 12, 1_61_01, 2_53_76, 1_38_07, 9, 2_22_05, 2_78_93, 16_35, 2, 0, 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, 21_18, 1_11_26, 5_65, 2_45_36, 80, 4_37_97, 48_78, 73_73, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 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_33, 78, 65, 16, 10, 37_24, 15_38, 3_31_83, 1_13_03, 4_37_97, 19_38, 4, 8_70, 2_41_65, 2_91_05, 5, 7_39, 3_26_44, 3_31_83, 1_13_03, 3_61_73, 88, 80, 6_50, 78_21, 4_59_40, 6, 52, 25_59, 5, 18_36, 9, 5, 73_97, 1_31_71, 31, 5, 18_36, 9, 3_26_44, 3_31_83, 1_13_03, 4, 2]
               ],
               """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]
               ],
               """attention_mask""": [
                   [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],
                   [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],
                   [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]
               ]
           }
           # fmt: on

           UpperCAmelCase_    :     Optional[Any]            =	[
               """ALBERT: A Lite BERT for Self-supervised Learning of Language Representations""",
               """ALBERT incorporates two parameter reduction techniques""",
               """The first one is a factorized embedding parameterization. By decomposing the large vocabulary"""
               """ embedding matrix into two small matrices, we separate the size of the hidden layers from the size of"""
               """ vocabulary embedding.""",
           ]

           self.assertDictEqual(encoding.data  , a_			)

           for expected, decoded in zip(a_  , a_			):
                self.assertEqual(a_  , a_			)



"""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 		UpperCAmelCase_						(unittest.TestCase    ):
 """simple docstring"""




 def 	a   (			self							:			str  , a_							:			int  , a_							:			int			)->       Optional[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	jnp.ones((batch_size, length)			) / length
      return scores




 def 	a   (			self							:			str			)->       int:



      """simple docstring"""

      UpperCAmelCase_    :     List[Any]            =	None
      UpperCAmelCase_    :     int            =	20

      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(batch_size=2  , length=a_			)

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

      # compute softmax
      UpperCAmelCase_    :     str            =	jax.nn.softmax(a_  , axis=-1			)

      UpperCAmelCase_    :     Union[str, Any]            =	FlaxTemperatureLogitsWarper(temperature=0.5			)
      UpperCAmelCase_    :     Any            =	FlaxTemperatureLogitsWarper(temperature=1.3			)

      UpperCAmelCase_    :     Optional[Any]            =	jax.nn.softmax(temp_dist_warper_sharper(a_  , scores.copy()  , cur_len=a_			)  , axis=-1			)
      UpperCAmelCase_    :     List[str]            =	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							:			Union[str, Any]			)->       Tuple:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	None
      UpperCAmelCase_    :     str            =	10
      UpperCAmelCase_    :     List[Any]            =	2

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

      UpperCAmelCase_    :     List[Any]            =	FlaxTopKLogitsWarper(3			)

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

      UpperCAmelCase_    :     Any            =	np.broadcast_to(np.arange(a_			)[None, :]  , (batch_size, length)			).copy()
      UpperCAmelCase_    :     Tuple            =	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							:			Dict			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Tuple            =	None
      UpperCAmelCase_    :     List[str]            =	10
      UpperCAmelCase_    :     List[Any]            =	2

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

      UpperCAmelCase_    :     str            =	FlaxTopPLogitsWarper(0.8			)
      UpperCAmelCase_    :     Union[str, Any]            =	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
      UpperCAmelCase_    :     List[Any]            =	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
      UpperCAmelCase_    :     Optional[Any]            =	np.broadcast_to(np.arange(a_			)[None, :]  , (batch_size, vocab_size)			).copy() - (
          vocab_size // 2
      )

      # make ramp_logits more extreme
      UpperCAmelCase_    :     Tuple            =	ramp_logits[1] * 100.0

      # make sure at least 2 tokens are kept
      UpperCAmelCase_    :     Any            =	FlaxTopPLogitsWarper(0.9  , min_tokens_to_keep=2  , filter_value=0.0			)
      UpperCAmelCase_    :     Dict            =	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							:			List[str]			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	20
      UpperCAmelCase_    :     List[str]            =	4
      UpperCAmelCase_    :     Dict            =	0

      UpperCAmelCase_    :     Union[str, Any]            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)

      # check that min length is applied at length 5
      UpperCAmelCase_    :     Union[str, Any]            =	ids_tensor((batch_size, 20)  , vocab_size=20			)
      UpperCAmelCase_    :     str            =	5
      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     List[Any]            =	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
      UpperCAmelCase_    :     str            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Optional[Any]            =	15
      UpperCAmelCase_    :     str            =	min_dist_processor(a_  , a_  , cur_len=a_			)
      self.assertFalse(jnp.isinf(a_			).any()			)




 def 	a   (			self							:			List[str]			)->       str:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	20
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     str            =	0

      UpperCAmelCase_    :     Optional[int]            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)

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




 def 	a   (			self							:			str			)->       List[str]:



      """simple docstring"""

      UpperCAmelCase_    :     Optional[int]            =	20
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     Union[str, Any]            =	0
      UpperCAmelCase_    :     Union[str, Any]            =	5

      UpperCAmelCase_    :     Optional[Any]            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      # 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=20			)
      UpperCAmelCase_    :     int            =	4
      UpperCAmelCase_    :     Optional[Any]            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	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
      UpperCAmelCase_    :     Any            =	3
      UpperCAmelCase_    :     Tuple            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     int            =	logits_processor(a_  , a_  , cur_len=a_			)
      self.assertFalse(jnp.isinf(a_			).any()			)




 def 	a   (			self							:			int			)->       List[Any]:



      """simple docstring"""

      UpperCAmelCase_    :     Dict            =	4
      UpperCAmelCase_    :     List[Any]            =	10
      UpperCAmelCase_    :     Tuple            =	15
      UpperCAmelCase_    :     Optional[Any]            =	2
      UpperCAmelCase_    :     List[Any]            =	1
      UpperCAmelCase_    :     Dict            =	15

      # dummy input_ids and scores
      UpperCAmelCase_    :     Dict            =	ids_tensor((batch_size, sequence_length)  , a_			)
      UpperCAmelCase_    :     Tuple            =	input_ids.copy()

      UpperCAmelCase_    :     int            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Optional[Any]            =	scores.copy()

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

      # instantiate all logits processors
      UpperCAmelCase_    :     Tuple            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)
      UpperCAmelCase_    :     Tuple            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      UpperCAmelCase_    :     Tuple            =	10

      # no processor list
      UpperCAmelCase_    :     Optional[Any]            =	temp_dist_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Optional[int]            =	top_k_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Optional[Any]            =	top_p_warp(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     List[Any]            =	min_dist_proc(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     Dict            =	bos_dist_proc(a_  , a_  , cur_len=a_			)
      UpperCAmelCase_    :     str            =	eos_dist_proc(a_  , a_  , cur_len=a_			)

      # with processor list
      UpperCAmelCase_    :     Tuple            =	FlaxLogitsProcessorList(
          [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]			)
      UpperCAmelCase_    :     Union[str, Any]            =	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							:			Optional[int]			)->       Any:



      """simple docstring"""

      UpperCAmelCase_    :     Any            =	4
      UpperCAmelCase_    :     List[Any]            =	10
      UpperCAmelCase_    :     Optional[int]            =	15
      UpperCAmelCase_    :     Optional[Any]            =	2
      UpperCAmelCase_    :     Dict            =	1
      UpperCAmelCase_    :     int            =	15

      # dummy input_ids and scores
      UpperCAmelCase_    :     List[str]            =	ids_tensor((batch_size, sequence_length)  , a_			)
      UpperCAmelCase_    :     Dict            =	input_ids.copy()

      UpperCAmelCase_    :     Dict            =	self._get_uniform_logits(a_  , a_			)
      UpperCAmelCase_    :     Tuple            =	scores.copy()

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

      # instantiate all logits processors
      UpperCAmelCase_    :     str            =	FlaxMinLengthLogitsProcessor(min_length=10  , eos_token_id=a_			)
      UpperCAmelCase_    :     Dict            =	FlaxForcedBOSTokenLogitsProcessor(bos_token_id=a_			)
      UpperCAmelCase_    :     Union[str, Any]            =	FlaxForcedEOSTokenLogitsProcessor(max_length=a_  , eos_token_id=a_			)

      UpperCAmelCase_    :     Dict            =	10

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

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

      UpperCAmelCase_    :     str            =	jax.jit(a_			)
      UpperCAmelCase_    :     Tuple            =	jax.jit(a_			)

      UpperCAmelCase_    :     Optional[int]            =	jitted_run_no_processor_list(a_  , a_  , a_			)
      UpperCAmelCase_    :     List[Any]            =	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()			)

