

'''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 		__SCREAMING_SNAKE_CASE       :







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







				def    lowerCamelCase_  (					self    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       Dict								=ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

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

										lowercase   :       Optional[int]								=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    lowerCamelCase_  (					self    :		Any		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Any		,				UpperCAmelCase__    :		Optional[int]      ):
										'''simple docstring'''
										lowercase   :       int								=TFDistilBertModel(config=UpperCAmelCase__      )
										lowercase   :       Dict								={'''input_ids''': input_ids, '''attention_mask''': input_mask}

										lowercase   :       List[str]								=model(UpperCAmelCase__      )

										lowercase   :       str								=[input_ids, input_mask]

										lowercase   :       str								=model(UpperCAmelCase__      )

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







				def    lowerCamelCase_  (					self    :		Union[str, Any]		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		List[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		int      ):
										'''simple docstring'''
										lowercase   :       str								=TFDistilBertForMaskedLM(config=UpperCAmelCase__      )
										lowercase   :       List[Any]								={'''input_ids''': input_ids, '''attention_mask''': input_mask}
										lowercase   :       int								=model(UpperCAmelCase__      )
										self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)      )







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







				def    lowerCamelCase_  (					self    :		List[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Tuple		,				UpperCAmelCase__    :		Union[str, Any]		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Union[str, Any]      ):
										'''simple docstring'''
										lowercase   :       Tuple								=self.num_choices
										lowercase   :       Optional[Any]								=TFDistilBertForMultipleChoice(UpperCAmelCase__      )
										lowercase   :       Union[str, Any]								=tf.tile(tf.expand_dims(UpperCAmelCase__		,				1      )		,				(1, self.num_choices, 1)      )
										lowercase   :       str								=tf.tile(tf.expand_dims(UpperCAmelCase__		,				1      )		,				(1, self.num_choices, 1)      )
										lowercase   :       Optional[int]								={
										    '''input_ids''': multiple_choice_inputs_ids,
										    '''attention_mask''': multiple_choice_input_mask,
										}
										lowercase   :       List[Any]								=model(UpperCAmelCase__      )
										self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.num_choices)      )







				def    lowerCamelCase_  (					self    :		Union[str, Any]		,				UpperCAmelCase__    :		Dict		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		str		,				UpperCAmelCase__    :		int		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		Optional[int]      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.num_labels
										lowercase   :       Dict								=TFDistilBertForTokenClassification(UpperCAmelCase__      )
										lowercase   :       int								={'''input_ids''': input_ids, '''attention_mask''': input_mask}
										lowercase   :       Dict								=model(UpperCAmelCase__      )
										self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels)      )







				def    lowerCamelCase_  (					self    :		Any      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.prepare_config_and_inputs()
										((lowercase)		,  (lowercase)		,  (lowercase)		,  (lowercase)		,  (lowercase)		,  (lowercase))   :       List[Any]								=config_and_inputs
										lowercase   :       Union[str, Any]								={'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict





@require_tf
class 		__SCREAMING_SNAKE_CASE       ( lowercase__	,       lowercase__	,       unittest.TestCase			):
				lowerCamelCase_							=						(
				    (
				        TFDistilBertModel,
				        TFDistilBertForMaskedLM,
				        TFDistilBertForQuestionAnswering,
				        TFDistilBertForSequenceClassification,
				        TFDistilBertForTokenClassification,
				        TFDistilBertForMultipleChoice,
				    )
				    if is_tf_available()
				    else None
				)
				lowerCamelCase_							=						(
				    {
				        'feature-extraction': TFDistilBertModel,
				        'fill-mask': TFDistilBertForMaskedLM,
				        'question-answering': TFDistilBertForQuestionAnswering,
				        'text-classification': TFDistilBertForSequenceClassification,
				        'token-classification': TFDistilBertForTokenClassification,
				        'zero-shot': TFDistilBertForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				lowerCamelCase_							=						False
				lowerCamelCase_							=						False







				def    lowerCamelCase_  (					self    :		Optional[Any]      ):
										'''simple docstring'''
										lowercase   :       Union[str, Any]								=TFDistilBertModelTester(self      )
										lowercase   :       str								=ConfigTester(self		,				config_class=UpperCAmelCase__		,				dim=37      )







				def    lowerCamelCase_  (					self    :		List[Any]      ):
										'''simple docstring'''
										self.config_tester.run_common_tests()







				def    lowerCamelCase_  (					self    :		Union[str, Any]      ):
										'''simple docstring'''
										lowercase   :       List[str]								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										lowercase   :       int								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		str      ):
										'''simple docstring'''
										lowercase   :       Tuple								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		int      ):
										'''simple docstring'''
										lowercase   :       str								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		List[str]      ):
										'''simple docstring'''
										lowercase   :       str								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__      )







				def    lowerCamelCase_  (					self    :		Optional[int]      ):
										'''simple docstring'''
										lowercase   :       Optional[int]								=self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__      )







				@slow
				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]      ):
																lowercase   :       Tuple								=TFDistilBertModel.from_pretrained(UpperCAmelCase__      )
																self.assertIsNotNone(UpperCAmelCase__      )





@require_tf
class 		__SCREAMING_SNAKE_CASE       ( unittest.TestCase			):







				@slow
				def    lowerCamelCase_  (					self    :		Dict      ):
										'''simple docstring'''
										lowercase   :       Tuple								=TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''      )
										lowercase   :       Optional[Any]								=tf.constant([[0, 1, 2, 3, 4, 5]]      )
										lowercase   :       Any								=model(UpperCAmelCase__      )[0]

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

										lowercase   :       int								=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]		,				UpperCAmelCase__		,				atol=1E-4      )








'''simple docstring'''





from typing import List, Optional, Tuple, Union

import torch

from ...schedulers import DDIMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class 		__SCREAMING_SNAKE_CASE       ( lowercase__			):







				def __init__(					self    :		Optional[Any]		,				UpperCAmelCase__    :		Optional[int]		,				UpperCAmelCase__    :		int      ):
										'''simple docstring'''
										super().__init__()

										# make sure scheduler can always be converted to DDIM
										lowercase   :       Any								=DDIMScheduler.from_config(scheduler.config      )

										self.register_modules(unet=UpperCAmelCase__		,				scheduler=UpperCAmelCase__      )







				@torch.no_grad()
				def __call__(					self    :		List[Any]		,				UpperCAmelCase__    :		int = 1		,				UpperCAmelCase__    :		Optional[Union[torch.Generator, List[torch.Generator]]] = None		,				UpperCAmelCase__    :		float = 0.0		,				UpperCAmelCase__    :		int = 50		,				UpperCAmelCase__    :		Optional[bool] = None		,				UpperCAmelCase__    :		Optional[str] = "pil"		,				UpperCAmelCase__    :		bool = True		,				):
										'''simple docstring'''

										# Sample gaussian noise to begin loop
										if isinstance(self.unet.config.sample_size		,				UpperCAmelCase__      ):
																lowercase   :       Optional[int]								=(
																    batch_size,
																    self.unet.config.in_channels,
																    self.unet.config.sample_size,
																    self.unet.config.sample_size,
																)
										else:
																lowercase   :       Optional[int]								=(batch_size, self.unet.config.in_channels, *self.unet.config.sample_size)

										if isinstance(UpperCAmelCase__		,				UpperCAmelCase__      ) and len(UpperCAmelCase__      ) != batch_size:
																raise ValueError(
																    F'''You have passed a list of generators of length {len(UpperCAmelCase__      )}, but requested an effective batch'''
																    F''' size of {batch_size}. Make sure the batch size matches the length of the generators.'''      )

										lowercase   :       str								=randn_tensor(UpperCAmelCase__		,				generator=UpperCAmelCase__		,				device=self.device		,				dtype=self.unet.dtype      )

										# set step values
										self.scheduler.set_timesteps(UpperCAmelCase__      )

										for t in self.progress_bar(self.scheduler.timesteps      ):
																# 1. predict noise model_output
																lowercase   :       Dict								=self.unet(UpperCAmelCase__		,				UpperCAmelCase__      ).sample

																# 2. predict previous mean of image x_t-1 and add variance depending on eta
																# eta corresponds to η in paper and should be between [0, 1]
																# do x_t -> x_t-1
																lowercase   :       Dict								=self.scheduler.step(
																    UpperCAmelCase__		,				UpperCAmelCase__		,				UpperCAmelCase__		,				eta=UpperCAmelCase__		,				use_clipped_model_output=UpperCAmelCase__		,				generator=UpperCAmelCase__      ).prev_sample

										lowercase   :       Optional[Any]								=(image / 2 + 0.5).clamp(0		,				1      )
										lowercase   :       Optional[Any]								=image.cpu().permute(0		,				2		,				3		,				1      ).numpy()
										if output_type == "pil":
																lowercase   :       List[str]								=self.numpy_to_pil(UpperCAmelCase__      )

										if not return_dict:
																return (image,)

										return ImagePipelineOutput(images=UpperCAmelCase__      )





