



import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


lowercase_													=      get_logger(__name__)


lowercase_													=      r'''
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

'''






class     __a			:



	@add_start_docstrings(snake_case_)
	def __call__(    self				:     List[Any]					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray)->				jnp.ndarray:
						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")






class     __a			:



	@add_start_docstrings(snake_case_)
	def __call__(    self				:     Optional[int]					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray)->				jnp.ndarray:
						raise NotImplementedError(
						    F"""{self.__class__} is an abstract class. Only classes inheriting this class can be called.""")






class     __a			(  SCREAMING_SNAKE_CASE					):



	@add_start_docstrings(snake_case_)
	def __call__(    self				:     int					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int					,						**snake_case_				:     str)->				jnp.ndarray:
						for processor in self:
											__lowerCAmelCase						 =inspect.signature(processor.__call__).parameters
											if len(snake_case_) > 3:
																if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
																					raise ValueError(
																					    F"""Make sure that all the required parameters: {list(function_args.keys())} for """
																					    F"""{processor.__class__} are passed to the logits processor.""")
																__lowerCAmelCase						 =processor(snake_case_					,						snake_case_					,						snake_case_					,						**snake_case_)
											else:
																__lowerCAmelCase						 =processor(snake_case_					,						snake_case_					,						snake_case_)
						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[Any]					,						snake_case_				:     float)->				Tuple:
						if not isinstance(snake_case_					,						snake_case_) or not (temperature > 0):
											raise ValueError(F"""`temperature` has to be a strictly positive float, but is {temperature}""")

						__lowerCAmelCase						 =temperature



	def __call__(    self				:     Dict					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase						 =scores / self.temperature
						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[int]					,						snake_case_				:     float					,						snake_case_				:     float = -float("""Inf""")					,						snake_case_				:     int = 1)->				Optional[Any]:
						if not isinstance(snake_case_					,						snake_case_) or (top_p < 0 or top_p > 1.0):
											raise ValueError(F"""`top_p` has to be a float > 0 and < 1, but is {top_p}""")
						if not isinstance(snake_case_					,						snake_case_) or (min_tokens_to_keep < 1):
											raise ValueError(F"""`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}""")

						__lowerCAmelCase						 =top_p
						__lowerCAmelCase						 =filter_value
						__lowerCAmelCase						 =min_tokens_to_keep



	def __call__(    self				:     int					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase ,						__lowerCAmelCase						 =lax.top_k(snake_case_					,						scores.shape[-1])

						__lowerCAmelCase						 =jnp.full_like(snake_case_					,						self.filter_value)
						__lowerCAmelCase						 =jax.nn.softmax(snake_case_					,						axis=-1).cumsum(axis=-1)
						__lowerCAmelCase						 =cumulative_probs < self.top_p

						# include the token that is higher than top_p as well
						__lowerCAmelCase						 =jnp.roll(snake_case_					,						1)
						score_mask |= score_mask.at[:, 0].set(snake_case_)

						# min tokens to keep
						__lowerCAmelCase						 =score_mask.at[:, : self.min_tokens_to_keep].set(snake_case_)

						__lowerCAmelCase						 =jnp.where(snake_case_					,						snake_case_					,						snake_case_)
						__lowerCAmelCase						 =jax.lax.sort_key_val(snake_case_					,						snake_case_)[-1]

						return next_scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Any					,						snake_case_				:     int					,						snake_case_				:     float = -float("""Inf""")					,						snake_case_				:     int = 1)->				Optional[int]:
						if not isinstance(snake_case_					,						snake_case_) or top_k <= 0:
											raise ValueError(F"""`top_k` has to be a strictly positive integer, but is {top_k}""")

						__lowerCAmelCase						 =max(snake_case_					,						snake_case_)
						__lowerCAmelCase						 =filter_value



	def __call__(    self				:     Any					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase ,						__lowerCAmelCase						 =scores.shape
						__lowerCAmelCase						 =jnp.full(batch_size * vocab_size					,						self.filter_value)

						__lowerCAmelCase						 =min(self.top_k					,						scores.shape[-1])  # Safety check
						__lowerCAmelCase ,						__lowerCAmelCase						 =lax.top_k(snake_case_					,						snake_case_)
						__lowerCAmelCase						 =jnp.broadcast_to((jnp.arange(snake_case_) * vocab_size)[:, None]					,						(batch_size, topk)).flatten()
						__lowerCAmelCase						 =topk_scores.flatten()
						__lowerCAmelCase						 =topk_indices.flatten() + shift

						__lowerCAmelCase						 =next_scores_flat.at[topk_indices_flat].set(snake_case_)
						__lowerCAmelCase						 =next_scores_flat.reshape(snake_case_					,						snake_case_)
						return next_scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     int					,						snake_case_				:     int)->				str:
						__lowerCAmelCase						 =bos_token_id



	def __call__(    self				:     List[str]					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase						 =jnp.full(scores.shape					,						-float("""inf"""))

						__lowerCAmelCase						 =1 - jnp.bool_(cur_len - 1)

						__lowerCAmelCase						 =jnp.where(snake_case_					,						new_scores.at[:, self.bos_token_id].set(0)					,						snake_case_)

						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[Any]					,						snake_case_				:     int					,						snake_case_				:     int)->				Optional[Any]:
						__lowerCAmelCase						 =max_length
						__lowerCAmelCase						 =eos_token_id



	def __call__(    self				:     str					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase						 =jnp.full(scores.shape					,						-float("""inf"""))

						__lowerCAmelCase						 =1 - jnp.bool_(cur_len - self.max_length + 1)

						__lowerCAmelCase						 =jnp.where(snake_case_					,						new_scores.at[:, self.eos_token_id].set(0)					,						snake_case_)

						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Union[str, Any]					,						snake_case_				:     int					,						snake_case_				:     int)->				Any:
						if not isinstance(snake_case_					,						snake_case_) or min_length < 0:
											raise ValueError(F"""`min_length` has to be a positive integer, but is {min_length}""")

						if not isinstance(snake_case_					,						snake_case_) or eos_token_id < 0:
											raise ValueError(F"""`eos_token_id` has to be a positive integer, but is {eos_token_id}""")

						__lowerCAmelCase						 =min_length
						__lowerCAmelCase						 =eos_token_id



	def __call__(    self				:     Any					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						# create boolean flag to decide if min length penalty should be applied
						__lowerCAmelCase						 =1 - jnp.clip(cur_len - self.min_length					,						0					,						1)

						__lowerCAmelCase						 =jnp.where(snake_case_					,						scores.at[:, self.eos_token_id].set(-float("""inf"""))					,						snake_case_)

						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[int]					,						snake_case_				:     Dict					,						snake_case_				:     str)->				List[Any]:
						__lowerCAmelCase						 =list(snake_case_)
						__lowerCAmelCase						 =begin_index



	def __call__(    self				:     List[Any]					,						snake_case_				:     Any					,						snake_case_				:     List[Any]					,						snake_case_				:     int)->				str:
						__lowerCAmelCase						 =1 - jnp.bool_(cur_len - self.begin_index)

						__lowerCAmelCase						 =jnp.where(snake_case_					,						scores.at[:, self.begin_suppress_tokens].set(-float("""inf"""))					,						snake_case_)

						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[Any]					,						snake_case_				:     list)->				Any:
						__lowerCAmelCase						 =list(snake_case_)



	def __call__(    self				:     Optional[Any]					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						__lowerCAmelCase						 =scores.at[..., self.suppress_tokens].set(-float("""inf"""))

						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Any					,						snake_case_				:     str)->				List[str]:
						__lowerCAmelCase						 =dict(snake_case_)
						# Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
						# index of the array corresponds to the index of the token to be forced, for XLA compatibility.
						# Indexes without forced tokens will have a negative value.
						__lowerCAmelCase						 =jnp.ones((max(force_token_map.keys()) + 1)					,						dtype=jnp.intaa) * -1
						for index, token in force_token_map.items():
											if token is not None:
																__lowerCAmelCase						 =force_token_array.at[index].set(snake_case_)
						__lowerCAmelCase						 =jnp.intaa(snake_case_)



	def __call__(    self				:     Tuple					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     jnp.ndarray					,						snake_case_				:     int)->				jnp.ndarray:
						def _force_token(snake_case_				:     Optional[Any]):
											__lowerCAmelCase						 =scores.shape[0]
											__lowerCAmelCase						 =self.force_token_array[generation_idx]

											__lowerCAmelCase						 =jnp.ones_like(snake_case_					,						dtype=scores.dtype) * -float("""inf""")
											__lowerCAmelCase						 =jnp.zeros((batch_size, 1)					,						dtype=scores.dtype)
											__lowerCAmelCase						 =lax.dynamic_update_slice(snake_case_					,						snake_case_					,						(0, current_token))
											return new_scores

						__lowerCAmelCase						 =lax.cond(
						    cur_len >= self.force_token_array.shape[0]					,						lambda: scores					,						lambda: lax.cond(
						        self.force_token_array[cur_len] >= 0					,						lambda: _force_token(snake_case_)					,						lambda: scores					,						)					,						)
						return scores






