





import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
  from transformers.models.esm.modeling_esm import (
      ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
      EsmEmbeddings,
      create_position_ids_from_input_ids,
  )



class 					__snake_case	:


    '''simple docstring'''







    def __init__(       self   ,				A_   ,				A_=13   ,				A_=7   ,				A_=False   ,				A_=True   ,				A_=False   ,				A_=True   ,				A_=33   ,				A_=32   ,				A_=5   ,				A_=4   ,				A_=37   ,				A_="gelu"   ,				A_=0.1   ,				A_=0.1   ,				A_=5_12   ,				A_=16   ,				A_=2   ,				A_=0.02   ,				A_=3   ,				A_=4   ,				A_=None   ,				):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			parent
        SCREAMING_SNAKE_CASE__											=			batch_size
        SCREAMING_SNAKE_CASE__											=			seq_length
        SCREAMING_SNAKE_CASE__											=			is_training
        SCREAMING_SNAKE_CASE__											=			use_input_mask
        SCREAMING_SNAKE_CASE__											=			use_token_type_ids
        SCREAMING_SNAKE_CASE__											=			use_labels
        SCREAMING_SNAKE_CASE__											=			vocab_size
        SCREAMING_SNAKE_CASE__											=			hidden_size
        SCREAMING_SNAKE_CASE__											=			num_hidden_layers
        SCREAMING_SNAKE_CASE__											=			num_attention_heads
        SCREAMING_SNAKE_CASE__											=			intermediate_size
        SCREAMING_SNAKE_CASE__											=			hidden_act
        SCREAMING_SNAKE_CASE__											=			hidden_dropout_prob
        SCREAMING_SNAKE_CASE__											=			attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE__											=			max_position_embeddings
        SCREAMING_SNAKE_CASE__											=			type_vocab_size
        SCREAMING_SNAKE_CASE__											=			type_sequence_label_size
        SCREAMING_SNAKE_CASE__											=			initializer_range
        SCREAMING_SNAKE_CASE__											=			num_labels
        SCREAMING_SNAKE_CASE__											=			num_choices
        SCREAMING_SNAKE_CASE__											=			scope



    def  lowercase_					(       self     ):





        '''simple docstring'''

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

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

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

        SCREAMING_SNAKE_CASE__											=			self.get_config()

        return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



    def  lowercase_					(       self     ):





        '''simple docstring'''

        return EsmConfig(
            vocab_size=self.vocab_size   ,				hidden_size=self.hidden_size   ,				pad_token_id=1   ,				num_hidden_layers=self.num_hidden_layers   ,				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   ,				)



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			EsmModel(config=A_     )
        model.to(A_     )
        model.eval()
        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_     )
        SCREAMING_SNAKE_CASE__											=			model(A_     )
        SCREAMING_SNAKE_CASE__											=			model(A_     )

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



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			EsmForMaskedLM(config=A_     )
        model.to(A_     )
        model.eval()
        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_   ,				labels=A_     )
        self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.vocab_size)     )



    def  lowercase_					(       self   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_   ,				A_     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.num_labels
        SCREAMING_SNAKE_CASE__											=			EsmForTokenClassification(config=A_     )
        model.to(A_     )
        model.eval()
        SCREAMING_SNAKE_CASE__											=			model(A_   ,				attention_mask=A_   ,				labels=A_     )
        self.parent.assertEqual(result.logits.shape   ,				(self.batch_size, self.seq_length, self.num_labels)     )




    def  lowercase_					(       self     ):





        '''simple docstring'''

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





@require_torch
class 					__snake_case	(  __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     unittest.TestCase    ):


    '''simple docstring'''







    lowerCamelCase__    :		List[str]				=	False

    lowerCamelCase__    :		Dict				=	(
        (
            EsmForMaskedLM,
            EsmModel,
            EsmForSequenceClassification,
            EsmForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowerCamelCase__    :		Dict				=	()
    lowerCamelCase__    :		Optional[int]				=	(
        {
            """feature-extraction""": EsmModel,
            """fill-mask""": EsmForMaskedLM,
            """text-classification""": EsmForSequenceClassification,
            """token-classification""": EsmForTokenClassification,
            """zero-shot""": EsmForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowerCamelCase__    :		List[Any]				=	True



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			EsmModelTester(self     )
        SCREAMING_SNAKE_CASE__											=			ConfigTester(self   ,				config_class=A_   ,				hidden_size=37     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        self.config_tester.run_common_tests()



    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
            SCREAMING_SNAKE_CASE__											=			type
            self.model_tester.create_and_check_model(*A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    def  lowercase_					(       self     ):





        '''simple docstring'''

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



    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            SCREAMING_SNAKE_CASE__											=			EsmModel.from_pretrained(A_     )
            self.assertIsNotNone(A_     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()[0]
        SCREAMING_SNAKE_CASE__											=			EsmEmbeddings(config=A_     )

        SCREAMING_SNAKE_CASE__											=			torch.as_tensor([[12, 31, 13, model.padding_idx]]     )
        SCREAMING_SNAKE_CASE__											=			torch.as_tensor(
            [
                [
                    0 + model.padding_idx + 1,
                    1 + model.padding_idx + 1,
                    2 + model.padding_idx + 1,
                    model.padding_idx,
                ]
            ]     )
        SCREAMING_SNAKE_CASE__											=			create_position_ids_from_input_ids(A_   ,				model.padding_idx     )
        self.assertEqual(position_ids.shape   ,				expected_positions.shape     )
        self.assertTrue(torch.all(torch.eq(A_   ,				A_     )     )     )



    def  lowercase_					(       self     ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__											=			self.model_tester.prepare_config_and_inputs()[0]
        SCREAMING_SNAKE_CASE__											=			EsmEmbeddings(config=A_     )

        SCREAMING_SNAKE_CASE__											=			torch.empty(2   ,				4   ,				30     )
        SCREAMING_SNAKE_CASE__											=			[
            0 + embeddings.padding_idx + 1,
            1 + embeddings.padding_idx + 1,
            2 + embeddings.padding_idx + 1,
            3 + embeddings.padding_idx + 1,
        ]
        SCREAMING_SNAKE_CASE__											=			torch.as_tensor([expected_single_positions, expected_single_positions]     )
        SCREAMING_SNAKE_CASE__											=			embeddings.create_position_ids_from_inputs_embeds(A_     )
        self.assertEqual(position_ids.shape   ,				expected_positions.shape     )
        self.assertTrue(torch.all(torch.eq(A_   ,				A_     )     )     )



    @unittest.skip('''Esm does not support embedding resizing'''     )
    def  lowercase_					(       self     ):





        '''simple docstring'''

        pass



    @unittest.skip('''Esm does not support embedding resizing'''     )
    def  lowercase_					(       self     ):





        '''simple docstring'''

        pass




    @unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''     )
    def  lowercase_					(       self     ):





        '''simple docstring'''

        pass





@require_torch
class 					__snake_case	(  __SCREAMING_SNAKE_CASE    ):


    '''simple docstring'''







    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        with torch.no_grad():
            SCREAMING_SNAKE_CASE__											=			EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''     )
            model.eval()
            SCREAMING_SNAKE_CASE__											=			torch.tensor([[0, 1, 2, 3, 4, 5]]     )
            SCREAMING_SNAKE_CASE__											=			model(A_     )[0]

            SCREAMING_SNAKE_CASE__											=			33

            SCREAMING_SNAKE_CASE__											=			torch.Size((1, 6, vocab_size)     )
            self.assertEqual(output.shape   ,				A_     )

            SCREAMING_SNAKE_CASE__											=			torch.tensor(
                [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]     )
            self.assertTrue(torch.allclose(output[:, :3, :3]   ,				A_   ,				atol=1E-4     )     )




    @slow
    def  lowercase_					(       self     ):





        '''simple docstring'''

        with torch.no_grad():
            SCREAMING_SNAKE_CASE__											=			EsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''     )
            model.eval()

            SCREAMING_SNAKE_CASE__											=			torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]     )
            SCREAMING_SNAKE_CASE__											=			model(A_     )[0]
            # compare the actual values for a slice.
            SCREAMING_SNAKE_CASE__											=			torch.tensor(
                [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]     )
            self.assertTrue(torch.allclose(output[:, :3, :3]   ,				A_   ,				atol=1E-4     )     )




from ...processing_utils import ProcessorMixin







class 	__lowercase    (__SCREAMING_SNAKE_CASE  ):





							"""simple docstring"""







							_UpperCAmelCase      =						"""WhisperFeatureExtractor"""
							_UpperCAmelCase      =						"""WhisperTokenizer"""






							def __init__(      self			,						lowerCAmelCase__			,						lowerCAmelCase__       ):





										"""simple docstring"""




										super().__init__(lowerCAmelCase__			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Union[str, Any]					=					self.feature_extractor
										SCREAMING_SNAKE_CASE_		:     List[Any]					=					False






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__=None			,						lowerCAmelCase__=None			,						lowerCAmelCase__=True       ):





										"""simple docstring"""




										return self.tokenizer.get_decoder_prompt_ids(task=lowerCAmelCase__			,						language=lowerCAmelCase__			,						no_timestamps=lowerCAmelCase__       )






							def __call__(      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




										if self._in_target_context_manager:
													return self.current_processor(*lowerCAmelCase__			,						**lowerCAmelCase__       )

										SCREAMING_SNAKE_CASE_		:     Tuple					=					kwargs.pop('audio'			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     str					=					kwargs.pop('sampling_rate'			,						lowerCAmelCase__       )
										SCREAMING_SNAKE_CASE_		:     Optional[Any]					=					kwargs.pop('text'			,						lowerCAmelCase__       )
										if len(lowerCAmelCase__       ) > 0:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					args[0]
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					args[1:]

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

										if audio is not None:
													SCREAMING_SNAKE_CASE_		:     List[Any]					=					self.feature_extractor(lowerCAmelCase__			,						*lowerCAmelCase__			,						sampling_rate=lowerCAmelCase__			,						**lowerCAmelCase__       )
										if text is not None:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					self.tokenizer(lowerCAmelCase__			,						**lowerCAmelCase__       )

										if text is None:
													return inputs

										elif audio is None:
													return encodings
										else:
													SCREAMING_SNAKE_CASE_		:     Optional[int]					=					encodings['input_ids']
													return inputs






							def 					UpperCamelCase__     (      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self			,						*lowerCAmelCase__			,						**lowerCAmelCase__       ):





										"""simple docstring"""




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






							def 					UpperCamelCase__     (      self			,						lowerCAmelCase__			,						lowerCAmelCase__="np"       ):





										"""simple docstring"""




										return self.tokenizer.get_prompt_ids(lowerCAmelCase__			,						return_tensors=lowerCAmelCase__       )





