


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


  '''simple docstring'''







  def __init__(							self :       str	,_a :       Dict	,):

        '''simple docstring'''




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






  def        _a							(							self :       str ):

        '''simple docstring'''




        A_						:  List[Any]          =				ids_tensor([self.batch_size, self.seq_length]	,self.vocab_size )

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

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

        A_						:  Any          =				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        _a							(							self :       List[str] ):

        '''simple docstring'''




        (
            (
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					
        )						:  str          =				self.prepare_config_and_inputs()

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

        '''simple docstring'''




        A_						:  str          =				TFEsmModel(config=_a )
        A_						:  Union[str, Any]          =				{"""input_ids""": input_ids, """attention_mask""": input_mask}
        A_						:  Tuple          =				model(_a )

        A_						:  Union[str, Any]          =				[input_ids, input_mask]
        A_						:  Dict          =				model(_a )

        A_						:  List[str]          =				model(_a )

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






  def        _a							(							self :       List[Any]	,_a :       str	,_a :       int	,_a :       List[str]	,_a :       Dict	,_a :       List[str]	,_a :       Optional[int]	,_a :       Any	,_a :       int	,):

        '''simple docstring'''




        A_						:  Dict          =				True

        A_						:  Tuple          =				TFEsmModel(config=_a )
        A_						:  Dict          =				{
            """input_ids""": input_ids,
            """attention_mask""": input_mask,
            """encoder_hidden_states""": encoder_hidden_states,
            """encoder_attention_mask""": encoder_attention_mask,
        }
        A_						:  Tuple          =				model(_a )

        A_						:  Union[str, Any]          =				[input_ids, input_mask]
        A_						:  Tuple          =				model(_a	,encoder_hidden_states=_a )

        # Also check the case where encoder outputs are not passed
        A_						:  Union[str, Any]          =				model(_a	,attention_mask=_a )

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






  def        _a							(							self :       str	,_a :       List[str]	,_a :       List[Any]	,_a :       Any	,_a :       List[Any]	,_a :       Dict	,_a :       Any ):

        '''simple docstring'''




        A_						:  Optional[Any]          =				TFEsmForMaskedLM(config=_a )
        A_						:  List[str]          =				model([input_ids, input_mask] )
        self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.vocab_size) )






  def        _a							(							self :       Optional[int]	,_a :       str	,_a :       int	,_a :       List[Any]	,_a :       Union[str, Any]	,_a :       Any	,_a :       List[Any] ):

        '''simple docstring'''




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






  def        _a							(							self :       str ):

        '''simple docstring'''




        A_						:  str          =				self.prepare_config_and_inputs()
        (
            (
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					(
            A_
        )   ,					
        )						:  Dict          =				config_and_inputs
        A_						:  List[Any]          =				{"""input_ids""": input_ids, """attention_mask""": input_mask}
        return config, inputs_dict







@require_tf
class 	__lowerCAmelCase		(							__SCREAMING_SNAKE_CASE  , __SCREAMING_SNAKE_CASE  , unittest.TestCase    ):


  '''simple docstring'''







  a_  					=							(
      (
          TFEsmModel,
          TFEsmForMaskedLM,
          TFEsmForSequenceClassification,
          TFEsmForTokenClassification,
      )
      if is_tf_available()
      else ()
  )
  a_  					=							(
      {
          """feature-extraction""": TFEsmModel,
          """fill-mask""": TFEsmForMaskedLM,
          """text-classification""": TFEsmForSequenceClassification,
          """token-classification""": TFEsmForTokenClassification,
          """zero-shot""": TFEsmForSequenceClassification,
      }
      if is_tf_available()
      else {}
  )
  a_  					=							False
  a_  					=							False






  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        A_						:  Tuple          =				TFEsmModelTester(self )
        A_						:  Union[str, Any]          =				ConfigTester(self	,config_class=_a	,hidden_size=37 )






  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        self.config_tester.run_common_tests()






  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        A_						:  Union[str, Any]          =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*_a )






  def        _a							(							self :       Any ):

        '''simple docstring'''




        A_						:  int          =				self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*_a )






  def        _a							(							self :       Optional[Any] ):

        '''simple docstring'''




        A_						:  Tuple          =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*_a )






  def        _a							(							self :       Any ):

        '''simple docstring'''




        A_						:  Optional[Any]          =				self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*_a )






  @slow
  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              A_						:  Union[str, Any]          =				TFEsmModel.from_pretrained(_a )
              self.assertIsNotNone(_a )






  @unittest.skip("""Protein models do not support embedding resizing.""" )
  def        _a							(							self :       str ):

        '''simple docstring'''




        pass






  @unittest.skip("""Protein models do not support embedding resizing.""" )
  def        _a							(							self :       Optional[int] ):

        '''simple docstring'''




        pass






  def        _a							(							self :       Union[str, Any] ):

        '''simple docstring'''




        A_   ,					A_						:  List[Any]          =				self.model_tester.prepare_config_and_inputs_for_common()

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







