


"""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       __A				(       _SCREAMING_SNAKE_CASE,    _SCREAMING_SNAKE_CASE,    unittest.TestCase			):






 """simple docstring"""




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


 def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A=False						)					->     Union[str, Any]:
     a 		=super()._prepare_for_class(__A ,    __A ,    return_labels=__A						)

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

     return inputs_dict






 class       __A				(       _SCREAMING_SNAKE_CASE			):






  """simple docstring"""




  def __init__(		self ,    __A ,    __A=13 ,    __A=7 ,    __A=True ,    __A=True ,    __A=True ,    __A=True ,    __A=99 ,    __A=32 ,    __A=32 ,    __A=2 ,    __A=4 ,    __A=37 ,    __A="gelu" ,    __A=0.1 ,    __A=0.1 ,    __A=512 ,    __A=16 ,    __A=2 ,    __A=0.02 ,    __A=3 ,    __A=4 ,    __A=None ,    )					->     Optional[int]:
      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						)					->     Dict:
      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 ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     Any:
      a 		=TFMobileBertModel(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)

      a 		=[input_ids, input_mask]
      a 		=model(__A						)

      a 		=model(__A						)

      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 ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     Optional[Any]:
      a 		=TFMobileBertForMaskedLM(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      self.parent.assertEqual(result.logits.shape ,    (self.batch_size, self.seq_length, self.vocab_size)						)


  def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     List[Any]:
      a 		=TFMobileBertForNextSentencePrediction(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      self.parent.assertEqual(result.logits.shape ,    (self.batch_size, 2)						)


  def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     Optional[Any]:
      a 		=TFMobileBertForPreTraining(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      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 ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     Optional[int]:
      a 		=self.num_labels
      a 		=TFMobileBertForSequenceClassification(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      self.parent.assertEqual(result.logits.shape ,    (self.batch_size, self.num_labels)						)


  def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     List[str]:
      a 		=self.num_choices
      a 		=TFMobileBertForMultipleChoice(config=__A						)
      a 		=tf.tile(tf.expand_dims(__A ,    1						) ,    (1, self.num_choices, 1)						)
      a 		=tf.tile(tf.expand_dims(__A ,    1						) ,    (1, self.num_choices, 1)						)
      a 		=tf.tile(tf.expand_dims(__A ,    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(__A						)
      self.parent.assertEqual(result.logits.shape ,    (self.batch_size, self.num_choices)						)


  def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     List[str]:
      a 		=self.num_labels
      a 		=TFMobileBertForTokenClassification(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      self.parent.assertEqual(result.logits.shape ,    (self.batch_size, self.seq_length, self.num_labels)						)


  def       SCREAMING_SNAKE_CASE							(		self ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A ,    __A						)					->     List[Any]:
      a 		=TFMobileBertForQuestionAnswering(config=__A						)
      a 		={'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
      a 		=model(__A						)
      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						)					->     Tuple:
      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						)					->     List[str]:
     a 		=TFMobileBertModelTest.TFMobileBertModelTester(self						)
     a 		=ConfigTester(self ,    config_class=__A ,    hidden_size=37						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Tuple:
     self.config_tester.run_common_tests()


 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_model(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_masked_lm(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     int:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_multiple_choice(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     List[Any]:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_pretraining(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_question_answering(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Dict:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_sequence_classification(*__A						)


 def       SCREAMING_SNAKE_CASE							(		self						)					->     Union[str, Any]:
     a 		=self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_mobilebert_for_token_classification(*__A						)





 @slow
 def       SCREAMING_SNAKE_CASE							(		self						)					->     Dict:
     # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     for model_name in ["google/mobilebert-uncased"]:
         a 		=TFMobileBertModel.from_pretrained(__A						)
         self.assertIsNotNone(__A						)






@require_tf
class       __A				(       unittest.TestCase			):






 """simple docstring"""




 @slow
 def       SCREAMING_SNAKE_CASE							(		self						)					->     str:
     a 		=TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''						)
     a 		=tf.constant([[0, 1, 2, 3, 4, 5]]						)
     a 		=model(__A						)[0]

     a 		=[1, 6, 3_0522]
     self.assertEqual(output.shape ,    __A						)

     a 		=tf.constant(
         [
             [
                 [-4.5_919_547, -9.248_295, -9.645_256],
                 [-6.7_306_175, -6.440_284, -6.6_052_837],
                 [-7.2_743_506, -6.7_847_915, -6.024_673],
             ]
         ]						)
     tf.debugging.assert_near(output[:, :3, :3] ,    __A ,    atol=1E-4						)





# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from argparse import ArgumentParser

from accelerate.commands.config import get_config_parser
from accelerate.commands.env import env_command_parser
from accelerate.commands.launch import launch_command_parser
from accelerate.commands.test import test_command_parser
from accelerate.commands.tpu import tpu_command_parser


def 						lowerCamelCase__    (		):







	'''simple docstring'''





	snake_case_		    =     ArgumentParser("Accelerate CLI tool"     ,	usage="accelerate <command> [<args>]"     ,	allow_abbrev=_A							)
	snake_case_		    =     parser.add_subparsers(help="accelerate command helpers"							)

	# Register commands
	get_config_parser(subparsers=_A							)
	env_command_parser(subparsers=_A							)
	launch_command_parser(subparsers=_A							)
	tpu_command_parser(subparsers=_A							)
	test_command_parser(subparsers=_A							)

	# Let's go
	snake_case_		    =     parser.parse_args()

	if not hasattr(_A     ,	"func"							):
		parser.print_help()
		exit(1							)

	# Run
	args.func(_A							)


if __name__ == "__main__":
		main()
