





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  A__       :





      def __init__(			self	,							A_	,							):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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

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

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

       UpperCamelCase     : List[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 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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

       UpperCamelCase     : Optional[Any]	     =       True
       UpperCamelCase     : Union[str, Any]	     =       floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
       UpperCamelCase     : Tuple	     =       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_       ):
       '''simple docstring'''




       UpperCamelCase     : List[Any]	     =       TFEsmModel(config=A_       )
       UpperCamelCase     : int	     =       {"input_ids": input_ids, "attention_mask": input_mask}
       UpperCamelCase     : Tuple	     =       model(A_       )

       UpperCamelCase     : str	     =       [input_ids, input_mask]
       UpperCamelCase     : Optional[Any]	     =       model(A_       )

       UpperCamelCase     : int	     =       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_	,							):
       '''simple docstring'''




       UpperCamelCase     : Any	     =       True

       UpperCamelCase     : Optional[Any]	     =       TFEsmModel(config=A_       )
       UpperCamelCase     : Any	     =       {
           "input_ids": input_ids,
           "attention_mask": input_mask,
           "encoder_hidden_states": encoder_hidden_states,
           "encoder_attention_mask": encoder_attention_mask,
       }
       UpperCamelCase     : Tuple	     =       model(A_       )

       UpperCamelCase     : Any	     =       [input_ids, input_mask]
       UpperCamelCase     : Optional[int]	     =       model(A_	,							encoder_hidden_states=A_       )

       # Also check the case where encoder outputs are not passed
       UpperCamelCase     : Dict	     =       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_       ):
       '''simple docstring'''




       UpperCamelCase     : Dict	     =       TFEsmForMaskedLM(config=A_       )
       UpperCamelCase     : str	     =       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_       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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



@require_tf
class  A__       (	__snake_case	,       __snake_case	,       unittest.TestCase      ):
      _UpperCAmelCase     :int						   =       (
          (
              TFEsmModel,
              TFEsmForMaskedLM,
              TFEsmForSequenceClassification,
              TFEsmForTokenClassification,
          )
          if is_tf_available()
          else ()
      )
      _UpperCAmelCase     :Union[str, Any]						   =       (
          {
              'feature-extraction': TFEsmModel,
              'fill-mask': TFEsmForMaskedLM,
              'text-classification': TFEsmForSequenceClassification,
              'token-classification': TFEsmForTokenClassification,
              'zero-shot': TFEsmForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      _UpperCAmelCase     :int						   =       False
      _UpperCAmelCase     :str						   =       False





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       self.config_tester.run_common_tests()





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[Any]	     =       self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*A_       )





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : List[Any]	     =       self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*A_       )





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




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





      @unittest.skip("Protein models do not support embedding resizing."       )
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       pass





      @unittest.skip("Protein models do not support embedding resizing."       )
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       pass





      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     ,   UpperCamelCase     : Any	     =       self.model_tester.prepare_config_and_inputs_for_common()

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



@require_tf
class  A__       (	unittest.TestCase      ):





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : Dict	     =       TFEsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"       )

       UpperCamelCase     : str	     =       tf.constant([[0, 1, 2, 3, 4, 5]]       )
       UpperCamelCase     : Dict	     =       model(A_       )[0]
       UpperCamelCase     : Dict	     =       [1, 6, 33]
       self.assertEqual(list(output.numpy().shape       )	,							A_       )
       # compare the actual values for a slice.
       UpperCamelCase     : Dict	     =       tf.constant(
           [
               [
                   [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
                   [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
                   [-7.78_12_47, -13.95_15_57, -3.74_05_92],
               ]
           ]       )
       self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()	,							expected_slice.numpy()	,							atol=1e-2       )       )





      @slow
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       UpperCamelCase     : str	     =       TFEsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"       )

       UpperCamelCase     : Optional[Any]	     =       tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
       UpperCamelCase     : Any	     =       model(A_       )[0]
       # compare the actual values for a slice.
       UpperCamelCase     : str	     =       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 time
