



"""simple docstring"""
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 				_UpperCAmelCase				:





     '''simple docstring'''



     def __init__(   self			,      A			,      )				->   Any:
         _UpperCAmelCase				:    Optional[int]											=       parent
         _UpperCAmelCase				:    str											=       1_3
         _UpperCAmelCase				:    Dict											=       7
         _UpperCAmelCase				:    Dict											=       True
         _UpperCAmelCase				:    Dict											=       True
         _UpperCAmelCase				:    Optional[Any]											=       True
         _UpperCAmelCase				:    Optional[Any]											=       9_9
         _UpperCAmelCase				:    Optional[int]											=       3_2
         _UpperCAmelCase				:    Optional[int]											=       2
         _UpperCAmelCase				:    Optional[Any]											=       4
         _UpperCAmelCase				:    int											=       3_7
         _UpperCAmelCase				:    int											=       "gelu"
         _UpperCAmelCase				:    Any											=       0.1
         _UpperCAmelCase				:    Any											=       0.1
         _UpperCAmelCase				:    List[Any]											=       5_1_2
         _UpperCAmelCase				:    List[str]											=       1_6
         _UpperCAmelCase				:    str											=       2
         _UpperCAmelCase				:    List[str]											=       0.02
         _UpperCAmelCase				:    int											=       3
         _UpperCAmelCase				:    List[str]											=       4
         _UpperCAmelCase				:    Optional[Any]											=       None





     def 					__lowerCAmelCase      (   self						)				->   Optional[Any]:
         _UpperCAmelCase				:    Dict											=       ids_tensor([self.batch_size, self.seq_length]			,      self.vocab_size						)

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

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

         _UpperCAmelCase				:    Any											=       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 					__lowerCAmelCase      (   self						)				->   List[Any]:
         (
             _UpperCAmelCase      
         )				:    List[str]											=       self.prepare_config_and_inputs()

         _UpperCAmelCase				:    List[str]											=       True
         _UpperCAmelCase				:    List[Any]											=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
         _UpperCAmelCase				:    List[str]											=       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 					__lowerCAmelCase      (   self			,      A			,      A			,      A			,      A			,      A			,      A						)				->   Tuple:
         _UpperCAmelCase				:    Tuple											=       TFEsmModel(config=lowerCAmelCase_						)
         _UpperCAmelCase				:    Any											=       {"input_ids": input_ids, "attention_mask": input_mask}
         _UpperCAmelCase				:    Any											=       model(lowerCAmelCase_						)

         _UpperCAmelCase				:    Union[str, Any]											=       [input_ids, input_mask]
         _UpperCAmelCase				:    Dict											=       model(lowerCAmelCase_						)

         _UpperCAmelCase				:    List[Any]											=       model(lowerCAmelCase_						)

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





     def 					__lowerCAmelCase      (   self			,      A			,      A			,      A			,      A			,      A			,      A			,      A			,      A			,      )				->   Union[str, Any]:
         _UpperCAmelCase				:    Optional[Any]											=       True

         _UpperCAmelCase				:    Optional[int]											=       TFEsmModel(config=lowerCAmelCase_						)
         _UpperCAmelCase				:    Union[str, Any]											=       {
             "input_ids": input_ids,
             "attention_mask": input_mask,
             "encoder_hidden_states": encoder_hidden_states,
             "encoder_attention_mask": encoder_attention_mask,
         }
         _UpperCAmelCase				:    Optional[Any]											=       model(lowerCAmelCase_						)

         _UpperCAmelCase				:    Tuple											=       [input_ids, input_mask]
         _UpperCAmelCase				:    Optional[int]											=       model(lowerCAmelCase_			,      encoder_hidden_states=lowerCAmelCase_						)

         # Also check the case where encoder outputs are not passed
         _UpperCAmelCase				:    List[Any]											=       model(lowerCAmelCase_			,      attention_mask=lowerCAmelCase_						)

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





     def 					__lowerCAmelCase      (   self			,      A			,      A			,      A			,      A			,      A			,      A						)				->   Dict:
         _UpperCAmelCase				:    Dict											=       TFEsmForMaskedLM(config=lowerCAmelCase_						)
         _UpperCAmelCase				:    Optional[int]											=       model([input_ids, input_mask]						)
         self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.vocab_size)						)





     def 					__lowerCAmelCase      (   self			,      A			,      A			,      A			,      A			,      A			,      A						)				->   Union[str, Any]:
         _UpperCAmelCase				:    Any											=       self.num_labels
         _UpperCAmelCase				:    List[Any]											=       TFEsmForTokenClassification(config=lowerCAmelCase_						)
         _UpperCAmelCase				:    str											=       {"input_ids": input_ids, "attention_mask": input_mask}
         _UpperCAmelCase				:    int											=       model(lowerCAmelCase_						)
         self.parent.assertEqual(result.logits.shape			,      (self.batch_size, self.seq_length, self.num_labels)						)





     def 					__lowerCAmelCase      (   self						)				->   Dict:
         _UpperCAmelCase				:    int											=       self.prepare_config_and_inputs()
         (
             _UpperCAmelCase      
         )				:    Optional[int]											=       config_and_inputs
         _UpperCAmelCase				:    List[str]											=       {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict



@require_tf
class 				_UpperCAmelCase				(					__A						,__A						,unittest.TestCase ):





     '''simple docstring'''



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





     def 					__lowerCAmelCase      (   self						)				->   Optional[int]:
         _UpperCAmelCase				:    Any											=       TFEsmModelTester(self						)
         _UpperCAmelCase				:    Optional[Any]											=       ConfigTester(self			,      config_class=lowerCAmelCase_			,      hidden_size=3_7						)





     def 					__lowerCAmelCase      (   self						)				->   List[Any]:
         self.config_tester.run_common_tests()





     def 					__lowerCAmelCase      (   self						)				->   Union[str, Any]:
         _UpperCAmelCase				:    Tuple											=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*lowerCAmelCase_						)





     def 					__lowerCAmelCase      (   self						)				->   List[str]:
         _UpperCAmelCase				:    List[str]											=       self.model_tester.prepare_config_and_inputs_for_decoder()
         self.model_tester.create_and_check_model_as_decoder(*lowerCAmelCase_						)





     def 					__lowerCAmelCase      (   self						)				->   Tuple:
         _UpperCAmelCase				:    int											=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*lowerCAmelCase_						)





     def 					__lowerCAmelCase      (   self						)				->   List[Any]:
         _UpperCAmelCase				:    Union[str, Any]											=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*lowerCAmelCase_						)





     @slow
     def 					__lowerCAmelCase      (   self						)				->   Optional[Any]:
         for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _UpperCAmelCase				:    Optional[int]											=       TFEsmModel.from_pretrained(lowerCAmelCase_						)
             self.assertIsNotNone(lowerCAmelCase_						)





     @unittest.skip('''Protein models do not support embedding resizing.'''						)
     def 					__lowerCAmelCase      (   self						)				->   List[Any]:
         pass





     @unittest.skip('''Protein models do not support embedding resizing.'''						)
     def 					__lowerCAmelCase      (   self						)				->   Optional[Any]:
         pass





     def 					__lowerCAmelCase      (   self						)				->   int:
         _UpperCAmelCase				:    Union[str, Any]											=       self.model_tester.prepare_config_and_inputs_for_common()

         for model_class in self.all_model_classes:
             _UpperCAmelCase				:    Any											=       model_class(lowerCAmelCase_						)
             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
                 _UpperCAmelCase				:    Tuple											=       model.get_bias()
                 assert isinstance(lowerCAmelCase_			,      lowerCAmelCase_						)
                 for k, v in name.items():
                     assert isinstance(lowerCAmelCase_			,      tf.Variable						)
             else:
                 _UpperCAmelCase				:    Optional[Any]											=       model.get_output_embeddings()
                 assert x is None
                 _UpperCAmelCase				:    List[Any]											=       model.get_bias()
                 assert name is None



