


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





							'''simple docstring'''


							def __init__(    self    ,					__snake_case    ,					__snake_case=13    ,					__snake_case=7    ,					__snake_case=False    ,					__snake_case=True    ,					__snake_case=False    ,					__snake_case=True    ,					__snake_case=33    ,					__snake_case=32    ,					__snake_case=5    ,					__snake_case=4    ,					__snake_case=37    ,					__snake_case="gelu"    ,					__snake_case=0.1    ,					__snake_case=0.1    ,					__snake_case=512    ,					__snake_case=16    ,					__snake_case=2    ,					__snake_case=0.02    ,					__snake_case=3    ,					__snake_case=4    ,					__snake_case=None    ,					):
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      parent
											_SCREAMING_SNAKE_CASE :       List[str]     		=      batch_size
											_SCREAMING_SNAKE_CASE :       Any     		=      seq_length
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      is_training
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      use_input_mask
											_SCREAMING_SNAKE_CASE :       List[str]     		=      use_token_type_ids
											_SCREAMING_SNAKE_CASE :       List[str]     		=      use_labels
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      vocab_size
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      hidden_size
											_SCREAMING_SNAKE_CASE :       str     		=      num_hidden_layers
											_SCREAMING_SNAKE_CASE :       int     		=      num_attention_heads
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      intermediate_size
											_SCREAMING_SNAKE_CASE :       List[str]     		=      hidden_act
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      hidden_dropout_prob
											_SCREAMING_SNAKE_CASE :       List[str]     		=      attention_probs_dropout_prob
											_SCREAMING_SNAKE_CASE :       List[str]     		=      max_position_embeddings
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      type_vocab_size
											_SCREAMING_SNAKE_CASE :       int     		=      type_sequence_label_size
											_SCREAMING_SNAKE_CASE :       Tuple     		=      initializer_range
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      num_labels
											_SCREAMING_SNAKE_CASE :       Tuple     		=      num_choices
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      scope






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

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

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

											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      self.get_config()

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






							def UpperCAmelCase_	(    self  ):
											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 UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case  ):
											_SCREAMING_SNAKE_CASE :       Dict     		=      EsmModel(config=__snake_case  )
											model.to(__snake_case  )
											model.eval()
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      model(__snake_case    ,					attention_mask=__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      model(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Any     		=      model(__snake_case  )

											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 UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case  ):
											_SCREAMING_SNAKE_CASE :       Tuple     		=      EsmForMaskedLM(config=__snake_case  )
											model.to(__snake_case  )
											model.eval()
											_SCREAMING_SNAKE_CASE :       List[str]     		=      model(__snake_case    ,					attention_mask=__snake_case    ,					labels=__snake_case  )
											self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)  )






							def UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case    ,					__snake_case  ):
											_SCREAMING_SNAKE_CASE :       Any     		=      self.num_labels
											_SCREAMING_SNAKE_CASE :       List[str]     		=      EsmForTokenClassification(config=__snake_case  )
											model.to(__snake_case  )
											model.eval()
											_SCREAMING_SNAKE_CASE :       List[str]     		=      model(__snake_case    ,					attention_mask=__snake_case    ,					labels=__snake_case  )
											self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.num_labels)  )






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





