





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

       _lowerCamelCase      							=   '0.12'  # assumed parallelism: 8







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




      """simple docstring"""







      @classmethod
      def 					__snake_case      (		cls):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		TOKEN
             HfFolder.save_token(__snake_case)


      @classmethod
      def 					__snake_case      (		cls):



             """simple docstring"""





             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 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		BertConfig(
                 vocab_size=99   ,				hidden_size=32   ,				num_hidden_layers=5   ,				num_attention_heads=4   ,				intermediate_size=37)
             _lowerCamelCase				:  List[Any]						=		FlaxBertModel(__snake_case)
             model.push_to_hub('''test-model-flax'''   ,				use_auth_token=self._token)

             _lowerCamelCase				:  Optional[int]						=		FlaxBertModel.from_pretrained(F"""{USER}/test-model-flax""")

             _lowerCamelCase				:  Optional[int]						=		flatten_dict(unfreeze(model.params))
             _lowerCamelCase				:  Union[str, Any]						=		flatten_dict(unfreeze(new_model.params))

             for key in base_params.keys():
                    _lowerCamelCase				:  int						=		(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)

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

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

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



      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		BertConfig(
                 vocab_size=99   ,				hidden_size=32   ,				num_hidden_layers=5   ,				num_attention_heads=4   ,				intermediate_size=37)
             _lowerCamelCase				:  Optional[int]						=		FlaxBertModel(__snake_case)
             model.push_to_hub('''valid_org/test-model-flax-org'''   ,				use_auth_token=self._token)

             _lowerCamelCase				:  Tuple						=		FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org''')

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

             for key in base_params.keys():
                    _lowerCamelCase				:  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)

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

             _lowerCamelCase				:  int						=		flatten_dict(unfreeze(model.params))
             _lowerCamelCase				:  str						=		flatten_dict(unfreeze(new_model.params))

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




def  __UpperCAmelCase(   lowercase_			,     lowercase_   ):
       _lowerCamelCase				:  List[Any]						=		True
       _lowerCamelCase				:  Union[str, Any]						=		flatten_dict(modela.params   )
       _lowerCamelCase				:  Union[str, Any]						=		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:
                     _lowerCamelCase				:  Union[str, Any]						=		False

       return models_are_equal







@require_flax
class       __A							(   unittest.TestCase	):




      """simple docstring"""







      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  int						=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''')
             _lowerCamelCase				:  Tuple						=		FlaxBertModel(__snake_case)

             _lowerCamelCase				:  Tuple						=		"""bert"""
             with tempfile.TemporaryDirectory() as tmp_dir:
                    model.save_pretrained(os.path.join(__snake_case   ,				__snake_case))

                    with self.assertRaises(__snake_case):
                           _lowerCamelCase				:  Union[str, Any]						=		FlaxBertModel.from_pretrained(__snake_case)

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

             self.assertTrue(check_models_equal(__snake_case   ,				__snake_case))


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[Any]						=		BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only''')
             _lowerCamelCase				:  Optional[int]						=		FlaxBertModel(__snake_case)

             _lowerCamelCase				:  str						=		"""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):
                           _lowerCamelCase				:  Tuple						=		FlaxBertModel.from_pretrained(__snake_case)

                    _lowerCamelCase				:  Optional[int]						=		FlaxBertModel.from_pretrained(__snake_case   ,				subfolder=__snake_case)

             self.assertTrue(check_models_equal(__snake_case   ,				__snake_case))


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[int]						=		"""bert"""
             _lowerCamelCase				:  Tuple						=		"""hf-internal-testing/tiny-random-bert-subfolder"""

             with self.assertRaises(__snake_case):
                    _lowerCamelCase				:  Dict						=		FlaxBertModel.from_pretrained(__snake_case)

             _lowerCamelCase				:  int						=		FlaxBertModel.from_pretrained(__snake_case   ,				subfolder=__snake_case)

             self.assertIsNotNone(__snake_case)



      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[Any]						=		"""bert"""
             _lowerCamelCase				:  Union[str, Any]						=		"""hf-internal-testing/tiny-random-bert-sharded-subfolder"""
             with self.assertRaises(__snake_case):
                    _lowerCamelCase				:  Union[str, Any]						=		FlaxBertModel.from_pretrained(__snake_case)

             _lowerCamelCase				:  Dict						=		FlaxBertModel.from_pretrained(__snake_case   ,				subfolder=__snake_case)

             self.assertIsNotNone(__snake_case)





'''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_mbart import MBartTokenizer
else:
				UpperCAmelCase_						:  List[str]				  =						None


UpperCAmelCase_						:  Tuple				  =						logging.get_logger(__name__)


