






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

import tqdm
from filelock import FileLock

from transformers import (
    BartTokenizer,
    BartTokenizerFast,
    DataProcessor,
    PreTrainedTokenizer,
    RobertaTokenizer,
    RobertaTokenizerFast,
    XLMRobertaTokenizer,
    is_tf_available,
    is_torch_available,
)


UpperCamelCase			    =     logging.getLogger(__name__)



@dataclass(frozen=lowerCamelCase__   )
class 				__UpperCAmelCase		:
			__snake_case    :				Optional[Any]									=							42
			__snake_case    :				Dict									=							42
			__snake_case    :				Union[str, Any]									=							None
			__snake_case    :				str									=							None
			__snake_case    :				Tuple									=							None



@dataclass(frozen=lowerCamelCase__   )
class 				__UpperCAmelCase		:
			__snake_case    :				Optional[int]									=							42
			__snake_case    :				Any									=							None
			__snake_case    :				int									=							None
			__snake_case    :				Optional[Any]									=							None
			__snake_case    :				List[Any]									=							None


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



		class 				__UpperCAmelCase		(lowerCamelCase__   ):
					__snake_case    :				Optional[Any]									=							42







					def __init__(     self:  Dict				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  PreTrainedTokenizer				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[int] = None				,     UpperCAmelCase_:  Any=False				,     UpperCAmelCase_:  bool = False				,     ):







											'''simple docstring'''




											_SCREAMING_SNAKE_CASE     				=  hans_processors[task]()

											_SCREAMING_SNAKE_CASE     				=  os.path.join(
											    UpperCAmelCase_				,     """cached_{}_{}_{}_{}""".format(
											        """dev""" if evaluate else """train"""				,     tokenizer.__class__.__name__				,     str(UpperCAmelCase_					)				,     UpperCAmelCase_				,     )				,     )
											_SCREAMING_SNAKE_CASE     				=  processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																	# HACK(label indices are swapped in RoBERTa pretrained model)
																	_SCREAMING_SNAKE_CASE     				=  label_list[2], label_list[1]
											_SCREAMING_SNAKE_CASE     				=  label_list

											# Make sure only the first process in distributed training processes the dataset,
											# and the others will use the cache.
											_SCREAMING_SNAKE_CASE     				=  cached_features_file + ".lock"
											with FileLock(UpperCAmelCase_					):
																	if os.path.exists(UpperCAmelCase_					) and not overwrite_cache:
																							logger.info(F'Loading features from cached file {cached_features_file}'					)
																							_SCREAMING_SNAKE_CASE     				=  torch.load(UpperCAmelCase_					)
																	else:
																							logger.info(F'Creating features from dataset file at {data_dir}'					)

																							_SCREAMING_SNAKE_CASE     				=  (
																							    processor.get_dev_examples(UpperCAmelCase_					) if evaluate else processor.get_train_examples(UpperCAmelCase_					)
																							)

																							logger.info("""Training examples: %s"""				,     len(UpperCAmelCase_					)					)
																							_SCREAMING_SNAKE_CASE     				=  hans_convert_examples_to_features(UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_					)
																							logger.info("""Saving features into cached file %s"""				,     UpperCAmelCase_					)
																							torch.save(self.features				,     UpperCAmelCase_					)







					def __len__(     self:  str					):







											'''simple docstring'''




											return len(self.features					)







					def __getitem__(     self:  List[str]				,     UpperCAmelCase_:  Optional[Any]					):







											'''simple docstring'''




											return self.features[i]







					def 							UpperCamelCase				(     self:  Optional[int]					):







											'''simple docstring'''




											return self.label_list


