




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     _SCREAMING_SNAKE_CASE    (							lowercase			: Any  ):


      '''simple docstring'''



      lowerCamelCase_	  =					ASTConfig()

      if "10-10" in model_name:
            pass
      elif "speech-commands" in model_name:
            lowerCamelCase_	  =					1_28
      elif "12-12" in model_name:
            lowerCamelCase_	  =					12
            lowerCamelCase_	  =					12
      elif "14-14" in model_name:
            lowerCamelCase_	  =					14
            lowerCamelCase_	  =					14
      elif "16-16" in model_name:
            lowerCamelCase_	  =					16
            lowerCamelCase_	  =					16
      else:
            raise ValueError('Model not supported'  )

      lowerCamelCase_	  =					'huggingface/label-files'
      if "speech-commands" in model_name:
            lowerCamelCase_	  =					35
            lowerCamelCase_	  =					'speech-commands-v2-id2label.json'
      else:
            lowerCamelCase_	  =					5_27
            lowerCamelCase_	  =					'audioset-id2label.json'

      lowerCamelCase_	  =					json.load(open(hf_hub_download(lowercase			,							lowercase			,							repo_type='dataset'  )			,							'r'  )  )
      lowerCamelCase_	  =					{int(lowercase  ): v for k, v in idalabel.items()}
      lowerCamelCase_	  =					idalabel
      lowerCamelCase_	  =					{v: k for k, v in idalabel.items()}

      return config

def     _SCREAMING_SNAKE_CASE    (							lowercase			: str  ):


      '''simple docstring'''



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

      return name

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[Any]			,							lowercase			: Any  ):


      '''simple docstring'''



      for key in orig_state_dict.copy().keys():
            lowerCamelCase_	  =					orig_state_dict.pop(lowercase  )

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

      return orig_state_dict

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Union[str, Any]  ):


      '''simple docstring'''



      lowerCamelCase_	  =					[
          '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(lowercase			,							lowercase  )

@torch.no_grad()
def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[int]			,							lowercase			: Optional[int]			,							lowercase			: Any=False  ):


      '''simple docstring'''



      lowerCamelCase_	  =					get_audio_spectrogram_transformer_config(lowercase  )

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

      # load 🤗 model
      lowerCamelCase_	  =					ASTForAudioClassification(lowercase  )
      model.eval()

      model.load_state_dict(lowercase  )

      # verify outputs on dummy input
      # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
      lowerCamelCase_	  =					-4.267_7393 if 'speech-commands' not in model_name else -6.84_5978
      lowerCamelCase_	  =					4.568_9974 if 'speech-commands' not in model_name else 5.565_4526
      lowerCamelCase_	  =					10_24 if 'speech-commands' not in model_name else 1_28
      lowerCamelCase_	  =					ASTFeatureExtractor(mean=lowercase			,							std=lowercase			,							max_length=lowercase  )

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

            lowerCamelCase_  ,       lowerCamelCase_	  =					torchaudio.load(lowercase  )
            lowerCamelCase_	  =					waveform.squeeze().numpy()

      lowerCamelCase_	  =					feature_extractor(lowercase			,							sampling_rate=1_60_00			,							return_tensors='pt'  )

      # forward pass
      lowerCamelCase_	  =					model(**lowercase  )
      lowerCamelCase_	  =					outputs.logits

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

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

      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       :	List[str]									=     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       :	List[Any]									=     parser.parse_args()
 convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)









import argparse
import json
import logging
import os
import sys
from unittest.mock import patch

from transformers.testing_utils import TestCasePlus, get_gpu_count, slow


lowerCamelCase       :	List[Any]									=     [
    os.path.join(os.path.dirname(__file__), dirname)
    for dirname in [
        "text-classification",
        "language-modeling",
        "summarization",
        "token-classification",
        "question-answering",
    ]
]
sys.path.extend(SRC_DIRS)


if SRC_DIRS is not None:
 import run_clm_flax
 import run_flax_glue
 import run_flax_ner
 import run_mlm_flax
 import run_qa
 import run_summarization_flax
 import run_ta_mlm_flax


logging.basicConfig(level=logging.DEBUG)

lowerCamelCase       :	Tuple									=     logging.getLogger()

def     _SCREAMING_SNAKE_CASE    (							):


      '''simple docstring'''



      lowerCamelCase_	  =					argparse.ArgumentParser()
      parser.add_argument('-f'  )
      lowerCamelCase_	  =					parser.parse_args()
      return args.f

