


"""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__						:


     """simple docstring"""

     def __init__(						self:    Tuple  ,		__a:    int  ,		__a:    Any=13  ,		__a:    Dict=7  ,		__a:    Optional[Any]=False  ,		__a:    List[str]=True  ,		__a:    List[Any]=False  ,		__a:    str=True  ,		__a:    Optional[Any]=33  ,		__a:    int=32  ,		__a:    List[str]=5  ,		__a:    List[str]=4  ,		__a:    int=37  ,		__a:    List[str]="gelu"  ,		__a:    Dict=0.1  ,		__a:    Dict=0.1  ,		__a:    Union[str, Any]=512  ,		__a:    Any=16  ,		__a:    Optional[int]=2  ,		__a:    List[Any]=0.02  ,		__a:    Optional[int]=3  ,		__a:    int=4  ,		__a:    Any=None  ,		)->     Dict:
           lowerCamelCase       :       int   =						parent
           lowerCamelCase       :       Dict   =						batch_size
           lowerCamelCase       :       Dict   =						seq_length
           lowerCamelCase       :       Tuple   =						is_training
           lowerCamelCase       :       List[str]   =						use_input_mask
           lowerCamelCase       :       Optional[Any]   =						use_token_type_ids
           lowerCamelCase       :       Tuple   =						use_labels
           lowerCamelCase       :       Optional[Any]   =						vocab_size
           lowerCamelCase       :       List[str]   =						hidden_size
           lowerCamelCase       :       int   =						num_hidden_layers
           lowerCamelCase       :       Dict   =						num_attention_heads
           lowerCamelCase       :       Tuple   =						intermediate_size
           lowerCamelCase       :       int   =						hidden_act
           lowerCamelCase       :       str   =						hidden_dropout_prob
           lowerCamelCase       :       str   =						attention_probs_dropout_prob
           lowerCamelCase       :       List[str]   =						max_position_embeddings
           lowerCamelCase       :       Optional[int]   =						type_vocab_size
           lowerCamelCase       :       Optional[int]   =						type_sequence_label_size
           lowerCamelCase       :       Any   =						initializer_range
           lowerCamelCase       :       str   =						num_labels
           lowerCamelCase       :       str   =						num_choices
           lowerCamelCase       :       str   =						scope

     def 		a__    (						self:    Union[str, Any]						)->     Optional[int]:
           lowerCamelCase       :       Tuple   =						ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size						)

           lowerCamelCase       :       Union[str, Any]   =						None
           if self.use_input_mask:
                 lowerCamelCase       :       Optional[Any]   =						random_attention_mask([self.batch_size, self.seq_length]						)

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

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

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

     def 		a__    (						self:    Optional[Any]						)->     Optional[int]:
           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:    Optional[int]  ,		__a:    Optional[Any]  ,		__a:    List[str]  ,		__a:    str  ,		__a:    Tuple  ,		__a:    Optional[int]  ,		__a:    Union[str, Any]						)->     Tuple:
           lowerCamelCase       :       List[str]   =						EsmModel(config=SCREAMING_SNAKE_CASE_						)
           model.to(SCREAMING_SNAKE_CASE_						)
           model.eval()
           lowerCamelCase       :       Optional[int]   =						model(SCREAMING_SNAKE_CASE_  ,		attention_mask=SCREAMING_SNAKE_CASE_						)
           lowerCamelCase       :       Optional[Any]   =						model(SCREAMING_SNAKE_CASE_						)
           lowerCamelCase       :       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 		a__    (						self:    str  ,		__a:    Union[str, Any]  ,		__a:    Union[str, Any]  ,		__a:    List[str]  ,		__a:    Optional[int]  ,		__a:    str  ,		__a:    Dict						)->     Tuple:
           lowerCamelCase       :       List[Any]   =						EsmForMaskedLM(config=SCREAMING_SNAKE_CASE_						)
           model.to(SCREAMING_SNAKE_CASE_						)
           model.eval()
           lowerCamelCase       :       Dict   =						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 		a__    (						self:    List[Any]  ,		__a:    Tuple  ,		__a:    str  ,		__a:    Dict  ,		__a:    Tuple  ,		__a:    Optional[Any]  ,		__a:    str						)->     str:
           lowerCamelCase       :       int   =						self.num_labels
           lowerCamelCase       :       Optional[int]   =						EsmForTokenClassification(config=SCREAMING_SNAKE_CASE_						)
           model.to(SCREAMING_SNAKE_CASE_						)
           model.eval()
           lowerCamelCase       :       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.num_labels)						)







     def 		a__    (						self:    Dict						)->     Dict:
           lowerCamelCase       :       Optional[int]   =						self.prepare_config_and_inputs()
           (
               (
               lowerCamelCase       
           )  ,      (
               lowerCamelCase       
           )  ,      (
               lowerCamelCase       
           )  ,      (
               lowerCamelCase       
           )  ,      (
               lowerCamelCase       
           )  ,      (
               lowerCamelCase       
           )  ,             
           )       :       Optional[int]   =						config_and_inputs
           lowerCamelCase       :       int   =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
           return config, inputs_dict


@require_torch
class       A__						(    lowercase__			,  lowercase__			,  unittest.TestCase):


     """simple docstring"""
     snake_case__						:       Union[str, Any]				    =False

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

     def 		a__    (						self:    Any						)->     Dict:
           lowerCamelCase       :       List[str]   =						EsmModelTester(self						)
           lowerCamelCase       :       int   =						ConfigTester(self  ,		config_class=SCREAMING_SNAKE_CASE_  ,		hidden_size=37						)

     def 		a__    (						self:    Union[str, Any]						)->     Tuple:
           self.config_tester.run_common_tests()

     def 		a__    (						self:    Tuple						)->     List[str]:
           lowerCamelCase       :       Dict   =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_						)

     def 		a__    (						self:    List[str]						)->     List[str]:
           lowerCamelCase       :       Union[str, Any]   =						self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                 lowerCamelCase       :       List[str]   =						type
                 self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_						)

     def 		a__    (						self:    Any						)->     List[str]:
           lowerCamelCase       :       Dict   =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_						)

     def 		a__    (						self:    Union[str, Any]						)->     int:
           lowerCamelCase       :       List[Any]   =						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_						)

     @slow
     def 		a__    (						self:    Any						)->     str:


           for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                 lowerCamelCase       :       Optional[Any]   =						EsmModel.from_pretrained(SCREAMING_SNAKE_CASE_						)
                 self.assertIsNotNone(SCREAMING_SNAKE_CASE_						)

     def 		a__    (						self:    List[Any]						)->     Tuple:
           lowerCamelCase       :       List[Any]   =						self.model_tester.prepare_config_and_inputs()[0]
           lowerCamelCase       :       List[Any]   =						EsmEmbeddings(config=SCREAMING_SNAKE_CASE_						)

           lowerCamelCase       :       List[Any]   =						torch.as_tensor([[12, 31, 13, model.padding_idx]]						)
           lowerCamelCase       :       Optional[Any]   =						torch.as_tensor(
               [
                   [
                       0 + model.padding_idx + 1,
                       1 + model.padding_idx + 1,
                       2 + model.padding_idx + 1,
                       model.padding_idx,
                   ]
               ]						)
           lowerCamelCase       :       Dict   =						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 		a__    (						self:    Dict						)->     Any:
           lowerCamelCase       :       Optional[int]   =						self.model_tester.prepare_config_and_inputs()[0]
           lowerCamelCase       :       Dict   =						EsmEmbeddings(config=SCREAMING_SNAKE_CASE_						)

           lowerCamelCase       :       str   =						torch.empty(2  ,		4  ,		30						)
           lowerCamelCase       :       Dict   =						[
               0 + embeddings.padding_idx + 1,
               1 + embeddings.padding_idx + 1,
               2 + embeddings.padding_idx + 1,
               3 + embeddings.padding_idx + 1,
           ]
           lowerCamelCase       :       int   =						torch.as_tensor([expected_single_positions, expected_single_positions]						)
           lowerCamelCase       :       str   =						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 		a__    (						self:    Optional[int]						)->     Dict:
           pass

     @unittest.skip("""Esm does not support embedding resizing"""						)
     def 		a__    (						self:    List[Any]						)->     Union[str, Any]:
           pass







     @unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""						)
     def 		a__    (						self:    Tuple						)->     Union[str, Any]:
           pass


@require_torch
class       A__						(    lowercase__):


     """simple docstring"""

     @slow
     def 		a__    (						self:    Any						)->     Tuple:


           with torch.no_grad():
                 lowerCamelCase       :       str   =						EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""						)
                 model.eval()
                 lowerCamelCase       :       Tuple   =						torch.tensor([[0, 1, 2, 3, 4, 5]]						)
                 lowerCamelCase       :       List[Any]   =						model(SCREAMING_SNAKE_CASE_						)[0]

                 lowerCamelCase       :       Tuple   =						33

                 lowerCamelCase       :       Any   =						torch.Size((1, 6, vocab_size)						)
                 self.assertEqual(output.shape  ,		SCREAMING_SNAKE_CASE_						)

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







     @slow
     def 		a__    (						self:    Union[str, Any]						)->     List[Any]:


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

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






"""simple docstring"""






from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase        =				logging.get_logger(__name__)


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

					from transformers import Trainer



					class 			lowercase					( lowercase__		):






								def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	int = 101			)       ->	List[str]:





									"""simple docstring"""


									lowerCAmelCase												=  length






								def __len__(self  :	Any			)       ->	List[Any]:





									"""simple docstring"""


									return self.length






								def __getitem__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	int:





									"""simple docstring"""


									return i





					class 			lowercase					:






								def __call__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





									"""simple docstring"""


									return {"input_ids": torch.tensor(SCREAMING_SNAKE_CASE_			), "labels": torch.tensor(SCREAMING_SNAKE_CASE_			)}





					class 			lowercase					( nn.Module		):






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





									"""simple docstring"""


									super().__init__()
									# Add some (unused) params otherwise DDP will complain.
									lowerCAmelCase												=  nn.Linear(120      ,80			)






								def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None			)       ->	Union[str, Any]:





									"""simple docstring"""


									if labels is not None:
										return torch.tensor(0.0      ,device=input_ids.device			), input_ids
									else:
										return input_ids



class 			lowercase					( lowercase__		):






			@require_torch_neuroncore
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call





class 			lowercase					( lowercase__		):






			@require_torch_multi_gpu
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
					# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
					#
					# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

					UpperCAmelCase        =				HfArgumentParser((TrainingArguments,))
					UpperCAmelCase        =				parser.parse_args_into_dataclasses()[0]

					logger.warning(
					    f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, """
					    f"""distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"""
					)

					# Essentially, what we want to verify in the distributed case is that we get all samples back,
					# in the right order. (this is crucial for prediction for instance)
					for dataset_length in [1_0_1, 4_0, 7]:
										UpperCAmelCase        =				DummyDataset(dataset_length)




										def __magic_name__	(				_lowerCamelCase:		EvalPrediction			)					->						Dict:

											'''simple docstring'''




											lowerCAmelCase												=  list(range(len(_lowerCamelCase			)			)			)
											lowerCAmelCase												=  p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
											if not success and training_args.local_rank == 0:
												logger.warning(
												    '''Predictions and/or labels do not match expected results:\n  - predictions: '''
												    F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""			)
											return {"success": success}

										UpperCAmelCase        =				Trainer(
										    model=DummyModel(),
										    args=training_args,
										    data_collator=DummyDataCollator(),
										    eval_dataset=dataset,
										    compute_metrics=compute_metrics,
										)
										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)

										UpperCAmelCase        =				2

										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)



										UpperCAmelCase        =				None


