



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



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_      )  ->							Optional[int]:
           lowercase__		       =						jnp.ones((batch_size, length)      ) / length
           return scores



       def 							_a      (				self      )  ->							List[str]:
           lowercase__		       =						None
           lowercase__		       =						20

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

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

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

           lowercase__		       =						FlaxTemperatureLogitsWarper(temperature=0.5      )
           lowercase__		       =						FlaxTemperatureLogitsWarper(temperature=1.3      )

           lowercase__		       =						jax.nn.softmax(temp_dist_warper_sharper(UpperCAmelCase_     ,scores.copy()     ,cur_len=UpperCAmelCase_      )     ,axis=-1      )
           lowercase__		       =						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 							_a      (				self      )  ->							List[Any]:
           lowercase__		       =						None
           lowercase__		       =						10
           lowercase__		       =						2

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

           lowercase__		       =						FlaxTopKLogitsWarper(3      )

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

           lowercase__		       =						np.broadcast_to(np.arange(UpperCAmelCase_      )[None, :]     ,(batch_size, length)      ).copy()
           lowercase__		       =						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 							_a      (				self      )  ->							str:
           lowercase__		       =						None
           lowercase__		       =						10
           lowercase__		       =						2

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

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

           # make ramp_logits more extreme
           lowercase__		       =						ramp_logits[1] * 100.0

           # make sure at least 2 tokens are kept
           lowercase__		       =						FlaxTopPLogitsWarper(0.9     ,min_tokens_to_keep=2     ,filter_value=0.0      )
           lowercase__		       =						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 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						20
           lowercase__		       =						4
           lowercase__		       =						0

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

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



       def 							_a      (				self      )  ->							Optional[int]:
           lowercase__		       =						20
           lowercase__		       =						4
           lowercase__		       =						0

           lowercase__		       =						FlaxForcedBOSTokenLogitsProcessor(bos_token_id=UpperCAmelCase_      )

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



       def 							_a      (				self      )  ->							int:
           lowercase__		       =						20
           lowercase__		       =						4
           lowercase__		       =						0
           lowercase__		       =						5

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

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



       def 							_a      (				self      )  ->							Tuple:
           lowercase__		       =						4
           lowercase__		       =						10
           lowercase__		       =						15
           lowercase__		       =						2
           lowercase__		       =						1
           lowercase__		       =						15

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

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

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

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

           lowercase__		       =						10

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

           # with processor list
           lowercase__		       =						FlaxLogitsProcessorList(
               [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc]      )
           lowercase__		       =						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 							_a      (				self      )  ->							str:
           lowercase__		       =						4
           lowercase__		       =						10
           lowercase__		       =						15
           lowercase__		       =						2
           lowercase__		       =						1
           lowercase__		       =						15

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

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

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

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

           lowercase__		       =						10

           # no processor list
           def run_no_processor_list(UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      ):
               lowercase__		       =						temp_dist_warp(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               lowercase__		       =						top_k_warp(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               lowercase__		       =						top_p_warp(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               lowercase__		       =						min_dist_proc(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               lowercase__		       =						bos_dist_proc(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               lowercase__		       =						eos_dist_proc(UpperCAmelCase_     ,UpperCAmelCase_     ,cur_len=UpperCAmelCase_      )
               return scores

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

           lowercase__		       =						jax.jit(UpperCAmelCase_      )
           lowercase__		       =						jax.jit(UpperCAmelCase_      )

           lowercase__		       =						jitted_run_no_processor_list(UpperCAmelCase_     ,UpperCAmelCase_     ,UpperCAmelCase_      )
           lowercase__		       =						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()      )









'''simple docstring'''





from typing import List, Optional

from tokenizers import ByteLevelBPETokenizer

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_blenderbot_small import BlenderbotSmallTokenizer


SCREAMING_SNAKE_CASE__      					=     logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__      					=     {
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
    "tokenizer_config_file": "tokenizer_config.json",
}

SCREAMING_SNAKE_CASE__      					=     {
    "vocab_file": {
        "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/vocab.json"
    },
    "merges_file": {
        "facebook/blenderbot_small-90M": "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/merges.txt"
    },
    "tokenizer_config_file": {
        "facebook/blenderbot_small-90M": (
            "https://huggingface.co/facebook/blenderbot_small-90M/resolve/main/tokenizer_config.json"
        )
    },
}

SCREAMING_SNAKE_CASE__      					=     {
    "facebook/blenderbot_small-90M": 512,
}


class 				snake_case (UpperCamelCase   ):
       lowerCAmelCase__  :Union[str, Any]						   =					VOCAB_FILES_NAMES
       lowerCAmelCase__  :Any						   =					PRETRAINED_VOCAB_FILES_MAP
       lowerCAmelCase__  :str						   =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       lowerCAmelCase__  :List[Any]						   =					BlenderbotSmallTokenizer



       def __init__(				self     ,UpperCAmelCase_=None     ,UpperCAmelCase_=None     ,UpperCAmelCase_="<|endoftext|>"     ,UpperCAmelCase_="<|endoftext|>"     ,UpperCAmelCase_="<|endoftext|>"     ,UpperCAmelCase_=False     ,UpperCAmelCase_=True     ,**UpperCAmelCase_     ,)  ->							Optional[Any]:
           super().__init__(
               ByteLevelBPETokenizer(
                   vocab=UpperCAmelCase_     ,merges=UpperCAmelCase_     ,add_prefix_space=UpperCAmelCase_     ,trim_offsets=UpperCAmelCase_     ,)     ,bos_token=UpperCAmelCase_     ,eos_token=UpperCAmelCase_     ,unk_token=UpperCAmelCase_     ,**UpperCAmelCase_     ,)
           lowercase__		       =						add_prefix_space



       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_=None      )  ->							Optional[int]:
           lowercase__		       =						[self.bos_token_id] + token_ids_a + [self.eos_token_id]
           if token_ids_a is None:
               return output

           return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]







       def 							_a      (				self     ,UpperCAmelCase_     ,UpperCAmelCase_ = None      )  ->							List[int]:
           lowercase__		       =						[self.sep_token_id]
           lowercase__		       =						[self.cls_token_id]

           if token_ids_a is None:
               return len(cls + token_ids_a + sep      ) * [0]
           return len(cls + token_ids_a + sep + sep + token_ids_a + sep      ) * [0]




