




'''simple docstring'''





import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

_SCREAMING_SNAKE_CASE				:  Optional[Any] 						=					[
    """word_embeddings_layernorm.weight""",
    """word_embeddings_layernorm.bias""",
    """input_layernorm.weight""",
    """input_layernorm.bias""",
    """post_attention_layernorm.weight""",
    """post_attention_layernorm.bias""",
    """self_attention.dense.bias""",
    """mlp.dense_4h_to_h.bias""",
    """ln_f.weight""",
    """ln_f.bias""",
]

_SCREAMING_SNAKE_CASE				:  int 						=					[
    """mlp.dense_4h_to_h.weight""",
    """self_attention.dense.weight""",
]


def 	UpperCamelCase_(      snake_case				:			Optional[Any]				,							snake_case				:			Union[str, Any]	):

							'''simple docstring'''






							snake_case_         = {
							    "word_embeddings.weight": "word_embeddings.weight",
							    "word_embeddings.norm.weight": "word_embeddings_layernorm.weight",
							    "word_embeddings.norm.bias": "word_embeddings_layernorm.bias",
							    "weight": "ln_f.weight",
							    "bias": "ln_f.bias",
							}

							if key in layer_rename_map:
														return layer_rename_map[key]

							# Handle transformer blocks
							snake_case_         = int(re.match(r".*layer_(\d*).*"				,							lowercase_	)[1]	)
							layer_number -= 3
							return f'h.{layer_number}.' + key


def 	UpperCamelCase_(      snake_case				:			List[str]	):

							'''simple docstring'''






							if dtype == torch.bool:
														return 1 / 8
							snake_case_         = re.search(r"[^\d](\d+)$"				,							str(lowercase_	)	)
							if bit_search is None:
														raise ValueError(f'`dtype` is not a valid dtype: {dtype}.'	)
							snake_case_         = int(bit_search.groups()[0]	)
							return bit_size // 8