class     __a			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[int]					,						snake_case_				:     Any					,						snake_case_				:     Dict					,						snake_case_				:     Union[str, Any])->				List[Any]:
						__lowerCAmelCase						 =generate_config.eos_token_id
						__lowerCAmelCase						 =generate_config.no_timestamps_token_id
						__lowerCAmelCase						 =generate_config.no_timestamps_token_id + 1

						__lowerCAmelCase						 =decoder_input_length + 1

						if generate_config.is_multilingual:
											# room for language token and task token
											self.begin_index += 2
						if hasattr(snake_case_					,						"""max_initial_timestamp_index"""):
											__lowerCAmelCase						 =generate_config.max_initial_timestamp_index
						else:
											__lowerCAmelCase						 =model_config.vocab_size
						if self.max_initial_timestamp_index is None:
											__lowerCAmelCase						 =model_config.vocab_size



	def __call__(    self				:     List[str]					,						snake_case_				:     int					,						snake_case_				:     int					,						snake_case_				:     Union[str, Any])->				List[str]:
						# suppress <|notimestamps|> which is handled by without_timestamps
						__lowerCAmelCase						 =scores.at[:, self.no_timestamps_token_id].set(-float("""inf"""))

						def handle_pairs(snake_case_				:     Union[str, Any]					,						snake_case_				:     str):
											__lowerCAmelCase						 =jnp.where((cur_len - self.begin_index) >= 1					,						snake_case_					,						snake_case_)
											__lowerCAmelCase						 =jnp.where(
											    input_ids_k[cur_len - 1] >= self.timestamp_begin					,						True and last_was_timestamp					,						snake_case_					,						)

											__lowerCAmelCase						 =jnp.where((cur_len - self.begin_index) < 2					,						snake_case_					,						snake_case_)
											__lowerCAmelCase						 =jnp.where(
											    input_ids_k[cur_len - 2] >= self.timestamp_begin					,						snake_case_					,						snake_case_					,						)

											return jnp.where(
											    snake_case_					,						jnp.where(
											        penultimate_was_timestamp > 0					,						scores_k.at[self.timestamp_begin :].set(-float("""inf"""))					,						scores_k.at[: self.eos_token_id].set(-float("""inf"""))					,						)					,						snake_case_					,						)

						__lowerCAmelCase						 =jax.vmap(snake_case_)(snake_case_					,						snake_case_)

						__lowerCAmelCase						 =jnp.where(cur_len == self.begin_index					,						snake_case_					,						snake_case_)
						__lowerCAmelCase						 =jnp.where(
						    self.max_initial_timestamp_index is not None					,						True and apply_max_initial_timestamp					,						snake_case_					,						)

						__lowerCAmelCase						 =self.timestamp_begin + self.max_initial_timestamp_index

						__lowerCAmelCase						 =jnp.where(
						    snake_case_					,						scores.at[:, last_allowed + 1 :].set(-float("""inf"""))					,						snake_case_					,						)

						# if sum of probability over timestamps is above any other token, sample timestamp
						__lowerCAmelCase						 =jax.nn.log_softmax(snake_case_					,						axis=-1)

						def handle_cumulative_probs(snake_case_				:     Any					,						snake_case_				:     Tuple):
											__lowerCAmelCase						 =jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]					,						axis=-1)
											__lowerCAmelCase						 =jnp.max(logprobs_k[: self.timestamp_begin])
											return jnp.where(
											    timestamp_logprob > max_text_token_logprob					,						scores_k.at[: self.timestamp_begin].set(-float("""inf"""))					,						snake_case_					,						)

						__lowerCAmelCase						 =jax.vmap(snake_case_)(snake_case_					,						snake_case_)

						return scores









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 BatchEncoding, PreTrainedTokenizer
from ...utils import logging


