


'''simple docstring'''







import argparse
import gc
import json
import os

import torch
from datasets import load_dataset
from torch.optim import AdamW
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed

from accelerate import Accelerator, DistributedType
from accelerate.utils.deepspeed import DummyOptim, DummyScheduler


A__:    Dict					    =       16
A__:    Tuple					    =       32


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[int]  )  ->			List[str]:
       return int(x / 2**20  )







class A__       :






      def __enter__( self			:str	)       ->					List[str]:




             '''simple docstring'''
             gc.collect()
             torch.cuda.empty_cache()
             torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
             _a    : str     					=torch.cuda.memory_allocated()
             return self






      def __exit__( self			:Any     ,						*SCREAMING_SNAKE_CASE			:List[str]	)       ->					List[Any]:




             '''simple docstring'''
             gc.collect()
             torch.cuda.empty_cache()
             _a    : str     					=torch.cuda.memory_allocated()
             _a    : Tuple     					=torch.cuda.max_memory_allocated()
             _a    : Optional[Any]     					=bamb(self.end - self.begin	)
             _a    : List[Any]     					=bamb(self.peak - self.begin	)
             # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Accelerator    ,_UpperCAmelCase	:	int = 16    ,_UpperCAmelCase	:	str = "bert-base-cased"    ,_UpperCAmelCase	:	int = 320    ,_UpperCAmelCase	:	int = 160    ,)  ->			Union[str, Any]:
       _a    : List[Any]     					=AutoTokenizer.from_pretrained(_UpperCAmelCase  )
       _a    : str     					=load_dataset(
           """glue"""    ,"""mrpc"""    ,split={"""train""": F"train[:{n_train}]", """validation""": F"validation[:{n_val}]"}  )

       def tokenize_function(_UpperCAmelCase	:	Any  ):
              # max_length=None => use the model max length (it's actually the default)
              _a    : int     					=tokenizer(examples["""sentence1"""]    ,examples["""sentence2"""]    ,truncation=_UpperCAmelCase    ,max_length=_UpperCAmelCase  )
              return outputs

       # Apply the method we just defined to all the examples in all the splits of the dataset
       _a    : Union[str, Any]     					=datasets.map(
           _UpperCAmelCase    ,batched=_UpperCAmelCase    ,remove_columns=["""idx""", """sentence1""", """sentence2"""]    ,load_from_cache_file=_UpperCAmelCase  )

       # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
       # transformers library
       _a    : Any     					=tokenized_datasets.rename_column("""label"""    ,"""labels"""  )

       def collate_fn(_UpperCAmelCase	:	Optional[Any]  ):
              # On TPU it's best to pad everything to the same length or training will be very slow.
              if accelerator.distributed_type == DistributedType.TPU:
                     return tokenizer.pad(_UpperCAmelCase    ,padding="""max_length"""    ,max_length=128    ,return_tensors="""pt"""  )
              return tokenizer.pad(_UpperCAmelCase    ,padding="""longest"""    ,return_tensors="""pt"""  )

       # Instantiate dataloaders.
       _a    : Dict     					=DataLoader(
           tokenized_datasets["""train"""]    ,shuffle=_UpperCAmelCase    ,collate_fn=_UpperCAmelCase    ,batch_size=_UpperCAmelCase  )
       _a    : Dict     					=DataLoader(
           tokenized_datasets["""validation"""]    ,shuffle=_UpperCAmelCase    ,collate_fn=_UpperCAmelCase    ,batch_size=_UpperCAmelCase  )

       return train_dataloader, eval_dataloader


