


import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


snake_case_    :      Union[str, Any]				    =				logging.get_logger(__name__)

snake_case_    :      Optional[int]				    =				{
    '''google/pix2struct-textcaps-base''': (
        '''https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'''
    ),
}

class     A__				(		a_	):
						UpperCAmelCase   							=    "pix2struct_text_model"
						UpperCAmelCase   							=    ["past_key_values"]
						UpperCAmelCase   							=    {
						    "hidden_size": "hidden_size",
						    "num_attention_heads": "num_heads",
						    "num_hidden_layers": "num_layers",
						}




						def __init__(							self       :   Optional[Any]					,     _a       :   Optional[int]=5_0244					,     _a       :   Union[str, Any]=768					,     _a       :   Optional[int]=64					,     _a       :   List[Any]=2048					,     _a       :   List[Any]=12					,     _a       :   Union[str, Any]=12					,     _a       :   Tuple=32					,     _a       :   List[str]=128					,     _a       :   List[str]=0.1					,     _a       :   Union[str, Any]=1E-6					,     _a       :   List[Any]=1.0					,     _a       :   str="gelu_new"					,     _a       :   List[str]=0					,     _a       :   Tuple=False					,     _a       :   Tuple=0					,     _a       :   Any=1					,     _a       :   Tuple=False					,     _a       :   Optional[int]=True					,     **_a       :   Optional[Any]					,     )				->  List[str]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=vocab_size
										_SCREAMING_SNAKE_CASE      	=hidden_size
										_SCREAMING_SNAKE_CASE      	=d_kv
										_SCREAMING_SNAKE_CASE      	=d_ff
										_SCREAMING_SNAKE_CASE      	=num_layers
										_SCREAMING_SNAKE_CASE      	=num_heads
										_SCREAMING_SNAKE_CASE      	=relative_attention_num_buckets
										_SCREAMING_SNAKE_CASE      	=relative_attention_max_distance
										_SCREAMING_SNAKE_CASE      	=dropout_rate
										_SCREAMING_SNAKE_CASE      	=layer_norm_epsilon
										_SCREAMING_SNAKE_CASE      	=initializer_factor
										_SCREAMING_SNAKE_CASE      	=use_cache

										_SCREAMING_SNAKE_CASE      	=eos_token_id
										_SCREAMING_SNAKE_CASE      	=decoder_start_token_id

										# for backwards compatibility
										_SCREAMING_SNAKE_CASE      	=dense_act_fn

										super().__init__(
										    pad_token_id=lowercase_					,     eos_token_id=lowercase_					,     decoder_start_token_id=lowercase_					,     tie_word_embeddings=lowercase_					,     is_decoder=lowercase_					,     **lowercase_					,     )




						@classmethod
						def 				__UpperCamelCase					(							cls       :   int					,     _a       :   Union[str, os.PathLike]					,     **_a       :   Optional[int]	)				->  Union[str, Any]:


										"""simple docstring"""





										cls._set_token_in_kwargs(lowercase_	)

										_SCREAMING_SNAKE_CASE      	=cls.get_config_dict(lowercase_					,     **lowercase_	)

										# get the text config dict if we are loading from Pix2StructConfig
										if config_dict.get('''model_type'''	) == "pix2struct":
														_SCREAMING_SNAKE_CASE      	=config_dict['''text_config''']

										if "model_type" in config_dict and hasattr(cls					,     '''model_type'''	) and config_dict["model_type"] != cls.model_type:
														logger.warning(
														    f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
														    f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."	)

										return cls.from_dict(lowercase_					,     **lowercase_	)

class     A__				(		a_	):
						UpperCAmelCase   							=    "pix2struct_vision_model"




						def __init__(							self       :   int					,     _a       :   Optional[int]=768					,     _a       :   Dict=768					,     _a       :   Dict=2048					,     _a       :   Optional[Any]=64					,     _a       :   List[Any]=12					,     _a       :   Optional[Any]=12					,     _a       :   Union[str, Any]="gelu_new"					,     _a       :   str=1E-6					,     _a       :   Any=0.0					,     _a       :   Any=0.0					,     _a       :   Optional[int]=1E-10					,     _a       :   str=1.0					,     _a       :   str=4096					,     _a       :   Union[str, Any]=32					,     _a       :   int=128					,     **_a       :   Tuple					,     )				->  Tuple:


										"""simple docstring"""





										super().__init__(**lowercase_	)

										_SCREAMING_SNAKE_CASE      	=hidden_size
										_SCREAMING_SNAKE_CASE      	=patch_embed_hidden_size
										_SCREAMING_SNAKE_CASE      	=d_ff
										_SCREAMING_SNAKE_CASE      	=dropout_rate
										_SCREAMING_SNAKE_CASE      	=num_hidden_layers
										_SCREAMING_SNAKE_CASE      	=num_attention_heads
										_SCREAMING_SNAKE_CASE      	=initializer_range
										_SCREAMING_SNAKE_CASE      	=initializer_factor
										_SCREAMING_SNAKE_CASE      	=attention_dropout
										_SCREAMING_SNAKE_CASE      	=layer_norm_eps
										_SCREAMING_SNAKE_CASE      	=dense_act_fn
										_SCREAMING_SNAKE_CASE      	=seq_len
										_SCREAMING_SNAKE_CASE      	=relative_attention_num_buckets
										_SCREAMING_SNAKE_CASE      	=relative_attention_max_distance
										_SCREAMING_SNAKE_CASE      	=d_kv




						@classmethod
						def 				__UpperCamelCase					(							cls       :   Union[str, Any]					,     _a       :   Union[str, os.PathLike]					,     **_a       :   int	)				->  Any:


										"""simple docstring"""





										cls._set_token_in_kwargs(lowercase_	)

										_SCREAMING_SNAKE_CASE      	=cls.get_config_dict(lowercase_					,     **lowercase_	)

										# get the vision config dict if we are loading from Pix2StructConfig
										if config_dict.get('''model_type'''	) == "pix2struct":
														_SCREAMING_SNAKE_CASE      	=config_dict['''vision_config''']

										if "model_type" in config_dict and hasattr(cls					,     '''model_type'''	) and config_dict["model_type"] != cls.model_type:
														logger.warning(
														    f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
														    f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."	)

										return cls.from_dict(lowercase_					,     **lowercase_	)

class     A__				(		a_	):
						UpperCAmelCase   							=    "pix2struct"
						UpperCAmelCase   							=    True




						def __init__(							self       :   Tuple					,     _a       :   Optional[int]=None					,     _a       :   Dict=None					,     _a       :   Dict=1.0					,     _a       :   Union[str, Any]=0.02					,     _a       :   Union[str, Any]=False					,     _a       :   int=False					,     _a       :   Optional[Any]=True					,     **_a       :   List[str]					,     )				->  Optional[int]:


										"""simple docstring"""





										super().__init__(tie_word_embeddings=lowercase_					,     is_encoder_decoder=lowercase_					,     **lowercase_	)

										if text_config is None:
														_SCREAMING_SNAKE_CASE      	={}
														logger.info('''text_config is None. Initializing the Pix2StructTextConfig with default values.'''	)

										if vision_config is None:
														_SCREAMING_SNAKE_CASE      	={}
														logger.info('''vision_config is None. Initializing the Pix2StructVisionConfig with default values.'''	)

										_SCREAMING_SNAKE_CASE      	=PixaStructTextConfig(**lowercase_	)
										_SCREAMING_SNAKE_CASE      	=PixaStructVisionConfig(**lowercase_	)

										_SCREAMING_SNAKE_CASE      	=self.text_config.decoder_start_token_id
										_SCREAMING_SNAKE_CASE      	=self.text_config.pad_token_id
										_SCREAMING_SNAKE_CASE      	=self.text_config.eos_token_id

										_SCREAMING_SNAKE_CASE      	=initializer_factor
										_SCREAMING_SNAKE_CASE      	=initializer_range

										_SCREAMING_SNAKE_CASE      	=self.initializer_range
										_SCREAMING_SNAKE_CASE      	=self.initializer_range

										_SCREAMING_SNAKE_CASE      	=is_vqa




						@classmethod
						def 				__UpperCamelCase					(							cls       :   Tuple					,     _a       :   PixaStructTextConfig					,     _a       :   PixaStructVisionConfig					,     **_a       :   str	)				->  Dict:


										"""simple docstring"""





										return cls(text_config=text_config.to_dict()					,     vision_config=vision_config.to_dict()					,     **lowercase_	)




						def 				__UpperCamelCase					(							self       :   List[str]	)				->  Optional[int]:


										"""simple docstring"""





										_SCREAMING_SNAKE_CASE      	=copy.deepcopy(self.__dict__	)
										_SCREAMING_SNAKE_CASE      	=self.text_config.to_dict()
										_SCREAMING_SNAKE_CASE      	=self.vision_config.to_dict()
										_SCREAMING_SNAKE_CASE      	=self.__class__.model_type
										return output


"""simple docstring"""





import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
							import torch

if is_vision_available():
							from PIL import Image

							from transformers import LevitImageProcessor







class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					def __init__(			self     :							Any  ,lowercase_     :							Dict  ,lowercase_     :							List[str]=7  ,lowercase_     :							Tuple=3  ,lowercase_     :							List[str]=1_8  ,lowercase_     :							Optional[Any]=3_0  ,lowercase_     :							List[Any]=4_0_0  ,lowercase_     :							List[Any]=True  ,lowercase_     :							Any=None  ,lowercase_     :							Optional[Any]=True  ,lowercase_     :							str=None  ,lowercase_     :							List[Any]=True  ,lowercase_     :							Dict=[0.5, 0.5, 0.5]  ,lowercase_     :							Dict=[0.5, 0.5, 0.5]  ,):
						lowerCAmelCase__      :		Any    							=      size if size is not None else {'''shortest_edge''': 1_8}
						lowerCAmelCase__      :		Union[str, Any]    							=      crop_size if crop_size is not None else {'''height''': 1_8, '''width''': 1_8}
						lowerCAmelCase__      :		Dict    							=      parent
						lowerCAmelCase__      :		Dict    							=      batch_size
						lowerCAmelCase__      :		List[str]    							=      num_channels
						lowerCAmelCase__      :		Any    							=      image_size
						lowerCAmelCase__      :		Union[str, Any]    							=      min_resolution
						lowerCAmelCase__      :		Dict    							=      max_resolution
						lowerCAmelCase__      :		List[str]    							=      do_resize
						lowerCAmelCase__      :		Optional[Any]    							=      size
						lowerCAmelCase__      :		Tuple    							=      do_center_crop
						lowerCAmelCase__      :		Optional[int]    							=      crop_size
						lowerCAmelCase__      :		List[str]    							=      do_normalize
						lowerCAmelCase__      :		Tuple    							=      image_mean
						lowerCAmelCase__      :		int    							=      image_std







					def 			__lowerCAmelCase					(			self     :							List[str]						):
						return {
						    "image_mean": self.image_mean,
						    "image_std": self.image_std,
						    "do_normalize": self.do_normalize,
						    "do_resize": self.do_resize,
						    "do_center_crop": self.do_center_crop,
						    "size": self.size,
						    "crop_size": self.crop_size,
						}







@require_torch
@require_vision
class 							SCREAMING_SNAKE_CASE   (						a_       ,				unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			LevitImageProcessor if is_vision_available() else None







					def 			__lowerCAmelCase					(			self     :							List[str]						):
						lowerCAmelCase__      :		Optional[Any]    							=      LevitImageProcessingTester(self						)







					@property
					def 			__lowerCAmelCase					(			self     :							Any						):
						return self.image_processor_tester.prepare_image_processor_dict()







					def 			__lowerCAmelCase					(			self     :							str						):
						lowerCAmelCase__      :		Optional[int]    							=      self.image_processing_class(**self.image_processor_dict						)
						self.assertTrue(hasattr(lowercase_  ,'''image_mean'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''image_std'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_normalize'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_resize'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''do_center_crop'''						)						)
						self.assertTrue(hasattr(lowercase_  ,'''size'''						)						)







					def 			__lowerCAmelCase					(			self     :							Any						):
						lowerCAmelCase__      :		int    							=      self.image_processing_class.from_dict(self.image_processor_dict						)
						self.assertEqual(image_processor.size  ,{'''shortest_edge''': 1_8}						)
						self.assertEqual(image_processor.crop_size  ,{'''height''': 1_8, '''width''': 1_8}						)

						lowerCAmelCase__      :		List[str]    							=      self.image_processing_class.from_dict(self.image_processor_dict  ,size=4_2  ,crop_size=8_4						)
						self.assertEqual(image_processor.size  ,{'''shortest_edge''': 4_2}						)
						self.assertEqual(image_processor.crop_size  ,{'''height''': 8_4, '''width''': 8_4}						)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						pass







					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						# Initialize image_processing
						lowerCAmelCase__      :		List[str]    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random PIL images
						lowerCAmelCase__      :		Union[str, Any]    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,Image.Image						)

						# Test not batched input
						lowerCAmelCase__      :		Any    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Optional[int]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)







					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						# Initialize image_processing
						lowerCAmelCase__      :		str    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random numpy tensors
						lowerCAmelCase__      :		Optional[Any]    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_  ,numpify=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,np.ndarray						)

						# Test not batched input
						lowerCAmelCase__      :		Tuple    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Union[str, Any]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)







					def 			__lowerCAmelCase					(			self     :							int						):
						# Initialize image_processing
						lowerCAmelCase__      :		Tuple    							=      self.image_processing_class(**self.image_processor_dict						)
						# create random PyTorch tensors
						lowerCAmelCase__      :		Dict    							=      prepare_image_inputs(self.image_processor_tester  ,equal_resolution=lowercase_  ,torchify=lowercase_						)
						for image in image_inputs:
							self.assertIsInstance(lowercase_  ,torch.Tensor						)

						# Test not batched input
						lowerCAmelCase__      :		Dict    							=      image_processing(image_inputs[0]  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)

						# Test batched
						lowerCAmelCase__      :		Optional[int]    							=      image_processing(lowercase_  ,return_tensors='''pt'''						).pixel_values
						self.assertEqual(
						    encoded_images.shape  ,(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size['''height'''],
						        self.image_processor_tester.crop_size['''width'''],
						    )  ,)




