


import unittest

from transformers import SqueezeBertConfig, is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers import (
	    SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
	    SqueezeBertForMaskedLM,
	    SqueezeBertForMultipleChoice,
	    SqueezeBertForQuestionAnswering,
	    SqueezeBertForSequenceClassification,
	    SqueezeBertForTokenClassification,
	    SqueezeBertModel,
	)



class 				A       (			UpperCAmelCase__				):
					'''simple docstring'''






					def __init__(self       :		Any ,   _UpperCAmelCase       :		Any ,   _UpperCAmelCase       :		Tuple=13 ,   _UpperCAmelCase       :		str=7 ,   _UpperCAmelCase       :		Optional[int]=True ,   _UpperCAmelCase       :		Union[str, Any]=True ,   _UpperCAmelCase       :		Union[str, Any]=False ,   _UpperCAmelCase       :		Optional[int]=True ,   _UpperCAmelCase       :		Any=99 ,   _UpperCAmelCase       :		Optional[Any]=32 ,   _UpperCAmelCase       :		Optional[Any]=5 ,   _UpperCAmelCase       :		Optional[Any]=4 ,   _UpperCAmelCase       :		Optional[Any]=64 ,   _UpperCAmelCase       :		Any="gelu" ,   _UpperCAmelCase       :		Optional[int]=0.1 ,   _UpperCAmelCase       :		Optional[int]=0.1 ,   _UpperCAmelCase       :		Optional[Any]=512 ,   _UpperCAmelCase       :		Any=16 ,   _UpperCAmelCase       :		int=2 ,   _UpperCAmelCase       :		Union[str, Any]=0.02 ,   _UpperCAmelCase       :		Tuple=3 ,   _UpperCAmelCase       :		Optional[int]=4 ,   _UpperCAmelCase       :		int=None ,   _UpperCAmelCase       :		Tuple=2 ,   _UpperCAmelCase       :		Tuple=2 ,   _UpperCAmelCase       :		str=2 ,   _UpperCAmelCase       :		Dict=2 ,   _UpperCAmelCase       :		str=4 ,   _UpperCAmelCase       :		Tuple=1 ,   )		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						parent
							lowercase__	     =						batch_size
							lowercase__	     =						seq_length
							lowercase__	     =						is_training
							lowercase__	     =						use_input_mask
							lowercase__	     =						use_token_type_ids
							lowercase__	     =						use_labels
							lowercase__	     =						vocab_size
							lowercase__	     =						hidden_size
							lowercase__	     =						num_hidden_layers
							lowercase__	     =						num_attention_heads
							lowercase__	     =						intermediate_size
							lowercase__	     =						hidden_act
							lowercase__	     =						hidden_dropout_prob
							lowercase__	     =						attention_probs_dropout_prob
							lowercase__	     =						max_position_embeddings
							lowercase__	     =						type_vocab_size
							lowercase__	     =						type_sequence_label_size
							lowercase__	     =						initializer_range
							lowercase__	     =						num_labels
							lowercase__	     =						num_choices
							lowercase__	     =						scope
							lowercase__	     =						q_groups
							lowercase__	     =						k_groups
							lowercase__	     =						v_groups
							lowercase__	     =						post_attention_groups
							lowercase__	     =						intermediate_groups
							lowercase__	     =						output_groups

					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  Union[str, Any]:




							"""simple docstring"""


							lowercase__	     =						ids_tensor([self.batch_size, self.seq_length] ,   self.vocab_size						)

							lowercase__	     =						None
							if self.use_input_mask:
									lowercase__	     =						random_attention_mask([self.batch_size, self.seq_length]						)

							lowercase__	     =						None
							lowercase__	     =						None
							lowercase__	     =						None
							if self.use_labels:
									lowercase__	     =						ids_tensor([self.batch_size] ,   self.type_sequence_label_size						)
									lowercase__	     =						ids_tensor([self.batch_size, self.seq_length] ,   self.num_labels						)
									lowercase__	     =						ids_tensor([self.batch_size] ,   self.num_choices						)

							lowercase__	     =						self.get_config()

							return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels

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




							"""simple docstring"""


							return SqueezeBertConfig(
							    embedding_size=self.hidden_size ,   vocab_size=self.vocab_size ,   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 ,   attention_probs_dropout_prob=self.hidden_dropout_prob ,   attention_dropout=self.attention_probs_dropout_prob ,   max_position_embeddings=self.max_position_embeddings ,   initializer_range=self.initializer_range ,   q_groups=self.q_groups ,   k_groups=self.k_groups ,   v_groups=self.v_groups ,   post_attention_groups=self.post_attention_groups ,   intermediate_groups=self.intermediate_groups ,   output_groups=self.output_groups ,   )

					def     lowerCamelCase__	(self       :		List[str] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		List[Any] ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Any ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		List[Any]						)		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						SqueezeBertModel(config=_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()
							lowercase__	     =						model(_UpperCAmelCase ,   _UpperCAmelCase						)
							lowercase__	     =						model(_UpperCAmelCase						)
							self.parent.assertEqual(result.last_hidden_state.shape ,   (self.batch_size, self.seq_length, self.hidden_size)						)

					def     lowerCamelCase__	(self       :		List[str] ,   _UpperCAmelCase       :		List[str] ,   _UpperCAmelCase       :		Optional[Any] ,   _UpperCAmelCase       :		List[Any] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		Dict ,   _UpperCAmelCase       :		Tuple						)		->  Tuple:




							"""simple docstring"""


							lowercase__	     =						SqueezeBertForMaskedLM(config=_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()
							lowercase__	     =						model(_UpperCAmelCase ,   attention_mask=_UpperCAmelCase ,   labels=_UpperCAmelCase						)
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.vocab_size)						)

					def     lowerCamelCase__	(self       :		str ,   _UpperCAmelCase       :		Dict ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		List[str] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		int						)		->  Union[str, Any]:




							"""simple docstring"""


							lowercase__	     =						SqueezeBertForQuestionAnswering(config=_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()
							lowercase__	     =						model(
							    _UpperCAmelCase ,   attention_mask=_UpperCAmelCase ,   start_positions=_UpperCAmelCase ,   end_positions=_UpperCAmelCase						)
							self.parent.assertEqual(result.start_logits.shape ,   (self.batch_size, self.seq_length)						)
							self.parent.assertEqual(result.end_logits.shape ,   (self.batch_size, self.seq_length)						)

					def     lowerCamelCase__	(self       :		Optional[Any] ,   _UpperCAmelCase       :		Optional[Any] ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		List[Any]						)		->  List[str]:




							"""simple docstring"""


							lowercase__	     =						self.num_labels
							lowercase__	     =						SqueezeBertForSequenceClassification(_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()
							lowercase__	     =						model(_UpperCAmelCase ,   attention_mask=_UpperCAmelCase ,   labels=_UpperCAmelCase						)
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_labels)						)

					def     lowerCamelCase__	(self       :		Optional[Any] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		Optional[int] ,   _UpperCAmelCase       :		str ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Tuple						)		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						self.num_labels
							lowercase__	     =						SqueezeBertForTokenClassification(config=_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()

							lowercase__	     =						model(_UpperCAmelCase ,   attention_mask=_UpperCAmelCase ,   labels=_UpperCAmelCase						)
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.seq_length, self.num_labels)						)

					def     lowerCamelCase__	(self       :		Tuple ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		Any ,   _UpperCAmelCase       :		int ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		Union[str, Any] ,   _UpperCAmelCase       :		List[Any]						)		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						self.num_choices
							lowercase__	     =						SqueezeBertForMultipleChoice(config=_UpperCAmelCase						)
							model.to(_UpperCAmelCase						)
							model.eval()
							lowercase__	     =						input_ids.unsqueeze(1						).expand(-1 ,   self.num_choices ,   -1						).contiguous()
							lowercase__	     =						input_mask.unsqueeze(1						).expand(-1 ,   self.num_choices ,   -1						).contiguous()
							lowercase__	     =						model(
							    _UpperCAmelCase ,   attention_mask=_UpperCAmelCase ,   labels=_UpperCAmelCase ,   )
							self.parent.assertEqual(result.logits.shape ,   (self.batch_size, self.num_choices)						)


					def     lowerCamelCase__	(self       :		str						)		->  str:




							"""simple docstring"""


							lowercase__	     =						self.prepare_config_and_inputs()
							((lowercase__)     ,      (lowercase__)     ,      (lowercase__)     ,      (lowercase__)     ,      (lowercase__)     ,      (lowercase__))	     =						config_and_inputs
							lowercase__	     =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
							return config, inputs_dict






