


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







     '''simple docstring'''







     def __init__(				self    :							Optional[int]			,				_UpperCamelCase    :							Dict			,				):
      _lowercase:       Any						=				parent
      _lowercase:       Tuple						=				13
      _lowercase:       int						=				7
      _lowercase:       Any						=				True
      _lowercase:       List[str]						=				True
      _lowercase:       Optional[int]						=				True
      _lowercase:       List[Any]						=				99
      _lowercase:       Any						=				32
      _lowercase:       Optional[Any]						=				2
      _lowercase:       List[Any]						=				4
      _lowercase:       Dict						=				37
      _lowercase:       List[Any]						=				"gelu"
      _lowercase:       Any						=				0.1
      _lowercase:       Dict						=				0.1
      _lowercase:       Optional[int]						=				512
      _lowercase:       Any						=				16
      _lowercase:       Tuple						=				2
      _lowercase:       Optional[Any]						=				0.0_2
      _lowercase:       Any						=				3
      _lowercase:       Dict						=				4
      _lowercase:       str						=				None

     def 					UpperCAmelCase__       (				self    :							List[str]):
      _lowercase:       Dict						=				ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size)

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

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

      _lowercase:       Union[str, 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 					UpperCAmelCase__       (				self    :							List[Any]):
      (
          (
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	 
      ):       Optional[Any]						=				self.prepare_config_and_inputs()

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

      _lowercase:       Optional[int]						=				[input_ids, input_mask]
      _lowercase:       int						=				model(_UpperCamelCase)

      _lowercase:       Tuple						=				model(_UpperCamelCase)

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

     def 					UpperCAmelCase__       (				self    :							Any			,				_UpperCamelCase    :							Optional[int]			,				_UpperCamelCase    :							Optional[Any]			,				_UpperCamelCase    :							str			,				_UpperCamelCase    :							List[str]			,				_UpperCamelCase    :							Optional[Any]			,				_UpperCamelCase    :							int			,				_UpperCamelCase    :							Tuple			,				_UpperCamelCase    :							Any			,				):
      _lowercase:       Dict						=				True

      _lowercase:       Optional[Any]						=				TFEsmModel(config=_UpperCamelCase)
      _lowercase:       Any						=				{
          "input_ids": input_ids,
          "attention_mask": input_mask,
          "encoder_hidden_states": encoder_hidden_states,
          "encoder_attention_mask": encoder_attention_mask,
      }
      _lowercase:       Optional[int]						=				model(_UpperCamelCase)

      _lowercase:       Optional[int]						=				[input_ids, input_mask]
      _lowercase:       List[Any]						=				model(_UpperCamelCase			,				encoder_hidden_states=_UpperCamelCase)

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

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

     def 					UpperCAmelCase__       (				self    :							List[str]			,				_UpperCamelCase    :							str			,				_UpperCamelCase    :							Optional[Any]			,				_UpperCamelCase    :							Any			,				_UpperCamelCase    :							int			,				_UpperCamelCase    :							Tuple			,				_UpperCamelCase    :							int):
      _lowercase:       int						=				TFEsmForMaskedLM(config=_UpperCamelCase)
      _lowercase:       List[str]						=				model([input_ids, input_mask])
      self.parent.assertEqual(result.logits.shape			,				(self.batch_size, self.seq_length, self.vocab_size))

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






     def 					UpperCAmelCase__       (				self    :							Tuple):
      _lowercase:       Tuple						=				self.prepare_config_and_inputs()
      (
          (
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	(
          _lowercase 
      )					,	 
      ):       List[Any]						=				config_and_inputs
      _lowercase:       str						=				{"input_ids": input_ids, "attention_mask": input_mask}
      return config, inputs_dict






@require_tf
class        A     (		lowerCamelCase_	, lowerCamelCase_	, unittest.TestCase	):







     '''simple docstring'''







     lowerCamelCase :						Optional[Any]            =   (
         (
             TFEsmModel,
             TFEsmForMaskedLM,
             TFEsmForSequenceClassification,
             TFEsmForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     lowerCamelCase :						Tuple            =   (
         {
             'feature-extraction': TFEsmModel,
             'fill-mask': TFEsmForMaskedLM,
             'text-classification': TFEsmForSequenceClassification,
             'token-classification': TFEsmForTokenClassification,
             'zero-shot': TFEsmForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     lowerCamelCase :						str            =   False
     lowerCamelCase :						Tuple            =   False

     def 					UpperCAmelCase__       (				self    :							Tuple):
      _lowercase:       List[str]						=				TFEsmModelTester(self)
      _lowercase:       str						=				ConfigTester(self			,				config_class=_UpperCamelCase			,				hidden_size=37)

     def 					UpperCAmelCase__       (				self    :							int):
      self.config_tester.run_common_tests()

     def 					UpperCAmelCase__       (				self    :							Tuple):
      _lowercase:       int						=				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*_UpperCamelCase)

     def 					UpperCAmelCase__       (				self    :							Any):
      _lowercase:       List[str]						=				self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*_UpperCamelCase)

     def 					UpperCAmelCase__       (				self    :							Union[str, Any]):
      _lowercase:       Optional[Any]						=				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase)

     def 					UpperCAmelCase__       (				self    :							Tuple):
      _lowercase:       Union[str, Any]						=				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase)

     @slow
     def 					UpperCAmelCase__       (				self    :							Optional[int]):
      for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       _lowercase:       Tuple						=				TFEsmModel.from_pretrained(_UpperCamelCase)
       self.assertIsNotNone(_UpperCamelCase)

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

     @unittest.skip("Protein models do not support embedding resizing.")
     def 					UpperCAmelCase__       (				self    :							Tuple):
      pass






     def 					UpperCAmelCase__       (				self    :							Optional[Any]):
      _lowercase					,	_lowercase:       Optional[Any]						=				self.model_tester.prepare_config_and_inputs_for_common()

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






@require_tf
class        A     (		unittest.TestCase	):







     '''simple docstring'''







     @slow
     def 					UpperCAmelCase__       (				self    :							Any):
      _lowercase:       Any						=				TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")

      _lowercase:       List[str]						=				tf.constant([[0, 1, 2, 3, 4, 5]])
      _lowercase:       Optional[Any]						=				model(_UpperCamelCase)[0]
      _lowercase:       Union[str, Any]						=				[1, 6, 33]
      self.assertEqual(list(output.numpy().shape)			,				_UpperCamelCase)
      # compare the actual values for a slice.
      _lowercase:       Optional[int]						=				tf.constant(
          [
              [
                  [8.9_2_1_5_1_8, -10.589_814, -6.4_6_7_1_3_0_7],
                  [-6.3_9_6_7_1_5_6, -13.911_377, -1.1_2_1_1_9_1_5],
                  [-7.7_8_1_2_4_7, -13.951_557, -3.7_4_0_5_9_2],
              ]
          ])
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,				expected_slice.numpy()			,				atol=1e-2))






     @slow
     def 					UpperCAmelCase__       (				self    :							Optional[Any]):
      _lowercase:       List[str]						=				TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D")

      _lowercase:       Tuple						=				tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
      _lowercase:       Union[str, Any]						=				model(_UpperCamelCase)[0]
      # compare the actual values for a slice.
      _lowercase:       Any						=				tf.constant(
          [
              [
                  [0.1_4_4_4_3_0_9_2, 0.5_4_1_2_5_3_2_7, 0.3_2_4_7_7_3_9],
                  [0.3_0_3_4_0_4_8_4, 0.0_0_5_2_6_6_7_6, 0.3_1_0_7_7_7_2_2],
                  [0.3_2_2_7_8_0_4_3, -0.2_4_9_8_7_0_9_6, 0.3_4_1_4_6_2_8],
              ]
          ])
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()			,				expected_slice.numpy()			,				atol=1e-4))



