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







				def __init__(       self						,				A_						,				)->		Optional[int]:







											'''simple docstring'''






											UpperCamelCase   			=     parent
											UpperCamelCase   			=     13
											UpperCamelCase   			=     7
											UpperCamelCase   			=     True
											UpperCamelCase   			=     True
											UpperCamelCase   			=     True
											UpperCamelCase   			=     99
											UpperCamelCase   			=     32
											UpperCamelCase   			=     2
											UpperCamelCase   			=     4
											UpperCamelCase   			=     37
											UpperCamelCase   			=     'gelu'
											UpperCamelCase   			=     0.1
											UpperCamelCase   			=     0.1
											UpperCamelCase   			=     512
											UpperCamelCase   			=     16
											UpperCamelCase   			=     2
											UpperCamelCase   			=     0.02
											UpperCamelCase   			=     3
											UpperCamelCase   			=     4
											UpperCamelCase   			=     None







				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											UpperCamelCase   			=     ids_tensor([self.batch_size, self.seq_length]						,				self.vocab_size     )

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

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

											UpperCamelCase   			=     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]:







											'''simple docstring'''






											(
											    (
											    UpperCamelCase    
											)			,	(
											    UpperCamelCase    
											)			,	(
											    UpperCamelCase    
											)			,	(
											    UpperCamelCase    
											)			,	(
											    UpperCamelCase    
											)			,	(
											    UpperCamelCase    
											)			,	    
											)   			=     self.prepare_config_and_inputs()

											UpperCamelCase   			=     True
											UpperCamelCase   			=     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]     )
											UpperCamelCase   			=     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						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_     )->		Union[str, Any]:







											'''simple docstring'''






											UpperCamelCase   			=     TFEsmModel(config=A_     )
											UpperCamelCase   			=     {'input_ids': input_ids, 'attention_mask': input_mask}
											UpperCamelCase   			=     model(A_     )

											UpperCamelCase   			=     [input_ids, input_mask]
											UpperCamelCase   			=     model(A_     )

											UpperCamelCase   			=     model(A_     )

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







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_						,				)->		Tuple:







											'''simple docstring'''






											UpperCamelCase   			=     True

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

											UpperCamelCase   			=     [input_ids, input_mask]
											UpperCamelCase   			=     model(A_						,				encoder_hidden_states=A_     )

											# Also check the case where encoder outputs are not passed
											UpperCamelCase   			=     model(A_						,				attention_mask=A_     )

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







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_     )->		Optional[int]:







											'''simple docstring'''






											UpperCamelCase   			=     TFEsmForMaskedLM(config=A_     )
											UpperCamelCase   			=     model([input_ids, input_mask]     )
											self.parent.assertEqual(result.logits.shape						,				(self.batch_size, self.seq_length, self.vocab_size)     )







				def 						UpperCAmelCase_    (       self						,				A_						,				A_						,				A_						,				A_						,				A_						,				A_     )->		Any:







											'''simple docstring'''






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







				def 						UpperCAmelCase_    (       self     )->		Optional[Any]:







											'''simple docstring'''






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


@require_tf
class 			SCREAMING_SNAKE_CASE__      (     UpperCAmelCase__							,       UpperCAmelCase__							,       unittest.TestCase):
				lowerCAmelCase_     			=							(
				    (
				        TFEsmModel,
				        TFEsmForMaskedLM,
				        TFEsmForSequenceClassification,
				        TFEsmForTokenClassification,
				    )
				    if is_tf_available()
				    else ()
				)
				lowerCAmelCase_     			=							(
				    {
				        "feature-extraction": TFEsmModel,
				        "fill-mask": TFEsmForMaskedLM,
				        "text-classification": TFEsmForSequenceClassification,
				        "token-classification": TFEsmForTokenClassification,
				        "zero-shot": TFEsmForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)
				lowerCAmelCase_     			=							False
				lowerCAmelCase_     			=							False







				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase   			=     TFEsmModelTester(self     )
											UpperCamelCase   			=     ConfigTester(self						,				config_class=A_						,				hidden_size=37     )







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







											'''simple docstring'''






											self.config_tester.run_common_tests()







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*A_     )







				def 						UpperCAmelCase_    (       self     )->		int:







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*A_     )







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







											'''simple docstring'''






											UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*A_     )







				@slow
				def 						UpperCAmelCase_    (       self     )->		Any:







											'''simple docstring'''






											for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																		UpperCamelCase   			=     TFEsmModel.from_pretrained(A_     )
																		self.assertIsNotNone(A_     )







				@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[Any]:







											'''simple docstring'''






											pass







				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase			,	UpperCamelCase   			=     self.model_tester.prepare_config_and_inputs_for_common()

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


@require_tf
class 			SCREAMING_SNAKE_CASE__      (     unittest.TestCase):







				@slow
				def 						UpperCAmelCase_    (       self     )->		Tuple:







											'''simple docstring'''






											UpperCamelCase   			=     TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'     )

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







											'''simple docstring'''






											UpperCamelCase   			=     TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'     )

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





'''simple docstring'''




import re

from flax.core.frozen_dict import freeze
from flax.traverse_util import flatten_dict, unflatten_dict
from jax.experimental import PartitionSpec as P


# Sentinels
__UpperCAmelCase					=object()

# For specifying empty leaf dict `{}`
__UpperCAmelCase					=object()






def __lowerCAmelCase					(				UpperCamelCase__							,    UpperCamelCase__							)      ->       List[Any]:
    __lowerCamelCase            =    tuple((re.compile(x + '''$'''							) for x in qs)							)
    for i in range(len(UpperCamelCase__							) - len(UpperCamelCase__							) + 1							):
        __lowerCamelCase            =    [x.match(UpperCamelCase__							) for x, y in zip(UpperCamelCase__							,    ks[i:]							)]
        if matches and all(UpperCamelCase__							):
            return True
    return False






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       Optional[int]:
    def replace(UpperCamelCase__							,    UpperCamelCase__							):
        for rule, replacement in rules:
            if _match(UpperCamelCase__							,    UpperCamelCase__							):
                return replacement
        return val

    return replace






def __lowerCAmelCase					(				)      ->       Dict:
    return [
        # embeddings
        (("transformer", "wpe", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        (("transformer", "wte", "embedding"), P('''mp'''							,    UpperCamelCase__							)),
        # atention
        (("attention", "(q_proj|k_proj|v_proj)", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("attention", "out_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("attention", "out_proj", "bias"), None),
        # mlp
        (("mlp", "c_fc", "kernel"), P(UpperCamelCase__							,    '''mp'''							)),
        (("mlp", "c_fc", "bias"), P('''mp'''							)),
        (("mlp", "c_proj", "kernel"), P('''mp'''							,    UpperCamelCase__							)),
        (("mlp", "c_proj", "bias"), None),
        # layer norms
        ((r"ln_\d+", "bias"), None),
        ((r"\d+", r"ln_\d+", "scale"), None),
        (("ln_f", "bias"), None),
        (("ln_f", "scale"), None),
    ]






def __lowerCAmelCase					(				UpperCamelCase__							)      ->       int:
    __lowerCamelCase            =    _get_partition_rules()
    __lowerCamelCase            =    _replacement_rules(UpperCamelCase__							)
    __lowerCamelCase            =    {k: _unmatched for k in flatten_dict(UpperCamelCase__							)}
    __lowerCamelCase            =    {k: replace(UpperCamelCase__							,    UpperCamelCase__							) for k, v in initd.items()}
    assert _unmatched not in result.values(), "Incomplete partition spec."
    return freeze(unflatten_dict(UpperCamelCase__							)							)
