






from __future__ import annotations

import os
import tempfile
import unittest

from transformers import ConvBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import tensorflow as tf

		from transformers import (
		    TFConvBertForMaskedLM,
		    TFConvBertForMultipleChoice,
		    TFConvBertForQuestionAnswering,
		    TFConvBertForSequenceClassification,
		    TFConvBertForTokenClassification,
		    TFConvBertModel,
		)



class 				__UpperCAmelCase		:







			def __init__(     self:  str				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]=13				,     UpperCAmelCase_:  Dict=7				,     UpperCAmelCase_:  Dict=True				,     UpperCAmelCase_:  Tuple=True				,     UpperCAmelCase_:  Any=True				,     UpperCAmelCase_:  Tuple=True				,     UpperCAmelCase_:  Tuple=99				,     UpperCAmelCase_:  Any=32				,     UpperCAmelCase_:  Union[str, Any]=2				,     UpperCAmelCase_:  List[Any]=4				,     UpperCAmelCase_:  str=37				,     UpperCAmelCase_:  Any="gelu"				,     UpperCAmelCase_:  Union[str, Any]=0.1				,     UpperCAmelCase_:  Any=0.1				,     UpperCAmelCase_:  Tuple=512				,     UpperCAmelCase_:  Dict=16				,     UpperCAmelCase_:  int=2				,     UpperCAmelCase_:  str=0.02				,     UpperCAmelCase_:  str=3				,     UpperCAmelCase_:  List[Any]=4				,     UpperCAmelCase_:  Optional[int]=None				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  13
									_SCREAMING_SNAKE_CASE     				=  7
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  99
									_SCREAMING_SNAKE_CASE     				=  384
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  37
									_SCREAMING_SNAKE_CASE     				=  """gelu"""
									_SCREAMING_SNAKE_CASE     				=  0.1
									_SCREAMING_SNAKE_CASE     				=  0.1
									_SCREAMING_SNAKE_CASE     				=  512
									_SCREAMING_SNAKE_CASE     				=  16
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  0.02
									_SCREAMING_SNAKE_CASE     				=  3
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  128
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  9
									_SCREAMING_SNAKE_CASE     				=  1
									_SCREAMING_SNAKE_CASE     				=  None







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_input_mask:
															_SCREAMING_SNAKE_CASE     				=  random_attention_mask([self.batch_size, self.seq_length]					)

									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_token_type_ids:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.type_vocab_size					)

									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									_SCREAMING_SNAKE_CASE     				=  None
									if self.use_labels:
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.type_sequence_label_size					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     self.num_labels					)
															_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size]				,     self.num_choices					)

									_SCREAMING_SNAKE_CASE     				=  ConvBertConfig(
									    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     num_attention_heads=self.num_attention_heads				,     intermediate_size=self.intermediate_size				,     hidden_act=self.hidden_act				,     hidden_dropout_prob=self.hidden_dropout_prob				,     attention_probs_dropout_prob=self.attention_probs_dropout_prob				,     max_position_embeddings=self.max_position_embeddings				,     type_vocab_size=self.type_vocab_size				,     initializer_range=self.initializer_range				,     return_dict=UpperCAmelCase_				,     )

									return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertModel(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}

									_SCREAMING_SNAKE_CASE     				=  [input_ids, input_mask]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size)					)







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertForMaskedLM(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size)					)







			def 							UpperCamelCase				(     self:  Union[str, Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.num_labels
									_SCREAMING_SNAKE_CASE     				=  TFConvBertForSequenceClassification(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}

									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_labels)					)







			def 							UpperCamelCase				(     self:  List[str]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.num_choices
									_SCREAMING_SNAKE_CASE     				=  TFConvBertForMultipleChoice(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  tf.tile(tf.expand_dims(UpperCAmelCase_				,     1					)				,     (1, self.num_choices, 1)					)
									_SCREAMING_SNAKE_CASE     				=  tf.tile(tf.expand_dims(UpperCAmelCase_				,     1					)				,     (1, self.num_choices, 1)					)
									_SCREAMING_SNAKE_CASE     				=  tf.tile(tf.expand_dims(UpperCAmelCase_				,     1					)				,     (1, self.num_choices, 1)					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": multiple_choice_inputs_ids,
									    """attention_mask""": multiple_choice_input_mask,
									    """token_type_ids""": multiple_choice_token_type_ids,
									}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.num_choices)					)







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.num_labels
									_SCREAMING_SNAKE_CASE     				=  TFConvBertForTokenClassification(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels)					)







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Optional[Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertForQuestionAnswering(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									}

									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.start_logits.shape				,     (self.batch_size, self.seq_length)					)
									self.parent.assertEqual(result.end_logits.shape				,     (self.batch_size, self.seq_length)					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.prepare_config_and_inputs()
									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  config_and_inputs
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict



@require_tf
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				Union[str, Any]									=							(
			    (
			        TFConvBertModel,
			        TFConvBertForMaskedLM,
			        TFConvBertForQuestionAnswering,
			        TFConvBertForSequenceClassification,
			        TFConvBertForTokenClassification,
			        TFConvBertForMultipleChoice,
			    )
			    if is_tf_available()
			    else ()
			)
			__snake_case    :				Union[str, Any]									=							(
			    {
			        "feature-extraction": TFConvBertModel,
			        "fill-mask": TFConvBertForMaskedLM,
			        "question-answering": TFConvBertForQuestionAnswering,
			        "text-classification": TFConvBertForSequenceClassification,
			        "token-classification": TFConvBertForTokenClassification,
			        "zero-shot": TFConvBertForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			__snake_case    :				Optional[int]									=							False
			__snake_case    :				List[Any]									=							False
			__snake_case    :				Dict									=							False







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True

									if hasattr(UpperCAmelCase_				,     """use_cache"""					):
															_SCREAMING_SNAKE_CASE     				=  True

									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """encoder_seq_length"""				,     self.model_tester.seq_length					)
									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """key_length"""				,     UpperCAmelCase_					)

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  self._prepare_for_class(UpperCAmelCase_				,     UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  len(model(UpperCAmelCase_					)					)

															with tempfile.TemporaryDirectory() as tmpdirname:
																					model.save_pretrained(UpperCAmelCase_				,     saved_model=UpperCAmelCase_					)
																					_SCREAMING_SNAKE_CASE     				=  os.path.join(UpperCAmelCase_				,     """saved_model"""				,     """1"""					)
																					_SCREAMING_SNAKE_CASE     				=  tf.keras.models.load_model(UpperCAmelCase_					)
																					_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

																					if self.is_encoder_decoder:
																											_SCREAMING_SNAKE_CASE     				=  outputs["""encoder_hidden_states"""]
																											_SCREAMING_SNAKE_CASE     				=  outputs["""encoder_attentions"""]
																					else:
																											_SCREAMING_SNAKE_CASE     				=  outputs["""hidden_states"""]
																											_SCREAMING_SNAKE_CASE     				=  outputs["""attentions"""]

																					self.assertEqual(len(UpperCAmelCase_					)				,     UpperCAmelCase_					)

																					_SCREAMING_SNAKE_CASE     				=  getattr(
																					    self.model_tester				,     """expected_num_hidden_layers"""				,     self.model_tester.num_hidden_layers + 1					)

																					self.assertEqual(len(UpperCAmelCase_					)				,     UpperCAmelCase_					)
																					self.assertListEqual(
																					    list(output_hidden_states[0].shape[-2:]					)				,     [self.model_tester.seq_length, self.model_tester.hidden_size]				,     )

																					self.assertEqual(len(UpperCAmelCase_					)				,     self.model_tester.num_hidden_layers					)
																					self.assertListEqual(
																					    list(output_attentions[0].shape[-3:]					)				,     [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]				,     )







			@slow
			def 							UpperCamelCase				(     self:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertModel.from_pretrained("""YituTech/conv-bert-base"""					)
									self.assertIsNotNone(UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """decoder_seq_length"""				,     self.model_tester.seq_length					)
									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """encoder_seq_length"""				,     self.model_tester.seq_length					)
									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """key_length"""				,     UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  getattr(self.model_tester				,     """key_length"""				,     UpperCAmelCase_					)

									def check_decoder_attentions_output(UpperCAmelCase_:  List[str]					):
															_SCREAMING_SNAKE_CASE     				=  len(UpperCAmelCase_					)
															self.assertEqual(out_len % 2				,     0					)
															_SCREAMING_SNAKE_CASE     				=  outputs.decoder_attentions
															self.assertEqual(len(UpperCAmelCase_					)				,     self.model_tester.num_hidden_layers					)
															self.assertListEqual(
															    list(decoder_attentions[0].shape[-3:]					)				,     [self.model_tester.num_attention_heads / 2, decoder_seq_length, decoder_key_length]				,     )

									def check_encoder_attentions_output(UpperCAmelCase_:  Union[str, Any]					):
															_SCREAMING_SNAKE_CASE     				=  [
															    t.numpy() for t in (outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions)
															]
															self.assertEqual(len(UpperCAmelCase_					)				,     self.model_tester.num_hidden_layers					)
															self.assertListEqual(
															    list(attentions[0].shape[-3:]					)				,     [self.model_tester.num_attention_heads / 2, encoder_seq_length, encoder_key_length]				,     )

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  True
															_SCREAMING_SNAKE_CASE     				=  False
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  model(self._prepare_for_class(UpperCAmelCase_				,     UpperCAmelCase_					)					)
															_SCREAMING_SNAKE_CASE     				=  len(UpperCAmelCase_					)
															self.assertEqual(config.output_hidden_states				,     UpperCAmelCase_					)
															check_encoder_attentions_output(UpperCAmelCase_					)

															if self.is_encoder_decoder:
																					_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
																					_SCREAMING_SNAKE_CASE     				=  model(self._prepare_for_class(UpperCAmelCase_				,     UpperCAmelCase_					)					)
																					self.assertEqual(config.output_hidden_states				,     UpperCAmelCase_					)
																					check_decoder_attentions_output(UpperCAmelCase_					)

															# Check that output attentions can also be changed via the config
															del inputs_dict["output_attentions"]
															_SCREAMING_SNAKE_CASE     				=  True
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  model(self._prepare_for_class(UpperCAmelCase_				,     UpperCAmelCase_					)					)
															self.assertEqual(config.output_hidden_states				,     UpperCAmelCase_					)
															check_encoder_attentions_output(UpperCAmelCase_					)

															# Check attention is always last and order is fine
															_SCREAMING_SNAKE_CASE     				=  True
															_SCREAMING_SNAKE_CASE     				=  True
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															_SCREAMING_SNAKE_CASE     				=  model(self._prepare_for_class(UpperCAmelCase_				,     UpperCAmelCase_					)					)

															self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)				,     len(UpperCAmelCase_					)					)
															self.assertEqual(model.config.output_hidden_states				,     UpperCAmelCase_					)
															check_encoder_attentions_output(UpperCAmelCase_					)



