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

										UpperCamelCase__     :		int	 =					self._get_uniform_logits(batch_size=2  ,    length=_UpperCamelCase)

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

										# compute softmax
										UpperCamelCase__     :		Tuple	 =					jax.nn.softmax(_UpperCamelCase  ,    axis=-1)

										UpperCamelCase__     :		Optional[Any]	 =					FlaxTemperatureLogitsWarper(temperature=0.5)
										UpperCamelCase__     :		List[str]	 =					FlaxTemperatureLogitsWarper(temperature=1.3)

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

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

										UpperCamelCase__     :		List[Any]	 =					FlaxTopKLogitsWarper(3)

										UpperCamelCase__     :		List[Any]	 =					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
										UpperCamelCase__     :		List[Any]	 =					5
										UpperCamelCase__     :		Dict	 =					FlaxTopKLogitsWarper(top_k=1  ,    filter_value=0.0  ,    min_tokens_to_keep=3)

										UpperCamelCase__     :		int	 =					np.broadcast_to(np.arange(_UpperCamelCase)[None, :]  ,    (batch_size, length)).copy()
										UpperCamelCase__     :		Tuple	 =					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 							lowerCAmelCase__			(					self)			->						int:
										UpperCamelCase__     :		str	 =					None
										UpperCamelCase__     :		int	 =					10
										UpperCamelCase__     :		List[Any]	 =					2

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

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

										# make ramp_logits more extreme
										UpperCamelCase__     :		Any	 =					ramp_logits[1] * 100.0

										# make sure at least 2 tokens are kept
										UpperCamelCase__     :		Union[str, Any]	 =					FlaxTopPLogitsWarper(0.9  ,    min_tokens_to_keep=2  ,    filter_value=0.0)
										UpperCamelCase__     :		Union[str, Any]	 =					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 							lowerCAmelCase__			(					self)			->						Union[str, Any]:
										UpperCamelCase__     :		Dict	 =					20
										UpperCamelCase__     :		Union[str, Any]	 =					4
										UpperCamelCase__     :		Any	 =					0

										UpperCamelCase__     :		List[str]	 =					FlaxMinLengthLogitsProcessor(min_length=10  ,    eos_token_id=_UpperCamelCase)

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

										UpperCamelCase__     :		Tuple	 =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCamelCase)

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

										UpperCamelCase__     :		List[Any]	 =					FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCamelCase  ,    eos_token_id=_UpperCamelCase)

										# check that all scores are -inf except the eos_token_id when max_length is reached
										UpperCamelCase__     :		Dict	 =					ids_tensor((batch_size, 4)  ,    vocab_size=20)
										UpperCamelCase__     :		Any	 =					4
										UpperCamelCase__     :		List[str]	 =					self._get_uniform_logits(_UpperCamelCase  ,    _UpperCamelCase)
										UpperCamelCase__     :		List[Any]	 =					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
										UpperCamelCase__     :		List[Any]	 =					3
										UpperCamelCase__     :		List[Any]	 =					self._get_uniform_logits(_UpperCamelCase  ,    _UpperCamelCase)
										UpperCamelCase__     :		List[Any]	 =					logits_processor(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										self.assertFalse(jnp.isinf(_UpperCamelCase).any())
							def 							lowerCAmelCase__			(					self)			->						Optional[Any]:
										UpperCamelCase__     :		List[str]	 =					4
										UpperCamelCase__     :		List[Any]	 =					10
										UpperCamelCase__     :		List[str]	 =					15
										UpperCamelCase__     :		Tuple	 =					2
										UpperCamelCase__     :		Any	 =					1
										UpperCamelCase__     :		Tuple	 =					15

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

										UpperCamelCase__     :		Optional[Any]	 =					self._get_uniform_logits(_UpperCamelCase  ,    _UpperCamelCase)
										UpperCamelCase__     :		Any	 =					scores.copy()

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

										# instantiate all logits processors
										UpperCamelCase__     :		List[Any]	 =					FlaxMinLengthLogitsProcessor(min_length=10  ,    eos_token_id=_UpperCamelCase)
										UpperCamelCase__     :		List[Any]	 =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCamelCase)
										UpperCamelCase__     :		List[str]	 =					FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCamelCase  ,    eos_token_id=_UpperCamelCase)

										UpperCamelCase__     :		Any	 =					10

										# no processor list
										UpperCamelCase__     :		Dict	 =					temp_dist_warp(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										UpperCamelCase__     :		str	 =					top_k_warp(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										UpperCamelCase__     :		int	 =					top_p_warp(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										UpperCamelCase__     :		Optional[int]	 =					min_dist_proc(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										UpperCamelCase__     :		str	 =					bos_dist_proc(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
										UpperCamelCase__     :		Dict	 =					eos_dist_proc(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)

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

										# dummy input_ids and scores
										UpperCamelCase__     :		Optional[int]	 =					ids_tensor((batch_size, sequence_length)  ,    _UpperCamelCase)
										UpperCamelCase__     :		str	 =					input_ids.copy()

										UpperCamelCase__     :		List[Any]	 =					self._get_uniform_logits(_UpperCamelCase  ,    _UpperCamelCase)
										UpperCamelCase__     :		List[str]	 =					scores.copy()

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

										# instantiate all logits processors
										UpperCamelCase__     :		Optional[int]	 =					FlaxMinLengthLogitsProcessor(min_length=10  ,    eos_token_id=_UpperCamelCase)
										UpperCamelCase__     :		str	 =					FlaxForcedBOSTokenLogitsProcessor(bos_token_id=_UpperCamelCase)
										UpperCamelCase__     :		Dict	 =					FlaxForcedEOSTokenLogitsProcessor(max_length=_UpperCamelCase  ,    eos_token_id=_UpperCamelCase)

										UpperCamelCase__     :		List[str]	 =					10

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

										# with processor list
										def run_processor_list(UpperCamelCase  ,    UpperCamelCase  ,    UpperCamelCase):
													UpperCamelCase__     :		Union[str, Any]	 =					FlaxLogitsProcessorList(
													    [temp_dist_warp, top_k_warp, top_p_warp, min_dist_proc, bos_dist_proc, eos_dist_proc])
													UpperCamelCase__     :		int	 =					processor(_UpperCamelCase  ,    _UpperCamelCase  ,    cur_len=_UpperCamelCase)
													return scores

										UpperCamelCase__     :		Union[str, Any]	 =					jax.jit(_UpperCamelCase)
										UpperCamelCase__     :		List[Any]	 =					jax.jit(_UpperCamelCase)

										UpperCamelCase__     :		Any	 =					jitted_run_no_processor_list(_UpperCamelCase  ,    _UpperCamelCase  ,    _UpperCamelCase)
										UpperCamelCase__     :		Union[str, Any]	 =					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())






import json
import os
from pathlib import Path
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple, Union

import sentencepiece

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


UpperCamelCase__		 =     logging.get_logger(__name__)

UpperCamelCase__		 =     '''▁'''

UpperCamelCase__		 =     {
    '''vocab_file''': '''vocab.json''',
    '''spm_file''': '''sentencepiece.bpe.model''',
}

UpperCamelCase__		 =     {
    '''vocab_file''': {
        '''facebook/s2t-small-librispeech-asr''': (
            '''https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/vocab.json'''
        ),
    },
    '''spm_file''': {
        '''facebook/s2t-small-librispeech-asr''': (
            '''https://huggingface.co/facebook/s2t-small-librispeech-asr/resolve/main/sentencepiece.bpe.model'''
        )
    },
}

UpperCamelCase__		 =     {
    '''facebook/s2t-small-librispeech-asr''': 1024,
}

UpperCamelCase__		 =     ['''pt''', '''fr''', '''ru''', '''nl''', '''ro''', '''it''', '''es''', '''de''']

UpperCamelCase__		 =     {'''mustc''': MUSTC_LANGS}

class     __snake_case    (				snake_case__  ):
       """simple docstring"""







       __SCREAMING_SNAKE_CASE				    =							VOCAB_FILES_NAMES
       __SCREAMING_SNAKE_CASE				    =							PRETRAINED_VOCAB_FILES_MAP
       __SCREAMING_SNAKE_CASE				    =							MAX_MODEL_INPUT_SIZES
       __SCREAMING_SNAKE_CASE				    =							["input_ids", "attention_mask"]

       __SCREAMING_SNAKE_CASE				    =							[]


       def __init__(		self      , _UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase="<s>"      , _UpperCamelCase="</s>"      , _UpperCamelCase="<pad>"      , _UpperCamelCase="<unk>"      , _UpperCamelCase=False      , _UpperCamelCase=False      , _UpperCamelCase=None      , _UpperCamelCase=None      , _UpperCamelCase = None      , **_UpperCamelCase      , )      ->							None:






              """simple docstring"""


              __snake_case       	=							{} if sp_model_kwargs is None else sp_model_kwargs

              super().__init__(
                  bos_token=_UpperCamelCase      , eos_token=_UpperCamelCase      , unk_token=_UpperCamelCase      , pad_token=_UpperCamelCase      , do_upper_case=_UpperCamelCase      , do_lower_case=_UpperCamelCase      , tgt_lang=_UpperCamelCase      , lang_codes=_UpperCamelCase      , sp_model_kwargs=self.sp_model_kwargs      , **_UpperCamelCase      , )
              __snake_case       	=							do_upper_case
              __snake_case       	=							do_lower_case

              __snake_case       	=							load_json(_UpperCamelCase   )
              __snake_case       	=							{v: k for k, v in self.encoder.items()}
              __snake_case       	=							spm_file
              __snake_case       	=							load_spm(_UpperCamelCase      , self.sp_model_kwargs   )

              if lang_codes is not None:
                     __snake_case       	=							lang_codes
                     __snake_case       	=							LANGUAGES[lang_codes]
                     __snake_case       	=							[F'<lang:{lang}>' for lang in self.langs]
                     __snake_case       	=							{lang: self.sp_model.PieceToId(F'<lang:{lang}>'   ) for lang in self.langs}

                     __snake_case       	=							self.lang_tokens
                     __snake_case       	=							tgt_lang if tgt_lang is not None else self.langs[0]

                     self.set_tgt_lang_special_tokens(self._tgt_lang   )
              else:
                     __snake_case       	=							{}


       @property
       def 			a		(		self   )      ->							int:






              """simple docstring"""


              return len(self.encoder   )


       @property
       def 			a		(		self   )      ->							str:






              """simple docstring"""


              return self._tgt_lang


       @tgt_lang.setter
       def 			a		(		self      , _UpperCamelCase   )      ->							None:






              """simple docstring"""


              __snake_case       	=							new_tgt_lang
              self.set_tgt_lang_special_tokens(_UpperCamelCase   )


       def 			a		(		self      , _UpperCamelCase   )      ->							None:






              """simple docstring"""


              __snake_case       	=							self.lang_code_to_id[tgt_lang]
              __snake_case       	=							[lang_code_id]


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






              """simple docstring"""


              return self.sp_model.encode(_UpperCamelCase      , out_type=_UpperCamelCase   )


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






              """simple docstring"""


              return self.encoder.get(_UpperCamelCase      , self.encoder[self.unk_token]   )


       def 			a		(		self      , _UpperCamelCase   )      ->							str:






              """simple docstring"""


              return self.decoder.get(_UpperCamelCase      , self.unk_token   )


       def 			a		(		self      , _UpperCamelCase   )      ->							str:






              """simple docstring"""


              __snake_case       	=							[]
              __snake_case       	=							""""""
              for token in tokens:
                     # make sure that special tokens are not decoded using sentencepiece model
                     if token in self.all_special_tokens:
                            __snake_case       	=							self.sp_model.decode(_UpperCamelCase   )
                            out_string += (decoded.upper() if self.do_upper_case else decoded) + token + " "
                            __snake_case       	=							[]
                     else:
                            current_sub_tokens.append(_UpperCamelCase   )
              __snake_case       	=							self.sp_model.decode(_UpperCamelCase   )
              out_string += decoded.upper() if self.do_upper_case else decoded
              return out_string.strip()


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase=None   )      ->							List[int]:






              """simple docstring"""


              if token_ids_a is None:
                     return self.prefix_tokens + token_ids_a + [self.eos_token_id]
              # 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.eos_token_id]


       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None      , _UpperCamelCase = False   )      ->							List[int]:






              """simple docstring"""


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

              __snake_case       	=							[1] * len(self.prefix_tokens   )
              __snake_case       	=							[1]
              if token_ids_a is None:
                     return prefix_ones + ([0] * len(_UpperCamelCase   )) + suffix_ones
              return prefix_ones + ([0] * len(_UpperCamelCase   )) + ([0] * len(_UpperCamelCase   )) + suffix_ones


       def 			a		(		self   )      ->							Dict:






              """simple docstring"""


              __snake_case       	=							self.encoder.copy()
              vocab.update(self.added_tokens_encoder   )
              return vocab


       def __getstate__(		self   )      ->							Dict:






              """simple docstring"""


              __snake_case       	=							self.__dict__.copy()
              __snake_case       	=							None
              return state


       def __setstate__(		self      , _UpperCamelCase   )      ->							None:






              """simple docstring"""


              __snake_case       	=							d

              # for backward compatibility
              if not hasattr(self      , """sp_model_kwargs"""   ):
                     __snake_case       	=							{}

              __snake_case       	=							load_spm(self.spm_file      , self.sp_model_kwargs   )






       def 			a		(		self      , _UpperCamelCase      , _UpperCamelCase = None   )      ->							Tuple[str]:






              """simple docstring"""


              __snake_case       	=							Path(_UpperCamelCase   )
              assert save_dir.is_dir(), F'{save_directory} should be a directory'
              __snake_case       	=							save_dir / (
                  (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
              )
              __snake_case       	=							save_dir / (
                  (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
              )

              save_json(self.encoder      , _UpperCamelCase   )

              if os.path.abspath(self.spm_file   ) != os.path.abspath(_UpperCamelCase   ) and os.path.isfile(self.spm_file   ):
                     copyfile(self.spm_file      , _UpperCamelCase   )
              elif not os.path.isfile(self.spm_file   ):
                     with open(_UpperCamelCase      , """wb"""   ) as fi:
                            __snake_case       	=							self.sp_model.serialized_model_proto()
                            fi.write(_UpperCamelCase   )

              return (str(_UpperCamelCase   ), str(_UpperCamelCase   ))







def 						lowerCamelCase__      ( __A    :str	,__A    :Dict[str, Any]		):


       """simple docstring"""




       __snake_case       	=							sentencepiece.SentencePieceProcessor(**__A		)
       spm.Load(str(__A		)		)
       return spm







def 						lowerCamelCase__      ( __A    :str		):


       """simple docstring"""




       with open(__A	,"""r"""		) as f:
              return json.load(__A		)







def 						lowerCamelCase__      ( __A    :Optional[Any]	,__A    :str		):


       """simple docstring"""




       with open(__A	,"""w"""		) as f:
              json.dump(__A	,__A	,indent=2		)

