
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[str] 				=					"""0.12"""  # assumed parallelism: 8

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






			'''simple docstring'''

			@classmethod
			def  lowerCAmelCase__		(    cls	:    Tuple ):
							UpperCamelCase_:							str		=  TOKEN
							HfFolder.save_token(lowercase_ )

			@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	:    List[str] ):
							UpperCamelCase_:							Dict		=  BertConfig(
							    vocab_size=99  ,		hidden_size=32  ,		num_hidden_layers=5  ,		num_attention_heads=4  ,		intermediate_size=37 )
							UpperCamelCase_:							int		=  FlaxBertModel(lowercase_ )
							model.push_to_hub("""test-model-flax"""  ,		use_auth_token=self._token )

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

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

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

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

							UpperCamelCase_:							Optional[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(lowercase_  ,		1e-3  ,		msg=f'''{key} not identical''' )




			def  lowerCAmelCase__		(    self	:    Tuple ):
							UpperCamelCase_:							Tuple		=  BertConfig(
							    vocab_size=99  ,		hidden_size=32  ,		num_hidden_layers=5  ,		num_attention_heads=4  ,		intermediate_size=37 )
							UpperCamelCase_:							Union[str, Any]		=  FlaxBertModel(lowercase_ )
							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_:							Any		=  flatten_dict(unfreeze(model.params ) )
							UpperCamelCase_:							List[Any]		=  flatten_dict(unfreeze(new_model.params ) )

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

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

							UpperCamelCase_:							int		=  flatten_dict(unfreeze(model.params ) )
							UpperCamelCase_:							Dict		=  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(lowercase_  ,		1e-3  ,		msg=f'''{key} not identical''' )







def     A__   (  lowerCamelCase       ,      lowerCamelCase       )					->		str:
				UpperCamelCase_:							Optional[int]		=  True
				UpperCamelCase_:							Union[str, Any]		=  flatten_dict(modela.params       )
				UpperCamelCase_:							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:
												UpperCamelCase_:							int		=  False

				return models_are_equal





@require_flax
class     _UpperCamelCase			(   unittest.TestCase					):






			'''simple docstring'''

			def  lowerCAmelCase__		(    self	:    Any ):
							UpperCamelCase_:							Optional[int]		=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only""" )
							UpperCamelCase_:							Any		=  FlaxBertModel(lowercase_ )

							UpperCamelCase_:							Optional[Any]		=  "bert"
							with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(lowercase_  ,		lowercase_ ) )

											with self.assertRaises(lowercase_ ):
															UpperCamelCase_:							Any		=  FlaxBertModel.from_pretrained(lowercase_ )

											UpperCamelCase_:							int		=  FlaxBertModel.from_pretrained(lowercase_  ,		subfolder=lowercase_ )

							self.assertTrue(check_models_equal(lowercase_  ,		lowercase_ ) )

			def  lowerCAmelCase__		(    self	:    Tuple ):
							UpperCamelCase_:							Union[str, Any]		=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only""" )
							UpperCamelCase_:							int		=  FlaxBertModel(lowercase_ )

							UpperCamelCase_:							Tuple		=  "bert"
							with tempfile.TemporaryDirectory() as tmp_dir:
											model.save_pretrained(os.path.join(lowercase_  ,		lowercase_ )  ,		max_shard_size="""10KB""" )

											with self.assertRaises(lowercase_ ):
															UpperCamelCase_:							Optional[Any]		=  FlaxBertModel.from_pretrained(lowercase_ )

											UpperCamelCase_:							Tuple		=  FlaxBertModel.from_pretrained(lowercase_  ,		subfolder=lowercase_ )

							self.assertTrue(check_models_equal(lowercase_  ,		lowercase_ ) )

			def  lowerCAmelCase__		(    self	:    List[Any] ):
							UpperCamelCase_:							List[Any]		=  "bert"
							UpperCamelCase_:							List[str]		=  "hf-internal-testing/tiny-random-bert-subfolder"

							with self.assertRaises(lowercase_ ):
											UpperCamelCase_:							Optional[Any]		=  FlaxBertModel.from_pretrained(lowercase_ )

							UpperCamelCase_:							int		=  FlaxBertModel.from_pretrained(lowercase_  ,		subfolder=lowercase_ )

							self.assertIsNotNone(lowercase_ )




			def  lowerCAmelCase__		(    self	:    Any ):
							UpperCamelCase_:							Optional[Any]		=  "bert"
							UpperCamelCase_:							str		=  "hf-internal-testing/tiny-random-bert-sharded-subfolder"
							with self.assertRaises(lowercase_ ):
											UpperCamelCase_:							Optional[int]		=  FlaxBertModel.from_pretrained(lowercase_ )

							UpperCamelCase_:							Optional[int]		=  FlaxBertModel.from_pretrained(lowercase_  ,		subfolder=lowercase_ )

							self.assertIsNotNone(lowercase_ )








def     A__   (  lowerCamelCase  )					->		list:
				return [
				    txt[:a] + txt[a].upper() + txt[a + 1 :]
				    for a in range(len(lowerCamelCase  )  )
				    if txt[a].isalpha()
				]


if __name__ == "__main__":
							__import__("""doctest""").testmod()






