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





     def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE=False	) ->				List[str]:






         """simple docstring"""






         A       :   Tuple					=						super()._prepare_for_class(SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			return_labels=SCREAMING_SNAKE_CASE	)

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

         return inputs_dict





     class  A				(							__snake_case						):





          def __init__(							self						,			SCREAMING_SNAKE_CASE						,			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=32						,			SCREAMING_SNAKE_CASE=32						,			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=None						,			) ->				str:






              """simple docstring"""






              A       :   List[str]					=						parent
              A       :   str					=						batch_size
              A       :   Optional[Any]					=						seq_length
              A       :   List[str]					=						is_training
              A       :   List[Any]					=						use_input_mask
              A       :   Optional[Any]					=						use_token_type_ids
              A       :   Optional[Any]					=						use_labels
              A       :   List[str]					=						vocab_size
              A       :   Dict					=						hidden_size
              A       :   Union[str, Any]					=						num_hidden_layers
              A       :   Tuple					=						num_attention_heads
              A       :   Dict					=						intermediate_size
              A       :   Tuple					=						hidden_act
              A       :   List[Any]					=						hidden_dropout_prob
              A       :   Tuple					=						attention_probs_dropout_prob
              A       :   int					=						max_position_embeddings
              A       :   int					=						type_vocab_size
              A       :   str					=						type_sequence_label_size
              A       :   int					=						initializer_range
              A       :   Optional[Any]					=						num_labels
              A       :   Optional[int]					=						num_choices
              A       :   Tuple					=						scope
              A       :   Dict					=						embedding_size





          def 				__lowerCAmelCase							(							self	) ->				List[str]:






              """simple docstring"""






              A       :   List[Any]					=						ids_tensor([self.batch_size, self.seq_length]						,			self.vocab_size	)

              A       :   Union[str, Any]					=						None
              if self.use_input_mask:
                  A       :   Optional[int]					=						random_attention_mask([self.batch_size, self.seq_length]	)

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

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

              A       :   Optional[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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[Any]:






              """simple docstring"""






              A       :   int					=						TFMobileBertModel(config=SCREAMING_SNAKE_CASE	)
              A       :   Union[str, Any]					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   List[str]					=						model(SCREAMING_SNAKE_CASE	)

              A       :   str					=						[input_ids, input_mask]
              A       :   List[str]					=						model(SCREAMING_SNAKE_CASE	)

              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)

              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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






              """simple docstring"""






              A       :   Optional[int]					=						TFMobileBertForMaskedLM(config=SCREAMING_SNAKE_CASE	)
              A       :   Any					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.vocab_size)	)





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






              """simple docstring"""






              A       :   Union[str, Any]					=						TFMobileBertForNextSentencePrediction(config=SCREAMING_SNAKE_CASE	)
              A       :   Tuple					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Optional[Any]					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, 2)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






              """simple docstring"""






              A       :   str					=						TFMobileBertForPreTraining(config=SCREAMING_SNAKE_CASE	)
              A       :   int					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   str					=						model(SCREAMING_SNAKE_CASE	)
              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 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Optional[int]:






              """simple docstring"""






              A       :   str					=						self.num_labels
              A       :   Tuple					=						TFMobileBertForSequenceClassification(config=SCREAMING_SNAKE_CASE	)
              A       :   Optional[Any]					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Dict					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.num_labels)	)





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Dict:






              """simple docstring"""






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





          def 				__lowerCAmelCase							(							self						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE						,			SCREAMING_SNAKE_CASE	) ->				Tuple:






              """simple docstring"""






              A       :   Dict					=						self.num_labels
              A       :   Optional[Any]					=						TFMobileBertForTokenClassification(config=SCREAMING_SNAKE_CASE	)
              A       :   Any					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   str					=						model(SCREAMING_SNAKE_CASE	)
              self.parent.assertEqual(result.logits.shape						,			(self.batch_size, self.seq_length, self.num_labels)	)





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






              """simple docstring"""






              A       :   Union[str, Any]					=						TFMobileBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE	)
              A       :   Dict					=						{'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
              A       :   Any					=						model(SCREAMING_SNAKE_CASE	)
              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 				__lowerCAmelCase							(							self	) ->				List[str]:






              """simple docstring"""






              A       :   List[str]					=						self.prepare_config_and_inputs()
              (
                  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  (
                  A
              ),  
              )       :   Any					=						config_and_inputs
              A       :   List[str]					=						{'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
              return config, inputs_dict





     def 				__lowerCAmelCase							(							self	) ->				Tuple:






         """simple docstring"""






         A       :   Optional[int]					=						TFMobileBertModelTest.TFMobileBertModelTester(self	)
         A       :   Union[str, Any]					=						ConfigTester(self						,			config_class=SCREAMING_SNAKE_CASE						,			hidden_size=37	)





     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         self.config_tester.run_common_tests()





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






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_model(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   Optional[int]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_masked_lm(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Tuple					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_multiple_choice(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   str					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				str:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_pretraining(*SCREAMING_SNAKE_CASE	)





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






         """simple docstring"""






         A       :   List[Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_question_answering(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				Optional[Any]:






         """simple docstring"""






         A       :   Optional[int]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_sequence_classification(*SCREAMING_SNAKE_CASE	)





     def 				__lowerCAmelCase							(							self	) ->				List[str]:






         """simple docstring"""






         A       :   Union[str, Any]					=						self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_mobilebert_for_token_classification(*SCREAMING_SNAKE_CASE	)





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Any:






         """simple docstring"""






         for model_name in ["google/mobilebert-uncased"]:
             A       :   Optional[int]					=						TFMobileBertModel.from_pretrained(SCREAMING_SNAKE_CASE	)
             self.assertIsNotNone(SCREAMING_SNAKE_CASE	)





@require_tf
class  A				(							unittest.TestCase						):





     @slow
     def 				__lowerCAmelCase							(							self	) ->				Dict:






         """simple docstring"""






         A       :   Optional[Any]					=						TFMobileBertForPreTraining.from_pretrained('''google/mobilebert-uncased'''	)
         A       :   Tuple					=						tf.constant([[0, 1, 2, 3, 4, 5]]	)
         A       :   str					=						model(SCREAMING_SNAKE_CASE	)[0]

         A       :   Dict					=						[1, 6, 30522]
         self.assertEqual(output.shape						,			SCREAMING_SNAKE_CASE	)

         A       :   Union[str, Any]					=						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]						,			SCREAMING_SNAKE_CASE						,			atol=1e-4	)











import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows





def lowerCAmelCase__(__snake_case							)				->      int:  # picklable for multiprocessing







	'''simple docstring'''







	return i + 1





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def lowerCAmelCase__()				->      Any:







	'''simple docstring'''







	with parallel_backend('''spark'''							):
		assert ParallelBackendConfig.backend_name == "spark"

	lowerCamelCase__								=		[1, 2, 3]
	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=2							)

	with pytest.raises(__snake_case							):
		with parallel_backend('''unsupported backend'''							):
			map_nested(__snake_case       ,__snake_case       ,num_proc=-1							)





@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize('''num_proc'''       ,[2, -1]							)
def lowerCAmelCase__(__snake_case							)				->      Tuple:







	'''simple docstring'''







	lowerCamelCase__								=		[1, 2]
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2}
	lowerCamelCase__								=		{'''a''': [1, 2], '''b''': [3, 4]}
	lowerCamelCase__								=		{'''a''': {'''1''': 1}, '''b''': 2}
	lowerCamelCase__								=		{'''a''': 1, '''b''': 2, '''c''': 3, '''d''': 4}
	lowerCamelCase__								=		[2, 3]
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3}
	lowerCamelCase__								=		{'''a''': [2, 3], '''b''': [4, 5]}
	lowerCamelCase__								=		{'''a''': {'''1''': 2}, '''b''': 3}
	lowerCamelCase__								=		{'''a''': 2, '''b''': 3, '''c''': 4, '''d''': 5}

	with parallel_backend('''spark'''							):
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa
		assert map_nested(__snake_case       ,__snake_case       ,num_proc=__snake_case							) == expected_map_nested_sa



