

"""simple docstring"""


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     :

				def __init__(     self							:    Any					,  UpperCamelCase__							:    Any					,  UpperCamelCase__							:    Union[str, Any]=1_3					,  UpperCamelCase__							:    Tuple=7					,  UpperCamelCase__							:    List[str]=False					,  UpperCamelCase__							:    Optional[int]=True					,  UpperCamelCase__							:    Optional[int]=False					,  UpperCamelCase__							:    Any=True					,  UpperCamelCase__							:    Any=3_3					,  UpperCamelCase__							:    List[str]=3_2					,  UpperCamelCase__							:    Dict=5					,  UpperCamelCase__							:    Any=4					,  UpperCamelCase__							:    Union[str, Any]=3_7					,  UpperCamelCase__							:    str="gelu"					,  UpperCamelCase__							:    str=0.1					,  UpperCamelCase__							:    Union[str, Any]=0.1					,  UpperCamelCase__							:    Optional[Any]=5_1_2					,  UpperCamelCase__							:    Optional[int]=1_6					,  UpperCamelCase__							:    Optional[int]=2					,  UpperCamelCase__							:    Tuple=0.02					,  UpperCamelCase__							:    Any=3					,  UpperCamelCase__							:    Dict=4					,  UpperCamelCase__							:    int=None					,  )->       Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[Any]									=      parent
						__lowerCAmelCase:	str									=      batch_size
						__lowerCAmelCase:	str									=      seq_length
						__lowerCAmelCase:	List[str]									=      is_training
						__lowerCAmelCase:	List[str]									=      use_input_mask
						__lowerCAmelCase:	Tuple									=      use_token_type_ids
						__lowerCAmelCase:	str									=      use_labels
						__lowerCAmelCase:	List[Any]									=      vocab_size
						__lowerCAmelCase:	Any									=      hidden_size
						__lowerCAmelCase:	List[str]									=      num_hidden_layers
						__lowerCAmelCase:	Optional[int]									=      num_attention_heads
						__lowerCAmelCase:	Union[str, Any]									=      intermediate_size
						__lowerCAmelCase:	Optional[int]									=      hidden_act
						__lowerCAmelCase:	List[Any]									=      hidden_dropout_prob
						__lowerCAmelCase:	str									=      attention_probs_dropout_prob
						__lowerCAmelCase:	Tuple									=      max_position_embeddings
						__lowerCAmelCase:	Union[str, Any]									=      type_vocab_size
						__lowerCAmelCase:	List[str]									=      type_sequence_label_size
						__lowerCAmelCase:	List[Any]									=      initializer_range
						__lowerCAmelCase:	str									=      num_labels
						__lowerCAmelCase:	Dict									=      num_choices
						__lowerCAmelCase:	str									=      scope

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



						'''simple docstring'''


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

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

						__lowerCAmelCase:	str									=      None
						__lowerCAmelCase:	Dict									=      None
						__lowerCAmelCase:	Union[str, Any]									=      None
						if self.use_labels:
								__lowerCAmelCase:	Tuple									=      ids_tensor([self.batch_size]					,  self.type_sequence_label_size)
								__lowerCAmelCase:	Dict									=      ids_tensor([self.batch_size, self.seq_length]					,  self.num_labels)
								__lowerCAmelCase:	List[str]									=      ids_tensor([self.batch_size]					,  self.num_choices)

						__lowerCAmelCase:	List[str]									=      self.get_config()

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

				def 							lowercase_					(     self							:    Dict)->       List[str]:



						'''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							:    str					,  UpperCamelCase__							:    List[Any]					,  UpperCamelCase__							:    Dict					,  UpperCamelCase__							:    int					,  UpperCamelCase__							:    Optional[int]					,  UpperCamelCase__							:    List[Any]					,  UpperCamelCase__							:    Union[str, Any])->       Optional[Any]:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[Any]									=      EsmModel(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:	Optional[int]									=      model(UpperCamelCase__					,  attention_mask=UpperCamelCase__)
						__lowerCAmelCase:	List[str]									=      model(UpperCamelCase__)
						__lowerCAmelCase:	Dict									=      model(UpperCamelCase__)

						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							:    int					,  UpperCamelCase__							:    Dict					,  UpperCamelCase__							:    List[str]					,  UpperCamelCase__							:    Optional[Any]					,  UpperCamelCase__							:    int					,  UpperCamelCase__							:    Any					,  UpperCamelCase__							:    Union[str, Any])->       Any:



						'''simple docstring'''


						__lowerCAmelCase:	Tuple									=      EsmForMaskedLM(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:	Optional[int]									=      model(UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.vocab_size))

				def 							lowercase_					(     self							:    str					,  UpperCamelCase__							:    Optional[Any]					,  UpperCamelCase__							:    List[str]					,  UpperCamelCase__							:    int					,  UpperCamelCase__							:    Any					,  UpperCamelCase__							:    List[Any]					,  UpperCamelCase__							:    List[Any])->       List[str]:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[int]									=      self.num_labels
						__lowerCAmelCase:	Tuple									=      EsmForTokenClassification(config=UpperCamelCase__)
						model.to(UpperCamelCase__)
						model.eval()
						__lowerCAmelCase:	int									=      model(UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  labels=UpperCamelCase__)
						self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.num_labels))




				def 							lowercase_					(     self							:    List[str])->       Dict:



						'''simple docstring'''


						__lowerCAmelCase:	Tuple									=      self.prepare_config_and_inputs()
						(
						    __lowerCAmelCase			
						):	int									=      config_and_inputs
						__lowerCAmelCase:	Dict									=      {"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class    snake_case     (    __snake_case,							__snake_case,							unittest.TestCase  ):
				SCREAMING_SNAKE_CASE_       :				List[Any]       				=  False

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

				def 							lowercase_					(     self							:    List[str])->       Optional[int]:



						'''simple docstring'''


						__lowerCAmelCase:	str									=      EsmModelTester(self)
						__lowerCAmelCase:	str									=      ConfigTester(self					,  config_class=UpperCamelCase__					,  hidden_size=3_7)

				def 							lowercase_					(     self							:    Tuple)->       Tuple:



						'''simple docstring'''


						self.config_tester.run_common_tests()

				def 							lowercase_					(     self							:    int)->       Any:



						'''simple docstring'''


						__lowerCAmelCase:	List[Any]									=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_					(     self							:    Optional[int])->       Dict:



						'''simple docstring'''


						__lowerCAmelCase:	Any									=      self.model_tester.prepare_config_and_inputs()
						for type in ["absolute", "relative_key", "relative_key_query"]:
								__lowerCAmelCase:	Dict									=      type
								self.model_tester.create_and_check_model(*UpperCamelCase__)

				def 							lowercase_					(     self							:    str)->       Tuple:



						'''simple docstring'''


						__lowerCAmelCase:	Any									=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__)

				def 							lowercase_					(     self							:    List[Any])->       str:



						'''simple docstring'''


						__lowerCAmelCase:	Tuple									=      self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__)

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



						'''simple docstring'''


						for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
								__lowerCAmelCase:	List[str]									=      EsmModel.from_pretrained(UpperCamelCase__)
								self.assertIsNotNone(UpperCamelCase__)

				def 							lowercase_					(     self							:    Any)->       int:



						'''simple docstring'''


						__lowerCAmelCase:	Optional[Any]									=      self.model_tester.prepare_config_and_inputs()[0]
						__lowerCAmelCase:	str									=      EsmEmbeddings(config=UpperCamelCase__)

						__lowerCAmelCase:	Union[str, Any]									=      torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]])
						__lowerCAmelCase:	int									=      torch.as_tensor(
						    [
						        [
						            0 + model.padding_idx + 1,
						            1 + model.padding_idx + 1,
						            2 + model.padding_idx + 1,
						            model.padding_idx,
						        ]
						    ])
						__lowerCAmelCase:	Tuple									=      create_position_ids_from_input_ids(UpperCamelCase__					,  model.padding_idx)
						self.assertEqual(position_ids.shape					,  expected_positions.shape)
						self.assertTrue(torch.all(torch.eq(UpperCamelCase__					,  UpperCamelCase__)))

				def 							lowercase_					(     self							:    str)->       int:



						'''simple docstring'''


						__lowerCAmelCase:	List[str]									=      self.model_tester.prepare_config_and_inputs()[0]
						__lowerCAmelCase:	str									=      EsmEmbeddings(config=UpperCamelCase__)

						__lowerCAmelCase:	Optional[int]									=      torch.empty(2					,  4					,  3_0)
						__lowerCAmelCase:	Any									=      [
						    0 + embeddings.padding_idx + 1,
						    1 + embeddings.padding_idx + 1,
						    2 + embeddings.padding_idx + 1,
						    3 + embeddings.padding_idx + 1,
						]
						__lowerCAmelCase:	Any									=      torch.as_tensor([expected_single_positions, expected_single_positions])
						__lowerCAmelCase:	Tuple									=      embeddings.create_position_ids_from_inputs_embeds(UpperCamelCase__)
						self.assertEqual(position_ids.shape					,  expected_positions.shape)
						self.assertTrue(torch.all(torch.eq(UpperCamelCase__					,  UpperCamelCase__)))

				@unittest.skip("Esm does not support embedding resizing")
				def 							lowercase_					(     self							:    str)->       Optional[int]:



						'''simple docstring'''


						pass

				@unittest.skip("Esm does not support embedding resizing")
				def 							lowercase_					(     self							:    Optional[int])->       Optional[Any]:



						'''simple docstring'''


						pass




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



						'''simple docstring'''


						pass





@require_torch
class    snake_case     (    __snake_case  ):

				@slow
				def 							lowercase_					(     self							:    str)->       Union[str, Any]:



						'''simple docstring'''


						with torch.no_grad():
								__lowerCAmelCase:	Any									=      EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D")
								model.eval()
								__lowerCAmelCase:	Any									=      torch.tensor([[0, 1, 2, 3, 4, 5]])
								__lowerCAmelCase:	Optional[int]									=      model(UpperCamelCase__)[0]

								__lowerCAmelCase:	str									=      3_3

								__lowerCAmelCase:	Any									=      torch.Size((1, 6, vocab_size))
								self.assertEqual(output.shape					,  UpperCamelCase__)

								__lowerCAmelCase:	Optional[int]									=      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]					,  UpperCamelCase__					,  atol=1e-4))




				@slow
				def 							lowercase_					(     self							:    Any)->       Tuple:



						'''simple docstring'''


						with torch.no_grad():
								__lowerCAmelCase:	str									=      EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D")
								model.eval()

								__lowerCAmelCase:	List[str]									=      torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]])
								__lowerCAmelCase:	List[Any]									=      model(UpperCamelCase__)[0]
								# compare the actual values for a slice.
								__lowerCAmelCase:	Dict									=      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]					,  UpperCamelCase__					,  atol=1e-4))





"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_torch_available,
)


__A						       =      {
    "configuration_speecht5": [
        "SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "SPEECHT5_PRETRAINED_HIFIGAN_CONFIG_ARCHIVE_MAP",
        "SpeechT5Config",
        "SpeechT5HifiGanConfig",
    ],
    "feature_extraction_speecht5": ["SpeechT5FeatureExtractor"],
    "processing_speecht5": ["SpeechT5Processor"],
}

try:
			if not is_sentencepiece_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			__A						       =      ["SpeechT5Tokenizer"]

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			__A						       =      [
			    "SPEECHT5_PRETRAINED_MODEL_ARCHIVE_LIST",
			    "SpeechT5ForSpeechToText",
			    "SpeechT5ForSpeechToSpeech",
			    "SpeechT5ForTextToSpeech",
			    "SpeechT5Model",
			    "SpeechT5PreTrainedModel",
			    "SpeechT5HifiGan",
			]

if TYPE_CHECKING:
			from .configuration_speechta import (
			    SPEECHT5_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    SPEECHT5_PRETRAINED_HIFIGAN_CONFIG_ARCHIVE_MAP,
			    SpeechTaConfig,
			    SpeechTaHifiGanConfig,
			)
			from .feature_extraction_speechta import SpeechTaFeatureExtractor
			from .processing_speechta import SpeechTaProcessor

			try:
						if not is_sentencepiece_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .tokenization_speechta import SpeechTaTokenizer

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_speechta import (
						    SPEECHT5_PRETRAINED_MODEL_ARCHIVE_LIST,
						    SpeechTaForSpeechToSpeech,
						    SpeechTaForSpeechToText,
						    SpeechTaForTextToSpeech,
						    SpeechTaHifiGan,
						    SpeechTaModel,
						    SpeechTaPreTrainedModel,
						)

else:
			import sys

			__A						       =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


