





from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
     import numpy
     import tensorflow as tf

     from transformers.models.esm.modeling_tf_esm import (
         TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
         TFEsmModel,
     )



class   _SCREAMING_SNAKE_CASE		:







 '''simple docstring'''

 def __init__(self							:  Union[str, Any]			,					UpperCAmelCase_							:  Union[str, Any]			,					)			->int:



   '''simple docstring'''



   lowerCamelCase__: Optional[Any]         =parent
   lowerCamelCase__: Optional[int]         =13
   lowerCamelCase__: Any         =7
   lowerCamelCase__: Tuple         =True
   lowerCamelCase__: str         =True
   lowerCamelCase__: Optional[Any]         =True
   lowerCamelCase__: Any         =99
   lowerCamelCase__: List[Any]         =32
   lowerCamelCase__: Union[str, Any]         =2
   lowerCamelCase__: str         =4
   lowerCamelCase__: Optional[int]         =37
   lowerCamelCase__: List[str]         ="gelu"
   lowerCamelCase__: Tuple         =0.1
   lowerCamelCase__: Optional[Any]         =0.1
   lowerCamelCase__: List[Any]         =512
   lowerCamelCase__: Tuple         =16
   lowerCamelCase__: List[Any]         =2
   lowerCamelCase__: Optional[int]         =0.02
   lowerCamelCase__: Any         =3
   lowerCamelCase__: Any         =4
   lowerCamelCase__: List[Any]         =None

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



   '''simple docstring'''



   lowerCamelCase__: List[str]         =ids_tensor([self.batch_size, self.seq_length]			,					self.vocab_size)

   lowerCamelCase__: int         =None
   if self.use_input_mask:
     lowerCamelCase__: Any         =random_attention_mask([self.batch_size, self.seq_length])

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

   lowerCamelCase__: str         =EsmConfig(
       vocab_size=self.vocab_size			,					hidden_size=self.hidden_size			,					num_hidden_layers=self.num_hidden_layers			,					pad_token_id=1			,					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			,					)

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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Any:



   '''simple docstring'''



   (
       (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,      
   ): int         =self.prepare_config_and_inputs()

   lowerCamelCase__: List[Any]         =True
   lowerCamelCase__: Tuple         =floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
   lowerCamelCase__: Any         =ids_tensor([self.batch_size, self.seq_length]			,					vocab_size=2)

   return (
       config,
       input_ids,
       input_mask,
       sequence_labels,
       token_labels,
       choice_labels,
       encoder_hidden_states,
       encoder_attention_mask,
   )

 def 				SCREAMING_SNAKE_CASE_   (self							:  str			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  Dict			,					UpperCAmelCase_							:  Optional[int]			,					UpperCAmelCase_							:  List[Any]			,					UpperCAmelCase_							:  Any			,					UpperCAmelCase_							:  Optional[Any])			->str:



   '''simple docstring'''



   lowerCamelCase__: Any         =TFEsmModel(config=__UpperCamelCase)
   lowerCamelCase__: List[Any]         ={"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase__: Tuple         =model(__UpperCamelCase)

   lowerCamelCase__: Optional[int]         =[input_ids, input_mask]
   lowerCamelCase__: Optional[Any]         =model(__UpperCamelCase)

   lowerCamelCase__: Optional[int]         =model(__UpperCamelCase)

   self.parent.assertEqual(result.last_hidden_state.shape			,					(self.batch_size, self.seq_length, self.hidden_size))

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



   '''simple docstring'''



   lowerCamelCase__: int         =True

   lowerCamelCase__: Union[str, Any]         =TFEsmModel(config=__UpperCamelCase)
   lowerCamelCase__: str         ={
       "input_ids": input_ids,
       "attention_mask": input_mask,
       "encoder_hidden_states": encoder_hidden_states,
       "encoder_attention_mask": encoder_attention_mask,
   }
   lowerCamelCase__: List[Any]         =model(__UpperCamelCase)

   lowerCamelCase__: int         =[input_ids, input_mask]
   lowerCamelCase__: Union[str, Any]         =model(__UpperCamelCase			,					encoder_hidden_states=__UpperCamelCase)

   # Also check the case where encoder outputs are not passed
   lowerCamelCase__: Dict         =model(__UpperCamelCase			,					attention_mask=__UpperCamelCase)

   self.parent.assertEqual(result.last_hidden_state.shape			,					(self.batch_size, self.seq_length, self.hidden_size))

 def 				SCREAMING_SNAKE_CASE_   (self							:  str			,					UpperCAmelCase_							:  int			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  str			,					UpperCAmelCase_							:  Union[str, Any])			->Union[str, Any]:



   '''simple docstring'''



   lowerCamelCase__: Optional[int]         =TFEsmForMaskedLM(config=__UpperCamelCase)
   lowerCamelCase__: str         =model([input_ids, input_mask])
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.vocab_size))

 def 				SCREAMING_SNAKE_CASE_   (self							:  int			,					UpperCAmelCase_							:  Dict			,					UpperCAmelCase_							:  List[str]			,					UpperCAmelCase_							:  List[str]			,					UpperCAmelCase_							:  Union[str, Any]			,					UpperCAmelCase_							:  Optional[Any]			,					UpperCAmelCase_							:  Any)			->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__: List[str]         =self.num_labels
   lowerCamelCase__: Tuple         =TFEsmForTokenClassification(config=__UpperCamelCase)
   lowerCamelCase__: Dict         ={"input_ids": input_ids, "attention_mask": input_mask}
   lowerCamelCase__: Union[str, Any]         =model(__UpperCamelCase)
   self.parent.assertEqual(result.logits.shape			,					(self.batch_size, self.seq_length, self.num_labels))







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



   '''simple docstring'''



   lowerCamelCase__: Any         =self.prepare_config_and_inputs()
   (
       (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,   (
       lowerCamelCase__   
   )    ,      
   ): List[Any]         =config_and_inputs
   lowerCamelCase__: Optional[int]         ={"input_ids": input_ids, "attention_mask": input_mask}
   return config, inputs_dict



@require_tf
class   _SCREAMING_SNAKE_CASE		(      _lowerCamelCase       ,				_lowerCamelCase       ,				unittest.TestCase       ):







 '''simple docstring'''
 lowercase_						=     (
     (
         TFEsmModel,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 lowercase_						=     (
     {
         'feature-extraction': TFEsmModel,
         'fill-mask': TFEsmForMaskedLM,
         'text-classification': TFEsmForSequenceClassification,
         'token-classification': TFEsmForTokenClassification,
         'zero-shot': TFEsmForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 lowercase_						=     False
 lowercase_						=     False

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



   '''simple docstring'''



   lowerCamelCase__: int         =TFEsmModelTester(self)
   lowerCamelCase__: Any         =ConfigTester(self			,					config_class=__UpperCamelCase			,					hidden_size=37)

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



   '''simple docstring'''



   self.config_tester.run_common_tests()

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



   '''simple docstring'''



   lowerCamelCase__: List[Any]         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__UpperCamelCase)

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[str])			->List[str]:



   '''simple docstring'''



   lowerCamelCase__: Dict         =self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(*__UpperCamelCase)

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->int:



   '''simple docstring'''



   lowerCamelCase__: int         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase)

 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict)			->str:



   '''simple docstring'''



   lowerCamelCase__: List[Any]         =self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase)

 @slow
 def 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->str:



   '''simple docstring'''



   for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     lowerCamelCase__: List[Any]         =TFEsmModel.from_pretrained(__UpperCamelCase)
     self.assertIsNotNone(__UpperCamelCase)

 @unittest.skip("Protein models do not support embedding resizing.")
 def 				SCREAMING_SNAKE_CASE_   (self							:  int)			->Union[str, Any]:



   '''simple docstring'''



   pass

 @unittest.skip("Protein models do not support embedding resizing.")
 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->Optional[int]:



   '''simple docstring'''



   pass







 def 				SCREAMING_SNAKE_CASE_   (self							:  Dict)			->Tuple:



   '''simple docstring'''



   lowerCamelCase__    ,   lowerCamelCase__: List[Any]         =self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     lowerCamelCase__: Union[str, Any]         =model_class(__UpperCamelCase)
     assert isinstance(model.get_input_embeddings()			,					tf.keras.layers.Layer)
     if model_class is TFEsmForMaskedLM:
       # Output embedding test differs from the main test because they're a matrix, not a layer
       lowerCamelCase__: List[str]         =model.get_bias()
       assert isinstance(__UpperCamelCase			,					__UpperCamelCase)
       for k, v in name.items():
         assert isinstance(__UpperCamelCase			,					tf.Variable)
     else:
       lowerCamelCase__: Optional[int]         =model.get_output_embeddings()
       assert x is None
       lowerCamelCase__: List[str]         =model.get_bias()
       assert name is None