def 	SCREAMING_SNAKE_CASE_					(  _UpperCAmelCase	:	Optional[Any]    ,_UpperCAmelCase	:	str  )  ->			int:
       # Initialize accelerator
       _a    : Optional[Any]     					=Accelerator()

       # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
       _a    : Optional[Any]     					=config["""lr"""]
       _a    : List[Any]     					=int(config["""num_epochs"""]  )
       _a    : Any     					=int(config["""seed"""]  )
       _a    : List[str]     					=int(config["""batch_size"""]  )
       _a    : Dict     					=args.model_name_or_path

       set_seed(_UpperCAmelCase  )
       _a					,    _a    : Optional[int]     					=get_dataloaders(_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,args.n_train    ,args.n_val  )

       # Instantiate the model (we build the model here so that the seed also control new weights initialization)
       _a    : Optional[Any]     					=AutoModelForSequenceClassification.from_pretrained(_UpperCAmelCase    ,return_dict=_UpperCAmelCase  )

       # Instantiate optimizer
       _a    : Tuple     					=(
           AdamW
           if accelerator.state.deepspeed_plugin is None
           or """optimizer""" not in accelerator.state.deepspeed_plugin.deepspeed_config
           else DummyOptim
       )
       _a    : int     					=optimizer_cls(params=model.parameters()    ,lr=_UpperCAmelCase  )

       if accelerator.state.deepspeed_plugin is not None:
              _a    : str     					=accelerator.state.deepspeed_plugin.deepspeed_config[
                  """gradient_accumulation_steps"""
              ]
       else:
              _a    : Dict     					=1
       _a    : List[str]     					=(len(_UpperCAmelCase  ) * num_epochs) // gradient_accumulation_steps

       # Instantiate scheduler
       if (
           accelerator.state.deepspeed_plugin is None
           or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
       ):
              _a    : int     					=get_linear_schedule_with_warmup(
                  optimizer=_UpperCAmelCase    ,num_warmup_steps=0    ,num_training_steps=_UpperCAmelCase    ,)
       else:
              _a    : List[Any]     					=DummyScheduler(_UpperCAmelCase    ,total_num_steps=_UpperCAmelCase    ,warmup_num_steps=0  )

       # Prepare everything
       # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
       # prepare method.
       _a					,    _a					,    _a					,    _a					,    _a    : List[str]     					=accelerator.prepare(
           _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase  )

       # We need to keep track of how many total steps we have iterated over
       _a    : Union[str, Any]     					=0
       # We also need to keep track of the stating epoch so files are named properly
       _a    : Optional[Any]     					=0

       # Now we train the model
       _a    : Tuple     					={}
       for epoch in range(_UpperCAmelCase    ,_UpperCAmelCase  ):
              with TorchTracemalloc() as tracemalloc:
                     model.train()
                     for step, batch in enumerate(_UpperCAmelCase  ):
                            _a    : Tuple     					=model(**_UpperCAmelCase  )
                            _a    : Union[str, Any]     					=outputs.loss
                            _a    : Dict     					=loss / gradient_accumulation_steps
                            accelerator.backward(_UpperCAmelCase  )
                            if step % gradient_accumulation_steps == 0:
                                   optimizer.step()
                                   lr_scheduler.step()
                                   optimizer.zero_grad()

                            overall_step += 1

        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
              accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin  )  )  )
              accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used  )  )
              accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked  )  )
              accelerator.print(
                  """Total Peak Memory consumed during the train (max): {}""".format(
                      tracemalloc.peaked + bamb(tracemalloc.begin  )  )  )
              _a    : Any     					=tracemalloc.peaked + bamb(tracemalloc.begin  )
              if args.peak_memory_upper_bound is not None:
                     assert (
                         train_total_peak_memory[F"epoch-{epoch}"] <= args.peak_memory_upper_bound
                     ), "Peak memory usage exceeded the upper bound"

       accelerator.wait_for_everyone()
       if accelerator.is_main_process:
              with open(os.path.join(args.output_dir    ,"""peak_memory_utilization.json"""  )    ,"""w"""  ) as f:
                     json.dump(_UpperCAmelCase    ,_UpperCAmelCase  )


def 	SCREAMING_SNAKE_CASE_					(  )  ->			Any:
       _a    : Dict     					=argparse.ArgumentParser(description="""Simple example of training script tracking peak GPU memory usage."""  )
       parser.add_argument(
           """--model_name_or_path"""    ,type=_UpperCAmelCase    ,default="""bert-base-cased"""    ,help="""Path to pretrained model or model identifier from huggingface.co/models."""    ,required=_UpperCAmelCase    ,)
       parser.add_argument(
           """--output_dir"""    ,type=_UpperCAmelCase    ,default="""."""    ,help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""    ,)
       parser.add_argument(
           """--peak_memory_upper_bound"""    ,type=_UpperCAmelCase    ,default=_UpperCAmelCase    ,help="""The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value."""    ,)
       parser.add_argument(
           """--n_train"""    ,type=_UpperCAmelCase    ,default=320    ,help="""Number of training examples to use."""    ,)
       parser.add_argument(
           """--n_val"""    ,type=_UpperCAmelCase    ,default=160    ,help="""Number of validation examples to use."""    ,)
       parser.add_argument(
           """--num_epochs"""    ,type=_UpperCAmelCase    ,default=1    ,help="""Number of train epochs."""    ,)
       _a    : Any     					=parser.parse_args()
       _a    : Union[str, Any]     					={"""lr""": 2e-5, """num_epochs""": args.num_epochs, """seed""": 42, """batch_size""": 16}
       training_function(_UpperCAmelCase    ,_UpperCAmelCase  )


if __name__ == "__main__":
   main()






'''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	)	)


