





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_      							=     """0.12"""  # assumed parallelism: 8
@require_flax
@is_staging_test
class 			UpperCamelCase_		(  unittest.TestCase					):




					@classmethod
					def 	lowerCAmelCase			(		cls  )		->							List[str]:
										_snake_case									=				TOKEN
										HfFolder.save_token(lowerCAmelCase_  )




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

										_snake_case									=				FlaxBertModel.from_pretrained(F'''{USER}/test-model-flax'''  )

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

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

										_snake_case									=				FlaxBertModel.from_pretrained(F'''{USER}/test-model-flax'''  )

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

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






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

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

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

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

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

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

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





def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[int]       ,    UpperCamelCase__			:       Tuple						)			-> int:



					'''simple docstring'''




					_snake_case									=				True
					_snake_case									=				flatten_dict(modela.params						)
					_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:
															_snake_case									=				False

					return models_are_equal




@require_flax
class 			UpperCamelCase_		(  unittest.TestCase					):




					def 	lowerCAmelCase			(		self  )		->							str:
										_snake_case									=				BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'  )
										_snake_case									=				FlaxBertModel(lowerCAmelCase_  )

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

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

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

										self.assertTrue(check_models_equal(lowerCAmelCase_				,						lowerCAmelCase_  )  )




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'  )
										_snake_case									=				FlaxBertModel(lowerCAmelCase_  )

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

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

										self.assertTrue(check_models_equal(lowerCAmelCase_				,						lowerCAmelCase_  )  )




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				'bert'
										_snake_case									=				'hf-internal-testing/tiny-random-bert-subfolder'

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

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

										self.assertIsNotNone(lowerCAmelCase_  )






					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				'bert'
										_snake_case									=				'hf-internal-testing/tiny-random-bert-sharded-subfolder'
										with self.assertRaises(lowerCAmelCase_  ):
															_snake_case									=				FlaxBertModel.from_pretrained(lowerCAmelCase_  )

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

										self.assertIsNotNone(lowerCAmelCase_  )










from collections import deque

from .hash_table import HashTable
class 			UpperCamelCase_		(  _lowerCamelCase					):




					def __init__(		self				,						*lowerCAmelCase_				,						**lowerCAmelCase_  )		->							Any:
										super().__init__(*lowerCAmelCase_				,						**lowerCAmelCase_  )




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Optional[int]:
										_snake_case									=				deque([]  ) if self.values[key] is None else self.values[key]
										self.values[key].appendleft(lowerCAmelCase_  )
										_snake_case									=				self.values[key]




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										return (
										    sum(self.charge_factor - len(lowerCAmelCase_  ) for slot in self.values  )
										    / self.size_table
										    * self.charge_factor
										)






					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_=None  )		->							str:
										if not (
										    len(self.values[key]  ) == self.charge_factor and self.values.count(lowerCAmelCase_  ) == 0
										):
															return key
										return super()._collision_resolution(lowerCAmelCase_				,						lowerCAmelCase_  )



