

"""simple docstring"""





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 							SCREAMING_SNAKE_CASE   :






					"""simple docstring"""



					def 			__lowerCAmelCase					(			self     :							List[str]  ,lowercase_     :							List[str]  ,lowercase_     :							Any  ,lowercase_     :							Any						):
						return None







class 							SCREAMING_SNAKE_CASE   :






					"""simple docstring"""



					def 			__lowerCAmelCase					(			self     :							Optional[Any]  ,lowercase_     :							Union[str, Any]  ,lowercase_     :							Union[str, Any]  ,lowercase_     :							int  ,lowercase_     :							str						):
						return None







class 							SCREAMING_SNAKE_CASE   (						unittest.TestCase			):






					"""simple docstring"""



					lowercase__												=			[
					    # (model_name, model_kwargs)
					    ("bert-base-cased", {}),
					    ("gpt2", {"use_cache": False}),  # We don't support exporting GPT2 past keys anymore
					]







					@require_tf
					@slow
					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
							self._test_export(__snake_case  ,'''tf'''  ,1_2  ,**__snake_case						)







					@require_torch
					@slow
					def 			__lowerCAmelCase					(			self     :							List[str]						):
						for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
							self._test_export(__snake_case  ,'''pt'''  ,1_2  ,**__snake_case						)







					@require_torch
					@slow
					def 			__lowerCAmelCase					(			self     :							Optional[int]						):
						from transformers import BertModel

						lowerCAmelCase__      :		Optional[Any]    							=      ["""[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()
							lowerCAmelCase__      :		str    							=      BertTokenizerFast(vocab_file.name						)

						with TemporaryDirectory() as bert_save_dir:
							lowerCAmelCase__      :		Optional[int]    							=      BertModel(BertConfig(vocab_size=len(__snake_case						)						)						)
							model.save_pretrained(__snake_case						)
							self._test_export(__snake_case  ,'''pt'''  ,1_2  ,__snake_case						)







					@require_tf
					@slow
					def 			__lowerCAmelCase					(			self     :							Tuple						):
						for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
							lowerCAmelCase__      :		int    							=      self._test_export(__snake_case  ,'''tf'''  ,1_2  ,**__snake_case						)
							lowerCAmelCase__      :		List[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 			__lowerCAmelCase					(			self     :							List[Any]						):
						for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
							lowerCAmelCase__      :		Dict    							=      self._test_export(__snake_case  ,'''pt'''  ,1_2  ,**__snake_case						)
							lowerCAmelCase__      :		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 			__lowerCAmelCase					(			self     :							Tuple  ,lowercase_     :							Optional[Any]  ,lowercase_     :							Optional[Any]  ,lowercase_     :							Any  ,lowercase_     :							Tuple=None  ,**lowercase_     :							int						):
						try:
							# Compute path
							with TemporaryDirectory() as tempdir:
								lowerCAmelCase__      :		Tuple    							=      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 			__lowerCAmelCase					(			self     :							List[str]						):
						from transformers import BertModel

						lowerCAmelCase__      :		str    							=      BertModel(BertConfig.from_pretrained('''lysandre/tiny-bert-random'''						)						)
						lowerCAmelCase__      :		int    							=      BertTokenizerFast.from_pretrained('''lysandre/tiny-bert-random'''						)
						self._test_infer_dynamic_axis(__snake_case  ,__snake_case  ,'''pt'''						)







					@require_tf
					@require_tokenizers
					@slow
					def 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						from transformers import TFBertModel

						lowerCAmelCase__      :		str    							=      TFBertModel(BertConfig.from_pretrained('''lysandre/tiny-bert-random'''						)						)
						lowerCAmelCase__      :		List[Any]    							=      BertTokenizerFast.from_pretrained('''lysandre/tiny-bert-random'''						)
						self._test_infer_dynamic_axis(__snake_case  ,__snake_case  ,'''tf'''						)







					def 			__lowerCAmelCase					(			self     :							str  ,lowercase_     :							Union[str, Any]  ,lowercase_     :							int  ,lowercase_     :							int						):
						lowerCAmelCase__      :		List[Any]    							=      FeatureExtractionPipeline(__snake_case  ,__snake_case						)

						lowerCAmelCase__      :		Optional[int]    							=      ["""input_ids""", """token_type_ids""", """attention_mask""", """output_0""", """output_1"""]
						lowerCAmelCase__      :		List[str]    							=      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 			__lowerCAmelCase					(			self     :							int						):
						lowerCAmelCase__      :		int    							=      ["""input_ids""", """attention_mask""", """token_type_ids"""]
						lowerCAmelCase__      :		str    							=      {"""input_ids""": [1, 2, 3, 4], """attention_mask""": [0, 0, 0, 0], """token_type_ids""": [1, 1, 1, 1]}
						lowerCAmelCase__      :		int    							=      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)
						lowerCAmelCase__      :		List[Any]    							=      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 			__lowerCAmelCase					(			self     :							Union[str, Any]						):
						lowerCAmelCase__      :		Optional[int]    							=      generate_identified_filename(Path('''/home/something/my_fake_model.onnx'''						)  ,'''-test'''						)
						self.assertEqual('''/home/something/my_fake_model-test.onnx'''  ,generated.as_posix()						)












"""simple docstring"""





import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, 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.models.esm.modeling_esmfold import EsmForProteinFolding



class       __A		:

					def __init__(			self  :      str     ,		__snake_case  :      int     ,		__snake_case  :      Dict=1_3     ,		__snake_case  :      Dict=7     ,		__snake_case  :      str=False     ,		__snake_case  :      Dict=True     ,		__snake_case  :      Tuple=False     ,		__snake_case  :      List[Any]=False     ,		__snake_case  :      Dict=1_9     ,		__snake_case  :      Dict=3_2     ,		__snake_case  :      Union[str, Any]=5     ,		__snake_case  :      List[Any]=4     ,		__snake_case  :      List[Any]=3_7     ,		__snake_case  :      Union[str, Any]="gelu"     ,		__snake_case  :      List[str]=0.1     ,		__snake_case  :      Union[str, Any]=0.1     ,		__snake_case  :      List[str]=5_1_2     ,		__snake_case  :      List[str]=1_6     ,		__snake_case  :      Tuple=2     ,		__snake_case  :      Optional[Any]=0.02     ,		__snake_case  :      str=3     ,		__snake_case  :      List[str]=4     ,		__snake_case  :      Optional[int]=None     ,		)			->							Optional[int]:
											__magic_name__:	Any							=       parent
											__magic_name__:	Optional[Any]							=       batch_size
											__magic_name__:	Optional[Any]							=       seq_length
											__magic_name__:	Optional[Any]							=       is_training
											__magic_name__:	str							=       use_input_mask
											__magic_name__:	Optional[Any]							=       use_token_type_ids
											__magic_name__:	List[str]							=       use_labels
											__magic_name__:	List[str]							=       vocab_size
											__magic_name__:	Union[str, Any]							=       hidden_size
											__magic_name__:	List[Any]							=       num_hidden_layers
											__magic_name__:	Dict							=       num_attention_heads
											__magic_name__:	int							=       intermediate_size
											__magic_name__:	List[Any]							=       hidden_act
											__magic_name__:	int							=       hidden_dropout_prob
											__magic_name__:	List[str]							=       attention_probs_dropout_prob
											__magic_name__:	Tuple							=       max_position_embeddings
											__magic_name__:	Any							=       type_vocab_size
											__magic_name__:	Union[str, Any]							=       type_sequence_label_size
											__magic_name__:	List[str]							=       initializer_range
											__magic_name__:	Optional[Any]							=       num_labels
											__magic_name__:	List[str]							=       num_choices
											__magic_name__:	Tuple							=       scope

					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Any:
											__magic_name__:	Union[str, Any]							=       ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size					)

											__magic_name__:	Dict							=       None
											if self.use_input_mask:
																	__magic_name__:	Optional[int]							=       random_attention_mask([self.batch_size, self.seq_length]					)

											__magic_name__:	str							=       None
											__magic_name__:	Optional[Any]							=       None
											__magic_name__:	int							=       None
											if self.use_labels:
																	__magic_name__:	List[Any]							=       ids_tensor([self.batch_size]     ,		self.type_sequence_label_size					)
																	__magic_name__:	str							=       ids_tensor([self.batch_size, self.seq_length]     ,		self.num_labels					)
																	__magic_name__:	Dict							=       ids_tensor([self.batch_size]     ,		self.num_choices					)

											__magic_name__:	Dict							=       self.get_config()

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

					def   lowerCamelCase__   (			self  :      int					)			->							Any:
											__magic_name__:	List[Any]							=       EsmConfig(
											    vocab_size=3_3     ,		hidden_size=self.hidden_size     ,		pad_token_id=1     ,		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     ,		max_position_embeddings=self.max_position_embeddings     ,		type_vocab_size=self.type_vocab_size     ,		initializer_range=self.initializer_range     ,		is_folding_model=__snake_case     ,		esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}     ,		)
											return config

					def   lowerCamelCase__   (			self  :      Optional[int]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Optional[Any]     ,		__snake_case  :      Any     ,		__snake_case  :      List[Any]     ,		__snake_case  :      Any     ,		__snake_case  :      Union[str, Any]					)			->							Optional[Any]:
											__magic_name__:	str							=       EsmForProteinFolding(config=__snake_case					).float()
											model.to(__snake_case					)
											model.eval()
											__magic_name__:	str							=       model(__snake_case     ,		attention_mask=__snake_case					)
											__magic_name__:	Any							=       model(__snake_case					)
											__magic_name__:	Dict							=       model(__snake_case					)

											self.parent.assertEqual(result.positions.shape     ,		(8, self.batch_size, self.seq_length, 1_4, 3)					)
											self.parent.assertEqual(result.angles.shape     ,		(8, self.batch_size, self.seq_length, 7, 2)					)






					def   lowerCamelCase__   (			self  :      Any					)			->							int:
											__magic_name__:	str							=       self.prepare_config_and_inputs()
											(
											    (
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			(
											    __magic_name__
											),			
											):	Union[str, Any]							=       config_and_inputs
											__magic_name__:	Any							=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
											return config, inputs_dict



@require_torch
class       __A		(   SCREAMING_SNAKE_CASE_     ,SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):
					UpperCAmelCase__								=				False

					UpperCAmelCase__								=				(EsmForProteinFolding,) if is_torch_available() else ()
					UpperCAmelCase__								=				()
					UpperCAmelCase__								=				{} if is_torch_available() else {}
					UpperCAmelCase__								=				False

					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Optional[Any]:
											__magic_name__:	Optional[int]							=       EsmFoldModelTester(self					)
											__magic_name__:	Optional[Any]							=       ConfigTester(self     ,		config_class=__snake_case     ,		hidden_size=3_7					)

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

					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											__magic_name__:	Tuple							=       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case					)

					@unittest.skip("""Does not support attention outputs"""					)
					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											pass

					@unittest.skip
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Tuple:
											pass

					@unittest.skip("""Esm does not support embedding resizing"""					)
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Optional[int]:
											pass

					@unittest.skip("""Esm does not support embedding resizing"""					)
					def   lowerCamelCase__   (			self  :      str					)			->							str:
											pass

					@unittest.skip("""ESMFold does not support passing input embeds!"""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							List[str]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Tuple:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold does not support head pruning."""					)
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold does not output hidden states in the normal way."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[Any]:
											pass

					@unittest.skip("""ESMfold does not output hidden states in the normal way."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							int:
											pass

					@unittest.skip("""ESMFold only has one output format."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							str:
											pass

					@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""					)
					def   lowerCamelCase__   (			self  :      List[Any]					)			->							Dict:
											pass

					@unittest.skip("""ESMFold does not support input chunking."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Any:
											pass

					@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""					)
					def   lowerCamelCase__   (			self  :      int					)			->							Union[str, Any]:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Any					)			->							Optional[Any]:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Optional[int]					)			->							int:
											pass

					@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							List[str]:
											pass

					@unittest.skip("""ESMFold doesn't support data parallel."""					)
					def   lowerCamelCase__   (			self  :      List[str]					)			->							str:
											pass






					@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
					def   lowerCamelCase__   (			self  :      Union[str, Any]					)			->							Dict:
											pass



@require_torch
class       __A		(   SCREAMING_SNAKE_CASE_			):

					@slow
					def   lowerCamelCase__   (			self  :      Optional[Any]					)			->							List[Any]:
											__magic_name__:	Any							=       EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""					).float()
											model.eval()
											__magic_name__:	Optional[int]							=       torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]					)
											__magic_name__:	Dict							=       model(__snake_case					)["""positions"""]
											__magic_name__:	Tuple							=       torch.tensor([2.5828, 0.7993, -10.9334]     ,		dtype=torch.floataa					)
											self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]     ,		__snake_case     ,		atol=1E-4					)					)



