





import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
lowerCAmelCase							:       Optional[Any]					 =	logging.get_logger(__name__)






def 		lowerCAmelCase  (	UpperCamelCase__   :   int					)      ->    str:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Dict           =   ASTConfig()

						if "10-10" in model_name:
												pass
						elif "speech-commands" in model_name:
												__SCREAMING_SNAKE_CASE:				int           =   128
						elif "12-12" in model_name:
												__SCREAMING_SNAKE_CASE:				int           =   12
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   12
						elif "14-14" in model_name:
												__SCREAMING_SNAKE_CASE:				List[str]           =   14
												__SCREAMING_SNAKE_CASE:				str           =   14
						elif "16-16" in model_name:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   16
												__SCREAMING_SNAKE_CASE:				Dict           =   16
						else:
												raise ValueError('''Model not supported'''					)

						__SCREAMING_SNAKE_CASE:				Optional[int]           =   '''huggingface/label-files'''
						if "speech-commands" in model_name:
												__SCREAMING_SNAKE_CASE:				List[str]           =   35
												__SCREAMING_SNAKE_CASE:				int           =   '''speech-commands-v2-id2label.json'''
						else:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   527
												__SCREAMING_SNAKE_CASE:				List[Any]           =   '''audioset-id2label.json'''

						__SCREAMING_SNAKE_CASE:				List[Any]           =   json.load(open(hf_hub_download(UpperCamelCase__       ,			UpperCamelCase__       ,			repo_type='''dataset'''					)       ,			'''r'''					)					)
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   {int(UpperCamelCase__					): v for k, v in idalabel.items()}
						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   idalabel
						__SCREAMING_SNAKE_CASE:				str           =   {v: k for k, v in idalabel.items()}

						return config






def 		lowerCAmelCase  (	UpperCamelCase__   :   Dict					)      ->    Tuple:

						"""simple docstring"""





						if "module.v" in name:
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   name.replace('''module.v'''       ,			'''audio_spectrogram_transformer'''					)
						if "cls_token" in name:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   name.replace('''cls_token'''       ,			'''embeddings.cls_token'''					)
						if "dist_token" in name:
												__SCREAMING_SNAKE_CASE:				List[str]           =   name.replace('''dist_token'''       ,			'''embeddings.distillation_token'''					)
						if "pos_embed" in name:
												__SCREAMING_SNAKE_CASE:				Optional[int]           =   name.replace('''pos_embed'''       ,			'''embeddings.position_embeddings'''					)
						if "patch_embed.proj" in name:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   name.replace('''patch_embed.proj'''       ,			'''embeddings.patch_embeddings.projection'''					)
						# transformer blocks
						if "blocks" in name:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   name.replace('''blocks'''       ,			'''encoder.layer'''					)
						if "attn.proj" in name:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   name.replace('''attn.proj'''       ,			'''attention.output.dense'''					)
						if "attn" in name:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   name.replace('''attn'''       ,			'''attention.self'''					)
						if "norm1" in name:
												__SCREAMING_SNAKE_CASE:				Tuple           =   name.replace('''norm1'''       ,			'''layernorm_before'''					)
						if "norm2" in name:
												__SCREAMING_SNAKE_CASE:				Any           =   name.replace('''norm2'''       ,			'''layernorm_after'''					)
						if "mlp.fc1" in name:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   name.replace('''mlp.fc1'''       ,			'''intermediate.dense'''					)
						if "mlp.fc2" in name:
												__SCREAMING_SNAKE_CASE:				Dict           =   name.replace('''mlp.fc2'''       ,			'''output.dense'''					)
						# final layernorm
						if "audio_spectrogram_transformer.norm" in name:
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   name.replace('''audio_spectrogram_transformer.norm'''       ,			'''audio_spectrogram_transformer.layernorm'''					)
						# classifier head
						if "module.mlp_head.0" in name:
												__SCREAMING_SNAKE_CASE:				int           =   name.replace('''module.mlp_head.0'''       ,			'''classifier.layernorm'''					)
						if "module.mlp_head.1" in name:
												__SCREAMING_SNAKE_CASE:				List[Any]           =   name.replace('''module.mlp_head.1'''       ,			'''classifier.dense'''					)

						return name






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

						"""simple docstring"""





						for key in orig_state_dict.copy().keys():
												__SCREAMING_SNAKE_CASE:				Any           =   orig_state_dict.pop(UpperCamelCase__					)

												if "qkv" in key:
																		__SCREAMING_SNAKE_CASE:				Tuple           =   key.split('''.'''					)
																		__SCREAMING_SNAKE_CASE:				str           =   int(key_split[3]					)
																		__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   config.hidden_size
																		if "weight" in key:
																								__SCREAMING_SNAKE_CASE:				Optional[int]           =   val[:dim, :]
																								__SCREAMING_SNAKE_CASE:				List[Any]           =   val[dim : dim * 2, :]
																								__SCREAMING_SNAKE_CASE:				List[str]           =   val[-dim:, :]
																		else:
																								__SCREAMING_SNAKE_CASE:				str           =   val[:dim]
																								__SCREAMING_SNAKE_CASE:				Optional[int]           =   val[dim : dim * 2]
																								__SCREAMING_SNAKE_CASE:				Any           =   val[-dim:]
												else:
																		__SCREAMING_SNAKE_CASE:				List[str]           =   val

						return orig_state_dict






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

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   [
						    '''module.v.head.weight''',
						    '''module.v.head.bias''',
						    '''module.v.head_dist.weight''',
						    '''module.v.head_dist.bias''',
						]
						for k in ignore_keys:
												state_dict.pop(UpperCamelCase__       ,			UpperCamelCase__					)






