



from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import DeiTConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import (
						    TFDeiTForImageClassification,
						    TFDeiTForImageClassificationWithTeacher,
						    TFDeiTForMaskedImageModeling,
						    TFDeiTModel,
						)
						from transformers.models.deit.modeling_tf_deit import TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import DeiTImageProcessor






class    SCREAMING_SNAKE_CASE			:





							def __init__(   self		:		Optional[int]      ,				__lowercase		:		str      ,				__lowercase		:		Optional[int]=13      ,				__lowercase		:		str=30      ,				__lowercase		:		Optional[int]=2      ,				__lowercase		:		List[Any]=3      ,				__lowercase		:		List[Any]=True      ,				__lowercase		:		str=True      ,				__lowercase		:		Optional[Any]=32      ,				__lowercase		:		int=2      ,				__lowercase		:		Optional[Any]=4      ,				__lowercase		:		str=37      ,				__lowercase		:		Union[str, Any]="gelu"      ,				__lowercase		:		Union[str, Any]=0.1      ,				__lowercase		:		Dict=0.1      ,				__lowercase		:		Tuple=10      ,				__lowercase		:		List[Any]=0.02      ,				__lowercase		:		int=3      ,				__lowercase		:		List[Any]=None      ,				__lowercase		:		Union[str, Any]=2      ,				):

											'''simple docstring'''
											__a       =		parent
											__a       =		batch_size
											__a       =		image_size
											__a       =		patch_size
											__a       =		num_channels
											__a       =		is_training
											__a       =		use_labels
											__a       =		hidden_size
											__a       =		num_hidden_layers
											__a       =		num_attention_heads
											__a       =		intermediate_size
											__a       =		hidden_act
											__a       =		hidden_dropout_prob
											__a       =		attention_probs_dropout_prob
											__a       =		type_sequence_label_size
											__a       =		initializer_range
											__a       =		scope
											__a       =		encoder_stride

											# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
											__a       =		(image_size // patch_size) ** 2
											__a       =		num_patches + 2





							def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

											'''simple docstring'''
											__a       =		floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

											__a       =		None
											if self.use_labels:
															__a       =		ids_tensor([self.batch_size]      ,				self.type_sequence_label_size	)

											__a       =		self.get_config()

											return config, pixel_values, labels





							def 						UpperCamelCase_					(   self		:		List[Any]	):

											'''simple docstring'''
											return DeiTConfig(
											    image_size=self.image_size      ,				patch_size=self.patch_size      ,				num_channels=self.num_channels      ,				hidden_size=self.hidden_size      ,				num_hidden_layers=self.num_hidden_layers      ,				num_attention_heads=self.num_attention_heads      ,				intermediate_size=self.intermediate_size      ,				hidden_act=self.hidden_act      ,				hidden_dropout_prob=self.hidden_dropout_prob      ,				attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,				is_decoder=__lowercase      ,				initializer_range=self.initializer_range      ,				encoder_stride=self.encoder_stride      ,				)





							def 						UpperCamelCase_					(   self		:		str      ,				__lowercase		:		Any      ,				__lowercase		:		Union[str, Any]      ,				__lowercase		:		Dict	):

											'''simple docstring'''
											__a       =		TFDeiTModel(config=__lowercase	)
											__a       =		model(__lowercase	)
											self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)	)





							def 						UpperCamelCase_					(   self		:		int      ,				__lowercase		:		Union[str, Any]      ,				__lowercase		:		Any      ,				__lowercase		:		List[str]	):

											'''simple docstring'''
											__a       =		TFDeiTForMaskedImageModeling(config=__lowercase	)
											__a       =		model(__lowercase	)
											self.parent.assertEqual(
											    result.reconstruction.shape      ,				(self.batch_size, self.num_channels, self.image_size, self.image_size)	)

											# test greyscale images
											__a       =		1
											__a       =		TFDeiTForMaskedImageModeling(__lowercase	)

											__a       =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
											__a       =		model(__lowercase	)
											self.parent.assertEqual(result.reconstruction.shape      ,				(self.batch_size, 1, self.image_size, self.image_size)	)





							def 						UpperCamelCase_					(   self		:		Optional[int]      ,				__lowercase		:		Any      ,				__lowercase		:		Any      ,				__lowercase		:		Union[str, Any]	):

											'''simple docstring'''
											__a       =		self.type_sequence_label_size
											__a       =		TFDeiTForImageClassification(__lowercase	)
											__a       =		model(__lowercase      ,				labels=__lowercase	)
											self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.type_sequence_label_size)	)

											# test greyscale images
											__a       =		1
											__a       =		TFDeiTForImageClassification(__lowercase	)

											__a       =		floats_tensor([self.batch_size, 1, self.image_size, self.image_size]	)
											__a       =		model(__lowercase      ,				labels=__lowercase	)
											self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.type_sequence_label_size)	)







							def 						UpperCamelCase_					(   self		:		int	):

											'''simple docstring'''
											__a       =		self.prepare_config_and_inputs()
											__a     ,					__a     ,					__a       =		config_and_inputs
											__a       =		{"""pixel_values""": pixel_values}
											return config, inputs_dict






