






"""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()
_lowerCAmelCase	:	str				=      logging.get_logger(__name__)


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Any				)->		int:






     '''simple docstring'''



     UpperCAmelCase__						:    Any							=							ASTConfig()

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

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

     UpperCAmelCase__						:    Union[str, Any]							=							json.load(open(hf_hub_download(snake_case			,					snake_case			,					repo_type="dataset"				)			,					"r"				)				)
     UpperCAmelCase__						:    Tuple							=							{int(snake_case				): v for k, v in idalabel.items()}
     UpperCAmelCase__						:    Union[str, Any]							=							idalabel
     UpperCAmelCase__						:    Any							=							{v: k for k, v in idalabel.items()}

     return config


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : int				)->		Union[str, Any]:






     '''simple docstring'''



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

     return name


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Dict			,					snake_case       : str				)->		Dict:






     '''simple docstring'''



     for key in orig_state_dict.copy().keys():
          UpperCAmelCase__						:    Dict							=							orig_state_dict.pop(snake_case				)

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

     return orig_state_dict


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Union[str, Any]				)->		Tuple:






     '''simple docstring'''



     UpperCAmelCase__						:    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(snake_case			,					snake_case				)


@torch.no_grad()
def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Any			,					snake_case       : int			,					snake_case       : Any=False				)->		Optional[Any]:






     '''simple docstring'''



     UpperCAmelCase__						:    Optional[int]							=							get_audio_spectrogram_transformer_config(snake_case				)

     UpperCAmelCase__						:    Dict							=							{
         "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__						:    Optional[int]							=							model_name_to_url[model_name]
     UpperCAmelCase__						:    int							=							torch.hub.load_state_dict_from_url(snake_case			,					map_location="cpu"				)
     # remove some keys
     remove_keys(snake_case				)
     # rename some keys
     UpperCAmelCase__						:    int							=							convert_state_dict(snake_case			,					snake_case				)

     # load 🤗 model
     UpperCAmelCase__						:    Dict							=							ASTForAudioClassification(snake_case				)
     model.eval()

     model.load_state_dict(snake_case				)

     # verify outputs on dummy input
     # source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
     UpperCAmelCase__						:    Union[str, Any]							=							-4.267_7393 if "speech-commands" not in model_name else -6.84_5978
     UpperCAmelCase__						:    int							=							4.568_9974 if "speech-commands" not in model_name else 5.565_4526
     UpperCAmelCase__						:    Any							=							1024 if "speech-commands" not in model_name else 128
     UpperCAmelCase__						:    Optional[int]							=							ASTFeatureExtractor(mean=snake_case			,					std=snake_case			,					max_length=snake_case				)

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

          UpperCAmelCase__							,				UpperCAmelCase__						:    int							=							torchaudio.load(snake_case				)
          UpperCAmelCase__						:    Tuple							=							waveform.squeeze().numpy()

     UpperCAmelCase__						:    List[Any]							=							feature_extractor(snake_case			,					sampling_rate=1_6000			,					return_tensors="pt"				)

     # forward pass
     UpperCAmelCase__						:    Optional[int]							=							model(**snake_case				)
     UpperCAmelCase__						:    Union[str, Any]							=							outputs.logits

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

     if pytorch_dump_folder_path is not None:
          Path(snake_case				).mkdir(exist_ok=snake_case				)
          print(f'Saving model {model_name} to {pytorch_dump_folder_path}'				)
          model.save_pretrained(snake_case				)
          print(f'Saving feature extractor to {pytorch_dump_folder_path}'				)
          feature_extractor.save_pretrained(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__":
 _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	:	int				=      parser.parse_args()
 convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)










"""simple docstring"""



import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


_lowerCAmelCase	:	Optional[int]				=      {
    """169M""": 12,
    """430M""": 24,
    """1B5""": 24,
    """3B""": 32,
    """7B""": 32,
    """14B""": 40,
}

_lowerCAmelCase	:	Optional[Any]				=      {
    """169M""": 768,
    """430M""": 1_024,
    """1B5""": 2_048,
    """3B""": 2_560,
    """7B""": 4_096,
    """14B""": 5_120,
}


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Any				)->		Tuple:






     '''simple docstring'''



     UpperCAmelCase__						:    str							=							list(state_dict.keys()				)
     for name in state_dict_keys:
          UpperCAmelCase__						:    Dict							=							state_dict.pop(snake_case				)
          # emb -> embedding
          if name.startswith("emb."				):
               UpperCAmelCase__						:    str							=							name.replace("emb."			,					"embeddings."				)
          # ln_0 -> pre_ln (only present at block 0)
          if name.startswith("blocks.0.ln0"				):
               UpperCAmelCase__						:    List[str]							=							name.replace("blocks.0.ln0"			,					"blocks.0.pre_ln"				)
          # att -> attention
          UpperCAmelCase__						:    Optional[int]							=							re.sub(r"blocks\.(\d+)\.att"			,					r"blocks.\1.attention"			,					snake_case				)
          # ffn -> feed_forward
          UpperCAmelCase__						:    Any							=							re.sub(r"blocks\.(\d+)\.ffn"			,					r"blocks.\1.feed_forward"			,					snake_case				)
          # time_mix_k -> time_mix_key and reshape
          if name.endswith(".time_mix_k"				):
               UpperCAmelCase__						:    List[Any]							=							name.replace(".time_mix_k"			,					".time_mix_key"				)
          # time_mix_v -> time_mix_value and reshape
          if name.endswith(".time_mix_v"				):
               UpperCAmelCase__						:    int							=							name.replace(".time_mix_v"			,					".time_mix_value"				)
          # time_mix_r -> time_mix_key and reshape
          if name.endswith(".time_mix_r"				):
               UpperCAmelCase__						:    Optional[Any]							=							name.replace(".time_mix_r"			,					".time_mix_receptance"				)

          if name != "head.weight":
               UpperCAmelCase__						:    int							=							"rwkv." + name

          UpperCAmelCase__						:    Dict							=							weight
     return state_dict


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : Optional[Any]			,					snake_case       : Optional[int]			,					snake_case       : Optional[Any]			,					snake_case       : Dict=None			,					snake_case       : List[Any]=None			,					snake_case       : List[str]=False			,					snake_case       : str=None				)->		int:






     '''simple docstring'''



     if tokenizer_file is None:
          print("No `--tokenizer_file` provided, we will use the default tokenizer."				)
          UpperCAmelCase__						:    str							=							5_0277
          UpperCAmelCase__						:    List[Any]							=							AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"				)
     else:
          UpperCAmelCase__						:    Union[str, Any]							=							PreTrainedTokenizerFast(tokenizer_file=snake_case				)
          UpperCAmelCase__						:    Tuple							=							len(snake_case				)
     tokenizer.save_pretrained(snake_case				)

     # 2. Build the config
     UpperCAmelCase__						:    Dict							=							list(NUM_HIDDEN_LAYERS_MAPPING.keys()				)
     if size is None:
          # Try to infer size from the checkpoint name
          for candidate in possible_sizes:
               if candidate in checkpoint_file:
                    UpperCAmelCase__						:    List[Any]							=							candidate
                    break
          if size is None:
               raise ValueError("Could not infer the size, please provide it with the `--size` argument."				)
     if size not in possible_sizes:
          raise ValueError(f'`size` should be one of {possible_sizes}, got {size}.'				)

     UpperCAmelCase__						:    int							=							RwkvConfig(
         vocab_size=snake_case			,					num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]			,					hidden_size=HIDEN_SIZE_MAPPING[size]			,					)
     config.save_pretrained(snake_case				)

     # 3. Download model file then convert state_dict
     UpperCAmelCase__						:    Optional[int]							=							hf_hub_download(snake_case			,					snake_case				)
     UpperCAmelCase__						:    Optional[Any]							=							torch.load(snake_case			,					map_location="cpu"				)
     UpperCAmelCase__						:    Tuple							=							convert_state_dict(snake_case				)

     # 4. Split in shards and save
     UpperCAmelCase__							,				UpperCAmelCase__						:    List[Any]							=							shard_checkpoint(snake_case				)
     for shard_file, shard in shards.items():
          torch.save(snake_case			,					os.path.join(snake_case			,					snake_case				)				)

     if index is not None:
          UpperCAmelCase__						:    Tuple							=							os.path.join(snake_case			,					snake_case				)
          # Save the index as well
          with open(snake_case			,					"w"			,					encoding="utf-8"				) as f:
               UpperCAmelCase__						:    Any							=							json.dumps(snake_case			,					indent=2			,					sort_keys=snake_case				) + "\n"
               f.write(snake_case				)

          # 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
          print(
              "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."				)
          UpperCAmelCase__						:    List[str]							=							list(shards.keys()				)

          del state_dict
          del shards
          gc.collect()

          for shard_file in shard_files:
               UpperCAmelCase__						:    Any							=							torch.load(os.path.join(snake_case			,					snake_case				)				)
               torch.save({k: v.cpu().clone() for k, v in state_dict.items()}			,					os.path.join(snake_case			,					snake_case				)				)

     del state_dict
     gc.collect()

     if push_to_hub:
          if model_name is None:
               raise ValueError("Please provide a `model_name` to push the model to the Hub."				)
          UpperCAmelCase__						:    Optional[Any]							=							AutoModelForCausalLM.from_pretrained(snake_case				)
          model.push_to_hub(snake_case			,					max_shard_size="2GB"				)
          tokenizer.push_to_hub(snake_case				)


if __name__ == "__main__":
 _lowerCAmelCase	:	Optional[Any]				=      argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     """--repo_id""", default=None, type=str, required=True, help="""Repo ID from which to pull the checkpoint."""
 )
 parser.add_argument(
     """--checkpoint_file""", default=None, type=str, required=True, help="""Name of the checkpoint file in the repo."""
 )
 parser.add_argument(
     """--output_dir""", default=None, type=str, required=True, help="""Where to save the converted model."""
 )
 parser.add_argument(
     """--tokenizer_file""",
     default=None,
     type=str,
     help="""Path to the tokenizer file to use (if not provided, only the model is converted).""",
 )
 parser.add_argument(
     """--size""",
     default=None,
     type=str,
     help="""Size of the model. Will be inferred from the `checkpoint_file` if not passed.""",
 )
 parser.add_argument(
     """--push_to_hub""",
     action="""store_true""",
     help="""Push to the Hub the converted model.""",
 )
 parser.add_argument(
     """--model_name""",
     default=None,
     type=str,
     help="""Name of the pushed model on the Hub, including the username / organization.""",
 )

 _lowerCAmelCase	:	Dict				=      parser.parse_args()
 convert_rmkv_checkpoint_to_hf_format(
     args.repo_id,
     args.checkpoint_file,
     args.output_dir,
     size=args.size,
     tokenizer_file=args.tokenizer_file,
     push_to_hub=args.push_to_hub,
     model_name=args.model_name,
 )