@require_tf
class 				_UpperCAmelCase				(					unittest.TestCase ):





     '''simple docstring'''



     @slow
     def 					__lowerCAmelCase      (   self						)				->   List[Any]:
         _UpperCAmelCase				:    Any											=       TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''						)

         _UpperCAmelCase				:    int											=       tf.constant([[0, 1, 2, 3, 4, 5]]						)
         _UpperCAmelCase				:    int											=       model(lowerCAmelCase_						)[0]
         _UpperCAmelCase				:    Optional[Any]											=       [1, 6, 3_3]
         self.assertEqual(list(output.numpy().shape						)			,      lowerCAmelCase_						)
         # compare the actual values for a slice.
         _UpperCAmelCase				:    int											=       tf.constant(
             [
                 [
                     [8.921_518, -1_0.5_8_9_8_1_4, -6.4_671_307],
                     [-6.3_967_156, -1_3.9_1_1_3_7_7, -1.1_211_915],
                     [-7.781_247, -1_3.9_5_1_5_5_7, -3.740_592],
                 ]
             ]						)
         self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,      expected_slice.numpy()			,      atol=1E-2						)						)





     @slow
     def 					__lowerCAmelCase      (   self						)				->   int:
         _UpperCAmelCase				:    Tuple											=       TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''						)

         _UpperCAmelCase				:    List[Any]											=       tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]						)
         _UpperCAmelCase				:    Optional[int]											=       model(lowerCAmelCase_						)[0]
         # compare the actual values for a slice.
         _UpperCAmelCase				:    Optional[Any]											=       tf.constant(
             [
                 [
                     [0.14_443_092, 0.54_125_327, 0.3_247_739],
                     [0.30_340_484, 0.00_526_676, 0.31_077_722],
                     [0.32_278_043, -0.24_987_096, 0.3_414_628],
                 ]
             ]						)
         self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,      expected_slice.numpy()			,      atol=1E-4						)						)










