






'''simple docstring'''
from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   if not inplace:
      lowercase_		:      int          =       ma + ma
   else:
      ma += ma

   return ma




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   lowercase_		:      Any          =       -1 * len(_UpperCamelCase							)
   lowercase_		:      Any          =       list(range(len(tensor.shape[:zero_index]							)							)							)
   return tensor.permute(first_inds + [zero_index + i for i in inds]							)




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   return t.reshape(t.shape[:-no_dims] + (-1,)							)




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase = 1e-4							):


   """simple docstring"""




   lowercase_		:      Dict          =       mask.expand(*value.shape							)
   return torch.sum(mask * value  ,						dim=_UpperCamelCase							) / (eps + torch.sum(_UpperCamelCase  ,						dim=_UpperCamelCase							))




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase = 2.3125  ,						_UpperCamelCase = 21.6875  ,						_UpperCamelCase = 64							):


   """simple docstring"""




   lowercase_		:      Dict          =       torch.linspace(_UpperCamelCase  ,						_UpperCamelCase  ,						no_bins - 1  ,						device=pts.device							)
   lowercase_		:      Optional[int]          =       torch.sqrt(torch.sum((pts.unsqueeze(-2							) - pts.unsqueeze(-3							)) ** 2  ,						dim=-1							)							)
   return torch.bucketize(_UpperCamelCase  ,						_UpperCamelCase							)




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   lowercase_		:      List[str]          =       dicts[0]
   lowercase_		:      Optional[Any]          =       {}
   for k, v in first.items():
      lowercase_		:      int          =       [d[k] for d in dicts]
      if isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
         lowercase_		:      int          =       dict_multimap(_UpperCamelCase  ,						_UpperCamelCase							)
      else:
         lowercase_		:      Any          =       fn(_UpperCamelCase							)

   return new_dict




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   lowercase_		:      List[Any]          =       v_bins.view(((1,) * len(x.shape							)) + (len(_UpperCamelCase							),)							)
   lowercase_		:      Any          =       x[..., None] - reshaped_bins
   lowercase_		:      List[str]          =       torch.argmin(torch.abs(_UpperCamelCase							)  ,						dim=-1							)
   return nn.functional.one_hot(_UpperCamelCase  ,						num_classes=len(_UpperCamelCase							)							).float()




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase = 0  ,						_UpperCamelCase = 0							):


   """simple docstring"""




   lowercase_		:      List[Union[slice, torch.Tensor]]          =       []
   for i, s in enumerate(data.shape[:no_batch_dims]							):
      lowercase_		:      Optional[int]          =       torch.arange(_UpperCamelCase							)
      lowercase_		:      Tuple          =       r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape							) - i - 1)))							)
      ranges.append(_UpperCamelCase							)

   lowercase_		:      List[Union[slice, torch.Tensor]]          =       [slice(_UpperCamelCase							) for _ in range(len(data.shape							) - no_batch_dims							)]
   lowercase_		:      List[Any]          =       inds
   ranges.extend(_UpperCamelCase							)
   # Matt note: Editing this to get around the behaviour of using a list as an array index changing
   # in recent Numpy versions
   return data[tuple(_UpperCamelCase							)]


UpperCamelCase__	=     TypeVar('T')




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   lowercase_		:      Dict[Any, Union[dict, list, tuple, Any]]          =       {}
   for k, v in dic.items():
      if isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
         lowercase_		:      Optional[Any]          =       dict_map(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							)
      else:
         lowercase_		:      Tuple          =       tree_map(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							)

   return new_dict




@overload
def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   ...




@overload
def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   ...




@overload
def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   ...




@overload
def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   ...




def 						__SCREAMING_SNAKE_CASE   (					_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							):


   """simple docstring"""




   if isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
      return dict_map(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							)
   elif isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
      return [tree_map(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							) for x in tree]
   elif isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
      return tuple(tree_map(_UpperCamelCase  ,						_UpperCamelCase  ,						_UpperCamelCase							) for x in tree							)
   elif isinstance(_UpperCamelCase  ,						_UpperCamelCase							):
      return fn(_UpperCamelCase							)
   else:
      print(type(_UpperCamelCase							)							)
      raise ValueError("Not supported"							)





UpperCamelCase__	=     partial(tree_map, leaf_type=torch.Tensor)












'''simple docstring'''
import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
 from .tokenization_mbartaa import MBartaaTokenizer
else:
 UpperCamelCase__	=     None


UpperCamelCase__	=     logging.get_logger(__name__)

UpperCamelCase__	=     {'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

UpperCamelCase__	=     {
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/tokenizer.json'
        ),
    },
}

