

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


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


 def _UpperCamelCase	(						self		: int    )    ->      Any:
   UpperCAmelCase					      =  ids_tensor([self.batch_size, self.seq_length]			,						self.vocab_size    )

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

   UpperCAmelCase					      =  None
   UpperCAmelCase					      =  None
   UpperCAmelCase					      =  None
   if self.use_labels:
     UpperCAmelCase					      =  ids_tensor([self.batch_size]			,						self.type_sequence_label_size    )
     UpperCAmelCase					      =  ids_tensor([self.batch_size, self.seq_length]			,						self.num_labels    )
     UpperCAmelCase					      =  ids_tensor([self.batch_size]			,						self.num_choices    )

   UpperCAmelCase					      =  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 _UpperCamelCase	(						self		: int    )    ->      List[Any]:
   (
       UpperCAmelCase 
   )					      =  self.prepare_config_and_inputs()

   UpperCAmelCase					      =  True
   UpperCAmelCase					      =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]    )
   UpperCAmelCase					      =  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 _UpperCamelCase	(						self		: List[Any]			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: List[Any]			,						lowerCAmelCase__		: Optional[int]    )    ->      Optional[Any]:
   UpperCAmelCase					      =  TFEsmModel(config=__lowerCamelCase    )
   UpperCAmelCase					      =  {"input_ids": input_ids, "attention_mask": input_mask}
   UpperCAmelCase					      =  model(__lowerCamelCase    )

   UpperCAmelCase					      =  [input_ids, input_mask]
   UpperCAmelCase					      =  model(__lowerCamelCase    )

   UpperCAmelCase					      =  model(__lowerCamelCase    )

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


 def _UpperCamelCase	(						self		: List[str]			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: int			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: List[str]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: List[str]			,						)    ->      List[Any]:
   UpperCAmelCase					      =  True

   UpperCAmelCase					      =  TFEsmModel(config=__lowerCamelCase    )
   UpperCAmelCase					      =  {
       "input_ids": input_ids,
       "attention_mask": input_mask,
       "encoder_hidden_states": encoder_hidden_states,
       "encoder_attention_mask": encoder_attention_mask,
   }
   UpperCAmelCase					      =  model(__lowerCamelCase    )

   UpperCAmelCase					      =  [input_ids, input_mask]
   UpperCAmelCase					      =  model(__lowerCamelCase			,						encoder_hidden_states=__lowerCamelCase    )

   # Also check the case where encoder outputs are not passed
   UpperCAmelCase					      =  model(__lowerCamelCase			,						attention_mask=__lowerCamelCase    )

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


 def _UpperCamelCase	(						self		: str			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Optional[Any]			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Union[str, Any]			,						lowerCAmelCase__		: Union[str, Any]    )    ->      Dict:
   UpperCAmelCase					      =  TFEsmForMaskedLM(config=__lowerCamelCase    )
   UpperCAmelCase					      =  model([input_ids, input_mask]    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.vocab_size)    )


 def _UpperCamelCase	(						self		: List[Any]			,						lowerCAmelCase__		: Dict			,						lowerCAmelCase__		: str			,						lowerCAmelCase__		: Tuple			,						lowerCAmelCase__		: Optional[int]			,						lowerCAmelCase__		: int			,						lowerCAmelCase__		: int    )    ->      int:
   UpperCAmelCase					      =  self.num_labels
   UpperCAmelCase					      =  TFEsmForTokenClassification(config=__lowerCamelCase    )
   UpperCAmelCase					      =  {"input_ids": input_ids, "attention_mask": input_mask}
   UpperCAmelCase					      =  model(__lowerCamelCase    )
   self.parent.assertEqual(result.logits.shape			,						(self.batch_size, self.seq_length, self.num_labels)    )


 def _UpperCamelCase	(						self		: Dict    )    ->      int:
   UpperCAmelCase					      =  self.prepare_config_and_inputs()
   (
       UpperCAmelCase 
   )					      =  config_and_inputs
   UpperCAmelCase					      =  {"input_ids": input_ids, "attention_mask": input_mask}
   return config, inputs_dict






