





'''simple docstring'''


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


lowerCAmelCase					:     Optional[int]												=get_logger(__name__)


lowerCAmelCase					:     Dict												=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(lowercase						)
     def __call__( self		:					Dict	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray						):





        """simple docstring"""




        raise NotImplementedError(
            F'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'						)







class       a_   :





     @add_start_docstrings(lowercase						)
     def __call__( self		:					Optional[Any]	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray						):





        """simple docstring"""




        raise NotImplementedError(
            F'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'						)







class       a_   (				_lowerCAmelCase							):





     @add_start_docstrings(lowercase						)
     def __call__( self		:					Optional[Any]	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int	,      **lowercase		:					Dict						):





        """simple docstring"""




        for processor in self:
           lowercase_			:Any    			= inspect.signature(processor.__call__						).parameters
           if len(lowercase						) > 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.'						)
              lowercase_			:str    			= processor(lowercase	,      lowercase	,      lowercase	,      **lowercase						)
           else:
              lowercase_			:Union[str, Any]    			= processor(lowercase	,      lowercase	,      lowercase						)
        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Dict	,      lowercase		:					float						):





        """simple docstring"""




        if not isinstance(lowercase	,      lowercase						) or not (temperature > 0):
           raise ValueError(F'`temperature` has to be a strictly positive float, but is {temperature}'						)

        lowercase_			:Optional[int]    			= temperature







     def __call__( self		:					str	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= scores / self.temperature
        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					str	,      lowercase		:					float	,      lowercase		:					float = -float("Inf"						)	,      lowercase		:					int = 1						):





        """simple docstring"""




        if not isinstance(lowercase	,      lowercase						) 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(lowercase	,      lowercase						) 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}'						)

        lowercase_			:Union[str, Any]    			= top_p
        lowercase_			:Tuple    			= filter_value
        lowercase_			:Optional[int]    			= min_tokens_to_keep







     def __call__( self		:					int	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			,							lowercase_			:Any    			= lax.top_k(lowercase	,      scores.shape[-1]						)

        lowercase_			:List[str]    			= jnp.full_like(lowercase	,      self.filter_value						)
        lowercase_			:Optional[Any]    			= jax.nn.softmax(lowercase	,      axis=-1						).cumsum(axis=-1						)
        lowercase_			:Any    			= cumulative_probs < self.top_p

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

        # min tokens to keep
        lowercase_			:List[Any]    			= score_mask.at[:, : self.min_tokens_to_keep].set(lowercase						)

        lowercase_			:Dict    			= jnp.where(lowercase	,      lowercase	,      lowercase						)
        lowercase_			:Optional[Any]    			= jax.lax.sort_key_val(lowercase	,      lowercase						)[-1]

        return next_scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Optional[Any]	,      lowercase		:					int	,      lowercase		:					float = -float("Inf"						)	,      lowercase		:					int = 1						):





        """simple docstring"""




        if not isinstance(lowercase	,      lowercase						) or top_k <= 0:
           raise ValueError(F'`top_k` has to be a strictly positive integer, but is {top_k}'						)

        lowercase_			:List[str]    			= max(lowercase	,      lowercase						)
        lowercase_			:Optional[int]    			= filter_value







     def __call__( self		:					int	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			,							lowercase_			:int    			= scores.shape
        lowercase_			:str    			= jnp.full(batch_size * vocab_size	,      self.filter_value						)

        lowercase_			:Union[str, Any]    			= min(self.top_k	,      scores.shape[-1]						)  # Safety check
        lowercase_			,							lowercase_			:Any    			= lax.top_k(lowercase	,      lowercase						)
        lowercase_			:List[Any]    			= jnp.broadcast_to((jnp.arange(lowercase						) * vocab_size)[:, None]	,      (batch_size, topk)						).flatten()
        lowercase_			:Tuple    			= topk_scores.flatten()
        lowercase_			:Any    			= topk_indices.flatten() + shift

        lowercase_			:Union[str, Any]    			= next_scores_flat.at[topk_indices_flat].set(lowercase						)
        lowercase_			:Optional[int]    			= next_scores_flat.reshape(lowercase	,      lowercase						)
        return next_scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Union[str, Any]	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:List[str]    			= bos_token_id







     def __call__( self		:					List[str]	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= jnp.full(scores.shape	,      -float("inf"						)						)

        lowercase_			:str    			= 1 - jnp.bool_(cur_len - 1						)

        lowercase_			:List[str]    			= jnp.where(lowercase	,      new_scores.at[:, self.bos_token_id].set(0						)	,      lowercase						)

        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Optional[Any]	,      lowercase		:					int	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Tuple    			= max_length
        lowercase_			:Optional[int]    			= eos_token_id







     def __call__( self		:					Tuple	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Any    			= jnp.full(scores.shape	,      -float("inf"						)						)

        lowercase_			:Union[str, Any]    			= 1 - jnp.bool_(cur_len - self.max_length + 1						)

        lowercase_			:List[Any]    			= jnp.where(lowercase	,      new_scores.at[:, self.eos_token_id].set(0						)	,      lowercase						)

        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Any	,      lowercase		:					int	,      lowercase		:					int						):





        """simple docstring"""




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

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

        lowercase_			:int    			= min_length
        lowercase_			:str    			= eos_token_id







     def __call__( self		:					str	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:List[str]    			= 1 - jnp.clip(cur_len - self.min_length	,      0	,      1						)

        lowercase_			:Union[str, Any]    			= jnp.where(lowercase	,      scores.at[:, self.eos_token_id].set(-float("inf"						)						)	,      lowercase						)

        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Dict	,      lowercase		:					Any	,      lowercase		:					List[str]						):





        """simple docstring"""




        lowercase_			:List[str]    			= list(lowercase						)
        lowercase_			:Optional[int]    			= begin_index







     def __call__( self		:					Tuple	,      lowercase		:					str	,      lowercase		:					Dict	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= 1 - jnp.bool_(cur_len - self.begin_index						)

        lowercase_			:List[Any]    			= jnp.where(lowercase	,      scores.at[:, self.begin_suppress_tokens].set(-float("inf"						)						)	,      lowercase						)

        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					List[Any]	,      lowercase		:					list						):





        """simple docstring"""




        lowercase_			:Dict    			= list(lowercase						)







     def __call__( self		:					Any	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:str    			= scores.at[..., self.suppress_tokens].set(-float("inf"						)						)

        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					List[str]	,      lowercase		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= dict(lowercase						)
        # 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.
        lowercase_			:Union[str, Any]    			= 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:
              lowercase_			:List[str]    			= force_token_array.at[index].set(lowercase						)
        lowercase_			:Union[str, Any]    			= jnp.intaa(lowercase						)







     def __call__( self		:					List[Any]	,      lowercase		:					jnp.ndarray	,      lowercase		:					jnp.ndarray	,      lowercase		:					int						):





        """simple docstring"""




        def _force_token(lowercase		:					Optional[int]						):
           lowercase_			:List[Any]    			= scores.shape[0]
           lowercase_			:Optional[Any]    			= self.force_token_array[generation_idx]

           lowercase_			:List[str]    			= jnp.ones_like(lowercase	,      dtype=scores.dtype						) * -float("inf"						)
           lowercase_			:Optional[int]    			= jnp.zeros((batch_size, 1)	,      dtype=scores.dtype						)
           lowercase_			:Optional[Any]    			= lax.dynamic_update_slice(lowercase	,      lowercase	,      (0, current_token)						)
           return new_scores

        lowercase_			:List[Any]    			= 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(lowercase						)	,      lambda: scores	,      )	,      )
        return scores







class       a_   (				_lowerCAmelCase							):





     def __init__( self		:					Dict	,      lowercase		:					str	,      lowercase		:					Any	,      lowercase		:					Any						):





        """simple docstring"""




        lowercase_			:str    			= generate_config.eos_token_id
        lowercase_			:Dict    			= generate_config.no_timestamps_token_id
        lowercase_			:int    			= generate_config.no_timestamps_token_id + 1

        lowercase_			:int    			= decoder_input_length + 1

        if generate_config.is_multilingual:
           # room for language token and task token
           self.begin_index += 2
        if hasattr(lowercase	,      "max_initial_timestamp_index"						):
           lowercase_			:Union[str, Any]    			= generate_config.max_initial_timestamp_index
        else:
           lowercase_			:str    			= model_config.vocab_size
        if self.max_initial_timestamp_index is None:
           lowercase_			:Union[str, Any]    			= model_config.vocab_size







     def __call__( self		:					Dict	,      lowercase		:					Tuple	,      lowercase		:					Tuple	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= scores.at[:, self.no_timestamps_token_id].set(-float("inf"						)						)

        def handle_pairs(lowercase		:					str	,      lowercase		:					int						):
           lowercase_			:List[str]    			= jnp.where((cur_len - self.begin_index) >= 1	,      lowercase	,      lowercase						)
           lowercase_			:List[Any]    			= jnp.where(
               input_ids_k[cur_len - 1] >= self.timestamp_begin	,      True and last_was_timestamp	,      lowercase	,      )

           lowercase_			:Optional[Any]    			= jnp.where((cur_len - self.begin_index) < 2	,      lowercase	,      lowercase						)
           lowercase_			:Optional[int]    			= jnp.where(
               input_ids_k[cur_len - 2] >= self.timestamp_begin	,      lowercase	,      lowercase	,      )

           return jnp.where(
               lowercase	,      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"						)						)	,      )	,      lowercase	,      )

        lowercase_			:Optional[Any]    			= jax.vmap(lowercase						)(lowercase	,      lowercase						)

        lowercase_			:Any    			= jnp.where(cur_len == self.begin_index	,      lowercase	,      lowercase						)
        lowercase_			:int    			= jnp.where(
            self.max_initial_timestamp_index is not None	,      True and apply_max_initial_timestamp	,      lowercase	,      )

        lowercase_			:List[str]    			= self.timestamp_begin + self.max_initial_timestamp_index

        lowercase_			:List[str]    			= jnp.where(
            lowercase	,      scores.at[:, last_allowed + 1 :].set(-float("inf"						)						)	,      lowercase	,      )

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

        def handle_cumulative_probs(lowercase		:					Dict	,      lowercase		:					Union[str, Any]						):
           lowercase_			:Union[str, Any]    			= jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]	,      axis=-1						)
           lowercase_			:Optional[Any]    			= 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"						)						)	,      lowercase	,      )

        lowercase_			:Dict    			= jax.vmap(lowercase						)(lowercase	,      lowercase						)

        return scores









'''simple docstring'''


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


lowerCAmelCase					:     int												=logging.get_logger(__name__)

lowerCAmelCase					:     List[str]												='''▁'''

lowerCAmelCase					:     List[str]												={
    '''vocab_file''': '''vocab.json''',
    '''spm_file''': '''sentencepiece.bpe.model''',
    '''tokenizer_config_file''': '''tokenizer_config.json''',
}

lowerCAmelCase					:     Optional[Any]												={
    '''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''',
    },
}

