


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							:     List[str]							,	lowerCamelCase__							:     Tuple							,	)  ->      Tuple:





						'''simple docstring'''






						UpperCamelCase__					:    Optional[int]          =   parent
						UpperCamelCase__					:    int          =   13
						UpperCamelCase__					:    Tuple          =   7
						UpperCamelCase__					:    Optional[Any]          =   True
						UpperCamelCase__					:    Dict          =   True
						UpperCamelCase__					:    List[str]          =   True
						UpperCamelCase__					:    Optional[Any]          =   99
						UpperCamelCase__					:    Dict          =   32
						UpperCamelCase__					:    Any          =   2
						UpperCamelCase__					:    int          =   4
						UpperCamelCase__					:    Optional[int]          =   37
						UpperCamelCase__					:    Optional[Any]          =   '''gelu'''
						UpperCamelCase__					:    Union[str, Any]          =   0.1
						UpperCamelCase__					:    Union[str, Any]          =   0.1
						UpperCamelCase__					:    Dict          =   512
						UpperCamelCase__					:    int          =   16
						UpperCamelCase__					:    str          =   2
						UpperCamelCase__					:    Dict          =   0.02
						UpperCamelCase__					:    Optional[int]          =   3
						UpperCamelCase__					:    Optional[int]          =   4
						UpperCamelCase__					:    str          =   None


					def      UpperCAmelCase__	(      self							:     str			)  ->      Any:





						'''simple docstring'''






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

						UpperCamelCase__					:    int          =   None
						if self.use_input_mask:
							UpperCamelCase__					:    Tuple          =   random_attention_mask([self.batch_size, self.seq_length]			)

						UpperCamelCase__					:    List[str]          =   None
						UpperCamelCase__					:    List[str]          =   None
						UpperCamelCase__					:    Union[str, Any]          =   None
						if self.use_labels:
							UpperCamelCase__					:    Optional[Any]          =   ids_tensor([self.batch_size]							,	self.type_sequence_label_size			)
							UpperCamelCase__					:    int          =   ids_tensor([self.batch_size, self.seq_length]							,	self.num_labels			)
							UpperCamelCase__					:    str          =   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							:     Optional[Any]			)  ->      List[Any]:





						'''simple docstring'''






						(
						    UpperCamelCase__				
						)					:    List[Any]          =   self.prepare_config_and_inputs()

						UpperCamelCase__					:    Dict          =   True
						UpperCamelCase__					:    int          =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
						UpperCamelCase__					:    Optional[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							:     Any							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Union[str, Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Dict			)  ->      List[Any]:





						'''simple docstring'''






						UpperCamelCase__					:    Union[str, Any]          =   TFEsmModel(config=lowerCamelCase__			)
						UpperCamelCase__					:    int          =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						UpperCamelCase__					:    int          =   model(lowerCamelCase__			)

						UpperCamelCase__					:    Any          =   [input_ids, input_mask]
						UpperCamelCase__					:    Dict          =   model(lowerCamelCase__			)

						UpperCamelCase__					:    str          =   model(lowerCamelCase__			)

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


					def      UpperCAmelCase__	(      self							:     List[Any]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     str							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Tuple							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Optional[Any]							,	)  ->      Dict:





						'''simple docstring'''






						UpperCamelCase__					:    List[Any]          =   True

						UpperCamelCase__					:    Optional[Any]          =   TFEsmModel(config=lowerCamelCase__			)
						UpperCamelCase__					:    Optional[Any]          =   {
						    '''input_ids''': input_ids,
						    '''attention_mask''': input_mask,
						    '''encoder_hidden_states''': encoder_hidden_states,
						    '''encoder_attention_mask''': encoder_attention_mask,
						}
						UpperCamelCase__					:    Tuple          =   model(lowerCamelCase__			)

						UpperCamelCase__					:    List[str]          =   [input_ids, input_mask]
						UpperCamelCase__					:    Optional[Any]          =   model(lowerCamelCase__							,	encoder_hidden_states=lowerCamelCase__			)

						# Also check the case where encoder outputs are not passed
						UpperCamelCase__					:    List[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							:     Optional[int]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     int							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     List[Any]			)  ->      Any:





						'''simple docstring'''






						UpperCamelCase__					:    int          =   TFEsmForMaskedLM(config=lowerCamelCase__			)
						UpperCamelCase__					:    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							:     Dict							,	lowerCamelCase__							:     List[Any]							,	lowerCamelCase__							:     Any							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     Dict							,	lowerCamelCase__							:     Optional[Any]							,	lowerCamelCase__							:     str			)  ->      str:





						'''simple docstring'''






						UpperCamelCase__					:    Optional[Any]          =   self.num_labels
						UpperCamelCase__					:    List[Any]          =   TFEsmForTokenClassification(config=lowerCamelCase__			)
						UpperCamelCase__					:    Any          =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						UpperCamelCase__					:    Any          =   model(lowerCamelCase__			)
						self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.seq_length, self.num_labels)			)


					def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      Tuple:





						'''simple docstring'''






						UpperCamelCase__					:    Any          =   self.prepare_config_and_inputs()
						(
						    UpperCamelCase__				
						)					:    int          =   config_and_inputs
						UpperCamelCase__					:    Tuple          =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						return config, inputs_dict