@require_torch
class 				A       (			UpperCAmelCase__					,    UpperCAmelCase__					,    unittest.TestCase				):
					'''simple docstring'''






					A__     =					(
					    (
					        SqueezeBertModel,
					        SqueezeBertForMaskedLM,
					        SqueezeBertForMultipleChoice,
					        SqueezeBertForQuestionAnswering,
					        SqueezeBertForSequenceClassification,
					        SqueezeBertForTokenClassification,
					    )
					    if is_torch_available()
					    else None
					)
					A__     =					(
					    {
					        '''feature-extraction''': SqueezeBertModel,
					        '''fill-mask''': SqueezeBertForMaskedLM,
					        '''question-answering''': SqueezeBertForQuestionAnswering,
					        '''text-classification''': SqueezeBertForSequenceClassification,
					        '''token-classification''': SqueezeBertForTokenClassification,
					        '''zero-shot''': SqueezeBertForSequenceClassification,
					    }
					    if is_torch_available()
					    else {}
					)
					A__     =					False
					A__     =					True
					A__     =					False

					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  List[Any]:




							"""simple docstring"""


							lowercase__	     =						SqueezeBertModelTester(self						)
							lowercase__	     =						ConfigTester(self ,   config_class=_UpperCAmelCase ,   dim=37						)

					def     lowerCamelCase__	(self       :		Any						)		->  Any:




							"""simple docstring"""


							self.config_tester.run_common_tests()

					def     lowerCamelCase__	(self       :		Union[str, Any]						)		->  Optional[int]:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_model(*_UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  int:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_masked_lm(*_UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		str						)		->  Any:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_question_answering(*_UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		str						)		->  Tuple:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_sequence_classification(*_UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  Optional[int]:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_token_classification(*_UpperCAmelCase						)

					def     lowerCamelCase__	(self       :		Tuple						)		->  str:




							"""simple docstring"""


							lowercase__	     =						self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_squeezebert_for_multiple_choice(*_UpperCAmelCase						)


					@slow
					def     lowerCamelCase__	(self       :		Optional[Any]						)		->  List[str]:




							"""simple docstring"""


							for model_name in SQUEEZEBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase__	     =						SqueezeBertModel.from_pretrained(_UpperCAmelCase						)
									self.assertIsNotNone(_UpperCAmelCase						)






@require_sentencepiece
@require_tokenizers
@require_torch
class 				A       (			unittest.TestCase				):
					'''simple docstring'''






					@slow
					def     lowerCamelCase__	(self       :		int						)		->  Dict:




							"""simple docstring"""


							lowercase__	     =						SqueezeBertForSequenceClassification.from_pretrained("""squeezebert/squeezebert-mnli"""						)

							lowercase__	     =						torch.tensor([[1, 2_9414, 232, 328, 740, 1140, 1_2695, 69, 13, 1588, 2]]						)
							lowercase__	     =						model(_UpperCAmelCase						)[0]
							lowercase__	     =						torch.Size((1, 3)						)
							self.assertEqual(output.shape ,   _UpperCAmelCase						)
							lowercase__	     =						torch.tensor([[0.6_401, -0.0_349, -0.6_041]]						)
							self.assertTrue(torch.allclose(_UpperCAmelCase ,   _UpperCAmelCase ,   atol=1E-4						)						)




'''simple docstring'''







import copy
import os
from typing import Union

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


__lowerCamelCase									=     logging.get_logger(__name__)

__lowerCamelCase									=     {
    '''google/pix2struct-textcaps-base''': (
        '''https://huggingface.co/google/pix2struct-textcaps-base/resolve/main/config.json'''
    ),
}


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



    def __init__(       self		,					UpperCamelCase__=50244		,					UpperCamelCase__=768		,					UpperCamelCase__=64		,					UpperCamelCase__=2048		,					UpperCamelCase__=12		,					UpperCamelCase__=12		,					UpperCamelCase__=32		,					UpperCamelCase__=128		,					UpperCamelCase__=0.1		,					UpperCamelCase__=1e-6		,					UpperCamelCase__=1.0		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=0		,					UpperCamelCase__=False		,					UpperCamelCase__=0		,					UpperCamelCase__=1		,					UpperCamelCase__=False		,					UpperCamelCase__=True		,					**UpperCamelCase__		,					)    ->   List[str]:





           '''simple docstring'''
           A_						       =   vocab_size
           A_						       =   hidden_size
           A_						       =   d_kv
           A_						       =   d_ff
           A_						       =   num_layers
           A_						       =   num_heads
           A_						       =   relative_attention_num_buckets
           A_						       =   relative_attention_max_distance
           A_						       =   dropout_rate
           A_						       =   layer_norm_epsilon
           A_						       =   initializer_factor
           A_						       =   use_cache

           A_						       =   eos_token_id
           A_						       =   decoder_start_token_id

           # for backwards compatibility
           A_						       =   dense_act_fn

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







    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   "PretrainedConfig":





           '''simple docstring'''
           cls._set_token_in_kwargs(UpperCamelCase__							)

           A_       ,      A_						       =   cls.get_config_dict(UpperCamelCase__		,					**UpperCamelCase__							)

           # get the text config dict if we are loading from Pix2StructConfig
           if config_dict.get("""model_type"""							) == "pix2struct":
                  A_						       =   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(UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):
    lowercase						       =		"pix2struct_vision_model"



    def __init__(       self		,					UpperCamelCase__=768		,					UpperCamelCase__=768		,					UpperCamelCase__=2048		,					UpperCamelCase__=64		,					UpperCamelCase__=12		,					UpperCamelCase__=12		,					UpperCamelCase__="gelu_new"		,					UpperCamelCase__=1e-6		,					UpperCamelCase__=0.0		,					UpperCamelCase__=0.0		,					UpperCamelCase__=1e-1_0		,					UpperCamelCase__=1.0		,					UpperCamelCase__=4096		,					UpperCamelCase__=32		,					UpperCamelCase__=128		,					**UpperCamelCase__		,					)    ->   int:





           '''simple docstring'''
           super().__init__(**UpperCamelCase__							)

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







    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   "PretrainedConfig":





           '''simple docstring'''
           cls._set_token_in_kwargs(UpperCamelCase__							)

           A_       ,      A_						       =   cls.get_config_dict(UpperCamelCase__		,					**UpperCamelCase__							)

           # get the vision config dict if we are loading from Pix2StructConfig
           if config_dict.get("""model_type"""							) == "pix2struct":
                  A_						       =   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(UpperCamelCase__		,					**UpperCamelCase__							)


class        A__   (   _snake_case      ):
    lowercase						       =		"pix2struct"
    lowercase						       =		True



    def __init__(       self		,					UpperCamelCase__=None		,					UpperCamelCase__=None		,					UpperCamelCase__=1.0		,					UpperCamelCase__=0.02		,					UpperCamelCase__=False		,					UpperCamelCase__=False		,					UpperCamelCase__=True		,					**UpperCamelCase__		,					)    ->   str:





           '''simple docstring'''
           super().__init__(tie_word_embeddings=UpperCamelCase__		,					is_encoder_decoder=UpperCamelCase__		,					**UpperCamelCase__							)

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

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

           A_						       =   PixaStructTextConfig(**UpperCamelCase__							)
           A_						       =   PixaStructVisionConfig(**UpperCamelCase__							)

           A_						       =   self.text_config.decoder_start_token_id
           A_						       =   self.text_config.pad_token_id
           A_						       =   self.text_config.eos_token_id

           A_						       =   initializer_factor
           A_						       =   initializer_range

           A_						       =   self.initializer_range
           A_						       =   self.initializer_range

           A_						       =   is_vqa



    @classmethod
    def        snake_case_      (       cls		,					UpperCamelCase__		,					UpperCamelCase__		,					**UpperCamelCase__							)    ->   Optional[int]:





           '''simple docstring'''
           return cls(text_config=text_config.to_dict()		,					vision_config=vision_config.to_dict()		,					**UpperCamelCase__							)







    def        snake_case_      (       self							)    ->   List[Any]:





           '''simple docstring'''
           A_						       =   copy.deepcopy(self.__dict__							)
           A_						       =   self.text_config.to_dict()
           A_						       =   self.vision_config.to_dict()
           A_						       =   self.__class__.model_type
           return output