def     _SCREAMING_SNAKE_CASE    (							lowercase			: Optional[Any]			,							lowercase			: Dict="eval"  ):


      '''simple docstring'''



      lowerCamelCase_	  =					os.path.join(lowercase			,							f"""{split}_results.json"""  )
      if os.path.exists(lowercase  ):
            with open(lowercase			,							'r'  ) as f:
                  return json.load(lowercase  )
      raise ValueError(f"""can't find {path}"""  )


lowerCamelCase       :	str									=     logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 						A(		UpperCamelCase     ):



 '''simple docstring'''


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

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_glue.py
            --model_name_or_path distilbert-base-uncased
            --output_dir {tmp_dir}
            --train_file ./tests/fixtures/tests_samples/MRPC/train.csv
            --validation_file ./tests/fixtures/tests_samples/MRPC/dev.csv
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --learning_rate=1e-4
            --eval_steps=2
            --warmup_steps=2
            --seed=42
            --max_seq_length=128
            """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_flax_glue.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertGreaterEqual(result['eval_accuracy']		,							0.75						)

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

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_clm_flax.py
            --model_name_or_path distilgpt2
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --block_size 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_clm_flax.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertLess(result['eval_perplexity']		,							100						)

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

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_summarization.py
            --model_name_or_path t5-small
            --train_file tests/fixtures/tests_samples/xsum/sample.json
            --validation_file tests/fixtures/tests_samples/xsum/sample.json
            --test_file tests/fixtures/tests_samples/xsum/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=8
            --do_train
            --do_eval
            --do_predict
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
            --predict_with_generate
        """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_summarization_flax.main()
             lowerCamelCase_	  =					get_results(A_		,							split='test'						)
             self.assertGreaterEqual(result['test_rouge1']		,							10						)
             self.assertGreaterEqual(result['test_rouge2']		,							2						)
             self.assertGreaterEqual(result['test_rougeL']		,							7						)
             self.assertGreaterEqual(result['test_rougeLsum']		,							7						)

 @slow
 def        a__				(       self      :						Optional[int]						)     ->			str:

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_mlm.py
            --model_name_or_path distilroberta-base
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --logging_steps 2 --eval_steps 2
            --do_train
            --do_eval
            --num_train_epochs=1
        """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_mlm_flax.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertLess(result['eval_perplexity']		,							42						)

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

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_t5_mlm_flax.py
            --model_name_or_path t5-small
            --train_file ./tests/fixtures/sample_text.txt
            --validation_file ./tests/fixtures/sample_text.txt
            --do_train
            --do_eval
            --max_seq_length 128
            --per_device_train_batch_size 4
            --per_device_eval_batch_size 4
            --num_train_epochs 2
            --logging_steps 2 --eval_steps 2
            --output_dir {tmp_dir}
            --overwrite_output_dir
            """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_ta_mlm_flax.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertGreaterEqual(result['eval_accuracy']		,							0.42						)

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

       """simple docstring"""






       lowerCamelCase_	  =					7 if get_gpu_count() > 1 else 2

       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_flax_ner.py
            --model_name_or_path bert-base-uncased
            --train_file tests/fixtures/tests_samples/conll/sample.json
            --validation_file tests/fixtures/tests_samples/conll/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --do_train
            --do_eval
            --warmup_steps=2
            --learning_rate=2e-4
            --logging_steps 2 --eval_steps 2
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=2
            --num_train_epochs={epochs}
            --seed 7
        """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_flax_ner.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertGreaterEqual(result['eval_accuracy']		,							0.75						)
             self.assertGreaterEqual(result['eval_f1']		,							0.3						)





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

       """simple docstring"""






       lowerCamelCase_	  =					self.get_auto_remove_tmp_dir()
       lowerCamelCase_	  =					f"""
            run_qa.py
            --model_name_or_path bert-base-uncased
            --version_2_with_negative
            --train_file tests/fixtures/tests_samples/SQUAD/sample.json
            --validation_file tests/fixtures/tests_samples/SQUAD/sample.json
            --output_dir {tmp_dir}
            --overwrite_output_dir
            --num_train_epochs=3
            --warmup_steps=2
            --do_train
            --do_eval
            --logging_steps 2 --eval_steps 2
            --learning_rate=2e-4
            --per_device_train_batch_size=2
            --per_device_eval_batch_size=1
        """.split()

       with patch.object(A_		,							'argv'		,							A_						):
             run_qa.main()
             lowerCamelCase_	  =					get_results(A_						)
             self.assertGreaterEqual(result['eval_f1']		,							30						)
             self.assertGreaterEqual(result['eval_exact']		,							30						)



