



"""simple docstring"""





import unittest

from transformers import MraConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
  import torch

  from transformers import (
      MraForMaskedLM,
      MraForMultipleChoice,
      MraForQuestionAnswering,
      MraForSequenceClassification,
      MraForTokenClassification,
      MraModel,
  )
  from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST
class  lowercase__  :







       '''simple docstring'''


       def __init__(				self     :    List[Any]						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Tuple=2						,							_UpperCAmelCase     :    Union[str, Any]=8						,							_UpperCAmelCase     :    Union[str, Any]=True						,							_UpperCAmelCase     :    Optional[int]=True						,							_UpperCAmelCase     :    Tuple=True						,							_UpperCAmelCase     :    Any=True						,							_UpperCAmelCase     :    Optional[int]=99						,							_UpperCAmelCase     :    Tuple=16						,							_UpperCAmelCase     :    Union[str, Any]=5						,							_UpperCAmelCase     :    Optional[Any]=2						,							_UpperCAmelCase     :    Dict=36						,							_UpperCAmelCase     :    List[str]="gelu"						,							_UpperCAmelCase     :    Optional[int]=0.0						,							_UpperCAmelCase     :    Dict=0.0						,							_UpperCAmelCase     :    Tuple=512						,							_UpperCAmelCase     :    Optional[int]=16						,							_UpperCAmelCase     :    List[Any]=2						,							_UpperCAmelCase     :    str=0.02						,							_UpperCAmelCase     :    Dict=3						,							_UpperCAmelCase     :    List[str]=4						,							_UpperCAmelCase     :    Optional[Any]=None						,							)					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				parent
             UpperCAmelCase_     					=				batch_size
             UpperCAmelCase_     					=				seq_length
             UpperCAmelCase_     					=				is_training
             UpperCAmelCase_     					=				use_input_mask
             UpperCAmelCase_     					=				use_token_type_ids
             UpperCAmelCase_     					=				use_labels
             UpperCAmelCase_     					=				vocab_size
             UpperCAmelCase_     					=				hidden_size
             UpperCAmelCase_     					=				num_hidden_layers
             UpperCAmelCase_     					=				num_attention_heads
             UpperCAmelCase_     					=				intermediate_size
             UpperCAmelCase_     					=				hidden_act
             UpperCAmelCase_     					=				hidden_dropout_prob
             UpperCAmelCase_     					=				attention_probs_dropout_prob
             UpperCAmelCase_     					=				max_position_embeddings
             UpperCAmelCase_     					=				type_vocab_size
             UpperCAmelCase_     					=				type_sequence_label_size
             UpperCAmelCase_     					=				initializer_range
             UpperCAmelCase_     					=				num_labels
             UpperCAmelCase_     					=				num_choices
             UpperCAmelCase_     					=				scope



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


             '''simple docstring'''




             UpperCAmelCase_     					=				ids_tensor([self.batch_size, self.seq_length]						,							self.vocab_size       )

             UpperCAmelCase_     					=				None
             if self.use_input_mask:
                   UpperCAmelCase_     					=				random_attention_mask([self.batch_size, self.seq_length]       )

             UpperCAmelCase_     					=				None
             if self.use_token_type_ids:
                   UpperCAmelCase_     					=				ids_tensor([self.batch_size, self.seq_length]						,							self.type_vocab_size       )

             UpperCAmelCase_     					=				None
             UpperCAmelCase_     					=				None
             UpperCAmelCase_     					=				None
             if self.use_labels:
                   UpperCAmelCase_     					=				ids_tensor([self.batch_size]						,							self.type_sequence_label_size       )
                   UpperCAmelCase_     					=				ids_tensor([self.batch_size, self.seq_length]						,							self.num_labels       )
                   UpperCAmelCase_     					=				ids_tensor([self.batch_size]						,							self.num_choices       )

             UpperCAmelCase_     					=				self.get_config()

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



       def 	lowercase__							(				self     :    int       )					->		Dict:


             '''simple docstring'''




             return MraConfig(
                 vocab_size=self.vocab_size						,							hidden_size=self.hidden_size						,							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						,							is_decoder=_UpperCAmelCase						,							initializer_range=self.initializer_range						,							)



       def 	lowercase__							(				self     :    Dict       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.get_config()
             UpperCAmelCase_     					=				300
             return config



       def 	lowercase__							(				self     :    int       )					->		List[Any]:


             '''simple docstring'''




             (
                 (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       
             )     					=				self.prepare_config_and_inputs()

             UpperCAmelCase_     					=				True
             UpperCAmelCase_     					=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
             UpperCAmelCase_     					=				ids_tensor([self.batch_size, self.seq_length]						,							vocab_size=2       )

             return (
                 config,
                 input_ids,
                 token_type_ids,
                 input_mask,
                 sequence_labels,
                 token_labels,
                 choice_labels,
                 encoder_hidden_states,
                 encoder_attention_mask,
             )



       def 	lowercase__							(				self     :    int						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Tuple       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraModel(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase       )
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase       )
             UpperCAmelCase_     					=				model(_UpperCAmelCase       )
             self.parent.assertEqual(result.last_hidden_state.shape						,							(self.batch_size, self.seq_length, self.hidden_size)       )



       def 	lowercase__							(				self     :    int						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    Dict						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Optional[int]						,							)					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				True
             UpperCAmelCase_     					=				MraModel(_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(
                 _UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							encoder_hidden_states=_UpperCAmelCase						,							encoder_attention_mask=_UpperCAmelCase						,							)
             UpperCAmelCase_     					=				model(
                 _UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							encoder_hidden_states=_UpperCAmelCase						,							)
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase       )
             self.parent.assertEqual(result.last_hidden_state.shape						,							(self.batch_size, self.seq_length, self.hidden_size)       )



       def 	lowercase__							(				self     :    List[Any]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    str						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    str						,							_UpperCAmelCase     :    Optional[int]       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraForMaskedLM(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							labels=_UpperCAmelCase       )
             self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.vocab_size)       )



       def 	lowercase__							(				self     :    str						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    str						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    Optional[int]						,							_UpperCAmelCase     :    Tuple       )					->		Union[str, Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraForQuestionAnswering(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(
                 _UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							start_positions=_UpperCAmelCase						,							end_positions=_UpperCAmelCase						,							)
             self.parent.assertEqual(result.start_logits.shape						,							(self.batch_size, self.seq_length)       )
             self.parent.assertEqual(result.end_logits.shape						,							(self.batch_size, self.seq_length)       )



       def 	lowercase__							(				self     :    Any						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Optional[Any]						,							_UpperCAmelCase     :    str						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    Dict       )					->		Any:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.num_labels
             UpperCAmelCase_     					=				MraForSequenceClassification(_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							labels=_UpperCAmelCase       )
             self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.num_labels)       )



       def 	lowercase__							(				self     :    Dict						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[str]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    List[Any]       )					->		int:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.num_labels
             UpperCAmelCase_     					=				MraForTokenClassification(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				model(_UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							labels=_UpperCAmelCase       )
             self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.num_labels)       )



       def 	lowercase__							(				self     :    Optional[int]						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    Tuple						,							_UpperCAmelCase     :    Union[str, Any]						,							_UpperCAmelCase     :    int						,							_UpperCAmelCase     :    List[Any]						,							_UpperCAmelCase     :    Any						,							_UpperCAmelCase     :    List[str]       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.num_choices
             UpperCAmelCase_     					=				MraForMultipleChoice(config=_UpperCAmelCase       )
             model.to(_UpperCAmelCase       )
             model.eval()
             UpperCAmelCase_     					=				input_ids.unsqueeze(1       ).expand(-1						,							self.num_choices						,							-1       ).contiguous()
             UpperCAmelCase_     					=				token_type_ids.unsqueeze(1       ).expand(-1						,							self.num_choices						,							-1       ).contiguous()
             UpperCAmelCase_     					=				input_mask.unsqueeze(1       ).expand(-1						,							self.num_choices						,							-1       ).contiguous()
             UpperCAmelCase_     					=				model(
                 _UpperCAmelCase						,							attention_mask=_UpperCAmelCase						,							token_type_ids=_UpperCAmelCase						,							labels=_UpperCAmelCase						,							)
             self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.num_choices)       )







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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.prepare_config_and_inputs()
             (
                 (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       (
                 UpperCAmelCase_
             )					,       
             )     					=				config_and_inputs
             UpperCAmelCase_     					=				{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
             return config, inputs_dict



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







       '''simple docstring'''


       UpperCamelCase		=	(
           (
               MraModel,
               MraForMaskedLM,
               MraForMultipleChoice,
               MraForQuestionAnswering,
               MraForSequenceClassification,
               MraForTokenClassification,
           )
           if is_torch_available()
           else ()
       )
       UpperCamelCase		=	False
       UpperCamelCase		=	False
       UpperCamelCase		=	False
       UpperCamelCase		=	False

       UpperCamelCase		=	()



       def 	lowercase__							(				self     :    Optional[int]       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraModelTester(self       )
             UpperCAmelCase_     					=				ConfigTester(self						,							config_class=_UpperCAmelCase						,							hidden_size=37       )



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


             '''simple docstring'''




             self.config_tester.run_common_tests()



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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*_UpperCAmelCase       )



       def 	lowercase__							(				self     :    Tuple       )					->		Tuple:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             for type in ["absolute", "relative_key", "relative_key_query"]:
                   UpperCAmelCase_     					=				type
                   self.model_tester.create_and_check_model(*_UpperCAmelCase       )



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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase       )



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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_multiple_choice(*_UpperCAmelCase       )



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


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_question_answering(*_UpperCAmelCase       )



       def 	lowercase__							(				self     :    Optional[int]       )					->		str:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_sequence_classification(*_UpperCAmelCase       )



       def 	lowercase__							(				self     :    str       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase       )



       @slow
       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Tuple:


             '''simple docstring'''




             for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                   UpperCAmelCase_     					=				MraModel.from_pretrained(_UpperCAmelCase       )
                   self.assertIsNotNone(_UpperCAmelCase       )







       @unittest.skip(reason="MRA does not output attentions"       )
       def 	lowercase__							(				self     :    Optional[Any]       )					->		Any:


             '''simple docstring'''




             return



@require_torch
class  lowercase__  (       unittest.TestCase ):







       '''simple docstring'''


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


             '''simple docstring'''




             UpperCAmelCase_     					=				MraModel.from_pretrained("uw-madison/mra-base-512-4"       )
             UpperCAmelCase_     					=				torch.arange(256       ).unsqueeze(0       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(_UpperCAmelCase       )[0]

             UpperCAmelCase_     					=				torch.Size((1, 256, 768)       )
             self.assertEqual(output.shape						,							_UpperCAmelCase       )

             UpperCAmelCase_     					=				torch.tensor(
                 [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]       )

             self.assertTrue(torch.allclose(output[:, :3, :3]						,							_UpperCAmelCase						,							atol=1e-4       )       )



       @slow
       def 	lowercase__							(				self     :    Union[str, Any]       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4"       )
             UpperCAmelCase_     					=				torch.arange(256       ).unsqueeze(0       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(_UpperCAmelCase       )[0]

             UpperCAmelCase_     					=				50265

             UpperCAmelCase_     					=				torch.Size((1, 256, vocab_size)       )
             self.assertEqual(output.shape						,							_UpperCAmelCase       )

             UpperCAmelCase_     					=				torch.tensor(
                 [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]       )

             self.assertTrue(torch.allclose(output[:, :3, :3]						,							_UpperCAmelCase						,							atol=1e-4       )       )







       @slow
       def 	lowercase__							(				self     :    Tuple       )					->		Optional[Any]:


             '''simple docstring'''




             UpperCAmelCase_     					=				MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3"       )
             UpperCAmelCase_     					=				torch.arange(4096       ).unsqueeze(0       )

             with torch.no_grad():
                   UpperCAmelCase_     					=				model(_UpperCAmelCase       )[0]

             UpperCAmelCase_     					=				50265

             UpperCAmelCase_     					=				torch.Size((1, 4096, vocab_size)       )
             self.assertEqual(output.shape						,							_UpperCAmelCase       )

             UpperCAmelCase_     					=				torch.tensor(
                 [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]       )

             self.assertTrue(torch.allclose(output[:, :3, :3]						,							_UpperCAmelCase						,							atol=1e-4       )       )






"""simple docstring"""





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()
UpperCamelCase_       :    Any						      =			logging.get_logger(__name__)
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    ASTConfig()

					if "10-10" in model_name:
										pass
					elif "speech-commands" in model_name:
										A_								=    128
					elif "12-12" in model_name:
										A_								=    12
										A_								=    12
					elif "14-14" in model_name:
										A_								=    14
										A_								=    14
					elif "16-16" in model_name:
										A_								=    16
										A_								=    16
					else:
										raise ValueError("Model not supported"					)

					A_								=    "huggingface/label-files"
					if "speech-commands" in model_name:
										A_								=    35
										A_								=    "speech-commands-v2-id2label.json"
					else:
										A_								=    527
										A_								=    "audioset-id2label.json"

					A_								=    json.load(open(hf_hub_download(__a					, __a					, repo_type="dataset"					)					, "r"					)					)
					A_								=    {int(__a					): v for k, v in idalabel.items()}
					A_								=    idalabel
					A_								=    {v: k for k, v in idalabel.items()}

					return config
def 				A_  (__a					):

					'''simple docstring'''





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

					return name
def 				A_  (__a					, __a					):

					'''simple docstring'''





					for key in orig_state_dict.copy().keys():
										A_								=    orig_state_dict.pop(__a					)

										if "qkv" in key:
															A_								=    key.split("."					)
															A_								=    int(key_split[3]					)
															A_								=    config.hidden_size
															if "weight" in key:
																				A_								=    val[:dim, :]
																				A_								=    val[dim : dim * 2, :]
																				A_								=    val[-dim:, :]
															else:
																				A_								=    val[:dim]
																				A_								=    val[dim : dim * 2]
																				A_								=    val[-dim:]
										else:
															A_								=    val

					return orig_state_dict
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    [
					    "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(__a					, __a					)
@torch.no_grad()
def 				A_  (__a					, __a					, __a=False					):

					'''simple docstring'''





					A_								=    get_audio_spectrogram_transformer_config(__a					)

					A_								=    {
					    "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
					A_								=    model_name_to_url[model_name]
					A_								=    torch.hub.load_state_dict_from_url(__a					, map_location="cpu"					)
					# remove some keys
					remove_keys(__a					)
					# rename some keys
					A_								=    convert_state_dict(__a					, __a					)

					# load 🤗 model
					A_								=    ASTForAudioClassification(__a					)
					model.eval()

					model.load_state_dict(__a					)

					# verify outputs on dummy input
					# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
					A_								=    -4.2677393 if "speech-commands" not in model_name else -6.845978
					A_								=    4.5689974 if "speech-commands" not in model_name else 5.5654526
					A_								=    1024 if "speech-commands" not in model_name else 128
					A_								=    ASTFeatureExtractor(mean=__a					, std=__a					, max_length=__a					)

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

										A_			,     A_								=    torchaudio.load(__a					)
										A_								=    waveform.squeeze().numpy()

					A_								=    feature_extractor(__a					, sampling_rate=1_6000					, return_tensors="pt"					)

					# forward pass
					A_								=    model(**__a					)
					A_								=    outputs.logits

					if model_name == "ast-finetuned-audioset-10-10-0.4593":
										A_								=    torch.tensor([-0.8760, -7.0042, -8.6602]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.450":
										A_								=    torch.tensor([-1.1986, -7.0903, -8.2718]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448":
										A_								=    torch.tensor([-2.6128, -8.0080, -9.4344]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
										A_								=    torch.tensor([-1.5080, -7.4534, -8.8917]					)
					elif model_name == "ast-finetuned-audioset-12-12-0.447":
										A_								=    torch.tensor([-0.5050, -6.5833, -8.0843]					)
					elif model_name == "ast-finetuned-audioset-14-14-0.443":
										A_								=    torch.tensor([-0.3826, -7.0336, -8.2413]					)
					elif model_name == "ast-finetuned-audioset-16-16-0.442":
										A_								=    torch.tensor([-1.2113, -6.9101, -8.3470]					)
					elif model_name == "ast-finetuned-speech-commands-v2":
										A_								=    torch.tensor([6.1589, -8.0566, -8.7984]					)
					else:
										raise ValueError("Unknown model name"					)
					if not torch.allclose(logits[0, :3]					, __a					, atol=1e-4					):
										raise ValueError("Logits don't match"					)
					print("Looks ok!"					)

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

					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__":
		UpperCamelCase_       :    Dict						      =			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.'''
		)

		UpperCamelCase_       :    List[str]						      =			parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)