UpperCamelCase__	=     {
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
UpperCamelCase__	=     ['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']


class     _UpperCAmelCase  (						snake_case     ):
    __lowerCamelCase:				List[str]						    =      VOCAB_FILES_NAMES
    __lowerCamelCase:				int						    =      PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    __lowerCamelCase:				List[str]						    =      PRETRAINED_VOCAB_FILES_MAP
    __lowerCamelCase:				List[str]						    =      ['input_ids', 'attention_mask']
    __lowerCamelCase:				str						    =      MBartaaTokenizer

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




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






       '''simple docstring'''






       lowercase_		:      int          =       AddedToken(a					,       lstrip=a					,       rstrip=a     ) if isinstance(a					,       a     ) else mask_token

       lowercase_		:      List[Any]          =       kwargs.get("additional_special_tokens"					,       []     )
       kwargs["additional_special_tokens"] += [
           code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
       ]

       super().__init__(
           a					,       src_lang=a					,       tgt_lang=a					,       tokenizer_file=a					,       eos_token=a					,       sep_token=a					,       cls_token=a					,       unk_token=a					,       pad_token=a					,       mask_token=a					,       **a					,       )

       lowercase_		:      Any          =       vocab_file
       lowercase_		:      int          =       False if not self.vocab_file else True

       lowercase_		:      Optional[Any]          =       {
           lang_code: self.convert_tokens_to_ids(a     ) for lang_code in FAIRSEQ_LANGUAGE_CODES
       }

       lowercase_		:      int          =       src_lang if src_lang is not None else "en_XX"
       lowercase_		:      Tuple          =       tgt_lang
       lowercase_		:      List[Any]          =       self.lang_code_to_id[self._src_lang]
       self.set_src_lang_special_tokens(self._src_lang     )




    @property
    def 			lowerCAmelCase__			(     self					:    List[Any]     ):






       '''simple docstring'''






       return self._src_lang




    @src_lang.setter
    def 			lowerCAmelCase__			(     self					:    Union[str, Any]					,       a					:    str     ):






       '''simple docstring'''






       lowercase_		:      Optional[Any]          =       new_src_lang
       self.set_src_lang_special_tokens(self._src_lang     )




    def 			lowerCAmelCase__			(     self					:    int					,       a					:    List[int]					,       a					:    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 			lowerCAmelCase__			(     self					:    Any					,       a					:    List[str]					,       a					:    str = "en_XX"					,       a					:    Optional[List[str]] = None					,       a					:    str = "ro_RO"					,       **a					:    Union[str, Any]					,       ):






       '''simple docstring'''






       lowercase_		:      Dict          =       src_lang
       lowercase_		:      Dict          =       tgt_lang
       return super().prepare_seqaseq_batch(a					,       a					,       **a     )




    def 			lowerCAmelCase__			(     self					:    Dict     ):






       '''simple docstring'''






       return self.set_src_lang_special_tokens(self.src_lang     )




    def 			lowerCAmelCase__			(     self					:    int     ):






       '''simple docstring'''






       return self.set_tgt_lang_special_tokens(self.tgt_lang     )




    def 			lowerCAmelCase__			(     self					:    str					,       a					:    str     ):






       '''simple docstring'''






       lowercase_		:      Optional[int]          =       self.convert_tokens_to_ids(a     )
       lowercase_		:      Optional[int]          =       [self.cur_lang_code_id]
       lowercase_		:      List[Any]          =       [self.eos_token_id]

       lowercase_		:      Optional[Any]          =       self.convert_ids_to_tokens(self.prefix_tokens     )
       lowercase_		:      Union[str, Any]          =       self.convert_ids_to_tokens(self.suffix_tokens     )

       lowercase_		:      Optional[Any]          =       processors.TemplateProcessing(
           single=prefix_tokens_str + ["$A"] + suffix_tokens_str					,       pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str					,       special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str					,       self.prefix_tokens + self.suffix_tokens     )     )					,       )




    def 			lowerCAmelCase__			(     self					:    Tuple					,       a					:    str     ):






       '''simple docstring'''






       lowercase_		:      Optional[Any]          =       self.convert_tokens_to_ids(a     )
       lowercase_		:      str          =       [self.cur_lang_code_id]
       lowercase_		:      List[str]          =       [self.eos_token_id]

       lowercase_		:      Dict          =       self.convert_ids_to_tokens(self.prefix_tokens     )
       lowercase_		:      Dict          =       self.convert_ids_to_tokens(self.suffix_tokens     )

       lowercase_		:      int          =       processors.TemplateProcessing(
           single=prefix_tokens_str + ["$A"] + suffix_tokens_str					,       pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str					,       special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str					,       self.prefix_tokens + self.suffix_tokens     )     )					,       )




    def 			lowerCAmelCase__			(     self					:    Any					,       a					:    int					,       a					:    str					,       a					:    Optional[str]					,       a					:    Optional[str]					,       **a					:    str     ):






       '''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_		:      Any          =       src_lang
       lowercase_		:      int          =       self(a					,       add_special_tokens=a					,       return_tensors=a					,       **a     )
       lowercase_		:      Any          =       self.convert_tokens_to_ids(a     )
       lowercase_		:      Optional[Any]          =       tgt_lang_id
       return inputs





    def 			lowerCAmelCase__			(     self					:    Any					,       a					:    str					,       a					:    Optional[str] = None     ):






       '''simple docstring'''






       if not self.can_save_slow_tokenizer:
          raise ValueError(
              "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
              "tokenizer."     )

       if not os.path.isdir(a     ):
          logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""     )
          return
       lowercase_		:      Any          =       os.path.join(
           a					,       (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]     )

       if os.path.abspath(self.vocab_file     ) != os.path.abspath(a     ):
          copyfile(self.vocab_file					,       a     )

       return (out_vocab_file,)