@require_tf
class 						__magic_name__    (       _A	,						_A	,						unittest.TestCase    ):
 UpperCAmelCase							    =						(
     (
         TFEsmModel,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 UpperCAmelCase							    =						(
     {
         """feature-extraction""": TFEsmModel,
         """fill-mask""": TFEsmForMaskedLM,
         """text-classification""": TFEsmForSequenceClassification,
         """token-classification""": TFEsmForTokenClassification,
         """zero-shot""": TFEsmForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 UpperCAmelCase							    =						False
 UpperCAmelCase							    =						False


 def _UpperCamelCase	(						self		: int    )    ->      Optional[Any]:
   UpperCAmelCase					      =  TFEsmModelTester(self    )
   UpperCAmelCase					      =  ConfigTester(self			,						config_class=__lowerCamelCase			,						hidden_size=3_7    )


 def _UpperCamelCase	(						self		: List[str]    )    ->      Dict:
   self.config_tester.run_common_tests()


 def _UpperCamelCase	(						self		: int    )    ->      Union[str, Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*__lowerCamelCase    )


 def _UpperCamelCase	(						self		: Optional[int]    )    ->      Optional[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs_for_decoder()
   self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase    )


 def _UpperCamelCase	(						self		: int    )    ->      str:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase    )


 def _UpperCamelCase	(						self		: Optional[Any]    )    ->      List[Any]:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase    )


 @slow
 def _UpperCamelCase	(						self		: str    )    ->      int:
   for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     UpperCAmelCase					      =  TFEsmModel.from_pretrained(__lowerCamelCase    )
     self.assertIsNotNone(__lowerCamelCase    )


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


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


 def _UpperCamelCase	(						self		: List[Any]    )    ->      str:
   UpperCAmelCase					      =  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     UpperCAmelCase					      =  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					      =  model.get_bias()
       assert isinstance(__lowerCamelCase			,						__lowerCamelCase    )
       for k, v in name.items():
         assert isinstance(__lowerCamelCase			,						tf.Variable    )
     else:
       UpperCAmelCase					      =  model.get_output_embeddings()
       assert x is None
       UpperCAmelCase					      =  model.get_bias()
       assert name is None






@require_tf
class 						__magic_name__    (       unittest.TestCase    ):


 @slow
 def _UpperCamelCase	(						self		: Union[str, Any]    )    ->      Optional[int]:
   UpperCAmelCase					      =  TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"    )

   UpperCAmelCase					      =  tf.constant([[0, 1, 2, 3, 4, 5]]    )
   UpperCAmelCase					      =  model(__lowerCamelCase    )[0]
   UpperCAmelCase					      =  [1, 6, 3_3]
   self.assertEqual(list(output.numpy().shape    )			,						__lowerCamelCase    )
   # compare the actual values for a slice.
   UpperCAmelCase					      =  tf.constant(
       [
           [
               [8.921_518, -10.589_814, -6.4_671_307],
               [-6.3_967_156, -13.911_377, -1.1_211_915],
               [-7.781_247, -13.951_557, -3.740_592],
           ]
       ]    )
   self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,						expected_slice.numpy()			,						atol=1e-2    )    )


 @slow
 def _UpperCamelCase	(						self		: Optional[int]    )    ->      int:
   UpperCAmelCase					      =  TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"    )

   UpperCAmelCase					      =  tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]    )
   UpperCAmelCase					      =  model(__lowerCamelCase    )[0]
   # compare the actual values for a slice.
   UpperCAmelCase					      =  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    )    )






import sys
import tempfile
import unittest
import unittest.mock as mock
from pathlib import Path

from huggingface_hub import HfFolder, delete_repo
from requests.exceptions import HTTPError

from transformers import AutoFeatureExtractor, WavaVecaFeatureExtractor
from transformers.testing_utils import TOKEN, USER, get_tests_dir, is_staging_test


sys.path.append(str(Path(__file__).parent.parent / "utils"))

