


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       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									return None



class 						__a :
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			List[str]:




									"""simple docstring"""





									return None



class 						__a (							unittest.TestCase     ):
		_a      :    List[Any]           =      [
		    # (model_name, model_kwargs)
		    ('''bert-base-cased''', {}),
		    ('''gpt2''', {'''use_cache''': False}),  # We don't support exporting GPT2 past keys anymore
		]
		@require_tf
		@slow
		def       UpperCAmelCase__							(     self		)					->			Optional[Any]:




									"""simple docstring"""





									for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
																self._test_export(_lowercase     ,       'tf'     ,       12     ,       **_lowercase		)
		@require_torch
		@slow
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
																self._test_export(_lowercase     ,       'pt'     ,       12     ,       **_lowercase		)
		@require_torch
		@slow
		def       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									from transformers import BertModel

									_UpperCAmelCase											=     ["""[UNK]""", """[SEP]""", """[CLS]""", """[PAD]""", """[MASK]""", """some""", """other""", """words"""]
									with NamedTemporaryFile(mode='w+t'		) as vocab_file:
																vocab_file.write('\n'.join(_lowercase		)		)
																vocab_file.flush()
																_UpperCAmelCase											=     BertTokenizerFast(vocab_file.name		)

									with TemporaryDirectory() as bert_save_dir:
																_UpperCAmelCase											=     BertModel(BertConfig(vocab_size=len(_lowercase		)		)		)
																model.save_pretrained(_lowercase		)
																self._test_export(_lowercase     ,       'pt'     ,       12     ,       _lowercase		)
		@require_tf
		@slow
		def       UpperCAmelCase__							(     self		)					->			Any:




									"""simple docstring"""





									for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
																_UpperCAmelCase											=     self._test_export(_lowercase     ,       'tf'     ,       12     ,       **_lowercase		)
																_UpperCAmelCase											=     quantize(Path(_lowercase		)		)

																# Ensure the actual quantized model is not bigger than the original one
																if quantized_path.stat().st_size >= Path(_lowercase		).stat().st_size:
																							self.fail('Quantized model is bigger than initial ONNX model'		)
		@require_torch
		@slow
		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
																_UpperCAmelCase											=     self._test_export(_lowercase     ,       'pt'     ,       12     ,       **_lowercase		)
																_UpperCAmelCase											=     quantize(_lowercase		)

																# Ensure the actual quantized model is not bigger than the original one
																if quantized_path.stat().st_size >= Path(_lowercase		).stat().st_size:
																							self.fail('Quantized model is bigger than initial ONNX model'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE=None     ,       **_SCREAMING_SNAKE_CASE		)					->			Optional[Any]:




									"""simple docstring"""





									try:
																# Compute path
																with TemporaryDirectory() as tempdir:
																							_UpperCAmelCase											=     Path(_lowercase		).joinpath('model.onnx'		)

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

																# Export
																convert(_lowercase     ,       _lowercase     ,       _lowercase     ,       _lowercase     ,       _lowercase     ,       **_lowercase		)

																return path
									except Exception as e:
																self.fail(_lowercase		)
		@require_torch
		@require_tokenizers
		@slow
		def       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									from transformers import BertModel

									_UpperCAmelCase											=     BertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'		)		)
									_UpperCAmelCase											=     BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'		)
									self._test_infer_dynamic_axis(_lowercase     ,       _lowercase     ,       'pt'		)
		@require_tf
		@require_tokenizers
		@slow
		def       UpperCAmelCase__							(     self		)					->			Optional[int]:




									"""simple docstring"""





									from transformers import TFBertModel

									_UpperCAmelCase											=     TFBertModel(BertConfig.from_pretrained('lysandre/tiny-bert-random'		)		)
									_UpperCAmelCase											=     BertTokenizerFast.from_pretrained('lysandre/tiny-bert-random'		)
									self._test_infer_dynamic_axis(_lowercase     ,       _lowercase     ,       'tf'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Dict:




									"""simple docstring"""





									_UpperCAmelCase											=     FeatureExtractionPipeline(_lowercase     ,       _lowercase		)

									_UpperCAmelCase											=     ["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
									_UpperCAmelCase											=     infer_shapes(_lowercase     ,       _lowercase		)

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

									# 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       UpperCAmelCase__							(     self		)					->			List[str]:




									"""simple docstring"""





									_UpperCAmelCase											=     ["""input_ids""", """attention_mask""", """token_type_ids"""]
									_UpperCAmelCase											=     {"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
									_UpperCAmelCase											=     ensure_valid_input(FuncContiguousArgs()     ,       _lowercase     ,       _lowercase		)

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

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

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

									# Generated args are interleaved with another args (for instance parameter "past" in GPT2)
									_UpperCAmelCase											=     ensure_valid_input(FuncNonContiguousArgs()     ,       _lowercase     ,       _lowercase		)

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

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






		def       UpperCAmelCase__							(     self		)					->			str:




									"""simple docstring"""





									_UpperCAmelCase											=     generate_identified_filename(Path('/home/something/my_fake_model.onnx'		)     ,       '-test'		)
									self.assertEqual('/home/something/my_fake_model-test.onnx'     ,       generated.as_posix()		)



import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
	import torch

	from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM



@require_torch
class        A__:
						"""simple docstring"""






						def __init__(    self  ,      _lowercase  ,      _lowercase=99  ,      _lowercase=13  ,      _lowercase=16  ,      _lowercase=7  ,      _lowercase=True  ,      _lowercase=True  ,      _lowercase=True  ,      _lowercase=False  ,      _lowercase=True  ,      _lowercase=2  ,      _lowercase=32  ,      _lowercase=4  ,      _lowercase=4  ,      _lowercase=30  ,      _lowercase=0  ,      _lowercase=1  ,      _lowercase=2  ,      _lowercase=None  ,      )     -> List[str]:
									a_						:						Union[str, Any] 							= parent
									a_						:						Dict 							= batch_size
									a_						:						Any 							= decoder_seq_length
									# For common tests
									a_						:						Any 							= self.decoder_seq_length
									a_						:						int 							= is_training
									a_						:						List[Any] 							= use_attention_mask
									a_						:						List[str] 							= use_labels

									a_						:						Optional[Any] 							= vocab_size
									a_						:						Dict 							= d_model
									a_						:						Any 							= d_model
									a_						:						List[Any] 							= decoder_layers
									a_						:						List[str] 							= decoder_layers
									a_						:						Union[str, Any] 							= decoder_ffn_dim
									a_						:						Optional[Any] 							= decoder_attention_heads
									a_						:						Optional[Any] 							= decoder_attention_heads
									a_						:						Tuple 							= eos_token_id
									a_						:						Union[str, Any] 							= bos_token_id
									a_						:						Dict 							= pad_token_id
									a_						:						Union[str, Any] 							= decoder_start_token_id
									a_						:						List[Any] 							= use_cache
									a_						:						List[str] 							= max_position_embeddings

									a_						:						int 							= None
									a_						:						Any 							= decoder_seq_length
									a_						:						Union[str, Any] 							= 2
									a_						:						Optional[int] 							= 1





						def 			UpperCamelCase__					(    self    )     -> int:
									a_						:						Optional[int] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      self.vocab_size    )

									a_						:						Optional[int] 							= None
									if self.use_attention_mask:
												a_						:						List[Any] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      vocab_size=2    )

									a_						:						List[Any] 							= None
									if self.use_labels:
												a_						:						Union[str, Any] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      self.vocab_size    )

									a_						:						Any 							= TrOCRConfig(
									    vocab_size=self.vocab_size  ,      d_model=self.d_model  ,      decoder_layers=self.decoder_layers  ,      decoder_ffn_dim=self.decoder_ffn_dim  ,      decoder_attention_heads=self.decoder_attention_heads  ,      eos_token_id=self.eos_token_id  ,      bos_token_id=self.bos_token_id  ,      use_cache=self.use_cache  ,      pad_token_id=self.pad_token_id  ,      decoder_start_token_id=self.decoder_start_token_id  ,      max_position_embeddings=self.max_position_embeddings  ,      )

									return (config, input_ids, attention_mask, lm_labels)





						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      )     -> List[Any]:
									a_						:						Optional[int] 							= True
									a_						:						List[str] 							= TrOCRDecoder(config=_lowercase    ).to(_lowercase    ).eval()
									a_						:						str 							= input_ids[:2]

									input_ids[input_ids == 0] += 1
									# first forward pass
									a_						:						Dict 							= model(_lowercase  ,      use_cache=_lowercase    )
									a_						:						Union[str, Any] 							= model(_lowercase    )
									a_						:						Union[str, Any] 							= model(_lowercase  ,      use_cache=_lowercase    )

									self.parent.assertTrue(len(_lowercase    ) == len(_lowercase    )    )
									self.parent.assertTrue(len(_lowercase    ) == len(_lowercase    ) + 1    )

									a_						:						Dict 							= outputs["""past_key_values"""]

									# create hypothetical next token and extent to next_input_ids
									a_						:						List[str] 							= ids_tensor((2, 1)  ,      config.vocab_size - 1    ) + 1

									# append to next input_ids and
									a_						:						int 							= torch.cat([input_ids, next_tokens]  ,      dim=-1    )

									a_						:						Dict 							= model(_lowercase    )["""last_hidden_state"""]
									a_						:						List[str] 							= model(_lowercase  ,      past_key_values=_lowercase    )["""last_hidden_state"""]

									# select random slice
									a_						:						Optional[Any] 							= ids_tensor((1,)  ,      output_from_past.shape[-1]    ).item()
									a_						:						Union[str, Any] 							= output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
									a_						:						Any 							= output_from_past[:, 0, random_slice_idx].detach()

									# test that outputs are equal for slice
									assert torch.allclose(_lowercase  ,      _lowercase  ,      atol=1e-3    )





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						List[Any] 							= self.prepare_config_and_inputs()
									a_							,					a_							,					a_							,					a_						:						List[Any] 							= config_and_inputs

									a_						:						str 							= {"""input_ids""": input_ids, """attention_mask""": attention_mask}
									return config, inputs_dict







@require_torch
class        A__(a_,  a_,  a_,  unittest.TestCase							):
						"""simple docstring"""






						_A    :				Optional[Any]       					=			(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
						_A    :				Optional[int]       					=			(TrOCRForCausalLM,) if is_torch_available() else ()
						_A    :				Tuple       					=			{'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
						_A    :				List[str]       					=			True
						_A    :				List[str]       					=			False





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						str 							= TrOCRStandaloneDecoderModelTester(self  ,      is_training=_lowercase    )
									a_						:						str 							= ConfigTester(self  ,      config_class=_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									pass





						def 			UpperCamelCase__					(    self    )     -> int:
									pass





						def 			UpperCamelCase__					(    self    )     -> str:
									pass





						def 			UpperCamelCase__					(    self    )     -> Dict:
									self.config_tester.run_common_tests()





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									a_						:						int 							= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_decoder_model_past(*_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									return





						@unittest.skip("""The model doesn't support left padding"""    )  # and it's not used enough to be worth fixing :)
						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									pass


