





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					,			snake_case					,			)				->  int:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





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

      a__	:						Optional[Any]      		=  None
      if self.use_input_mask:
         a__	:						Optional[int]      		=  random_attention_mask([self.batch_size, self.seq_length]      )

      a__	:						Tuple      		=  None
      a__	:						Union[str, Any]      		=  None
      a__	:						List[str]      		=  None
      if self.use_labels:
         a__	:						List[Any]      		=  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__	:						int      		=  ids_tensor([self.batch_size]					,			self.num_choices      )

      a__	:						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 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





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

      a__	:						str      		=  True
      a__	:						int      		=  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 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Dict:
      """simple docstring"""





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

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

      a__	:						Optional[int]      		=  model(lowerCamelCase_      )

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





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			)				->  str:
      """simple docstring"""





      a__	:						Dict      		=  True

      a__	:						Tuple      		=  TFEsmModel(config=lowerCamelCase_      )
      a__	:						Dict      		=  {
          '''input_ids''': input_ids,
          '''attention_mask''': input_mask,
          '''encoder_hidden_states''': encoder_hidden_states,
          '''encoder_attention_mask''': encoder_attention_mask,
      }
      a__	:						Any      		=  model(lowerCamelCase_      )

      a__	:						List[str]      		=  [input_ids, input_mask]
      a__	:						Any      		=  model(lowerCamelCase_					,			encoder_hidden_states=lowerCamelCase_      )

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

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





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  Any:
      """simple docstring"""





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





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case					,			snake_case      )				->  int:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Union[str, Any]:
      """simple docstring"""





      a__	:						Any      		=  self.prepare_config_and_inputs()
      (
          a__					
      )	:						List[Any]      		=  config_and_inputs
      a__	:						List[Any]      		=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
      return config, inputs_dict







@require_tf
class   __lowerCAmelCase			(					_UpperCAmelCase		,_UpperCAmelCase		,unittest.TestCase  ):
   _UpperCamelCase :						int     		=				(
       (
           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 :						List[str]     		=				False
   _UpperCamelCase :						Optional[Any]     		=				False





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						List[str]      		=  TFEsmModelTester(self      )
      a__	:						Any      		=  ConfigTester(self					,			config_class=lowerCamelCase_					,			hidden_size=37      )





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      self.config_tester.run_common_tests()





   def 	_snake_case ( self      )				->  Optional[int]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCamelCase_      )





   def 	_snake_case ( self      )				->  List[str]:
      """simple docstring"""





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





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						List[Any]      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_masked_lm(*lowerCamelCase_      )





   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						str      		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_token_classification(*lowerCamelCase_      )





   @slow
   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





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





   @unittest.skip("Protein models do not support embedding resizing."      )
   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      pass





   @unittest.skip("Protein models do not support embedding resizing."      )
   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      pass





   def 	_snake_case ( self      )				->  List[Any]:
      """simple docstring"""





      a__	:						Tuple      		=  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__	:						Dict      		=  model.get_bias()
            assert isinstance(lowerCamelCase_					,			lowerCamelCase_      )
            for k, v in name.items():
               assert isinstance(lowerCamelCase_					,			tf.Variable      )
         else:
            a__	:						List[Any]      		=  model.get_output_embeddings()
            assert x is None
            a__	:						str      		=  model.get_bias()
            assert name is None







@require_tf
class   __lowerCAmelCase			(					unittest.TestCase  ):





   @slow
   def 	_snake_case ( self      )				->  Optional[Any]:
      """simple docstring"""





      a__	:						int      		=  TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"      )

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





      a__	:						Tuple      		=  TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"      )

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










import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_lxmert import LxmertTokenizer


SCREAMING_SNAKE_CASE__       :			Optional[Any]			      =						{"""vocab_file""": """vocab.txt""", """tokenizer_file""": """tokenizer.json"""}

SCREAMING_SNAKE_CASE__       :			List[Any]			      =						{
    """vocab_file""": {
        """unc-nlp/lxmert-base-uncased""": """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/vocab.txt""",
    },
    """tokenizer_file""": {
        """unc-nlp/lxmert-base-uncased""": (
            """https://huggingface.co/unc-nlp/lxmert-base-uncased/resolve/main/tokenizer.json"""
        ),
    },
}

SCREAMING_SNAKE_CASE__       :			Any			      =						{
    """unc-nlp/lxmert-base-uncased""": 5_1_2,
}

SCREAMING_SNAKE_CASE__       :			Optional[int]			      =						{
    """unc-nlp/lxmert-base-uncased""": {"""do_lower_case""": True},
}







class   __lowerCAmelCase			(					_UpperCamelCase  ):
   _UpperCamelCase :						List[str]     		=				VOCAB_FILES_NAMES
   _UpperCamelCase :						str     		=				PRETRAINED_VOCAB_FILES_MAP
   _UpperCamelCase :						List[str]     		=				PRETRAINED_INIT_CONFIGURATION
   _UpperCamelCase :						Optional[int]     		=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   _UpperCamelCase :						Optional[Any]     		=				LxmertTokenizer





   def __init__( self					,			snake_case=None					,			snake_case=None					,			snake_case=True					,			snake_case="[UNK]"					,			snake_case="[SEP]"					,			snake_case="[PAD]"					,			snake_case="[CLS]"					,			snake_case="[MASK]"					,			snake_case=True					,			snake_case=None					,			**snake_case					,			)				->  Any:
      """simple docstring"""





      super().__init__(
          snake_case					,			tokenizer_file=snake_case					,			do_lower_case=snake_case					,			unk_token=snake_case					,			sep_token=snake_case					,			pad_token=snake_case					,			cls_token=snake_case					,			mask_token=snake_case					,			tokenize_chinese_chars=snake_case					,			strip_accents=snake_case					,			**snake_case					,			)

      a__	:						Optional[Any]      		=  json.loads(self.backend_tokenizer.normalizer.__getstate__()      )
      if (
          normalizer_state.get("lowercase"					,			snake_case      ) != do_lower_case
          or normalizer_state.get("strip_accents"					,			snake_case      ) != strip_accents
          or normalizer_state.get("handle_chinese_chars"					,			snake_case      ) != tokenize_chinese_chars
      ):
         a__	:						str      		=  getattr(snake_case					,			normalizer_state.pop("type"      )      )
         a__	:						Tuple      		=  do_lower_case
         a__	:						Union[str, Any]      		=  strip_accents
         a__	:						str      		=  tokenize_chinese_chars
         a__	:						List[str]      		=  normalizer_class(**snake_case      )

      a__	:						str      		=  do_lower_case





   def 	_snake_case ( self					,			snake_case					,			snake_case=None      )				->  List[str]:
      """simple docstring"""





      a__	:						Optional[int]      		=  [self.cls_token_id] + token_ids_a + [self.sep_token_id]

      if token_ids_a:
         output += token_ids_a + [self.sep_token_id]

      return output





   def 	_snake_case ( self					,			snake_case					,			snake_case = None      )				->  List[int]:
      """simple docstring"""





      a__	:						Tuple      		=  [self.sep_token_id]
      a__	:						str      		=  [self.cls_token_id]
      if token_ids_a is None:
         return len(cls + token_ids_a + sep      ) * [0]
      return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]





   def 	_snake_case ( self					,			snake_case					,			snake_case = None      )				->  Tuple[str]:
      """simple docstring"""





      a__	:						Optional[int]      		=  self._tokenizer.model.save(snake_case					,			name=snake_case      )
      return tuple(snake_case      )

