


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      :					Optional[int]    							=logging.get_logger(__name__)

__lowerCAmelCase      :					Dict    							='▁'

__lowerCAmelCase      :					Optional[Any]    							={
    'vocab_file': 'vocab.json',
    'spm_file': 'sentencepiece.bpe.model',
    'tokenizer_config_file': 'tokenizer_config.json',
}

__lowerCAmelCase      :					Tuple    							={
    '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      :					List[Any]    							={
    'facebook/m2m100_418M': 1_0_2_4,
}

# fmt: off
__lowerCAmelCase      :					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 			_lowercase (     SCREAMING_SNAKE_CASE_			):






				'''simple docstring'''



				SCREAMING_SNAKE_CASE__							:     int         =						VOCAB_FILES_NAMES
				SCREAMING_SNAKE_CASE__							:     Union[str, Any]         =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				SCREAMING_SNAKE_CASE__							:     str         =						PRETRAINED_VOCAB_FILES_MAP
				SCREAMING_SNAKE_CASE__							:     str         =						['''input_ids''', '''attention_mask''']

				SCREAMING_SNAKE_CASE__							:     Any         =						[]
				SCREAMING_SNAKE_CASE__							:     List[Any]         =						[]


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

										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						language_codes
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						FAIRSEQ_LANGUAGE_CODES[language_codes]
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						{lang_code: f'''__{lang_code}__''' for lang_code in fairseq_language_code}

										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						kwargs.get('''additional_special_tokens'''     ,     []      )
										kwargs["additional_special_tokens"] += [
										    self.get_lang_token(__SCREAMING_SNAKE_CASE      )
										    for lang_code in fairseq_language_code
										    if self.get_lang_token(__SCREAMING_SNAKE_CASE      ) not in kwargs["additional_special_tokens"]
										]

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

										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						vocab_file
										__SCREAMING_SNAKE_CASE		:	Dict						     =						load_json(__SCREAMING_SNAKE_CASE      )
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						{v: k for k, v in self.encoder.items()}
										__SCREAMING_SNAKE_CASE		:	Any						     =						spm_file
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						load_spm(__SCREAMING_SNAKE_CASE     ,     self.sp_model_kwargs      )

										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						len(self.encoder      )

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

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

										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						num_madeup_words


				@property
				def 						__magic_name__(			self     :List[str]      )						->     int:
										return len(self.encoder      ) + len(self.lang_token_to_id      )


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


				@src_lang.setter
				def 						__magic_name__(			self     :List[Any]     ,     lowerCAmelCase__     :str      )						->     None:
										__SCREAMING_SNAKE_CASE		:	str						     =						new_src_lang
										self.set_src_lang_special_tokens(self._src_lang      )


				def 						__magic_name__(			self     :Optional[Any]     ,     lowerCAmelCase__     :str      )						->     List[str]:
										return self.sp_model.encode(__SCREAMING_SNAKE_CASE     ,     out_type=__SCREAMING_SNAKE_CASE      )


				def 						__magic_name__(			self     :Union[str, Any]     ,     lowerCAmelCase__     :Tuple      )						->     Union[str, Any]:
										if token in self.lang_token_to_id:
																return self.lang_token_to_id[token]
										return self.encoder.get(__SCREAMING_SNAKE_CASE     ,     self.encoder[self.unk_token]      )


				def 						__magic_name__(			self     :str     ,     lowerCAmelCase__     :int      )						->     str:
										if index in self.id_to_lang_token:
																return self.id_to_lang_token[index]
										return self.decoder.get(__SCREAMING_SNAKE_CASE     ,     self.unk_token      )


				def 						__magic_name__(			self     :Any     ,     lowerCAmelCase__     :List[str]      )						->     Optional[Any]:
										__SCREAMING_SNAKE_CASE		:	Any						     =						[]
										__SCREAMING_SNAKE_CASE		:	Optional[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(__SCREAMING_SNAKE_CASE      ) + token
																						__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						[]
																else:
																						current_sub_tokens.append(__SCREAMING_SNAKE_CASE      )
										out_string += self.sp_model.decode(__SCREAMING_SNAKE_CASE      )
										return out_string.strip()


				def 						__magic_name__(			self     :int     ,     lowerCAmelCase__     :List[int]     ,     lowerCAmelCase__     :Optional[List[int]] = None     ,     lowerCAmelCase__     :bool = False      )						->     List[int]:

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

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


				def 						__magic_name__(			self     :int     ,     lowerCAmelCase__     :List[int]     ,     lowerCAmelCase__     :Optional[List[int]] = None      )						->     List[int]:
										if token_ids_a is None:
																return self.prefix_tokens + token_ids_a + self.suffix_tokens
										# We don't expect to process pairs, but leave the pair logic for API consistency
										return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens


				def 						__magic_name__(			self     :Optional[int]      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						{self.convert_ids_to_tokens(__SCREAMING_SNAKE_CASE      ): i for i in range(self.vocab_size      )}
										vocab.update(self.added_tokens_encoder      )
										return vocab


				def __getstate__(			self     :int      )						->     Dict:
										__SCREAMING_SNAKE_CASE		:	Optional[Any]						     =						self.__dict__.copy()
										__SCREAMING_SNAKE_CASE		:	Any						     =						None
										return state


				def __setstate__(			self     :str     ,     lowerCAmelCase__     :Dict      )						->     None:
										__SCREAMING_SNAKE_CASE		:	int						     =						d

										# for backward compatibility
										if not hasattr(self     ,     '''sp_model_kwargs'''      ):
																__SCREAMING_SNAKE_CASE		:	Any						     =						{}

										__SCREAMING_SNAKE_CASE		:	Tuple						     =						load_spm(self.spm_file     ,     self.sp_model_kwargs      )


				def 						__magic_name__(			self     :List[Any]     ,     lowerCAmelCase__     :str     ,     lowerCAmelCase__     :Optional[str] = None      )						->     Tuple[str]:
										__SCREAMING_SNAKE_CASE		:	Dict						     =						Path(__SCREAMING_SNAKE_CASE      )
										if not save_dir.is_dir():
																raise OSError(f'''{save_directory} should be a directory'''      )
										__SCREAMING_SNAKE_CASE		:	int						     =						save_dir / (
										    (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''vocab_file''']
										)
										__SCREAMING_SNAKE_CASE		:	int						     =						save_dir / (
										    (filename_prefix + '''-''' if filename_prefix else '''''') + self.vocab_files_names['''spm_file''']
										)

										save_json(self.encoder     ,     __SCREAMING_SNAKE_CASE      )

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

										return (str(__SCREAMING_SNAKE_CASE      ), str(__SCREAMING_SNAKE_CASE      ))


				def 						__magic_name__(			self     :Union[str, Any]     ,     lowerCAmelCase__     :List[str]     ,     lowerCAmelCase__     :str = "en"     ,     lowerCAmelCase__     :Optional[List[str]] = None     ,     lowerCAmelCase__     :str = "ro"     ,     **lowerCAmelCase__     :Optional[Any]     ,     )						->     BatchEncoding:
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						src_lang
										__SCREAMING_SNAKE_CASE		:	Dict						     =						tgt_lang
										self.set_src_lang_special_tokens(self.src_lang      )
										return super().prepare_seqaseq_batch(__SCREAMING_SNAKE_CASE     ,     __SCREAMING_SNAKE_CASE     ,     **__SCREAMING_SNAKE_CASE      )


				def 						__magic_name__(			self     :Optional[Any]     ,     lowerCAmelCase__     :Optional[Any]     ,     lowerCAmelCase__     :Optional[str]     ,     lowerCAmelCase__     :Optional[str]     ,     **lowerCAmelCase__     :Optional[Any]      )						->     Optional[Any]:
										if src_lang is None or tgt_lang is None:
																raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''      )
										__SCREAMING_SNAKE_CASE		:	Any						     =						src_lang
										__SCREAMING_SNAKE_CASE		:	Union[str, Any]						     =						self(__SCREAMING_SNAKE_CASE     ,     add_special_tokens=__SCREAMING_SNAKE_CASE     ,     **__SCREAMING_SNAKE_CASE      )
										__SCREAMING_SNAKE_CASE		:	Any						     =						self.get_lang_id(__SCREAMING_SNAKE_CASE      )
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						tgt_lang_id
										return inputs


				def 						__magic_name__(			self     :List[str]      )						->     Optional[Any]:
										self.set_src_lang_special_tokens(self.src_lang      )


				def 						__magic_name__(			self     :int      )						->     Optional[Any]:
										self.set_tgt_lang_special_tokens(self.tgt_lang      )


				def 						__magic_name__(			self     :Optional[Any]     ,     lowerCAmelCase__     :str      )						->     None:
										__SCREAMING_SNAKE_CASE		:	int						     =						self.get_lang_token(__SCREAMING_SNAKE_CASE      )
										__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						self.lang_token_to_id[lang_token]
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						[self.cur_lang_id]
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						[self.eos_token_id]


				def 						__magic_name__(			self     :Dict     ,     lowerCAmelCase__     :str      )						->     None:
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						self.get_lang_token(__SCREAMING_SNAKE_CASE      )
										__SCREAMING_SNAKE_CASE		:	Tuple						     =						self.lang_token_to_id[lang_token]
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						[self.cur_lang_id]
										__SCREAMING_SNAKE_CASE		:	List[Any]						     =						[self.eos_token_id]


				def 						__magic_name__(			self     :Optional[int]     ,     lowerCAmelCase__     :str      )						->     str:
										return self.lang_code_to_token[lang]



				def 						__magic_name__(			self     :Union[str, Any]     ,     lowerCAmelCase__     :str      )						->     int:
										__SCREAMING_SNAKE_CASE		:	List[str]						     =						self.get_lang_token(__SCREAMING_SNAKE_CASE      )
										return self.lang_token_to_id[lang_token]




def 			_UpperCamelCase				( lowercase__       ,			lowercase__ ):
						__SCREAMING_SNAKE_CASE		:	Optional[int]						     =						sentencepiece.SentencePieceProcessor(**__A )
						spm.Load(str(__A ) )
						return spm


def 			_UpperCamelCase				( lowercase__ ):
						with open(__A       ,			'''r''' ) as f:
												return json.load(__A )


def 			_UpperCamelCase				( lowercase__       ,			lowercase__ ):
						with open(__A       ,			'''w''' ) as f:
												json.dump(__A       ,			__A       ,			indent=2 )








'''simple docstring'''







import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__lowerCAmelCase									=   object()

# For specifying empty leaf dict `{}`
__lowerCAmelCase									=   object()





def 		_UpperCAmelCase (				__A				: List[str]    ,		__A				: Tuple   ):
				a_		:	List[Any]									=   tuple((re.compile(x + '''$'''   ) for x in qs)   )
				for i in range(len(__A   ) - len(__A   ) + 1   ):
								a_		:	Union[str, Any]									=   [x.match(__A   ) for x, y in zip(__A    ,		ks[i:]   )]
								if matches and all(__A   ):
												return True
				return False





def 		_UpperCAmelCase (				__A				: List[str]   ):
				def replace(__A				: int    ,		__A				: Union[str, Any]   ):
								for rule, replacement in rules:
												if _match(__A    ,		__A   ):
																return replacement
								return val

				return replace





def 		_UpperCAmelCase (				):
				return [
				    # embeddings
				    (("transformer", "wpe", "embedding"), P('''mp'''    ,		__A   )),
				    (("transformer", "wte", "embedding"), P('''mp'''    ,		__A   )),
				    # atention
				    (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(__A    ,		'''mp'''   )),
				    (("attention", "out_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("attention", "out_proj", "bias"), None),
				    # mlp
				    (("mlp", "c_fc", "kernel"), P(__A    ,		'''mp'''   )),
				    (("mlp", "c_fc", "bias"), P('''mp'''   )),
				    (("mlp", "c_proj", "kernel"), P('''mp'''    ,		__A   )),
				    (("mlp", "c_proj", "bias"), None),
				    # layer norms
				    ((r"ln_\d+", "bias"), None),
				    ((r"\d+", r"ln_\d+", "scale"), None),
				    (("ln_f", "bias"), None),
				    (("ln_f", "scale"), None),
				]





def 		_UpperCAmelCase (				__A				: Union[str, Any]   ):
				a_		:	Tuple									=   _get_partition_rules()
				a_		:	Tuple									=   _replacement_rules(__A   )
				a_		:	Optional[Any]									=   {k: _unmatched for k in flatten_dict(__A   )}
				a_		:	Optional[Any]									=   {k: replace(__A    ,		__A   ) for k, v in initd.items()}
				assert _unmatched not in result.values(), "Incomplete partition spec."
				return freeze(unflatten_dict(__A   )   )




