




import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow




class 		__A      :

						def     A__						(		self   :Optional[int]    , __snake_case   :Any    , __snake_case   :str    , __snake_case   :Optional[int] ):
										'''simple docstring'''

										return None




class 		__A      :

						def     A__						(		self   :Dict    , __snake_case   :Union[str, Any]    , __snake_case   :Union[str, Any]    , __snake_case   :Optional[Any]    , __snake_case   :List[str] ):
										'''simple docstring'''

										return None




class 		__A      (			unittest.TestCase			):
						UpperCamelCase     						=					[
						    # (model_name, model_kwargs)
						    ('''bert-base-cased''', {}),
						    ('''gpt2''', {'''use_cache''': False}),  # We don't support exporting GPT2 past keys anymore
						]

						@require_tf
						@slow
						def     A__						(		self   :List[str] ):
										'''simple docstring'''

										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														self._test_export(__snake_case    , """tf"""    , 12    , **__snake_case )

						@require_torch
						@slow
						def     A__						(		self   :int ):
										'''simple docstring'''

										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														self._test_export(__snake_case    , """pt"""    , 12    , **__snake_case )

						@require_torch
						@slow
						def     A__						(		self   :Union[str, Any] ):
										'''simple docstring'''

										from transformers import BertModel

										__magic_name__	:	List[str]								=["""[UNK]""", """[SEP]""", """[CLS]""", """[PAD]""", """[MASK]""", """some""", """other""", """words"""]
										with NamedTemporaryFile(mode="""w+t""" ) as vocab_file:
														vocab_file.write("""\n""".join(__snake_case ) )
														vocab_file.flush()
														__magic_name__	:	Optional[Any]								=BertTokenizerFast(vocab_file.name )

										with TemporaryDirectory() as bert_save_dir:
														__magic_name__	:	str								=BertModel(BertConfig(vocab_size=len(__snake_case ) ) )
														model.save_pretrained(__snake_case )
														self._test_export(__snake_case    , """pt"""    , 12    , __snake_case )

						@require_tf
						@slow
						def     A__						(		self   :List[Any] ):
										'''simple docstring'''

										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														__magic_name__	:	Dict								=self._test_export(__snake_case    , """tf"""    , 12    , **__snake_case )
														__magic_name__	:	str								=quantize(Path(__snake_case ) )

														# Ensure the actual quantized model is not bigger than the original one
														if quantized_path.stat().st_size >= Path(__snake_case ).stat().st_size:
																		self.fail("""Quantized model is bigger than initial ONNX model""" )

						@require_torch
						@slow
						def     A__						(		self   :Union[str, Any] ):
										'''simple docstring'''

										for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
														__magic_name__	:	Tuple								=self._test_export(__snake_case    , """pt"""    , 12    , **__snake_case )
														__magic_name__	:	Tuple								=quantize(__snake_case )

														# Ensure the actual quantized model is not bigger than the original one
														if quantized_path.stat().st_size >= Path(__snake_case ).stat().st_size:
																		self.fail("""Quantized model is bigger than initial ONNX model""" )

						def     A__						(		self   :List[str]    , __snake_case   :List[Any]    , __snake_case   :Dict    , __snake_case   :Union[str, Any]    , __snake_case   :str=None    , **__snake_case   :Tuple ):
										'''simple docstring'''

										try:
														# Compute path
														with TemporaryDirectory() as tempdir:
																		__magic_name__	:	Union[str, Any]								=Path(__snake_case ).joinpath("""model.onnx""" )

														# Remove folder if exists
														if path.parent.exists():
																		path.parent.rmdir()

														# Export
														convert(__snake_case    , __snake_case    , __snake_case    , __snake_case    , __snake_case    , **__snake_case )

														return path
										except Exception as e:
														self.fail(__snake_case )

						@require_torch
						@require_tokenizers
						@slow
						def     A__						(		self   :int ):
										'''simple docstring'''

										from transformers import BertModel

										__magic_name__	:	Any								=BertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random""" ) )
										__magic_name__	:	int								=BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random""" )
										self._test_infer_dynamic_axis(__snake_case    , __snake_case    , """pt""" )

						@require_tf
						@require_tokenizers
						@slow
						def     A__						(		self   :Optional[int] ):
										'''simple docstring'''

										from transformers import TFBertModel

										__magic_name__	:	str								=TFBertModel(BertConfig.from_pretrained("""lysandre/tiny-bert-random""" ) )
										__magic_name__	:	int								=BertTokenizerFast.from_pretrained("""lysandre/tiny-bert-random""" )
										self._test_infer_dynamic_axis(__snake_case    , __snake_case    , """tf""" )

						def     A__						(		self   :int    , __snake_case   :List[str]    , __snake_case   :Tuple    , __snake_case   :int ):
										'''simple docstring'''

										__magic_name__	:	int								=FeatureExtractionPipeline(__snake_case    , __snake_case )

										__magic_name__	:	List[Any]								=["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
										__magic_name__	:	List[Any]								=infer_shapes(__snake_case    , __snake_case )

										# Assert all variables are present
										self.assertEqual(len(__snake_case )    , len(__snake_case ) )
										self.assertTrue(all(var_name in shapes for var_name in variable_names ) )
										self.assertSequenceEqual(variable_names[:3]    , __snake_case )
										self.assertSequenceEqual(variable_names[3:]    , __snake_case )

										# Assert inputs are {0: batch, 1: sequence}
										for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
														self.assertDictEqual(shapes[var_name]    , {0: """batch""", 1: """sequence"""} )

										# Assert outputs are {0: batch, 1: sequence} and {0: batch}
										self.assertDictEqual(shapes["""output_0"""]    , {0: """batch""", 1: """sequence"""} )
										self.assertDictEqual(shapes["""output_1"""]    , {0: """batch"""} )

						def     A__						(		self   :Optional[int] ):
										'''simple docstring'''

										__magic_name__	:	int								=["""input_ids""", """attention_mask""", """token_type_ids"""]
										__magic_name__	:	str								={"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
										__magic_name__	:	str								=ensure_valid_input(FuncContiguousArgs()    , __snake_case    , __snake_case )

										# Should have exactly the same number of args (all are valid)
										self.assertEqual(len(__snake_case )    , 3 )

										# Should have exactly the same input names
										self.assertEqual(set(__snake_case )    , set(__snake_case ) )

										# Parameter should be reordered according to their respective place in the function:
										# (input_ids, token_type_ids, attention_mask)
										self.assertEqual(__snake_case    , (tokens["""input_ids"""], tokens["""token_type_ids"""], tokens["""attention_mask"""]) )

										# Generated args are interleaved with another args (for instance parameter "past" in GPT2)
										__magic_name__	:	int								=ensure_valid_input(FuncNonContiguousArgs()    , __snake_case    , __snake_case )

										# Should have exactly the one arg (all before the one not provided "some_other_args")
										self.assertEqual(len(__snake_case )    , 1 )
										self.assertEqual(len(__snake_case )    , 1 )

										# Should have only "input_ids"
										self.assertEqual(inputs_args[0]    , tokens["""input_ids"""] )
										self.assertEqual(ordered_input_names[0]    , """input_ids""" )


						def     A__						(		self   :Dict ):
										'''simple docstring'''

										__magic_name__	:	List[Any]								=generate_identified_filename(Path("""/home/something/my_fake_model.onnx""" )    , """-test""" )
										self.assertEqual("""/home/something/my_fake_model-test.onnx"""    , generated.as_posix() )











import copy
import os
from typing import Union

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


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

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




class 		__A      (			UpperCamelCase__			):
						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   :List[Any]    , __snake_case   :Dict=5_02_44    , __snake_case   :Dict=7_68    , __snake_case   :Union[str, Any]=64    , __snake_case   :Union[str, Any]=20_48    , __snake_case   :List[Any]=12    , __snake_case   :Any=12    , __snake_case   :str=32    , __snake_case   :Dict=1_28    , __snake_case   :Optional[int]=0.1    , __snake_case   :Any=1E-6    , __snake_case   :Optional[int]=1.0    , __snake_case   :List[Any]="gelu_new"    , __snake_case   :Any=0    , __snake_case   :Tuple=False    , __snake_case   :Any=0    , __snake_case   :int=1    , __snake_case   :Dict=False    , __snake_case   :Optional[Any]=True    , **__snake_case   :Dict    , ):
										'''simple docstring'''

										__magic_name__	:	List[str]								=vocab_size
										__magic_name__	:	List[str]								=hidden_size
										__magic_name__	:	Optional[int]								=d_kv
										__magic_name__	:	Dict								=d_ff
										__magic_name__	:	Tuple								=num_layers
										__magic_name__	:	str								=num_heads
										__magic_name__	:	Optional[Any]								=relative_attention_num_buckets
										__magic_name__	:	List[Any]								=relative_attention_max_distance
										__magic_name__	:	str								=dropout_rate
										__magic_name__	:	Optional[Any]								=layer_norm_epsilon
										__magic_name__	:	Union[str, Any]								=initializer_factor
										__magic_name__	:	Tuple								=use_cache

										__magic_name__	:	Any								=eos_token_id
										__magic_name__	:	int								=decoder_start_token_id

										# for backwards compatibility
										__magic_name__	:	Union[str, Any]								=dense_act_fn

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


						@classmethod
						def     A__						(		cls   :Union[str, Any]    , __snake_case   :Union[str, os.PathLike]    , **__snake_case   :Any ):
										'''simple docstring'''

										cls._set_token_in_kwargs(__snake_case )

										__magic_name__				, __magic_name__	:	Optional[Any]								=cls.get_config_dict(__snake_case    , **__snake_case )

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




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					"""pix2struct_vision_model"""

						def __init__(		self   :Union[str, Any]    , __snake_case   :Tuple=7_68    , __snake_case   :Any=7_68    , __snake_case   :Tuple=20_48    , __snake_case   :Dict=64    , __snake_case   :List[str]=12    , __snake_case   :str=12    , __snake_case   :str="gelu_new"    , __snake_case   :Optional[int]=1E-6    , __snake_case   :Optional[int]=0.0    , __snake_case   :Any=0.0    , __snake_case   :Any=1E-10    , __snake_case   :List[str]=1.0    , __snake_case   :Tuple=40_96    , __snake_case   :Optional[int]=32    , __snake_case   :Union[str, Any]=1_28    , **__snake_case   :Optional[int]    , ):
										'''simple docstring'''

										super().__init__(**__snake_case )

										__magic_name__	:	Union[str, Any]								=hidden_size
										__magic_name__	:	Dict								=patch_embed_hidden_size
										__magic_name__	:	Tuple								=d_ff
										__magic_name__	:	str								=dropout_rate
										__magic_name__	:	str								=num_hidden_layers
										__magic_name__	:	Dict								=num_attention_heads
										__magic_name__	:	List[str]								=initializer_range
										__magic_name__	:	Optional[int]								=initializer_factor
										__magic_name__	:	str								=attention_dropout
										__magic_name__	:	Union[str, Any]								=layer_norm_eps
										__magic_name__	:	List[str]								=dense_act_fn
										__magic_name__	:	List[str]								=seq_len
										__magic_name__	:	str								=relative_attention_num_buckets
										__magic_name__	:	Optional[int]								=relative_attention_max_distance
										__magic_name__	:	List[str]								=d_kv


						@classmethod
						def     A__						(		cls   :Optional[Any]    , __snake_case   :Union[str, os.PathLike]    , **__snake_case   :Any ):
										'''simple docstring'''

										cls._set_token_in_kwargs(__snake_case )

										__magic_name__				, __magic_name__	:	Optional[int]								=cls.get_config_dict(__snake_case    , **__snake_case )

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




class 		__A      (			UpperCamelCase__			):
						UpperCamelCase     						=					"""pix2struct"""
						UpperCamelCase     						=					True

						def __init__(		self   :Union[str, Any]    , __snake_case   :Optional[int]=None    , __snake_case   :List[Any]=None    , __snake_case   :Optional[Any]=1.0    , __snake_case   :Tuple=0.02    , __snake_case   :str=False    , __snake_case   :List[str]=False    , __snake_case   :str=True    , **__snake_case   :List[str]    , ):
										'''simple docstring'''

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

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

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

										__magic_name__	:	Union[str, Any]								=PixaStructTextConfig(**__snake_case )
										__magic_name__	:	str								=PixaStructVisionConfig(**__snake_case )

										__magic_name__	:	int								=self.text_config.decoder_start_token_id
										__magic_name__	:	Optional[Any]								=self.text_config.pad_token_id
										__magic_name__	:	str								=self.text_config.eos_token_id

										__magic_name__	:	int								=initializer_factor
										__magic_name__	:	List[Any]								=initializer_range

										__magic_name__	:	Dict								=self.initializer_range
										__magic_name__	:	Union[str, Any]								=self.initializer_range

										__magic_name__	:	Tuple								=is_vqa

						@classmethod
						def     A__						(		cls   :Optional[int]    , __snake_case   :PixaStructTextConfig    , __snake_case   :PixaStructVisionConfig    , **__snake_case   :List[str] ):
										'''simple docstring'''

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


						def     A__						(		self   :Dict ):
										'''simple docstring'''

										__magic_name__	:	Union[str, Any]								=copy.deepcopy(self.__dict__ )
										__magic_name__	:	Any								=self.text_config.to_dict()
										__magic_name__	:	List[str]								=self.vision_config.to_dict()
										__magic_name__	:	Optional[int]								=self.__class__.model_type
										return output




