





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





					"""simple docstring"""


					def __init__( self	,  _lowerCAmelCase	,  _lowerCAmelCase=13	,  _lowerCAmelCase=30	,  _lowerCAmelCase=2	,  _lowerCAmelCase=3	,  _lowerCAmelCase=True	,  _lowerCAmelCase=True	,  _lowerCAmelCase=32	,  _lowerCAmelCase=2	,  _lowerCAmelCase=4	,  _lowerCAmelCase=37	,  _lowerCAmelCase="gelu"	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase=0.1	,  _lowerCAmelCase=10	,  _lowerCAmelCase=0.02	,  _lowerCAmelCase=3	,  _lowerCAmelCase=None	,  _lowerCAmelCase=2	,  ):
										lowerCamelCase__       =       parent
										lowerCamelCase__       =       batch_size
										lowerCamelCase__       =       image_size
										lowerCamelCase__       =       patch_size
										lowerCamelCase__       =       num_channels
										lowerCamelCase__       =       is_training
										lowerCamelCase__       =       use_labels
										lowerCamelCase__       =       hidden_size
										lowerCamelCase__       =       num_hidden_layers
										lowerCamelCase__       =       num_attention_heads
										lowerCamelCase__       =       intermediate_size
										lowerCamelCase__       =       hidden_act
										lowerCamelCase__       =       hidden_dropout_prob
										lowerCamelCase__       =       attention_probs_dropout_prob
										lowerCamelCase__       =       type_sequence_label_size
										lowerCamelCase__       =       initializer_range
										lowerCamelCase__       =       scope
										lowerCamelCase__       =       encoder_stride

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





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]	)

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

										lowerCamelCase__       =       self.get_config()

										return config, pixel_values, labels





					def 					__magic_name__ ( self	):
										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=_lowerCAmelCase	,  initializer_range=self.initializer_range	,  encoder_stride=self.encoder_stride	,  )





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       TFDeiTModel(config=_lowerCAmelCase	)
										lowerCamelCase__       =       model(_lowerCAmelCase	)
										self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size)	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       TFDeiTForMaskedImageModeling(config=_lowerCAmelCase	)
										lowerCamelCase__       =       model(_lowerCAmelCase	)
										self.parent.assertEqual(
										    result.reconstruction.shape	,  (self.batch_size, self.num_channels, self.image_size, self.image_size)	)

										# test greyscale images
										lowerCamelCase__       =       1
										lowerCamelCase__       =       TFDeiTForMaskedImageModeling(_lowerCAmelCase	)

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





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										lowerCamelCase__       =       self.type_sequence_label_size
										lowerCamelCase__       =       TFDeiTForImageClassification(_lowerCAmelCase	)
										lowerCamelCase__       =       model(_lowerCAmelCase	,  labels=_lowerCAmelCase	)
										self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.type_sequence_label_size)	)

										# test greyscale images
										lowerCamelCase__       =       1
										lowerCamelCase__       =       TFDeiTForImageClassification(_lowerCAmelCase	)

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






					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.prepare_config_and_inputs()
										lowerCamelCase__    ,		lowerCamelCase__    ,		lowerCamelCase__       =       config_and_inputs
										lowerCamelCase__       =       {"pixel_values": pixel_values}
										return config, inputs_dict

