






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       :	List[str]					, UpperCAmelCase_       :	Tuple					, ):
       SCREAMING_SNAKE_CASE       :      str							=    parent
       SCREAMING_SNAKE_CASE       :      int							=    13
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    7
       SCREAMING_SNAKE_CASE       :      str							=    True
       SCREAMING_SNAKE_CASE       :      List[Any]							=    True
       SCREAMING_SNAKE_CASE       :      Any							=    False
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    True
       SCREAMING_SNAKE_CASE       :      Tuple							=    99
       SCREAMING_SNAKE_CASE       :      str							=    32
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    2
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    4
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    37
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    "gelu"
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    0.1
       SCREAMING_SNAKE_CASE       :      Tuple							=    0.1
       SCREAMING_SNAKE_CASE       :      str							=    512
       SCREAMING_SNAKE_CASE       :      List[str]							=    16
       SCREAMING_SNAKE_CASE       :      Dict							=    2
       SCREAMING_SNAKE_CASE       :      int							=    0.02
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    3
       SCREAMING_SNAKE_CASE       :      List[str]							=    4
       SCREAMING_SNAKE_CASE       :      str							=    None


     def 				_A       (				self       :	List[str]      ):
       SCREAMING_SNAKE_CASE       :      int							=    ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size      )

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

       SCREAMING_SNAKE_CASE       :      Tuple							=    None
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    None
       SCREAMING_SNAKE_CASE       :      str							=    None
       if self.use_labels:
         SCREAMING_SNAKE_CASE       :      int							=    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       :      str							=    ids_tensor([self.batch_size]					, self.num_choices      )

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

       SCREAMING_SNAKE_CASE       :      str							=    model(UpperCamelCase__      )

       SCREAMING_SNAKE_CASE       :      int							=    [input_ids, input_mask]

       SCREAMING_SNAKE_CASE       :      int							=    model(UpperCamelCase__      )

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


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


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


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


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






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





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


     UpperCamelCase_							:		Optional[int]			=	(
         (
             TFDistilBertModel,
             TFDistilBertForMaskedLM,
             TFDistilBertForQuestionAnswering,
             TFDistilBertForSequenceClassification,
             TFDistilBertForTokenClassification,
             TFDistilBertForMultipleChoice,
         )
         if is_tf_available()
         else None
     )
     UpperCamelCase_							:		Union[str, 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[Any]			=	False
     UpperCamelCase_							:		int			=	False


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


     def 				_A       (				self       :	Optional[Any]      ):
       self.config_tester.run_common_tests()


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      Tuple							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_model(*UpperCamelCase__      )


     def 				_A       (				self       :	List[str]      ):
       SCREAMING_SNAKE_CASE       :      str							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCamelCase__      )


     def 				_A       (				self       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCamelCase__      )


     def 				_A       (				self       :	Optional[Any]      ):
       SCREAMING_SNAKE_CASE       :      str							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCamelCase__      )


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      Any							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCamelCase__      )


     def 				_A       (				self       :	Tuple      ):
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCamelCase__      )






     @slow
     def 				_A       (				self       :	Any      ):





       for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]      ):
         SCREAMING_SNAKE_CASE       :      Optional[Any]							=    TFDistilBertModel.from_pretrained(UpperCamelCase__      )
         self.assertIsNotNone(UpperCamelCase__      )





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


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

       SCREAMING_SNAKE_CASE       :      List[Any]							=    [1, 6, 768]
       self.assertEqual(output.shape					, UpperCamelCase__      )

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














import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()


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


     def 				_A       (				self       :	List[Any]      ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()


     def 				_A       (				self       :	str      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    StableDiffusionKDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4"      )
       SCREAMING_SNAKE_CASE       :      int							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_euler"      )

       SCREAMING_SNAKE_CASE       :      str							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Any							=    sd_pipe([prompt]					, generator=UpperCAmelCase_					, guidance_scale=9.0					, num_inference_steps=20					, output_type="np"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    output.images

       SCREAMING_SNAKE_CASE       :      str							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      int							=    np.array([0.0_447, 0.0_492, 0.0_468, 0.0_408, 0.0_383, 0.0_408, 0.0_354, 0.0_380, 0.0_339]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2


     def 				_A       (				self       :	List[Any]      ):
       SCREAMING_SNAKE_CASE       :      str							=    StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
       SCREAMING_SNAKE_CASE       :      str							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_euler"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Dict							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    sd_pipe([prompt]					, generator=UpperCAmelCase_					, guidance_scale=9.0					, num_inference_steps=20					, output_type="np"      )

       SCREAMING_SNAKE_CASE       :      Optional[int]							=    output.images

       SCREAMING_SNAKE_CASE       :      Dict							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      int							=    np.array([0.1_237, 0.1_320, 0.1_438, 0.1_359, 0.1_390, 0.1_132, 0.1_277, 0.1_175, 0.1_112]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 5E-1






     def 				_A       (				self       :	Optional[int]      ):
       SCREAMING_SNAKE_CASE       :      List[str]							=    StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base"      )
       SCREAMING_SNAKE_CASE       :      Any							=    sd_pipe.to(UpperCAmelCase_      )
       sd_pipe.set_progress_bar_config(disable=UpperCAmelCase_      )

       sd_pipe.set_scheduler("sample_dpmpp_2m"      )

       SCREAMING_SNAKE_CASE       :      List[str]							=    "A painting of a squirrel eating a burger"
       SCREAMING_SNAKE_CASE       :      Any							=    torch.manual_seed(0      )
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    sd_pipe(
           [prompt]					, generator=UpperCAmelCase_					, guidance_scale=7.5					, num_inference_steps=15					, output_type="np"					, use_karras_sigmas=UpperCAmelCase_					, )

       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    output.images

       SCREAMING_SNAKE_CASE       :      List[Any]							=    image[0, -3:, -3:, -1]

       assert image.shape == (1, 512, 512, 3)
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    np.array(
           [0.11_381_689, 0.12_112_921, 0.1_389_457, 0.12_549_606, 0.1_244_964, 0.10_831_517, 0.11_562_866, 0.10_867_816, 0.10_499_048]      )

       assert np.abs(image_slice.flatten() - expected_slice      ).max() < 1E-2






