




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__			:       List[Any]         =							'0.12'  # assumed parallelism: 8





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


 '''simple docstring'''







 @classmethod
 def 				lowerCAmelCase_  (							cls  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_     						=      TOKEN
       HfFolder.save_token(_lowerCAmelCase						)



 @classmethod
 def 				lowerCAmelCase_  (							cls  :		List[str]						):
       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  :		Tuple						):
       SCREAMING_SNAKE_CASE_     						=      BertConfig(
           vocab_size=99    ,    hidden_size=32    ,    num_hidden_layers=5    ,    num_attention_heads=4    ,    intermediate_size=37						)
       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel(_lowerCAmelCase						)
       model.push_to_hub('test-model-flax'    ,    use_auth_token=self._token						)

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

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

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

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

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

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




 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      BertConfig(
           vocab_size=99    ,    hidden_size=32    ,    num_hidden_layers=5    ,    num_attention_heads=4    ,    intermediate_size=37						)
       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel(_lowerCAmelCase						)
       model.push_to_hub('valid_org/test-model-flax-org'    ,    use_auth_token=self._token						)

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

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

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

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

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

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


def  UpperCAmelCase_							(				__UpperCAmelCase						:     Optional[Any]      ,      __UpperCAmelCase						:     Dict							) ->			List[Any]:
      SCREAMING_SNAKE_CASE_     						=      True
      SCREAMING_SNAKE_CASE_     						=      flatten_dict(modela.params							)
      SCREAMING_SNAKE_CASE_     						=      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:
                  SCREAMING_SNAKE_CASE_     						=      False

      return models_are_equal





@require_flax
class 	lowerCamelCase_							(				unittest.TestCase ):


 '''simple docstring'''







 def 				lowerCAmelCase_  (							self  :		List[str]						):
       SCREAMING_SNAKE_CASE_     						=      BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'						)
       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel(_lowerCAmelCase						)

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

             with self.assertRaises(_lowerCAmelCase						):
                   SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase						)

             SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase    ,    subfolder=_lowerCAmelCase						)

       self.assertTrue(check_models_equal(_lowerCAmelCase    ,    _lowerCAmelCase						)						)



 def 				lowerCAmelCase_  (							self  :		Optional[int]						):
       SCREAMING_SNAKE_CASE_     						=      BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'						)
       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      'bert'
       with tempfile.TemporaryDirectory() as tmp_dir:
             model.save_pretrained(os.path.join(_lowerCAmelCase    ,    _lowerCAmelCase						)    ,    max_shard_size='10KB'						)

             with self.assertRaises(_lowerCAmelCase						):
                   SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase						)

             SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase    ,    subfolder=_lowerCAmelCase						)

       self.assertTrue(check_models_equal(_lowerCAmelCase    ,    _lowerCAmelCase						)						)



 def 				lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_     						=      'bert'
       SCREAMING_SNAKE_CASE_     						=      'hf-internal-testing/tiny-random-bert-subfolder'

       with self.assertRaises(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase    ,    subfolder=_lowerCAmelCase						)

       self.assertIsNotNone(_lowerCAmelCase						)




 def 				lowerCAmelCase_  (							self  :		Union[str, Any]						):
       SCREAMING_SNAKE_CASE_     						=      'bert'
       SCREAMING_SNAKE_CASE_     						=      'hf-internal-testing/tiny-random-bert-sharded-subfolder'
       with self.assertRaises(_lowerCAmelCase						):
             SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase						)

       SCREAMING_SNAKE_CASE_     						=      FlaxBertModel.from_pretrained(_lowerCAmelCase    ,    subfolder=_lowerCAmelCase						)

       self.assertIsNotNone(_lowerCAmelCase						)





from abc import ABC, abstractmethod
from argparse import ArgumentParser





class 	lowerCamelCase_							(				_SCREAMING_SNAKE_CASE ):


 '''simple docstring'''







 @staticmethod
 @abstractmethod
 def 				lowerCAmelCase_  (							_lowerCAmelCase  :		ArgumentParser						):
       raise NotImplementedError()




 @abstractmethod
 def 				lowerCAmelCase_  (							self  :		Dict						):
       raise NotImplementedError()