

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 	A_		:
							def __init__(		self		:      Dict       ,SCREAMING_SNAKE_CASE__		:      Tuple       ,SCREAMING_SNAKE_CASE__		:      Any=1_3       ,SCREAMING_SNAKE_CASE__		:      str=3_0       ,SCREAMING_SNAKE_CASE__		:      List[str]=2       ,SCREAMING_SNAKE_CASE__		:      int=3       ,SCREAMING_SNAKE_CASE__		:      List[str]=True       ,SCREAMING_SNAKE_CASE__		:      Dict=True       ,SCREAMING_SNAKE_CASE__		:      Optional[int]=3_2       ,SCREAMING_SNAKE_CASE__		:      Optional[int]=2       ,SCREAMING_SNAKE_CASE__		:      Dict=4       ,SCREAMING_SNAKE_CASE__		:      Dict=3_7       ,SCREAMING_SNAKE_CASE__		:      Optional[Any]="gelu"       ,SCREAMING_SNAKE_CASE__		:      List[Any]=0.1       ,SCREAMING_SNAKE_CASE__		:      Dict=0.1       ,SCREAMING_SNAKE_CASE__		:      int=1_0       ,SCREAMING_SNAKE_CASE__		:      Dict=0.02       ,SCREAMING_SNAKE_CASE__		:      Any=3       ,SCREAMING_SNAKE_CASE__		:      List[str]=None       ,SCREAMING_SNAKE_CASE__		:      Optional[int]=2       ,):
								__lowerCamelCase    :    Tuple 	=			parent
								__lowerCamelCase    :    Any 	=			batch_size
								__lowerCamelCase    :    Any 	=			image_size
								__lowerCamelCase    :    int 	=			patch_size
								__lowerCamelCase    :    Any 	=			num_channels
								__lowerCamelCase    :    List[Any] 	=			is_training
								__lowerCamelCase    :    List[Any] 	=			use_labels
								__lowerCamelCase    :    Any 	=			hidden_size
								__lowerCamelCase    :    str 	=			num_hidden_layers
								__lowerCamelCase    :    Any 	=			num_attention_heads
								__lowerCamelCase    :    Optional[int] 	=			intermediate_size
								__lowerCamelCase    :    int 	=			hidden_act
								__lowerCamelCase    :    str 	=			hidden_dropout_prob
								__lowerCamelCase    :    Tuple 	=			attention_probs_dropout_prob
								__lowerCamelCase    :    List[Any] 	=			type_sequence_label_size
								__lowerCamelCase    :    int 	=			initializer_range
								__lowerCamelCase    :    Union[str, Any] 	=			scope
								__lowerCamelCase    :    Dict 	=			encoder_stride

								# in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
								__lowerCamelCase    :    Optional[Any] 	=			(image_size // patch_size) ** 2
								__lowerCamelCase    :    Optional[Any] 	=			num_patches + 2
							def    lowerCAmelCase					(		self		:      List[str]):
								__lowerCamelCase    :    Optional[Any] 	=			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

								__lowerCamelCase    :    Any 	=			None
								if self.use_labels:
									__lowerCamelCase    :    Optional[Any] 	=			ids_tensor([self.batch_size]       ,self.type_sequence_label_size)

								__lowerCamelCase    :    int 	=			self.get_config()

								return config, pixel_values, labels
							def    lowerCAmelCase					(		self		:      Optional[int]):
								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    lowerCAmelCase					(		self		:      int       ,SCREAMING_SNAKE_CASE__		:      List[Any]       ,SCREAMING_SNAKE_CASE__		:      Optional[Any]       ,SCREAMING_SNAKE_CASE__		:      str):
								__lowerCamelCase    :    Any 	=			TFDeiTModel(config=lowercase_)
								__lowerCamelCase    :    List[str] 	=			model(lowercase_)
								self.parent.assertEqual(result.last_hidden_state.shape       ,(self.batch_size, self.seq_length, self.hidden_size))
							def    lowerCAmelCase					(		self		:      List[str]       ,SCREAMING_SNAKE_CASE__		:      List[Any]       ,SCREAMING_SNAKE_CASE__		:      int       ,SCREAMING_SNAKE_CASE__		:      List[Any]):
								__lowerCamelCase    :    str 	=			TFDeiTForMaskedImageModeling(config=lowercase_)
								__lowerCamelCase    :    Dict 	=			model(lowercase_)
								self.parent.assertEqual(
								    result.reconstruction.shape       ,(self.batch_size, self.num_channels, self.image_size, self.image_size))

								# test greyscale images
								__lowerCamelCase    :    Optional[Any] 	=			1
								__lowerCamelCase    :    int 	=			TFDeiTForMaskedImageModeling(lowercase_)

								__lowerCamelCase    :    List[Any] 	=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
								__lowerCamelCase    :    Optional[Any] 	=			model(lowercase_)
								self.parent.assertEqual(result.reconstruction.shape       ,(self.batch_size, 1, self.image_size, self.image_size))
							def    lowerCAmelCase					(		self		:      List[Any]       ,SCREAMING_SNAKE_CASE__		:      int       ,SCREAMING_SNAKE_CASE__		:      Union[str, Any]       ,SCREAMING_SNAKE_CASE__		:      Any):
								__lowerCamelCase    :    Optional[int] 	=			self.type_sequence_label_size
								__lowerCamelCase    :    Union[str, Any] 	=			TFDeiTForImageClassification(lowercase_)
								__lowerCamelCase    :    Dict 	=			model(lowercase_       ,labels=lowercase_)
								self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.type_sequence_label_size))

								# test greyscale images
								__lowerCamelCase    :    Dict 	=			1
								__lowerCamelCase    :    List[str] 	=			TFDeiTForImageClassification(lowercase_)

								__lowerCamelCase    :    Any 	=			floats_tensor([self.batch_size, 1, self.image_size, self.image_size])
								__lowerCamelCase    :    Tuple 	=			model(lowercase_       ,labels=lowercase_)
								self.parent.assertEqual(result.logits.shape       ,(self.batch_size, self.type_sequence_label_size))


							def    lowerCAmelCase					(		self		:      str):
								__lowerCamelCase    :    List[str] 	=			self.prepare_config_and_inputs()
								__lowerCamelCase    :    Dict 	=			config_and_inputs
								__lowerCamelCase    :    Any 	=			{'''pixel_values''': pixel_values}
								return config, inputs_dict