lowercase_													=      logging.get_logger(__name__)

lowercase_													=      '''▁'''

lowercase_													=      {
    '''vocab_file''': '''vocab.json''',
    '''spm_file''': '''sentencepiece.bpe.model''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

lowercase_													=      {
    '''vocab_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/vocab.json''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/vocab.json''',
    },
    '''spm_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/sentencepiece.bpe.model''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/sentencepiece.bpe.model''',
    },
    '''tokenizer_config_file''': {
        '''facebook/m2m100_418M''': '''https://huggingface.co/facebook/m2m100_418M/resolve/main/tokenizer_config.json''',
        '''facebook/m2m100_1.2B''': '''https://huggingface.co/facebook/m2m100_1.2B/resolve/main/tokenizer_config.json''',
    },
}

lowercase_													=      {
    '''facebook/m2m100_418M''': 10_24,
}

# fmt: off
lowercase_													=      {
    '''m2m100''': ['''af''', '''am''', '''ar''', '''ast''', '''az''', '''ba''', '''be''', '''bg''', '''bn''', '''br''', '''bs''', '''ca''', '''ceb''', '''cs''', '''cy''', '''da''', '''de''', '''el''', '''en''', '''es''', '''et''', '''fa''', '''ff''', '''fi''', '''fr''', '''fy''', '''ga''', '''gd''', '''gl''', '''gu''', '''ha''', '''he''', '''hi''', '''hr''', '''ht''', '''hu''', '''hy''', '''id''', '''ig''', '''ilo''', '''is''', '''it''', '''ja''', '''jv''', '''ka''', '''kk''', '''km''', '''kn''', '''ko''', '''lb''', '''lg''', '''ln''', '''lo''', '''lt''', '''lv''', '''mg''', '''mk''', '''ml''', '''mn''', '''mr''', '''ms''', '''my''', '''ne''', '''nl''', '''no''', '''ns''', '''oc''', '''or''', '''pa''', '''pl''', '''ps''', '''pt''', '''ro''', '''ru''', '''sd''', '''si''', '''sk''', '''sl''', '''so''', '''sq''', '''sr''', '''ss''', '''su''', '''sv''', '''sw''', '''ta''', '''th''', '''tl''', '''tn''', '''tr''', '''uk''', '''ur''', '''uz''', '''vi''', '''wo''', '''xh''', '''yi''', '''yo''', '''zh''', '''zu'''],
    '''wmt21''': ['''en''', '''ha''', '''is''', '''ja''', '''cs''', '''ru''', '''zh''', '''de''']
}