@require_tf
class   _SCREAMING_SNAKE_CASE		(      unittest.TestCase       ):







 '''simple docstring'''

 @slow
 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__: str         =TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")

   lowerCamelCase__: Optional[int]         =tf.constant([[0, 1, 2, 3, 4, 5]])
   lowerCamelCase__: List[str]         =model(__UpperCamelCase)[0]
   lowerCamelCase__: Optional[int]         =[1, 6, 33]
   self.assertEqual(list(output.numpy().shape)			,					__UpperCamelCase)
   # compare the actual values for a slice.
   lowerCamelCase__: Union[str, Any]         =tf.constant(
       [
           [
               [8.92_1518, -10.589_814, -6.467_1307],
               [-6.396_7156, -13.911_377, -1.121_1915],
               [-7.78_1247, -13.951_557, -3.74_0592],
           ]
       ])
   self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,					expected_slice.numpy()			,					atol=1E-2))







 @slow
 def 				SCREAMING_SNAKE_CASE_   (self							:  Union[str, Any])			->List[Any]:



   '''simple docstring'''



   lowerCamelCase__: Dict         =TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D")

   lowerCamelCase__: int         =tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
   lowerCamelCase__: Tuple         =model(__UpperCamelCase)[0]
   # compare the actual values for a slice.
   lowerCamelCase__: int         =tf.constant(
       [
           [
               [0.1444_3092, 0.5412_5327, 0.324_7739],
               [0.3034_0484, 0.0052_6676, 0.3107_7722],
               [0.3227_8043, -0.2498_7096, 0.341_4628],
           ]
       ])
   self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,					expected_slice.numpy()			,					atol=1E-4))






