


"""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    :     Dict   						=		'0.12'  # assumed parallelism: 8





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







					"""simple docstring"""


					@classmethod
					def lowerCamelCase__							(       cls    : Dict  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[Any]      							=					TOKEN
												HfFolder.save_token(UpperCamelCase  )






					@classmethod
					def lowerCamelCase__							(       cls    : Any  ):
												'''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 lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Tuple      							=					BertConfig(
												    vocab_size=99 ,       hidden_size=32 ,       num_hidden_layers=5 ,       num_attention_heads=4 ,       intermediate_size=37  )
												__UpperCAmelCase   :						Any      							=					FlaxBertModel(UpperCamelCase  )
												model.push_to_hub("""test-model-flax""" ,       use_auth_token=self._token  )

												__UpperCAmelCase   :						Any      							=					FlaxBertModel.from_pretrained(f'''{USER}/test-model-flax'''  )

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

												for key in base_params.keys():
																			__UpperCAmelCase   :						List[str]      							=					(base_params[key] - new_params[key]).sum().item()
																			self.assertLessEqual(UpperCamelCase ,       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(UpperCamelCase ,       repo_id="""test-model-flax""" ,       push_to_hub=UpperCamelCase ,       use_auth_token=self._token  )

												__UpperCAmelCase   :						Union[str, Any]      							=					FlaxBertModel.from_pretrained(f'''{USER}/test-model-flax'''  )

												__UpperCAmelCase   :						List[str]      							=					flatten_dict(unfreeze(model.params  )  )
												__UpperCAmelCase   :						Union[str, Any]      							=					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(UpperCamelCase ,       1e-3 ,       msg=f'''{key} not identical'''  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






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

												__UpperCAmelCase   :						Any      							=					FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""  )

												__UpperCAmelCase   :						Tuple      							=					flatten_dict(unfreeze(model.params  )  )
												__UpperCAmelCase   :						Union[str, Any]      							=					flatten_dict(unfreeze(new_model.params  )  )

												for key in base_params.keys():
																			__UpperCAmelCase   :						Optional[int]      							=					(base_params[key] - new_params[key]).sum().item()
																			self.assertLessEqual(UpperCamelCase ,       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(
																			    UpperCamelCase ,       repo_id="""valid_org/test-model-flax-org""" ,       push_to_hub=UpperCamelCase ,       use_auth_token=self._token  )

												__UpperCAmelCase   :						Optional[Any]      							=					FlaxBertModel.from_pretrained("""valid_org/test-model-flax-org"""  )

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

												for key in base_params.keys():
																			__UpperCAmelCase   :						Optional[int]      							=					(base_params[key] - new_params[key]).sum().item()
																			self.assertLessEqual(UpperCamelCase ,       1e-3 ,       msg=f'''{key} not identical'''  )






def      lowerCamelCase       (     _UpperCamelCase		:      str       ,			_UpperCamelCase		:      int							)		->       Any:

							'''simple docstring'''



							__UpperCAmelCase   :						int      							=					True
							__UpperCAmelCase   :						Tuple      							=					flatten_dict(modela.params							)
							__UpperCAmelCase   :						str      							=					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      lowerCamelCase__   (     unittest.TestCase				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : Tuple  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Any      							=					BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""  )
												__UpperCAmelCase   :						Tuple      							=					FlaxBertModel(UpperCamelCase  )

												__UpperCAmelCase   :						int      							=					"""bert"""
												with tempfile.TemporaryDirectory() as tmp_dir:
																			model.save_pretrained(os.path.join(UpperCamelCase ,       UpperCamelCase  )  )

																			with self.assertRaises(UpperCamelCase  ):
																										__UpperCAmelCase   :						Any      							=					FlaxBertModel.from_pretrained(UpperCamelCase  )

																			__UpperCAmelCase   :						Dict      							=					FlaxBertModel.from_pretrained(UpperCamelCase ,       subfolder=UpperCamelCase  )

												self.assertTrue(check_models_equal(UpperCamelCase ,       UpperCamelCase  )  )






					def lowerCamelCase__							(       self    : Optional[int]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						List[str]      							=					BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""  )
												__UpperCAmelCase   :						Tuple      							=					FlaxBertModel(UpperCamelCase  )

												__UpperCAmelCase   :						Optional[Any]      							=					"""bert"""
												with tempfile.TemporaryDirectory() as tmp_dir:
																			model.save_pretrained(os.path.join(UpperCamelCase ,       UpperCamelCase  ) ,       max_shard_size="""10KB"""  )

																			with self.assertRaises(UpperCamelCase  ):
																										__UpperCAmelCase   :						List[str]      							=					FlaxBertModel.from_pretrained(UpperCamelCase  )

																			__UpperCAmelCase   :						Optional[int]      							=					FlaxBertModel.from_pretrained(UpperCamelCase ,       subfolder=UpperCamelCase  )

												self.assertTrue(check_models_equal(UpperCamelCase ,       UpperCamelCase  )  )






					def lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






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

												with self.assertRaises(UpperCamelCase  ):
																			__UpperCAmelCase   :						List[Any]      							=					FlaxBertModel.from_pretrained(UpperCamelCase  )

												__UpperCAmelCase   :						Union[str, Any]      							=					FlaxBertModel.from_pretrained(UpperCamelCase ,       subfolder=UpperCamelCase  )

												self.assertIsNotNone(UpperCamelCase  )






					def lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Tuple      							=					"""bert"""
												__UpperCAmelCase   :						Dict      							=					"""hf-internal-testing/tiny-random-bert-sharded-subfolder"""
												with self.assertRaises(UpperCamelCase  ):
																			__UpperCAmelCase   :						Any      							=					FlaxBertModel.from_pretrained(UpperCamelCase  )

												__UpperCAmelCase   :						str      							=					FlaxBertModel.from_pretrained(UpperCamelCase ,       subfolder=UpperCamelCase  )

												self.assertIsNotNone(UpperCamelCase  )




"""simple docstring"""






# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


UpperCAmelCase    :     Union[str, Any]   						=		{'configuration_mra': ['MRA_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MraConfig']}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCAmelCase    :     int   						=		[
							    'MRA_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'MraForMaskedLM',
							    'MraForMultipleChoice',
							    'MraForQuestionAnswering',
							    'MraForSequenceClassification',
							    'MraForTokenClassification',
							    'MraLayer',
							    'MraModel',
							    'MraPreTrainedModel',
							]


if TYPE_CHECKING:
							from .configuration_mra import MRA_PRETRAINED_CONFIG_ARCHIVE_MAP, MraConfig

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_mra import (
														    MRA_PRETRAINED_MODEL_ARCHIVE_LIST,
														    MraForMaskedLM,
														    MraForMultipleChoice,
														    MraForQuestionAnswering,
														    MraForSequenceClassification,
														    MraForTokenClassification,
														    MraLayer,
														    MraModel,
														    MraPreTrainedModel,
														)
else:
							import sys

							UpperCAmelCase    :     Tuple   						=		_LazyModule(__name__, globals()['__file__'], _import_structure)
