



import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__UpperCAmelCase	       =	logging.get_logger(__name__)

__UpperCAmelCase	       =	'''▁'''

__UpperCAmelCase	       =	{'''vocab_file''': '''sentencepiece.bpe.model'''}

__UpperCAmelCase	       =	{
    '''vocab_file''': {
        '''facebook/mbart-large-en-ro''': (
            '''https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model'''
        ),
        '''facebook/mbart-large-cc25''': (
            '''https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model'''
        ),
    }
}

__UpperCAmelCase	       =	{
    '''facebook/mbart-large-en-ro''': 1_024,
    '''facebook/mbart-large-cc25''': 1_024,
}

# fmt: off
__UpperCAmelCase	       =	['''ar_AR''', '''cs_CZ''', '''de_DE''', '''en_XX''', '''es_XX''', '''et_EE''', '''fi_FI''', '''fr_XX''', '''gu_IN''', '''hi_IN''', '''it_IT''', '''ja_XX''', '''kk_KZ''', '''ko_KR''', '''lt_LT''', '''lv_LV''', '''my_MM''', '''ne_NP''', '''nl_XX''', '''ro_RO''', '''ru_RU''', '''si_LK''', '''tr_TR''', '''vi_VN''', '''zh_CN''']



class     lowerCAmelCase_		(  a__							):
       UpperCAmelCase__    : Tuple       					=  VOCAB_FILES_NAMES
       UpperCAmelCase__    : List[Any]       					=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       UpperCAmelCase__    : List[str]       					=  PRETRAINED_VOCAB_FILES_MAP
       UpperCAmelCase__    : Optional[int]       					=  ["input_ids", "attention_mask"]

       UpperCAmelCase__    : List[int]       					=  []
       UpperCAmelCase__    : List[int]       					=  []
       def __init__(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_="<s>",      SCREAMING_SNAKE_CASE_="</s>",      SCREAMING_SNAKE_CASE_="</s>",      SCREAMING_SNAKE_CASE_="<s>",      SCREAMING_SNAKE_CASE_="<unk>",      SCREAMING_SNAKE_CASE_="<pad>",      SCREAMING_SNAKE_CASE_="<mask>",      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_=None,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_=None,      **SCREAMING_SNAKE_CASE_,      )					->		Optional[Any]:
          # Mask token behave like a normal word, i.e. include the space before it
          UpperCamelCase   :  List[Any]						     =     AddedToken(SCREAMING_SNAKE_CASE_,      lstrip=SCREAMING_SNAKE_CASE_,      rstrip=SCREAMING_SNAKE_CASE_	) if isinstance(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	) else mask_token

          UpperCamelCase   :  List[Any]						     =     {} if sp_model_kwargs is None else sp_model_kwargs

          super().__init__(
              bos_token=SCREAMING_SNAKE_CASE_,      eos_token=SCREAMING_SNAKE_CASE_,      unk_token=SCREAMING_SNAKE_CASE_,      sep_token=SCREAMING_SNAKE_CASE_,      cls_token=SCREAMING_SNAKE_CASE_,      pad_token=SCREAMING_SNAKE_CASE_,      mask_token=SCREAMING_SNAKE_CASE_,      tokenizer_file=SCREAMING_SNAKE_CASE_,      src_lang=SCREAMING_SNAKE_CASE_,      tgt_lang=SCREAMING_SNAKE_CASE_,      additional_special_tokens=SCREAMING_SNAKE_CASE_,      sp_model_kwargs=self.sp_model_kwargs,      **SCREAMING_SNAKE_CASE_,      )

          UpperCamelCase   :  Tuple						     =     spm.SentencePieceProcessor(**self.sp_model_kwargs	)
          self.sp_model.Load(str(SCREAMING_SNAKE_CASE_	)	)
          UpperCamelCase   :  int						     =     vocab_file

          # Original fairseq vocab and spm vocab must be "aligned":
          # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
          # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
          # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
          # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

          # Mimic fairseq token-to-id alignment for the first 4 token
          UpperCamelCase   :  Any						     =     {'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

          # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
          UpperCamelCase   :  List[Any]						     =     1

          UpperCamelCase   :  Optional[int]						     =     len(self.sp_model	)
          UpperCamelCase   :  int						     =     {
              code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(SCREAMING_SNAKE_CASE_	)
          }
          UpperCamelCase   :  List[Any]						     =     {v: k for k, v in self.lang_code_to_id.items()}
          UpperCamelCase   :  Tuple						     =     len(self.sp_model	) + len(self.lang_code_to_id	) + self.fairseq_offset

          self.fairseq_tokens_to_ids.update(self.lang_code_to_id	)
          UpperCamelCase   :  Union[str, Any]						     =     {v: k for k, v in self.fairseq_tokens_to_ids.items()}
          UpperCamelCase   :  str						     =     list(self.lang_code_to_id.keys()	)

          if additional_special_tokens is not None:
             # Only add those special tokens if they are not already there.
             self._additional_special_tokens.extend(
                 [t for t in additional_special_tokens if t not in self._additional_special_tokens]	)

          UpperCamelCase   :  Union[str, Any]						     =     src_lang if src_lang is not None else 'en_XX'
          UpperCamelCase   :  Union[str, Any]						     =     self.lang_code_to_id[self._src_lang]
          UpperCamelCase   :  List[str]						     =     tgt_lang
          self.set_src_lang_special_tokens(self._src_lang	)
       def __getstate__(      self	)					->		List[str]:
          UpperCamelCase   :  Union[str, Any]						     =     self.__dict__.copy()
          UpperCamelCase   :  Tuple						     =     None
          UpperCamelCase   :  List[Any]						     =     self.sp_model.serialized_model_proto()
          return state
       def __setstate__(      self,      SCREAMING_SNAKE_CASE_	)					->		Optional[int]:
          UpperCamelCase   :  Any						     =     d

          # for backward compatibility
          if not hasattr(self,      'sp_model_kwargs'	):
             UpperCamelCase   :  Optional[Any]						     =     {}

          UpperCamelCase   :  Optional[int]						     =     spm.SentencePieceProcessor(**self.sp_model_kwargs	)
          self.sp_model.LoadFromSerializedProto(self.sp_model_proto	)
       @property
       def  snake_case_	(      self	)					->		Optional[int]:
          return len(self.sp_model	) + len(self.lang_code_to_id	) + self.fairseq_offset + 1  # Plus 1 for the mask token
       @property
       def  snake_case_	(      self	)					->		str:
          return self._src_lang
       @src_lang.setter
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		None:
          UpperCamelCase   :  Union[str, Any]						     =     new_src_lang
          self.set_src_lang_special_tokens(self._src_lang	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = False	)					->		List[int]:

          if already_has_special_tokens:
             return super().get_special_tokens_mask(
                 token_ids_a=SCREAMING_SNAKE_CASE_,      token_ids_a=SCREAMING_SNAKE_CASE_,      already_has_special_tokens=SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  Any						     =     [1] * len(self.prefix_tokens	)
          UpperCamelCase   :  List[Any]						     =     [1] * len(self.suffix_tokens	)
          if token_ids_a is None:
             return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE_	)) + suffix_ones
          return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE_	)) + ([0] * len(SCREAMING_SNAKE_CASE_	)) + suffix_ones
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		List[int]:
          if token_ids_a is None:
             return self.prefix_tokens + token_ids_a + self.suffix_tokens
          # We don't expect to process pairs, but leave the pair logic for API consistency
          return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		List[int]:
          UpperCamelCase   :  List[str]						     =     [self.sep_token_id]
          UpperCamelCase   :  Tuple						     =     [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]
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_	)					->		Dict:
          if src_lang is None or tgt_lang is None:
             raise ValueError('Translation requires a `src_lang` and a `tgt_lang` for this model'	)
          UpperCamelCase   :  Optional[Any]						     =     src_lang
          UpperCamelCase   :  Any						     =     self(SCREAMING_SNAKE_CASE_,      add_special_tokens=SCREAMING_SNAKE_CASE_,      return_tensors=SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Tuple						     =     self.convert_tokens_to_ids(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Tuple						     =     tgt_lang_id
          return inputs
       def  snake_case_	(      self	)					->		Dict:
          UpperCamelCase   :  Union[str, Any]						     =     {self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE_	): i for i in range(self.vocab_size	)}
          vocab.update(self.added_tokens_encoder	)
          return vocab
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		List[str]:
          return self.sp_model.encode(SCREAMING_SNAKE_CASE_,      out_type=SCREAMING_SNAKE_CASE_	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		List[Any]:
          if token in self.fairseq_tokens_to_ids:
             return self.fairseq_tokens_to_ids[token]
          UpperCamelCase   :  List[Any]						     =     self.sp_model.PieceToId(SCREAMING_SNAKE_CASE_	)

          # Need to return unknown token if the SP model returned 0
          return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Union[str, Any]:
          if index in self.fairseq_ids_to_tokens:
             return self.fairseq_ids_to_tokens[index]
          return self.sp_model.IdToPiece(index - self.fairseq_offset	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		Dict:
          UpperCamelCase   :  Optional[int]						     =     ''.join(SCREAMING_SNAKE_CASE_	).replace(SCREAMING_SNAKE_CASE_,      ' '	).strip()
          return out_string
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = None	)					->		Tuple[str]:
          if not os.path.isdir(SCREAMING_SNAKE_CASE_	):
             logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""	)
             return
          UpperCamelCase   :  Union[str, Any]						     =     os.path.join(
              SCREAMING_SNAKE_CASE_,      (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']	)

          if os.path.abspath(self.vocab_file	) != os.path.abspath(SCREAMING_SNAKE_CASE_	) and os.path.isfile(self.vocab_file	):
             copyfile(self.vocab_file,      SCREAMING_SNAKE_CASE_	)
          elif not os.path.isfile(self.vocab_file	):
             with open(SCREAMING_SNAKE_CASE_,      'wb'	) as fi:
                UpperCamelCase   :  List[Any]						     =     self.sp_model.serialized_model_proto()
                fi.write(SCREAMING_SNAKE_CASE_	)

          return (out_vocab_file,)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_ = "en_XX",      SCREAMING_SNAKE_CASE_ = None,      SCREAMING_SNAKE_CASE_ = "ro_RO",      **SCREAMING_SNAKE_CASE_,      )					->		BatchEncoding:
          UpperCamelCase   :  Optional[Any]						     =     src_lang
          UpperCamelCase   :  Any						     =     tgt_lang
          return super().prepare_seqaseq_batch(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      **SCREAMING_SNAKE_CASE_	)
       def  snake_case_	(      self	)					->		Tuple:
          return self.set_src_lang_special_tokens(self.src_lang	)
       def  snake_case_	(      self	)					->		str:
          return self.set_tgt_lang_special_tokens(self.tgt_lang	)
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		None:
          UpperCamelCase   :  Union[str, Any]						     =     self.lang_code_to_id[src_lang]
          UpperCamelCase   :  Optional[Any]						     =     []
          UpperCamelCase   :  List[Any]						     =     [self.eos_token_id, self.cur_lang_code]






       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_	)					->		None:
          UpperCamelCase   :  Dict						     =     self.lang_code_to_id[lang]
          UpperCamelCase   :  Optional[Any]						     =     []
          UpperCamelCase   :  Any						     =     [self.eos_token_id, self.cur_lang_code]








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     lowerCAmelCase_		(  unittest.TestCase							):
       def  snake_case_	(      self,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)					->		Union[str, Any]:
          UpperCamelCase   :  List[str]						     =     jnp.ones((batch_size, length)	) / length
          return scores
       def  snake_case_	(      self	)					->		Optional[Any]:
          UpperCamelCase   :  Optional[Any]						     =     None
          UpperCamelCase   :  Optional[int]						     =     20

          UpperCamelCase   :  Optional[Any]						     =     self._get_uniform_logits(batch_size=2,      length=SCREAMING_SNAKE_CASE_	)

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

          # compute softmax
          UpperCamelCase   :  List[str]						     =     jax.nn.softmax(SCREAMING_SNAKE_CASE_,      axis=-1	)

          UpperCamelCase   :  List[Any]						     =     FlaxTemperatureLogitsWarper(temperature=0.5	)
          UpperCamelCase   :  int						     =     FlaxTemperatureLogitsWarper(temperature=1.3	)

          UpperCamelCase   :  Tuple						     =     jax.nn.softmax(temp_dist_warper_sharper(SCREAMING_SNAKE_CASE_,      scores.copy(),      cur_len=SCREAMING_SNAKE_CASE_	),      axis=-1	)
          UpperCamelCase   :  Any						     =     jax.nn.softmax(temp_dist_warper_smoother(SCREAMING_SNAKE_CASE_,      scores.copy(),      cur_len=SCREAMING_SNAKE_CASE_	),      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  snake_case_	(      self	)					->		Optional[Any]:
          UpperCamelCase   :  Dict						     =     None
          UpperCamelCase   :  Any						     =     10
          UpperCamelCase   :  Any						     =     2

          # create ramp distribution
          UpperCamelCase   :  List[Any]						     =     np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_	)[None, :],      (batch_size, vocab_size)	).copy()
          UpperCamelCase   :  Tuple						     =     ramp_logits[1:, : vocab_size // 2] + vocab_size

          UpperCamelCase   :  Dict						     =     FlaxTopKLogitsWarper(3	)

          UpperCamelCase   :  Tuple						     =     top_k_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)

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

          UpperCamelCase   :  Union[str, Any]						     =     np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_	)[None, :],      (batch_size, length)	).copy()
          UpperCamelCase   :  List[str]						     =     top_k_warp_safety_check(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)

          # min_tokens overwrites k: 3 tokens are kept => 2 tokens are nullified
          self.assertListEqual((scores == 0.0).sum(axis=-1	).tolist(),      [2, 2]	)
       def  snake_case_	(      self	)					->		Union[str, Any]:
          UpperCamelCase   :  int						     =     None
          UpperCamelCase   :  List[str]						     =     10
          UpperCamelCase   :  Optional[Any]						     =     2

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

          UpperCamelCase   :  Optional[Any]						     =     FlaxTopPLogitsWarper(0.8	)
          UpperCamelCase   :  int						     =     np.exp(top_p_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)	)

          # dist should be filtered to keep min num values so that sum is >= top_p
          # exp (-inf) => 0
          UpperCamelCase   :  Any						     =     np.array([[0.3, 0.0, 0.0, 0.5], [0.0, 0.3, 0.3, 0.25]]	)
          self.assertTrue(np.allclose(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      atol=1e-3	)	)

          # check edge cases with negative and extreme logits
          UpperCamelCase   :  Optional[Any]						     =     np.broadcast_to(np.arange(SCREAMING_SNAKE_CASE_	)[None, :],      (batch_size, vocab_size)	).copy() - (
              vocab_size // 2
          )

          # make ramp_logits more extreme
          UpperCamelCase   :  Tuple						     =     ramp_logits[1] * 1_00.0

          # make sure at least 2 tokens are kept
          UpperCamelCase   :  int						     =     FlaxTopPLogitsWarper(0.9,      min_tokens_to_keep=2,      filter_value=0.0	)
          UpperCamelCase   :  List[str]						     =     top_p_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)

          # 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  snake_case_	(      self	)					->		List[Any]:
          UpperCamelCase   :  Union[str, Any]						     =     20
          UpperCamelCase   :  Union[str, Any]						     =     4
          UpperCamelCase   :  Optional[int]						     =     0

          UpperCamelCase   :  Dict						     =     FlaxMinLengthLogitsProcessor(min_length=10,      eos_token_id=SCREAMING_SNAKE_CASE_	)

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

          UpperCamelCase   :  int						     =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_	)

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

          UpperCamelCase   :  Tuple						     =     FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_,      eos_token_id=SCREAMING_SNAKE_CASE_	)

          # 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   :  Tuple						     =     4
          UpperCamelCase   :  Union[str, Any]						     =     self._get_uniform_logits(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Any						     =     logits_processor(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          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   :  str						     =     3
          UpperCamelCase   :  List[Any]						     =     self._get_uniform_logits(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Union[str, Any]						     =     logits_processor(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          self.assertFalse(jnp.isinf(SCREAMING_SNAKE_CASE_	).any()	)
       def  snake_case_	(      self	)					->		int:
          UpperCamelCase   :  int						     =     4
          UpperCamelCase   :  Tuple						     =     10
          UpperCamelCase   :  str						     =     15
          UpperCamelCase   :  List[str]						     =     2
          UpperCamelCase   :  Any						     =     1
          UpperCamelCase   :  List[str]						     =     15

          # dummy input_ids and scores
          UpperCamelCase   :  Dict						     =     ids_tensor((batch_size, sequence_length),      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[Any]						     =     input_ids.copy()

          UpperCamelCase   :  Any						     =     self._get_uniform_logits(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Any						     =     scores.copy()

          # instantiate all dist processors
          UpperCamelCase   :  List[str]						     =     FlaxTemperatureLogitsWarper(temperature=0.5	)
          UpperCamelCase   :  Optional[Any]						     =     FlaxTopKLogitsWarper(3	)
          UpperCamelCase   :  Optional[int]						     =     FlaxTopPLogitsWarper(0.8	)

          # instantiate all logits processors
          UpperCamelCase   :  Optional[Any]						     =     FlaxMinLengthLogitsProcessor(min_length=10,      eos_token_id=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[Any]						     =     FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_,      eos_token_id=SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  List[str]						     =     10

          # no processor list
          UpperCamelCase   :  Any						     =     temp_dist_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     top_k_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     top_p_warp(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Dict						     =     min_dist_proc(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Union[str, Any]						     =     bos_dist_proc(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[str]						     =     eos_dist_proc(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)

          # with processor list
          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(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)

          # scores should be equal
          self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      atol=1e-3	)	)

          # input_ids should never be changed
          self.assertListEqual(input_ids.tolist(),      input_ids_comp.tolist()	)






       def  snake_case_	(      self	)					->		int:
          UpperCamelCase   :  Optional[Any]						     =     4
          UpperCamelCase   :  Tuple						     =     10
          UpperCamelCase   :  Union[str, Any]						     =     15
          UpperCamelCase   :  Union[str, Any]						     =     2
          UpperCamelCase   :  Optional[Any]						     =     1
          UpperCamelCase   :  int						     =     15

          # dummy input_ids and scores
          UpperCamelCase   :  Dict						     =     ids_tensor((batch_size, sequence_length),      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  str						     =     input_ids.copy()

          UpperCamelCase   :  Optional[int]						     =     self._get_uniform_logits(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Union[str, Any]						     =     scores.copy()

          # instantiate all dist processors
          UpperCamelCase   :  Dict						     =     FlaxTemperatureLogitsWarper(temperature=0.5	)
          UpperCamelCase   :  Optional[Any]						     =     FlaxTopKLogitsWarper(3	)
          UpperCamelCase   :  Union[str, Any]						     =     FlaxTopPLogitsWarper(0.8	)

          # instantiate all logits processors
          UpperCamelCase   :  str						     =     FlaxMinLengthLogitsProcessor(min_length=10,      eos_token_id=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Any						     =     FlaxForcedBOSTokenLogitsProcessor(bos_token_id=SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  Optional[int]						     =     FlaxForcedEOSTokenLogitsProcessor(max_length=SCREAMING_SNAKE_CASE_,      eos_token_id=SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  Any						     =     10

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

          # with processor list
          def run_processor_list(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	):
             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(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      cur_len=SCREAMING_SNAKE_CASE_	)
             return scores

          UpperCamelCase   :  Dict						     =     jax.jit(SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  List[str]						     =     jax.jit(SCREAMING_SNAKE_CASE_	)

          UpperCamelCase   :  int						     =     jitted_run_no_processor_list(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)
          UpperCamelCase   :  int						     =     jitted_run_processor_list(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_	)

          # scores should be equal
          self.assertTrue(jnp.allclose(SCREAMING_SNAKE_CASE_,      SCREAMING_SNAKE_CASE_,      atol=1e-3	)	)

          # input_ids should never be changed
          self.assertListEqual(input_ids.tolist(),      input_ids_comp.tolist()	)



