




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
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.models.distilbert.modeling_tf_distilbert import (
         TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFDistilBertForMaskedLM,
         TFDistilBertForMultipleChoice,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertModel,
     )





class 			__lowerCamelCase	:




    """simple docstring"""

    def __init__(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     )				->	int:
         lowerCAmelCase__          =			parent
         lowerCAmelCase__          =			13
         lowerCAmelCase__          =			7
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			False
         lowerCAmelCase__          =			True
         lowerCAmelCase__          =			99
         lowerCAmelCase__          =			32
         lowerCAmelCase__          =			2
         lowerCAmelCase__          =			4
         lowerCAmelCase__          =			37
         lowerCAmelCase__          =			"gelu"
         lowerCAmelCase__          =			0.1
         lowerCAmelCase__          =			0.1
         lowerCAmelCase__          =			512
         lowerCAmelCase__          =			16
         lowerCAmelCase__          =			2
         lowerCAmelCase__          =			0.02
         lowerCAmelCase__          =			3
         lowerCAmelCase__          =			4
         lowerCAmelCase__          =			None


    def        a				(	self					:					Dict       )				->	Dict:
         lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size       )

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

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

         lowerCAmelCase__          =			DistilBertConfig(
             vocab_size=self.vocab_size       ,     dim=self.hidden_size       ,     n_layers=self.num_hidden_layers       ,     n_heads=self.num_attention_heads       ,     hidden_dim=self.intermediate_size       ,     hidden_act=self.hidden_act       ,     dropout=self.hidden_dropout_prob       ,     attention_dropout=self.attention_probs_dropout_prob       ,     max_position_embeddings=self.max_position_embeddings       ,     initializer_range=self.initializer_range       ,     )

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


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Any       )				->	Optional[int]:
         lowerCAmelCase__          =			TFDistilBertModel(config=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}

         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			[input_ids, input_mask]

         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )

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


    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	int:
         lowerCAmelCase__          =			TFDistilBertForMaskedLM(config=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size)       )


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	Dict:
         lowerCAmelCase__          =			TFDistilBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			{
             "input_ids": input_ids,
             "attention_mask": input_mask,
         }
         lowerCAmelCase__          =			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        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			self.num_labels
         lowerCAmelCase__          =			TFDistilBertForSequenceClassification(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels)       )


    def        a				(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	str:
         lowerCAmelCase__          =			self.num_choices
         lowerCAmelCase__          =			TFDistilBertForMultipleChoice(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE__       ,     1       )       ,     (1, self.num_choices, 1)       )
         lowerCAmelCase__          =			tf.tile(tf.expand_dims(SCREAMING_SNAKE_CASE__       ,     1       )       ,     (1, self.num_choices, 1)       )
         lowerCAmelCase__          =			{
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
         }
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_choices)       )


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					List[str]       )				->	Any:
         lowerCAmelCase__          =			self.num_labels
         lowerCAmelCase__          =			TFDistilBertForTokenClassification(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )
         self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels)       )


    def        a				(	self					:					List[str]       )				->	int:
         lowerCAmelCase__          =			self.prepare_config_and_inputs()
         ((lowerCAmelCase__)     ,		(lowerCAmelCase__)     ,		(lowerCAmelCase__)     ,		(lowerCAmelCase__)     ,		(lowerCAmelCase__)     ,		(lowerCAmelCase__))          =			config_and_inputs
         lowerCAmelCase__          =			{"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict






@require_tf
class 			__lowerCamelCase	(   UpperCamelCase__     ,    UpperCamelCase__     ,    unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				(
        (
            TFDistilBertModel,
            TFDistilBertForMaskedLM,
            TFDistilBertForQuestionAnswering,
            TFDistilBertForSequenceClassification,
            TFDistilBertForTokenClassification,
            TFDistilBertForMultipleChoice,
        )
        if is_tf_available()
        else None
    )
    snake_case__ =				(
        {
            "feature-extraction": TFDistilBertModel,
            "fill-mask": TFDistilBertForMaskedLM,
            "question-answering": TFDistilBertForQuestionAnswering,
            "text-classification": TFDistilBertForSequenceClassification,
            "token-classification": TFDistilBertForTokenClassification,
            "zero-shot": TFDistilBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    snake_case__ =				False
    snake_case__ =				False


    def        a				(	self					:					Union[str, Any]       )				->	Tuple:
         lowerCAmelCase__          =			TFDistilBertModelTester(self       )
         lowerCAmelCase__          =			ConfigTester(self       ,     config_class=SCREAMING_SNAKE_CASE__       ,     dim=37       )


    def        a				(	self					:					Union[str, Any]       )				->	Any:
         self.config_tester.run_common_tests()


    def        a				(	self					:					Union[str, Any]       )				->	Optional[int]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Optional[Any]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					str       )				->	str:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					str       )				->	Dict:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					List[str]       )				->	Optional[Any]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Dict       )				->	List[Any]:
         lowerCAmelCase__          =			self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*SCREAMING_SNAKE_CASE__       )


    @slow
    def        a				(	self					:					int       )				->	List[Any]:
         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]       ):
              lowerCAmelCase__          =			TFDistilBertModel.from_pretrained(SCREAMING_SNAKE_CASE__       )
              self.assertIsNotNone(SCREAMING_SNAKE_CASE__       )






@require_tf
class 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    @slow
    def        a				(	self					:					str       )				->	Optional[Any]:
         lowerCAmelCase__          =			TFDistilBertModel.from_pretrained("distilbert-base-uncased"       )
         lowerCAmelCase__          =			tf.constant([[0, 1, 2, 3, 4, 5]]       )
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       )[0]

         lowerCAmelCase__          =			[1, 6, 768]
         self.assertEqual(output.shape       ,     SCREAMING_SNAKE_CASE__       )

         lowerCAmelCase__          =			tf.constant(
             [
                 [
                     [0.19_261_885, -0.13_732_955, 0.4_119_799],
                     [0.22_150_156, -0.07_422_661, 0.39_037_204],
                     [0.22_756_018, -0.0_896_414, 0.3_701_467],
                 ]
             ]       )
         tf.debugging.assert_near(output[:, :3, :3]       ,     SCREAMING_SNAKE_CASE__       ,     atol=1e-4       )






import fire

from transformers import AutoConfig, AutoModelForSeqaSeqLM, AutoTokenizer



def 							_A       (						lowerCAmelCase_					: str       ,    lowerCAmelCase_					: str       ,    **lowerCAmelCase_					: str      ):
     """simple docstring"""

     lowerCAmelCase__          =			AutoConfig.from_pretrained(lowerCAmelCase_       ,    **lowerCAmelCase_      )
     lowerCAmelCase__          =			AutoModelForSeqaSeqLM.from_config(lowerCAmelCase_      )
     model.save_pretrained(lowerCAmelCase_      )
     AutoTokenizer.from_pretrained(lowerCAmelCase_      ).save_pretrained(lowerCAmelCase_      )
     return model


if __name__ == "__main__":
     fire.Fire(save_randomly_initialized_version)