@require_tf
class 	A_		(	a_      ,							a_      ,							unittest.TestCase    ):
							_UpperCAmelCase   :  Optional[int]										=				(
							    (
							        TFDeiTModel,
							        TFDeiTForImageClassification,
							        TFDeiTForImageClassificationWithTeacher,
							        TFDeiTForMaskedImageModeling,
							    )
							    if is_tf_available()
							    else ()
							)
							_UpperCAmelCase   :  Dict										=				(
							    {
							        '''feature-extraction''': TFDeiTModel,
							        '''image-classification''': (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
							    }
							    if is_tf_available()
							    else {}
							)

							_UpperCAmelCase   :  int										=				False
							_UpperCAmelCase   :  str										=				False
							_UpperCAmelCase   :  List[Any]										=				False
							_UpperCAmelCase   :  Optional[Any]										=				False
							def    lowerCAmelCase					(		self		:      Union[str, Any]):
								__lowerCamelCase    :    List[Any] 	=			TFDeiTModelTester(self)
								__lowerCamelCase    :    int 	=			ConfigTester(self       ,config_class=lowercase_       ,has_text_modality=lowercase_       ,hidden_size=3_7)
							def    lowerCAmelCase					(		self		:      Union[str, Any]):
								self.config_tester.run_common_tests()
							@unittest.skip(reason='DeiT does not use inputs_embeds')
							def    lowerCAmelCase					(		self		:      List[Any]):
								pass
							def    lowerCAmelCase					(		self		:      Dict):
								__lowerCamelCase    :    Dict 	=			self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
									__lowerCamelCase    :    Optional[int] 	=			model_class(lowercase_)
									self.assertIsInstance(model.get_input_embeddings()       ,(tf.keras.layers.Layer))
									__lowerCamelCase    :    int 	=			model.get_output_embeddings()
									self.assertTrue(x is None or isinstance(lowercase_       ,tf.keras.layers.Dense))
							def    lowerCAmelCase					(		self		:      str):
								__lowerCamelCase    :    Dict 	=			self.model_tester.prepare_config_and_inputs_for_common()

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

									__lowerCamelCase    :    Tuple 	=			['''pixel_values''']
									self.assertListEqual(arg_names[:1]       ,lowercase_)
							def    lowerCAmelCase					(		self		:      Optional[Any]):
								__lowerCamelCase    :    int 	=			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*lowercase_)
							def    lowerCAmelCase					(		self		:      int):
								__lowerCamelCase    :    Optional[int] 	=			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_image_modeling(*lowercase_)
							def    lowerCAmelCase					(		self		:      int):
								__lowerCamelCase    :    str 	=			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_image_classification(*lowercase_)
							def    lowerCAmelCase					(		self		:      List[Any]       ,SCREAMING_SNAKE_CASE__		:      Union[str, Any]       ,SCREAMING_SNAKE_CASE__		:      Dict       ,SCREAMING_SNAKE_CASE__		:      Tuple=False):
								__lowerCamelCase    :    Tuple 	=			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    lowerCAmelCase					(		self		:      List[Any]):
								for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									__lowerCamelCase    :    Optional[Any] 	=			TFDeiTModel.from_pretrained(lowercase_)
									self.assertIsNotNone(lowercase_)