@require_tf
class       __magic_name__	(     __lowerCAmelCase					,				__lowerCAmelCase					,				unittest.TestCase):
					A:		Any			=      (
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					A:		Dict			=      (
					    {
					        "feature-extraction": TFEsmModel,
					        "fill-mask": TFEsmForMaskedLM,
					        "text-classification": TFEsmForSequenceClassification,
					        "token-classification": TFEsmForTokenClassification,
					        "zero-shot": TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					A:		Optional[Any]			=      False
					A:		Union[str, Any]			=      False


					def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      Any:





						'''simple docstring'''






						UpperCamelCase__					:    Union[str, Any]          =   TFEsmModelTester(self			)
						UpperCamelCase__					:    Optional[Any]          =   ConfigTester(self							,	config_class=lowerCamelCase__							,	hidden_size=37			)


					def      UpperCAmelCase__	(      self							:     Any			)  ->      Any:





						'''simple docstring'''






						self.config_tester.run_common_tests()


					def      UpperCAmelCase__	(      self							:     List[str]			)  ->      Optional[int]:





						'''simple docstring'''






						UpperCamelCase__					:    Optional[Any]          =   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowerCamelCase__			)


					def      UpperCAmelCase__	(      self							:     int			)  ->      Any:





						'''simple docstring'''






						UpperCamelCase__					:    List[str]          =   self.model_tester.prepare_config_and_inputs_for_decoder()
						self.model_tester.create_and_check_model_as_decoder(*lowerCamelCase__			)


					def      UpperCAmelCase__	(      self							:     Tuple			)  ->      Optional[int]:





						'''simple docstring'''






						UpperCamelCase__					:    Optional[int]          =   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase__			)


					def      UpperCAmelCase__	(      self							:     Dict			)  ->      List[Any]:





						'''simple docstring'''






						UpperCamelCase__					:    Union[str, Any]          =   self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*lowerCamelCase__			)


					@slow
					def      UpperCAmelCase__	(      self							:     Dict			)  ->      Optional[int]:





						'''simple docstring'''






						for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							UpperCamelCase__					:    Tuple          =   TFEsmModel.from_pretrained(lowerCamelCase__			)
							self.assertIsNotNone(lowerCamelCase__			)


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





						'''simple docstring'''






						pass


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





						'''simple docstring'''






						pass


					def      UpperCAmelCase__	(      self							:     str			)  ->      int:





						'''simple docstring'''






						UpperCamelCase__					:    str          =   self.model_tester.prepare_config_and_inputs_for_common()

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



@require_tf
class       __magic_name__	(     unittest.TestCase):


					@slow
					def      UpperCAmelCase__	(      self							:     Union[str, Any]			)  ->      Dict:





						'''simple docstring'''






						UpperCamelCase__					:    Dict          =   TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''			)

						UpperCamelCase__					:    Optional[int]          =   tf.constant([[0, 1, 2, 3, 4, 5]]			)
						UpperCamelCase__					:    List[Any]          =   model(lowerCamelCase__			)[0]
						UpperCamelCase__					:    Tuple          =   [1, 6, 33]
						self.assertEqual(list(output.numpy().shape			)							,	lowerCamelCase__			)
						# compare the actual values for a slice.
						UpperCamelCase__					:    List[str]          =   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							:     Optional[Any]			)  ->      str:





						'''simple docstring'''






						UpperCamelCase__					:    Dict          =   TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''			)

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






from typing import TYPE_CHECKING

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


__UpperCamelCase       :       Optional[int]						       =				{
    "configuration_rag": ["RagConfig"],
    "retrieval_rag": ["RagRetriever"],
    "tokenization_rag": ["RagTokenizer"],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase       :       Optional[int]						       =				[
           "RagModel",
           "RagPreTrainedModel",
           "RagSequenceForGeneration",
           "RagTokenForGeneration",
       ]

try:
       if not is_tf_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       __UpperCamelCase       :       List[Any]						       =				[
           "TFRagModel",
           "TFRagPreTrainedModel",
           "TFRagSequenceForGeneration",
           "TFRagTokenForGeneration",
       ]


if TYPE_CHECKING:
       from .configuration_rag import RagConfig
       from .retrieval_rag import RagRetriever
       from .tokenization_rag import RagTokenizer

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

       try:
              if not is_tf_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_tf_rag import (
                  TFRagModel,
                  TFRagPreTrainedModel,
                  TFRagSequenceForGeneration,
                  TFRagTokenForGeneration,
              )

else:
       import sys

       __UpperCamelCase       :       List[Any]						       =				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