from test_module.custom_feature_extraction import CustomFeatureExtractor  # noqa E402


lowerCAmelCase__       = get_tests_dir("fixtures")






class 						__magic_name__    (       unittest.TestCase    ):


 def _UpperCamelCase	(						self		: Tuple    )    ->      Union[str, Any]:
   # A mock response for an HTTP head request to emulate server down
   UpperCAmelCase					      =  mock.Mock()
   UpperCAmelCase					      =  5_0_0
   UpperCAmelCase					      =  {}
   UpperCAmelCase					      =  HTTPError
   UpperCAmelCase					      =  {}

   # Download this model to make sure it's in the cache.
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
   # Under the mock environment we get a 500 error when trying to reach the model.
   with mock.patch("requests.Session.request"			,						return_value=lowerCAmelCase__    ) as mock_head:
     UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("hf-internal-testing/tiny-random-wav2vec2"    )
     # This check we did call the fake head request
     mock_head.assert_called()


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Dict:
   # This test is for deprecated behavior and can be removed in v5
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(
       "https://huggingface.co/hf-internal-testing/tiny-random-wav2vec2/resolve/main/preprocessor_config.json"    )






@is_staging_test
class 						__magic_name__    (       unittest.TestCase    ):


 @classmethod
 def _UpperCamelCase	(						cls		: List[str]    )    ->      List[Any]:
   UpperCAmelCase					      =  TOKEN
   HfFolder.save_token(lowerCAmelCase__    )


 @classmethod
 def _UpperCamelCase	(						cls		: Optional[int]    )    ->      Union[str, Any]:
   try:
     delete_repo(token=cls._token			,						repo_id="test-feature-extractor"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="valid_org/test-feature-extractor-org"    )
   except HTTPError:
     pass

   try:
     delete_repo(token=cls._token			,						repo_id="test-dynamic-feature-extractor"    )
   except HTTPError:
     pass


 def _UpperCamelCase	(						self		: Any    )    ->      Any:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="test-feature-extractor"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(f"{USER}/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: List[Any]    )    ->      Tuple:
   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained(lowerCAmelCase__    )
   feature_extractor.push_to_hub("valid_org/test-feature-extractor"			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )

   # Reset repo
   delete_repo(token=self._token			,						repo_id="valid_org/test-feature-extractor"    )

   # Push to hub via save_pretrained
   with tempfile.TemporaryDirectory() as tmp_dir:
     feature_extractor.save_pretrained(
         lowerCAmelCase__			,						repo_id="valid_org/test-feature-extractor-org"			,						push_to_hub=lowerCAmelCase__			,						use_auth_token=self._token    )

   UpperCAmelCase					      =  WavaVecaFeatureExtractor.from_pretrained("valid_org/test-feature-extractor-org"    )
   for k, v in feature_extractor.__dict__.items():
     self.assertEqual(lowerCAmelCase__			,						getattr(lowerCAmelCase__			,						lowerCAmelCase__    )    )


 def _UpperCamelCase	(						self		: Dict    )    ->      List[str]:
   CustomFeatureExtractor.register_for_auto_class()
   UpperCAmelCase					      =  CustomFeatureExtractor.from_pretrained(lowerCAmelCase__    )

   feature_extractor.push_to_hub("test-dynamic-feature-extractor"			,						use_auth_token=self._token    )

   # This has added the proper auto_map field to the config
   self.assertDictEqual(
       feature_extractor.auto_map			,						{"AutoFeatureExtractor": "custom_feature_extraction.CustomFeatureExtractor"}			,						)

   UpperCAmelCase					      =  AutoFeatureExtractor.from_pretrained(
       f"{USER}/test-dynamic-feature-extractor"			,						trust_remote_code=lowerCAmelCase__    )
   # Can't make an isinstance check because the new_feature_extractor is from the CustomFeatureExtractor class of a dynamic module
   self.assertEqual(new_feature_extractor.__class__.__name__			,						"CustomFeatureExtractor"    )