if is_tf_available():
		import tensorflow as tf



		class 				__UpperCAmelCase		:
					__snake_case    :				List[str]									=							42







					def __init__(     self:  int				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  PreTrainedTokenizer				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[int] = 128				,     UpperCAmelCase_:  Dict=False				,     UpperCAmelCase_:  bool = False				,     ):







											'''simple docstring'''




											_SCREAMING_SNAKE_CASE     				=  hans_processors[task]()
											_SCREAMING_SNAKE_CASE     				=  processor.get_labels()
											if tokenizer.__class__ in (
											    RobertaTokenizer,
											    RobertaTokenizerFast,
											    XLMRobertaTokenizer,
											    BartTokenizer,
											    BartTokenizerFast,
											):
																	# HACK(label indices are swapped in RoBERTa pretrained model)
																	_SCREAMING_SNAKE_CASE     				=  label_list[2], label_list[1]
											_SCREAMING_SNAKE_CASE     				=  label_list

											_SCREAMING_SNAKE_CASE     				=  processor.get_dev_examples(UpperCAmelCase_					) if evaluate else processor.get_train_examples(UpperCAmelCase_					)
											_SCREAMING_SNAKE_CASE     				=  hans_convert_examples_to_features(UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_				,     UpperCAmelCase_					)

											def gen():
																	for ex_index, ex in tqdm.tqdm(enumerate(self.features					)				,     desc="""convert examples to features"""					):
																							if ex_index % 10_000 == 0:
																													logger.info("""Writing example %d of %d""" % (ex_index, len(UpperCAmelCase_					))					)

																							yield (
																							    {
																							        "example_id": 0,
																							        "input_ids": ex.input_ids,
																							        "attention_mask": ex.attention_mask,
																							        "token_type_ids": ex.token_type_ids,
																							    },
																							    ex.label,
																							)

											_SCREAMING_SNAKE_CASE     				=  tf.data.Dataset.from_generator(
											    UpperCAmelCase_				,     (
											        {
											            """example_id""": tf.intaa,
											            """input_ids""": tf.intaa,
											            """attention_mask""": tf.intaa,
											            """token_type_ids""": tf.intaa,
											        },
											        tf.intaa,
											    )				,     (
											        {
											            """example_id""": tf.TensorShape([]					),
											            """input_ids""": tf.TensorShape([None, None]					),
											            """attention_mask""": tf.TensorShape([None, None]					),
											            """token_type_ids""": tf.TensorShape([None, None]					),
											        },
											        tf.TensorShape([]					),
											    )				,     )







					def 							UpperCamelCase				(     self:  Optional[Any]					):







											'''simple docstring'''




											return self.dataset







					def __len__(     self:  int					):







											'''simple docstring'''




											return len(self.features					)







					def __getitem__(     self:  Optional[Any]				,     UpperCAmelCase_:  List[Any]					):







											'''simple docstring'''




											return self.features[i]







					def 							UpperCamelCase				(     self:  Any					):







											'''simple docstring'''




											return self.label_list



