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.models.esm.modeling_esmfold import EsmForProteinFolding







class __snake_case       :




   def __init__(					self    :       int ,			_lowercase    :       Tuple ,			_lowercase    :       Tuple=13 ,			_lowercase    :       Tuple=7 ,			_lowercase    :       Union[str, Any]=False ,			_lowercase    :       int=True ,			_lowercase    :       Tuple=False ,			_lowercase    :       int=False ,			_lowercase    :       int=19 ,			_lowercase    :       str=32 ,			_lowercase    :       Optional[Any]=5 ,			_lowercase    :       Union[str, Any]=4 ,			_lowercase    :       Any=37 ,			_lowercase    :       Tuple="gelu" ,			_lowercase    :       Union[str, Any]=0.1 ,			_lowercase    :       Dict=0.1 ,			_lowercase    :       Optional[int]=5_12 ,			_lowercase    :       Union[str, Any]=16 ,			_lowercase    :       Union[str, Any]=2 ,			_lowercase    :       Optional[int]=0.02 ,			_lowercase    :       str=3 ,			_lowercase    :       Optional[Any]=4 ,			_lowercase    :       Optional[Any]=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 			__a		(					self    :       List[str]   ):



     """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 			__a		(					self    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   EsmConfig(
         vocab_size=33 ,			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 ,			is_folding_model=_lowercase ,			esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False} ,			)
     return config




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



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   EsmForProteinFolding(config=_lowercase   ).float()
     model.to(_lowercase   )
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase ,			attention_mask=_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )

     self.parent.assertEqual(result.positions.shape ,			(8, self.batch_size, self.seq_length, 14, 3)   )
     self.parent.assertEqual(result.angles.shape ,			(8, self.batch_size, self.seq_length, 7, 2)   )




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



     """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       (		lowerCamelCase_     ,	lowerCamelCase_     ,	unittest.TestCase	):
   lowerCAmelCase_														=							False

   lowerCAmelCase_														=							(EsmForProteinFolding,) if is_torch_available() else ()
   lowerCAmelCase_														=							()
   lowerCAmelCase_														=							{} if is_torch_available() else {}
   lowerCAmelCase_														=							False




   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   EsmFoldModelTester(self   )
     SCREAMING_SNAKE_CASE__    							=   ConfigTester(self ,			config_class=_lowercase ,			hidden_size=37   )




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     self.config_tester.run_common_tests()




   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






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




   @unittest.skip("""Does not support attention outputs"""   )
   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     pass




   @unittest.skip
   def 			__a		(					self    :       Optional[int]   ):



     """simple docstring"""






     pass




   @unittest.skip("""Esm does not support embedding resizing"""   )
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     pass




   @unittest.skip("""Esm does not support embedding resizing"""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support passing input embeds!"""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support head pruning."""   )
   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support head pruning."""   )
   def 			__a		(					self    :       Union[str, Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support head pruning."""   )
   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support head pruning."""   )
   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support head pruning."""   )
   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not output hidden states in the normal way."""   )
   def 			__a		(					self    :       Tuple   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMfold does not output hidden states in the normal way."""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold only has one output format."""   )
   def 			__a		(					self    :       Optional[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""   )
   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold does not support input chunking."""   )
   def 			__a		(					self    :       Dict   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""   )
   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold doesn't support torchscript compilation."""   )
   def 			__a		(					self    :       List[Any]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold doesn't support torchscript compilation."""   )
   def 			__a		(					self    :       int   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold doesn't support torchscript compilation."""   )
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     pass




   @unittest.skip("""ESMFold doesn't support data parallel."""   )
   def 			__a		(					self    :       str   ):



     """simple docstring"""






     pass




   @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""   )
   def 			__a		(					self    :       List[str]   ):



     """simple docstring"""






     pass







@require_torch
class __snake_case       (		lowerCamelCase_	):




   @slow
   def 			__a		(					self    :       Any   ):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""   ).float()
     model.eval()
     SCREAMING_SNAKE_CASE__    							=   torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
     SCREAMING_SNAKE_CASE__    							=   model(_lowercase   )["""positions"""]
     SCREAMING_SNAKE_CASE__    							=   torch.tensor([2.58_28, 0.79_93, -10.93_34] ,			dtype=torch.floataa   )
     self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0] ,			_lowercase ,			atol=1E-4   )   )

from ...configuration_utils import PretrainedConfig
from ...utils import logging


__lowerCamelCase    :				Optional[int]							     =	logging.get_logger(__name__)

__lowerCamelCase    :				str							     =	{
    '''facebook/nllb-moe-54B''': '''https://huggingface.co/facebook/nllb-moe-54b/resolve/main/config.json''',
}







