'''simple docstring'''





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 	__A				:



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



    def  _lowercase						(self   :   str			):
         UpperCAmelCase_	   =       ids_tensor([self.batch_size, self.seq_length]       ,      self.vocab_size			)

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

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

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

         UpperCAmelCase_	   =       model(__a			)

         UpperCAmelCase_	   =       [input_ids, input_mask]

         UpperCAmelCase_	   =       model(__a			)

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



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



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



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



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






    def  _lowercase						(self   :   Optional[int]			):
         UpperCAmelCase_	   =       self.prepare_config_and_inputs()
         ((UpperCAmelCase_)   ,			(UpperCAmelCase_)   ,			(UpperCAmelCase_)   ,			(UpperCAmelCase_)   ,			(UpperCAmelCase_)   ,			(UpperCAmelCase_))	   =       config_and_inputs
         UpperCAmelCase_	   =       {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict




@require_tf
class 	__A				(			UpperCamelCase__		,     UpperCamelCase__		,     unittest.TestCase     ):
    a__	:						Any	       =  (
        (
            TFDistilBertModel,
            TFDistilBertForMaskedLM,
            TFDistilBertForQuestionAnswering,
            TFDistilBertForSequenceClassification,
            TFDistilBertForTokenClassification,
            TFDistilBertForMultipleChoice,
        )
        if is_tf_available()
        else None
    )
    a__	:						int	       =  (
        {
            """feature-extraction""": TFDistilBertModel,
            """fill-mask""": TFDistilBertForMaskedLM,
            """question-answering""": TFDistilBertForQuestionAnswering,
            """text-classification""": TFDistilBertForSequenceClassification,
            """token-classification""": TFDistilBertForTokenClassification,
            """zero-shot""": TFDistilBertForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    a__	:						int	       =  False
    a__	:						Dict	       =  False



    def  _lowercase						(self   :   Optional[int]			):
         UpperCAmelCase_	   =       TFDistilBertModelTester(self			)
         UpperCAmelCase_	   =       ConfigTester(self       ,      config_class=__a       ,      dim=37			)



    def  _lowercase						(self   :   Union[str, Any]			):
         self.config_tester.run_common_tests()



    def  _lowercase						(self   :   Dict			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_model(*__a			)



    def  _lowercase						(self   :   List[Any]			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_masked_lm(*__a			)



    def  _lowercase						(self   :   Union[str, Any]			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_question_answering(*__a			)



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_sequence_classification(*__a			)



    def  _lowercase						(self   :   Optional[int]			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_multiple_choice(*__a			)



    def  _lowercase						(self   :   List[Any]			):
         UpperCAmelCase_	   =       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_distilbert_for_token_classification(*__a			)






    @slow
    def  _lowercase						(self   :   Any			):
         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]			):
              UpperCAmelCase_	   =       TFDistilBertModel.from_pretrained(__a			)
              self.assertIsNotNone(__a			)




@require_tf
class 	__A				(			unittest.TestCase     ):



    @slow
    def  _lowercase						(self   :   Union[str, Any]			):
         UpperCAmelCase_	   =       TFDistilBertModel.from_pretrained("distilbert-base-uncased"			)
         UpperCAmelCase_	   =       tf.constant([[0, 1, 2, 3, 4, 5]]			)
         UpperCAmelCase_	   =       model(__a			)[0]

         UpperCAmelCase_	   =       [1, 6, 768]
         self.assertEqual(output.shape       ,      __a			)

         UpperCAmelCase_	   =       tf.constant(
             [
                 [
                     [0.19_26_18_85, -0.13_73_29_55, 0.4_11_97_99],
                     [0.22_15_01_56, -0.07_42_26_61, 0.39_03_72_04],
                     [0.22_75_60_18, -0.0_89_64_14, 0.3_70_14_67],
                 ]
             ]			)
         tf.debugging.assert_near(output[:, :3, :3]       ,      __a       ,      atol=1E-4			)





'''simple docstring'''





import unittest
from typing import Tuple

import torch

from diffusers.utils import floats_tensor, randn_tensor, torch_all_close, torch_device
from diffusers.utils.testing_utils import require_torch




@require_torch
class 	__A				:



    @property
    def  _lowercase						(self   :   int			):
         return self.get_dummy_input()



    @property
    def  _lowercase						(self   :   Any			):
         if self.block_type == "down":
              return (4, 32, 16, 16)
         elif self.block_type == "mid":
              return (4, 32, 32, 32)
         elif self.block_type == "up":
              return (4, 32, 64, 64)

         raise ValueError(f"""'{self.block_type}' is not a supported block_type. Set it to 'up', 'mid', or 'down'."""			)



    def  _lowercase						(self   :   Tuple       ,      __a   :   List[str]=True       ,      __a   :   Any=False       ,      __a   :   List[Any]=False       ,      __a   :   Any=False       ,      ):
         UpperCAmelCase_	   =       4
         UpperCAmelCase_	   =       32
         UpperCAmelCase_	   =       (32, 32)

         UpperCAmelCase_	   =       torch.manual_seed(0			)
         UpperCAmelCase_	   =       torch.device(__a			)
         UpperCAmelCase_	   =       (batch_size, num_channels) + sizes
         UpperCAmelCase_	   =       randn_tensor(__a       ,      generator=__a       ,      device=__a			)
         UpperCAmelCase_	   =       {"hidden_states": hidden_states}

         if include_temb:
              UpperCAmelCase_	   =       128
              UpperCAmelCase_	   =       randn_tensor((batch_size, temb_channels)       ,      generator=__a       ,      device=__a			)

         if include_res_hidden_states_tuple:
              UpperCAmelCase_	   =       torch.manual_seed(1			)
              UpperCAmelCase_	   =       (randn_tensor(__a       ,      generator=__a       ,      device=__a			),)

         if include_encoder_hidden_states:
              UpperCAmelCase_	   =       floats_tensor((batch_size, 32, 32)			).to(__a			)

         if include_skip_sample:
              UpperCAmelCase_	   =       randn_tensor(((batch_size, 3) + sizes)       ,      generator=__a       ,      device=__a			)

         return dummy_input



    def  _lowercase						(self   :   Tuple			):
         UpperCAmelCase_	   =       {
             "in_channels": 32,
             "out_channels": 32,
             "temb_channels": 128,
         }
         if self.block_type == "up":
              UpperCAmelCase_	   =       32

         if self.block_type == "mid":
              init_dict.pop("out_channels"			)

         UpperCAmelCase_	   =       self.dummy_input
         return init_dict, inputs_dict



    def  _lowercase						(self   :   Tuple       ,      __a   :   Any			):
         UpperCAmelCase_   ,			UpperCAmelCase_	   =       self.prepare_init_args_and_inputs_for_common()
         UpperCAmelCase_	   =       self.block_class(**__a			)
         unet_block.to(__a			)
         unet_block.eval()

         with torch.no_grad():
              UpperCAmelCase_	   =       unet_block(**__a			)

         if isinstance(__a       ,      __a			):
              UpperCAmelCase_	   =       output[0]

         self.assertEqual(output.shape       ,      self.output_shape			)

         UpperCAmelCase_	   =       output[0, -1, -3:, -3:]
         UpperCAmelCase_	   =       torch.tensor(__a			).to(__a			)
         assert torch_all_close(output_slice.flatten()       ,      __a       ,      atol=5E-3			)






    @unittest.skipIf(torch_device == "mps"       ,      "Training is not supported in mps"			)
    def  _lowercase						(self   :   Dict			):
         UpperCAmelCase_   ,			UpperCAmelCase_	   =       self.prepare_init_args_and_inputs_for_common()
         UpperCAmelCase_	   =       self.block_class(**__a			)
         model.to(__a			)
         model.train()
         UpperCAmelCase_	   =       model(**__a			)

         if isinstance(__a       ,      __a			):
              UpperCAmelCase_	   =       output[0]

         UpperCAmelCase_	   =       torch.device(__a			)
         UpperCAmelCase_	   =       randn_tensor(output.shape       ,      device=__a			)
         UpperCAmelCase_	   =       torch.nn.functional.mse_loss(__a       ,      __a			)
         loss.backward()




