




'''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

							a  :				List[str]									=  '0.12'  # assumed parallelism: 8






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


				@classmethod
				def   A_ (		cls    :					str					):
									snake_case_									=				TOKEN
									HfFolder.save_token(lowercase_					)


				@classmethod
				def   A_ (		cls    :					Optional[Any]					):
									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   A_ (		self    :					Tuple					):
									snake_case_									=				BertConfig(
									    vocab_size=99       ,				hidden_size=32       ,				num_hidden_layers=5       ,				num_attention_heads=4       ,				intermediate_size=37					)
									snake_case_									=				FlaxBertModel(lowercase_					)
									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(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					)

									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(lowercase_       ,				1e-3       ,				msg=F"{key} not identical"					)


				def   A_ (		self    :					Optional[Any]					):
									snake_case_									=				BertConfig(
									    vocab_size=99       ,				hidden_size=32       ,				num_hidden_layers=5       ,				num_attention_heads=4       ,				intermediate_size=37					)
									snake_case_									=				FlaxBertModel(lowercase_					)
									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(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					)

									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(lowercase_       ,				1e-3       ,				msg=F"{key} not identical"					)







def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					Union[str, Any]:

					'''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 						a				(   unittest.TestCase     ):


				def   A_ (		self    :					Union[str, Any]					):
									snake_case_									=				BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''					)
									snake_case_									=				FlaxBertModel(lowercase_					)

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

														with self.assertRaises(lowercase_					):
																			snake_case_									=				FlaxBertModel.from_pretrained(lowercase_					)

														snake_case_									=				FlaxBertModel.from_pretrained(lowercase_       ,				subfolder=lowercase_					)

									self.assertTrue(check_models_equal(lowercase_       ,				lowercase_					)					)


				def   A_ (		self    :					Union[str, Any]					):
									snake_case_									=				BertConfig.from_pretrained('''hf-internal-testing/tiny-bert-flax-only'''					)
									snake_case_									=				FlaxBertModel(lowercase_					)

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

														with self.assertRaises(lowercase_					):
																			snake_case_									=				FlaxBertModel.from_pretrained(lowercase_					)

														snake_case_									=				FlaxBertModel.from_pretrained(lowercase_       ,				subfolder=lowercase_					)

									self.assertTrue(check_models_equal(lowercase_       ,				lowercase_					)					)


				def   A_ (		self    :					str					):
									snake_case_									=				'''bert'''
									snake_case_									=				'''hf-internal-testing/tiny-random-bert-subfolder'''

									with self.assertRaises(lowercase_					):
														snake_case_									=				FlaxBertModel.from_pretrained(lowercase_					)

									snake_case_									=				FlaxBertModel.from_pretrained(lowercase_       ,				subfolder=lowercase_					)

									self.assertIsNotNone(lowercase_					)


				def   A_ (		self    :					Tuple					):
									snake_case_									=				'''bert'''
									snake_case_									=				'''hf-internal-testing/tiny-random-bert-sharded-subfolder'''
									with self.assertRaises(lowercase_					):
														snake_case_									=				FlaxBertModel.from_pretrained(lowercase_					)

									snake_case_									=				FlaxBertModel.from_pretrained(lowercase_       ,				subfolder=lowercase_					)

									self.assertIsNotNone(lowercase_					)




import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate


def   __SCREAMING_SNAKE_CASE					(SCREAMING_SNAKE_CASE__       ):
						return x + 2





class 	snake_case_   (		unittest.TestCase					):







			'''simple docstring'''

			def 				snake_case__(     self :				Optional[Any]							)      ->int:
									snake_case_   =	'''x = 3'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3}							)

									snake_case_   =	'''x = y'''
									snake_case_   =	{'''y''': 5}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 5, '''y''': 5}							)







			def 				snake_case__(     self :				Dict							)      ->Optional[int]:
									snake_case_   =	'''y = add_two(x)'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)

									# Won't work without the tool
									with CaptureStdout() as out:
															snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result is None
									assert "tried to execute add_two" in out.out







			def 				snake_case__(     self :				Union[str, Any]							)      ->Dict:
									snake_case_   =	'''x = 3'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3}							)







			def 				snake_case__(     self :				Optional[int]							)      ->Optional[int]:
									snake_case_   =	'''test_dict = {\'x\': x, \'y\': add_two(x)}'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}							)







			def 				snake_case__(     self :				Dict							)      ->str:
									snake_case_   =	'''x = 3\ny = 5'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 5}							)







			def 				snake_case__(     self :				str							)      ->Tuple:
									snake_case_   =	'''text = f\'This is x: {x}.\''''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == "This is x: 3."
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''text''': '''This is x: 3.'''}							)







			def 				snake_case__(     self :				Optional[Any]							)      ->List[str]:
									snake_case_   =	'''if x <= 3:\n    y = 2\nelse:\n    y = 5'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 2
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 2}							)

									snake_case_   =	{'''x''': 8}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									# evaluate returns the value of the last assignment.
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 8, '''y''': 5}							)







			def 				snake_case__(     self :				str							)      ->str:
									snake_case_   =	'''test_list = [x, add_two(x)]'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									self.assertListEqual(_UpperCamelCase							,   [3, 5]							)
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_list''': [3, 5]}							)







			def 				snake_case__(     self :				Any							)      ->List[Any]:
									snake_case_   =	'''y = x'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {}							,   state=_UpperCamelCase							)
									assert result == 3
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''y''': 3}							)







			def 				snake_case__(     self :				Optional[int]							)      ->Dict:
									snake_case_   =	'''test_list = [x, add_two(x)]\ntest_list[1]'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_list''': [3, 5]}							)

									snake_case_   =	'''test_dict = {\'x\': x, \'y\': add_two(x)}\ntest_dict[\'y\']'''
									snake_case_   =	{'''x''': 3}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''add_two''': add_two}							,   state=_UpperCamelCase							)
									assert result == 5
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 3, '''test_dict''': {'''x''': 3, '''y''': 5}}							)







			def 				snake_case__(     self :				Optional[Any]							)      ->int:
									snake_case_   =	'''x = 0\nfor i in range(3):\n    x = i'''
									snake_case_   =	{}
									snake_case_   =	evaluate(_UpperCamelCase							,   {'''range''': range}							,   state=_UpperCamelCase							)
									assert result == 2
									self.assertDictEqual(_UpperCamelCase							,   {'''x''': 2, '''i''': 2}							)