


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 _lowerCamelCase:
       def __init__(					self,      lowerCamelCase,      )       ->   int:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		parent
           _lowercase       :							Optional[int]      		=		13
           _lowercase       :							List[Any]      		=		7
           _lowercase       :							List[str]      		=		True
           _lowercase       :							Optional[Any]      		=		True
           _lowercase       :							Dict      		=		True
           _lowercase       :							Dict      		=		99
           _lowercase       :							Optional[int]      		=		32
           _lowercase       :							Optional[Any]      		=		2
           _lowercase       :							Union[str, Any]      		=		4
           _lowercase       :							Dict      		=		37
           _lowercase       :							str      		=		'gelu'
           _lowercase       :							Optional[Any]      		=		0.1
           _lowercase       :							Tuple      		=		0.1
           _lowercase       :							Optional[Any]      		=		5_12
           _lowercase       :							str      		=		16
           _lowercase       :							Any      		=		2
           _lowercase       :							str      		=		0.0_2
           _lowercase       :							str      		=		3
           _lowercase       :							List[str]      		=		4
           _lowercase       :							Optional[Any]      		=		None
       def 			UpperCamelCase     (					self)       ->   List[str]:




           """simple docstring"""

           _lowercase       :							Dict      		=		ids_tensor([self.batch_size, self.seq_length],      self.vocab_size)

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

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

           _lowercase       :							Optional[int]      		=		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)       ->   str:




           """simple docstring"""

           (
               (
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				
           )       :							List[Any]      		=		self.prepare_config_and_inputs()

           _lowercase       :							Any      		=		True
           _lowercase       :							Tuple      		=		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,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   str:




           """simple docstring"""

           _lowercase       :							str      		=		TFEsmModel(config=lowerCamelCase)
           _lowercase       :							Dict      		=		{'input_ids': input_ids, 'attention_mask': input_mask}
           _lowercase       :							Tuple      		=		model(lowerCamelCase)

           _lowercase       :							int      		=		[input_ids, input_mask]
           _lowercase       :							Tuple      		=		model(lowerCamelCase)

           _lowercase       :							Tuple      		=		model(lowerCamelCase)

           self.parent.assertEqual(result.last_hidden_state.shape,      (self.batch_size, self.seq_length, self.hidden_size))
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      )       ->   int:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		True

           _lowercase       :							List[Any]      		=		TFEsmModel(config=lowerCamelCase)
           _lowercase       :							str      		=		{
               'input_ids': input_ids,
               'attention_mask': input_mask,
               'encoder_hidden_states': encoder_hidden_states,
               'encoder_attention_mask': encoder_attention_mask,
           }
           _lowercase       :							Union[str, Any]      		=		model(lowerCamelCase)

           _lowercase       :							Optional[int]      		=		[input_ids, input_mask]
           _lowercase       :							Dict      		=		model(lowerCamelCase,      encoder_hidden_states=lowerCamelCase)

           # Also check the case where encoder outputs are not passed
           _lowercase       :							str      		=		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,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							Tuple      		=		TFEsmForMaskedLM(config=lowerCamelCase)
           _lowercase       :							Optional[int]      		=		model([input_ids, input_mask])
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.vocab_size))
       def 			UpperCamelCase     (					self,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase,      lowerCamelCase)       ->   Any:




           """simple docstring"""

           _lowercase       :							str      		=		self.num_labels
           _lowercase       :							Tuple      		=		TFEsmForTokenClassification(config=lowerCamelCase)
           _lowercase       :							Dict      		=		{'input_ids': input_ids, 'attention_mask': input_mask}
           _lowercase       :							int      		=		model(lowerCamelCase)
           self.parent.assertEqual(result.logits.shape,      (self.batch_size, self.seq_length, self.num_labels))




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




           """simple docstring"""

           _lowercase       :							Union[str, Any]      		=		self.prepare_config_and_inputs()
           (
               (
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				(
               _lowercase
           )		,				
           )       :							Optional[int]      		=		config_and_inputs
           _lowercase       :							List[Any]      		=		{'input_ids': input_ids, 'attention_mask': input_mask}
           return config, inputs_dict






@require_tf
class _lowerCamelCase(						_a, _a, unittest.TestCase   ):
       lowercase_ :	str     =			(
           (
               TFEsmModel,
               TFEsmForMaskedLM,
               TFEsmForSequenceClassification,
               TFEsmForTokenClassification,
           )
           if is_tf_available()
           else ()
       )
       lowercase_ :	Optional[int]     =			(
           {
               """feature-extraction""": TFEsmModel,
               """fill-mask""": TFEsmForMaskedLM,
               """text-classification""": TFEsmForSequenceClassification,
               """token-classification""": TFEsmForTokenClassification,
               """zero-shot""": TFEsmForSequenceClassification,
           }
           if is_tf_available()
           else {}
       )
       lowercase_ :	List[str]     =			False
       lowercase_ :	int     =			False
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Optional[Any]      		=		TFEsmModelTester(self)
           _lowercase       :							str      		=		ConfigTester(self,      config_class=lowerCamelCase,      hidden_size=37)
       def 			UpperCamelCase     (					self)       ->   Dict:




           """simple docstring"""

           self.config_tester.run_common_tests()
       def 			UpperCamelCase     (					self)       ->   Any:




           """simple docstring"""

           _lowercase       :							Tuple      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   Optional[Any]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		self.model_tester.prepare_config_and_inputs_for_decoder()
           self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   str:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase)
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							str      		=		self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*lowerCamelCase)
       @slow
       def 			UpperCamelCase     (					self)       ->   Dict:




           """simple docstring"""

           for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _lowercase       :							str      		=		TFEsmModel.from_pretrained(lowerCamelCase)
               self.assertIsNotNone(lowerCamelCase)
       @unittest.skip('Protein models do not support embedding resizing.')
       def 			UpperCamelCase     (					self)       ->   Dict:




           """simple docstring"""

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




           """simple docstring"""

           pass




       def 			UpperCamelCase     (					self)       ->   Optional[int]:




           """simple docstring"""

           _lowercase		,				_lowercase       :							int      		=		self.model_tester.prepare_config_and_inputs_for_common()

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