import argparse
from collections import defaultdict






def  lowerCamelCase_	(      lowerCAmelCase__	:	List[str]			,	lowerCAmelCase__	:	List[str]			,	lowerCAmelCase__	:	Tuple			,	lowerCAmelCase__	:	Optional[int]			,	lowerCAmelCase__	:	List[str]				)       ->							Union[str, Any]:

							'''simple docstring'''






							A      						=    F'''{file}_{class_name}_{test_name}'''
							done_test[_id] += 1

							with open(lowerCAmelCase__			,	'r'				) as f:
														A      						=    f.readlines()

							A      						=    F'''class {class_name}('''
							A      						=    F'''{4 * " "}def {test_name}('''
							A      						=    F'''{8 * " "}{correct_line.split()[0]}'''
							A      						=    F'''{16 * " "}{correct_line.split()[0]}'''
							A      						=    False
							A      						=    False
							A      						=    False
							A      						=    False
							A      						=    0
							A      						=    0

							A      						=    []
							for line in lines:
														if line.startswith(lowerCAmelCase__				):
																					A      						=    True
														elif in_class and line.startswith(lowerCAmelCase__				):
																					A      						=    True
														elif in_class and in_func and (line.startswith(lowerCAmelCase__				) or line.startswith(lowerCAmelCase__				)):
																					A      						=    len(line.split(correct_line.split()[0]				)[0]				)
																					count += 1

																					if count == done_test[_id]:
																												A      						=    True

														if in_class and in_func and in_line:
																					if ")" not in line:
																												continue
																					else:
																												A      						=    True

														if in_class and in_func and in_line and insert_line:
																					new_lines.append(F'''{spaces * " "}{correct_line}'''				)
																					A      						=    A      						=    A      						=    A      						=    False
														else:
																					new_lines.append(lowerCAmelCase__				)

							with open(lowerCAmelCase__			,	'w'				) as f:
														for line in new_lines:
																					f.write(lowerCAmelCase__				)






def  lowerCamelCase_	(      lowerCAmelCase__	:	str			,	lowerCAmelCase__	:	Tuple=None				)       ->							Union[str, Any]:

							'''simple docstring'''






							if fail is not None:
														with open(lowerCAmelCase__			,	'r'				) as f:
																					A      						=    {l.strip() for l in f.readlines()}
							else:
														A      						=    None

							with open(lowerCAmelCase__			,	'r'				) as f:
														A      						=    f.readlines()

							A      						=    defaultdict(lowerCAmelCase__				)
							for line in correct_lines:
														A						,   A						,   A						,   A      						=    line.split(';'				)
														if test_failures is None or "::".join([file, class_name, test_name]				) in test_failures:
																					overwrite_file(lowerCAmelCase__			,	lowerCAmelCase__			,	lowerCAmelCase__			,	lowerCAmelCase__			,	lowerCAmelCase__				)


if __name__ == "__main__":
				__snake_case     :List[Any]			       =argparse.ArgumentParser()
				parser.add_argument('--correct_filename', help='filename of tests with expected result')
				parser.add_argument('--fail_filename', help='filename of test failures', type=str, default=None)
				__snake_case     :Optional[int]			       =parser.parse_args()

				main(args.correct_filename, args.fail_filename)