


from typing import List, Optional, Tuple, Union

import torch
from torch import nn
from torch.nn import CrossEntropyLoss

from ... import AutoBackbone
from ...modeling_outputs import SemanticSegmenterOutput
from ...modeling_utils import PreTrainedModel
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, replace_return_docstrings
from ...utils.backbone_utils import BackboneMixin
from .configuration_upernet import UperNetConfig


_A		       =						[
    '''openmmlab/upernet-convnext-tiny''',
    # See all UperNet models at https://huggingface.co/models?filter=upernet
]

# General docstring
_A		       =						'''UperNetConfig'''






class     UpperCAmelCase__       (					nn.Module		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_					,   A_					,   A_ = 0					,   A_ = False					,   A_ = 1					,   )  ->       None:
         super().__init__()
         __UpperCamelCase    							=nn.Convad(
             in_channels=__lowerCAmelCase					,   out_channels=__lowerCAmelCase					,   kernel_size=__lowerCAmelCase					,   padding=__lowerCAmelCase					,   bias=__lowerCAmelCase					,   dilation=__lowerCAmelCase					,   )
         __UpperCamelCase    							=nn.BatchNormad(__lowerCAmelCase    )
         __UpperCamelCase    							=nn.ReLU()






    def _a					(			self					,   A_    )  ->       torch.Tensor:
         __UpperCamelCase    							=self.conv(__lowerCAmelCase    )
         __UpperCamelCase    							=self.batch_norm(__lowerCAmelCase    )
         __UpperCamelCase    							=self.activation(__lowerCAmelCase    )

         return output






class     UpperCAmelCase__       (					nn.Module		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_					,   A_    )  ->       None:
         super().__init__()
         __UpperCamelCase    							=[
             nn.AdaptiveAvgPoolad(__lowerCAmelCase    ),
             UperNetConvModule(__lowerCAmelCase					,   __lowerCAmelCase					,   kernel_size=1    ),
         ]
         for i, layer in enumerate(self.layers    ):
              self.add_module(str(__lowerCAmelCase    )					,   __lowerCAmelCase    )






    def _a					(			self					,   A_    )  ->       torch.Tensor:
         __UpperCamelCase    							=input
         for layer in self.layers:
              __UpperCamelCase    							=layer(__lowerCAmelCase    )
         return hidden_state






class     UpperCAmelCase__       (					nn.Module		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_					,   A_					,   A_    )  ->       None:
         super().__init__()
         __UpperCamelCase    							=pool_scales
         __UpperCamelCase    							=align_corners
         __UpperCamelCase    							=in_channels
         __UpperCamelCase    							=channels
         __UpperCamelCase    							=[]
         for i, pool_scale in enumerate(__lowerCAmelCase    ):
              __UpperCamelCase    							=UperNetPyramidPoolingBlock(pool_scale=__lowerCAmelCase					,   in_channels=__lowerCAmelCase					,   channels=__lowerCAmelCase    )
              self.blocks.append(__lowerCAmelCase    )
              self.add_module(str(__lowerCAmelCase    )					,   __lowerCAmelCase    )






    def _a					(			self					,   A_    )  ->       List[torch.Tensor]:
         __UpperCamelCase    							=[]
         for ppm in self.blocks:
              __UpperCamelCase    							=ppm(__lowerCAmelCase    )
              __UpperCamelCase    							=nn.functional.interpolate(
                  __lowerCAmelCase					,   size=x.size()[2:]					,   mode='bilinear'					,   align_corners=self.align_corners    )
              ppm_outs.append(__lowerCAmelCase    )
         return ppm_outs