@require_tf
class    SCREAMING_SNAKE_CASE			(							lowerCamelCase__			,      lowerCamelCase__			,      unittest.TestCase					):
							__lowerCamelCase	:	Dict    =(
							    (
							        TFDeiTModel,
							        TFDeiTForImageClassification,
							        TFDeiTForImageClassificationWithTeacher,
							        TFDeiTForMaskedImageModeling,
							    )
							    if is_tf_available()
							    else ()
							)
							__lowerCamelCase	:	Union[str, Any]    =(
							    {
							        'feature-extraction': TFDeiTModel,
							        'image-classification': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
							    }
							    if is_tf_available()
							    else {}
							)

							__lowerCamelCase	:	Any    =False
							__lowerCamelCase	:	Tuple    =False
							__lowerCamelCase	:	Dict    =False
							__lowerCamelCase	:	List[Any]    =False





							def 						UpperCamelCase_					(   self		:		Optional[Any]	):

											'''simple docstring'''
											__a       =		TFDeiTModelTester(self	)
											__a       =		ConfigTester(self      ,				config_class=__lowercase      ,				has_text_modality=__lowercase      ,				hidden_size=37	)





							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											self.config_tester.run_common_tests()





							@unittest.skip(reason="""DeiT does not use inputs_embeds"""	)
							def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

											'''simple docstring'''
											pass





							def 						UpperCamelCase_					(   self		:		Optional[int]	):

											'''simple docstring'''
											__a     ,					__a       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															__a       =		model_class(__lowercase	)
															self.assertIsInstance(model.get_input_embeddings()      ,				(tf.keras.layers.Layer)	)
															__a       =		model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(__lowercase      ,				tf.keras.layers.Dense	)	)





							def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

											'''simple docstring'''
											__a     ,					__a       =		self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															__a       =		model_class(__lowercase	)
															__a       =		inspect.signature(model.call	)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															__a       =		[*signature.parameters.keys()]

															__a       =		["""pixel_values"""]
															self.assertListEqual(arg_names[:1]      ,				__lowercase	)





							def 						UpperCamelCase_					(   self		:		int	):

											'''simple docstring'''
											__a       =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__lowercase	)





							def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

											'''simple docstring'''
											__a       =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_image_modeling(*__lowercase	)





							def 						UpperCamelCase_					(   self		:		List[Any]	):

											'''simple docstring'''
											__a       =		self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*__lowercase	)





							def 						UpperCamelCase_					(   self		:		int      ,				__lowercase		:		Dict      ,				__lowercase		:		Any      ,				__lowercase		:		Tuple=False	):

											'''simple docstring'''
											__a       =		super()._prepare_for_class(__lowercase      ,				__lowercase      ,				return_labels=__lowercase	)

											if return_labels:
															if "labels" in inputs_dict and "labels" not in inspect.signature(model_class.call	).parameters:
																			del inputs_dict["labels"]

											return inputs_dict







							@slow
							def 						UpperCamelCase_					(   self		:		Union[str, Any]	):

											'''simple docstring'''
											for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__a       =		TFDeiTModel.from_pretrained(__lowercase	)
															self.assertIsNotNone(__lowercase	)



def   lowerCAmelCase__						(							):
				"""simple docstring"""







				__a       =		Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""    )
				return image






@require_tf
@require_vision
class    SCREAMING_SNAKE_CASE			(							unittest.TestCase					):





							@cached_property
							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											return (
											    DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""	)
											    if is_vision_available()
											    else None
											)







							@slow
							def 						UpperCamelCase_					(   self		:		str	):

											'''simple docstring'''
											__a       =		TFDeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""	)

											__a       =		self.default_image_processor
											__a       =		prepare_img()
											__a       =		image_processor(images=__lowercase      ,				return_tensors="""tf"""	)

											# forward pass
											__a       =		model(**__lowercase	)

											# verify the logits
											__a       =		tf.TensorShape((1, 1000)	)
											self.assertEqual(outputs.logits.shape      ,				__lowercase	)

											__a       =		tf.constant([-1.0266, 0.1912, -1.2861]	)

											self.assertTrue(np.allclose(outputs.logits[0, :3]      ,				__lowercase      ,				atol=1E-4	)	)





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowerCamelCase__			= {"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

try:
						if not is_sentencepiece_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""MBartTokenizer"""]

try:
						if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= ["""MBartTokenizerFast"""]

try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= [
						    """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
						    """MBartForCausalLM""",
						    """MBartForConditionalGeneration""",
						    """MBartForQuestionAnswering""",
						    """MBartForSequenceClassification""",
						    """MBartModel""",
						    """MBartPreTrainedModel""",
						]

try:
						if not is_tf_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= [
						    """TFMBartForConditionalGeneration""",
						    """TFMBartModel""",
						    """TFMBartPreTrainedModel""",
						]

try:
						if not is_flax_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						lowerCamelCase__			= [
						    """FlaxMBartForConditionalGeneration""",
						    """FlaxMBartForQuestionAnswering""",
						    """FlaxMBartForSequenceClassification""",
						    """FlaxMBartModel""",
						    """FlaxMBartPreTrainedModel""",
						]


if TYPE_CHECKING:
						from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

						try:
												if not is_sentencepiece_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_mbart import MBartTokenizer

						try:
												if not is_tokenizers_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .tokenization_mbart_fast import MBartTokenizerFast

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_mbart import (
												    MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
												    MBartForCausalLM,
												    MBartForConditionalGeneration,
												    MBartForQuestionAnswering,
												    MBartForSequenceClassification,
												    MBartModel,
												    MBartPreTrainedModel,
												)

						try:
												if not is_tf_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

						try:
												if not is_flax_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_flax_mbart import (
												    FlaxMBartForConditionalGeneration,
												    FlaxMBartForQuestionAnswering,
												    FlaxMBartForSequenceClassification,
												    FlaxMBartModel,
												    FlaxMBartPreTrainedModel,
												)

else:
						import sys

						lowerCamelCase__			= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
