


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

 '''simple docstring'''

 def __init__(     self		,				_UpperCAmelCase		,				):
       snake_case_	=				parent
       snake_case_	=				13
       snake_case_	=				7
       snake_case_	=				True
       snake_case_	=				True
       snake_case_	=				True
       snake_case_	=				99
       snake_case_	=				32
       snake_case_	=				2
       snake_case_	=				4
       snake_case_	=				37
       snake_case_	=				'''gelu'''
       snake_case_	=				0.1
       snake_case_	=				0.1
       snake_case_	=				5_12
       snake_case_	=				16
       snake_case_	=				2
       snake_case_	=				0.02
       snake_case_	=				3
       snake_case_	=				4
       snake_case_	=				None



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size )

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

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

       snake_case_	=				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 ):
       (
           (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   
       )	=				self.prepare_config_and_inputs()

       snake_case_	=				True
       snake_case_	=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size] )
       snake_case_	=				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		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				TFEsmModel(config=_UpperCAmelCase )
       snake_case_	=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       snake_case_	=				model(_UpperCAmelCase )

       snake_case_	=				[input_ids, input_mask]
       snake_case_	=				model(_UpperCAmelCase )

       snake_case_	=				model(_UpperCAmelCase )

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



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				):
       snake_case_	=				True

       snake_case_	=				TFEsmModel(config=_UpperCAmelCase )
       snake_case_	=				{
           '''input_ids''': input_ids,
           '''attention_mask''': input_mask,
           '''encoder_hidden_states''': encoder_hidden_states,
           '''encoder_attention_mask''': encoder_attention_mask,
       }
       snake_case_	=				model(_UpperCAmelCase )

       snake_case_	=				[input_ids, input_mask]
       snake_case_	=				model(_UpperCAmelCase		,				encoder_hidden_states=_UpperCAmelCase )

       # Also check the case where encoder outputs are not passed
       snake_case_	=				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		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				TFEsmForMaskedLM(config=_UpperCAmelCase )
       snake_case_	=				model([input_ids, input_mask] )
       self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size) )



 def 						UpperCamelCase__					(     self		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase		,				_UpperCAmelCase ):
       snake_case_	=				self.num_labels
       snake_case_	=				TFEsmForTokenClassification(config=_UpperCAmelCase )
       snake_case_	=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       snake_case_	=				model(_UpperCAmelCase )
       self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.num_labels) )





 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.prepare_config_and_inputs()
       (
           (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   (
           snake_case_
       )	,   
       )	=				config_and_inputs
       snake_case_	=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
       return config, inputs_dict







@require_tf
class lowerCAmelCase_						(			lowerCamelCase__ ,						lowerCamelCase__ ,						unittest.TestCase	):

 '''simple docstring'''

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



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				TFEsmModelTester(self )
       snake_case_	=				ConfigTester(self		,				config_class=_UpperCAmelCase		,				hidden_size=37 )



 def 						UpperCamelCase__					(     self ):
       self.config_tester.run_common_tests()



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*_UpperCAmelCase )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase )



 def 						UpperCamelCase__					(     self ):
       snake_case_	=				self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase )



 @slow
 def 						UpperCamelCase__					(     self ):
       for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             snake_case_	=				TFEsmModel.from_pretrained(_UpperCAmelCase )
             self.assertIsNotNone(_UpperCAmelCase )



 @unittest.skip('''Protein models do not support embedding resizing.''' )
 def 						UpperCamelCase__					(     self ):
       pass



 @unittest.skip('''Protein models do not support embedding resizing.''' )
 def 						UpperCamelCase__					(     self ):
       pass





 def 						UpperCamelCase__					(     self ):
       snake_case_	,   snake_case_	=				self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class lowerCAmelCase_						(			unittest.TestCase	):

 '''simple docstring'''

 @slow
 def 						UpperCamelCase__					(     self ):
       snake_case_	=				TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )

       snake_case_	=				tf.constant([[0, 1, 2, 3, 4, 5]] )
       snake_case_	=				model(_UpperCAmelCase )[0]
       snake_case_	=				[1, 6, 33]
       self.assertEqual(list(output.numpy().shape )		,				_UpperCAmelCase )
       # compare the actual values for a slice.
       snake_case_	=				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 ):
       snake_case_	=				TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D''' )

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



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_sentencepiece_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase      =   {"""configuration_xglm""": ["""XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XGLMConfig"""]}

try:
   if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   ["""XGLMTokenizer"""]

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   ["""XGLMTokenizerFast"""]

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """XGLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """XGLMForCausalLM""",
       """XGLMModel""",
       """XGLMPreTrainedModel""",
   ]


try:
   if not is_flax_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """FlaxXGLMForCausalLM""",
       """FlaxXGLMModel""",
       """FlaxXGLMPreTrainedModel""",
   ]


try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCAmelCase      =   [
       """TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """TFXGLMForCausalLM""",
       """TFXGLMModel""",
       """TFXGLMPreTrainedModel""",
   ]


if TYPE_CHECKING:
   from .configuration_xglm import XGLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XGLMConfig

   try:
      if not is_sentencepiece_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_xglm import XGLMTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_xglm_fast import XGLMTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_xglm import XGLM_PRETRAINED_MODEL_ARCHIVE_LIST, XGLMForCausalLM, XGLMModel, XGLMPreTrainedModel

   try:
      if not is_flax_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_xglm import (
          TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFXGLMForCausalLM,
          TFXGLMModel,
          TFXGLMPreTrainedModel,
      )


else:
   import sys

   UpperCAmelCase      =   _LazyModule(__name__, globals()["""__file__"""], _import_structure)