@require_tf
class _lowerCamelCase(						unittest.TestCase   ):
       @slow
       def 			UpperCamelCase     (					self)       ->   int:




           """simple docstring"""

           _lowercase       :							Tuple      		=		TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D')

           _lowercase       :							Union[str, Any]      		=		tf.constant([[0, 1, 2, 3, 4, 5]])
           _lowercase       :							List[str]      		=		model(lowerCamelCase)[0]
           _lowercase       :							Tuple      		=		[1, 6, 33]
           self.assertEqual(list(output.numpy().shape),      lowerCamelCase)
           # compare the actual values for a slice.
           _lowercase       :							Optional[int]      		=		tf.constant(
               [
                   [
                       [8.9_2_1_5_1_8, -1_0.5_8_9_8_1_4, -6.4_6_7_1_3_0_7],
                       [-6.3_9_6_7_1_5_6, -1_3.9_1_1_3_7_7, -1.1_2_1_1_9_1_5],
                       [-7.7_8_1_2_4_7, -1_3.9_5_1_5_5_7, -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]:




           """simple docstring"""

           _lowercase       :							Optional[int]      		=		TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D')

           _lowercase       :							str      		=		tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]])
           _lowercase       :							Any      		=		model(lowerCamelCase)[0]
           # compare the actual values for a slice.
           _lowercase       :							Tuple      		=		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))






import inspect
import warnings
from typing import Any, Dict, Optional, Union

from packaging import version


def   _SCREAMING_SNAKE_CASE	( *SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE = None ,			SCREAMING_SNAKE_CASE=True ,			SCREAMING_SNAKE_CASE=2       ):
  from .. import __version__

  A_      :				Union[str, Any]		     =					take_from
  A_      :				Optional[Any]		     =					()
  if not isinstance(args[0] ,			SCREAMING_SNAKE_CASE       ):
    A_      :				List[str]		     =					(args,)

  for attribute, version_name, message in args:
    if version.parse(version.parse(SCREAMING_SNAKE_CASE       ).base_version       ) >= version.parse(SCREAMING_SNAKE_CASE       ):
      raise ValueError(
          f'''The deprecation tuple {(attribute, version_name, message)} should be removed since diffusers\''''
          f''' version {__version__} is >= {version_name}'''       )

    A_      :				List[Any]		     =					None
    if isinstance(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ) and attribute in deprecated_kwargs:
      values += (deprecated_kwargs.pop(SCREAMING_SNAKE_CASE       ),)
      A_      :				Optional[Any]		     =					f'''The `{attribute}` argument is deprecated and will be removed in version {version_name}.'''
    elif hasattr(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ):
      values += (getattr(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ),)
      A_      :				int		     =					f'''The `{attribute}` attribute is deprecated and will be removed in version {version_name}.'''
    elif deprecated_kwargs is None:
      A_      :				List[Any]		     =					f'''`{attribute}` is deprecated and will be removed in version {version_name}.'''

    if warning is not None:
      A_      :				Union[str, Any]		     =					warning + ''' ''' if standard_warn else ''''''
      warnings.warn(warning + message ,			SCREAMING_SNAKE_CASE ,			stacklevel=SCREAMING_SNAKE_CASE       )

  if isinstance(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE       ) and len(SCREAMING_SNAKE_CASE       ) > 0:
    A_      :				Dict		     =					inspect.getouterframes(inspect.currentframe()       )[1]
    A_      :				Optional[int]		     =					call_frame.filename
    A_      :				Optional[int]		     =					call_frame.lineno
    A_      :				str		     =					call_frame.function
    A_     ,       A_      :				List[str]		     =					next(iter(deprecated_kwargs.items()       )       )
    raise TypeError(f'''{function} in {filename} line {line_number-1} got an unexpected keyword argument `{key}`'''       )

  if len(SCREAMING_SNAKE_CASE       ) == 0:
    return
  elif len(SCREAMING_SNAKE_CASE       ) == 1:
    return values[0]
  return values