def 	UpperCamelCase_(      snake_case				:			str				,							snake_case				:			Optional[Any]				,							snake_case				:			Dict				,							snake_case				:			Union[str, Any]				,							snake_case				:			Optional[Any]	):

							'''simple docstring'''






							if bloom_config_file == "":
														snake_case_         = BloomConfig()
							else:
														snake_case_         = BloomConfig.from_json_file(lowercase_	)

							if shard_model:
														snake_case_         = os.listdir(lowercase_	)
														snake_case_         = sorted(filter(lambda snake_case	: s.startswith("layer"	) and "model_00" in s				,							lowercase_	)	)

														snake_case_         = {"weight_map": {}, "metadata": {}}
														snake_case_         = 0

														snake_case_         = None

														snake_case_         = BloomConfig()

														for j, file in enumerate(lowercase_	):
																					print("Processing file: {}".format(lowercase_	)	)
																					snake_case_         = None

																					for i in range(lowercase_	):
																												# load all TP files
																												snake_case_         = file.replace("model_00"				,							f'model_0{i}'	)
																												snake_case_         = torch.load(os.path.join(lowercase_				,							lowercase_	)				,							map_location="cpu"	)

																												# Rename keys in the transformers names
																												snake_case_         = list(temp.keys()	)
																												for key in keys:
																																			snake_case_         = temp.pop(lowercase_	)

																												if tensors is None:
																																			snake_case_         = temp
																												else:
																																			for key in tensors.keys():
																																										if any(key.endswith(lowercase_	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																																																	# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																																																	tensors[key] += temp[key]
																																										else:
																																																	# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																																	snake_case_         = 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	snake_case_         = torch.cat([tensors[key], temp[key]]				,							dim=lowercase_	)

            # Divide by the number of TP the weights we want to average
																					for key in tensors.keys():
																												if any(key.endswith(lowercase_	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																																			snake_case_         = tensors[key] / pretraining_tp
																					torch.save(
																					    lowercase_				,							os.path.join(
																					        lowercase_				,							"pytorch_model_{}-of-{}.bin".format(str(j + 1	).zfill(5	)				,							str(len(lowercase_	)	).zfill(5	)	)				,							)				,							)

																					for key in tensors.keys():
																												snake_case_         = tensors[key]
																												total_size += value.numel() * get_dtype_size(value.dtype	)
																												if key not in index_dict["weight_map"]:
																																			snake_case_         = "pytorch_model_{}-of-{}.bin".format(
																																			    str(j + 1	).zfill(5	)				,							str(len(lowercase_	)	).zfill(5	)	)

														snake_case_         = BloomConfig()
														snake_case_         = pytorch_dump_folder_path + "/" + CONFIG_NAME
														snake_case_         = total_size
														with open(lowercase_				,							"w"				,							encoding="utf-8"	) as f:
																					f.write(config.to_json_string()	)
														with open(os.path.join(lowercase_				,							WEIGHTS_NAME + ".index.json"	)				,							"w"				,							encoding="utf-8"	) as f:
																					snake_case_         = json.dumps(lowercase_				,							indent=2				,							sort_keys=lowercase_	) + "\n"
																					f.write(lowercase_	)
							else:
														snake_case_         = BloomModel(lowercase_	)

														snake_case_         = os.listdir(lowercase_	)
														snake_case_         = sorted(filter(lambda snake_case	: s.startswith("layer"	) and "model_00" in s				,							lowercase_	)	)

														snake_case_         = None
														for i, file in enumerate(lowercase_	):
																					snake_case_         = None
																					for i in range(lowercase_	):
																												# load all TP files
																												snake_case_         = file.replace("model_00"				,							f'model_0{i}'	)
																												snake_case_         = torch.load(os.path.join(lowercase_				,							lowercase_	)				,							map_location="cpu"	)

																												# Rename keys in the transformers names
																												snake_case_         = list(temp.keys()	)
																												for key in keys:
																																			snake_case_         = temp.pop(lowercase_	)

																												if tensors is None:
																																			snake_case_         = temp
																												else:
																																			for key in tensors.keys():
																																										# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																																										if any(key.endswith(lowercase_	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																																																	tensors[key] += temp[key]
																																										else:
																																																	# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																																																	snake_case_         = 1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN	) else 0
																																																	# We concatenate these weights accross TP ranks
																																																	snake_case_         = torch.cat([tensors[key], temp[key]]				,							dim=lowercase_	)

            # Divide by the number of TP the weights we want to average
																					for key in tensors.keys():
																												if any(key.endswith(lowercase_	) for end in WEIGHTS_TO_AVERAGE_ENDSWITH	):
																																			snake_case_         = tensors[key] / pretraining_tp

																					snake_case_         = model.load_state_dict(lowercase_				,							strict=lowercase_	)
																					assert not other_keys.unexpected_keys, f'The keys {other_keys.unexpected_keys} are unexpected'
																					if missing_keys is None:
																												snake_case_         = set(other_keys.missing_keys	)
																					else:
																												snake_case_         = missing_keys.intersection(set(other_keys.missing_keys	)	)

														assert not missing_keys, f'The keys {missing_keys} are missing'

														# Save pytorch-model
														os.makedirs(lowercase_				,							exist_ok=lowercase_	)
														snake_case_         = pytorch_dump_folder_path + "/" + WEIGHTS_NAME
														snake_case_         = pytorch_dump_folder_path + "/" + CONFIG_NAME
														print(f'Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}'	)
														if config.torch_dtype is not None:
																					snake_case_         = model.to(config.torch_dtype	)
														torch.save(model.state_dict()				,							lowercase_	)
														print(f'Save configuration file to {pytorch_config_dump_path}'	)
														with open(lowercase_				,							"w"				,							encoding="utf-8"	) as f:
																					f.write(config.to_json_string()	)


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE				:  List[Any] 						=					argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    "--bloom_checkpoint_path",
		    default=None,
		    type=str,
		    required=True,
		    help="Path to the Megatron-LM checkpoint path.",
		)
		parser.add_argument(
		    "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
		)
		parser.add_argument(
		    "--bloom_config_file",
		    default="",
		    type=str,
		    help=(
		        "An optional config json file corresponding to the pre-trained model. \n"
		        "This specifies the model architecture."
		    ),
		)
		parser.add_argument(
		    "--shard_model",
		    action="store_true",
		    help="An optional setting to shard the output model \nThis enables sharding the converted checkpoint",
		)
		parser.add_argument(
		    "--pretraining_tp",
		    default=4,
		    type=int,
		    help="Pretraining TP rank that has been used when training the model in Megatron-LM \n",
		)
		_SCREAMING_SNAKE_CASE				:  Any 						=					parser.parse_args()
		convert_bloom_checkpoint_to_pytorch(
		    args.bloom_checkpoint_path,
		    args.bloom_config_file,
		    args.pytorch_dump_folder_path,
		    args.shard_model,
		    args.pretraining_tp,
		)










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 							UpperCamelCase_       (					UpperCAmelCase__						,   UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = (
        (
            TFMobileBertModel,
            TFMobileBertForMaskedLM,
            TFMobileBertForNextSentencePrediction,
            TFMobileBertForPreTraining,
            TFMobileBertForQuestionAnswering,
            TFMobileBertForSequenceClassification,
            TFMobileBertForTokenClassification,
            TFMobileBertForMultipleChoice,
        )
        if is_tf_available()
        else ()
    )
    UpperCAmelCase__				    = (
        {
            '''feature-extraction''': TFMobileBertModel,
            '''fill-mask''': TFMobileBertForMaskedLM,
            '''question-answering''': TFMobileBertForQuestionAnswering,
            '''text-classification''': TFMobileBertForSequenceClassification,
            '''token-classification''': TFMobileBertForTokenClassification,
            '''zero-shot''': TFMobileBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    UpperCAmelCase__				    = False
    UpperCAmelCase__				    = False


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      str=False)		->Optional[Any]:





           '''simple docstring'''

           A__ 					=       super()._prepare_for_class(UpperCAmelCase__   ,					UpperCAmelCase__   ,					return_labels=UpperCAmelCase__)

           if return_labels:
                  if model_class in get_values(UpperCAmelCase__):
                         A__ 					=       tf.zeros(self.model_tester.batch_size   ,					dtype=tf.intaa)

           return inputs_dict


    class 							UpperCamelCase_       (					UpperCAmelCase__	):

        '''simple docstring'''






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





               '''simple docstring'''

               A__ 					=       parent
               A__ 					=       batch_size
               A__ 					=       seq_length
               A__ 					=       is_training
               A__ 					=       use_input_mask
               A__ 					=       use_token_type_ids
               A__ 					=       use_labels
               A__ 					=       vocab_size
               A__ 					=       hidden_size
               A__ 					=       num_hidden_layers
               A__ 					=       num_attention_heads
               A__ 					=       intermediate_size
               A__ 					=       hidden_act
               A__ 					=       hidden_dropout_prob
               A__ 					=       attention_probs_dropout_prob
               A__ 					=       max_position_embeddings
               A__ 					=       type_vocab_size
               A__ 					=       type_sequence_label_size
               A__ 					=       initializer_range
               A__ 					=       num_labels
               A__ 					=       num_choices
               A__ 					=       scope
               A__ 					=       embedding_size


        def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Tuple:





               '''simple docstring'''

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

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

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

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

               A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any]   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Optional[Any])		->Any:





               '''simple docstring'''

               A__ 					=       TFMobileBertModel(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)

               A__ 					=       [input_ids, input_mask]
               A__ 					=       model(UpperCAmelCase__)

               A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      List[str]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Tuple)		->Optional[Any]:





               '''simple docstring'''

               A__ 					=       TFMobileBertForMaskedLM(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.seq_length, self.vocab_size))


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any])		->int:





               '''simple docstring'''

               A__ 					=       TFMobileBertForNextSentencePrediction(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, 2))


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int)		->List[Any]:





               '''simple docstring'''

               A__ 					=       TFMobileBertForPreTraining(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Tuple   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      List[str]   ,					UpperCAmelCase__ :      Tuple)		->Dict:





               '''simple docstring'''

               A__ 					=       self.num_labels
               A__ 					=       TFMobileBertForSequenceClassification(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_labels))


        def 			SCREAMING_SNAKE_CASE				(     self :      Any   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      str   ,					UpperCAmelCase__ :      int)		->Dict:





               '''simple docstring'''

               A__ 					=       self.num_choices
               A__ 					=       TFMobileBertForMultipleChoice(config=UpperCAmelCase__)
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					1)   ,					(1, self.num_choices, 1))
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					1)   ,					(1, self.num_choices, 1))
               A__ 					=       tf.tile(tf.expand_dims(UpperCAmelCase__   ,					1)   ,					(1, self.num_choices, 1))
               A__ 					=       {
                   '''input_ids''': multiple_choice_inputs_ids,
                   '''attention_mask''': multiple_choice_input_mask,
                   '''token_type_ids''': multiple_choice_token_type_ids,
               }
               A__ 					=       model(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.num_choices))


        def 			SCREAMING_SNAKE_CASE				(     self :      Dict   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      Optional[int]   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      Optional[int])		->int:





               '''simple docstring'''

               A__ 					=       self.num_labels
               A__ 					=       TFMobileBertForTokenClassification(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       model(UpperCAmelCase__)
               self.parent.assertEqual(result.logits.shape   ,					(self.batch_size, self.seq_length, self.num_labels))


        def 			SCREAMING_SNAKE_CASE				(     self :      List[Any]   ,					UpperCAmelCase__ :      Optional[Any]   ,					UpperCAmelCase__ :      Dict   ,					UpperCAmelCase__ :      List[Any]   ,					UpperCAmelCase__ :      Tuple   ,					UpperCAmelCase__ :      Union[str, Any]   ,					UpperCAmelCase__ :      Any   ,					UpperCAmelCase__ :      Union[str, Any])		->Union[str, Any]:





               '''simple docstring'''

               A__ 					=       TFMobileBertForQuestionAnswering(config=UpperCAmelCase__)
               A__ 					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
               A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Any)		->str:





               '''simple docstring'''

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


    def 			SCREAMING_SNAKE_CASE				(     self :      Any)		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       TFMobileBertModelTest.TFMobileBertModelTester(self)
           A__ 					=       ConfigTester(self   ,					config_class=UpperCAmelCase__   ,					hidden_size=37)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_model(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Tuple:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Union[str, Any])		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Dict:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->Union[str, Any]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      int)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[int]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCAmelCase__)


    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->List[str]:





           '''simple docstring'''

           A__ 					=       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCAmelCase__)





    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->List[Any]:





           '''simple docstring'''

           for model_name in ["google/mobilebert-uncased"]:
                  A__ 					=       TFMobileBertModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)


@require_tf
class 							UpperCamelCase_       (					unittest.TestCase	):

    '''simple docstring'''






    @slow
    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->Any:





           '''simple docstring'''

           A__ 					=       TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased''')
           A__ 					=       tf.constant([[0, 1, 2, 3, 4, 5]])
           A__ 					=       model(UpperCAmelCase__)[0]

           A__ 					=       [1, 6, 30_522]
           self.assertEqual(output.shape   ,					UpperCAmelCase__)

           A__ 					=       tf.constant(
               [
                   [
                       [-4.5919547, -9.248295, -9.645256],
                       [-6.7306175, -6.440284, -6.6052837],
                       [-7.2743506, -6.7847915, -6.024673],
                   ]
               ])
           tf.debugging.assert_near(output[:, :3, :3]   ,					UpperCAmelCase__   ,					atol=1e-4)