class     UpperCAmelCase__       (					nn.Module		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_    )  ->       List[str]:
         super().__init__()

         __UpperCamelCase    							=config
         __UpperCamelCase    							=config.pool_scales  # e.g. (1, 2, 3, 6)
         __UpperCamelCase    							=in_channels
         __UpperCamelCase    							=config.hidden_size
         __UpperCamelCase    							=False
         __UpperCamelCase    							=nn.Convad(self.channels					,   config.num_labels					,   kernel_size=1    )

         # PSP Module
         __UpperCamelCase    							=UperNetPyramidPoolingModule(
             self.pool_scales					,   self.in_channels[-1]					,   self.channels					,   align_corners=self.align_corners					,   )
         __UpperCamelCase    							=UperNetConvModule(
             self.in_channels[-1] + len(self.pool_scales    ) * self.channels					,   self.channels					,   kernel_size=3					,   padding=1					,   )
         # FPN Module
         __UpperCamelCase    							=nn.ModuleList()
         __UpperCamelCase    							=nn.ModuleList()
         for in_channels in self.in_channels[:-1]:  # skip the top layer
              __UpperCamelCase    							=UperNetConvModule(__lowerCAmelCase					,   self.channels					,   kernel_size=1    )
              __UpperCamelCase    							=UperNetConvModule(self.channels					,   self.channels					,   kernel_size=3					,   padding=1    )
              self.lateral_convs.append(__lowerCAmelCase    )
              self.fpn_convs.append(__lowerCAmelCase    )

         __UpperCamelCase    							=UperNetConvModule(
             len(self.in_channels    ) * self.channels					,   self.channels					,   kernel_size=3					,   padding=1					,   )






    def _a					(			self    )  ->       Union[str, Any]:
         self.apply(self._init_weights    )






    def _a					(			self					,   A_    )  ->       Dict:






         if isinstance(__lowerCAmelCase					,   nn.Convad    ):
              module.weight.data.normal_(mean=0.0					,   std=self.config.initializer_range    )
              if module.bias is not None:
                   module.bias.data.zero_()






    def _a					(			self					,   A_    )  ->       Dict:
         __UpperCamelCase    							=inputs[-1]
         __UpperCamelCase    							=[x]
         psp_outs.extend(self.psp_modules(__lowerCAmelCase    )    )
         __UpperCamelCase    							=torch.cat(__lowerCAmelCase					,   dim=1    )
         __UpperCamelCase    							=self.bottleneck(__lowerCAmelCase    )

         return output






    def _a					(			self					,   A_    )  ->       torch.Tensor:
         __UpperCamelCase    							=[lateral_conv(encoder_hidden_states[i]    ) for i, lateral_conv in enumerate(self.lateral_convs    )]

         laterals.append(self.psp_forward(__lowerCAmelCase    )    )

         # build top-down path
         __UpperCamelCase    							=len(__lowerCAmelCase    )
         for i in range(used_backbone_levels - 1					,   0					,   -1    ):
              __UpperCamelCase    							=laterals[i - 1].shape[2:]
              __UpperCamelCase    							=laterals[i - 1] + nn.functional.interpolate(
                  laterals[i]					,   size=__lowerCAmelCase					,   mode='bilinear'					,   align_corners=self.align_corners    )

         # build outputs
         __UpperCamelCase    							=[self.fpn_convs[i](laterals[i]    ) for i in range(used_backbone_levels - 1    )]
         # append psp feature
         fpn_outs.append(laterals[-1]    )

         for i in range(used_backbone_levels - 1					,   0					,   -1    ):
              __UpperCamelCase    							=nn.functional.interpolate(
                  fpn_outs[i]					,   size=fpn_outs[0].shape[2:]					,   mode='bilinear'					,   align_corners=self.align_corners    )
         __UpperCamelCase    							=torch.cat(__lowerCAmelCase					,   dim=1    )
         __UpperCamelCase    							=self.fpn_bottleneck(__lowerCAmelCase    )
         __UpperCamelCase    							=self.classifier(__lowerCAmelCase    )

         return output