UpperCAmelCase_						:  List[Any]				  =						{'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

UpperCAmelCase_						:  Dict				  =						{
    'vocab_file': {
        'facebook/mbart-large-en-ro': (
            'https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/sentencepiece.bpe.model'
        ),
        'facebook/mbart-large-cc25': (
            'https://huggingface.co/facebook/mbart-large-cc25/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'facebook/mbart-large-en-ro': 'https://huggingface.co/facebook/mbart-large-en-ro/resolve/main/tokenizer.json',
        'facebook/mbart-large-cc25': 'https://huggingface.co/facebook/mbart-large-cc25/resolve/main/tokenizer.json',
    },
}

UpperCAmelCase_						:  str				  =						{
    'facebook/mbart-large-en-ro': 1024,
    'facebook/mbart-large-cc25': 1024,
}

# fmt: off
UpperCAmelCase_						:  List[str]				  =						['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']





class        lowercase__    (		_snake_case   ):





							'''simple docstring'''


							A_				: Any       =    VOCAB_FILES_NAMES
							A_				: List[str]       =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							A_				: Any       =    PRETRAINED_VOCAB_FILES_MAP
							A_				: Tuple       =    ["""input_ids""", """attention_mask"""]
							A_				: Dict       =    MBartTokenizer

							A_				: List[int]       =    []
							A_				: List[int]       =    []






							def __init__(       self       ,   __snake_case=None       ,   __snake_case=None       ,   __snake_case="<s>"       ,   __snake_case="</s>"       ,   __snake_case="</s>"       ,   __snake_case="<s>"       ,   __snake_case="<unk>"       ,   __snake_case="<pad>"       ,   __snake_case="<mask>"       ,   __snake_case=None       ,   __snake_case=None       ,   __snake_case=None       ,   **__snake_case       ,   ):
											# Mask token behave like a normal word, i.e. include the space before it
											_SCREAMING_SNAKE_CASE					:			List[Any]							= AddedToken(__snake_case       ,   lstrip=__snake_case       ,   rstrip=__snake_case    ) if isinstance(__snake_case       ,   __snake_case    ) else mask_token

											super().__init__(
											    vocab_file=__snake_case       ,   tokenizer_file=__snake_case       ,   bos_token=__snake_case       ,   eos_token=__snake_case       ,   sep_token=__snake_case       ,   cls_token=__snake_case       ,   unk_token=__snake_case       ,   pad_token=__snake_case       ,   mask_token=__snake_case       ,   src_lang=__snake_case       ,   tgt_lang=__snake_case       ,   additional_special_tokens=__snake_case       ,   **__snake_case       ,   )

											_SCREAMING_SNAKE_CASE					:			List[str]							= vocab_file
											_SCREAMING_SNAKE_CASE					:			Optional[int]							= False if not self.vocab_file else True

											_SCREAMING_SNAKE_CASE					:			int							= FAIRSEQ_LANGUAGE_CODES.copy()

											if additional_special_tokens is not None:
															# Only add those special tokens if they are not already there.
															_additional_special_tokens.extend(
															    [t for t in additional_special_tokens if t not in _additional_special_tokens]    )

											self.add_special_tokens({"""additional_special_tokens""": _additional_special_tokens}    )
											_SCREAMING_SNAKE_CASE					:			str							= {
											    lang_code: self.convert_tokens_to_ids(__snake_case    ) for lang_code in FAIRSEQ_LANGUAGE_CODES
											}

											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= src_lang if src_lang is not None else """en_XX"""
											_SCREAMING_SNAKE_CASE					:			str							= self.convert_tokens_to_ids(self._src_lang    )
											_SCREAMING_SNAKE_CASE					:			Any							= tgt_lang
											self.set_src_lang_special_tokens(self._src_lang    )






							@property
							def UpperCAmelCase_      (       self    ):
											return self._src_lang






							@src_lang.setter
							def UpperCAmelCase_      (       self       ,   __snake_case    ):
											_SCREAMING_SNAKE_CASE					:			Any							= new_src_lang
											self.set_src_lang_special_tokens(self._src_lang    )






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = None    ):
											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 UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = None    ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= [self.sep_token_id]
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= [self.cls_token_id]

											if token_ids_a is None:
															return len(cls + token_ids_a + sep    ) * [0]
											return len(cls + token_ids_a + sep + sep + token_ids_a + sep    ) * [0]






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   __snake_case       ,   **__snake_case    ):
											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					:			Dict							= src_lang
											_SCREAMING_SNAKE_CASE					:			List[Any]							= self(__snake_case       ,   add_special_tokens=__snake_case       ,   return_tensors=__snake_case       ,   **__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Dict							= self.convert_tokens_to_ids(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= tgt_lang_id
											return inputs






							def UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = "en_XX"       ,   __snake_case = None       ,   __snake_case = "ro_RO"       ,   **__snake_case       ,   ):
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= src_lang
											_SCREAMING_SNAKE_CASE					:			Union[str, Any]							= tgt_lang
											return super().prepare_seqaseq_batch(__snake_case       ,   __snake_case       ,   **__snake_case    )






							def UpperCAmelCase_      (       self    ):
											return self.set_src_lang_special_tokens(self.src_lang    )






							def UpperCAmelCase_      (       self    ):
											return self.set_tgt_lang_special_tokens(self.tgt_lang    )






							def UpperCAmelCase_      (       self       ,   __snake_case    ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= self.convert_tokens_to_ids(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			str							= []
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= [self.eos_token_id, self.cur_lang_code]

											_SCREAMING_SNAKE_CASE					:			int							= self.convert_ids_to_tokens(self.prefix_tokens    )
											_SCREAMING_SNAKE_CASE					:			Optional[Any]							= self.convert_ids_to_tokens(self.suffix_tokens    )

											_SCREAMING_SNAKE_CASE					:			Optional[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 UpperCAmelCase_      (       self       ,   __snake_case    ):
											_SCREAMING_SNAKE_CASE					:			List[str]							= self.convert_tokens_to_ids(__snake_case    )
											_SCREAMING_SNAKE_CASE					:			List[str]							= []
											_SCREAMING_SNAKE_CASE					:			Tuple							= [self.eos_token_id, self.cur_lang_code]

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

											_SCREAMING_SNAKE_CASE					:			Union[str, 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 UpperCAmelCase_      (       self       ,   __snake_case       ,   __snake_case = None    ):
											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(__snake_case    ):
															logger.error(f"""Vocabulary path ({save_directory}) should be a directory."""    )
															return
											_SCREAMING_SNAKE_CASE					:			List[str]							= os.path.join(
											    __snake_case       ,   (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]    )

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

											return (out_vocab_file,)




