






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

					__snake_case      :      Tuple             =       """0.12"""  # assumed parallelism: 8


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






						@classmethod
						def  UpperCamelCase__ (		cls				):






												"""simple docstring"""
												lowerCAmelCase__        =  TOKEN
												HfFolder.save_token(_UpperCamelCase				)






						@classmethod
						def  UpperCamelCase__ (		cls				):






												"""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  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  BertConfig(
												    vocab_size=99       ,  hidden_size=32       ,  num_hidden_layers=5       ,  num_attention_heads=4       ,  intermediate_size=37				)
												lowerCAmelCase__        =  FlaxBertModel(_UpperCamelCase				)
												model.push_to_hub('test-model-flax'       ,  use_auth_token=self._token				)

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

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

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

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

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

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







						def  UpperCamelCase__ (		self				):






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

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

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

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

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

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

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







def 		_UpperCamelCase				(      UpperCamelCase_    :   Optional[int]      ,    UpperCamelCase_    :   List[Any]						)					->				int:

						"""simple docstring"""



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

						return models_are_equal


@require_flax
class       __SCREAMING_SNAKE_CASE   (				unittest.TestCase):






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'				)
												lowerCAmelCase__        =  FlaxBertModel(_UpperCamelCase				)

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

																		with self.assertRaises(_UpperCamelCase				):
																								lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase				)

																		lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase       ,  subfolder=_UpperCamelCase				)

												self.assertTrue(check_models_equal(_UpperCamelCase       ,  _UpperCamelCase				)				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  BertConfig.from_pretrained('hf-internal-testing/tiny-bert-flax-only'				)
												lowerCAmelCase__        =  FlaxBertModel(_UpperCamelCase				)

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

																		with self.assertRaises(_UpperCamelCase				):
																								lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase				)

																		lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase       ,  subfolder=_UpperCamelCase				)

												self.assertTrue(check_models_equal(_UpperCamelCase       ,  _UpperCamelCase				)				)






						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  'bert'
												lowerCAmelCase__        =  'hf-internal-testing/tiny-random-bert-subfolder'

												with self.assertRaises(_UpperCamelCase				):
																		lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase				)

												lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase       ,  subfolder=_UpperCamelCase				)

												self.assertIsNotNone(_UpperCamelCase				)







						def  UpperCamelCase__ (		self				):






												"""simple docstring"""
												lowerCAmelCase__        =  'bert'
												lowerCAmelCase__        =  'hf-internal-testing/tiny-random-bert-sharded-subfolder'
												with self.assertRaises(_UpperCamelCase				):
																		lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase				)

												lowerCAmelCase__        =  FlaxBertModel.from_pretrained(_UpperCamelCase       ,  subfolder=_UpperCamelCase				)

												self.assertIsNotNone(_UpperCamelCase				)













from .glue import GlueDataset, GlueDataTrainingArguments
from .language_modeling import (
    LineByLineTextDataset,
    LineByLineWithRefDataset,
    LineByLineWithSOPTextDataset,
    TextDataset,
    TextDatasetForNextSentencePrediction,
)
from .squad import SquadDataset, SquadDataTrainingArguments





