


import warnings
from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType

class 						__a (							UpperCAmelCase     ):
		_a      :    List[str]           =      ['image_processor', 'tokenizer']
		_a      :    str           =      'FlavaImageProcessor'
		_a      :    Dict           =      ('BertTokenizer', 'BertTokenizerFast')
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       **_SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									_UpperCAmelCase											=     None
									if "feature_extractor" in kwargs:
																warnings.warn(
																    'The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'
																    ' instead.'     ,       _SCREAMING_SNAKE_CASE     ,       )
																_UpperCAmelCase											=     kwargs.pop('feature_extractor'		)

									_UpperCAmelCase											=     image_processor if image_processor is not None else feature_extractor
									if image_processor is None:
																raise ValueError('You need to specify an `image_processor`.'		)
									if tokenizer is None:
																raise ValueError('You need to specify a `tokenizer`.'		)

									super().__init__(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     self.image_processor
		def __call__(     self     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = True     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = 0     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = None     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = False     ,       _SCREAMING_SNAKE_CASE = True     ,       _SCREAMING_SNAKE_CASE = None     ,       **_SCREAMING_SNAKE_CASE     ,       )					->			List[str]:




									"""simple docstring"""





									if text is None and images is None:
																raise ValueError('You have to specify either text or images. Both cannot be none.'		)

									if text is not None:
																_UpperCAmelCase											=     self.tokenizer(
																    text=_SCREAMING_SNAKE_CASE     ,       add_special_tokens=_SCREAMING_SNAKE_CASE     ,       padding=_SCREAMING_SNAKE_CASE     ,       truncation=_SCREAMING_SNAKE_CASE     ,       max_length=_SCREAMING_SNAKE_CASE     ,       stride=_SCREAMING_SNAKE_CASE     ,       pad_to_multiple_of=_SCREAMING_SNAKE_CASE     ,       return_token_type_ids=_SCREAMING_SNAKE_CASE     ,       return_attention_mask=_SCREAMING_SNAKE_CASE     ,       return_overflowing_tokens=_SCREAMING_SNAKE_CASE     ,       return_special_tokens_mask=_SCREAMING_SNAKE_CASE     ,       return_offsets_mapping=_SCREAMING_SNAKE_CASE     ,       return_length=_SCREAMING_SNAKE_CASE     ,       verbose=_SCREAMING_SNAKE_CASE     ,       return_tensors=_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE     ,       )
									if images is not None:
																_UpperCAmelCase											=     self.image_processor(
																    _SCREAMING_SNAKE_CASE     ,       return_image_mask=_SCREAMING_SNAKE_CASE     ,       return_codebook_pixels=_SCREAMING_SNAKE_CASE     ,       return_tensors=_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE     ,       )

									if text is not None and images is not None:
																encoding.update(_SCREAMING_SNAKE_CASE		)
																return encoding
									elif text is not None:
																return encoding
									else:
																return BatchEncoding(data=dict(**_SCREAMING_SNAKE_CASE		)     ,       tensor_type=_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			Tuple:




									"""simple docstring"""





									return self.tokenizer.batch_decode(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									return self.tokenizer.decode(*_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)
		@property
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.tokenizer.model_input_names
									_UpperCAmelCase											=     self.image_processor.model_input_names
									return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names		)		)
		@property
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									warnings.warn(
									    '`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'     ,       _SCREAMING_SNAKE_CASE     ,       )
									return self.image_processor_class






		@property
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									warnings.warn(
									    '`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'     ,       _SCREAMING_SNAKE_CASE     ,       )
									return self.image_processor




import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
		import jax

		from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel

class 						__a (							unittest.TestCase     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=13     ,       _SCREAMING_SNAKE_CASE=30     ,       _SCREAMING_SNAKE_CASE=2     ,       _SCREAMING_SNAKE_CASE=3     ,       _SCREAMING_SNAKE_CASE=True     ,       _SCREAMING_SNAKE_CASE=True     ,       _SCREAMING_SNAKE_CASE=32     ,       _SCREAMING_SNAKE_CASE=5     ,       _SCREAMING_SNAKE_CASE=4     ,       _SCREAMING_SNAKE_CASE=37     ,       _SCREAMING_SNAKE_CASE="gelu"     ,       _SCREAMING_SNAKE_CASE=0.1     ,       _SCREAMING_SNAKE_CASE=0.1     ,       _SCREAMING_SNAKE_CASE=10     ,       _SCREAMING_SNAKE_CASE=0.02     ,       )					->			int:




									"""simple docstring"""





									_UpperCAmelCase											=     parent
									_UpperCAmelCase											=     batch_size
									_UpperCAmelCase											=     image_size
									_UpperCAmelCase											=     patch_size
									_UpperCAmelCase											=     num_channels
									_UpperCAmelCase											=     is_training
									_UpperCAmelCase											=     use_labels
									_UpperCAmelCase											=     hidden_size
									_UpperCAmelCase											=     num_hidden_layers
									_UpperCAmelCase											=     num_attention_heads
									_UpperCAmelCase											=     intermediate_size
									_UpperCAmelCase											=     hidden_act
									_UpperCAmelCase											=     hidden_dropout_prob
									_UpperCAmelCase											=     attention_probs_dropout_prob
									_UpperCAmelCase											=     type_sequence_label_size
									_UpperCAmelCase											=     initializer_range

									# in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
									_UpperCAmelCase											=     (image_size // patch_size) ** 2
									_UpperCAmelCase											=     num_patches + 1
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									_UpperCAmelCase											=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]		)

									_UpperCAmelCase											=     ViTConfig(
									    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=_SCREAMING_SNAKE_CASE     ,       initializer_range=self.initializer_range     ,       )

									return config, pixel_values
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     FlaxViTModel(config=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
									_UpperCAmelCase											=     (self.image_size, self.image_size)
									_UpperCAmelCase											=     (self.patch_size, self.patch_size)
									_UpperCAmelCase											=     (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
									self.parent.assertEqual(result.last_hidden_state.shape     ,       (self.batch_size, num_patches + 1, self.hidden_size)		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Any:




									"""simple docstring"""





									_UpperCAmelCase											=     self.type_sequence_label_size
									_UpperCAmelCase											=     FlaxViTForImageClassification(config=_SCREAMING_SNAKE_CASE		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)
									self.parent.assertEqual(result.logits.shape     ,       (self.batch_size, self.type_sequence_label_size)		)

									# test greyscale images
									_UpperCAmelCase											=     1
									_UpperCAmelCase											=     FlaxViTForImageClassification(_SCREAMING_SNAKE_CASE		)

									_UpperCAmelCase											=     floats_tensor([self.batch_size, 1, self.image_size, self.image_size]		)
									_UpperCAmelCase											=     model(_SCREAMING_SNAKE_CASE		)






		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     self.prepare_config_and_inputs()
									(
									    (
									    _UpperCAmelCase
									)						,			(
									    _UpperCAmelCase
									)						,			
									)											=     config_and_inputs
									_UpperCAmelCase											=     {'pixel_values': pixel_values}
									return config, inputs_dict



@require_flax
class 						__a (							UpperCAmelCase       , unittest.TestCase     ):
		_a      :    List[str]           =      (FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()
		def       UpperCAmelCase__							(     self		)					->			None:




									"""simple docstring"""





									_UpperCAmelCase											=     FlaxViTModelTester(self		)
									_UpperCAmelCase											=     ConfigTester(self     ,       config_class=_SCREAMING_SNAKE_CASE     ,       has_text_modality=_SCREAMING_SNAKE_CASE     ,       hidden_size=37		)
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									self.config_tester.run_common_tests()
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			List[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*_SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase						,			_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs_for_common()

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

																_UpperCAmelCase											=     ['pixel_values']
																self.assertListEqual(arg_names[:1]     ,       _SCREAMING_SNAKE_CASE		)
		def       UpperCAmelCase__							(     self		)					->			int:




									"""simple docstring"""





									_UpperCAmelCase						,			_UpperCAmelCase											=     self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
																with self.subTest(model_class.__name__		):
																							_UpperCAmelCase											=     self._prepare_for_class(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)
																							_UpperCAmelCase											=     model_class(_SCREAMING_SNAKE_CASE		)

																							@jax.jit
																							def model_jitted(_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		):
																														return model(pixel_values=_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)

																							with self.subTest('JIT Enabled'		):
																														_UpperCAmelCase											=     model_jitted(**_SCREAMING_SNAKE_CASE		).to_tuple()

																							with self.subTest('JIT Disabled'		):
																														with jax.disable_jit():
																																					_UpperCAmelCase											=     model_jitted(**_SCREAMING_SNAKE_CASE		).to_tuple()

																							self.assertEqual(len(_SCREAMING_SNAKE_CASE		)     ,       len(_SCREAMING_SNAKE_CASE		)		)
																							for jitted_output, output in zip(_SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		):
																														self.assertEqual(jitted_output.shape     ,       output.shape		)






		@slow
		def       UpperCAmelCase__							(     self		)					->			Dict:




									"""simple docstring"""





									for model_class_name in self.all_model_classes:
																_UpperCAmelCase											=     model_class_name.from_pretrained('google/vit-base-patch16-224'		)
																_UpperCAmelCase											=     model(np.ones((1, 3, 224, 224)		)		)
																self.assertIsNotNone(_SCREAMING_SNAKE_CASE		)
