



'''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


UpperCamelCase__:   Tuple								=			logging.get_logger(__name__)

UpperCamelCase__:   Optional[Any]								=			"▁"

UpperCamelCase__:   Tuple								=			{
    "vocab_file": "vocab.json",
    "spm_file": "sentencepiece.bpe.model",
    "tokenizer_config_file": "tokenizer_config.json",
}

UpperCamelCase__:   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",
    },
}

UpperCamelCase__:   Dict								=			{
    "facebook/m2m100_418M": 1024,
}

# fmt: off
UpperCamelCase__:   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 				SCREAMING_SNAKE_CASE(							A__			):







	"""simple docstring"""






	lowerCamelCase__			 =       VOCAB_FILES_NAMES
	lowerCamelCase__			 =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	lowerCamelCase__			 =       PRETRAINED_VOCAB_FILES_MAP
	lowerCamelCase__			 =       ["""input_ids""", """attention_mask"""]

	lowerCamelCase__			 =       []
	lowerCamelCase__			 =       []







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

								UpperCAmelCase :       Tuple 							=		language_codes
								UpperCAmelCase :       Union[str, Any] 							=		FAIRSEQ_LANGUAGE_CODES[language_codes]
								UpperCAmelCase :       Union[str, Any] 							=		{lang_code: F"""__{lang_code}__""" for lang_code in fairseq_language_code}

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

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

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

								UpperCAmelCase :       Optional[Any] 							=		len(self.encoder		)

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

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

								UpperCAmelCase :       Optional[int] 							=		num_madeup_words







	@property
	def      A	(						self			:     int		)			->		int:
								return len(self.encoder		) + len(self.lang_token_to_id		)







	@property
	def      A	(						self			:     List[str]		)			->		str:
								return self._src_lang







	@src_lang.setter
	def      A	(						self			:     str		,			__snake_case			:     str		)			->		None:
								UpperCAmelCase :       Union[str, Any] 							=		new_src_lang
								self.set_src_lang_special_tokens(self._src_lang		)







	def      A	(						self			:     List[str]		,			__snake_case			:     str		)			->		List[str]:
								return self.sp_model.encode(__snake_case		,			out_type=__snake_case		)







	def      A	(						self			:     str		,			__snake_case			:     Tuple		)			->		Tuple:
								if token in self.lang_token_to_id:
															return self.lang_token_to_id[token]
								return self.encoder.get(__snake_case		,			self.encoder[self.unk_token]		)







	def      A	(						self			:     Any		,			__snake_case			:     int		)			->		str:
								if index in self.id_to_lang_token:
															return self.id_to_lang_token[index]
								return self.decoder.get(__snake_case		,			self.unk_token		)







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     Tuple		)			->		Tuple:
								UpperCAmelCase :       Optional[int] 							=		[]
								UpperCAmelCase :       Dict 							=		''''''
								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(__snake_case		) + token
																						UpperCAmelCase :       Optional[Any] 							=		[]
															else:
																						current_sub_tokens.append(__snake_case		)
								out_string += self.sp_model.decode(__snake_case		)
								return out_string.strip()







	def      A	(						self			:     Dict		,			__snake_case			:     List[int]		,			__snake_case			:     Optional[List[int]] = None		,			__snake_case			:     bool = False		)			->		List[int]:

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

								UpperCAmelCase :       Optional[Any] 							=		[1] * len(self.prefix_tokens		)
								UpperCAmelCase :       Union[str, Any] 							=		[1] * len(self.suffix_tokens		)
								if token_ids_a is None:
															return prefix_ones + ([0] * len(__snake_case		)) + suffix_ones
								return prefix_ones + ([0] * len(__snake_case		)) + ([0] * len(__snake_case		)) + suffix_ones







	def      A	(						self			:     Tuple		,			__snake_case			:     List[int]		,			__snake_case			:     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      A	(						self			:     Optional[Any]		)			->		Dict:
								UpperCAmelCase :       Optional[int] 							=		{self.convert_ids_to_tokens(__snake_case		): i for i in range(self.vocab_size		)}
								vocab.update(self.added_tokens_encoder		)
								return vocab







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







	def __setstate__(						self			:     List[Any]		,			__snake_case			:     Dict		)			->		None:
								UpperCAmelCase :       Optional[int] 							=		d

								# for backward compatibility
								if not hasattr(self		,			'''sp_model_kwargs'''		):
															UpperCAmelCase :       Optional[int] 							=		{}

								UpperCAmelCase :       Optional[Any] 							=		load_spm(self.spm_file		,			self.sp_model_kwargs		)







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

								save_json(self.encoder		,			__snake_case		)

								if os.path.abspath(self.spm_file		) != os.path.abspath(__snake_case		) and os.path.isfile(self.spm_file		):
															copyfile(self.spm_file		,			__snake_case		)
								elif not os.path.isfile(self.spm_file		):
															with open(__snake_case		,			'''wb'''		) as fi:
																						UpperCAmelCase :       Dict 							=		self.sp_model.serialized_model_proto()
																						fi.write(__snake_case		)

								return (str(__snake_case		), str(__snake_case		))







	def      A	(						self			:     int		,			__snake_case			:     List[str]		,			__snake_case			:     str = "en"		,			__snake_case			:     Optional[List[str]] = None		,			__snake_case			:     str = "ro"		,			**__snake_case			:     Optional[int]		,			)			->		BatchEncoding:
								UpperCAmelCase :       List[Any] 							=		src_lang
								UpperCAmelCase :       List[Any] 							=		tgt_lang
								self.set_src_lang_special_tokens(self.src_lang		)
								return super().prepare_seqaseq_batch(__snake_case		,			__snake_case		,			**__snake_case		)







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







	def      A	(						self			:     str		)			->		int:
								self.set_src_lang_special_tokens(self.src_lang		)







	def      A	(						self			:     Any		)			->		int:
								self.set_tgt_lang_special_tokens(self.tgt_lang		)







	def      A	(						self			:     Any		,			__snake_case			:     str		)			->		None:
								UpperCAmelCase :       Dict 							=		self.get_lang_token(__snake_case		)
								UpperCAmelCase :       str 							=		self.lang_token_to_id[lang_token]
								UpperCAmelCase :       Any 							=		[self.cur_lang_id]
								UpperCAmelCase :       str 							=		[self.eos_token_id]







	def      A	(						self			:     Dict		,			__snake_case			:     str		)			->		None:
								UpperCAmelCase :       Union[str, Any] 							=		self.get_lang_token(__snake_case		)
								UpperCAmelCase :       str 							=		self.lang_token_to_id[lang_token]
								UpperCAmelCase :       Optional[int] 							=		[self.cur_lang_id]
								UpperCAmelCase :       int 							=		[self.eos_token_id]







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     str		)			->		str:
								return self.lang_code_to_token[lang]







	def      A	(						self			:     Any		,			__snake_case			:     str		)			->		int:
								UpperCAmelCase :       Any 							=		self.get_lang_token(__snake_case		)
								return self.lang_token_to_id[lang_token]


def  snake_case_     (					_lowerCAmelCase  :						str   ,      _lowerCAmelCase  :						Dict[str, Any]  )      ->   sentencepiece.SentencePieceProcessor:
							UpperCAmelCase :       Any 							=		sentencepiece.SentencePieceProcessor(**_lowerCAmelCase  )
							spm.Load(str(_lowerCAmelCase  )  )
							return spm


def  snake_case_     (					_lowerCAmelCase  :						str  )      ->   Union[Dict, List]:
							with open(_lowerCAmelCase   ,      '''r'''  ) as f:
														return json.load(_lowerCAmelCase  )




def  snake_case_     (					_lowerCAmelCase  :						List[str]   ,      _lowerCAmelCase  :						str  )      ->   None:
							with open(_lowerCAmelCase   ,      '''w'''  ) as f:
														json.dump(_lowerCAmelCase   ,      _lowerCAmelCase   ,      indent=2  )










'''simple docstring'''






import tempfile
import unittest

import numpy as np
from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import TOKEN, USER, is_staging_test, require_flax


if is_flax_available():
						import os

						from flax.core.frozen_dict import unfreeze
						from flax.traverse_util import flatten_dict

						from transformers import FlaxBertModel

						UpperCamelCase__:   Tuple								=			"0.12"  # assumed parallelism: 8

@require_flax
@is_staging_test
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	@classmethod
	def      A	(						cls			:     Union[str, Any]		)			->		int:
								UpperCAmelCase :       Optional[Any] 							=		TOKEN
								HfFolder.save_token(__snake_case		)







	@classmethod
	def      A	(						cls			:     List[str]		)			->		Tuple:
								try:
															delete_repo(token=cls._token		,			repo_id='''test-model-flax'''		)
								except HTTPError:
															pass

								try:
															delete_repo(token=cls._token		,			repo_id='''valid_org/test-model-flax-org'''		)
								except HTTPError:
															pass







	def      A	(						self			:     int		)			->		Tuple:
								UpperCAmelCase :       List[Any] 							=		BertConfig(
								    vocab_size=99		,			hidden_size=32		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=37		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)
								model.push_to_hub('''test-model-flax'''		,			use_auth_token=self._token		)

								UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""		)

								UpperCAmelCase :       List[Any] 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       Tuple 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Union[str, Any] 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)

								# Reset repo
								delete_repo(token=self._token		,			repo_id='''test-model-flax'''		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(__snake_case		,			repo_id='''test-model-flax'''		,			push_to_hub=__snake_case		,			use_auth_token=self._token		)

								UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax"""		)

								UpperCAmelCase :       Any 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       str 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Optional[Any] 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)







	def      A	(						self			:     Optional[Any]		)			->		Union[str, Any]:
								UpperCAmelCase :       Dict 							=		BertConfig(
								    vocab_size=99		,			hidden_size=32		,			num_hidden_layers=5		,			num_attention_heads=4		,			intermediate_size=37		)
								UpperCAmelCase :       Optional[Any] 							=		FlaxBertModel(__snake_case		)
								model.push_to_hub('''valid_org/test-model-flax-org'''		,			use_auth_token=self._token		)

								UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''		)

								UpperCAmelCase :       List[Any] 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       int 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       Any 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)

								# Reset repo
								delete_repo(token=self._token		,			repo_id='''valid_org/test-model-flax-org'''		)

								# Push to hub via save_pretrained
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(
															    __snake_case		,			repo_id='''valid_org/test-model-flax-org'''		,			push_to_hub=__snake_case		,			use_auth_token=self._token		)

								UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''		)

								UpperCAmelCase :       Any 							=		flatten_dict(unfreeze(model.params		)		)
								UpperCAmelCase :       Optional[Any] 							=		flatten_dict(unfreeze(new_model.params		)		)

								for key in base_params.keys():
															UpperCAmelCase :       int 							=		(base_params[key] - new_params[key]).sum().item()
															self.assertLessEqual(__snake_case		,			1E-3		,			msg=F"""{key} not identical"""		)


def  snake_case_     (					_lowerCAmelCase  :						Optional[Any]   ,      _lowerCAmelCase  :						Any  )      ->   Union[str, Any]:
							UpperCAmelCase :       str 							=		True
							UpperCAmelCase :       int 							=		flatten_dict(modela.params  )
							UpperCAmelCase :       Dict 							=		flatten_dict(modela.params  )
							for key in flat_params_a.keys():
														if np.sum(np.abs(flat_params_a[key] - flat_params_a[key]  )  ) > 1e-4:
																					UpperCAmelCase :       Dict 							=		False

							return models_are_equal




@require_flax
class 				SCREAMING_SNAKE_CASE(							unittest.TestCase			):







	"""simple docstring"""






	def      A	(						self			:     Tuple		)			->		Union[str, Any]:
								UpperCAmelCase :       List[Any] 							=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)

								UpperCAmelCase :       int 							=		'''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(__snake_case		,			__snake_case		)		)

															with self.assertRaises(__snake_case		):
																						UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(__snake_case		)

															UpperCAmelCase :       str 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertTrue(check_models_equal(__snake_case		,			__snake_case		)		)







	def      A	(						self			:     List[str]		)			->		Dict:
								UpperCAmelCase :       Dict 							=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''		)
								UpperCAmelCase :       Dict 							=		FlaxBertModel(__snake_case		)

								UpperCAmelCase :       Optional[int] 							=		'''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(__snake_case		,			__snake_case		)		,			max_shard_size='''10KB'''		)

															with self.assertRaises(__snake_case		):
																						UpperCAmelCase :       Any 							=		FlaxBertModel.from_pretrained(__snake_case		)

															UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertTrue(check_models_equal(__snake_case		,			__snake_case		)		)







	def      A	(						self			:     Optional[int]		)			->		str:
								UpperCAmelCase :       Dict 							=		'''bert'''
								UpperCAmelCase :       int 							=		'''hf-internal-testing/tiny-random-bert-subfolder'''

								with self.assertRaises(__snake_case		):
															UpperCAmelCase :       Optional[Any] 							=		FlaxBertModel.from_pretrained(__snake_case		)

								UpperCAmelCase :       Tuple 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertIsNotNone(__snake_case		)







	def      A	(						self			:     Dict		)			->		List[Any]:
								UpperCAmelCase :       Optional[int] 							=		'''bert'''
								UpperCAmelCase :       int 							=		'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
								with self.assertRaises(__snake_case		):
															UpperCAmelCase :       Dict 							=		FlaxBertModel.from_pretrained(__snake_case		)

								UpperCAmelCase :       Union[str, Any] 							=		FlaxBertModel.from_pretrained(__snake_case		,			subfolder=__snake_case		)

								self.assertIsNotNone(__snake_case		)





