






"""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()
A__					: Optional[int]=	logging.get_logger(__name__)







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE					)					->	Any:
    """simple docstring"""

    UpperCamelCase__         =    ASTConfig()

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

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

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

    return config







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE					)					->	int:
    """simple docstring"""

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

    return name







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE					)					->	Dict:
    """simple docstring"""

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

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

    return orig_state_dict







def   lowerCAmelCase_( SCREAMING_SNAKE_CASE					)					->	List[Any]:
    """simple docstring"""

    UpperCamelCase__         =    [
        '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(SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE					)







@torch.no_grad()
def   lowerCAmelCase_( SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE ,	SCREAMING_SNAKE_CASE=False					)					->	int:
    """simple docstring"""

    UpperCamelCase__         =    get_audio_spectrogram_transformer_config(SCREAMING_SNAKE_CASE					)

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

    # load 🤗 model
    UpperCamelCase__         =    ASTForAudioClassification(SCREAMING_SNAKE_CASE					)
    model.eval()

    model.load_state_dict(SCREAMING_SNAKE_CASE					)

    # verify outputs on dummy input
    # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
    UpperCamelCase__         =    -4.2677393 if 'speech-commands' not in model_name else -6.845978
    UpperCamelCase__         =    4.5689974 if 'speech-commands' not in model_name else 5.5654526
    UpperCamelCase__         =    10_24 if 'speech-commands' not in model_name else 1_28
    UpperCamelCase__         =    ASTFeatureExtractor(mean=SCREAMING_SNAKE_CASE ,	std=SCREAMING_SNAKE_CASE ,	max_length=SCREAMING_SNAKE_CASE					)

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

        UpperCamelCase__       ,						UpperCamelCase__         =    torchaudio.load(SCREAMING_SNAKE_CASE					)
        UpperCamelCase__         =    waveform.squeeze().numpy()

    UpperCamelCase__         =    feature_extractor(SCREAMING_SNAKE_CASE ,	sampling_rate=1_60_00 ,	return_tensors='pt'					)

    # forward pass
    UpperCamelCase__         =    model(**SCREAMING_SNAKE_CASE					)
    UpperCamelCase__         =    outputs.logits

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

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

    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__":
  A__					: Optional[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."""
  )

  A__					: str=	parser.parse_args()
  convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)












"""simple docstring"""





def   lowerCAmelCase_( SCREAMING_SNAKE_CASE = 50_00_00_00					)					->	int:
    """simple docstring"""

    UpperCamelCase__         =    set()
    UpperCamelCase__         =    int((limit - 24) ** (1 / 2)					)

    UpperCamelCase__         =    set(range(3 ,	prime_square_limit + 1 ,	2					)					)
    primes.add(2					)
    for p in range(3 ,	prime_square_limit + 1 ,	2					):
        if p not in primes:
            continue
        primes.difference_update(set(range(p * p ,	prime_square_limit + 1 ,	SCREAMING_SNAKE_CASE					)					)					)

    for primea in primes:
        UpperCamelCase__         =    primea * primea
        for primea in primes:
            UpperCamelCase__         =    primea * primea * primea
            if square + cube >= limit - 16:
                break
            for primea in primes:
                UpperCamelCase__         =    primea * primea * primea * primea
                UpperCamelCase__         =    square + cube + tetr
                if total >= limit:
                    break
                ret.add(SCREAMING_SNAKE_CASE					)

    return len(SCREAMING_SNAKE_CASE					)


if __name__ == "__main__":
  print(F"""{solution() = }""")