@require_tf
class 	__lowerCAmelCase		(							unittest.TestCase    ):


  '''simple docstring'''







  @slow
  def        _a							(							self :       Optional[int] ):

        '''simple docstring'''




        A_						:  List[str]          =				TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D""" )

        A_						:  Dict          =				tf.constant([[0, 1, 2, 3, 4, 5]] )
        A_						:  int          =				model(_a )[0]
        A_						:  int          =				[1, 6, 33]
        self.assertEqual(list(output.numpy().shape )	,_a )
        # compare the actual values for a slice.
        A_						:  List[Any]          =				tf.constant(
            [
                [
                    [8.921518, -10.589814, -6.4671307],
                    [-6.3967156, -13.911377, -1.1211915],
                    [-7.781247, -13.951557, -3.740592],
                ]
            ] )
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()	,expected_slice.numpy()	,atol=1e-2 ) )






  @slow
  def        _a							(							self :       List[str] ):

        '''simple docstring'''




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

        A_						:  str          =				tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]] )
        A_						:  Optional[int]          =				model(_a )[0]
        # compare the actual values for a slice.
        A_						:  Union[str, Any]          =				tf.constant(
            [
                [
                    [0.14443092, 0.54125327, 0.3247739],
                    [0.30340484, 0.00526676, 0.31077722],
                    [0.32278043, -0.24987096, 0.3414628],
                ]
            ] )
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()	,expected_slice.numpy()	,atol=1e-4 ) )




'''simple docstring'''


import warnings
from typing import List, Optional, Union

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType



class 	__lowerCAmelCase		(							__SCREAMING_SNAKE_CASE    ):


  '''simple docstring'''







  a_  					=							["""image_processor""", """tokenizer"""]
  a_  					=							"""ViltImageProcessor"""
  a_  					=							("""BertTokenizer""", """BertTokenizerFast""")






  def __init__(							self :       List[Any]	,_a :       Optional[Any]=None	,_a :       List[str]=None	,**_a :       Any ):

        '''simple docstring'''




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

        A_						:  List[Any]          =				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`.""" )

        super().__init__(_a	,_a )
        A_						:  Optional[Any]          =				self.image_processor






  def __call__(							self :       Any	,_a :       Tuple	,_a :       Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None	,_a :       bool = True	,_a :       Union[bool, str, PaddingStrategy] = False	,_a :       Union[bool, str, TruncationStrategy] = None	,_a :       Optional[int] = None	,_a :       int = 0	,_a :       Optional[int] = None	,_a :       Optional[bool] = None	,_a :       Optional[bool] = None	,_a :       bool = False	,_a :       bool = False	,_a :       bool = False	,_a :       bool = False	,_a :       bool = True	,_a :       Optional[Union[str, TensorType]] = None	,**_a :       Tuple	,):

        '''simple docstring'''




        A_						:  int          =				self.tokenizer(
            text=_a	,add_special_tokens=_a	,padding=_a	,truncation=_a	,max_length=_a	,stride=_a	,pad_to_multiple_of=_a	,return_token_type_ids=_a	,return_attention_mask=_a	,return_overflowing_tokens=_a	,return_special_tokens_mask=_a	,return_offsets_mapping=_a	,return_length=_a	,verbose=_a	,return_tensors=_a	,**_a	,)
        # add pixel_values + pixel_mask
        A_						:  Optional[int]          =				self.image_processor(_a	,return_tensors=_a )
        encoding.update(_a )

        return encoding






  def        _a							(							self :       List[Any]	,*_a :       Any	,**_a :       Any ):

        '''simple docstring'''




        return self.tokenizer.batch_decode(*_a	,**_a )






  def        _a							(							self :       int	,*_a :       int	,**_a :       Optional[int] ):

        '''simple docstring'''




        return self.tokenizer.decode(*_a	,**_a )






  @property
  def        _a							(							self :       List[Any] ):

        '''simple docstring'''




        A_						:  Optional[int]          =				self.tokenizer.model_input_names
        A_						:  str          =				self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names ) )






  @property
  def        _a							(							self :       str ):

        '''simple docstring'''




        warnings.warn(
            """`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead."""	,_a	,)
        return self.image_processor_class






  @property
  def        _a							(							self :       int ):

        '''simple docstring'''




        warnings.warn(
            """`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead."""	,_a	,)
        return self.image_processor
