




import json
import os
import sys
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from huggingface_hub import HfFolder, Repository, create_repo, delete_repo
from requests.exceptions import HTTPError

import transformers
from transformers import (
    CONFIG_MAPPING,
    FEATURE_EXTRACTOR_MAPPING,
    PROCESSOR_MAPPING,
    TOKENIZER_MAPPING,
    AutoConfig,
    AutoFeatureExtractor,
    AutoProcessor,
    AutoTokenizer,
    BertTokenizer,
    ProcessorMixin,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
)
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test
from transformers.tokenization_utils import TOKENIZER_CONFIG_FILE
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_tokenizers_available


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402
from test_module.custom_processing import CustomProcessor  # noqa E402
from test_module.custom_tokenization import CustomTokenizer  # noqa E402


_lowerCamelCase			:	Dict							 =     get_tests_dir("""fixtures/dummy_feature_extractor_config.json""")
_lowerCamelCase			:	List[Any]							 =     get_tests_dir("""fixtures/vocab.json""")
_lowerCamelCase			:	int							 =     get_tests_dir("""fixtures""")


class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   ['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''bla''', '''blou''']


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





           '''simple docstring'''

           A__   				=					0


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





           '''simple docstring'''

           A__   				=					AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h''')
           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


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





           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmpdirname:
                  A__   				=					WavaVecaConfig()
                  A__   				=					AutoProcessor.from_pretrained('''facebook/wav2vec2-base-960h''')

                  # save in new folder
                  model_config.save_pretrained(UpperCAmelCase__)
                  processor.save_pretrained(UpperCAmelCase__)

                  A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)

           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


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





           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmpdirname:
                  # copy relevant files
                  copyfile(UpperCAmelCase__      ,				os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__))
                  copyfile(UpperCAmelCase__      ,				os.path.join(UpperCAmelCase__      ,				'''vocab.json'''))

                  A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)

           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Tuple:





           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmpdirname:
                  A__   				=					WavaVecaFeatureExtractor()
                  A__   				=					AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h''')

                  A__   				=					WavaVecaProcessor(UpperCAmelCase__      ,				UpperCAmelCase__)

                  # save in new folder
                  processor.save_pretrained(UpperCAmelCase__)

                  # drop `processor_class` in tokenizer
                  with open(os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__)      ,				'''r''') as f:
                         A__   				=					json.load(UpperCAmelCase__)
                         config_dict.pop('''processor_class''')

                  with open(os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__)      ,				'''w''') as f:
                         f.write(json.dumps(UpperCAmelCase__))

                  A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)

           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


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





           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmpdirname:
                  A__   				=					WavaVecaFeatureExtractor()
                  A__   				=					AutoTokenizer.from_pretrained('''facebook/wav2vec2-base-960h''')

                  A__   				=					WavaVecaProcessor(UpperCAmelCase__      ,				UpperCAmelCase__)

                  # save in new folder
                  processor.save_pretrained(UpperCAmelCase__)

                  # drop `processor_class` in feature extractor
                  with open(os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__)      ,				'''r''') as f:
                         A__   				=					json.load(UpperCAmelCase__)
                         config_dict.pop('''processor_class''')

                  with open(os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__)      ,				'''w''') as f:
                         f.write(json.dumps(UpperCAmelCase__))

                  A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)

           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


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





           '''simple docstring'''

           with tempfile.TemporaryDirectory() as tmpdirname:
                  A__   				=					WavaVecaConfig(processor_class='''Wav2Vec2Processor''')
                  model_config.save_pretrained(UpperCAmelCase__)
                  # copy relevant files
                  copyfile(UpperCAmelCase__      ,				os.path.join(UpperCAmelCase__      ,				'''vocab.json'''))
                  # create emtpy sample processor
                  with open(os.path.join(UpperCAmelCase__      ,				UpperCAmelCase__)      ,				'''w''') as f:
                         f.write('''{}''')

                  A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)

           self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)


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





           '''simple docstring'''

           with self.assertRaises(UpperCAmelCase__):
                  A__   				=					AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor''')
           # If remote code is disabled, we can't load this config.
           with self.assertRaises(UpperCAmelCase__):
                  A__   				=					AutoProcessor.from_pretrained(
                      '''hf-internal-testing/test_dynamic_processor'''      ,				trust_remote_code=UpperCAmelCase__)

           A__   				=					AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor'''      ,				trust_remote_code=UpperCAmelCase__)
           self.assertTrue(processor.special_attribute_present)
           self.assertEqual(processor.__class__.__name__      ,				'''NewProcessor''')

           A__   				=					processor.feature_extractor
           self.assertTrue(feature_extractor.special_attribute_present)
           self.assertEqual(feature_extractor.__class__.__name__      ,				'''NewFeatureExtractor''')

           A__   				=					processor.tokenizer
           self.assertTrue(tokenizer.special_attribute_present)
           if is_tokenizers_available():
                  self.assertEqual(tokenizer.__class__.__name__      ,				'''NewTokenizerFast''')

                  # Test we can also load the slow version
                  A__   				=					AutoProcessor.from_pretrained(
                      '''hf-internal-testing/test_dynamic_processor'''      ,				trust_remote_code=UpperCAmelCase__      ,				use_fast=UpperCAmelCase__)
                  A__   				=					new_processor.tokenizer
                  self.assertTrue(new_tokenizer.special_attribute_present)
                  self.assertEqual(new_tokenizer.__class__.__name__      ,				'''NewTokenizer''')
           else:
                  self.assertEqual(tokenizer.__class__.__name__      ,				'''NewTokenizer''')


    def  SCREAMING_SNAKE_CASE							(    self   :			Any)		->Tuple:





           '''simple docstring'''

           try:
                  AutoConfig.register('''custom'''      ,				UpperCAmelCase__)
                  AutoFeatureExtractor.register(UpperCAmelCase__      ,				UpperCAmelCase__)
                  AutoTokenizer.register(UpperCAmelCase__      ,				slow_tokenizer_class=UpperCAmelCase__)
                  AutoProcessor.register(UpperCAmelCase__      ,				UpperCAmelCase__)
                  # Trying to register something existing in the Transformers library will raise an error
                  with self.assertRaises(UpperCAmelCase__):
                         AutoProcessor.register(UpperCAmelCase__      ,				UpperCAmelCase__)

                  # Now that the config is registered, it can be used as any other config with the auto-API
                  A__   				=					CustomFeatureExtractor.from_pretrained(UpperCAmelCase__)

                  with tempfile.TemporaryDirectory() as tmp_dir:
                         A__   				=					os.path.join(UpperCAmelCase__      ,				'''vocab.txt''')
                         with open(UpperCAmelCase__      ,				'''w'''      ,				encoding='''utf-8''') as vocab_writer:
                                vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]))
                         A__   				=					CustomTokenizer(UpperCAmelCase__)

                  A__   				=					CustomProcessor(UpperCAmelCase__      ,				UpperCAmelCase__)

                  with tempfile.TemporaryDirectory() as tmp_dir:
                         processor.save_pretrained(UpperCAmelCase__)
                         A__   				=					AutoProcessor.from_pretrained(UpperCAmelCase__)
                         self.assertIsInstance(UpperCAmelCase__      ,				UpperCAmelCase__)

           finally:
                  if "custom" in CONFIG_MAPPING._extra_content:
                         del CONFIG_MAPPING._extra_content["custom"]
                  if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
                         del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
                  if CustomConfig in TOKENIZER_MAPPING._extra_content:
                         del TOKENIZER_MAPPING._extra_content[CustomConfig]
                  if CustomConfig in PROCESSOR_MAPPING._extra_content:
                         del PROCESSOR_MAPPING._extra_content[CustomConfig]


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





           '''simple docstring'''

           class        UpperCamelCase_					(				UpperCAmelCase__  ):

               '''simple docstring'''






               UpperCAmelCase__													=   False


           class        UpperCamelCase_					(				UpperCAmelCase__  ):

               '''simple docstring'''






               UpperCAmelCase__													=   False


           class        UpperCamelCase_					(				UpperCAmelCase__  ):

               '''simple docstring'''






               UpperCAmelCase__													=   '''AutoFeatureExtractor'''
               UpperCAmelCase__													=   '''AutoTokenizer'''
               UpperCAmelCase__													=   False

           try:
                  AutoConfig.register('''custom'''      ,				UpperCAmelCase__)
                  AutoFeatureExtractor.register(UpperCAmelCase__      ,				UpperCAmelCase__)
                  AutoTokenizer.register(UpperCAmelCase__      ,				slow_tokenizer_class=UpperCAmelCase__)
                  AutoProcessor.register(UpperCAmelCase__      ,				UpperCAmelCase__)
                  # If remote code is not set, the default is to use local classes.
                  A__   				=					AutoProcessor.from_pretrained('''hf-internal-testing/test_dynamic_processor''')
                  self.assertEqual(processor.__class__.__name__      ,				'''NewProcessor''')
                  self.assertFalse(processor.special_attribute_present)
                  self.assertFalse(processor.feature_extractor.special_attribute_present)
                  self.assertFalse(processor.tokenizer.special_attribute_present)

                  # If remote code is disabled, we load the local ones.
                  A__   				=					AutoProcessor.from_pretrained(
                      '''hf-internal-testing/test_dynamic_processor'''      ,				trust_remote_code=UpperCAmelCase__)
                  self.assertEqual(processor.__class__.__name__      ,				'''NewProcessor''')
                  self.assertFalse(processor.special_attribute_present)
                  self.assertFalse(processor.feature_extractor.special_attribute_present)
                  self.assertFalse(processor.tokenizer.special_attribute_present)

                  # If remote is enabled, we load from the Hub.
                  A__   				=					AutoProcessor.from_pretrained(
                      '''hf-internal-testing/test_dynamic_processor'''      ,				trust_remote_code=UpperCAmelCase__)
                  self.assertEqual(processor.__class__.__name__      ,				'''NewProcessor''')
                  self.assertTrue(processor.special_attribute_present)
                  self.assertTrue(processor.feature_extractor.special_attribute_present)
                  self.assertTrue(processor.tokenizer.special_attribute_present)

           finally:
                  if "custom" in CONFIG_MAPPING._extra_content:
                         del CONFIG_MAPPING._extra_content["custom"]
                  if CustomConfig in FEATURE_EXTRACTOR_MAPPING._extra_content:
                         del FEATURE_EXTRACTOR_MAPPING._extra_content[CustomConfig]
                  if CustomConfig in TOKENIZER_MAPPING._extra_content:
                         del TOKENIZER_MAPPING._extra_content[CustomConfig]
                  if CustomConfig in PROCESSOR_MAPPING._extra_content:
                         del PROCESSOR_MAPPING._extra_content[CustomConfig]


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





           '''simple docstring'''

           A__   				=					AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-bert''')
           self.assertEqual(processor.__class__.__name__      ,				'''BertTokenizerFast''')





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





           '''simple docstring'''

           A__   				=					AutoProcessor.from_pretrained('''hf-internal-testing/tiny-random-convnext''')
           self.assertEqual(processor.__class__.__name__      ,				'''ConvNextImageProcessor''')


@is_staging_test
class        UpperCamelCase_					(				unittest.TestCase  ):

    '''simple docstring'''






    UpperCAmelCase__													=   ['''[UNK]''', '''[CLS]''', '''[SEP]''', '''[PAD]''', '''[MASK]''', '''bla''', '''blou''']


    @classmethod
    def  SCREAMING_SNAKE_CASE							(    cls   :			Tuple)		->Any:





           '''simple docstring'''

           A__   				=					TOKEN
           HfFolder.save_token(UpperCAmelCase__)


    @classmethod
    def  SCREAMING_SNAKE_CASE							(    cls   :			List[Any])		->List[str]:





           '''simple docstring'''

           try:
                  delete_repo(token=cls._token      ,				repo_id='''test-processor''')
           except HTTPError:
                  pass

           try:
                  delete_repo(token=cls._token      ,				repo_id='''valid_org/test-processor-org''')
           except HTTPError:
                  pass

           try:
                  delete_repo(token=cls._token      ,				repo_id='''test-dynamic-processor''')
           except HTTPError:
                  pass


    def  SCREAMING_SNAKE_CASE							(    self   :			str)		->Tuple:





           '''simple docstring'''

           A__   				=					WavaVecaProcessor.from_pretrained(UpperCAmelCase__)
           with tempfile.TemporaryDirectory() as tmp_dir:
                  processor.save_pretrained(
                      os.path.join(UpperCAmelCase__      ,				'''test-processor''')      ,				push_to_hub=UpperCAmelCase__      ,				use_auth_token=self._token)

                  A__   				=					WavaVecaProcessor.from_pretrained(f"""{USER}/test-processor""")
                  for k, v in processor.feature_extractor.__dict__.items():
                         self.assertEqual(UpperCAmelCase__      ,				getattr(new_processor.feature_extractor      ,				UpperCAmelCase__))
                  self.assertDictEqual(new_processor.tokenizer.get_vocab()      ,				processor.tokenizer.get_vocab())


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





           '''simple docstring'''

           A__   				=					WavaVecaProcessor.from_pretrained(UpperCAmelCase__)

           with tempfile.TemporaryDirectory() as tmp_dir:
                  processor.save_pretrained(
                      os.path.join(UpperCAmelCase__      ,				'''test-processor-org''')      ,				push_to_hub=UpperCAmelCase__      ,				use_auth_token=self._token      ,				organization='''valid_org'''      ,				)

                  A__   				=					WavaVecaProcessor.from_pretrained('''valid_org/test-processor-org''')
                  for k, v in processor.feature_extractor.__dict__.items():
                         self.assertEqual(UpperCAmelCase__      ,				getattr(new_processor.feature_extractor      ,				UpperCAmelCase__))
                  self.assertDictEqual(new_processor.tokenizer.get_vocab()      ,				processor.tokenizer.get_vocab())





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





           '''simple docstring'''

           CustomFeatureExtractor.register_for_auto_class()
           CustomTokenizer.register_for_auto_class()
           CustomProcessor.register_for_auto_class()

           A__   				=					CustomFeatureExtractor.from_pretrained(UpperCAmelCase__)

           with tempfile.TemporaryDirectory() as tmp_dir:
                  A__   				=					os.path.join(UpperCAmelCase__      ,				'''vocab.txt''')
                  with open(UpperCAmelCase__      ,				'''w'''      ,				encoding='''utf-8''') as vocab_writer:
                         vocab_writer.write(''''''.join([x + '''\n''' for x in self.vocab_tokens]))
                  A__   				=					CustomTokenizer(UpperCAmelCase__)

           A__   				=					CustomProcessor(UpperCAmelCase__      ,				UpperCAmelCase__)

           with tempfile.TemporaryDirectory() as tmp_dir:
                  create_repo(f"""{USER}/test-dynamic-processor"""      ,				token=self._token)
                  A__   				=					Repository(UpperCAmelCase__      ,				clone_from=f"""{USER}/test-dynamic-processor"""      ,				token=self._token)
                  processor.save_pretrained(UpperCAmelCase__)

                  # This has added the proper auto_map field to the feature extractor config
                  self.assertDictEqual(
                      processor.feature_extractor.auto_map      ,				{
                          '''AutoFeatureExtractor''': '''custom_feature_extraction.CustomFeatureExtractor''',
                          '''AutoProcessor''': '''custom_processing.CustomProcessor''',
                      }      ,				)

                  # This has added the proper auto_map field to the tokenizer config
                  with open(os.path.join(UpperCAmelCase__      ,				'''tokenizer_config.json''')) as f:
                         A__   				=					json.load(UpperCAmelCase__)
                  self.assertDictEqual(
                      tokenizer_config['''auto_map''']      ,				{
                          '''AutoTokenizer''': ['''custom_tokenization.CustomTokenizer''', None],
                          '''AutoProcessor''': '''custom_processing.CustomProcessor''',
                      }      ,				)

                  # The code has been copied from fixtures
                  self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase__      ,				'''custom_feature_extraction.py''')))
                  self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase__      ,				'''custom_tokenization.py''')))
                  self.assertTrue(os.path.isfile(os.path.join(UpperCAmelCase__      ,				'''custom_processing.py''')))

                  repo.push_to_hub()

           A__   				=					AutoProcessor.from_pretrained(f"""{USER}/test-dynamic-processor"""      ,				trust_remote_code=UpperCAmelCase__)
           # Can't make an isinstance check because the new_processor is from the CustomProcessor class of a dynamic module
           self.assertEqual(new_processor.__class__.__name__      ,				'''CustomProcessor''')








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())


