




import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


UpperCamelCase__    			=			logging.get_logger(__name__)

UpperCamelCase__    			=			"""▁"""

UpperCamelCase__    			=			{"""vocab_file""": """sentencepiece.bpe.model"""}

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"""
        ),
    }
}

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 			a__    (					snake_case__      ):
						_a							:    Tuple					=    VOCAB_FILES_NAMES
						_a							:    int					=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						_a							:    Optional[int]					=    PRETRAINED_VOCAB_FILES_MAP
						_a							:    Dict					=    ["""input_ids""", """attention_mask"""]

						_a							:    List[int]					=    []
						_a							:    List[int]					=    []






						def __init__(   self		,			_A		,			_A=None		,			_A=None		,			_A="</s>"		,			_A="</s>"		,			_A="<s>"		,			_A="<unk>"		,			_A="<pad>"		,			_A="<mask>"		,			_A = None		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       AddedToken(_A		,			lstrip=_A		,			rstrip=_A			) if isinstance(_A		,			_A			) else mask_token

											__lowerCAmelCase				   =       {} if sp_model_kwargs is None else sp_model_kwargs

											__lowerCAmelCase				   =       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__(
											    src_lang=_A		,			tgt_lang=_A		,			eos_token=_A		,			unk_token=_A		,			sep_token=_A		,			cls_token=_A		,			pad_token=_A		,			mask_token=_A		,			sp_model_kwargs=self.sp_model_kwargs		,			**_A		,			)

											__lowerCAmelCase				   =       spm.SentencePieceProcessor(**self.sp_model_kwargs			)
											self.sp_model.Load(str(_A			)			)
											__lowerCAmelCase				   =       vocab_file

											# Original fairseq vocab and spm vocab must be "aligned":
											# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
											# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
											# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
											# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

											# Mimic fairseq token-to-id alignment for the first 4 token
											__lowerCAmelCase				   =       {"<s>": 0, "<pad>": 1, "</s>": 2, "<unk>": 3}

											# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
											__lowerCAmelCase				   =       1

											__lowerCAmelCase				   =       len(self.sp_model			)
											__lowerCAmelCase				   =       {
											    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_A			)
											}
											__lowerCAmelCase				   =       {v: k for k, v in self.lang_code_to_id.items()}
											__lowerCAmelCase				   =       len(self.sp_model			) + len(self.lang_code_to_id			) + self.fairseq_offset

											self.fairseq_tokens_to_ids.update(self.lang_code_to_id			)
											__lowerCAmelCase				   =       {v: k for k, v in self.fairseq_tokens_to_ids.items()}

											__lowerCAmelCase				   =       src_lang if src_lang is not None else "en_XX"
											__lowerCAmelCase				   =       self.lang_code_to_id[self._src_lang]
											__lowerCAmelCase				   =       tgt_lang
											self.set_src_lang_special_tokens(self._src_lang			)






						@property
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return len(self.sp_model			) + len(self.lang_code_to_id			) + self.fairseq_offset + 1  # Plus 1 for the mask token






						@property
						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self._src_lang






						@src_lang.setter
						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       new_src_lang
											self.set_src_lang_special_tokens(self._src_lang			)






						def __getstate__(   self			):




											"""simple docstring"""



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






						def __setstate__(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       d

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

											__lowerCAmelCase				   =       spm.SentencePieceProcessor(**self.sp_model_kwargs			)
											self.sp_model.Load(self.vocab_file			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											if token in self.fairseq_tokens_to_ids:
																return self.fairseq_tokens_to_ids[token]
											__lowerCAmelCase				   =       self.sp_model.PieceToId(_A			)

											# Need to return unknown token if the SP model returned 0
											return spm_id + self.fairseq_offset if spm_id else self.unk_token_id






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											if index in self.fairseq_ids_to_tokens:
																return self.fairseq_ids_to_tokens[index]
											return self.sp_model.IdToPiece(index - self.fairseq_offset			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       []
											__lowerCAmelCase				   =       ""
											__lowerCAmelCase				   =       False
											for token in tokens:
																# make sure that special tokens are not decoded using sentencepiece model
																if token in self.all_special_tokens:
																					if not prev_is_special:
																										out_string += " "
																					out_string += self.sp_model.decode(_A			) + token
																					__lowerCAmelCase				   =       True
																					__lowerCAmelCase				   =       []
																else:
																					current_sub_tokens.append(_A			)
																					__lowerCAmelCase				   =       False
											out_string += self.sp_model.decode(_A			)
											return out_string.strip()






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = None			):




											"""simple docstring"""



											if not os.path.isdir(_A			):
																logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""			)
																return
											__lowerCAmelCase				   =       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			) and os.path.isfile(self.vocab_file			):
																copyfile(self.vocab_file		,			_A			)
											elif not os.path.isfile(self.vocab_file			):
																with open(_A		,			"wb"			) as fi:
																					__lowerCAmelCase				   =       self.sp_model.serialized_model_proto()
																					fi.write(_A			)

											return (out_vocab_file,)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = None		,			_A = False			):




											"""simple docstring"""



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

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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = 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 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A		,			_A		,			**_A			):




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






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A = "en_XX"		,			_A = None		,			_A = "ro_RO"		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       src_lang
											__lowerCAmelCase				   =       tgt_lang
											return super().prepare_seqaseq_batch(_A		,			_A		,			**_A			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self.set_src_lang_special_tokens(self.src_lang			)






						def 				__SCREAMING_SNAKE_CASE(   self			):




											"""simple docstring"""



											return self.set_tgt_lang_special_tokens(self.tgt_lang			)






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.lang_code_to_id[src_lang]
											__lowerCAmelCase				   =       [self.cur_lang_code_id]
											__lowerCAmelCase				   =       [self.eos_token_id]






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       self.lang_code_to_id[tgt_lang]
											__lowerCAmelCase				   =       [self.cur_lang_code_id]
											__lowerCAmelCase				   =       [self.eos_token_id]









import re
from typing import Callable, List, Optional, Union

import tensorflow as tf


try:
     from tensorflow.keras.optimizers.legacy import Adam
except ImportError:
     from tensorflow.keras.optimizers import Adam






class SCREAMING_SNAKE_CASE__     (		tf.keras.optimizers.schedules.LearningRateSchedule    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__ = 1.0,		lowerCamelCase__ = None,		):
        super().__init__()
        A :       Union[str, Any]          =   initial_learning_rate
        A :       List[Any]          =   warmup_steps
        A :       int          =   power
        A :       Optional[int]          =   decay_schedule_fn
        A :       int          =   name






    def __call__(       self,		lowerCamelCase__		):
        with tf.name_scope(self.name or """WarmUp"""		) as name:
            # Implements polynomial warmup. i.e., if global_step < warmup_steps, the
            # learning rate will be `global_step/num_warmup_steps * init_lr`.
            A :       str          =   tf.cast(lowerCamelCase__,		tf.floataa		)
            A :       List[Any]          =   tf.cast(self.warmup_steps,		tf.floataa		)
            A :       Dict          =   global_step_float / warmup_steps_float
            A :       Union[str, Any]          =   self.initial_learning_rate * tf.math.pow(lowerCamelCase__,		self.power		)
            return tf.cond(
                global_step_float < warmup_steps_float,		lambda: warmup_learning_rate,		lambda: self.decay_schedule_fn(step - self.warmup_steps		),		name=lowerCamelCase__,		)






    def 			_lowerCAmelCase					(       self		):
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_schedule_fn": self.decay_schedule_fn,
            "warmup_steps": self.warmup_steps,
            "power": self.power,
            "name": self.name,
        }




def 	__UpperCamelCase	(  _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase      ,     _lowerCAmelCase = 0.0      ,     _lowerCAmelCase = 0.9      ,     _lowerCAmelCase = 0.999      ,     _lowerCAmelCase = 1e-8      ,     _lowerCAmelCase = None      ,     _lowerCAmelCase = None      ,     _lowerCAmelCase = 0.0      ,     _lowerCAmelCase = 1.0      ,     _lowerCAmelCase = None      ,     )       ->  Union[str, Any]:



    """simple docstring"""






    A :       Optional[int]          =   tf.keras.optimizers.schedules.PolynomialDecay(
        initial_learning_rate=_lowerCAmelCase      ,     decay_steps=num_train_steps - num_warmup_steps      ,     end_learning_rate=init_lr * min_lr_ratio      ,     power=_lowerCAmelCase      ,     )
    if num_warmup_steps:
        A :       Dict          =   WarmUp(
            initial_learning_rate=_lowerCAmelCase      ,     decay_schedule_fn=_lowerCAmelCase      ,     warmup_steps=_lowerCAmelCase      ,     )
    if weight_decay_rate > 0.0:
        A :       str          =   AdamWeightDecay(
            learning_rate=_lowerCAmelCase      ,     weight_decay_rate=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     epsilon=_lowerCAmelCase      ,     clipnorm=_lowerCAmelCase      ,     global_clipnorm=_lowerCAmelCase      ,     exclude_from_weight_decay=["""LayerNorm""", """layer_norm""", """bias"""]      ,     include_in_weight_decay=_lowerCAmelCase      ,     )
    else:
        A :       Optional[int]          =   tf.keras.optimizers.Adam(
            learning_rate=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     beta_a=_lowerCAmelCase      ,     epsilon=_lowerCAmelCase      ,     clipnorm=_lowerCAmelCase      ,     global_clipnorm=_lowerCAmelCase      ,     )
    # We return the optimizer and the LR scheduler in order to better track the
    # evolution of the LR independently of the optimizer.
    return optimizer, lr_schedule






class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__    ):




    '''simple docstring'''






    def __init__(       self,		lowerCamelCase__ = 0.001,		lowerCamelCase__ = 0.9,		lowerCamelCase__ = 0.999,		lowerCamelCase__ = 1e-7,		lowerCamelCase__ = False,		lowerCamelCase__ = 0.0,		lowerCamelCase__ = None,		lowerCamelCase__ = None,		lowerCamelCase__ = "AdamWeightDecay",		**lowerCamelCase__,		):
        super().__init__(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)
        A :       int          =   weight_decay_rate
        A :       Any          =   include_in_weight_decay
        A :       Dict          =   exclude_from_weight_decay






    @classmethod
    def 			_lowerCAmelCase					(       cls,		lowerCamelCase__		):
        A :       Tuple          =   {"""WarmUp""": WarmUp}
        return super(lowerCamelCase__,		cls		).from_config(lowerCamelCase__,		custom_objects=lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        super(lowerCamelCase__,		self		)._prepare_local(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        A :       List[str]          =   tf.constant(
            self.weight_decay_rate,		name="""adam_weight_decay_rate"""		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        A :       Any          =   self._do_use_weight_decay(var.name		)
        if do_decay:
            return var.assign_sub(
                learning_rate * var * apply_state[(var.device, var.dtype.base_dtype)]["""weight_decay_rate"""],		use_locking=self._use_locking,		)
        return tf.no_op()






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__=None,		**lowerCamelCase__		):
        A     ,							A :       Dict          =   list(zip(*lowerCamelCase__		)		)
        return super(lowerCamelCase__,		self		).apply_gradients(zip(lowerCamelCase__,		lowerCamelCase__		),		name=lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		):
        if apply_state is None:
            return self._decayed_lr_t[var_dtype], {}

        A :       Union[str, Any]          =   apply_state or {}
        A :       Optional[int]          =   apply_state.get((var_device, var_dtype)		)
        if coefficients is None:
            A :       Dict          =   self._fallback_apply_state(lowerCamelCase__,		lowerCamelCase__		)
            A :       List[Any]          =   coefficients

        return coefficients["lr_t"], {"apply_state": apply_state}






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A     ,							A :       str          =   self._get_lr(var.device,		var.dtype.base_dtype,		lowerCamelCase__		)
        A :       Any          =   self._decay_weights_op(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        with tf.control_dependencies([decay]		):
            return super(lowerCamelCase__,		self		)._resource_apply_dense(lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__=None		):
        A     ,							A :       Tuple          =   self._get_lr(var.device,		var.dtype.base_dtype,		lowerCamelCase__		)
        A :       Optional[Any]          =   self._decay_weights_op(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__		)
        with tf.control_dependencies([decay]		):
            return super(lowerCamelCase__,		self		)._resource_apply_sparse(lowerCamelCase__,		lowerCamelCase__,		lowerCamelCase__,		**lowerCamelCase__		)






    def 			_lowerCAmelCase					(       self		):
        A :       Dict          =   super().get_config()
        config.update({"""weight_decay_rate""": self.weight_decay_rate}		)
        return config






    def 			_lowerCAmelCase					(       self,		lowerCamelCase__		):
        if self.weight_decay_rate == 0:
            return False

        if self._include_in_weight_decay:
            for r in self._include_in_weight_decay:
                if re.search(lowerCamelCase__,		lowerCamelCase__		) is not None:
                    return True

        if self._exclude_from_weight_decay:
            for r in self._exclude_from_weight_decay:
                if re.search(lowerCamelCase__,		lowerCamelCase__		) is not None:
                    return False
        return True






class SCREAMING_SNAKE_CASE__     (		SCREAMING_SNAKE_CASE__    ):




    '''simple docstring'''






    def __init__(       self		):
        A :       List[str]          =   []
        A :       List[str]          =   None






    @property
    def 			_lowerCAmelCase					(       self		):
        if self._accum_steps is None:
            A :       str          =   tf.Variable(
                tf.constant(0,		dtype=tf.intaa		),		trainable=lowerCamelCase__,		synchronization=tf.VariableSynchronization.ON_READ,		aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,		)

        return self._accum_steps.value()






    @property
    def 			_lowerCAmelCase					(       self		):
        if not self._gradients:
            raise ValueError("""The accumulator should be called first to initialize the gradients"""		)
        return [gradient.value() if gradient is not None else gradient for gradient in self._gradients]






    def __call__(       self,		lowerCamelCase__		):
        if not self._gradients:
            A :       int          =   self.step  # Create the step variable.
            self._gradients.extend(
                [
                    tf.Variable(
                        tf.zeros_like(lowerCamelCase__		),		trainable=lowerCamelCase__,		synchronization=tf.VariableSynchronization.ON_READ,		aggregation=tf.VariableAggregation.ONLY_FIRST_REPLICA,		)
                    if gradient is not None
                    else gradient
                    for gradient in gradients
                ]		)
        if len(lowerCamelCase__		) != len(self._gradients		):
            raise ValueError(f'''Expected {len(self._gradients		)} gradients, but got {len(lowerCamelCase__		)}'''		)

        for accum_gradient, gradient in zip(self._gradients,		lowerCamelCase__		):
            if accum_gradient is not None and gradient is not None:
                accum_gradient.assign_add(lowerCamelCase__		)

        self._accum_steps.assign_add(1		)






    def 			_lowerCAmelCase					(       self		):
        if not self._gradients:
            return
        self._accum_steps.assign(0		)
        for gradient in self._gradients:
            if gradient is not None:
                gradient.assign(tf.zeros_like(lowerCamelCase__		)		)


