


'''simple docstring'''






import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


__A	:	Union[str, Any]     =     logging.getLogger(__name__)
class 						__UpperCamelCase	(						lowercase__							):
							lowercase			:		str   		=  'token-classification'





							def __init__(  self		:Dict  ,_UpperCamelCase		:Union[str, Any]			):
									if type(_UpperCamelCase			) == dict:
											snake_case_		:							Union[str, Any]			       =						Namespace(**_UpperCamelCase			)
									snake_case_		:							List[Any]			       =						import_module("""tasks"""			)
									try:
											snake_case_		:							Union[str, Any]			       =						getattr(_UpperCamelCase  ,hparams.task_type			)
											snake_case_		:							TokenClassificationTask			       =						token_classification_task_clazz()
									except AttributeError:
											raise ValueError(
											    F'''Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. '''
											    F'''Available tasks classes are: {TokenClassificationTask.__subclasses__()}'''			)
									snake_case_		:							Optional[int]			       =						self.token_classification_task.get_labels(hparams.labels			)
									snake_case_		:							int			       =						CrossEntropyLoss().ignore_index
									super().__init__(_UpperCamelCase  ,len(self.labels			)  ,self.mode			)





							def 		a__   (  self		:Dict  ,**_UpperCamelCase		:List[str]			):
									return self.model(**_UpperCamelCase			)





							def 		a__   (  self		:Union[str, Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							List[Any]			       =						{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
									if self.config.model_type != "distilbert":
											snake_case_		:							List[Any]			       =						(
											    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
											)  # XLM and RoBERTa don"t use token_type_ids

									snake_case_		:							List[str]			       =						self(**_UpperCamelCase			)
									snake_case_		:							Any			       =						outputs[0]
									# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
									return {"loss": loss}





							def 		a__   (  self		:List[str]			):
									snake_case_		:							List[Any]			       =						self.hparams
									for mode in ["train", "dev", "test"]:
											snake_case_		:							Dict			       =						self._feature_file(_UpperCamelCase			)
											if os.path.exists(_UpperCamelCase			) and not args.overwrite_cache:
													logger.info("""Loading features from cached file %s"""  ,_UpperCamelCase			)
													snake_case_		:							Optional[Any]			       =						torch.load(_UpperCamelCase			)
											else:
													logger.info("""Creating features from dataset file at %s"""  ,args.data_dir			)
													snake_case_		:							str			       =						self.token_classification_task.read_examples_from_file(args.data_dir  ,_UpperCamelCase			)
													snake_case_		:							List[Any]			       =						self.token_classification_task.convert_examples_to_features(
													    _UpperCamelCase  ,self.labels  ,args.max_seq_length  ,self.tokenizer  ,cls_token_at_end=bool(self.config.model_type in ["""xlnet"""]			)  ,cls_token=self.tokenizer.cls_token  ,cls_token_segment_id=2 if self.config.model_type in ["""xlnet"""] else 0  ,sep_token=self.tokenizer.sep_token  ,sep_token_extra=_UpperCamelCase  ,pad_on_left=bool(self.config.model_type in ["""xlnet"""]			)  ,pad_token=self.tokenizer.pad_token_id  ,pad_token_segment_id=self.tokenizer.pad_token_type_id  ,pad_token_label_id=self.pad_token_label_id  ,)
													logger.info("""Saving features into cached file %s"""  ,_UpperCamelCase			)
													torch.save(_UpperCamelCase  ,_UpperCamelCase			)





							def 		a__   (  self		:List[Any]  ,_UpperCamelCase		:int  ,_UpperCamelCase		:int  ,_UpperCamelCase		:bool = False			):
									snake_case_		:							Tuple			       =						self._feature_file(_UpperCamelCase			)
									logger.info("""Loading features from cached file %s"""  ,_UpperCamelCase			)
									snake_case_		:							Any			       =						torch.load(_UpperCamelCase			)
									snake_case_		:							Optional[int]			       =						torch.tensor([f.input_ids for f in features]  ,dtype=torch.long			)
									snake_case_		:							Optional[int]			       =						torch.tensor([f.attention_mask for f in features]  ,dtype=torch.long			)
									if features[0].token_type_ids is not None:
											snake_case_		:							Any			       =						torch.tensor([f.token_type_ids for f in features]  ,dtype=torch.long			)
									else:
											snake_case_		:							List[Any]			       =						torch.tensor([0 for f in features]  ,dtype=torch.long			)
											# HACK(we will not use this anymore soon)
									snake_case_		:							List[Any]			       =						torch.tensor([f.label_ids for f in features]  ,dtype=torch.long			)
									return DataLoader(
									    TensorDataset(_UpperCamelCase  ,_UpperCamelCase  ,_UpperCamelCase  ,_UpperCamelCase			)  ,batch_size=_UpperCamelCase			)





							def 		a__   (  self		:Optional[int]  ,_UpperCamelCase		:Optional[int]  ,_UpperCamelCase		:str			):
									"""Compute validation""" ""
									snake_case_		:							Optional[int]			       =						{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
									if self.config.model_type != "distilbert":
											snake_case_		:							Optional[int]			       =						(
											    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
											)  # XLM and RoBERTa don"t use token_type_ids
									snake_case_		:							List[str]			       =						self(**_UpperCamelCase			)
									snake_case_       ,      snake_case_		:							Tuple			       =						outputs[:2]
									snake_case_		:							Tuple			       =						logits.detach().cpu().numpy()
									snake_case_		:							Optional[int]			       =						inputs["""labels"""].detach().cpu().numpy()
									return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}





							def 		a__   (  self		:Optional[int]  ,_UpperCamelCase		:Optional[Any]			):
									snake_case_		:							int			       =						torch.stack([x["""val_loss"""] for x in outputs]			).mean()
									snake_case_		:							List[str]			       =						np.concatenate([x["""pred"""] for x in outputs]  ,axis=0			)
									snake_case_		:							str			       =						np.argmax(_UpperCamelCase  ,axis=2			)
									snake_case_		:							List[Any]			       =						np.concatenate([x["""target"""] for x in outputs]  ,axis=0			)

									snake_case_		:							Any			       =						dict(enumerate(self.labels			)			)
									snake_case_		:							str			       =						[[] for _ in range(out_label_ids.shape[0]			)]
									snake_case_		:							int			       =						[[] for _ in range(out_label_ids.shape[0]			)]

									for i in range(out_label_ids.shape[0]			):
											for j in range(out_label_ids.shape[1]			):
													if out_label_ids[i, j] != self.pad_token_label_id:
															out_label_list[i].append(label_map[out_label_ids[i][j]]			)
															preds_list[i].append(label_map[preds[i][j]]			)

									snake_case_		:							Dict			       =						{
									    """val_loss""": val_loss_mean,
									    """accuracy_score""": accuracy_score(_UpperCamelCase  ,_UpperCamelCase			),
									    """precision""": precision_score(_UpperCamelCase  ,_UpperCamelCase			),
									    """recall""": recall_score(_UpperCamelCase  ,_UpperCamelCase			),
									    """f1""": fa_score(_UpperCamelCase  ,_UpperCamelCase			),
									}

									snake_case_		:							str			       =						dict(results.items()			)
									snake_case_		:							str			       =						results
									return ret, preds_list, out_label_list





							def 		a__   (  self		:List[str]  ,_UpperCamelCase		:Tuple			):
									# when stable
									snake_case_       ,      snake_case_       ,      snake_case_		:							Union[str, Any]			       =						self._eval_end(_UpperCamelCase			)
									snake_case_		:							int			       =						ret["""log"""]
									return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





							def 		a__   (  self		:str  ,_UpperCamelCase		:List[Any]			):
									# updating to test_epoch_end instead of deprecated test_end
									snake_case_       ,      snake_case_       ,      snake_case_		:							str			       =						self._eval_end(_UpperCamelCase			)

									# Converting to the dict required by pl
									# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
									# pytorch_lightning/trainer/logging.py#L139
									snake_case_		:							Tuple			       =						ret["""log"""]
									# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
									return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





							@staticmethod
							def 		a__   (  _UpperCamelCase		:Any  ,_UpperCamelCase		:Optional[int]			):
									# Add NER specific options
									BaseTransformer.add_model_specific_args(_UpperCamelCase  ,_UpperCamelCase			)
									parser.add_argument(
									    """--task_type"""  ,default="""NER"""  ,type=_UpperCamelCase  ,help="""Task type to fine tune in training (e.g. NER, POS, etc)"""			)
									parser.add_argument(
									    """--max_seq_length"""  ,default=1_2_8  ,type=_UpperCamelCase  ,help=(
									        """The maximum total input sequence length after tokenization. Sequences longer """
									        """than this will be truncated, sequences shorter will be padded."""
									    )  ,)

									parser.add_argument(
									    """--labels"""  ,default=""""""  ,type=_UpperCamelCase  ,help="""Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."""  ,)
									parser.add_argument(
									    """--gpus"""  ,default=0  ,type=_UpperCamelCase  ,help="""The number of GPUs allocated for this, it is by default 0 meaning none"""  ,)

									parser.add_argument(
									    """--overwrite_cache"""  ,action="""store_true"""  ,help="""Overwrite the cached training and evaluation sets"""			)

									return parser


if __name__ == "__main__":
							__A	:	Union[str, Any]     =     argparse.ArgumentParser()
							add_generic_args(parser, os.getcwd())
							__A	:	int     =     NERTransformer.add_model_specific_args(parser, os.getcwd())
							__A	:	Dict     =     parser.parse_args()
							__A	:	List[str]     =     NERTransformer(args)
							__A	:	str     =     generic_train(model, args)

							if args.do_predict:
														# See https://github.com/huggingface/transformers/issues/3159
														# pl use this default format to create a checkpoint:
														# https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
														# /pytorch_lightning/callbacks/model_checkpoint.py#L322
														__A	:	Union[str, Any]     =     sorted(glob.glob(os.path.join(args.output_dir, 'checkpoint-epoch=*.ckpt'), recursive=True))
														__A	:	Any     =     model.load_from_checkpoint(checkpoints[-1])
														trainer.test(model)



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





							def __init__(  self		:Union[str, Any]  ,_UpperCamelCase		:int  ,_UpperCamelCase		:str=1_3  ,_UpperCamelCase		:Tuple=7  ,_UpperCamelCase		:Tuple=False  ,_UpperCamelCase		:Tuple=True  ,_UpperCamelCase		:Union[str, Any]=False  ,_UpperCamelCase		:int=True  ,_UpperCamelCase		:List[str]=3_3  ,_UpperCamelCase		:Any=3_2  ,_UpperCamelCase		:Any=5  ,_UpperCamelCase		:List[str]=4  ,_UpperCamelCase		:Tuple=3_7  ,_UpperCamelCase		:Optional[Any]="gelu"  ,_UpperCamelCase		:Any=0.1  ,_UpperCamelCase		:List[Any]=0.1  ,_UpperCamelCase		:Any=5_1_2  ,_UpperCamelCase		:Tuple=1_6  ,_UpperCamelCase		:Any=2  ,_UpperCamelCase		:Optional[Any]=0.02  ,_UpperCamelCase		:List[str]=3  ,_UpperCamelCase		:Union[str, Any]=4  ,_UpperCamelCase		:Dict=None  ,):
									snake_case_		:							Tuple			       =						parent
									snake_case_		:							List[str]			       =						batch_size
									snake_case_		:							List[str]			       =						seq_length
									snake_case_		:							Any			       =						is_training
									snake_case_		:							List[Any]			       =						use_input_mask
									snake_case_		:							int			       =						use_token_type_ids
									snake_case_		:							Optional[int]			       =						use_labels
									snake_case_		:							List[str]			       =						vocab_size
									snake_case_		:							Dict			       =						hidden_size
									snake_case_		:							Optional[int]			       =						num_hidden_layers
									snake_case_		:							Optional[Any]			       =						num_attention_heads
									snake_case_		:							int			       =						intermediate_size
									snake_case_		:							Optional[Any]			       =						hidden_act
									snake_case_		:							int			       =						hidden_dropout_prob
									snake_case_		:							Optional[int]			       =						attention_probs_dropout_prob
									snake_case_		:							Optional[int]			       =						max_position_embeddings
									snake_case_		:							List[str]			       =						type_vocab_size
									snake_case_		:							int			       =						type_sequence_label_size
									snake_case_		:							Dict			       =						initializer_range
									snake_case_		:							Any			       =						num_labels
									snake_case_		:							Any			       =						num_choices
									snake_case_		:							Tuple			       =						scope





							def 		a__   (  self		:str			):
									snake_case_		:							Dict			       =						ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size			)

									snake_case_		:							str			       =						None
									if self.use_input_mask:
											snake_case_		:							Optional[int]			       =						random_attention_mask([self.batch_size, self.seq_length]			)

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

									snake_case_		:							Dict			       =						self.get_config()

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





							def 		a__   (  self		:Optional[Any]			):
									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 		a__   (  self		:str  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Optional[int]  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							Union[str, Any]			       =						EsmModel(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase			)
									snake_case_		:							str			       =						model(_UpperCamelCase			)
									snake_case_		:							Optional[int]			       =						model(_UpperCamelCase			)

									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 		a__   (  self		:Any  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Dict  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							str			       =						EsmForMaskedLM(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							Dict			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)			)





							def 		a__   (  self		:Optional[int]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Any			):
									snake_case_		:							List[Any]			       =						self.num_labels
									snake_case_		:							int			       =						EsmForTokenClassification(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.num_labels)			)





							def 		a__   (  self		:Union[str, Any]			):
									snake_case_		:							Dict			       =						self.prepare_config_and_inputs()
									(
									    (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      
									)		:							Optional[int]			       =						config_and_inputs
									snake_case_		:							List[str]			       =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict





@require_torch
class 						__UpperCamelCase	(						lowercase__						,  lowercase__						,  unittest.TestCase							):
							lowercase			:		Optional[int]   		=  False

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





							def 		a__   (  self		:Any			):
									snake_case_		:							Any			       =						EsmModelTester(self			)
									snake_case_		:							str			       =						ConfigTester(self  ,config_class=_UpperCamelCase  ,hidden_size=3_7			)





							def 		a__   (  self		:Optional[Any]			):
									self.config_tester.run_common_tests()





							def 		a__   (  self		:Optional[Any]			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
											snake_case_		:							int			       =						type
											self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase			)





							def 		a__   (  self		:List[str]			):
									snake_case_		:							List[str]			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase			)





							@slow
							def 		a__   (  self		:Union[str, Any]			):
									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											snake_case_		:							int			       =						EsmModel.from_pretrained(_UpperCamelCase			)
											self.assertIsNotNone(_UpperCamelCase			)





							def 		a__   (  self		:Tuple			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							Optional[int]			       =						EsmEmbeddings(config=_UpperCamelCase			)

									snake_case_		:							List[Any]			       =						torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]			)
									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,
									        ]
									    ]			)
									snake_case_		:							Optional[Any]			       =						create_position_ids_from_input_ids(_UpperCamelCase  ,model.padding_idx			)
									self.assertEqual(position_ids.shape  ,expected_positions.shape			)
									self.assertTrue(torch.all(torch.eq(_UpperCamelCase  ,_UpperCamelCase			)			)			)





							def 		a__   (  self		:List[Any]			):
									snake_case_		:							Union[str, Any]			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							List[Any]			       =						EsmEmbeddings(config=_UpperCamelCase			)

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





							@unittest.skip("""Esm does not support embedding resizing"""			)
							def 		a__   (  self		:Optional[Any]			):
									pass





							@unittest.skip("""Esm does not support embedding resizing"""			)
							def 		a__   (  self		:Optional[int]			):
									pass





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





@require_torch
class 						__UpperCamelCase	(						lowercase__							):





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

											snake_case_		:							Optional[int]			       =						3_3

											snake_case_		:							List[Any]			       =						torch.Size((1, 6, vocab_size)			)
											self.assertEqual(output.shape  ,_UpperCamelCase			)

											snake_case_		:							Union[str, Any]			       =						torch.tensor(
											    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)





							@slow
							def 		a__   (  self		:List[Any]			):
									with torch.no_grad():
											snake_case_		:							List[Any]			       =						EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""			)
											model.eval()

											snake_case_		:							str			       =						torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]			)
											snake_case_		:							Optional[Any]			       =						model(_UpperCamelCase			)[0]
											# compare the actual values for a slice.
											snake_case_		:							List[str]			       =						torch.tensor(
											    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)