class __snake_case       (		lowerCamelCase_	):
   lowerCAmelCase_														=							"nllb-moe"
   lowerCAmelCase_														=							["past_key_values"]
   lowerCAmelCase_														=							{"num_attention_heads": "encoder_attention_heads", "hidden_size": "d_model"}




   def __init__(					self    :       int ,			_lowercase    :       Tuple=12_81_12 ,			_lowercase    :       List[Any]=10_24 ,			_lowercase    :       Any=12 ,			_lowercase    :       List[Any]=40_96 ,			_lowercase    :       str=16 ,			_lowercase    :       str=12 ,			_lowercase    :       Optional[int]=40_96 ,			_lowercase    :       List[Any]=16 ,			_lowercase    :       str=0.05 ,			_lowercase    :       Tuple=0.05 ,			_lowercase    :       str=True ,			_lowercase    :       List[str]=True ,			_lowercase    :       Optional[Any]="relu" ,			_lowercase    :       str=10_24 ,			_lowercase    :       Tuple=0.1 ,			_lowercase    :       int=0.1 ,			_lowercase    :       Dict=0.0 ,			_lowercase    :       List[str]=0.02 ,			_lowercase    :       int=2 ,			_lowercase    :       Optional[Any]=True ,			_lowercase    :       List[Any]=False ,			_lowercase    :       List[str]="float32" ,			_lowercase    :       Optional[Any]=False ,			_lowercase    :       str=1_28 ,			_lowercase    :       int=64 ,			_lowercase    :       Optional[int]=4 ,			_lowercase    :       List[str]=4 ,			_lowercase    :       Union[str, Any]=0.0_01 ,			_lowercase    :       List[Any]=0.0_01 ,			_lowercase    :       List[str]="all" ,			_lowercase    :       Optional[Any]=False ,			_lowercase    :       int=False ,			_lowercase    :       Tuple=1.0 ,			_lowercase    :       Optional[int]=0.2 ,			_lowercase    :       Optional[int]=1 ,			_lowercase    :       List[Any]=0 ,			_lowercase    :       List[Any]=2 ,			_lowercase    :       int=False ,			**_lowercase    :       Union[str, Any] ,			):



     """simple docstring"""






     SCREAMING_SNAKE_CASE__    							=   vocab_size
     SCREAMING_SNAKE_CASE__    							=   max_position_embeddings
     SCREAMING_SNAKE_CASE__    							=   d_model
     SCREAMING_SNAKE_CASE__    							=   encoder_ffn_dim
     SCREAMING_SNAKE_CASE__    							=   encoder_layers
     SCREAMING_SNAKE_CASE__    							=   encoder_attention_heads
     SCREAMING_SNAKE_CASE__    							=   decoder_ffn_dim
     SCREAMING_SNAKE_CASE__    							=   decoder_layers
     SCREAMING_SNAKE_CASE__    							=   decoder_attention_heads
     SCREAMING_SNAKE_CASE__    							=   dropout
     SCREAMING_SNAKE_CASE__    							=   attention_dropout
     SCREAMING_SNAKE_CASE__    							=   activation_dropout
     SCREAMING_SNAKE_CASE__    							=   activation_function
     SCREAMING_SNAKE_CASE__    							=   init_std
     SCREAMING_SNAKE_CASE__    							=   encoder_layerdrop
     SCREAMING_SNAKE_CASE__    							=   decoder_layerdrop
     SCREAMING_SNAKE_CASE__    							=   use_cache
     SCREAMING_SNAKE_CASE__    							=   encoder_layers
     SCREAMING_SNAKE_CASE__    							=   scale_embedding  # scale factor will be sqrt(d_model) if True
     SCREAMING_SNAKE_CASE__    							=   router_z_loss_coef
     SCREAMING_SNAKE_CASE__    							=   router_aux_loss_coef
     SCREAMING_SNAKE_CASE__    							=   decoder_sparse_step
     SCREAMING_SNAKE_CASE__    							=   encoder_sparse_step
     SCREAMING_SNAKE_CASE__    							=   num_experts
     SCREAMING_SNAKE_CASE__    							=   expert_capacity
     SCREAMING_SNAKE_CASE__    							=   router_bias
     if router_dtype not in ["float32", "float16", "bfloat16"]:
       raise ValueError(f"""`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}"""   )
     SCREAMING_SNAKE_CASE__    							=   router_dtype

     SCREAMING_SNAKE_CASE__    							=   router_ignore_padding_tokens
     SCREAMING_SNAKE_CASE__    							=   batch_prioritized_routing
     SCREAMING_SNAKE_CASE__    							=   second_expert_policy
     SCREAMING_SNAKE_CASE__    							=   normalize_router_prob_before_dropping
     SCREAMING_SNAKE_CASE__    							=   moe_eval_capacity_token_fraction
     SCREAMING_SNAKE_CASE__    							=   moe_token_dropout
     SCREAMING_SNAKE_CASE__    							=   output_router_logits
     super().__init__(
         pad_token_id=_lowercase ,			bos_token_id=_lowercase ,			eos_token_id=_lowercase ,			is_encoder_decoder=_lowercase ,			decoder_start_token_id=_lowercase ,			**_lowercase ,			)