lowerCAmelCase					:     int												={
    '''facebook/m2m100_418M''': 1_024,
}

# fmt: off
lowerCAmelCase					:     str												={
    '''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_   (				_lowerCAmelCase							):
     __A      =			VOCAB_FILES_NAMES
     __A      =			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __A      =			PRETRAINED_VOCAB_FILES_MAP
     __A      =			["input_ids", "attention_mask"]

     __A      =			[]
     __A      =			[]





     def __init__( self		:					Any	,      lowercase		:					Any	,      lowercase		:					List[Any]	,      lowercase		:					int=None	,      lowercase		:					Optional[Any]=None	,      lowercase		:					Union[str, Any]="<s>"	,      lowercase		:					Any="</s>"	,      lowercase		:					Optional[int]="</s>"	,      lowercase		:					List[Any]="<pad>"	,      lowercase		:					Optional[int]="<unk>"	,      lowercase		:					Optional[int]="m2m100"	,      lowercase		:					Optional[Dict[str, Any]] = None	,      lowercase		:					Any=8	,      **lowercase		:					int	,      ):





        """simple docstring"""




        lowercase_			:Optional[int]    			= {} if sp_model_kwargs is None else sp_model_kwargs

        lowercase_			:Optional[Any]    			= language_codes
        lowercase_			:Tuple    			= FAIRSEQ_LANGUAGE_CODES[language_codes]
        lowercase_			:List[Any]    			= {lang_code: F'__{lang_code}__' for lang_code in fairseq_language_code}

        lowercase_			:Union[str, Any]    			= kwargs.get("additional_special_tokens"	,      []						)
        kwargs["additional_special_tokens"] += [
            self.get_lang_token(lowercase						)
            for lang_code in fairseq_language_code
            if self.get_lang_token(lowercase						) not in kwargs["additional_special_tokens"]
        ]

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

        lowercase_			:Optional[int]    			= vocab_file
        lowercase_			:Any    			= load_json(lowercase						)
        lowercase_			:Optional[Any]    			= {v: k for k, v in self.encoder.items()}
        lowercase_			:List[str]    			= spm_file
        lowercase_			:List[str]    			= load_spm(lowercase	,      self.sp_model_kwargs						)

        lowercase_			:Optional[int]    			= len(self.encoder						)

        lowercase_			:int    			= {
            self.get_lang_token(lowercase						): self.encoder_size + i for i, lang_code in enumerate(lowercase						)
        }
        lowercase_			:List[Any]    			= {lang_code: self.encoder_size + i for i, lang_code in enumerate(lowercase						)}
        lowercase_			:List[Any]    			= {v: k for k, v in self.lang_token_to_id.items()}

        lowercase_			:int    			= src_lang if src_lang is not None else "en"
        lowercase_			:Union[str, Any]    			= tgt_lang
        lowercase_			:List[Any]    			= self.get_lang_id(self._src_lang						)
        self.set_src_lang_special_tokens(self._src_lang						)

        lowercase_			:int    			= num_madeup_words





     @property
     def 							lowercase__      ( self		:					List[str]						):





        """simple docstring"""




        return len(self.encoder						) + len(self.lang_token_to_id						)





     @property
     def 							lowercase__      ( self		:					Any						):





        """simple docstring"""




        return self._src_lang





     @src_lang.setter
     def 							lowercase__      ( self		:					Optional[int]	,      lowercase		:					str						):





        """simple docstring"""




        lowercase_			:str    			= new_src_lang
        self.set_src_lang_special_tokens(self._src_lang						)





     def 							lowercase__      ( self		:					Dict	,      lowercase		:					str						):





        """simple docstring"""




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





     def 							lowercase__      ( self		:					Tuple	,      lowercase		:					Dict						):





        """simple docstring"""




        if token in self.lang_token_to_id:
           return self.lang_token_to_id[token]
        return self.encoder.get(lowercase	,      self.encoder[self.unk_token]						)





     def 							lowercase__      ( self		:					Any	,      lowercase		:					int						):





        """simple docstring"""




        if index in self.id_to_lang_token:
           return self.id_to_lang_token[index]
        return self.decoder.get(lowercase	,      self.unk_token						)





     def 							lowercase__      ( self		:					int	,      lowercase		:					int						):





        """simple docstring"""




        lowercase_			:Optional[Any]    			= []
        lowercase_			:Any    			= ""
        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(lowercase						) + token
              lowercase_			:str    			= []
           else:
              current_sub_tokens.append(lowercase						)
        out_string += self.sp_model.decode(lowercase						)
        return out_string.strip()





     def 							lowercase__      ( self		:					Any	,      lowercase		:					List[int]	,      lowercase		:					Optional[List[int]] = None	,      lowercase		:					bool = False						):





        """simple docstring"""




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

        lowercase_			:List[Any]    			= [1] * len(self.prefix_tokens						)
        lowercase_			:List[Any]    			= [1] * len(self.suffix_tokens						)
        if token_ids_a is None:
           return prefix_ones + ([0] * len(lowercase						)) + suffix_ones
        return prefix_ones + ([0] * len(lowercase						)) + ([0] * len(lowercase						)) + suffix_ones





     def 							lowercase__      ( self		:					Union[str, Any]	,      lowercase		:					List[int]	,      lowercase		:					Optional[List[int]] = None						):





        """simple docstring"""




        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 							lowercase__      ( self		:					Union[str, Any]						):





        """simple docstring"""




        lowercase_			:str    			= {self.convert_ids_to_tokens(lowercase						): i for i in range(self.vocab_size						)}
        vocab.update(self.added_tokens_encoder						)
        return vocab





     def __getstate__( self		:					str						):





        """simple docstring"""




        lowercase_			:Any    			= self.__dict__.copy()
        lowercase_			:str    			= None
        return state





     def __setstate__( self		:					Tuple	,      lowercase		:					Dict						):





        """simple docstring"""




        lowercase_			:int    			= d

        # for backward compatibility
        if not hasattr(self	,      "sp_model_kwargs"						):
           lowercase_			:List[str]    			= {}

        lowercase_			:List[Any]    			= load_spm(self.spm_file	,      self.sp_model_kwargs						)





     def 							lowercase__      ( self		:					str	,      lowercase		:					str	,      lowercase		:					Optional[str] = None						):





        """simple docstring"""




        lowercase_			:Dict    			= Path(lowercase						)
        if not save_dir.is_dir():
           raise OSError(F'{save_directory} should be a directory'						)
        lowercase_			:Dict    			= save_dir / (
            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
        )
        lowercase_			:Dict    			= save_dir / (
            (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
        )

        save_json(self.encoder	,      lowercase						)

        if os.path.abspath(self.spm_file						) != os.path.abspath(lowercase						) and os.path.isfile(self.spm_file						):
           copyfile(self.spm_file	,      lowercase						)
        elif not os.path.isfile(self.spm_file						):
           with open(lowercase	,      "wb"						) as fi:
              lowercase_			:List[str]    			= self.sp_model.serialized_model_proto()
              fi.write(lowercase						)

        return (str(lowercase						), str(lowercase						))





     def 							lowercase__      ( self		:					List[str]	,      lowercase		:					List[str]	,      lowercase		:					str = "en"	,      lowercase		:					Optional[List[str]] = None	,      lowercase		:					str = "ro"	,      **lowercase		:					Optional[int]	,      ):





        """simple docstring"""




        lowercase_			:int    			= src_lang
        lowercase_			:Optional[int]    			= tgt_lang
        self.set_src_lang_special_tokens(self.src_lang						)
        return super().prepare_seqaseq_batch(lowercase	,      lowercase	,      **lowercase						)





     def 							lowercase__      ( self		:					List[Any]	,      lowercase		:					Any	,      lowercase		:					Optional[str]	,      lowercase		:					Optional[str]	,      **lowercase		:					Union[str, Any]						):





        """simple docstring"""




        if src_lang is None or tgt_lang is None:
           raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"						)
        lowercase_			:List[str]    			= src_lang
        lowercase_			:Union[str, Any]    			= self(lowercase	,      add_special_tokens=lowercase	,      **lowercase						)
        lowercase_			:str    			= self.get_lang_id(lowercase						)
        lowercase_			:Union[str, Any]    			= tgt_lang_id
        return inputs





     def 							lowercase__      ( self		:					str						):





        """simple docstring"""




        self.set_src_lang_special_tokens(self.src_lang						)





     def 							lowercase__      ( self		:					Tuple						):





        """simple docstring"""




        self.set_tgt_lang_special_tokens(self.tgt_lang						)





     def 							lowercase__      ( self		:					str	,      lowercase		:					str						):





        """simple docstring"""




        lowercase_			:List[str]    			= self.get_lang_token(lowercase						)
        lowercase_			:List[str]    			= self.lang_token_to_id[lang_token]
        lowercase_			:List[Any]    			= [self.cur_lang_id]
        lowercase_			:str    			= [self.eos_token_id]





     def 							lowercase__      ( self		:					str	,      lowercase		:					str						):





        """simple docstring"""




        lowercase_			:Optional[int]    			= self.get_lang_token(lowercase						)
        lowercase_			:Tuple    			= self.lang_token_to_id[lang_token]
        lowercase_			:Dict    			= [self.cur_lang_id]
        lowercase_			:List[Any]    			= [self.eos_token_id]





     def 							lowercase__      ( self		:					Union[str, Any]	,      lowercase		:					str						):





        """simple docstring"""




        return self.lang_code_to_token[lang]







     def 							lowercase__      ( self		:					Dict	,      lowercase		:					str						):





        """simple docstring"""




        lowercase_			:Union[str, Any]    			= self.get_lang_token(lowercase						)
        return self.lang_token_to_id[lang_token]







def 							UpperCAmelCase_  (						__lowerCamelCase					:					str	,__lowerCamelCase					:					Dict[str, Any] ):
   lowercase_			:List[str]    			= sentencepiece.SentencePieceProcessor(**__lowerCamelCase )
   spm.Load(str(__lowerCamelCase ) )
   return spm
def 							UpperCAmelCase_  (						__lowerCamelCase					:					str ):
   with open(__lowerCamelCase	,"r" ) as f:
      return json.load(__lowerCamelCase )


def 							UpperCAmelCase_  (						__lowerCamelCase					:					int	,__lowerCamelCase					:					str ):
   with open(__lowerCamelCase	,"w" ) as f:
      json.dump(__lowerCamelCase	,__lowerCamelCase	,indent=2 )


