

'''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 	A_			:






     '''simple docstring'''



     def __init__(		self			:					Any					,      lowercase_			:					Optional[int]					,      )						->     int:
      UpperCAmelCase	:     Tuple						 =			parent
      UpperCAmelCase	:     Union[str, Any]						 =			13
      UpperCAmelCase	:     List[Any]						 =			7
      UpperCAmelCase	:     int						 =			True
      UpperCAmelCase	:     Dict						 =			True
      UpperCAmelCase	:     Union[str, Any]						 =			True
      UpperCAmelCase	:     List[str]						 =			99
      UpperCAmelCase	:     Optional[Any]						 =			32
      UpperCAmelCase	:     Tuple						 =			2
      UpperCAmelCase	:     Optional[int]						 =			4
      UpperCAmelCase	:     Optional[Any]						 =			37
      UpperCAmelCase	:     Optional[int]						 =			'gelu'
      UpperCAmelCase	:     Optional[int]						 =			0.1
      UpperCAmelCase	:     int						 =			0.1
      UpperCAmelCase	:     List[str]						 =			512
      UpperCAmelCase	:     Union[str, Any]						 =			16
      UpperCAmelCase	:     Dict						 =			2
      UpperCAmelCase	:     Dict						 =			0.02
      UpperCAmelCase	:     str						 =			3
      UpperCAmelCase	:     Tuple						 =			4
      UpperCAmelCase	:     List[str]						 =			None


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

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

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

      UpperCAmelCase	:     Dict						 =			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	)						->     Any:
      (
          (
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			
      )	:     int						 =			self.prepare_config_and_inputs()

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

      UpperCAmelCase	:     List[Any]						 =			[input_ids, input_mask]
      UpperCAmelCase	:     Optional[Any]						 =			model(lowercase_	)

      UpperCAmelCase	:     str						 =			model(lowercase_	)

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


     def        UpperCAmelCase_	(		self			:					Union[str, Any]					,      lowercase_			:					List[str]					,      lowercase_			:					int					,      lowercase_			:					List[Any]					,      lowercase_			:					Dict					,      lowercase_			:					Optional[Any]					,      lowercase_			:					Any					,      lowercase_			:					Optional[Any]					,      lowercase_			:					List[Any]					,      )						->     Dict:
      UpperCAmelCase	:     Union[str, Any]						 =			True

      UpperCAmelCase	:     int						 =			TFEsmModel(config=lowercase_	)
      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	:     str						 =			model(lowercase_	)

      UpperCAmelCase	:     Optional[Any]						 =			[input_ids, input_mask]
      UpperCAmelCase	:     Any						 =			model(lowercase_					,      encoder_hidden_states=lowercase_	)

      # Also check the case where encoder outputs are not passed
      UpperCAmelCase	:     Optional[Any]						 =			model(lowercase_					,      attention_mask=lowercase_	)

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


     def        UpperCAmelCase_	(		self			:					Optional[Any]					,      lowercase_			:					Any					,      lowercase_			:					Any					,      lowercase_			:					Union[str, Any]					,      lowercase_			:					Optional[int]					,      lowercase_			:					str					,      lowercase_			:					str	)						->     Tuple:
      UpperCAmelCase	:     List[Any]						 =			TFEsmForMaskedLM(config=lowercase_	)
      UpperCAmelCase	:     Optional[Any]						 =			model([input_ids, input_mask]	)
      self.parent.assertEqual(result.logits.shape					,      (self.batch_size, self.seq_length, self.vocab_size)	)


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




     def        UpperCAmelCase_	(		self			:					str	)						->     Optional[int]:
      UpperCAmelCase	:     Dict						 =			self.prepare_config_and_inputs()
      (
          (
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			(
          UpperCAmelCase
      ) ,			
      )	:     int						 =			config_and_inputs
      UpperCAmelCase	:     List[Any]						 =			{'input_ids': input_ids, 'attention_mask': input_mask}
      return config, inputs_dict




@require_tf
class 	A_			(    _snake_case	,       _snake_case	,       unittest.TestCase							):






     '''simple docstring'''



     UpperCAmelCase_						:     Any  		=  (
         (
             TFEsmModel,
             TFEsmForMaskedLM,
             TFEsmForSequenceClassification,
             TFEsmForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     UpperCAmelCase_						:     List[str]  		=  (
         {
             """feature-extraction""": TFEsmModel,
             """fill-mask""": TFEsmForMaskedLM,
             """text-classification""": TFEsmForSequenceClassification,
             """token-classification""": TFEsmForTokenClassification,
             """zero-shot""": TFEsmForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     UpperCAmelCase_						:     Union[str, Any]  		=  False
     UpperCAmelCase_						:     Dict  		=  False


     def        UpperCAmelCase_	(		self			:					Any	)						->     Tuple:
      UpperCAmelCase	:     Any						 =			TFEsmModelTester(self	)
      UpperCAmelCase	:     Optional[Any]						 =			ConfigTester(self					,      config_class=lowercase_					,      hidden_size=37	)


     def        UpperCAmelCase_	(		self			:					Optional[int]	)						->     Optional[Any]:
      self.config_tester.run_common_tests()


     def        UpperCAmelCase_	(		self			:					Optional[Any]	)						->     Dict:
      UpperCAmelCase	:     Any						 =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowercase_	)


     def        UpperCAmelCase_	(		self			:					List[Any]	)						->     List[str]:
      UpperCAmelCase	:     List[Any]						 =			self.model_tester.prepare_config_and_inputs_for_decoder()
      self.model_tester.create_and_check_model_as_decoder(*lowercase_	)


     def        UpperCAmelCase_	(		self			:					Optional[int]	)						->     Union[str, Any]:
      UpperCAmelCase	:     str						 =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*lowercase_	)


     def        UpperCAmelCase_	(		self			:					Any	)						->     List[str]:
      UpperCAmelCase	:     Optional[Any]						 =			self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*lowercase_	)


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


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


     @unittest.skip('Protein models do not support embedding resizing.'	)
     def        UpperCAmelCase_	(		self			:					Union[str, Any]	)						->     Dict:
      pass




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

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




@require_tf
class 	A_			(    unittest.TestCase							):






     '''simple docstring'''



     @slow
     def        UpperCAmelCase_	(		self			:					str	)						->     Any:
      UpperCAmelCase	:     int						 =			TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'	)

      UpperCAmelCase	:     Optional[int]						 =			tf.constant([[0, 1, 2, 3, 4, 5]]	)
      UpperCAmelCase	:     Optional[int]						 =			model(lowercase_	)[0]
      UpperCAmelCase	:     Dict						 =			[1, 6, 33]
      self.assertEqual(list(output.numpy().shape	)					,      lowercase_	)
      # compare the actual values for a slice.
      UpperCAmelCase	:     Optional[Any]						 =			tf.constant(
          [
              [
                  [8.92_1518, -10.58_9814, -6.467_1307],
                  [-6.396_7156, -13.91_1377, -1.121_1915],
                  [-7.78_1247, -13.95_1557, -3.74_0592],
              ]
          ]	)
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()					,      expected_slice.numpy()					,      atol=1E-2	)	)




     @slow
     def        UpperCAmelCase_	(		self			:					Any	)						->     str:
      UpperCAmelCase	:     Optional[int]						 =			TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'	)

      UpperCAmelCase	:     Union[str, Any]						 =			tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]	)
      UpperCAmelCase	:     List[Any]						 =			model(lowercase_	)[0]
      # compare the actual values for a slice.
      UpperCAmelCase	:     Dict						 =			tf.constant(
          [
              [
                  [0.1444_3092, 0.5412_5327, 0.324_7739],
                  [0.3034_0484, 0.0052_6676, 0.3107_7722],
                  [0.3227_8043, -0.2498_7096, 0.341_4628],
              ]
          ]	)
      self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()					,      expected_slice.numpy()					,      atol=1E-4	)	)







'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


lowercase__         =      {
    "configuration_xlm": ["XLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "XLMConfig", "XLMOnnxConfig"],
    "tokenization_xlm": ["XLMTokenizer"],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowercase__         =      [
     "XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
     "XLMForMultipleChoice",
     "XLMForQuestionAnswering",
     "XLMForQuestionAnsweringSimple",
     "XLMForSequenceClassification",
     "XLMForTokenClassification",
     "XLMModel",
     "XLMPreTrainedModel",
     "XLMWithLMHeadModel",
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 lowercase__         =      [
     "TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST",
     "TFXLMForMultipleChoice",
     "TFXLMForQuestionAnsweringSimple",
     "TFXLMForSequenceClassification",
     "TFXLMForTokenClassification",
     "TFXLMMainLayer",
     "TFXLMModel",
     "TFXLMPreTrainedModel",
     "TFXLMWithLMHeadModel",
 ]


if TYPE_CHECKING:
 from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
 from .tokenization_xlm import XLMTokenizer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_xlm import (
      XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
      XLMForMultipleChoice,
      XLMForQuestionAnswering,
      XLMForQuestionAnsweringSimple,
      XLMForSequenceClassification,
      XLMForTokenClassification,
      XLMModel,
      XLMPreTrainedModel,
      XLMWithLMHeadModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_xlm import (
      TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFXLMForMultipleChoice,
      TFXLMForQuestionAnsweringSimple,
      TFXLMForSequenceClassification,
      TFXLMForTokenClassification,
      TFXLMMainLayer,
      TFXLMModel,
      TFXLMPreTrainedModel,
      TFXLMWithLMHeadModel,
  )

else:
 import sys

 lowercase__         =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




