
import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
					from transformers.models.bert.modeling_flax_bert import (
					    FlaxBertForMaskedLM,
					    FlaxBertForMultipleChoice,
					    FlaxBertForNextSentencePrediction,
					    FlaxBertForPreTraining,
					    FlaxBertForQuestionAnswering,
					    FlaxBertForSequenceClassification,
					    FlaxBertForTokenClassification,
					    FlaxBertModel,
					)







class 			__A( unittest.TestCase			):




				def __init__(     self		,     _snake_case		,     _snake_case=13		,     _snake_case=7		,     _snake_case=True		,     _snake_case=True		,     _snake_case=True		,     _snake_case=True		,     _snake_case=99		,     _snake_case=32		,     _snake_case=5		,     _snake_case=4		,     _snake_case=37		,     _snake_case="gelu"		,     _snake_case=0.1		,     _snake_case=0.1		,     _snake_case=512		,     _snake_case=16		,     _snake_case=2		,     _snake_case=0.02		,     _snake_case=4		,     )    ->      Tuple:


								'''simple docstring'''




								__a										=    parent
								__a										=    batch_size
								__a										=    seq_length
								__a										=    is_training
								__a										=    use_attention_mask
								__a										=    use_token_type_ids
								__a										=    use_labels
								__a										=    vocab_size
								__a										=    hidden_size
								__a										=    num_hidden_layers
								__a										=    num_attention_heads
								__a										=    intermediate_size
								__a										=    hidden_act
								__a										=    hidden_dropout_prob
								__a										=    attention_probs_dropout_prob
								__a										=    max_position_embeddings
								__a										=    type_vocab_size
								__a										=    type_sequence_label_size
								__a										=    initializer_range
								__a										=    num_choices




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Any:


								'''simple docstring'''




								__a										=    ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size				)

								__a										=    None
								if self.use_attention_mask:
												__a										=    random_attention_mask([self.batch_size, self.seq_length]				)

								__a										=    None
								if self.use_token_type_ids:
												__a										=    ids_tensor([self.batch_size, self.seq_length]		,     self.type_vocab_size				)

								__a										=    BertConfig(
								    vocab_size=self.vocab_size		,     hidden_size=self.hidden_size		,     num_hidden_layers=self.num_hidden_layers		,     num_attention_heads=self.num_attention_heads		,     intermediate_size=self.intermediate_size		,     hidden_act=self.hidden_act		,     hidden_dropout_prob=self.hidden_dropout_prob		,     attention_probs_dropout_prob=self.attention_probs_dropout_prob		,     max_position_embeddings=self.max_position_embeddings		,     type_vocab_size=self.type_vocab_size		,     is_decoder=_snake_case		,     initializer_range=self.initializer_range		,     )

								return config, input_ids, token_type_ids, attention_mask




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Union[str, Any]:


								'''simple docstring'''




								__a										=    self.prepare_config_and_inputs()
								__a							,     __a							,     __a							,     __a										=    config_and_inputs
								__a										=    {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
								return config, inputs_dict




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Dict:


								'''simple docstring'''




								__a										=    self.prepare_config_and_inputs()
								__a							,     __a							,     __a							,     __a										=    config_and_inputs

								__a										=    True
								__a										=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
								__a										=    ids_tensor([self.batch_size, self.seq_length]		,     vocab_size=2				)

								return (
								    config,
								    input_ids,
								    attention_mask,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)







@require_flax
class 			__A( a  ,   unittest.TestCase			):
				snake_case_						   =    True

				snake_case_						   =    (
				    (
				        FlaxBertModel,
				        FlaxBertForPreTraining,
				        FlaxBertForMaskedLM,
				        FlaxBertForMultipleChoice,
				        FlaxBertForQuestionAnswering,
				        FlaxBertForNextSentencePrediction,
				        FlaxBertForSequenceClassification,
				        FlaxBertForTokenClassification,
				        FlaxBertForQuestionAnswering,
				    )
				    if is_flax_available()
				    else ()
				)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Tuple:


								'''simple docstring'''




								__a										=    FlaxBertModelTester(self				)




				@slow
				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      List[Any]:


								'''simple docstring'''




								__a										=    FlaxBertModel.from_pretrained('''bert-base-cased'''				)
								__a										=    model(np.ones((1, 1)				)				)
								self.assertIsNotNone(_snake_case				)

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

					A     :	int    = '0.12'  # assumed parallelism: 8







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




				@classmethod
				def 				SCREAMING_SNAKE_CASE_		(     cls				)    ->      Union[str, Any]:


								'''simple docstring'''




								__a										=    TOKEN
								HfFolder.save_token(_snake_case				)




				@classmethod
				def 				SCREAMING_SNAKE_CASE_		(     cls				)    ->      Union[str, 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 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    BertConfig(
								    vocab_size=99		,     hidden_size=32		,     num_hidden_layers=5		,     num_attention_heads=4		,     intermediate_size=37				)
								__a										=    FlaxBertModel(_snake_case				)
								model.push_to_hub('''test-model-flax'''		,     use_auth_token=self._token				)

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

								__a										=    flatten_dict(unfreeze(model.params				)				)
								__a										=    flatten_dict(unfreeze(new_model.params				)				)

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

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

								__a										=    flatten_dict(unfreeze(model.params				)				)
								__a										=    flatten_dict(unfreeze(new_model.params				)				)

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




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[int]:


								'''simple docstring'''




								__a										=    BertConfig(
								    vocab_size=99		,     hidden_size=32		,     num_hidden_layers=5		,     num_attention_heads=4		,     intermediate_size=37				)
								__a										=    FlaxBertModel(_snake_case				)
								model.push_to_hub('''valid_org/test-model-flax-org'''		,     use_auth_token=self._token				)

								__a										=    FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''				)

								__a										=    flatten_dict(unfreeze(model.params				)				)
								__a										=    flatten_dict(unfreeze(new_model.params				)				)

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

								__a										=    FlaxBertModel.from_pretrained('''valid_org/test-model-flax-org'''				)

								__a										=    flatten_dict(unfreeze(model.params				)				)
								__a										=    flatten_dict(unfreeze(new_model.params				)				)

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



def     __lowerCAmelCase			(			a__						,    a__  )       ->  str:
				__a										=    True
				__a										=    flatten_dict(modela.params  )
				__a										=    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:
												__a										=    False

				return models_are_equal







@require_flax
class 			__A( unittest.TestCase			):




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Union[str, Any]:


								'''simple docstring'''




								__a										=    BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''				)
								__a										=    FlaxBertModel(_snake_case				)

								__a										=    '''bert'''
								with tempfile.TemporaryDirectory() as tmp_dir:
												model.save_pretrained(os.path.join(_snake_case		,     _snake_case				)				)

												with self.assertRaises(_snake_case				):
																__a										=    FlaxBertModel.from_pretrained(_snake_case				)

												__a										=    FlaxBertModel.from_pretrained(_snake_case		,     subfolder=_snake_case				)

								self.assertTrue(check_models_equal(_snake_case		,     _snake_case				)				)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[Any]:


								'''simple docstring'''




								__a										=    BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''				)
								__a										=    FlaxBertModel(_snake_case				)

								__a										=    '''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				):
																__a										=    FlaxBertModel.from_pretrained(_snake_case				)

												__a										=    FlaxBertModel.from_pretrained(_snake_case		,     subfolder=_snake_case				)

								self.assertTrue(check_models_equal(_snake_case		,     _snake_case				)				)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      int:


								'''simple docstring'''




								__a										=    '''bert'''
								__a										=    '''hf-internal-testing/tiny-random-bert-subfolder'''

								with self.assertRaises(_snake_case				):
												__a										=    FlaxBertModel.from_pretrained(_snake_case				)

								__a										=    FlaxBertModel.from_pretrained(_snake_case		,     subfolder=_snake_case				)

								self.assertIsNotNone(_snake_case				)




				def 				SCREAMING_SNAKE_CASE_		(     self				)    ->      Optional[int]:


								'''simple docstring'''




								__a										=    '''bert'''
								__a										=    '''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
								with self.assertRaises(_snake_case				):
												__a										=    FlaxBertModel.from_pretrained(_snake_case				)

								__a										=    FlaxBertModel.from_pretrained(_snake_case		,     subfolder=_snake_case				)

								self.assertIsNotNone(_snake_case				)