





'''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       UpperCAmelCase (						SCREAMING_SNAKE_CASE_			,       SCREAMING_SNAKE_CASE_			,       unittest.TestCase  ):
  UpperCAmelCase					:	Optional[Any] 			=		(
      (
          TFMobileBertModel,
          TFMobileBertForMaskedLM,
          TFMobileBertForNextSentencePrediction,
          TFMobileBertForPreTraining,
          TFMobileBertForQuestionAnswering,
          TFMobileBertForSequenceClassification,
          TFMobileBertForTokenClassification,
          TFMobileBertForMultipleChoice,
      )
      if is_tf_available()
      else ()
  )
  UpperCAmelCase					:	Any 			=		(
      {
          '''feature-extraction''': TFMobileBertModel,
          '''fill-mask''': TFMobileBertForMaskedLM,
          '''question-answering''': TFMobileBertForQuestionAnswering,
          '''text-classification''': TFMobileBertForSequenceClassification,
          '''token-classification''': TFMobileBertForTokenClassification,
          '''zero-shot''': TFMobileBertForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  UpperCAmelCase					:	str 			=		False
  UpperCAmelCase					:	Union[str, Any] 			=		False




  def 				UpperCAmelCase__   (self      :	Union[str, Any]							,			A__      :	Optional[Any]							,			A__      :	Optional[Any]							,			A__      :	Tuple=False				)					-> Dict:
      lowercase         = super()._prepare_for_class(UpperCamelCase__							,			UpperCamelCase__							,			return_labels=UpperCamelCase__				)

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

      return inputs_dict

  class       UpperCAmelCase (						SCREAMING_SNAKE_CASE_  ):




    def __init__(self      :	str							,			A__      :	List[str]							,			A__      :	List[str]=1_3							,			A__      :	Optional[Any]=7							,			A__      :	Dict=True							,			A__      :	Optional[int]=True							,			A__      :	Optional[int]=True							,			A__      :	Union[str, Any]=True							,			A__      :	List[str]=9_9							,			A__      :	Union[str, Any]=3_2							,			A__      :	List[Any]=3_2							,			A__      :	str=2							,			A__      :	Tuple=4							,			A__      :	List[str]=3_7							,			A__      :	Union[str, Any]="gelu"							,			A__      :	Union[str, Any]=0.1							,			A__      :	Optional[int]=0.1							,			A__      :	Tuple=5_1_2							,			A__      :	int=1_6							,			A__      :	Optional[Any]=2							,			A__      :	Union[str, Any]=0.0_2							,			A__      :	List[Any]=3							,			A__      :	Union[str, Any]=4							,			A__      :	Dict=None							,			)					-> Tuple:
        lowercase         = parent
        lowercase         = batch_size
        lowercase         = seq_length
        lowercase         = is_training
        lowercase         = use_input_mask
        lowercase         = use_token_type_ids
        lowercase         = use_labels
        lowercase         = vocab_size
        lowercase         = hidden_size
        lowercase         = num_hidden_layers
        lowercase         = num_attention_heads
        lowercase         = intermediate_size
        lowercase         = hidden_act
        lowercase         = hidden_dropout_prob
        lowercase         = attention_probs_dropout_prob
        lowercase         = max_position_embeddings
        lowercase         = type_vocab_size
        lowercase         = type_sequence_label_size
        lowercase         = initializer_range
        lowercase         = num_labels
        lowercase         = num_choices
        lowercase         = scope
        lowercase         = embedding_size




    def 				UpperCAmelCase__   (self      :	List[Any]				)					-> List[Any]:
        lowercase         = ids_tensor([self.batch_size, self.seq_length]							,			self.vocab_size				)

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

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

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

        lowercase         = 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 				UpperCAmelCase__   (self      :	List[Any]							,			A__      :	Union[str, Any]							,			A__      :	Optional[Any]							,			A__      :	Optional[Any]							,			A__      :	List[str]							,			A__      :	List[str]							,			A__      :	Optional[int]							,			A__      :	int				)					-> List[Any]:
        lowercase         = TFMobileBertModel(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = model(UpperCamelCase__				)

        lowercase         = [input_ids, input_mask]
        lowercase         = model(UpperCamelCase__				)

        lowercase         = 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 				UpperCAmelCase__   (self      :	Any							,			A__      :	int							,			A__      :	Tuple							,			A__      :	List[str]							,			A__      :	Optional[Any]							,			A__      :	Optional[int]							,			A__      :	List[str]							,			A__      :	Optional[int]				)					-> List[str]:
        lowercase         = TFMobileBertForMaskedLM(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = model(UpperCamelCase__				)
        self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.seq_length, self.vocab_size)				)




    def 				UpperCAmelCase__   (self      :	int							,			A__      :	Dict							,			A__      :	int							,			A__      :	List[Any]							,			A__      :	List[str]							,			A__      :	Optional[Any]							,			A__      :	int							,			A__      :	str				)					-> int:
        lowercase         = TFMobileBertForNextSentencePrediction(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = model(UpperCamelCase__				)
        self.parent.assertEqual(result.logits.shape							,			(self.batch_size, 2)				)




    def 				UpperCAmelCase__   (self      :	Dict							,			A__      :	Union[str, Any]							,			A__      :	Optional[int]							,			A__      :	int							,			A__      :	int							,			A__      :	Optional[int]							,			A__      :	int							,			A__      :	Dict				)					-> Optional[int]:
        lowercase         = TFMobileBertForPreTraining(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = 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 				UpperCAmelCase__   (self      :	Union[str, Any]							,			A__      :	str							,			A__      :	Union[str, Any]							,			A__      :	str							,			A__      :	Tuple							,			A__      :	List[str]							,			A__      :	Tuple							,			A__      :	int				)					-> Optional[Any]:
        lowercase         = self.num_labels
        lowercase         = TFMobileBertForSequenceClassification(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = model(UpperCamelCase__				)
        self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.num_labels)				)




    def 				UpperCAmelCase__   (self      :	Union[str, Any]							,			A__      :	Union[str, Any]							,			A__      :	List[Any]							,			A__      :	List[str]							,			A__      :	Tuple							,			A__      :	int							,			A__      :	Dict							,			A__      :	int				)					-> Tuple:
        lowercase         = self.num_choices
        lowercase         = TFMobileBertForMultipleChoice(config=UpperCamelCase__				)
        lowercase         = tf.tile(tf.expand_dims(UpperCamelCase__							,			1				)							,			(1, self.num_choices, 1)				)
        lowercase         = tf.tile(tf.expand_dims(UpperCamelCase__							,			1				)							,			(1, self.num_choices, 1)				)
        lowercase         = tf.tile(tf.expand_dims(UpperCamelCase__							,			1				)							,			(1, self.num_choices, 1)				)
        lowercase         = {
            '''input_ids''': multiple_choice_inputs_ids,
            '''attention_mask''': multiple_choice_input_mask,
            '''token_type_ids''': multiple_choice_token_type_ids,
        }
        lowercase         = model(UpperCamelCase__				)
        self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.num_choices)				)




    def 				UpperCAmelCase__   (self      :	List[str]							,			A__      :	Optional[Any]							,			A__      :	Union[str, Any]							,			A__      :	Optional[int]							,			A__      :	int							,			A__      :	Optional[Any]							,			A__      :	List[Any]							,			A__      :	Tuple				)					-> Dict:
        lowercase         = self.num_labels
        lowercase         = TFMobileBertForTokenClassification(config=UpperCamelCase__				)
        lowercase         = {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
        lowercase         = model(UpperCamelCase__				)
        self.parent.assertEqual(result.logits.shape							,			(self.batch_size, self.seq_length, self.num_labels)				)




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




  def 				UpperCAmelCase__   (self      :	Optional[int]				)					-> int:
      lowercase         = TFMobileBertModelTest.TFMobileBertModelTester(self				)
      lowercase         = ConfigTester(self							,			config_class=UpperCamelCase__							,			hidden_size=3_7				)




  def 				UpperCAmelCase__   (self      :	Dict				)					-> str:
      self.config_tester.run_common_tests()




  def 				UpperCAmelCase__   (self      :	Optional[int]				)					-> Union[str, Any]:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_model(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	List[Any]				)					-> int:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_masked_lm(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	int				)					-> Any:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_multiple_choice(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	str				)					-> str:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	Union[str, Any]				)					-> List[str]:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_pretraining(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	str				)					-> List[str]:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_question_answering(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	Optional[int]				)					-> Tuple:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_sequence_classification(*UpperCamelCase__				)




  def 				UpperCAmelCase__   (self      :	Dict				)					-> Optional[int]:
      lowercase         = self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_mobilebert_for_token_classification(*UpperCamelCase__				)





  @slow
  def 				UpperCAmelCase__   (self      :	Dict				)					-> List[str]:
      # for model_name in TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      for model_name in ["google/mobilebert-uncased"]:
          lowercase         = TFMobileBertModel.from_pretrained(UpperCamelCase__				)
          self.assertIsNotNone(UpperCamelCase__				)




@require_tf
class       UpperCAmelCase (						unittest.TestCase  ):




  @slow
  def 				UpperCAmelCase__   (self      :	int				)					-> List[str]:
      lowercase         = TFMobileBertForPreTraining.from_pretrained("google/mobilebert-uncased"				)
      lowercase         = tf.constant([[0, 1, 2, 3, 4, 5]]				)
      lowercase         = model(UpperCamelCase__				)[0]

      lowercase         = [1, 6, 3_0_5_2_2]
      self.assertEqual(output.shape							,			UpperCamelCase__				)

      lowercase         = tf.constant(
          [
              [
                  [-4.5_9_1_9_5_4_7, -9.2_4_8_2_9_5, -9.6_4_5_2_5_6],
                  [-6.7_3_0_6_1_7_5, -6.4_4_0_2_8_4, -6.6_0_5_2_8_3_7],
                  [-7.2_7_4_3_5_0_6, -6.7_8_4_7_9_1_5, -6.0_2_4_6_7_3],
              ]
          ]				)
      tf.debugging.assert_near(output[:, :3, :3]							,			UpperCamelCase__							,			atol=1e-4				)











'''simple docstring'''

from argparse import ArgumentParser

from . import BaseTransformersCLICommand







def 							UpperCAmelCase_    (	lowerCAmelCase_							):

				"""simple docstring"""






				return DownloadCommand(args.model		,							args.cache_dir		,							args.force		,							args.trust_remote_code							)


class       UpperCAmelCase (						_lowercase  ):




		@staticmethod
		def 				UpperCAmelCase__   (A__      :	ArgumentParser				)					-> List[Any]:
						lowercase         = parser.add_parser("download"				)
						download_parser.add_argument(
						    "--cache-dir"							,			type=A__							,			default=A__							,			help="Path to location to store the models"				)
						download_parser.add_argument(
						    "--force"							,			action="store_true"							,			help="Force the model to be download even if already in cache-dir"				)
						download_parser.add_argument(
						    "--trust-remote-code"							,			action="store_true"							,			help="Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine"							,			)
						download_parser.add_argument("model"							,			type=A__							,			help="Name of the model to download"				)
						download_parser.set_defaults(func=A__				)




		def __init__(self      :	Tuple							,			A__      :	str							,			A__      :	str							,			A__      :	bool							,			A__      :	bool				)					-> str:
						lowercase         = model
						lowercase         = cache
						lowercase         = force
						lowercase         = trust_remote_code





		def 				UpperCAmelCase__   (self      :	Union[str, Any]				)					-> Dict:
						from ..models.auto import AutoModel, AutoTokenizer

						AutoModel.from_pretrained(
						    self._model							,			cache_dir=self._cache							,			force_download=self._force							,			trust_remote_code=self._trust_remote_code				)
						AutoTokenizer.from_pretrained(
						    self._model							,			cache_dir=self._cache							,			force_download=self._force							,			trust_remote_code=self._trust_remote_code				)