def        SCREAMING_SNAKE_CASE__	(  )  ->					str:
	__lowerCamelCase    :    Optional[Any] 	=			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'      )
	return image






@require_tf
@require_vision
class 	A_		(	unittest.TestCase    ):
							@cached_property
							def    lowerCAmelCase					(		self		:      Dict):
								return (
								    DeiTImageProcessor.from_pretrained('facebook/deit-base-distilled-patch16-224')
								    if is_vision_available()
								    else None
								)


							@slow
							def    lowerCAmelCase					(		self		:      List[str]):
								__lowerCamelCase    :    List[Any] 	=			TFDeiTForImageClassificationWithTeacher.from_pretrained('facebook/deit-base-distilled-patch16-224')

								__lowerCamelCase    :    int 	=			self.default_image_processor
								__lowerCamelCase    :    Union[str, Any] 	=			prepare_img()
								__lowerCamelCase    :    Dict 	=			image_processor(images=lowercase_       ,return_tensors='tf')

								# forward pass
								__lowerCamelCase    :    Union[str, Any] 	=			model(**lowercase_)

								# verify the logits
								__lowerCamelCase    :    int 	=			tf.TensorShape((1, 1_0_0_0))
								self.assertEqual(outputs.logits.shape       ,lowercase_)

								__lowerCamelCase    :    Tuple 	=			tf.constant([-1.0266, 0.1912, -1.2861])

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







# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse

from ...utils.dataclasses import (
    ComputeEnvironment,
    DistributedType,
    DynamoBackend,
    PrecisionType,
    SageMakerDistributedType,
)
from ..menu import BulletMenu


a		=[
    """EAGER""",
    """AOT_EAGER""",
    """INDUCTOR""",
    """NVFUSER""",
    """AOT_NVFUSER""",
    """AOT_CUDAGRAPHS""",
    """OFI""",
    """FX2TRT""",
    """ONNXRT""",
    """IPEX""",
]






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__     ,			lowerCamelCase__=None     ,			lowerCamelCase__=None     ,			lowerCamelCase__=None      )  ->					Optional[int]:
	__lowerCamelCase    :    int 	=			True
	while ask_again:
		__lowerCamelCase    :    Dict 	=			input(lowerCamelCase__      )
		try:
			if default is not None and len(lowerCamelCase__      ) == 0:
				return default
			return convert_value(lowerCamelCase__      ) if convert_value is not None else result
		except Exception:
			if error_message is not None:
				print(lowerCamelCase__      )






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__     ,			lowerCamelCase__=[]     ,			lowerCamelCase__=None     ,			lowerCamelCase__=0      )  ->					str:
	__lowerCamelCase    :    Union[str, Any] 	=			BulletMenu(lowerCamelCase__     ,			lowerCamelCase__      )
	__lowerCamelCase    :    Tuple 	=			menu.run(default_choice=lowerCamelCase__      )
	return convert_value(lowerCamelCase__      ) if convert_value is not None else result






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					Dict:
	__lowerCamelCase    :    List[str] 	=			int(lowerCamelCase__      )
	return ComputeEnvironment(['LOCAL_MACHINE', 'AMAZON_SAGEMAKER'][value]      )






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					Optional[int]:
	__lowerCamelCase    :    Union[str, Any] 	=			int(lowerCamelCase__      )
	return DistributedType(['NO', 'MULTI_CPU', 'MULTI_XPU', 'MULTI_GPU', 'MULTI_NPU', 'TPU'][value]      )






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					Union[str, Any]:
	__lowerCamelCase    :    Optional[Any] 	=			int(lowerCamelCase__      )
	return DynamoBackend(DYNAMO_BACKENDS[value]      ).value






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					int:
	__lowerCamelCase    :    Union[str, Any] 	=			int(lowerCamelCase__      )
	return PrecisionType(['no', 'fp16', 'bf16', 'fp8'][value]      )






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					str:
	__lowerCamelCase    :    Optional[Any] 	=			int(lowerCamelCase__      )
	return SageMakerDistributedType(['NO', 'DATA_PARALLEL', 'MODEL_PARALLEL'][value]      )






def        SCREAMING_SNAKE_CASE__	(  lowerCamelCase__      )  ->					Optional[Any]:
	return {"yes": True, "no": False}[value.lower()]






class 	A_		(	argparse.RawDescriptionHelpFormatter    ):
							def    lowerCAmelCase					(		self		:      Optional[int]       ,SCREAMING_SNAKE_CASE__		:      str       ,SCREAMING_SNAKE_CASE__		:      int       ,SCREAMING_SNAKE_CASE__		:      Union[str, Any]       ,SCREAMING_SNAKE_CASE__		:      str):
								__lowerCamelCase    :    int 	=			super()._format_usage(SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__       ,SCREAMING_SNAKE_CASE__)
								__lowerCamelCase    :    Union[str, Any] 	=			usage.replace('<command> [<args>] '       ,'')
								return usage




