import logging
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

from filelock import FileLock

from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available


lowercase__			:				int		=	logging.getLogger(__name__)




@dataclass
class    lowercase_  :



	"""simple docstring"""





	UpperCAmelCase_  :					str
	UpperCAmelCase_  :					List[str]
	UpperCAmelCase_  :					Optional[List[str]]




@dataclass
class    lowercase_  :



	"""simple docstring"""





	UpperCAmelCase_  :					List[int]
	UpperCAmelCase_  :					List[int]
	UpperCAmelCase_  :					Optional[List[int]]   =				None
	UpperCAmelCase_  :					Optional[List[int]]   =				None




class    lowercase_  (  UpperCAmelCase__						):



	"""simple docstring"""





	UpperCAmelCase_  :					str   =				"train"
	UpperCAmelCase_  :					Tuple   =				"dev"
	UpperCAmelCase_  :					str   =				"test"




class    lowercase_  :



	"""simple docstring"""





	@staticmethod
	def 			SCREAMING_SNAKE_CASE_				(							__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE				)						->List[InputExample]:
								raise NotImplementedError


	@staticmethod
	def 			SCREAMING_SNAKE_CASE_				(							__SCREAMING_SNAKE_CASE				)						->List[str]:
								raise NotImplementedError




	@staticmethod
	def 			SCREAMING_SNAKE_CASE_				(							__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE=False		,			__SCREAMING_SNAKE_CASE="[CLS]"		,			__SCREAMING_SNAKE_CASE=1		,			__SCREAMING_SNAKE_CASE="[SEP]"		,			__SCREAMING_SNAKE_CASE=False		,			__SCREAMING_SNAKE_CASE=False		,			__SCREAMING_SNAKE_CASE=0		,			__SCREAMING_SNAKE_CASE=0		,			__SCREAMING_SNAKE_CASE=-100		,			__SCREAMING_SNAKE_CASE=0		,			__SCREAMING_SNAKE_CASE=True		,			)						->List[InputFeatures]:
								lowerCAmelCase			=		{label: i for i, label in enumerate(__SCREAMING_SNAKE_CASE				)}

								lowerCAmelCase			=		[]
								for ex_index, example in enumerate(__SCREAMING_SNAKE_CASE				):
															if ex_index % 10000 == 0:
																						logger.info('''Writing example %d of %d'''		,			__SCREAMING_SNAKE_CASE		,			len(__SCREAMING_SNAKE_CASE				)				)

															lowerCAmelCase			=		[]
															lowerCAmelCase			=		[]
															for word, label in zip(example.words		,			example.labels				):
																						lowerCAmelCase			=		tokenizer.tokenize(__SCREAMING_SNAKE_CASE				)

																						# bert-base-multilingual-cased sometimes output "nothing ([]) when calling tokenize with just a space.
																						if len(__SCREAMING_SNAKE_CASE				) > 0:
																													tokens.extend(__SCREAMING_SNAKE_CASE				)
																													# Use the real label id for the first token of the word, and padding ids for the remaining tokens
																													label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(__SCREAMING_SNAKE_CASE				) - 1)				)

            # Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
															lowerCAmelCase			=		tokenizer.num_special_tokens_to_add()
															if len(__SCREAMING_SNAKE_CASE				) > max_seq_length - special_tokens_count:
																						lowerCAmelCase			=		tokens[: (max_seq_length - special_tokens_count)]
																						lowerCAmelCase			=		label_ids[: (max_seq_length - special_tokens_count)]

															# The convention in BERT is:
															# (a) For sequence pairs:
															#  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
															#  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1
															# (b) For single sequences:
															#  tokens:   [CLS] the dog is hairy . [SEP]
															#  type_ids:   0   0   0   0  0     0   0
															#
															# Where "type_ids" are used to indicate whether this is the first
															# sequence or the second sequence. The embedding vectors for `type=0` and
															# `type=1` were learned during pre-training and are added to the wordpiece
															# embedding vector (and position vector). This is not *strictly* necessary
															# since the [SEP] token unambiguously separates the sequences, but it makes
															# it easier for the model to learn the concept of sequences.
															#
															# For classification tasks, the first vector (corresponding to [CLS]) is
															# used as the "sentence vector". Note that this only makes sense because
															# the entire model is fine-tuned.
															tokens += [sep_token]
															label_ids += [pad_token_label_id]
															if sep_token_extra:
																						# roberta uses an extra separator b/w pairs of sentences
																						tokens += [sep_token]
																						label_ids += [pad_token_label_id]
															lowerCAmelCase			=		[sequence_a_segment_id] * len(__SCREAMING_SNAKE_CASE				)

															if cls_token_at_end:
																						tokens += [cls_token]
																						label_ids += [pad_token_label_id]
																						segment_ids += [cls_token_segment_id]
															else:
																						lowerCAmelCase			=		[cls_token] + tokens
																						lowerCAmelCase			=		[pad_token_label_id] + label_ids
																						lowerCAmelCase			=		[cls_token_segment_id] + segment_ids

															lowerCAmelCase			=		tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE				)

															# The mask has 1 for real tokens and 0 for padding tokens. Only real
															# tokens are attended to.
															lowerCAmelCase			=		[1 if mask_padding_with_zero else 0] * len(__SCREAMING_SNAKE_CASE				)

															# Zero-pad up to the sequence length.
															lowerCAmelCase			=		max_seq_length - len(__SCREAMING_SNAKE_CASE				)
															if pad_on_left:
																						lowerCAmelCase			=		([pad_token] * padding_length) + input_ids
																						lowerCAmelCase			=		([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
																						lowerCAmelCase			=		([pad_token_segment_id] * padding_length) + segment_ids
																						lowerCAmelCase			=		([pad_token_label_id] * padding_length) + label_ids
															else:
																						input_ids += [pad_token] * padding_length
																						input_mask += [0 if mask_padding_with_zero else 1] * padding_length
																						segment_ids += [pad_token_segment_id] * padding_length
																						label_ids += [pad_token_label_id] * padding_length

															assert len(__SCREAMING_SNAKE_CASE				) == max_seq_length
															assert len(__SCREAMING_SNAKE_CASE				) == max_seq_length
															assert len(__SCREAMING_SNAKE_CASE				) == max_seq_length
															assert len(__SCREAMING_SNAKE_CASE				) == max_seq_length

															if ex_index < 5:
																						logger.info('''*** Example ***'''				)
																						logger.info('''guid: %s'''		,			example.guid				)
																						logger.info('''tokens: %s'''		,			''' '''.join([str(__SCREAMING_SNAKE_CASE				) for x in tokens]				)				)
																						logger.info('''input_ids: %s'''		,			''' '''.join([str(__SCREAMING_SNAKE_CASE				) for x in input_ids]				)				)
																						logger.info('''input_mask: %s'''		,			''' '''.join([str(__SCREAMING_SNAKE_CASE				) for x in input_mask]				)				)
																						logger.info('''segment_ids: %s'''		,			''' '''.join([str(__SCREAMING_SNAKE_CASE				) for x in segment_ids]				)				)
																						logger.info('''label_ids: %s'''		,			''' '''.join([str(__SCREAMING_SNAKE_CASE				) for x in label_ids]				)				)

															if "token_type_ids" not in tokenizer.model_input_names:
																						lowerCAmelCase			=		None

															features.append(
															    InputFeatures(
															        input_ids=__SCREAMING_SNAKE_CASE		,			attention_mask=__SCREAMING_SNAKE_CASE		,			token_type_ids=__SCREAMING_SNAKE_CASE		,			label_ids=__SCREAMING_SNAKE_CASE				)				)
								return features


if is_torch_available():
	import torch
	from torch import nn
	from torch.utils.data import Dataset




	class    lowercase_  (  UpperCAmelCase__						):



		"""simple docstring"""





		UpperCAmelCase_  :					List[InputFeatures]
		UpperCAmelCase_  :					int   =				nn.CrossEntropyLoss().ignore_index


		def __init__(							self		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE = None		,			__SCREAMING_SNAKE_CASE=False		,			__SCREAMING_SNAKE_CASE = Split.train		,			)						->List[str]:
									lowerCAmelCase			=		os.path.join(
									    __SCREAMING_SNAKE_CASE		,			'''cached_{}_{}_{}'''.format(mode.value		,			tokenizer.__class__.__name__		,			str(__SCREAMING_SNAKE_CASE				)				)		,			)

									# Make sure only the first process in distributed training processes the dataset,
									# and the others will use the cache.
									lowerCAmelCase			=		cached_features_file + """.lock"""
									with FileLock(__SCREAMING_SNAKE_CASE				):
																if os.path.exists(__SCREAMING_SNAKE_CASE				) and not overwrite_cache:
																							logger.info(F"Loading features from cached file {cached_features_file}"				)
																							lowerCAmelCase			=		torch.load(__SCREAMING_SNAKE_CASE				)
																else:
																							logger.info(F"Creating features from dataset file at {data_dir}"				)
																							lowerCAmelCase			=		token_classification_task.read_examples_from_file(__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE				)
																							# TODO clean up all this to leverage built-in features of tokenizers
																							lowerCAmelCase			=		token_classification_task.convert_examples_to_features(
																							    __SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			cls_token_at_end=bool(model_type in ['''xlnet''']				)		,			cls_token=tokenizer.cls_token		,			cls_token_segment_id=2 if model_type in ['''xlnet'''] else 0		,			sep_token=tokenizer.sep_token		,			sep_token_extra=__SCREAMING_SNAKE_CASE		,			pad_on_left=bool(tokenizer.padding_side == '''left'''				)		,			pad_token=tokenizer.pad_token_id		,			pad_token_segment_id=tokenizer.pad_token_type_id		,			pad_token_label_id=self.pad_token_label_id		,			)
																							logger.info(F"Saving features into cached file {cached_features_file}"				)
																							torch.save(self.features		,			__SCREAMING_SNAKE_CASE				)


		def __len__(							self				)						->Union[str, Any]:
									return len(self.features				)




		def __getitem__(							self		,			__SCREAMING_SNAKE_CASE				)						->InputFeatures:
									return self.features[i]


if is_tf_available():
	import tensorflow as tf




	class    lowercase_  :



		"""simple docstring"""





		UpperCAmelCase_  :					List[InputFeatures]
		UpperCAmelCase_  :					int   =				-100


		def __init__(							self		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE = None		,			__SCREAMING_SNAKE_CASE=False		,			__SCREAMING_SNAKE_CASE = Split.train		,			)						->Any:
									lowerCAmelCase			=		token_classification_task.read_examples_from_file(__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE				)
									# TODO clean up all this to leverage built-in features of tokenizers
									lowerCAmelCase			=		token_classification_task.convert_examples_to_features(
									    __SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			__SCREAMING_SNAKE_CASE		,			cls_token_at_end=bool(model_type in ['''xlnet''']				)		,			cls_token=tokenizer.cls_token		,			cls_token_segment_id=2 if model_type in ['''xlnet'''] else 0		,			sep_token=tokenizer.sep_token		,			sep_token_extra=__SCREAMING_SNAKE_CASE		,			pad_on_left=bool(tokenizer.padding_side == '''left'''				)		,			pad_token=tokenizer.pad_token_id		,			pad_token_segment_id=tokenizer.pad_token_type_id		,			pad_token_label_id=self.pad_token_label_id		,			)

									def gen():
																for ex in self.features:
																							if ex.token_type_ids is None:
																														yield (
																														    {"input_ids": ex.input_ids, "attention_mask": ex.attention_mask},
																														    ex.label_ids,
																														)
																							else:
																														yield (
																														    {
																														        "input_ids": ex.input_ids,
																														        "attention_mask": ex.attention_mask,
																														        "token_type_ids": ex.token_type_ids,
																														    },
																														    ex.label_ids,
																														)

									if "token_type_ids" not in tokenizer.model_input_names:
																lowerCAmelCase			=		tf.data.Dataset.from_generator(
																    __SCREAMING_SNAKE_CASE		,			({'''input_ids''': tf.intaa, '''attention_mask''': tf.intaa}, tf.intaa)		,			(
																        {'''input_ids''': tf.TensorShape([None]				), '''attention_mask''': tf.TensorShape([None]				)},
																        tf.TensorShape([None]				),
																    )		,			)
									else:
																lowerCAmelCase			=		tf.data.Dataset.from_generator(
																    __SCREAMING_SNAKE_CASE		,			({'''input_ids''': tf.intaa, '''attention_mask''': tf.intaa, '''token_type_ids''': tf.intaa}, tf.intaa)		,			(
																        {
																            '''input_ids''': tf.TensorShape([None]				),
																            '''attention_mask''': tf.TensorShape([None]				),
																            '''token_type_ids''': tf.TensorShape([None]				),
																        },
																        tf.TensorShape([None]				),
																    )		,			)


		def 			SCREAMING_SNAKE_CASE_				(							self				)						->Any:
									lowerCAmelCase			=		self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features				)				)				)

									return self.dataset


		def __len__(							self				)						->Optional[int]:
									return len(self.features				)




		def __getitem__(							self		,			__SCREAMING_SNAKE_CASE				)						->InputFeatures:
									return self.features[i]







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






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




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






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




             '''simple docstring'''
             _a    : Optional[int]     					=ids_tensor([self.batch_size, self.seq_length]     ,						self.vocab_size	)

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

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

             _a    : List[Any]     					=self.get_config()

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






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




             '''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 __UpperCAmelCase       ( self			:List[str]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:List[str]     ,						SCREAMING_SNAKE_CASE			:int	)       ->					Tuple:




             '''simple docstring'''
             _a    : Any     					=EsmModel(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Optional[Any]     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE	)
             _a    : Union[str, Any]     					=model(SCREAMING_SNAKE_CASE	)
             _a    : str     					=model(SCREAMING_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			:str     ,						SCREAMING_SNAKE_CASE			:Any     ,						SCREAMING_SNAKE_CASE			:List[str]     ,						SCREAMING_SNAKE_CASE			:str     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:Optional[Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : str     					=EsmForMaskedLM(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Union[str, Any]     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.vocab_size)	)






      def __UpperCAmelCase       ( self			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Tuple     ,						SCREAMING_SNAKE_CASE			:Optional[Any]     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]     ,						SCREAMING_SNAKE_CASE			:Dict     ,						SCREAMING_SNAKE_CASE			:List[Any]     ,						SCREAMING_SNAKE_CASE			:Union[str, Any]	)       ->					Optional[Any]:




             '''simple docstring'''
             _a    : int     					=self.num_labels
             _a    : Tuple     					=EsmForTokenClassification(config=SCREAMING_SNAKE_CASE	)
             model.to(SCREAMING_SNAKE_CASE	)
             model.eval()
             _a    : Tuple     					=model(SCREAMING_SNAKE_CASE     ,						attention_mask=SCREAMING_SNAKE_CASE     ,						labels=SCREAMING_SNAKE_CASE	)
             self.parent.assertEqual(result.logits.shape     ,						(self.batch_size, self.seq_length, self.num_labels)	)






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




             '''simple docstring'''
             _a    : Optional[Any]     					=self.prepare_config_and_inputs()
             (
                 (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    (
                 _a
             )					,    
             )    : Any     					=config_and_inputs
             _a    : List[Any]     					={"""input_ids""": input_ids, """attention_mask""": input_mask}
             return config, inputs_dict







@require_torch
class A__       (						UpperCAmelCase__ ,     UpperCAmelCase__ ,     unittest.TestCase      ):
      __UpperCamelCase       :  Any	       =							False

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






      def __UpperCAmelCase       ( self			:Optional[int]	)       ->					int:




             '''simple docstring'''
             _a    : Dict     					=EsmModelTester(self	)
             _a    : Optional[Any]     					=ConfigTester(self     ,						config_class=SCREAMING_SNAKE_CASE     ,						hidden_size=3_7	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					List[Any]:




             '''simple docstring'''
             self.config_tester.run_common_tests()






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




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:List[Any]	)       ->					Dict:




             '''simple docstring'''
             _a    : List[str]     					=self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                    _a    : Dict     					=type
                    self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             _a    : Tuple     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE	)






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




             '''simple docstring'''
             _a    : str     					=self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE	)






      @slow
      def __UpperCAmelCase       ( self			:str	)       ->					Dict:




             '''simple docstring'''
             for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    _a    : Union[str, Any]     					=EsmModel.from_pretrained(SCREAMING_SNAKE_CASE	)
                    self.assertIsNotNone(SCREAMING_SNAKE_CASE	)






      def __UpperCAmelCase       ( self			:Tuple	)       ->					int:




             '''simple docstring'''
             _a    : Optional[Any]     					=self.model_tester.prepare_config_and_inputs()[0]
             _a    : Dict     					=EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

             _a    : Tuple     					=torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]	)
             _a    : Optional[Any]     					=torch.as_tensor(
                 [
                     [
                         0 + model.padding_idx + 1,
                         1 + model.padding_idx + 1,
                         2 + model.padding_idx + 1,
                         model.padding_idx,
                     ]
                 ]	)
             _a    : Any     					=create_position_ids_from_input_ids(SCREAMING_SNAKE_CASE     ,						model.padding_idx	)
             self.assertEqual(position_ids.shape     ,						expected_positions.shape	)
             self.assertTrue(torch.all(torch.eq(SCREAMING_SNAKE_CASE     ,						SCREAMING_SNAKE_CASE	)	)	)






      def __UpperCAmelCase       ( self			:Optional[Any]	)       ->					Tuple:




             '''simple docstring'''
             _a    : List[Any]     					=self.model_tester.prepare_config_and_inputs()[0]
             _a    : Optional[int]     					=EsmEmbeddings(config=SCREAMING_SNAKE_CASE	)

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






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




             '''simple docstring'''
             pass






      @unittest.skip("""Esm does not support embedding resizing"""	)
      def __UpperCAmelCase       ( self			:str	)       ->					Any:




             '''simple docstring'''
             pass






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




             '''simple docstring'''
             pass







@require_torch
class A__       (						UpperCAmelCase__      ):






      @slow
      def __UpperCAmelCase       ( self			:List[Any]	)       ->					str:




             '''simple docstring'''
             with torch.no_grad():
                    _a    : Optional[int]     					=EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
                    model.eval()
                    _a    : Any     					=torch.tensor([[0, 1, 2, 3, 4, 5]]	)
                    _a    : Tuple     					=model(SCREAMING_SNAKE_CASE	)[0]

                    _a    : int     					=3_3

                    _a    : Tuple     					=torch.Size((1, 6, vocab_size)	)
                    self.assertEqual(output.shape     ,						SCREAMING_SNAKE_CASE	)

                    _a    : Union[str, Any]     					=torch.tensor(
                        [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]	)
                    self.assertTrue(torch.allclose(output[:, :3, :3]     ,						SCREAMING_SNAKE_CASE     ,						atol=1e-4	)	)






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




             '''simple docstring'''
             with torch.no_grad():
                    _a    : Any     					=EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""	)
                    model.eval()

                    _a    : Any     					=torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]	)
                    _a    : int     					=model(SCREAMING_SNAKE_CASE	)[0]
                    # compare the actual values for a slice.
                    _a    : str     					=torch.tensor(
                        [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]	)
                    self.assertTrue(torch.allclose(output[:, :3, :3]     ,						SCREAMING_SNAKE_CASE     ,						atol=1e-4	)	)


