





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




								a__     : Optional[int]  =parent
								a__     : Optional[Any]  =1_3
								a__     : List[Any]  =7
								a__     : str  =True
								a__     : Tuple  =True
								a__     : List[Any]  =True
								a__     : Dict  =9_9
								a__     : Dict  =3_2
								a__     : Any  =2
								a__     : Tuple  =4
								a__     : List[Any]  =3_7
								a__     : Optional[Any]  ='''gelu'''
								a__     : Tuple  =0.1
								a__     : str  =0.1
								a__     : Optional[Any]  =5_1_2
								a__     : Optional[Any]  =1_6
								a__     : Optional[int]  =2
								a__     : int  =0.02
								a__     : int  =3
								a__     : Optional[Any]  =4
								a__     : Optional[Any]  =None







						def        _lowercase  (   self							)       ->						int:



								'''simple docstring'''




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

								a__     : Tuple  =None
								if self.use_input_mask:
										a__     : Union[str, Any]  =random_attention_mask([self.batch_size, self.seq_length]							)

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

								a__     : List[str]  =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        _lowercase  (   self							)       ->						List[str]:



								'''simple docstring'''




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

								a__     : List[str]  =True
								a__     : List[str]  =floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
								a__     : List[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        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Any:



								'''simple docstring'''




								a__     : List[str]  =TFEsmModel(config=_lowerCamelCase							)
								a__     : int  ={'''input_ids''': input_ids, '''attention_mask''': input_mask}
								a__     : List[str]  =model(_lowerCamelCase							)

								a__     : Any  =[input_ids, input_mask]
								a__     : int  =model(_lowerCamelCase							)

								a__     : List[Any]  =model(_lowerCamelCase							)

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







						def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				)       ->						Optional[int]:



								'''simple docstring'''




								a__     : List[str]  =True

								a__     : Tuple  =TFEsmModel(config=_lowerCamelCase							)
								a__     : Tuple  ={
								    '''input_ids''': input_ids,
								    '''attention_mask''': input_mask,
								    '''encoder_hidden_states''': encoder_hidden_states,
								    '''encoder_attention_mask''': encoder_attention_mask,
								}
								a__     : str  =model(_lowerCamelCase							)

								a__     : Dict  =[input_ids, input_mask]
								a__     : Union[str, Any]  =model(_lowerCamelCase		,				encoder_hidden_states=_lowerCamelCase							)

								# Also check the case where encoder outputs are not passed
								a__     : List[Any]  =model(_lowerCamelCase		,				attention_mask=_lowerCamelCase							)

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







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



								'''simple docstring'''




								a__     : Any  =TFEsmForMaskedLM(config=_lowerCamelCase							)
								a__     : int  =model([input_ids, input_mask]							)
								self.parent.assertEqual(result.logits.shape		,				(self.batch_size, self.seq_length, self.vocab_size)							)







						def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Optional[int]:



								'''simple docstring'''




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







						def        _lowercase  (   self							)       ->						Any:



								'''simple docstring'''




								a__     : Union[str, Any]  =self.prepare_config_and_inputs()
								(
								    a__    
								)     : List[str]  =config_and_inputs
								a__     : Optional[int]  ={'''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    : Any            =      (
						    {
						        'feature-extraction': TFEsmModel,
						        'fill-mask': TFEsmForMaskedLM,
						        'text-classification': TFEsmForSequenceClassification,
						        'token-classification': TFEsmForTokenClassification,
						        'zero-shot': TFEsmForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						_lowercase    : Tuple            =      False
						_lowercase    : Dict            =      False







						def        _lowercase  (   self							)       ->						Tuple:



								'''simple docstring'''




								a__     : List[Any]  =TFEsmModelTester(self							)
								a__     : Dict  =ConfigTester(self		,				config_class=_lowerCamelCase		,				hidden_size=3_7							)







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



								'''simple docstring'''




								self.config_tester.run_common_tests()







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



								'''simple docstring'''




								a__     : List[Any]  =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*_lowerCamelCase							)







						def        _lowercase  (   self							)       ->						int:



								'''simple docstring'''




								a__     : Dict  =self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase							)







						def        _lowercase  (   self							)       ->						str:



								'''simple docstring'''




								a__     : str  =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*_lowerCamelCase							)







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



								'''simple docstring'''




								a__     : List[str]  =self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*_lowerCamelCase							)







						@slow
						def        _lowercase  (   self							)       ->						str:



								'''simple docstring'''




								for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										a__     : List[str]  =TFEsmModel.from_pretrained(_lowerCamelCase							)
										self.assertIsNotNone(_lowerCamelCase							)







						@unittest.skip("Protein models do not support embedding resizing."							)
						def        _lowercase  (   self							)       ->						Tuple:



								'''simple docstring'''




								pass







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



								'''simple docstring'''




								pass







						def        _lowercase  (   self							)       ->						int:



								'''simple docstring'''




								a__     : List[str]  =self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class     __lowerCAmelCase	(   unittest.TestCase):







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



								'''simple docstring'''




								a__     : Optional[Any]  =TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)

								a__     : Union[str, Any]  =tf.constant([[0, 1, 2, 3, 4, 5]]							)
								a__     : Union[str, Any]  =model(_lowerCamelCase							)[0]
								a__     : List[Any]  =[1, 6, 3_3]
								self.assertEqual(list(output.numpy().shape							)		,				_lowerCamelCase							)
								# compare the actual values for a slice.
								a__     : Any  =tf.constant(
								    [
								        [
								            [8.92_15_18, -1_0.5_8_9_8_1_4, -6.4_67_13_07],
								            [-6.3_96_71_56, -1_3.9_1_1_3_7_7, -1.1_21_19_15],
								            [-7.78_12_47, -1_3.9_5_1_5_5_7, -3.74_05_92],
								        ]
								    ]							)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,				expected_slice.numpy()		,				atol=1E-2							)							)







						@slow
						def        _lowercase  (   self							)       ->						Dict:



								'''simple docstring'''




								a__     : List[Any]  =TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)

								a__     : Tuple  =tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
								a__     : Any  =model(_lowerCamelCase							)[0]
								# compare the actual values for a slice.
								a__     : int  =tf.constant(
								    [
								        [
								            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
								            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
								            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
								        ]
								    ]							)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,				expected_slice.numpy()		,				atol=1E-4							)							)













import warnings

from transformers import AutoTokenizer
from transformers.utils import is_torch_available
from transformers.utils.generic import ExplicitEnum

from ...processing_utils import ProcessorMixin


if is_torch_available():
     import torch







class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : Any            =      """char"""
      _lowercase    : str            =      """bpe"""
      _lowercase    : List[Any]            =      """wp"""
UpperCAmelCase     :     List[Any]							=	(DecodeType.CHARACTER, DecodeType.BPE, DecodeType.WORDPIECE)







class     __lowerCAmelCase	(   UpperCamelCase__):
      _lowercase    : Tuple            =      ["""image_processor""", """char_tokenizer"""]
      _lowercase    : Any            =      """ViTImageProcessor"""
      _lowercase    : Optional[Any]            =      """MgpstrTokenizer"""







      def __init__(   self		,				lowerCAmelCase__=None		,				lowerCAmelCase__=None		,				**lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : Optional[int]  =None
        if "feature_extractor" in kwargs:
          warnings.warn(
              "The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`"
              " instead."		,				lowerCAmelCase__		,				)
          a__     : List[str]  =kwargs.pop("feature_extractor"							)

        a__     : List[str]  =image_processor if image_processor is not None else feature_extractor
        if image_processor is None:
          raise ValueError("You need to specify an `image_processor`."							)
        if tokenizer is None:
          raise ValueError("You need to specify a `tokenizer`."							)

        a__     : str  =tokenizer
        a__     : List[str]  =AutoTokenizer.from_pretrained("gpt2"							)
        a__     : Optional[int]  =AutoTokenizer.from_pretrained("bert-base-uncased"							)

        super().__init__(lowerCAmelCase__		,				lowerCAmelCase__							)







      def __call__(   self		,				lowerCAmelCase__=None		,				lowerCAmelCase__=None		,				lowerCAmelCase__=None		,				**lowerCAmelCase__							)       ->						Union[str, Any]:



        '''simple docstring'''




        if images is None and text is None:
          raise ValueError("You need to specify either an `images` or `text` input to process."							)

        if images is not None:
          a__     : Union[str, Any]  =self.image_processor(lowerCAmelCase__		,				return_tensors=lowerCAmelCase__		,				**lowerCAmelCase__							)
        if text is not None:
          a__     : int  =self.char_tokenizer(lowerCAmelCase__		,				return_tensors=lowerCAmelCase__		,				**lowerCAmelCase__							)

        if text is None:
          return inputs
        elif images is None:
          return encodings
        else:
          a__     : Tuple  =encodings["input_ids"]
          return inputs







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        a__ ,							a__ ,							a__     : Any  =sequences
        a__     : Union[str, Any]  =char_preds.size(0							)

        a__ ,							a__     : Dict  =self._decode_helper(lowerCAmelCase__		,				"char"							)
        a__ ,							a__     : List[Any]  =self._decode_helper(lowerCAmelCase__		,				"bpe"							)
        a__ ,							a__     : Optional[int]  =self._decode_helper(lowerCAmelCase__		,				"wp"							)

        a__     : List[Any]  =[]
        a__     : Dict  =[]
        for i in range(lowerCAmelCase__							):
          a__     : int  =[char_scores[i], bpe_scores[i], wp_scores[i]]
          a__     : Tuple  =[char_strs[i], bpe_strs[i], wp_strs[i]]
          a__     : Any  =scores.index(max(lowerCAmelCase__							)							)
          final_strs.append(strs[max_score_index]							)
          final_scores.append(scores[max_score_index]							)

        a__     : Dict  ={}
        a__     : str  =final_strs
        a__     : Optional[int]  =final_scores
        a__     : Union[str, Any]  =char_strs
        a__     : List[str]  =bpe_strs
        a__     : Union[str, Any]  =wp_strs
        return out







      def        _lowercase  (   self		,				lowerCAmelCase__		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        if format == DecodeType.CHARACTER:
          a__     : Optional[Any]  =self.char_decode
          a__     : Dict  =1
          a__     : Tuple  ="[s]"
        elif format == DecodeType.BPE:
          a__     : str  =self.bpe_decode
          a__     : Dict  =2
          a__     : Optional[int]  ="#"
        elif format == DecodeType.WORDPIECE:
          a__     : Union[str, Any]  =self.wp_decode
          a__     : List[Any]  =1_0_2
          a__     : Dict  ="[SEP]"
        else:
          raise ValueError(F'''Format {format} is not supported.'''							)

        a__ ,							a__     : Any  =[], []
        a__     : str  =pred_logits.size(0							)
        a__     : Optional[Any]  =pred_logits.size(1							)
        a__ ,							a__     : Optional[int]  =pred_logits.topk(1		,				dim=-1		,				largest=lowerCAmelCase__		,				sorted=lowerCAmelCase__							)
        a__     : Optional[Any]  =preds_index.view(-1		,				lowerCAmelCase__							)[:, 1:]
        a__     : Dict  =decoder(lowerCAmelCase__							)
        a__ ,							a__     : Any  =torch.nn.functional.softmax(lowerCAmelCase__		,				dim=2							).max(dim=2							)
        a__     : int  =preds_max_prob[:, 1:]

        for index in range(lowerCAmelCase__							):
          a__     : Optional[Any]  =preds_str[index].find(lowerCAmelCase__							)
          a__     : Optional[int]  =preds_str[index][:pred_eos]
          a__     : List[Any]  =preds_index[index].cpu().tolist()
          a__     : List[Any]  =pred_index.index(lowerCAmelCase__							) if eos_token in pred_index else -1
          a__     : Union[str, Any]  =preds_max_prob[index][: pred_eos_index + 1]
          a__     : List[Any]  =pred_max_prob.cumprod(dim=0							)[-1] if pred_max_prob.nelement() != 0 else 0.0
          dec_strs.append(lowerCAmelCase__							)
          conf_scores.append(lowerCAmelCase__							)

        return dec_strs, conf_scores







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						str:



        '''simple docstring'''




        a__     : int  =[seq.replace(" "		,				""							) for seq in self.char_tokenizer.batch_decode(lowerCAmelCase__							)]
        return decode_strs







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Tuple:



        '''simple docstring'''




        return self.bpe_tokenizer.batch_decode(lowerCAmelCase__							)







      def        _lowercase  (   self		,				lowerCAmelCase__							)       ->						Optional[Any]:



        '''simple docstring'''




        a__     : List[Any]  =[seq.replace(" "		,				""							) for seq in self.wp_tokenizer.batch_decode(lowerCAmelCase__							)]
        return decode_strs