"""simple docstring"""
import argparse
import logging
import sys
from unittest.mock import patch

import run_glue_deebert

from transformers.testing_utils import TestCasePlus, get_gpu_count, require_torch_non_multi_gpu, slow


logging.basicConfig(level=logging.DEBUG)

lowerCamelCase_							      = logging.getLogger()

def 							snake_case (  ):
			UpperCAmelCase_	:					str     	=   argparse.ArgumentParser()
			parser.add_argument("-f"		)
			UpperCAmelCase_	:					Optional[Any]     	=   parser.parse_args()
			return args.f







class      UpperCamelCase_			(__A	):





	def 	_SCREAMING_SNAKE_CASE ( self   :					str		)		->     None:
				UpperCAmelCase_	:					List[str]     	=   logging.StreamHandler(sys.stdout		)
				logger.addHandler(lowerCAmelCase_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					List[str]		)		->     Tuple:
				UpperCAmelCase_	:					Optional[int]     	=   get_gpu_count()

				if n_gpu > 1:
							pass
							# XXX: doesn't quite work with n_gpu > 1 https://github.com/huggingface/transformers/issues/10560
							# script = f"{self.examples_dir_str}/research_projects/deebert/run_glue_deebert.py"
							# distributed_args = f"-m torch.distributed.launch --nproc_per_node={n_gpu} {script}".split()
							# cmd = [sys.executable] + distributed_args + args
							# execute_subprocess_async(cmd, env=self.get_env())
							# XXX: test the results - need to save them first into .json file
				else:
							args.insert(0      ,      "run_glue_deebert.py"		)
							with patch.object(lowerCAmelCase_      ,      "argv"      ,      lowerCAmelCase_		):
										UpperCAmelCase_	:					Dict     	=   run_glue_deebert.main()
										for value in result.values():
													self.assertGreaterEqual(lowerCAmelCase_      ,      0.6_6_6		)





	@slow
	@require_torch_non_multi_gpu
	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     int:
				UpperCAmelCase_	:					Tuple     	=   "\n            --model_type roberta\n            --model_name_or_path roberta-base\n            --task_name MRPC\n            --do_train\n            --do_eval\n            --do_lower_case\n            --data_dir ./tests/fixtures/tests_samples/MRPC/\n            --max_seq_length 128\n            --per_gpu_eval_batch_size=1\n            --per_gpu_train_batch_size=8\n            --learning_rate 2e-4\n            --num_train_epochs 3\n            --overwrite_output_dir\n            --seed 42\n            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage\n            --plot_data_dir ./examples/deebert/results/\n            --save_steps 0\n            --overwrite_cache\n            --eval_after_first_stage\n            ".split()
				self.run_and_check(lowerCAmelCase_		)

				UpperCAmelCase_	:					Union[str, Any]     	=   "\n            --model_type roberta\n            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage\n            --task_name MRPC\n            --do_eval\n            --do_lower_case\n            --data_dir ./tests/fixtures/tests_samples/MRPC/\n            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage\n            --plot_data_dir ./examples/deebert/results/\n            --max_seq_length 128\n            --eval_each_highway\n            --eval_highway\n            --overwrite_cache\n            --per_gpu_eval_batch_size=1\n            ".split()
				self.run_and_check(lowerCAmelCase_		)

				UpperCAmelCase_	:					str     	=   "\n            --model_type roberta\n            --model_name_or_path ./examples/deebert/saved_models/roberta-base/MRPC/two_stage\n            --task_name MRPC\n            --do_eval\n            --do_lower_case\n            --data_dir ./tests/fixtures/tests_samples/MRPC/\n            --output_dir ./examples/deebert/saved_models/roberta-base/MRPC/two_stage\n            --plot_data_dir ./examples/deebert/results/\n            --max_seq_length 128\n            --early_exit_entropy 0.1\n            --eval_highway\n            --overwrite_cache\n            --per_gpu_eval_batch_size=1\n            ".split()
				self.run_and_check(lowerCAmelCase_		)