class 				__UpperCAmelCase		(lowerCamelCase__   ):







			def 							UpperCamelCase				(     self:  Any				,     UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									return self._create_examples(self._read_tsv(os.path.join(UpperCAmelCase_				,     """heuristics_train_set.txt"""					)					)				,     """train"""					)







			def 							UpperCamelCase				(     self:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]					):







									'''simple docstring'''




									return self._create_examples(self._read_tsv(os.path.join(UpperCAmelCase_				,     """heuristics_evaluation_set.txt"""					)					)				,     """dev"""					)







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







									'''simple docstring'''




									return ["contradiction", "entailment", "neutral"]







			def 							UpperCamelCase				(     self:  Optional[int]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  []
									for i, line in enumerate(UpperCAmelCase_					):
															if i == 0:
																					continue
															_SCREAMING_SNAKE_CASE     				=  "%s-%s" % (set_type, line[0])
															_SCREAMING_SNAKE_CASE     				=  line[5]
															_SCREAMING_SNAKE_CASE     				=  line[6]
															_SCREAMING_SNAKE_CASE     				=  line[7][2:] if line[7].startswith("""ex"""					) else line[7]
															_SCREAMING_SNAKE_CASE     				=  line[0]
															examples.append(InputExample(guid=UpperCAmelCase_				,     text_a=UpperCAmelCase_				,     text_b=UpperCAmelCase_				,     label=UpperCAmelCase_				,     pairID=UpperCAmelCase_					)					)
									return examples

def 	__lowerCamelCase      (		snake_case__	,snake_case__	,snake_case__	,snake_case__	,)							->  Optional[int]:






						"""simple docstring"""





						_SCREAMING_SNAKE_CASE     				=  {label: i for i, label in enumerate(_lowerCAmelCase		)}

						_SCREAMING_SNAKE_CASE     				=  []
						for ex_index, example in tqdm.tqdm(enumerate(_lowerCAmelCase		)	,desc="""convert examples to features"""		):
												if ex_index % 1_00_00 == 0:
																		logger.info("""Writing example %d""" % (ex_index)		)

												_SCREAMING_SNAKE_CASE     				=  tokenizer(
												    example.text_a	,example.text_b	,add_special_tokens=_lowerCAmelCase	,max_length=_lowerCAmelCase	,padding="""max_length"""	,truncation=_lowerCAmelCase	,return_overflowing_tokens=_lowerCAmelCase	,)

												_SCREAMING_SNAKE_CASE     				=  label_map[example.label] if example.label in label_map else 0

												_SCREAMING_SNAKE_CASE     				=  int(example.pairID		)

												features.append(InputFeatures(**_lowerCAmelCase	,label=_lowerCAmelCase	,pairID=_lowerCAmelCase		)		)

						for i, example in enumerate(examples[:5]		):
												logger.info("""*** Example ***"""		)
												logger.info(F'guid: {example}'		)
												logger.info(F'features: {features[i]}'		)

						return features


UpperCamelCase			    =     {
    '''hans''': 3,
}




UpperCamelCase			    =     {
    '''hans''': HansProcessor,
}









from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
		import numpy
		import tensorflow as tf

		from transformers.models.esm.modeling_tf_esm import (
		    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFEsmForMaskedLM,
		    TFEsmForSequenceClassification,
		    TFEsmForTokenClassification,
		    TFEsmModel,
		)



class 				__UpperCAmelCase		:







			def __init__(     self:  Tuple				,     UpperCAmelCase_:  Tuple				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  parent
									_SCREAMING_SNAKE_CASE     				=  13
									_SCREAMING_SNAKE_CASE     				=  7
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  99
									_SCREAMING_SNAKE_CASE     				=  32
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  37
									_SCREAMING_SNAKE_CASE     				=  """gelu"""
									_SCREAMING_SNAKE_CASE     				=  0.1
									_SCREAMING_SNAKE_CASE     				=  0.1
									_SCREAMING_SNAKE_CASE     				=  512
									_SCREAMING_SNAKE_CASE     				=  16
									_SCREAMING_SNAKE_CASE     				=  2
									_SCREAMING_SNAKE_CASE     				=  0.02
									_SCREAMING_SNAKE_CASE     				=  3
									_SCREAMING_SNAKE_CASE     				=  4
									_SCREAMING_SNAKE_CASE     				=  None







			def 							UpperCamelCase				(     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     				=  EsmConfig(
									    vocab_size=self.vocab_size				,     hidden_size=self.hidden_size				,     num_hidden_layers=self.num_hidden_layers				,     pad_token_id=1				,     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				,     )

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







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									(
									    (
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					(
									    _SCREAMING_SNAKE_CASE
									)				,					
									)     				=  self.prepare_config_and_inputs()

									_SCREAMING_SNAKE_CASE     				=  True
									_SCREAMING_SNAKE_CASE     				=  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]					)
									_SCREAMING_SNAKE_CASE     				=  ids_tensor([self.batch_size, self.seq_length]				,     vocab_size=2					)

									return (
									    config,
									    input_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    encoder_hidden_states,
									    encoder_attention_mask,
									)







			def 							UpperCamelCase				(     self:  int				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Dict					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModel(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  [input_ids, input_mask]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size)					)







			def 							UpperCamelCase				(     self:  Optional[Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Dict				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  Any				,     UpperCAmelCase_:  List[str]				,     ):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  True

									_SCREAMING_SNAKE_CASE     				=  TFEsmModel(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """encoder_hidden_states""": encoder_hidden_states,
									    """encoder_attention_mask""": encoder_attention_mask,
									}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)

									_SCREAMING_SNAKE_CASE     				=  [input_ids, input_mask]
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     encoder_hidden_states=UpperCAmelCase_					)

									# Also check the case where encoder outputs are not passed
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_				,     attention_mask=UpperCAmelCase_					)

									self.parent.assertEqual(result.last_hidden_state.shape				,     (self.batch_size, self.seq_length, self.hidden_size)					)







			def 							UpperCamelCase				(     self:  Dict				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  List[Any]				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  List[str]				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  List[str]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmForMaskedLM(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  model([input_ids, input_mask]					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.vocab_size)					)







			def 							UpperCamelCase				(     self:  List[Any]				,     UpperCAmelCase_:  Tuple				,     UpperCAmelCase_:  str				,     UpperCAmelCase_:  Optional[int]				,     UpperCAmelCase_:  Union[str, Any]				,     UpperCAmelCase_:  int				,     UpperCAmelCase_:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.num_labels
									_SCREAMING_SNAKE_CASE     				=  TFEsmForTokenClassification(config=UpperCAmelCase_					)
									_SCREAMING_SNAKE_CASE     				=  {"""input_ids""": input_ids, """attention_mask""": input_mask}
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)
									self.parent.assertEqual(result.logits.shape				,     (self.batch_size, self.seq_length, self.num_labels)					)







			def 							UpperCamelCase				(     self:  Dict					):







									'''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_tf
class 				__UpperCAmelCase		(_UpperCAmelCase	,_UpperCAmelCase	,unittest.TestCase   ):
			__snake_case    :				List[Any]									=							(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			__snake_case    :				Tuple									=							(
			    {
			        "feature-extraction": TFEsmModel,
			        "fill-mask": TFEsmForMaskedLM,
			        "text-classification": TFEsmForSequenceClassification,
			        "token-classification": TFEsmForTokenClassification,
			        "zero-shot": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			__snake_case    :				List[str]									=							False
			__snake_case    :				Union[str, Any]									=							False







			def 							UpperCamelCase				(     self:  Optional[int]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModelTester(self					)
									_SCREAMING_SNAKE_CASE     				=  ConfigTester(self				,     config_class=UpperCAmelCase_				,     hidden_size=37					)







			def 							UpperCamelCase				(     self:  List[Any]					):







									'''simple docstring'''




									self.config_tester.run_common_tests()







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







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_decoder()
									self.model_tester.create_and_check_model_as_decoder(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase_					)







			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase_					)







			@slow
			def 							UpperCamelCase				(     self:  Tuple					):







									'''simple docstring'''




									for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_SCREAMING_SNAKE_CASE     				=  TFEsmModel.from_pretrained(UpperCAmelCase_					)
															self.assertIsNotNone(UpperCAmelCase_					)







			@unittest.skip("""Protein models do not support embedding resizing."""					)
			def 							UpperCamelCase				(     self:  Dict					):







									'''simple docstring'''




									pass







			@unittest.skip("""Protein models do not support embedding resizing."""					)
			def 							UpperCamelCase				(     self:  Optional[Any]					):







									'''simple docstring'''




									pass







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







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE				,					_SCREAMING_SNAKE_CASE     				=  self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
															_SCREAMING_SNAKE_CASE     				=  model_class(UpperCAmelCase_					)
															assert isinstance(model.get_input_embeddings()				,     tf.keras.layers.Layer					)
															if model_class is TFEsmForMaskedLM:
																					# Output embedding test differs from the main test because they're a matrix, not a layer
																					_SCREAMING_SNAKE_CASE     				=  model.get_bias()
																					assert isinstance(UpperCAmelCase_				,     UpperCAmelCase_					)
																					for k, v in name.items():
																											assert isinstance(UpperCAmelCase_				,     tf.Variable					)
															else:
																					_SCREAMING_SNAKE_CASE     				=  model.get_output_embeddings()
																					assert x is None
																					_SCREAMING_SNAKE_CASE     				=  model.get_bias()
																					assert name is None



@require_tf
class 				__UpperCAmelCase		(unittest.TestCase   ):







			@slow
			def 							UpperCamelCase				(     self:  Any					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant([[0, 1, 2, 3, 4, 5]]					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]
									_SCREAMING_SNAKE_CASE     				=  [1, 6, 33]
									self.assertEqual(list(output.numpy().shape					)				,     UpperCAmelCase_					)
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
									            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
									            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,     expected_slice.numpy()				,     atol=1E-2					)					)







			@slow
			def 							UpperCamelCase				(     self:  str					):







									'''simple docstring'''




									_SCREAMING_SNAKE_CASE     				=  TFEsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""					)

									_SCREAMING_SNAKE_CASE     				=  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
									_SCREAMING_SNAKE_CASE     				=  model(UpperCAmelCase_					)[0]
									# compare the actual values for a slice.
									_SCREAMING_SNAKE_CASE     				=  tf.constant(
									    [
									        [
									            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
									            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
									            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
									        ]
									    ]					)
									self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,     expected_slice.numpy()				,     atol=1E-4					)					)