@require_tf
class SCREAMING_SNAKE_CASE_ (					lowercase_							,       lowercase_							,       unittest.TestCase   ):





					"""simple docstring"""


					A__				 =    (
					    (
					        TFDeiTModel,
					        TFDeiTForImageClassification,
					        TFDeiTForImageClassificationWithTeacher,
					        TFDeiTForMaskedImageModeling,
					    )
					    if is_tf_available()
					    else ()
					)
					A__				 =    (
					    {
					        "feature-extraction": TFDeiTModel,
					        "image-classification": (TFDeiTForImageClassification, TFDeiTForImageClassificationWithTeacher),
					    }
					    if is_tf_available()
					    else {}
					)

					A__				 =    False
					A__				 =    False
					A__				 =    False
					A__				 =    False





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       TFDeiTModelTester(self	)
										lowerCamelCase__       =       ConfigTester(self	,  config_class=_lowerCAmelCase	,  has_text_modality=_lowerCAmelCase	,  hidden_size=37	)





					def 					__magic_name__ ( self	):
										self.config_tester.run_common_tests()





					@unittest.skip(reason="DeiT does not use inputs_embeds"	)
					def 					__magic_name__ ( self	):
										pass





					def 					__magic_name__ ( self	):
										lowerCamelCase__    ,		lowerCamelCase__       =       self.model_tester.prepare_config_and_inputs_for_common()

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





					def 					__magic_name__ ( self	):
										lowerCamelCase__    ,		lowerCamelCase__       =       self.model_tester.prepare_config_and_inputs_for_common()

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

															lowerCamelCase__       =       ["pixel_values"]
															self.assertListEqual(arg_names[:1]	,  _lowerCAmelCase	)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCAmelCase	)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_image_modeling(*_lowerCAmelCase	)





					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase	)





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	,  _lowerCAmelCase=False	):
										lowerCamelCase__       =       super()._prepare_for_class(_lowerCAmelCase	,  _lowerCAmelCase	,  return_labels=_lowerCAmelCase	)

										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 					__magic_name__ ( self	):
										for model_name in TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															lowerCamelCase__       =       TFDeiTModel.from_pretrained(_lowerCAmelCase	)
															self.assertIsNotNone(_lowerCAmelCase	)





def    __UpperCamelCase      (    ) ->int:
					lowerCamelCase__       =       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
					return image







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





					"""simple docstring"""


					@cached_property
					def 					__magic_name__ ( self	):
										return (
										    DeiTImageProcessor.from_pretrained("facebook/deit-base-distilled-patch16-224"	)
										    if is_vision_available()
										    else None
										)






					@slow
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       TFDeiTForImageClassificationWithTeacher.from_pretrained("facebook/deit-base-distilled-patch16-224"	)

										lowerCamelCase__       =       self.default_image_processor
										lowerCamelCase__       =       prepare_img()
										lowerCamelCase__       =       image_processor(images=_lowerCAmelCase	,  return_tensors="tf"	)

										# forward pass
										lowerCamelCase__       =       model(**_lowerCAmelCase	)

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

										lowerCamelCase__       =       tf.constant([-1.0_266, 0.1_912, -1.2_861]	)

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








import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam

class SCREAMING_SNAKE_CASE_ (					datasets.BeamBasedBuilder   ):





					"""simple docstring"""


					def 					__magic_name__ ( self	):
										return datasets.DatasetInfo(
										    features=datasets.Features({"content": datasets.Value("string"	)}	)	,  supervised_keys=_lowerCAmelCase	,  )





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										return [datasets.SplitGenerator(name=datasets.Split.TRAIN	,  gen_kwargs={"examples": get_test_dummy_examples()}	)]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										import apache_beam as beam

										return pipeline | "Load Examples" >> beam.Create(_lowerCAmelCase	)