@torch.no_grad()
def 		lowerCAmelCase  (	UpperCamelCase__   :   int       ,			UpperCamelCase__   :   Optional[Any]       ,			UpperCamelCase__   :   List[str]=False					)      ->    List[Any]:

						"""simple docstring"""





						__SCREAMING_SNAKE_CASE:				List[str]           =   get_audio_spectrogram_transformer_config(UpperCamelCase__					)

						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   {
						    '''ast-finetuned-audioset-10-10-0.4593''': (
						        '''https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-10-10-0.450''': (
						        '''https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-10-10-0.448''': (
						        '''https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-10-10-0.448-v2''': (
						        '''https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-12-12-0.447''': (
						        '''https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-14-14-0.443''': (
						        '''https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1'''
						    ),
						    '''ast-finetuned-audioset-16-16-0.442''': (
						        '''https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1'''
						    ),
						    '''ast-finetuned-speech-commands-v2''': (
						        '''https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1'''
						    ),
						}

						# load original state_dict
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model_name_to_url[model_name]
						__SCREAMING_SNAKE_CASE:				str           =   torch.hub.load_state_dict_from_url(UpperCamelCase__       ,			map_location='''cpu'''					)
						# remove some keys
						remove_keys(UpperCamelCase__					)
						# rename some keys
						__SCREAMING_SNAKE_CASE:				List[str]           =   convert_state_dict(UpperCamelCase__       ,			UpperCamelCase__					)

						# load 🤗 model
						__SCREAMING_SNAKE_CASE:				Optional[int]           =   ASTForAudioClassification(UpperCamelCase__					)
						model.eval()

						model.load_state_dict(UpperCamelCase__					)

						# verify outputs on dummy input
						# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
						__SCREAMING_SNAKE_CASE:				Any           =   -4.2_67_73_93 if '''speech-commands''' not in model_name else -6.84_59_78
						__SCREAMING_SNAKE_CASE:				List[Any]           =   4.5_68_99_74 if '''speech-commands''' not in model_name else 5.5_65_45_26
						__SCREAMING_SNAKE_CASE:				str           =   1_024 if '''speech-commands''' not in model_name else 128
						__SCREAMING_SNAKE_CASE:				int           =   ASTFeatureExtractor(mean=UpperCamelCase__       ,			std=UpperCamelCase__       ,			max_length=UpperCamelCase__					)

						if "speech-commands" in model_name:
												__SCREAMING_SNAKE_CASE:				Tuple           =   load_dataset('''speech_commands'''       ,			'''v0.02'''       ,			split='''validation'''					)
												__SCREAMING_SNAKE_CASE:				List[str]           =   dataset[0]['''audio''']['''array''']
						else:
												__SCREAMING_SNAKE_CASE:				Dict           =   hf_hub_download(
												    repo_id='''nielsr/audio-spectogram-transformer-checkpoint'''       ,			filename='''sample_audio.flac'''       ,			repo_type='''dataset'''       ,			)

												__SCREAMING_SNAKE_CASE       ,__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torchaudio.load(UpperCamelCase__					)
												__SCREAMING_SNAKE_CASE:				str           =   waveform.squeeze().numpy()

						__SCREAMING_SNAKE_CASE:				Any           =   feature_extractor(UpperCamelCase__       ,			sampling_rate=16_000       ,			return_tensors='''pt'''					)

						# forward pass
						__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model(**UpperCamelCase__					)
						__SCREAMING_SNAKE_CASE:				str           =   outputs.logits

						if model_name == "ast-finetuned-audioset-10-10-0.4593":
												__SCREAMING_SNAKE_CASE:				Any           =   torch.tensor([-0.87_60, -7.00_42, -8.66_02]					)
						elif model_name == "ast-finetuned-audioset-10-10-0.450":
												__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.tensor([-1.19_86, -7.09_03, -8.27_18]					)
						elif model_name == "ast-finetuned-audioset-10-10-0.448":
												__SCREAMING_SNAKE_CASE:				str           =   torch.tensor([-2.61_28, -8.00_80, -9.43_44]					)
						elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
												__SCREAMING_SNAKE_CASE:				List[str]           =   torch.tensor([-1.50_80, -7.45_34, -8.89_17]					)
						elif model_name == "ast-finetuned-audioset-12-12-0.447":
												__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.tensor([-0.50_50, -6.58_33, -8.08_43]					)
						elif model_name == "ast-finetuned-audioset-14-14-0.443":
												__SCREAMING_SNAKE_CASE:				Tuple           =   torch.tensor([-0.38_26, -7.03_36, -8.24_13]					)
						elif model_name == "ast-finetuned-audioset-16-16-0.442":
												__SCREAMING_SNAKE_CASE:				str           =   torch.tensor([-1.21_13, -6.91_01, -8.34_70]					)
						elif model_name == "ast-finetuned-speech-commands-v2":
												__SCREAMING_SNAKE_CASE:				int           =   torch.tensor([6.15_89, -8.05_66, -8.79_84]					)
						else:
												raise ValueError('''Unknown model name'''					)
						if not torch.allclose(logits[0, :3]       ,			UpperCamelCase__       ,			atol=1E-4					):
												raise ValueError('''Logits don\'t match'''					)
						print('''Looks ok!'''					)

						if pytorch_dump_folder_path is not None:
												Path(UpperCamelCase__					).mkdir(exist_ok=UpperCamelCase__					)
												print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""					)
												model.save_pretrained(UpperCamelCase__					)
												print(F"""Saving feature extractor to {pytorch_dump_folder_path}"""					)
												feature_extractor.save_pretrained(UpperCamelCase__					)

						if push_to_hub:
												print('''Pushing model and feature extractor to the hub...'''					)
												model.push_to_hub(F"""MIT/{model_name}"""					)
												feature_extractor.push_to_hub(F"""MIT/{model_name}"""					)


if __name__ == "__main__":
							lowerCAmelCase							:       int					 =	argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    """--model_name""",
							    default="""ast-finetuned-audioset-10-10-0.4593""",
							    type=str,
							    help="""Name of the Audio Spectrogram Transformer model you'd like to convert.""",
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
							)
							parser.add_argument(
							    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
							)

							lowerCAmelCase							:       Optional[int]					 =	parser.parse_args()
							convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










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			, _lowerCAmelCase			, _lowerCAmelCase=13			, _lowerCAmelCase=7			, _lowerCAmelCase=False			, _lowerCAmelCase=True			, _lowerCAmelCase=False			, _lowerCAmelCase=True			, _lowerCAmelCase=33			, _lowerCAmelCase=32			, _lowerCAmelCase=5			, _lowerCAmelCase=4			, _lowerCAmelCase=37			, _lowerCAmelCase="gelu"			, _lowerCAmelCase=0.1			, _lowerCAmelCase=0.1			, _lowerCAmelCase=512			, _lowerCAmelCase=16			, _lowerCAmelCase=2			, _lowerCAmelCase=0.02			, _lowerCAmelCase=3			, _lowerCAmelCase=4			, _lowerCAmelCase=None			, ):

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   ids_tensor([self.batch_size, self.seq_length]			, self.vocab_size					)

										__SCREAMING_SNAKE_CASE:				int           =   None
										if self.use_input_mask:
																__SCREAMING_SNAKE_CASE:				Tuple           =   random_attention_mask([self.batch_size, self.seq_length]					)

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

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

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






				def 			snake_case_    (  self					):

										"""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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmModel(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				Tuple           =   model(_lowerCAmelCase					)
										__SCREAMING_SNAKE_CASE:				int           =   model(_lowerCAmelCase					)

										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 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmForMaskedLM(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase			, labels=_lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.vocab_size)					)






				def 			snake_case_    (  self			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase			, _lowerCAmelCase					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Any           =   self.num_labels
										__SCREAMING_SNAKE_CASE:				List[Any]           =   EsmForTokenClassification(config=_lowerCAmelCase					)
										model.to(_lowerCAmelCase					)
										model.eval()
										__SCREAMING_SNAKE_CASE:				int           =   model(_lowerCAmelCase			, attention_mask=_lowerCAmelCase			, labels=_lowerCAmelCase					)
										self.parent.assertEqual(result.logits.shape			, (self.batch_size, self.seq_length, self.num_labels)					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				str           =   self.prepare_config_and_inputs()
										(
										    (
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,(
										    __SCREAMING_SNAKE_CASE
										)       ,
										):				int           =   config_and_inputs
										__SCREAMING_SNAKE_CASE:				Tuple           =   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
										return config, inputs_dict





@require_torch
class 				a			(      __lowercase   ,__lowercase   ,unittest.TestCase			):
				SCREAMING_SNAKE_CASE__     :   Tuple		      =     False

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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   EsmModelTester(self					)
										__SCREAMING_SNAKE_CASE:				List[str]           =   ConfigTester(self			, config_class=_lowerCAmelCase			, hidden_size=37					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										self.config_tester.run_common_tests()






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Dict           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




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






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*_lowerCAmelCase					)






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   EsmModel.from_pretrained(_lowerCAmelCase					)
																self.assertIsNotNone(_lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   self.model_tester.prepare_config_and_inputs()[0]
										__SCREAMING_SNAKE_CASE:				int           =   EsmEmbeddings(config=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   torch.as_tensor([[12, 31, 13, model.padding_idx]]					)
										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   torch.as_tensor(
										    [
										        [
										            0 + model.padding_idx + 1,
										            1 + model.padding_idx + 1,
										            2 + model.padding_idx + 1,
										            model.padding_idx,
										        ]
										    ]					)
										__SCREAMING_SNAKE_CASE:				Any           =   create_position_ids_from_input_ids(_lowerCAmelCase			, model.padding_idx					)
										self.assertEqual(position_ids.shape			, expected_positions.shape					)
										self.assertTrue(torch.all(torch.eq(_lowerCAmelCase			, _lowerCAmelCase					)					)					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Tuple           =   self.model_tester.prepare_config_and_inputs()[0]
										__SCREAMING_SNAKE_CASE:				Dict           =   EsmEmbeddings(config=_lowerCAmelCase					)

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






				@unittest.skip('''Esm does not support embedding resizing'''					)
				def 			snake_case_    (  self					):

										"""simple docstring"""




										pass






				@unittest.skip('''Esm does not support embedding resizing'''					)
				def 			snake_case_    (  self					):

										"""simple docstring"""




										pass






				@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''					)
				def 			snake_case_    (  self					):

										"""simple docstring"""




										pass





@require_torch
class 				a			(      __lowercase			):






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




										with torch.no_grad():
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   EsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''					)
																model.eval()
																__SCREAMING_SNAKE_CASE:				Tuple           =   torch.tensor([[0, 1, 2, 3, 4, 5]]					)
																__SCREAMING_SNAKE_CASE:				str           =   model(_lowerCAmelCase					)[0]

																__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   33

																__SCREAMING_SNAKE_CASE:				Dict           =   torch.Size((1, 6, vocab_size)					)
																self.assertEqual(output.shape			, _lowerCAmelCase					)

																__SCREAMING_SNAKE_CASE:				Optional[int]           =   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]			, _lowerCAmelCase			, atol=1e-4					)					)






				@slow
				def 			snake_case_    (  self					):

										"""simple docstring"""




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

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



