



'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import MobileBertConfig, is_tf_available
from transformers.models.auto import get_values
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 (
							    TF_MODEL_FOR_PRETRAINING_MAPPING,
							    TFMobileBertForMaskedLM,
							    TFMobileBertForMultipleChoice,
							    TFMobileBertForNextSentencePrediction,
							    TFMobileBertForPreTraining,
							    TFMobileBertForQuestionAnswering,
							    TFMobileBertForSequenceClassification,
							    TFMobileBertForTokenClassification,
							    TFMobileBertModel,
							)

@require_tf
class       __magic_name__		(       _UpperCamelCase     ,		_UpperCamelCase     ,		unittest.TestCase  ):
							lowerCAmelCase			:       Union[str, Any]  							=   (
							    (
							        TFMobileBertModel,
							        TFMobileBertForMaskedLM,
							        TFMobileBertForNextSentencePrediction,
							        TFMobileBertForPreTraining,
							        TFMobileBertForQuestionAnswering,
							        TFMobileBertForSequenceClassification,
							        TFMobileBertForTokenClassification,
							        TFMobileBertForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							lowerCAmelCase			:       Tuple  							=   (
							    {
							        'feature-extraction': TFMobileBertModel,
							        'fill-mask': TFMobileBertForMaskedLM,
							        'question-answering': TFMobileBertForQuestionAnswering,
							        'text-classification': TFMobileBertForSequenceClassification,
							        'token-classification': TFMobileBertForTokenClassification,
							        'zero-shot': TFMobileBertForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							lowerCAmelCase			:       Dict  							=   False
							lowerCAmelCase			:       str  							=   False







							def 				__lowercase     (				self   :			int      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			Optional[Any]      ,_UpperCAmelCase   :			int=False						):
									_a	: Any								=				super()._prepare_for_class(_UpperCAmelCase      ,_UpperCAmelCase      ,return_labels=_UpperCAmelCase						)

									if return_labels:
											if model_class in get_values(_UpperCAmelCase						):
													_a	: int								=				tf.zeros(self.model_tester.batch_size      ,dtype=tf.intaa						)

									return inputs_dict

							class       __magic_name__		(       _UpperCamelCase  ):







														def __init__(				self   :			Tuple      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			Tuple=13      ,_UpperCAmelCase   :			List[Any]=7      ,_UpperCAmelCase   :			int=True      ,_UpperCAmelCase   :			List[str]=True      ,_UpperCAmelCase   :			Optional[Any]=True      ,_UpperCAmelCase   :			Dict=True      ,_UpperCAmelCase   :			Dict=99      ,_UpperCAmelCase   :			List[Any]=32      ,_UpperCAmelCase   :			Optional[int]=32      ,_UpperCAmelCase   :			Tuple=2      ,_UpperCAmelCase   :			Optional[Any]=4      ,_UpperCAmelCase   :			List[Any]=37      ,_UpperCAmelCase   :			Tuple="gelu"      ,_UpperCAmelCase   :			Dict=0.1      ,_UpperCAmelCase   :			List[str]=0.1      ,_UpperCAmelCase   :			Tuple=512      ,_UpperCAmelCase   :			List[str]=16      ,_UpperCAmelCase   :			str=2      ,_UpperCAmelCase   :			Any=0.02      ,_UpperCAmelCase   :			Optional[Any]=3      ,_UpperCAmelCase   :			Union[str, Any]=4      ,_UpperCAmelCase   :			str=None      ,):
																_a	: str								=				parent
																_a	: Any								=				batch_size
																_a	: List[Any]								=				seq_length
																_a	: Dict								=				is_training
																_a	: Any								=				use_input_mask
																_a	: str								=				use_token_type_ids
																_a	: List[Any]								=				use_labels
																_a	: Tuple								=				vocab_size
																_a	: Optional[Any]								=				hidden_size
																_a	: Tuple								=				num_hidden_layers
																_a	: Optional[Any]								=				num_attention_heads
																_a	: Dict								=				intermediate_size
																_a	: Optional[int]								=				hidden_act
																_a	: Dict								=				hidden_dropout_prob
																_a	: str								=				attention_probs_dropout_prob
																_a	: int								=				max_position_embeddings
																_a	: Any								=				type_vocab_size
																_a	: Tuple								=				type_sequence_label_size
																_a	: Any								=				initializer_range
																_a	: str								=				num_labels
																_a	: Tuple								=				num_choices
																_a	: int								=				scope
																_a	: List[str]								=				embedding_size







														def 				__lowercase     (				self   :			Any						):
																_a	: Union[str, Any]								=				ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size						)

																_a	: List[Any]								=				None
																if self.use_input_mask:
																		_a	: Union[str, Any]								=				random_attention_mask([self.batch_size, self.seq_length]						)

																_a	: Dict								=				None
																if self.use_token_type_ids:
																		_a	: List[str]								=				ids_tensor([self.batch_size, self.seq_length]      ,self.type_vocab_size						)

																_a	: Dict								=				None
																_a	: Tuple								=				None
																_a	: Dict								=				None
																if self.use_labels:
																		_a	: int								=				ids_tensor([self.batch_size]      ,self.type_sequence_label_size						)
																		_a	: List[str]								=				ids_tensor([self.batch_size, self.seq_length]      ,self.num_labels						)
																		_a	: Optional[Any]								=				ids_tensor([self.batch_size]      ,self.num_choices						)

																_a	: List[Any]								=				MobileBertConfig(
																    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      ,embedding_size=self.embedding_size      ,)

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







														def 				__lowercase     (				self   :			int      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			int      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			int      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			List[str]      ,_UpperCAmelCase   :			List[str]						):
																_a	: Any								=				TFMobileBertModel(config=_UpperCAmelCase						)
																_a	: int								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: int								=				model(_UpperCAmelCase						)

																_a	: Optional[Any]								=				[input_ids, input_mask]
																_a	: Tuple								=				model(_UpperCAmelCase						)

																_a	: List[str]								=				model(_UpperCAmelCase						)

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







														def 				__lowercase     (				self   :			Dict      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			Tuple      ,_UpperCAmelCase   :			Optional[Any]      ,_UpperCAmelCase   :			str						):
																_a	: Any								=				TFMobileBertForMaskedLM(config=_UpperCAmelCase						)
																_a	: str								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: Any								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)						)







														def 				__lowercase     (				self   :			str      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			int      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			str						):
																_a	: List[str]								=				TFMobileBertForNextSentencePrediction(config=_UpperCAmelCase						)
																_a	: str								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: Any								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(result.logits.shape      ,(self.batch_size, 2)						)







														def 				__lowercase     (				self   :			Union[str, Any]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			List[Any]						):
																_a	: Tuple								=				TFMobileBertForPreTraining(config=_UpperCAmelCase						)
																_a	: Optional[int]								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: Optional[Any]								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(
																    result.prediction_logits.shape      ,(self.batch_size, self.seq_length, self.vocab_size)						)
																self.parent.assertEqual(result.seq_relationship_logits.shape      ,(self.batch_size, 2)						)







														def 				__lowercase     (				self   :			Any      ,_UpperCAmelCase   :			List[str]      ,_UpperCAmelCase   :			Optional[Any]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			List[Any]						):
																_a	: Any								=				self.num_labels
																_a	: int								=				TFMobileBertForSequenceClassification(config=_UpperCAmelCase						)
																_a	: List[Any]								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: str								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels)						)







														def 				__lowercase     (				self   :			Optional[int]      ,_UpperCAmelCase   :			List[str]      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			Any      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			Optional[Any]      ,_UpperCAmelCase   :			Optional[Any]						):
																_a	: str								=				self.num_choices
																_a	: Union[str, Any]								=				TFMobileBertForMultipleChoice(config=_UpperCAmelCase						)
																_a	: Optional[Any]								=				tf.tile(tf.expand_dims(_UpperCAmelCase      ,1						)      ,(1, self.num_choices, 1)						)
																_a	: Union[str, Any]								=				tf.tile(tf.expand_dims(_UpperCAmelCase      ,1						)      ,(1, self.num_choices, 1)						)
																_a	: Union[str, Any]								=				tf.tile(tf.expand_dims(_UpperCAmelCase      ,1						)      ,(1, self.num_choices, 1)						)
																_a	: Dict								=				{
																    'input_ids': multiple_choice_inputs_ids,
																    'attention_mask': multiple_choice_input_mask,
																    'token_type_ids': multiple_choice_token_type_ids,
																}
																_a	: List[str]								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_choices)						)







														def 				__lowercase     (				self   :			List[Any]      ,_UpperCAmelCase   :			List[Any]      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			str      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			List[Any]						):
																_a	: Union[str, Any]								=				self.num_labels
																_a	: Optional[Any]								=				TFMobileBertForTokenClassification(config=_UpperCAmelCase						)
																_a	: Tuple								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: Any								=				model(_UpperCAmelCase						)
																self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.seq_length, self.num_labels)						)







														def 				__lowercase     (				self   :			List[Any]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Tuple      ,_UpperCAmelCase   :			Optional[int]      ,_UpperCAmelCase   :			Union[str, Any]      ,_UpperCAmelCase   :			Dict      ,_UpperCAmelCase   :			Optional[Any]						):
																_a	: Any								=				TFMobileBertForQuestionAnswering(config=_UpperCAmelCase						)
																_a	: Optional[int]								=				{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
																_a	: Dict								=				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 				__lowercase     (				self   :			List[Any]						):
																_a	: Tuple								=				self.prepare_config_and_inputs()
																(
																    (
																    _a
																)			,   (
																    _a
																)			,   (
																    _a
																)			,   (
																    _a
																)			,   (
																    _a
																)			,   (
																    _a
																)			,   (
																    _a
																)			,   
																)	: Any								=				config_and_inputs
																_a	: str								=				{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
																return config, inputs_dict







							def 				__lowercase     (				self   :			str						):
									_a	: str								=				TFMobileBertModelTest.TFMobileBertModelTester(self						)
									_a	: Tuple								=				ConfigTester(self      ,config_class=_UpperCAmelCase      ,hidden_size=37						)







							def 				__lowercase     (				self   :			Tuple						):
									self.config_tester.run_common_tests()







							def 				__lowercase     (				self   :			Optional[int]						):
									_a	: int								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_model(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			str						):
									_a	: Optional[Any]								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_masked_lm(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Union[str, Any]						):
									_a	: Optional[Any]								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_multiple_choice(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Union[str, Any]						):
									_a	: Dict								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Optional[int]						):
									_a	: Tuple								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_pretraining(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			str						):
									_a	: Optional[Any]								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_question_answering(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Optional[Any]						):
									_a	: Optional[Any]								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_sequence_classification(*_UpperCAmelCase						)







							def 				__lowercase     (				self   :			Dict						):
									_a	: str								=				self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_mobilebert_for_token_classification(*_UpperCAmelCase						)







							@slow
							def 				__lowercase     (				self   :			Tuple						):
									# for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									for model_name in ["google/mobilebert-uncased"]:
											_a	: List[str]								=				TFMobileBertModel.from_pretrained(_UpperCAmelCase						)
											self.assertIsNotNone(_UpperCAmelCase						)

@require_tf
class       __magic_name__		(       unittest.TestCase  ):







							@slow
							def 				__lowercase     (				self   :			Optional[int]						):
									_a	: Tuple								=				TFMobileBertForPreTraining.from_pretrained('google/mobilebert-uncased'						)
									_a	: str								=				tf.constant([[0, 1, 2, 3, 4, 5]]						)
									_a	: List[Any]								=				model(_UpperCAmelCase						)[0]

									_a	: Any								=				[1, 6, 30522]
									self.assertEqual(output.shape      ,_UpperCAmelCase						)

									_a	: List[str]								=				tf.constant(
									    [
									        [
									            [-4.5_91_95_47, -9.24_82_95, -9.64_52_56],
									            [-6.7_30_61_75, -6.44_02_84, -6.6_05_28_37],
									            [-7.2_74_35_06, -6.7_84_79_15, -6.02_46_73],
									        ]
									    ]						)
									tf.debugging.assert_near(output[:, :3, :3]      ,_UpperCAmelCase      ,atol=1E-4						)









'''simple docstring'''

import argparse

from transformers import BigBirdConfig, BigBirdForPreTraining, BigBirdForQuestionAnswering, load_tf_weights_in_big_bird
from transformers.utils import logging


logging.set_verbosity_info()



def 	__lowerCamelCase	( lowerCAmelCase_		,     lowerCAmelCase_		,     lowerCAmelCase_		,     lowerCAmelCase_    )     -> Any:
		# Initialise PyTorch model
		_a	: int								=				BigBirdConfig.from_json_file(lowerCAmelCase_    )
		print(f"""Building PyTorch model from configuration: {config}"""    )

		if is_trivia_qa:
				_a	: Union[str, Any]								=				BigBirdForQuestionAnswering(lowerCAmelCase_    )
		else:
				_a	: Any								=				BigBirdForPreTraining(lowerCAmelCase_    )

		# Load weights from tf checkpoint
		load_tf_weights_in_big_bird(lowerCAmelCase_		,     lowerCAmelCase_		,     is_trivia_qa=lowerCAmelCase_    )

		# Save pytorch-model
		print(f"""Save PyTorch model to {pytorch_dump_path}"""    )
		model.save_pretrained(lowerCAmelCase_    )


if __name__ == "__main__":
							__lowerCAmelCase		     =  argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
							)
							parser.add_argument(
							    '''--big_bird_config_file''',
							    default=None,
							    type=str,
							    required=True,
							    help=(
							        '''The config json file corresponding to the pre-trained BERT model. \n'''
							        '''This specifies the model architecture.'''
							    ),
							)
							parser.add_argument(
							    '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
							)
							parser.add_argument(
							    '''--is_trivia_qa''', action='''store_true''', help='''Whether to convert a model with a trivia_qa head.'''
							)
							__lowerCAmelCase		     =  parser.parse_args()
							convert_tf_checkpoint_to_pytorch(
							    args.tf_checkpoint_path, args.big_bird_config_file, args.pytorch_dump_path, args.is_trivia_qa
							)