"""simple docstring"""


import copy
import tempfile
import unittest

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

from transformers import AutoConfig, GenerationConfig
from transformers.testing_utils import TOKEN, USER, is_staging_test




class  __lowerCamelCase  ( unittest.TestCase			):




				@parameterized.expand([(None,), ('''foo.json''',)]   )
				def        UpperCAmelCase__						(				self							,						UpperCAmelCase   ):
							lowerCamelCase_					=  GenerationConfig(
							    do_sample=UpperCAmelCase							,						temperature=0.7							,						length_penalty=1.0							,						bad_words_ids=[[1, 2, 3], [4, 5]]							,						)
							with tempfile.TemporaryDirectory() as tmp_dir:
										config.save_pretrained(UpperCAmelCase							,						config_name=UpperCAmelCase   )
										lowerCamelCase_					=  GenerationConfig.from_pretrained(UpperCAmelCase							,						config_name=UpperCAmelCase   )

							# Checks parameters that were specified
							self.assertEqual(loaded_config.do_sample							,						UpperCAmelCase   )
							self.assertEqual(loaded_config.temperature							,						0.7   )
							self.assertEqual(loaded_config.length_penalty							,						1.0   )
							self.assertEqual(loaded_config.bad_words_ids							,						[[1, 2, 3], [4, 5]]   )

							# Checks parameters that were not specified (defaults)
							self.assertEqual(loaded_config.top_k							,						50   )
							self.assertEqual(loaded_config.max_length							,						20   )
							self.assertEqual(loaded_config.max_time							,						UpperCAmelCase   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  AutoConfig.from_pretrained('''gpt2'''   )
							lowerCamelCase_					=  GenerationConfig.from_model_config(UpperCAmelCase   )
							lowerCamelCase_					=  GenerationConfig()

							# The generation config has loaded a few non-default parameters from the model config
							self.assertNotEqual(UpperCAmelCase							,						UpperCAmelCase   )

							# One of those parameters is eos_token_id -- check if it matches
							self.assertNotEqual(generation_config_from_model.eos_token_id							,						default_generation_config.eos_token_id   )
							self.assertEqual(generation_config_from_model.eos_token_id							,						model_config.eos_token_id   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  GenerationConfig()
							lowerCamelCase_					=  {
							    '''max_new_tokens''': 1024,
							    '''foo''': '''bar''',
							}
							lowerCamelCase_					=  copy.deepcopy(UpperCAmelCase   )
							lowerCamelCase_					=  generation_config.update(**UpperCAmelCase   )

							# update_kwargs was not modified (no side effects)
							self.assertEqual(UpperCAmelCase							,						UpperCAmelCase   )

							# update_kwargs was used to update the config on valid attributes
							self.assertEqual(generation_config.max_new_tokens							,						1024   )

							# `.update()` returns a dictionary of unused kwargs
							self.assertEqual(UpperCAmelCase							,						{'''foo''': '''bar'''}   )




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  GenerationConfig()
							lowerCamelCase_					=  '''bar'''

							with tempfile.TemporaryDirectory('''test-generation-config'''   ) as tmp_dir:
										generation_config.save_pretrained(UpperCAmelCase   )

										lowerCamelCase_					=  GenerationConfig.from_pretrained(UpperCAmelCase   )
							# update_kwargs was used to update the config on valid attributes
							self.assertEqual(new_config.foo							,						'''bar'''   )

							lowerCamelCase_					=  GenerationConfig.from_model_config(UpperCAmelCase   )
							assert not hasattr(UpperCAmelCase							,						'''foo'''   )  # no new kwargs should be initialized if from config





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  GenerationConfig()
							self.assertEqual(default_config.temperature							,						1.0   )
							self.assertEqual(default_config.do_sample							,						UpperCAmelCase   )
							self.assertEqual(default_config.num_beams							,						1   )

							lowerCamelCase_					=  GenerationConfig(
							    do_sample=UpperCAmelCase							,						temperature=0.7							,						length_penalty=1.0							,						bad_words_ids=[[1, 2, 3], [4, 5]]							,						)
							self.assertEqual(config.temperature							,						0.7   )
							self.assertEqual(config.do_sample							,						UpperCAmelCase   )
							self.assertEqual(config.num_beams							,						1   )

							with tempfile.TemporaryDirectory() as tmp_dir:
										config.save_pretrained(UpperCAmelCase   )
										lowerCamelCase_					=  GenerationConfig.from_pretrained(UpperCAmelCase							,						temperature=1.0   )

							self.assertEqual(loaded_config.temperature							,						1.0   )
							self.assertEqual(loaded_config.do_sample							,						UpperCAmelCase   )
							self.assertEqual(loaded_config.num_beams							,						1   )  # default value




@is_staging_test
class  __lowerCamelCase  ( unittest.TestCase			):




				@classmethod
				def        UpperCAmelCase__						(				cls   ):
							lowerCamelCase_					=  TOKEN
							HfFolder.save_token(UpperCAmelCase   )




				@classmethod
				def        UpperCAmelCase__						(				cls   ):
							try:
										delete_repo(token=cls._token							,						repo_id='''test-generation-config'''   )
							except HTTPError:
										pass

							try:
										delete_repo(token=cls._token							,						repo_id='''valid_org/test-generation-config-org'''   )
							except HTTPError:
										pass




				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  GenerationConfig(
							    do_sample=UpperCAmelCase							,						temperature=0.7							,						length_penalty=1.0							,						)
							config.push_to_hub('''test-generation-config'''							,						use_auth_token=self._token   )

							lowerCamelCase_					=  GenerationConfig.from_pretrained(f"{USER}/test-generation-config"   )
							for k, v in config.to_dict().items():
										if k != "transformers_version":
													self.assertEqual(UpperCAmelCase							,						getattr(UpperCAmelCase							,						UpperCAmelCase   )   )

        # Reset repo
							delete_repo(token=self._token							,						repo_id='''test-generation-config'''   )

							# Push to hub via save_pretrained
							with tempfile.TemporaryDirectory() as tmp_dir:
										config.save_pretrained(
										    UpperCAmelCase							,						repo_id='''test-generation-config'''							,						push_to_hub=UpperCAmelCase							,						use_auth_token=self._token   )

							lowerCamelCase_					=  GenerationConfig.from_pretrained(f"{USER}/test-generation-config"   )
							for k, v in config.to_dict().items():
										if k != "transformers_version":
													self.assertEqual(UpperCAmelCase							,						getattr(UpperCAmelCase							,						UpperCAmelCase   )   )





				def        UpperCAmelCase__						(				self   ):
							lowerCamelCase_					=  GenerationConfig(
							    do_sample=UpperCAmelCase							,						temperature=0.7							,						length_penalty=1.0							,						)
							config.push_to_hub('''valid_org/test-generation-config-org'''							,						use_auth_token=self._token   )

							lowerCamelCase_					=  GenerationConfig.from_pretrained('''valid_org/test-generation-config-org'''   )
							for k, v in config.to_dict().items():
										if k != "transformers_version":
													self.assertEqual(UpperCAmelCase							,						getattr(UpperCAmelCase							,						UpperCAmelCase   )   )

        # Reset repo
							delete_repo(token=self._token							,						repo_id='''valid_org/test-generation-config-org'''   )

							# Push to hub via save_pretrained
							with tempfile.TemporaryDirectory() as tmp_dir:
										config.save_pretrained(
										    UpperCAmelCase							,						repo_id='''valid_org/test-generation-config-org'''							,						push_to_hub=UpperCAmelCase							,						use_auth_token=self._token   )

							lowerCamelCase_					=  GenerationConfig.from_pretrained('''valid_org/test-generation-config-org'''   )
							for k, v in config.to_dict().items():
										if k != "transformers_version":
													self.assertEqual(UpperCAmelCase							,						getattr(UpperCAmelCase							,						UpperCAmelCase   )   )



