





import argparse
import os
import re

import tensorflow as tf
import torch

from transformers import BertConfig, BertModel
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase							:       List[Any]					 =	logging.get_logger(__name__)






def 		lowerCAmelCase  (	UpperCamelCase__   :   str       ,			UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   List[str]					)      ->    List[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Tuple           =   os.path.abspath(UpperCamelCase__					)
						logger.info(F"""Converting TensorFlow checkpoint from {tf_path}"""					)
						# Load weights from TF model
						__SCREAMING_SNAKE_CASE:				Tuple           =   tf.train.list_variables(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   []
						__SCREAMING_SNAKE_CASE:				Any           =   []
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   []
						for full_name, shape in init_vars:
												# logger.info(f"Loading TF weight {name} with shape {shape}")
												__SCREAMING_SNAKE_CASE:				Any           =   full_name.split('''/'''					)
												if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
																		logger.info(F"""Skipping non-model layer {full_name}"""					)
																		continue
												if "optimizer" in full_name:
																		logger.info(F"""Skipping optimization layer {full_name}"""					)
																		continue
												if name[0] == "model":
																		# ignore initial 'model'
																		__SCREAMING_SNAKE_CASE:				Dict           =   name[1:]
												# figure out how many levels deep the name is
												__SCREAMING_SNAKE_CASE:				List[Any]           =   0
												for _name in name:
																		if _name.startswith('''layer_with_weights'''					):
																								depth += 1
																		else:
																								break
												layer_depth.append(UpperCamelCase__					)
												# read data
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   tf.train.load_variable(UpperCamelCase__       ,			UpperCamelCase__					)
												names.append('''/'''.join(UpperCamelCase__					)					)
												arrays.append(UpperCamelCase__					)
						logger.info(F"""Read a total of {len(UpperCamelCase__					):,} layers"""					)

						# Sanity check
						if len(set(UpperCamelCase__					)					) != 1:
												raise ValueError(F"""Found layer names with different depths (layer depth {list(set(UpperCamelCase__					)					)})"""					)
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   list(set(UpperCamelCase__					)					)[0]
						if layer_depth != 1:
												raise ValueError(
												    '''The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP'''
												    ''' heads.'''					)

						# convert layers
						logger.info('''Converting weights...'''					)
						for full_name, array in zip(UpperCamelCase__       ,			UpperCamelCase__					):
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   full_name.split('''/'''					)
												__SCREAMING_SNAKE_CASE:				Dict           =   model
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   []
												for i, m_name in enumerate(UpperCamelCase__					):
																		if m_name == ".ATTRIBUTES":
																								# variable names end with .ATTRIBUTES/VARIABLE_VALUE
																								break
																		if m_name.startswith('''layer_with_weights'''					):
																								__SCREAMING_SNAKE_CASE:				Dict           =   int(m_name.split('''-'''					)[-1]					)
																								if layer_num <= 2:
																														# embedding layers
																														# layer_num 0: word_embeddings
																														# layer_num 1: position_embeddings
																														# layer_num 2: token_type_embeddings
																														continue
																								elif layer_num == 3:
																														# embedding LayerNorm
																														trace.extend(['''embeddings''', '''LayerNorm''']					)
																														__SCREAMING_SNAKE_CASE:				int           =   getattr(UpperCamelCase__       ,			'''embeddings'''					)
																														__SCREAMING_SNAKE_CASE:				Optional[int]           =   getattr(UpperCamelCase__       ,			'''LayerNorm'''					)
																								elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
																														# encoder layers
																														trace.extend(['''encoder''', '''layer''', str(layer_num - 4					)]					)
																														__SCREAMING_SNAKE_CASE:				Dict           =   getattr(UpperCamelCase__       ,			'''encoder'''					)
																														__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   getattr(UpperCamelCase__       ,			'''layer'''					)
																														__SCREAMING_SNAKE_CASE:				Dict           =   pointer[layer_num - 4]
																								elif layer_num == config.num_hidden_layers + 4:
																														# pooler layer
																														trace.extend(['''pooler''', '''dense''']					)
																														__SCREAMING_SNAKE_CASE:				Tuple           =   getattr(UpperCamelCase__       ,			'''pooler'''					)
																														__SCREAMING_SNAKE_CASE:				int           =   getattr(UpperCamelCase__       ,			'''dense'''					)
																		elif m_name == "embeddings":
																								trace.append('''embeddings'''					)
																								__SCREAMING_SNAKE_CASE:				Optional[int]           =   getattr(UpperCamelCase__       ,			'''embeddings'''					)
																								if layer_num == 0:
																														trace.append('''word_embeddings'''					)
																														__SCREAMING_SNAKE_CASE:				int           =   getattr(UpperCamelCase__       ,			'''word_embeddings'''					)
																								elif layer_num == 1:
																														trace.append('''position_embeddings'''					)
																														__SCREAMING_SNAKE_CASE:				Dict           =   getattr(UpperCamelCase__       ,			'''position_embeddings'''					)
																								elif layer_num == 2:
																														trace.append('''token_type_embeddings'''					)
																														__SCREAMING_SNAKE_CASE:				Optional[Any]           =   getattr(UpperCamelCase__       ,			'''token_type_embeddings'''					)
																								else:
																														raise ValueError(F"""Unknown embedding layer with name {full_name}"""					)
																								trace.append('''weight'''					)
																								__SCREAMING_SNAKE_CASE:				Dict           =   getattr(UpperCamelCase__       ,			'''weight'''					)
																		elif m_name == "_attention_layer":
																								# self-attention layer
																								trace.extend(['''attention''', '''self''']					)
																								__SCREAMING_SNAKE_CASE:				Optional[Any]           =   getattr(UpperCamelCase__       ,			'''attention'''					)
																								__SCREAMING_SNAKE_CASE:				List[str]           =   getattr(UpperCamelCase__       ,			'''self'''					)
																		elif m_name == "_attention_layer_norm":
																								# output attention norm
																								trace.extend(['''attention''', '''output''', '''LayerNorm''']					)
																								__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   getattr(UpperCamelCase__       ,			'''attention'''					)
																								__SCREAMING_SNAKE_CASE:				Optional[Any]           =   getattr(UpperCamelCase__       ,			'''output'''					)
																								__SCREAMING_SNAKE_CASE:				str           =   getattr(UpperCamelCase__       ,			'''LayerNorm'''					)
																		elif m_name == "_attention_output_dense":
																								# output attention dense
																								trace.extend(['''attention''', '''output''', '''dense''']					)
																								__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   getattr(UpperCamelCase__       ,			'''attention'''					)
																								__SCREAMING_SNAKE_CASE:				List[str]           =   getattr(UpperCamelCase__       ,			'''output'''					)
																								__SCREAMING_SNAKE_CASE:				List[Any]           =   getattr(UpperCamelCase__       ,			'''dense'''					)
																		elif m_name == "_output_dense":
																								# output dense
																								trace.extend(['''output''', '''dense''']					)
																								__SCREAMING_SNAKE_CASE:				str           =   getattr(UpperCamelCase__       ,			'''output'''					)
																								__SCREAMING_SNAKE_CASE:				Optional[int]           =   getattr(UpperCamelCase__       ,			'''dense'''					)
																		elif m_name == "_output_layer_norm":
																								# output dense
																								trace.extend(['''output''', '''LayerNorm''']					)
																								__SCREAMING_SNAKE_CASE:				Optional[Any]           =   getattr(UpperCamelCase__       ,			'''output'''					)
																								__SCREAMING_SNAKE_CASE:				List[str]           =   getattr(UpperCamelCase__       ,			'''LayerNorm'''					)
																		elif m_name == "_key_dense":
																								# attention key
																								trace.append('''key'''					)
																								__SCREAMING_SNAKE_CASE:				Optional[Any]           =   getattr(UpperCamelCase__       ,			'''key'''					)
																		elif m_name == "_query_dense":
																								# attention query
																								trace.append('''query'''					)
																								__SCREAMING_SNAKE_CASE:				Dict           =   getattr(UpperCamelCase__       ,			'''query'''					)
																		elif m_name == "_value_dense":
																								# attention value
																								trace.append('''value'''					)
																								__SCREAMING_SNAKE_CASE:				Any           =   getattr(UpperCamelCase__       ,			'''value'''					)
																		elif m_name == "_intermediate_dense":
																								# attention intermediate dense
																								trace.extend(['''intermediate''', '''dense''']					)
																								__SCREAMING_SNAKE_CASE:				Any           =   getattr(UpperCamelCase__       ,			'''intermediate'''					)
																								__SCREAMING_SNAKE_CASE:				Any           =   getattr(UpperCamelCase__       ,			'''dense'''					)
																		elif m_name == "_output_layer_norm":
																								# output layer norm
																								trace.append('''output'''					)
																								__SCREAMING_SNAKE_CASE:				Any           =   getattr(UpperCamelCase__       ,			'''output'''					)
																		# weights & biases
																		elif m_name in ["bias", "beta"]:
																								trace.append('''bias'''					)
																								__SCREAMING_SNAKE_CASE:				Optional[int]           =   getattr(UpperCamelCase__       ,			'''bias'''					)
																		elif m_name in ["kernel", "gamma"]:
																								trace.append('''weight'''					)
																								__SCREAMING_SNAKE_CASE:				Tuple           =   getattr(UpperCamelCase__       ,			'''weight'''					)
																		else:
																								logger.warning(F"""Ignored {m_name}"""					)
        # for certain layers reshape is necessary
												__SCREAMING_SNAKE_CASE:				Tuple           =   '''.'''.join(UpperCamelCase__					)
												if re.match(R'''(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)'''       ,			UpperCamelCase__					) or re.match(
												    R'''(\S+)\.attention\.output\.dense\.weight'''       ,			UpperCamelCase__					):
																		__SCREAMING_SNAKE_CASE:				Optional[Any]           =   array.reshape(pointer.data.shape					)
												if "kernel" in full_name:
																		__SCREAMING_SNAKE_CASE:				List[Any]           =   array.transpose()
												if pointer.shape == array.shape:
																		__SCREAMING_SNAKE_CASE:				Any           =   torch.from_numpy(UpperCamelCase__					)
												else:
																		raise ValueError(
																		    F"""Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:"""
																		    F""" {array.shape}"""					)
												logger.info(F"""Successfully set variable {full_name} to PyTorch layer {trace}"""					)
						return model






def 		lowerCAmelCase  (	UpperCamelCase__   :   str       ,			UpperCamelCase__   :   List[Any]       ,			UpperCamelCase__   :   int					)      ->    List[Any]:

						"""simple docstring"""





						logger.info(F"""Loading model based on config from {config_path}..."""					)
						__SCREAMING_SNAKE_CASE:				Tuple           =   BertConfig.from_json_file(UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				int           =   BertModel(UpperCamelCase__					)

						# Load weights from checkpoint
						logger.info(F"""Loading weights from checkpoint {tf_checkpoint_path}..."""					)
						load_tfa_weights_in_bert(UpperCamelCase__       ,			UpperCamelCase__       ,			UpperCamelCase__					)

						# Save pytorch-model
						logger.info(F"""Saving PyTorch model to {pytorch_dump_path}..."""					)
						torch.save(model.state_dict()       ,			UpperCamelCase__					)


if __name__ == "__main__":
							lowerCAmelCase							:       List[str]					 =	argparse.ArgumentParser()
							parser.add_argument(
							    """--tf_checkpoint_path""", type=str, required=True, help="""Path to the TensorFlow 2.x checkpoint path."""
							)
							parser.add_argument(
							    """--bert_config_file""",
							    type=str,
							    required=True,
							    help="""The config json file corresponding to the BERT model. This specifies the model architecture.""",
							)
							parser.add_argument(
							    """--pytorch_dump_path""",
							    type=str,
							    required=True,
							    help="""Path to the output PyTorch model (must include filename).""",
							)
							lowerCAmelCase							:       Optional[Any]					 =	parser.parse_args()
							convert_tfa_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)










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






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

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




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

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

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

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.get_config()

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






				def 			snake_case_    (  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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmModel(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   model(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				int           =   model(_lowerCAmelCase					)

										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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmForMaskedLM(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase			, labels=_lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.vocab_size)					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Any           =   self.num_labels
										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmForTokenClassification(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				int           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase			, labels=_lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.num_labels)					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				str           =   self.prepare_config_and_inputs()
										(
										    (
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,
										):				int           =   config_and_inputs
										__SCREAMING_SNAKE_CASE:				Tuple           =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict





@require_torch
class 				a			(      __lowercase   ,__lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   Tuple		      =     False

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   EsmModelTester(self					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   ConfigTester(self			, config_class=_lowerCAmelCase			, hidden_size=37					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										self.config_tester.run_common_tests()






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																__SCREAMING_SNAKE_CASE:				Any           =   type
																self.model_tester.create_and_check_model(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase					)






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   EsmModel.from_pretrained(_lowerCAmelCase					)
																self.assertIsNotNone(_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   self.model_tester.prepare_config_and_inputs()[0]
										__SCREAMING_SNAKE_CASE:				int           =   EsmEmbeddings(config=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
										__SCREAMING_SNAKE_CASE:				Optional[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:				Any           =   create_position_ids_from_input_ids(_lowerCAmelCase			, model.padding_idx					)
										self.assertEqual(position_ids.shape			, expected_positions.shape					)
										self.assertTrue(torch.all(torch.eq(_lowerCAmelCase			, _lowerCAmelCase					)					)					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   self.model_tester.prepare_config_and_inputs()[0]
										__SCREAMING_SNAKE_CASE:				Dict           =   EsmEmbeddings(config=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Optional[int]           =   torch.empty(2			, 4			, 30					)
										__SCREAMING_SNAKE_CASE:				List[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:				Optional[Any]           =   embeddings.create_position_ids_from_inputs_embeds(_lowerCAmelCase					)
										self.assertEqual(position_ids.shape			, expected_positions.shape					)
										self.assertTrue(torch.all(torch.eq(_lowerCAmelCase			, _lowerCAmelCase					)					)					)






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

										"""simple docstring"""




										pass






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

										"""simple docstring"""




										pass






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

										"""simple docstring"""




										pass





@require_torch
class 				a			(      __lowercase			):






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										with torch.no_grad():
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''					)
																model.eval()
																__SCREAMING_SNAKE_CASE:				Tuple           =   torch.tensor([[0, 1, 2, 3, 4, 5]]					)
																__SCREAMING_SNAKE_CASE:				str           =   model(_lowerCAmelCase					)[0]

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   33

																__SCREAMING_SNAKE_CASE:				Dict           =   torch.Size((1, 6, vocab_size)					)
																self.assertEqual(output.shape			, _lowerCAmelCase					)

																__SCREAMING_SNAKE_CASE:				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]			, _lowerCAmelCase			, atol=1e-4					)					)






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




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

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