class     __a			(  SCREAMING_SNAKE_CASE					):
	SCREAMING_SNAKE_CASE  	=							VOCAB_FILES_NAMES
	SCREAMING_SNAKE_CASE  	=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	SCREAMING_SNAKE_CASE  	=							PRETRAINED_VOCAB_FILES_MAP
	SCREAMING_SNAKE_CASE  	=							["input_ids", "attention_mask"]

	SCREAMING_SNAKE_CASE  	=							[]
	SCREAMING_SNAKE_CASE  	=							[]



	def __init__(    self				:     int					,						snake_case_				:     List[str]					,						snake_case_				:     Any					,						snake_case_				:     Union[str, Any]=None					,						snake_case_				:     Tuple=None					,						snake_case_				:     List[str]="<s>"					,						snake_case_				:     List[str]="</s>"					,						snake_case_				:     Optional[Any]="</s>"					,						snake_case_				:     Dict="<pad>"					,						snake_case_				:     Any="<unk>"					,						snake_case_				:     Tuple="m2m100"					,						snake_case_				:     Optional[Dict[str, Any]] = None					,						snake_case_				:     Dict=8					,						**snake_case_				:     List[Any]					,						)->				None:
						__lowerCAmelCase						 ={} if sp_model_kwargs is None else sp_model_kwargs

						__lowerCAmelCase						 =language_codes
						__lowerCAmelCase						 =FAIRSEQ_LANGUAGE_CODES[language_codes]
						__lowerCAmelCase						 ={lang_code: F"""__{lang_code}__""" for lang_code in fairseq_language_code}

						__lowerCAmelCase						 =kwargs.get("""additional_special_tokens"""					,						[])
						kwargs["additional_special_tokens"] += [
						    self.get_lang_token(snake_case_)
						    for lang_code in fairseq_language_code
						    if self.get_lang_token(snake_case_) not in kwargs["additional_special_tokens"]
						]

						super().__init__(
						    src_lang=snake_case_					,						tgt_lang=snake_case_					,						bos_token=snake_case_					,						eos_token=snake_case_					,						sep_token=snake_case_					,						unk_token=snake_case_					,						pad_token=snake_case_					,						language_codes=snake_case_					,						sp_model_kwargs=self.sp_model_kwargs					,						num_madeup_words=snake_case_					,						**snake_case_					,						)

						__lowerCAmelCase						 =vocab_file
						__lowerCAmelCase						 =load_json(snake_case_)
						__lowerCAmelCase						 ={v: k for k, v in self.encoder.items()}
						__lowerCAmelCase						 =spm_file
						__lowerCAmelCase						 =load_spm(snake_case_					,						self.sp_model_kwargs)

						__lowerCAmelCase						 =len(self.encoder)

						__lowerCAmelCase						 ={
						    self.get_lang_token(snake_case_): self.encoder_size + i for i, lang_code in enumerate(snake_case_)
						}
						__lowerCAmelCase						 ={lang_code: self.encoder_size + i for i, lang_code in enumerate(snake_case_)}
						__lowerCAmelCase						 ={v: k for k, v in self.lang_token_to_id.items()}

						__lowerCAmelCase						 =src_lang if src_lang is not None else """en"""
						__lowerCAmelCase						 =tgt_lang
						__lowerCAmelCase						 =self.get_lang_id(self._src_lang)
						self.set_src_lang_special_tokens(self._src_lang)

						__lowerCAmelCase						 =num_madeup_words



	@property
	def 		UpperCamelCase				(    self				:     int)->				int:
						return len(self.encoder) + len(self.lang_token_to_id)



	@property
	def 		UpperCamelCase				(    self				:     Optional[Any])->				str:
						return self._src_lang



	@src_lang.setter
	def 		UpperCamelCase				(    self				:     List[Any]					,						snake_case_				:     str)->				None:
						__lowerCAmelCase						 =new_src_lang
						self.set_src_lang_special_tokens(self._src_lang)



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     str)->				List[str]:
						return self.sp_model.encode(snake_case_					,						out_type=snake_case_)



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     List[str])->				str:
						if token in self.lang_token_to_id:
											return self.lang_token_to_id[token]
						return self.encoder.get(snake_case_					,						self.encoder[self.unk_token])



	def 		UpperCamelCase				(    self				:     Dict					,						snake_case_				:     int)->				str:
						if index in self.id_to_lang_token:
											return self.id_to_lang_token[index]
						return self.decoder.get(snake_case_					,						self.unk_token)



	def 		UpperCamelCase				(    self				:     Union[str, Any]					,						snake_case_				:     Tuple)->				str:
						__lowerCAmelCase						 =[]
						__lowerCAmelCase						 =""""""
						for token in tokens:
											# make sure that special tokens are not decoded using sentencepiece model
											if token in self.all_special_tokens:
																out_string += self.sp_model.decode(snake_case_) + token
																__lowerCAmelCase						 =[]
											else:
																current_sub_tokens.append(snake_case_)
						out_string += self.sp_model.decode(snake_case_)
						return out_string.strip()



	def 		UpperCamelCase				(    self				:     Union[str, Any]					,						snake_case_				:     List[int]					,						snake_case_				:     Optional[List[int]] = None					,						snake_case_				:     bool = False)->				List[int]:

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

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



	def 		UpperCamelCase				(    self				:     Any					,						snake_case_				:     List[int]					,						snake_case_				:     Optional[List[int]] = 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 		UpperCamelCase				(    self				:     Optional[Any])->				Dict:
						__lowerCAmelCase						 ={self.convert_ids_to_tokens(snake_case_): i for i in range(self.vocab_size)}
						vocab.update(self.added_tokens_encoder)
						return vocab



	def __getstate__(    self				:     Tuple)->				Dict:
						__lowerCAmelCase						 =self.__dict__.copy()
						__lowerCAmelCase						 =None
						return state



	def __setstate__(    self				:     List[Any]					,						snake_case_				:     Dict)->				None:
						__lowerCAmelCase						 =d

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

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



	def 		UpperCamelCase				(    self				:     Union[str, Any]					,						snake_case_				:     str					,						snake_case_				:     Optional[str] = None)->				Tuple[str]:
						__lowerCAmelCase						 =Path(snake_case_)
						if not save_dir.is_dir():
											raise OSError(F"""{save_directory} should be a directory""")
						__lowerCAmelCase						 =save_dir / (
						    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""vocab_file"""]
						)
						__lowerCAmelCase						 =save_dir / (
						    (filename_prefix + """-""" if filename_prefix else """""") + self.vocab_files_names["""spm_file"""]
						)

						save_json(self.encoder					,						snake_case_)

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

						return (str(snake_case_), str(snake_case_))



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     List[str]					,						snake_case_				:     str = "en"					,						snake_case_				:     Optional[List[str]] = None					,						snake_case_				:     str = "ro"					,						**snake_case_				:     int					,						)->				BatchEncoding:
						__lowerCAmelCase						 =src_lang
						__lowerCAmelCase						 =tgt_lang
						self.set_src_lang_special_tokens(self.src_lang)
						return super().prepare_seqaseq_batch(snake_case_					,						snake_case_					,						**snake_case_)



	def 		UpperCamelCase				(    self				:     Optional[int]					,						snake_case_				:     int					,						snake_case_				:     Optional[str]					,						snake_case_				:     Optional[str]					,						**snake_case_				:     List[str])->				List[Any]:
						if src_lang is None or tgt_lang is None:
											raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model""")
						__lowerCAmelCase						 =src_lang
						__lowerCAmelCase						 =self(snake_case_					,						add_special_tokens=snake_case_					,						**snake_case_)
						__lowerCAmelCase						 =self.get_lang_id(snake_case_)
						__lowerCAmelCase						 =tgt_lang_id
						return inputs



	def 		UpperCamelCase				(    self				:     Optional[int])->				Union[str, Any]:
						self.set_src_lang_special_tokens(self.src_lang)



	def 		UpperCamelCase				(    self				:     Dict)->				Dict:
						self.set_tgt_lang_special_tokens(self.tgt_lang)



	def 		UpperCamelCase				(    self				:     Union[str, Any]					,						snake_case_				:     str)->				None:
						__lowerCAmelCase						 =self.get_lang_token(snake_case_)
						__lowerCAmelCase						 =self.lang_token_to_id[lang_token]
						__lowerCAmelCase						 =[self.cur_lang_id]
						__lowerCAmelCase						 =[self.eos_token_id]



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     str)->				None:
						__lowerCAmelCase						 =self.get_lang_token(snake_case_)
						__lowerCAmelCase						 =self.lang_token_to_id[lang_token]
						__lowerCAmelCase						 =[self.cur_lang_id]
						__lowerCAmelCase						 =[self.eos_token_id]



	def 		UpperCamelCase				(    self				:     int					,						snake_case_				:     str)->				str:
						return self.lang_code_to_token[lang]



	def 		UpperCamelCase				(    self				:     List[str]					,						snake_case_				:     str)->				int:
						__lowerCAmelCase						 =self.get_lang_token(snake_case_)
						return self.lang_token_to_id[lang_token]






def        __lowerCAmelCase		(	__lowerCamelCase						:     str					, __lowerCamelCase						:     Dict[str, Any]						)	->    sentencepiece.SentencePieceProcessor:
					__lowerCAmelCase						 =sentencepiece.SentencePieceProcessor(**__lowerCamelCase						)
					spm.Load(str(__lowerCamelCase						)						)
					return spm






def        __lowerCAmelCase		(	__lowerCamelCase						:     str						)	->    Union[Dict, List]:
					with open(__lowerCamelCase					, """r"""						) as f:
										return json.load(__lowerCamelCase						)






def        __lowerCAmelCase		(	__lowerCamelCase						:     Any					, __lowerCamelCase						:     str						)	->    None:
					with open(__lowerCamelCase					, """w"""						) as f:
										json.dump(__lowerCamelCase					, __lowerCamelCase					, indent=2						)