class SCREAMING_SNAKE_CASE_ (					datasets.BeamBasedBuilder   ):





					"""simple docstring"""


					def 					__magic_name__ ( self	):
										return datasets.DatasetInfo(
										    features=datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"	)}	)}	)	,  supervised_keys=_lowerCAmelCase	,  )





					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										return [
										    datasets.SplitGenerator(name=datasets.Split.TRAIN	,  gen_kwargs={"examples": get_test_nested_examples()}	)
										]






					def 					__magic_name__ ( self	,  _lowerCAmelCase	,  _lowerCAmelCase	):
										import apache_beam as beam

										return pipeline | "Load Examples" >> beam.Create(_lowerCAmelCase	)





def    __UpperCamelCase      (    ) ->Optional[Any]:
					return [(i, {"content": content}) for i, content in enumerate(["foo", "bar", "foobar"])]





def    __UpperCamelCase      (    ) ->Optional[Any]:
					return [(i, {"a": {"b": [content]}}) for i, content in enumerate(["foo", "bar", "foobar"])]







class SCREAMING_SNAKE_CASE_ (					lowercase_   ):





					"""simple docstring"""


					@require_beam
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       len(get_test_dummy_examples()	)
										with tempfile.TemporaryDirectory() as tmp_cache_dir:
															lowerCamelCase__       =       DummyBeamDataset(cache_dir=_lowerCAmelCase	,  beam_runner="DirectRunner"	)
															builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(_lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  F"{builder.name}-train.arrow"	)	)	)
															self.assertDictEqual(builder.info.features	,  datasets.Features({"content": datasets.Value("string"	)}	)	)
															lowerCamelCase__       =       builder.as_dataset()
															self.assertEqual(dset["train"].num_rows	,  _lowerCAmelCase	)
															self.assertEqual(dset["train"].info.splits["train"].num_examples	,  _lowerCAmelCase	)
															self.assertDictEqual(dset["train"][0]	,  get_test_dummy_examples()[0][1]	)
															self.assertDictEqual(
															    dset["train"][expected_num_examples - 1]	,  get_test_dummy_examples()[expected_num_examples - 1][1]	)
															self.assertTrue(
															    os.path.exists(os.path.join(_lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  "dataset_info.json"	)	)	)
															del dset





					@require_beam
					def 					__magic_name__ ( self	):
										import apache_beam as beam

										lowerCamelCase__       =       beam.io.parquetio.WriteToParquet

										lowerCamelCase__       =       len(get_test_dummy_examples()	)
										with tempfile.TemporaryDirectory() as tmp_cache_dir:
															lowerCamelCase__       =       DummyBeamDataset(cache_dir=_lowerCAmelCase	,  beam_runner="DirectRunner"	)
															with patch("apache_beam.io.parquetio.WriteToParquet"	) as write_parquet_mock:
																				lowerCamelCase__       =       partial(_lowerCAmelCase	,  num_shards=2	)
																				builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(
															            _lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  F"{builder.name}-train-00000-of-00002.arrow"	)	)	)
															self.assertTrue(
															    os.path.exists(
															        os.path.join(
															            _lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  F"{builder.name}-train-00000-of-00002.arrow"	)	)	)
															self.assertDictEqual(builder.info.features	,  datasets.Features({"content": datasets.Value("string"	)}	)	)
															lowerCamelCase__       =       builder.as_dataset()
															self.assertEqual(dset["train"].num_rows	,  _lowerCAmelCase	)
															self.assertEqual(dset["train"].info.splits["train"].num_examples	,  _lowerCAmelCase	)
															# Order is not preserved when sharding, so we just check that all the elements are there
															self.assertListEqual(sorted(dset["train"]["content"]	)	,  sorted(["foo", "bar", "foobar"]	)	)
															self.assertTrue(
															    os.path.exists(os.path.join(_lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  "dataset_info.json"	)	)	)
															del dset





					@require_beam
					def 					__magic_name__ ( self	):
										with tempfile.TemporaryDirectory() as tmp_cache_dir:
															lowerCamelCase__       =       DummyBeamDataset(cache_dir=_lowerCAmelCase	)
															self.assertRaises(datasets.builder.MissingBeamOptions	,  builder.download_and_prepare	)






					@require_beam
					def 					__magic_name__ ( self	):
										lowerCamelCase__       =       len(get_test_nested_examples()	)
										with tempfile.TemporaryDirectory() as tmp_cache_dir:
															lowerCamelCase__       =       NestedBeamDataset(cache_dir=_lowerCAmelCase	,  beam_runner="DirectRunner"	)
															builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(_lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  F"{builder.name}-train.arrow"	)	)	)
															self.assertDictEqual(
															    builder.info.features	,  datasets.Features({"a": datasets.Sequence({"b": datasets.Value("string"	)}	)}	)	)
															lowerCamelCase__       =       builder.as_dataset()
															self.assertEqual(dset["train"].num_rows	,  _lowerCAmelCase	)
															self.assertEqual(dset["train"].info.splits["train"].num_examples	,  _lowerCAmelCase	)
															self.assertDictEqual(dset["train"][0]	,  get_test_nested_examples()[0][1]	)
															self.assertDictEqual(
															    dset["train"][expected_num_examples - 1]	,  get_test_nested_examples()[expected_num_examples - 1][1]	)
															self.assertTrue(
															    os.path.exists(os.path.join(_lowerCAmelCase	,  builder.name	,  "default"	,  "0.0.0"	,  "dataset_info.json"	)	)	)
															del dset