@require_tf
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFConvBertModel.from_pretrained("""YituTech/conv-bert-base"""					)
									_SCREAMING_SNAKE_CASE     				=  tf.constant([[0, 1, 2, 3, 4, 5]]					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]

									_SCREAMING_SNAKE_CASE     				=  [1, 6, 768]
									self.assertEqual(output.shape				,     UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            [-0.03_47_54_93, -0.4_68_60_34, -0.30_63_88_32],
									            [0.22_63_72_48, -0.26_98_86_46, -0.7_42_34_24],
									            [0.10_32_48_68, -0.45_01_35_08, -0.58_28_07_84],
									        ]
									    ]					)
									tf.debugging.assert_near(output[:, :3, :3]				,     UpperCAmelCase_				,     atol=1E-4					)









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 							UpperCamelCase				(     cls:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TOKEN
									HfFolder.save_token(UpperCAmelCase_					)







			@classmethod
			def 							UpperCamelCase				(     cls:  Union[str, Any]					):







									'''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:  str					):







									'''simple docstring'''




									_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(UpperCAmelCase_					)
									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(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					)

									_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(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)







			def 							UpperCamelCase				(     self:  int					):







									'''simple docstring'''




									_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(UpperCAmelCase_					)
									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(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					)

									_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(UpperCAmelCase_				,     1E-3				,     msg=F'{key} not identical'					)

def 	__lowerCamelCase      (		snake_case__	,snake_case__		)							->  Tuple:






						"""simple docstring"""





						_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 				__UpperCAmelCase		(unittest.TestCase   ):







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  """bert"""
									with tempfile.TemporaryDirectory() as tmp_dir:
															model.save_pretrained(os.path.join(UpperCAmelCase_				,     UpperCAmelCase_					)					)

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  BertConfig.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""					)
									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel(UpperCAmelCase_					)

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

															with self.assertRaises(UpperCAmelCase_					):
																					_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertTrue(check_models_equal(UpperCAmelCase_				,     UpperCAmelCase_					)					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-subfolder"""

									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  """bert"""
									_SCREAMING_SNAKE_CASE     				=  """hf-internal-testing/tiny-random-bert-sharded-subfolder"""
									with self.assertRaises(UpperCAmelCase_					):
															_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  FlaxBertModel.from_pretrained(UpperCAmelCase_				,     subfolder=UpperCAmelCase_					)

									self.assertIsNotNone(UpperCAmelCase_					)

