






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 		lowerCAmelCase     :

					'''simple docstring'''

					def __init__(  self				:	List[str]		,     __a				:	Union[str, Any]		,     __a				:	int=13		,     __a				:	str=7		,     __a				:	Dict=False		,     __a				:	List[Any]=True		,     __a				:	Dict=False		,     __a				:	str=True		,     __a				:	int=33		,     __a				:	List[str]=32		,     __a				:	int=5		,     __a				:	Dict=4		,     __a				:	str=37		,     __a				:	Optional[Any]="gelu"		,     __a				:	Any=0.1		,     __a				:	List[str]=0.1		,     __a				:	List[str]=512		,     __a				:	Dict=16		,     __a				:	Optional[Any]=2		,     __a				:	str=0.02		,     __a				:	Union[str, Any]=3		,     __a				:	Optional[Any]=4		,     __a				:	List[Any]=None		,     )			->  Optional[Any]:


												"""simple docstring"""

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



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


												"""simple docstring"""

												__lowercase      :							Dict				=     ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size							)

												__lowercase      :							Dict				=     None
												if self.use_input_mask:
																			__lowercase      :							str				=     random_attention_mask([self.batch_size, self.seq_length]							)

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

												__lowercase      :							str				=     self.get_config()

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



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


												"""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       lowerCAmelCase				(  self				:	Tuple		,     __a				:	Any		,     __a				:	Tuple		,     __a				:	Optional[Any]		,     __a				:	Union[str, Any]		,     __a				:	Any		,     __a				:	List[Any]							)			->  str:


												"""simple docstring"""

												__lowercase      :							Dict				=     EsmModel(config=__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							Optional[int]				=     model(__a		,     attention_mask=__a							)
												__lowercase      :							str				=     model(__a							)
												__lowercase      :							str				=     model(__a							)

												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       lowerCAmelCase				(  self				:	Tuple		,     __a				:	Tuple		,     __a				:	Dict		,     __a				:	Union[str, Any]		,     __a				:	str		,     __a				:	Union[str, Any]		,     __a				:	Union[str, Any]							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Tuple				=     EsmForMaskedLM(config=__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							int				=     model(__a		,     attention_mask=__a		,     labels=__a							)
												self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.vocab_size)							)



					def       lowerCAmelCase				(  self				:	str		,     __a				:	Optional[int]		,     __a				:	int		,     __a				:	str		,     __a				:	Tuple		,     __a				:	str		,     __a				:	Tuple							)			->  str:


												"""simple docstring"""

												__lowercase      :							Any				=     self.num_labels
												__lowercase      :							Union[str, Any]				=     EsmForTokenClassification(config=__a							)
												model.to(__a							)
												model.eval()
												__lowercase      :							str				=     model(__a		,     attention_mask=__a		,     labels=__a							)
												self.parent.assertEqual(result.logits.shape		,     (self.batch_size, self.seq_length, self.num_labels)							)



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


												"""simple docstring"""

												__lowercase      :							str				=     self.prepare_config_and_inputs()
												(
												    (
												    __lowercase
												)			,			(
												    __lowercase
												)			,			(
												    __lowercase
												)			,			(
												    __lowercase
												)			,			(
												    __lowercase
												)			,			(
												    __lowercase
												)			,			
												)      :							Any				=     config_and_inputs
												__lowercase      :							Dict				=     {"""input_ids""": input_ids, """attention_mask""": input_mask}
												return config, inputs_dict







@require_torch
class 		lowerCAmelCase     ( __a      ,       __a      ,       unittest.TestCase			):

					'''simple docstring'''

					_A						:			int									=      False

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



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


												"""simple docstring"""

												__lowercase      :							Any				=     EsmModelTester(self							)
												__lowercase      :							str				=     ConfigTester(self		,     config_class=__a		,     hidden_size=37							)



					def       lowerCAmelCase				(  self				:	Any							)			->  Optional[int]:


												"""simple docstring"""

												self.config_tester.run_common_tests()



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


												"""simple docstring"""

												__lowercase      :							Optional[int]				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_model(*__a							)



					def       lowerCAmelCase				(  self				:	Tuple							)			->  Tuple:


												"""simple docstring"""

												__lowercase      :							Any				=     self.model_tester.prepare_config_and_inputs()
												for type in ["absolute", "relative_key", "relative_key_query"]:
																			__lowercase      :							Optional[Any]				=     type
																			self.model_tester.create_and_check_model(*__a							)



					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  Dict:


												"""simple docstring"""

												__lowercase      :							str				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_masked_lm(*__a							)



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


												"""simple docstring"""

												__lowercase      :							Dict				=     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_for_token_classification(*__a							)



					@slow
					def       lowerCAmelCase				(  self				:	Optional[Any]							)			->  Dict:


												"""simple docstring"""

												for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			__lowercase      :							int				=     EsmModel.from_pretrained(__a							)
																			self.assertIsNotNone(__a							)



					def       lowerCAmelCase				(  self				:	Dict							)			->  Any:


												"""simple docstring"""

												__lowercase      :							Union[str, Any]				=     self.model_tester.prepare_config_and_inputs()[0]
												__lowercase      :							Optional[Any]				=     EsmEmbeddings(config=__a							)

												__lowercase      :							Optional[int]				=     torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
												__lowercase      :							List[str]				=     torch.as_tensor(
												    [
												        [
												            0 + model.padding_idx + 1,
												            1 + model.padding_idx + 1,
												            2 + model.padding_idx + 1,
												            model.padding_idx,
												        ]
												    ]							)
												__lowercase      :							List[str]				=     create_position_ids_from_input_ids(__a		,     model.padding_idx							)
												self.assertEqual(position_ids.shape		,     expected_positions.shape							)
												self.assertTrue(torch.all(torch.eq(__a		,     __a							)							)							)



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


												"""simple docstring"""

												__lowercase      :							Tuple				=     self.model_tester.prepare_config_and_inputs()[0]
												__lowercase      :							List[Any]				=     EsmEmbeddings(config=__a							)

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



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


												"""simple docstring"""

												pass



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


												"""simple docstring"""

												pass



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


												"""simple docstring"""

												pass







@require_torch
class 		lowerCAmelCase     ( __a			):

					'''simple docstring'''

					@slow
					def       lowerCAmelCase				(  self				:	Optional[int]							)			->  List[Any]:


												"""simple docstring"""

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

																			__lowercase      :							str				=     33

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

																			__lowercase      :							List[str]				=     torch.tensor(
																			    [[[8.9215, -10.5898, -6.4671], [-6.3967, -13.9114, -1.1212], [-7.7812, -13.9516, -3.7406]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3]		,     __a		,     atol=1E-4							)							)



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


												"""simple docstring"""

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

																			__lowercase      :							Optional[int]				=     torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
																			__lowercase      :							str				=     model(__a							)[0]
																			# compare the actual values for a slice.
																			__lowercase      :							Optional[int]				=     torch.tensor(
																			    [[[0.1444, 0.5413, 0.3248], [0.3034, 0.0053, 0.3108], [0.3228, -0.2499, 0.3415]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3]		,     __a		,     atol=1E-4							)							)







# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


lowerCamelCase : int						   =       '''Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine'''
def       snake_case_		(     ):
							__lowercase      :							List[Any]				=     _ask_options(
							    """In which compute environment are you running?"""       ,       ["""This machine""", """AWS (Amazon SageMaker)"""]       ,       _convert_compute_environment       ,       )
							if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
														__lowercase      :							Union[str, Any]				=     get_sagemaker_input()
							else:
														__lowercase      :							str				=     get_cluster_input()
							return config
def       snake_case_		(     lowerCAmelCase_       :						List[str]=None      ):
							if subparsers is not None:
														__lowercase      :							Optional[int]				=     subparsers.add_parser("""config"""       ,       description=lowerCAmelCase_      )
							else:
														__lowercase      :							List[str]				=     argparse.ArgumentParser("""Accelerate config command"""       ,       description=lowerCAmelCase_      )

							parser.add_argument(
							    """--config_file"""       ,       default=lowerCAmelCase_       ,       help=(
							        """The path to use to store the config file. Will default to a file named default_config.yaml in the cache """
							        """location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have """
							        """such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed """
							        """with 'huggingface'."""
							    )       ,       )

							if subparsers is not None:
														parser.set_defaults(func=lowerCAmelCase_      )
							return parser
def       snake_case_		(     lowerCAmelCase_       :						Tuple      ):
							__lowercase      :							Union[str, Any]				=     get_user_input()
							if args.config_file is not None:
														__lowercase      :							List[Any]				=     args.config_file
							else:
														if not os.path.isdir(lowerCAmelCase_      ):
																					os.makedirs(lowerCAmelCase_      )
														__lowercase      :							Any				=     default_yaml_config_file

							if config_file.endswith(""".json"""      ):
														config.to_json_file(lowerCAmelCase_      )
							else:
														config.to_yaml_file(lowerCAmelCase_      )
							print(F"accelerate configuration saved at {config_file}"      )
def       snake_case_		(     ):
							__lowercase      :							str				=     config_command_parser()
							__lowercase      :							str				=     parser.parse_args()
							config_command(lowerCAmelCase_      )


if __name__ == "__main__":
	main()