import warnings
from abc import ABC
from copy import deepcopy
from typing import Optional

import torch

from ..utils import add_start_docstrings, logging


__lowerCamelCase	:					List[str]			  =					logging.get_logger(__name__)


__lowerCamelCase	:					Any			  =					r"""
    Args:
        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be scores for each vocabulary token before SoftMax
            or scores for each vocabulary token after SoftMax.
        kwargs (`Dict[str, Any]`, *optional*):
            Additional stopping criteria specific kwargs.

    Return:
        `bool`. `False` indicates we should continue, `True` indicates we should stop.

"""



class  A__       (	__snake_case      ):





      @add_start_docstrings(A_       )
      def __call__(			self	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       raise NotImplementedError("StoppingCriteria needs to be subclassed"       )



class  A__       (	__snake_case      ):





      def __init__(			self	,							A_	,							A_ = None       ):
       '''simple docstring'''




       UpperCamelCase     : Union[str, Any]	     =       max_length
       UpperCamelCase     : Dict	     =       max_position_embeddings





      @add_start_docstrings(A_       )
      def __call__(			self	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       UpperCamelCase     : int	     =       input_ids.shape[-1]
       UpperCamelCase     : str	     =       cur_len >= self.max_length
       if self.max_position_embeddings is not None and not is_done and cur_len >= self.max_position_embeddings:
        logger.warning_once(
            "This is a friendly reminder - the current text generation call will exceed the model's predefined "
            F"""maximum length ({self.max_position_embeddings}). Depending on the model, you may observe """
            "exceptions, performance degradation, or nothing at all."       )
       return is_done



class  A__       (	__snake_case      ):





      def __init__(			self	,							A_	,							A_       ):
       '''simple docstring'''




       warnings.warn(
           "The class `MaxNewTokensCriteria` is deprecated. "
           F"""Please use `MaxLengthCriteria(max_length={start_length + max_new_tokens})` """
           "with `max_length = start_length + max_new_tokens` instead."	,							A_	,							)
       UpperCamelCase     : Union[str, Any]	     =       start_length
       UpperCamelCase     : List[str]	     =       max_new_tokens
       UpperCamelCase     : Tuple	     =       start_length + max_new_tokens





      @add_start_docstrings(A_       )
      def __call__(			self	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       return input_ids.shape[-1] >= self.max_length



class  A__       (	__snake_case      ):





      def __init__(			self	,							A_	,							A_ = None       ):
       '''simple docstring'''




       UpperCamelCase     : Optional[int]	     =       max_time
       UpperCamelCase     : Dict	     =       time.time() if initial_timestamp is None else initial_timestamp





      @add_start_docstrings(A_       )
      def __call__(			self	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       return time.time() - self.initial_timestamp > self.max_time



class  A__       (	__snake_case      ):





      @add_start_docstrings(A_       )
      def __call__(			self	,							A_	,							A_	,							**A_       ):
       '''simple docstring'''




       return any(criteria(A_	,							A_       ) for criteria in self       )





      @property
      def 	__UpperCamelCase(			self       ):
       '''simple docstring'''




       for stopping_criterium in self:
        if isinstance(A_	,							A_       ):
         return stopping_criterium.max_length
        elif isinstance(A_	,							A_       ):
         return stopping_criterium.max_length
       return None



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase					)						->			StoppingCriteriaList:
 UpperCamelCase     : Tuple	     =       stopping_criteria.max_length
 UpperCamelCase     : Union[str, Any]	     =       deepcopy(_lowerCAmelCase					)
 if stopping_max_length is not None and stopping_max_length != max_length:
  warnings.warn("You set different `max_length` for stopping criteria and `max_length` parameter"							,       _lowerCAmelCase					)
 elif stopping_max_length is None:
  new_stopping_criteria.append(MaxLengthCriteria(max_length=_lowerCAmelCase					)					)
 return new_stopping_criteria
