






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  SCREAMING_SNAKE_CASE			:
     '''simple docstring'''


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


     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size      )

       SCREAMING_SNAKE_CASE       :      Any							=    None
       if self.use_input_mask:
         SCREAMING_SNAKE_CASE       :      Optional[Any]							=    random_attention_mask([self.batch_size, self.seq_length]      )

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

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    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					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Optional[int]      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    TFDistilBertModel(config=lowerCAmelCase__      )
       SCREAMING_SNAKE_CASE       :      List[Any]							=    {"input_ids": input_ids, "attention_mask": input_mask}

       SCREAMING_SNAKE_CASE       :      Tuple							=    model(lowerCAmelCase__      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    [input_ids, input_mask]

       SCREAMING_SNAKE_CASE       :      Tuple							=    model(lowerCAmelCase__      )

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


     def 				_A       (				self       :	str					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    TFDistilBertForMaskedLM(config=lowerCAmelCase__      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    {"input_ids": input_ids, "attention_mask": input_mask}
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    model(lowerCAmelCase__      )
       self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size)      )


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


     def 				_A       (				self       :	str					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	str					, UpperCAmelCase_       :	List[str]					, UpperCAmelCase_       :	Optional[int]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      Dict							=    self.num_choices
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    TFDistilBertForMultipleChoice(lowerCAmelCase__      )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    tf.tile(tf.expand_dims(lowerCAmelCase__					, 1      )					, (1, self.num_choices, 1)      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    tf.tile(tf.expand_dims(lowerCAmelCase__					, 1      )					, (1, self.num_choices, 1)      )
       SCREAMING_SNAKE_CASE       :      List[Any]							=    {
           "input_ids": multiple_choice_inputs_ids,
           "attention_mask": multiple_choice_input_mask,
       }
       SCREAMING_SNAKE_CASE       :      List[str]							=    model(lowerCAmelCase__      )
       self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_choices)      )


     def 				_A       (				self       :	str					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	Union[str, Any]					, UpperCAmelCase_       :	Any					, UpperCAmelCase_       :	str					, UpperCAmelCase_       :	Tuple					, UpperCAmelCase_       :	Union[str, Any]      ):
       SCREAMING_SNAKE_CASE       :      Dict							=    self.num_labels
       SCREAMING_SNAKE_CASE       :      Tuple							=    TFDistilBertForTokenClassification(lowerCAmelCase__      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    {"input_ids": input_ids, "attention_mask": input_mask}
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    model(lowerCAmelCase__      )
       self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels)      )






     def 				_A       (				self       :	Optional[Any]      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    self.prepare_config_and_inputs()
       (SCREAMING_SNAKE_CASE)       :      Union[str, Any]							=    config_and_inputs
       SCREAMING_SNAKE_CASE       :      List[str]							=    {"input_ids": input_ids, "attention_mask": input_mask}
       return config, inputs_dict





@require_tf
class  SCREAMING_SNAKE_CASE			(			a__      ,		a__      ,		unittest.TestCase					):
     '''simple docstring'''


     UpperCamelCase_							:		Optional[Any]			=	(
         (
             TFDistilBertModel,
             TFDistilBertForMaskedLM,
             TFDistilBertForQuestionAnswering,
             TFDistilBertForSequenceClassification,
             TFDistilBertForTokenClassification,
             TFDistilBertForMultipleChoice,
         )
         if is_tf_available()
         else None
     )
     UpperCamelCase_							:		Any			=	(
         {
             """feature-extraction""": TFDistilBertModel,
             """fill-mask""": TFDistilBertForMaskedLM,
             """question-answering""": TFDistilBertForQuestionAnswering,
             """text-classification""": TFDistilBertForSequenceClassification,
             """token-classification""": TFDistilBertForTokenClassification,
             """zero-shot""": TFDistilBertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCamelCase_							:		List[str]			=	False
     UpperCamelCase_							:		Optional[int]			=	False


     def 				_A       (				self       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    TFDistilBertModelTester(self      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    ConfigTester(self					, config_class=lowerCAmelCase__					, dim=37      )


     def 				_A       (				self       :	List[str]      ):
       self.config_tester.run_common_tests()


     def 				_A       (				self       :	Optional[Any]      ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*lowerCAmelCase__      )


     def 				_A       (				self       :	Dict      ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*lowerCAmelCase__      )


     def 				_A       (				self       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      str							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*lowerCAmelCase__      )


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      Dict							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowerCAmelCase__      )


     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowerCAmelCase__      )


     def 				_A       (				self       :	Optional[Any]      ):
       SCREAMING_SNAKE_CASE       :      Any							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*lowerCAmelCase__      )






     @slow
     def 				_A       (				self       :	str      ):
       for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]      ):
         SCREAMING_SNAKE_CASE       :      Tuple							=    TFDistilBertModel.from_pretrained(lowerCAmelCase__      )
         self.assertIsNotNone(lowerCAmelCase__      )





@require_tf
class  SCREAMING_SNAKE_CASE			(			unittest.TestCase					):
     '''simple docstring'''


     @slow
     def 				_A       (				self       :	Any      ):
       SCREAMING_SNAKE_CASE       :      List[Any]							=    TFDistilBertModel.from_pretrained("distilbert-base-uncased"      )
       SCREAMING_SNAKE_CASE       :      Dict							=    tf.constant([[0, 1, 2, 3, 4, 5]]      )
       SCREAMING_SNAKE_CASE       :      int							=    model(lowerCAmelCase__      )[0]

       SCREAMING_SNAKE_CASE       :      int							=    [1, 6, 768]
       self.assertEqual(output.shape					, lowerCAmelCase__      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    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]					, lowerCAmelCase__					, atol=1E-4      )














from ...utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_torch_available,
    is_transformers_available,
)


try:
   if not (is_transformers_available() and is_torch_available()):
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
   from .multicontrolnet import MultiControlNetModel
   from .pipeline_controlnet import StableDiffusionControlNetPipeline
   from .pipeline_controlnet_imgaimg import StableDiffusionControlNetImgaImgPipeline
   from .pipeline_controlnet_inpaint import StableDiffusionControlNetInpaintPipeline


if is_transformers_available() and is_flax_available():
   from .pipeline_flax_controlnet import FlaxStableDiffusionControlNetPipeline