@require_torch
class     lowercase__  (						_snake_case     ,     _snake_case     ,     unittest.TestCase					):





							'''simple docstring'''


							A_							:      str      	=  False

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






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Any     		=      EsmModelTester(self  )
											_SCREAMING_SNAKE_CASE :       str     		=      ConfigTester(self    ,					config_class=__snake_case    ,					hidden_size=37  )






							def UpperCAmelCase_	(    self  ):
											self.config_tester.run_common_tests()






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       int     		=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Any     		=      self.model_tester.prepare_config_and_inputs()
											for type in ["absolute", "relative_key", "relative_key_query"]:
															_SCREAMING_SNAKE_CASE :       List[str]     		=      type
															self.model_tester.create_and_check_model(*__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_masked_lm(*__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[Any]     		=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_token_classification(*__snake_case  )






							@slow
							def UpperCAmelCase_	(    self  ):
											for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      EsmModel.from_pretrained(__snake_case  )
															self.assertIsNotNone(__snake_case  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.model_tester.prepare_config_and_inputs()[0]
											_SCREAMING_SNAKE_CASE :       str     		=      EsmEmbeddings(config=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch.as_tensor([[12, 31, 13, model.padding_idx]]  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      torch.as_tensor(
											    [
											        [
											            0 + model.padding_idx + 1,
											            1 + model.padding_idx + 1,
											            2 + model.padding_idx + 1,
											            model.padding_idx,
											        ]
											    ]  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      create_position_ids_from_input_ids(__snake_case    ,					model.padding_idx  )
											self.assertEqual(position_ids.shape    ,					expected_positions.shape  )
											self.assertTrue(torch.all(torch.eq(__snake_case    ,					__snake_case  )  )  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       str     		=      self.model_tester.prepare_config_and_inputs()[0]
											_SCREAMING_SNAKE_CASE :       Dict     		=      EsmEmbeddings(config=__snake_case  )

											_SCREAMING_SNAKE_CASE :       Dict     		=      torch.empty(2    ,					4    ,					30  )
											_SCREAMING_SNAKE_CASE :       str     		=      [
											    0 + embeddings.padding_idx + 1,
											    1 + embeddings.padding_idx + 1,
											    2 + embeddings.padding_idx + 1,
											    3 + embeddings.padding_idx + 1,
											]
											_SCREAMING_SNAKE_CASE :       str     		=      torch.as_tensor([expected_single_positions, expected_single_positions]  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      embeddings.create_position_ids_from_inputs_embeds(__snake_case  )
											self.assertEqual(position_ids.shape    ,					expected_positions.shape  )
											self.assertTrue(torch.all(torch.eq(__snake_case    ,					__snake_case  )  )  )






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






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






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





@require_torch
class     lowercase__  (						_snake_case					):





							'''simple docstring'''


							@slow
							def UpperCAmelCase_	(    self  ):
											with torch.no_grad():
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""  )
															model.eval()
															_SCREAMING_SNAKE_CASE :       Any     		=      torch.tensor([[0, 1, 2, 3, 4, 5]]  )
															_SCREAMING_SNAKE_CASE :       List[Any]     		=      model(__snake_case  )[0]

															_SCREAMING_SNAKE_CASE :       Dict     		=      33

															_SCREAMING_SNAKE_CASE :       int     		=      torch.Size((1, 6, vocab_size)  )
															self.assertEqual(output.shape    ,					__snake_case  )

															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      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]    ,					__snake_case    ,					atol=1e-4  )  )






							@slow
							def UpperCAmelCase_	(    self  ):
											with torch.no_grad():
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""  )
															model.eval()

															_SCREAMING_SNAKE_CASE :       List[Any]     		=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      model(__snake_case  )[0]
															# compare the actual values for a slice.
															_SCREAMING_SNAKE_CASE :       Tuple     		=      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]    ,					__snake_case    ,					atol=1e-4  )  )








'''simple docstring'''


from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       Tuple     		=      []
				_SCREAMING_SNAKE_CASE :       Dict     		=      []
				_SCREAMING_SNAKE_CASE :       str     		=      []

				for rt in rc.restypes:
								_SCREAMING_SNAKE_CASE :       Optional[int]     		=      rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names]     )
								_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      {name: i for i, name in enumerate(SCREAMING_SNAKE_CASE__     )}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types]     )

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names]     )

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14     )
				restype_atomaa_to_atomaa_list.append([0] * 37     )
				restype_atomaa_mask_list.append([0.0] * 14     )

				_SCREAMING_SNAKE_CASE :       Dict     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.intaa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.intaa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       str     		=      torch.tensor(
				    SCREAMING_SNAKE_CASE__							,      dtype=torch.floataa							,      device=protein["""aatype"""].device							,      )
				_SCREAMING_SNAKE_CASE :       int     		=      protein["""aatype"""].to(torch.long     )

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				_SCREAMING_SNAKE_CASE :       List[str]     		=      restype_atomaa_to_atomaa[protein_aatype]
				_SCREAMING_SNAKE_CASE :       str     		=      restype_atomaa_mask[protein_aatype]

				_SCREAMING_SNAKE_CASE :       List[Any]     		=      residx_atomaa_mask
				_SCREAMING_SNAKE_CASE :       List[str]     		=      residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				_SCREAMING_SNAKE_CASE :       int     		=      restype_atomaa_to_atomaa[protein_aatype]
				_SCREAMING_SNAKE_CASE :       Dict     		=      residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				_SCREAMING_SNAKE_CASE :       str     		=      torch.zeros([21, 37]							,      dtype=torch.floataa							,      device=protein["""aatype"""].device     )
				for restype, restype_letter in enumerate(rc.restypes     ):
								_SCREAMING_SNAKE_CASE :       int     		=      rc.restype_atoa[restype_letter]
								_SCREAMING_SNAKE_CASE :       Dict     		=      rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												_SCREAMING_SNAKE_CASE :       List[Any]     		=      rc.atom_order[atom_name]
												_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      1

				_SCREAMING_SNAKE_CASE :       int     		=      restype_atomaa_mask[protein_aatype]
				_SCREAMING_SNAKE_CASE :       int     		=      residx_atomaa_mask

				return protein







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      tree_map(lambda SCREAMING_SNAKE_CASE__     : torch.tensor(SCREAMING_SNAKE_CASE__							,      device=batch["""aatype"""].device     )							,      SCREAMING_SNAKE_CASE__							,      np.ndarray     )
				_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      tensor_tree_map(lambda SCREAMING_SNAKE_CASE__     : np.array(SCREAMING_SNAKE_CASE__     )							,      make_atomaa_masks(SCREAMING_SNAKE_CASE__     )     )
				return out




