




"""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							:			Tuple     							=							logging.get_logger(__name__)

lowerCAmelCase							:			Any     							=							"""▁"""

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

lowerCAmelCase							:			List[str]     							=							{
    """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							:			Optional[Any]     							=							{
    """facebook/m2m100_418M""": 1024,
}

# 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     __magic_name__							(       UpperCAmelCase__				):






				'''simple docstring'''
				__UpperCamelCase			  =						VOCAB_FILES_NAMES
				__UpperCamelCase			  =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				__UpperCamelCase			  =						PRETRAINED_VOCAB_FILES_MAP
				__UpperCamelCase			  =						["input_ids", "attention_mask"]

				__UpperCamelCase			  =						[]
				__UpperCamelCase			  =						[]





				def __init__(       self		, _a		, _a		, _a=None		, _a=None		, _a="<s>"		, _a="</s>"		, _a="</s>"		, _a="<pad>"		, _a="<unk>"		, _a="m2m100"		, _a = None		, _a=8		, **_a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					{} if sp_model_kwargs is None else sp_model_kwargs

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

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

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

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

							lowerCamelCase			      =					len(self.encoder )

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

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

							lowerCamelCase			      =					num_madeup_words





				@property
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							return len(self.encoder ) + len(self.lang_token_to_id )





				@property
				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							return self._src_lang





				@src_lang.setter
				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					new_src_lang
							self.set_src_lang_special_tokens(self._src_lang )





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							return self.sp_model.encode(_a		, out_type=_a )





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							if token in self.lang_token_to_id:
										return self.lang_token_to_id[token]
							return self.encoder.get(_a		, self.encoder[self.unk_token] )





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							if index in self.id_to_lang_token:
										return self.id_to_lang_token[index]
							return self.decoder.get(_a		, self.unk_token )





				def    _lowerCAmelCase				(       self		, _a ):





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





				def    _lowerCAmelCase				(       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    _lowerCAmelCase				(       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    _lowerCAmelCase				(       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 __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			      =					load_spm(self.spm_file		, self.sp_model_kwargs )





				def    _lowerCAmelCase				(       self		, _a		, _a = None ):





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

							save_json(self.encoder		, _a )

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

							return (str(_a ), str(_a ))





				def    _lowerCAmelCase				(       self		, _a		, _a = "en"		, _a = None		, _a = "ro"		, **_a		, ):





							"""simple docstring"""
							lowerCamelCase			      =					src_lang
							lowerCamelCase			      =					tgt_lang
							self.set_src_lang_special_tokens(self.src_lang )
							return super().prepare_seqaseq_batch(_a		, _a		, **_a )





				def    _lowerCAmelCase				(       self		, _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		, **_a )
							lowerCamelCase			      =					self.get_lang_id(_a )
							lowerCamelCase			      =					tgt_lang_id
							return inputs





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							self.set_src_lang_special_tokens(self.src_lang )





				def    _lowerCAmelCase				(       self ):





							"""simple docstring"""
							self.set_tgt_lang_special_tokens(self.tgt_lang )





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.get_lang_token(_a )
							lowerCamelCase			      =					self.lang_token_to_id[lang_token]
							lowerCamelCase			      =					[self.cur_lang_id]
							lowerCamelCase			      =					[self.eos_token_id]





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.get_lang_token(_a )
							lowerCamelCase			      =					self.lang_token_to_id[lang_token]
							lowerCamelCase			      =					[self.cur_lang_id]
							lowerCamelCase			      =					[self.eos_token_id]





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							return self.lang_code_to_token[lang]





				def    _lowerCAmelCase				(       self		, _a ):





							"""simple docstring"""
							lowerCamelCase			      =					self.get_lang_token(_a )
							return self.lang_token_to_id[lang_token]




def     a__ (				snake_case__       , snake_case__				)						->					sentencepiece.SentencePieceProcessor:
			lowerCamelCase			      =					sentencepiece.SentencePieceProcessor(**snake_case__				)
			spm.Load(str(snake_case__				)				)
			return spm



def     a__ (				snake_case__				)						->					Union[Dict, List]:
			with open(snake_case__       , """r"""				) as f:
						return json.load(snake_case__				)



def     a__ (				snake_case__       , snake_case__				)						->					None:
			with open(snake_case__       , """w"""				) as f:
						json.dump(snake_case__       , snake_case__       , indent=2				)









"""simple docstring"""




from pickle import UnpicklingError

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict

from ..utils import logging


lowerCAmelCase							:			Tuple     							=							logging.get_logger(__name__)



def     a__ (				snake_case__       , snake_case__				)						->					Tuple:
			try:
						with open(snake_case__       , """rb"""				) as flax_state_f:
									lowerCamelCase			      =					from_bytes(snake_case__       , flax_state_f.read()				)
			except UnpicklingError as e:
						try:
									with open(snake_case__				) as f:
												if f.read().startswith("""version"""				):
															raise OSError(
															    """You seem to have cloned a repository without having git-lfs installed. Please"""
															    """ install git-lfs and run `git lfs install` followed by `git lfs pull` in the"""
															    """ folder you cloned."""				)
												else:
															raise ValueError from e
						except (UnicodeDecodeError, ValueError):
									raise EnvironmentError(F'Unable to convert {model_file} to Flax deserializable object. '				)

			return load_flax_weights_in_pytorch_model(snake_case__       , snake_case__				)



def     a__ (				snake_case__       , snake_case__				)						->					Tuple:

			try:
						import torch  # noqa: F401
			except ImportError:
						logger.error(
						    """Loading Flax weights in PyTorch requires both PyTorch and Flax to be installed. Please see"""
						    """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
						    """ instructions."""				)
						raise

			# check if we have bf16 weights
			lowerCamelCase			      =					flatten_dict(jax.tree_util.tree_map(lambda snake_case__				: x.dtype == jnp.bfloataa       , snake_case__				)				).values()
			if any(snake_case__				):
						# convert all weights to fp32 if they are bf16 since torch.from_numpy can-not handle bf16

						# and bf16 is not fully supported in PT yet.
						logger.warning(
						    """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
						    """before loading those in PyTorch model."""				)
						lowerCamelCase			      =					jax.tree_util.tree_map(
						    lambda snake_case__				: params.astype(np.floataa				) if params.dtype == jnp.bfloataa else params       , snake_case__				)

			lowerCamelCase			      =					""""""

			lowerCamelCase			      =					flatten_dict(snake_case__       , sep="""."""				)
			lowerCamelCase			      =					pt_model.state_dict()

			# keep track of unexpected & missing keys
			lowerCamelCase			      =					[]
			lowerCamelCase			      =					set(pt_model_dict.keys()				)

			for flax_key_tuple, flax_tensor in flax_state_dict.items():
						lowerCamelCase			      =					flax_key_tuple.split("""."""				)

						if flax_key_tuple_array[-1] == "kernel" and flax_tensor.ndim == 4:
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]
									lowerCamelCase			      =					jnp.transpose(snake_case__       , (3, 2, 0, 1)				)
						elif flax_key_tuple_array[-1] == "kernel":
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]
									lowerCamelCase			      =					flax_tensor.T
						elif flax_key_tuple_array[-1] == "scale":
									lowerCamelCase			      =					flax_key_tuple_array[:-1] + ["""weight"""]

						if "time_embedding" not in flax_key_tuple_array:
									for i, flax_key_tuple_string in enumerate(snake_case__				):
												lowerCamelCase			      =					(
												    flax_key_tuple_string.replace("""_0"""       , """.0"""				)
												    .replace("""_1"""       , """.1"""				)
												    .replace("""_2"""       , """.2"""				)
												    .replace("""_3"""       , """.3"""				)
												    .replace("""_4"""       , """.4"""				)
												    .replace("""_5"""       , """.5"""				)
												    .replace("""_6"""       , """.6"""				)
												    .replace("""_7"""       , """.7"""				)
												    .replace("""_8"""       , """.8"""				)
												    .replace("""_9"""       , """.9"""				)
												)

						lowerCamelCase			      =					""".""".join(snake_case__				)

						if flax_key in pt_model_dict:
									if flax_tensor.shape != pt_model_dict[flax_key].shape:
												raise ValueError(
												    F'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
												    F'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'				)
									else:
												# add weight to pytorch dict
												lowerCamelCase			      =					np.asarray(snake_case__				) if not isinstance(snake_case__       , np.ndarray				) else flax_tensor
												lowerCamelCase			      =					torch.from_numpy(snake_case__				)
												# remove from missing keys
												missing_keys.remove(snake_case__				)
						else:
									# weight is not expected by PyTorch model
									unexpected_keys.append(snake_case__				)

			pt_model.load_state_dict(snake_case__				)

			# re-transform missing_keys to list
			lowerCamelCase			      =					list(snake_case__				)

			if len(snake_case__				) > 0:
						logger.warning(
						    """Some weights of the Flax model were not used when initializing the PyTorch model"""
						    F' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
						    F' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
						    """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
						    F' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
						    """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
						    """ FlaxBertForSequenceClassification model)."""				)
			if len(snake_case__				) > 0:
						logger.warning(
						    F'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
						    F' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
						    """ use it for predictions and inference."""				)

			return pt_model