class     UpperCAmelCase__       (					nn.Module		):



    """simple docstring"""






    def __init__(			self					,   A_					,   A_ = 2					,   A_ = 3					,   A_ = 1    )  ->       None:
         super().__init__()

         __UpperCamelCase    							=config
         __UpperCamelCase    							=config.auxiliary_in_channels
         __UpperCamelCase    							=config.auxiliary_channels
         __UpperCamelCase    							=config.auxiliary_num_convs
         __UpperCamelCase    							=config.auxiliary_concat_input
         __UpperCamelCase    							=in_index

         __UpperCamelCase    							=(kernel_size // 2) * dilation
         __UpperCamelCase    							=[]
         convs.append(
             UperNetConvModule(
                 self.in_channels					,   self.channels					,   kernel_size=__lowerCAmelCase					,   padding=__lowerCAmelCase					,   dilation=__lowerCAmelCase    )    )
         for i in range(self.num_convs - 1    ):
              convs.append(
                  UperNetConvModule(
                      self.channels					,   self.channels					,   kernel_size=__lowerCAmelCase					,   padding=__lowerCAmelCase					,   dilation=__lowerCAmelCase    )    )
         if self.num_convs == 0:
              __UpperCamelCase    							=nn.Identity()
         else:
              __UpperCamelCase    							=nn.Sequential(*__lowerCAmelCase    )
         if self.concat_input:
              __UpperCamelCase    							=UperNetConvModule(
                  self.in_channels + self.channels					,   self.channels					,   kernel_size=__lowerCAmelCase					,   padding=kernel_size // 2    )

         __UpperCamelCase    							=nn.Convad(self.channels					,   config.num_labels					,   kernel_size=1    )






    def _a					(			self    )  ->       Tuple:
         self.apply(self._init_weights    )






    def _a					(			self					,   A_    )  ->       Union[str, Any]:






         if isinstance(__lowerCAmelCase					,   nn.Convad    ):
              module.weight.data.normal_(mean=0.0					,   std=self.config.initializer_range    )
              if module.bias is not None:
                   module.bias.data.zero_()






    def _a					(			self					,   A_    )  ->       torch.Tensor:
         __UpperCamelCase    							=encoder_hidden_states[self.in_index]
         __UpperCamelCase    							=self.convs(__lowerCAmelCase    )
         if self.concat_input:
              __UpperCamelCase    							=self.conv_cat(torch.cat([hidden_states, output]					,   dim=1    )    )
         __UpperCamelCase    							=self.classifier(__lowerCAmelCase    )
         return output






class     UpperCAmelCase__       (					A_		):



    """simple docstring"""






    UpperCAmelCase__						:		str      =  UperNetConfig
    UpperCAmelCase__						:		List[Any]      =  '''pixel_values'''
    UpperCAmelCase__						:		Optional[Any]      =  True






    def _a					(			self					,   A_    )  ->       List[str]:






         if isinstance(__lowerCAmelCase					,   __lowerCAmelCase    ):
              module.backbone.init_weights()
              module.decode_head.init_weights()
              module.auxiliary_head.init_weights()






    def _a					(			self    )  ->       Tuple:
         self.backbone.init_weights()
         self.decode_head.init_weights()
         self.auxiliary_head.init_weights()






    def _a					(			self					,   A_					,   A_=False    )  ->       Any:






         if isinstance(__lowerCAmelCase					,   __lowerCAmelCase    ):
              __UpperCamelCase    							=value
_A		       =						R'''
    Parameters:
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.
        config ([`UperNetConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

_A		       =						R'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Padding will be ignored by default should you provide it. Pixel values can be obtained using
            [`AutoImageProcessor`]. See [`SegformerImageProcessor.__call__`] for details.
        output_attentions (`bool`, *optional*):
            Whether or not to return the attentions tensors of all attention layers in case the backbone has them. See
            `attentions` under returned tensors for more detail.
        output_hidden_states (`bool`, *optional*):
            Whether or not to return the hidden states of all layers of the backbone. See `hidden_states` under
            returned tensors for more detail.
        return_dict (`bool`, *optional*):
            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
'''






@add_start_docstrings(
    "UperNet framework leveraging any vision backbone e.g. for ADE20k, CityScapes."   ,  A_   ,  )
class     UpperCAmelCase__       (					A_		):



    """simple docstring"""






    def __init__(			self					,   A_    )  ->       List[Any]:
         super().__init__(__lowerCAmelCase    )

         __UpperCamelCase    							=AutoBackbone.from_config(config.backbone_config    )

         # Semantic segmentation head(s)
         __UpperCamelCase    							=UperNetHead(__lowerCAmelCase					,   in_channels=self.backbone.channels    )
         __UpperCamelCase    							=UperNetFCNHead(__lowerCAmelCase    ) if config.use_auxiliary_head else None

         # Initialize weights and apply final processing
         self.post_init()






    @add_start_docstrings_to_model_forward(UPERNET_INPUTS_DOCSTRING.format('batch_size, sequence_length'    )    )
    @replace_return_docstrings(output_type=__lowerCAmelCase					,   config_class=_CONFIG_FOR_DOC    )
    def _a					(			self					,   A_ = None					,   A_ = None					,   A_ = None					,   A_ = None					,   A_ = None					,   )  ->       Union[tuple, SemanticSegmenterOutput]:
         __UpperCamelCase    							=return_dict if return_dict is not None else self.config.use_return_dict
         __UpperCamelCase    							=(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         __UpperCamelCase    							=output_attentions if output_attentions is not None else self.config.output_attentions

         __UpperCamelCase    							=self.backbone.forward_with_filtered_kwargs(
             __lowerCAmelCase					,   output_hidden_states=__lowerCAmelCase					,   output_attentions=__lowerCAmelCase    )
         __UpperCamelCase    							=outputs.feature_maps

         __UpperCamelCase    							=self.decode_head(__lowerCAmelCase    )
         __UpperCamelCase    							=nn.functional.interpolate(__lowerCAmelCase					,   size=pixel_values.shape[2:]					,   mode='bilinear'					,   align_corners=__lowerCAmelCase    )

         __UpperCamelCase    							=None
         if self.auxiliary_head is not None:
              __UpperCamelCase    							=self.auxiliary_head(__lowerCAmelCase    )
              __UpperCamelCase    							=nn.functional.interpolate(
                  __lowerCAmelCase					,   size=pixel_values.shape[2:]					,   mode='bilinear'					,   align_corners=__lowerCAmelCase    )

         __UpperCamelCase    							=None
         if labels is not None:
              if self.config.num_labels == 1:
                   raise ValueError('The number of labels should be greater than one'    )
              else:
                   # compute weighted loss
                   __UpperCamelCase    							=CrossEntropyLoss(ignore_index=self.config.loss_ignore_index    )
                   __UpperCamelCase    							=loss_fct(__lowerCAmelCase					,   __lowerCAmelCase    )
                   __UpperCamelCase    							=loss_fct(__lowerCAmelCase					,   __lowerCAmelCase    )
                   __UpperCamelCase    							=main_loss + self.config.auxiliary_loss_weight * auxiliary_loss

         if not return_dict:
              if output_hidden_states:
                   __UpperCamelCase    							=(logits,) + outputs[1:]
              else:
                   __UpperCamelCase    							=(logits,) + outputs[2:]
              return ((loss,) + output) if loss is not None else output

         return SemanticSegmenterOutput(
             loss=__lowerCAmelCase					,   logits=__lowerCAmelCase					,   hidden_states=outputs.hidden_states					,   attentions=outputs.attentions					,   )












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


A    :					Optional[Any] =      logging.get_logger(__name__)

A    :					str =      '''▁'''

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

A    :					List[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''',
    },
}

A    :					Tuple =      {
    '''facebook/m2m100_418M''': 1_0_2_4,
}

# fmt: off
A    :					Optional[int] =      {
    '''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    ):
			'''simple docstring'''




			__lowerCamelCase	:       Union[str, Any]							=    VOCAB_FILES_NAMES
			__lowerCamelCase	:       List[str]							=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			__lowerCamelCase	:       Optional[int]							=    PRETRAINED_VOCAB_FILES_MAP
			__lowerCamelCase	:       Dict							=    ['''input_ids''', '''attention_mask''']

			__lowerCamelCase	:       List[int]							=    []
			__lowerCamelCase	:       List[int]							=    []






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







					"""simple docstring"""






					A__        =   {} if sp_model_kwargs is None else sp_model_kwargs

					A__        =   language_codes
					A__        =   FAIRSEQ_LANGUAGE_CODES[language_codes]
					A__        =   {lang_code: f'__{lang_code}__' for lang_code in fairseq_language_code}

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

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

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

					A__        =   len(self.encoder   )

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

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

					A__        =   num_madeup_words






			@property
			def    a_       (			self   :    Optional[int]   )	->	int:







					"""simple docstring"""






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






			@property
			def    a_       (			self   :    Optional[Any]   )	->	str:







					"""simple docstring"""






					return self._src_lang






			@src_lang.setter
			def    a_       (			self   :    List[Any]							,	__lowerCAmelCase   :    str   )	->	None:







					"""simple docstring"""






					A__        =   new_src_lang
					self.set_src_lang_special_tokens(self._src_lang   )






			def    a_       (			self   :    Optional[int]							,	__lowerCAmelCase   :    str   )	->	List[str]:







					"""simple docstring"""






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






			def    a_       (			self   :    Optional[Any]							,	__lowerCAmelCase   :    Dict   )	->	Optional[Any]:







					"""simple docstring"""






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






			def    a_       (			self   :    Optional[int]							,	__lowerCAmelCase   :    int   )	->	str:







					"""simple docstring"""






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






			def    a_       (			self   :    Optional[int]							,	__lowerCAmelCase   :    Dict   )	->	str:







					"""simple docstring"""






					A__        =   []
					A__        =   """"""
					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(__lowerCAmelCase   ) + token
									A__        =   []
							else:
									current_sub_tokens.append(__lowerCAmelCase   )
					out_string += self.sp_model.decode(__lowerCAmelCase   )
					return out_string.strip()






			def    a_       (			self   :    List[str]							,	__lowerCAmelCase   :    List[int]							,	__lowerCAmelCase   :    Optional[List[int]] = None							,	__lowerCAmelCase   :    bool = False   )	->	List[int]:







					"""simple docstring"""






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

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






			def    a_       (			self   :    Tuple							,	__lowerCAmelCase   :    List[int]							,	__lowerCAmelCase   :    Optional[List[int]] = None   )	->	List[int]:







					"""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    a_       (			self   :    int   )	->	Dict:







					"""simple docstring"""






					A__        =   {self.convert_ids_to_tokens(__lowerCAmelCase   ): i for i in range(self.vocab_size   )}
					vocab.update(self.added_tokens_encoder   )
					return vocab






			def __getstate__(			self   :    Union[str, Any]   )	->	Dict:







					"""simple docstring"""






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






			def __setstate__(			self   :    str							,	__lowerCAmelCase   :    Dict   )	->	None:







					"""simple docstring"""






					A__        =   d

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

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






			def    a_       (			self   :    List[str]							,	__lowerCAmelCase   :    str							,	__lowerCAmelCase   :    Optional[str] = None   )	->	Tuple[str]:







					"""simple docstring"""






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

					save_json(self.encoder							,	__lowerCAmelCase   )

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

					return (str(__lowerCAmelCase   ), str(__lowerCAmelCase   ))






			def    a_       (			self   :    str							,	__lowerCAmelCase   :    List[str]							,	__lowerCAmelCase   :    str = "en"							,	__lowerCAmelCase   :    Optional[List[str]] = None							,	__lowerCAmelCase   :    str = "ro"							,	**__lowerCAmelCase   :    List[Any]							,	)	->	BatchEncoding:







					"""simple docstring"""






					A__        =   src_lang
					A__        =   tgt_lang
					self.set_src_lang_special_tokens(self.src_lang   )
					return super().prepare_seqaseq_batch(__lowerCAmelCase							,	__lowerCAmelCase							,	**__lowerCAmelCase   )






			def    a_       (			self   :    Optional[int]							,	__lowerCAmelCase   :    Any							,	__lowerCAmelCase   :    Optional[str]							,	__lowerCAmelCase   :    Optional[str]							,	**__lowerCAmelCase   :    Tuple   )	->	Tuple:







					"""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"""   )
					A__        =   src_lang
					A__        =   self(__lowerCAmelCase							,	add_special_tokens=__lowerCAmelCase							,	**__lowerCAmelCase   )
					A__        =   self.get_lang_id(__lowerCAmelCase   )
					A__        =   tgt_lang_id
					return inputs






			def    a_       (			self   :    Dict   )	->	int:







					"""simple docstring"""






					self.set_src_lang_special_tokens(self.src_lang   )






			def    a_       (			self   :    Union[str, Any]   )	->	Union[str, Any]:







					"""simple docstring"""






					self.set_tgt_lang_special_tokens(self.tgt_lang   )






			def    a_       (			self   :    str							,	__lowerCAmelCase   :    str   )	->	None:







					"""simple docstring"""






					A__        =   self.get_lang_token(__lowerCAmelCase   )
					A__        =   self.lang_token_to_id[lang_token]
					A__        =   [self.cur_lang_id]
					A__        =   [self.eos_token_id]






			def    a_       (			self   :    Tuple							,	__lowerCAmelCase   :    str   )	->	None:







					"""simple docstring"""






					A__        =   self.get_lang_token(__lowerCAmelCase   )
					A__        =   self.lang_token_to_id[lang_token]
					A__        =   [self.cur_lang_id]
					A__        =   [self.eos_token_id]






			def    a_       (			self   :    Union[str, Any]							,	__lowerCAmelCase   :    str   )	->	str:







					"""simple docstring"""






					return self.lang_code_to_token[lang]







			def    a_       (			self   :    Union[str, Any]							,	__lowerCAmelCase   :    str   )	->	int:







					"""simple docstring"""






					A__        =   self.get_lang_token(__lowerCAmelCase   )
					return self.lang_token_to_id[lang_token]







def    __lowerCamelCase					(							__a   :str    ,     __a   :Dict[str, Any]    )		->				sentencepiece.SentencePieceProcessor:






		"""simple docstring"""

		A__        =   sentencepiece.SentencePieceProcessor(**__a    )
		spm.Load(str(__a    )    )
		return spm







def    __lowerCamelCase					(							__a   :str    )		->				Union[Dict, List]:






		"""simple docstring"""

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







def    __lowerCamelCase					(							__a   :List[Any]    ,     __a   :str    )		->				None:






		"""simple docstring"""

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



