



'''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						      =       logging.get_logger(__name__)



def 					__lowerCamelCase					(				lowerCAmelCase_  )      ->				Any:
		_a					:				Optional[int]				=							ASTConfig()

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

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

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

		return config



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

		return name



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_  )      ->				int:
		for key in orig_state_dict.copy().keys():
				_a					:				Tuple				=							orig_state_dict.pop(lowerCAmelCase_  )

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

		return orig_state_dict



def 					__lowerCamelCase					(				lowerCAmelCase_  )      ->				Dict:
		_a					:				str				=							[
		    '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(lowerCAmelCase_						,      lowerCAmelCase_  )



@torch.no_grad()
def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_=False  )      ->				Optional[int]:
		_a					:				Optional[int]				=							get_audio_spectrogram_transformer_config(lowerCAmelCase_  )

		_a					:				str				=							{
		    '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					:				List[Any]				=							model_name_to_url[model_name]
		_a					:				Tuple				=							torch.hub.load_state_dict_from_url(lowerCAmelCase_						,      map_location='cpu'  )
		# remove some keys
		remove_keys(lowerCAmelCase_  )
		# rename some keys
		_a					:				Any				=							convert_state_dict(lowerCAmelCase_						,      lowerCAmelCase_  )

		# load 🤗 model
		_a					:				Union[str, Any]				=							ASTForAudioClassification(lowerCAmelCase_  )
		model.eval()

		model.load_state_dict(lowerCAmelCase_  )

		# verify outputs on dummy input
		# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
		_a					:				Any				=							-4.2_677_393 if 'speech-commands' not in model_name else -6.845_978
		_a					:				List[Any]				=							4.5_689_974 if 'speech-commands' not in model_name else 5.5_654_526
		_a					:				int				=							1024 if 'speech-commands' not in model_name else 128
		_a					:				Tuple				=							ASTFeatureExtractor(mean=lowerCAmelCase_						,      std=lowerCAmelCase_						,      max_length=lowerCAmelCase_  )

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

				_a	, _a					:				Any				=							torchaudio.load(lowerCAmelCase_  )
				_a					:				List[str]				=							waveform.squeeze().numpy()

		_a					:				str				=							feature_extractor(lowerCAmelCase_						,      sampling_rate=16000						,      return_tensors='pt'  )

		# forward pass
		_a					:				List[str]				=							model(**lowerCAmelCase_  )
		_a					:				Union[str, Any]				=							outputs.logits

		if model_name == "ast-finetuned-audioset-10-10-0.4593":
				_a					:				Optional[Any]				=							torch.tensor([-0.8_760, -7.0_042, -8.6_602]  )
		elif model_name == "ast-finetuned-audioset-10-10-0.450":
				_a					:				Optional[Any]				=							torch.tensor([-1.1_986, -7.0_903, -8.2_718]  )
		elif model_name == "ast-finetuned-audioset-10-10-0.448":
				_a					:				List[Any]				=							torch.tensor([-2.6_128, -8.0_080, -9.4_344]  )
		elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
				_a					:				Union[str, Any]				=							torch.tensor([-1.5_080, -7.4_534, -8.8_917]  )
		elif model_name == "ast-finetuned-audioset-12-12-0.447":
				_a					:				Optional[int]				=							torch.tensor([-0.5_050, -6.5_833, -8.0_843]  )
		elif model_name == "ast-finetuned-audioset-14-14-0.443":
				_a					:				Optional[int]				=							torch.tensor([-0.3_826, -7.0_336, -8.2_413]  )
		elif model_name == "ast-finetuned-audioset-16-16-0.442":
				_a					:				Union[str, Any]				=							torch.tensor([-1.2_113, -6.9_101, -8.3_470]  )
		elif model_name == "ast-finetuned-speech-commands-v2":
				_a					:				Optional[Any]				=							torch.tensor([6.1_589, -8.0_566, -8.7_984]  )
		else:
				raise ValueError('Unknown model name'  )
		if not torch.allclose(logits[0, :3]						,      lowerCAmelCase_						,      atol=1E-4  ):
				raise ValueError('Logits don\'t match'  )
		print('Looks ok!'  )

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

		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						      =       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						      =       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 os
from pathlib import Path

import torch
from bark.generation import _load_model as _bark_load_model
from huggingface_hub import hf_hub_download

from transformers import EncodecConfig, EncodecModel, set_seed
from transformers.models.bark.configuration_bark import (
    BarkCoarseConfig,
    BarkConfig,
    BarkFineConfig,
    BarkSemanticConfig,
)
from transformers.models.bark.generation_configuration_bark import (
    BarkCoarseGenerationConfig,
    BarkFineGenerationConfig,
    BarkGenerationConfig,
    BarkSemanticGenerationConfig,
)
from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCAmelCase						      =       logging.get_logger(__name__)

set_seed(770)


__lowerCAmelCase						      =       {
    '''c_attn''': '''att_proj''',
    '''c_proj''': '''out_proj''',
    '''c_fc''': '''in_proj''',
    '''transformer.''': '''''',
    '''h.''': '''layers.''',
    '''ln_1''': '''layernorm_1''',
    '''ln_2''': '''layernorm_2''',
    '''ln_f''': '''layernorm_final''',
    '''wpe''': '''position_embeds_layer''',
    '''wte''': '''input_embeds_layer''',
}


__lowerCAmelCase						      =       {
    '''text_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''text.pt''',
    },
    '''coarse_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''coarse.pt''',
    },
    '''fine_small''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''fine.pt''',
    },
    '''text''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''text_2.pt''',
    },
    '''coarse''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''coarse_2.pt''',
    },
    '''fine''': {
        '''repo_id''': '''suno/bark''',
        '''file_name''': '''fine_2.pt''',
    },
}

__lowerCAmelCase						      =       os.path.dirname(os.path.abspath(__file__))
__lowerCAmelCase						      =       os.path.join(os.path.expanduser('''~'''), '''.cache''')
__lowerCAmelCase						      =       os.path.join(os.getenv('''XDG_CACHE_HOME''', default_cache_dir), '''suno''', '''bark_v0''')



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_=False  )      ->				Optional[int]:
		_a					:				int				=							model_type
		if use_small:
				key += "_small"
		return os.path.join(lowerCAmelCase_						,      REMOTE_MODEL_PATHS[key]['file_name']  )



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_  )      ->				int:
		os.makedirs(lowerCAmelCase_						,      exist_ok=lowerCAmelCase_  )
		hf_hub_download(repo_id=lowerCAmelCase_						,      filename=lowerCAmelCase_						,      local_dir=lowerCAmelCase_  )



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_=False						,      lowerCAmelCase_="text"  )      ->				List[str]:
		if model_type == "text":
				_a					:				List[str]				=							BarkSemanticModel
				_a					:				Optional[Any]				=							BarkSemanticConfig
				_a					:				Any				=							BarkSemanticGenerationConfig
		elif model_type == "coarse":
				_a					:				Tuple				=							BarkCoarseModel
				_a					:				str				=							BarkCoarseConfig
				_a					:				str				=							BarkCoarseGenerationConfig
		elif model_type == "fine":
				_a					:				List[str]				=							BarkFineModel
				_a					:				Optional[Any]				=							BarkFineConfig
				_a					:				str				=							BarkFineGenerationConfig
		else:
				raise NotImplementedError()
		_a					:				Dict				=							f"""{model_type}_small""" if use_small else model_type
		_a					:				Union[str, Any]				=							REMOTE_MODEL_PATHS[model_key]
		if not os.path.exists(lowerCAmelCase_  ):
				logger.info(f"""{model_type} model not found, downloading into `{CACHE_DIR}`."""  )
				_download(model_info['repo_id']						,      model_info['file_name']  )
		_a					:				int				=							torch.load(lowerCAmelCase_						,      map_location=lowerCAmelCase_  )
		# this is a hack
		_a					:				List[Any]				=							checkpoint['model_args']
		if "input_vocab_size" not in model_args:
				_a					:				Dict				=							model_args['vocab_size']
				_a					:				Dict				=							model_args['vocab_size']
				del model_args["vocab_size"]

		# convert Bark model arguments to HF Bark model arguments
		_a					:				List[Any]				=							model_args.pop('n_head'  )
		_a					:				Any				=							model_args.pop('n_embd'  )
		_a					:				List[Any]				=							model_args.pop('n_layer'  )

		_a					:				Optional[int]				=							ConfigClass(**checkpoint['model_args']  )
		_a					:				List[str]				=							ModelClass(config=lowerCAmelCase_  )
		_a					:				Tuple				=							GenerationConfigClass()

		_a					:				Optional[Any]				=							model_generation_config
		_a					:				Optional[Any]				=							checkpoint['model']
		# fixup checkpoint
		_a					:				int				=							'_orig_mod.'
		for k, v in list(state_dict.items()  ):
				if k.startswith(lowerCAmelCase_  ):
						# replace part of the key with corresponding layer name in HF implementation
						_a					:				str				=							k[len(lowerCAmelCase_  ) :]
						for old_layer_name in new_layer_name_dict:
								_a					:				List[Any]				=							new_k.replace(lowerCAmelCase_						,      new_layer_name_dict[old_layer_name]  )

						_a					:				List[Any]				=							state_dict.pop(lowerCAmelCase_  )

		_a					:				List[Any]				=							set(state_dict.keys()  ) - set(model.state_dict().keys()  )
		_a					:				Tuple				=							{k for k in extra_keys if not k.endswith('.attn.bias'  )}
		_a					:				Tuple				=							set(model.state_dict().keys()  ) - set(state_dict.keys()  )
		_a					:				Optional[Any]				=							{k for k in missing_keys if not k.endswith('.attn.bias'  )}
		if len(lowerCAmelCase_  ) != 0:
				raise ValueError(f"""extra keys found: {extra_keys}"""  )
		if len(lowerCAmelCase_  ) != 0:
				raise ValueError(f"""missing keys: {missing_keys}"""  )
		model.load_state_dict(lowerCAmelCase_						,      strict=lowerCAmelCase_  )
		_a					:				Dict				=							model.num_parameters(exclude_embeddings=lowerCAmelCase_  )
		_a					:				Tuple				=							checkpoint['best_val_loss'].item()
		logger.info(f"""model loaded: {round(n_params/1E6						,      1  )}M params, {round(lowerCAmelCase_						,      3  )} loss"""  )
		model.eval()
		model.to(lowerCAmelCase_  )
		del checkpoint, state_dict

		return model



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_=False						,      lowerCAmelCase_="text"  )      ->				List[Any]:
		if model_type not in ("text", "coarse", "fine"):
				raise NotImplementedError()

		_a					:				Optional[int]				=							'cpu'  # do conversion on cpu

		_a					:				Tuple				=							_get_ckpt_path(lowerCAmelCase_						,      use_small=lowerCAmelCase_  )
		_a					:				List[Any]				=							_load_model(lowerCAmelCase_						,      lowerCAmelCase_						,      model_type=lowerCAmelCase_						,      use_small=lowerCAmelCase_  )

		# load bark initial model
		_a					:				Any				=							_bark_load_model(lowerCAmelCase_						,      'cpu'						,      model_type=lowerCAmelCase_						,      use_small=lowerCAmelCase_  )

		if model_type == "text":
				_a					:				int				=							bark_model['model']

		if model.num_parameters(exclude_embeddings=lowerCAmelCase_  ) != bark_model.get_num_params():
				raise ValueError('initial and new models don\'t have the same number of parameters'  )

		# check if same output as the bark model
		_a					:				Any				=							5
		_a					:				List[str]				=							10

		if model_type in ["text", "coarse"]:
				_a					:				Dict				=							torch.randint(256						,      (batch_size, sequence_length)						,      dtype=torch.int  )
				_a					:				Dict				=							bark_model(lowerCAmelCase_  )[0]

				_a					:				Tuple				=							model(lowerCAmelCase_  )

				# take last logits
				_a					:				Optional[int]				=							output_new_model_total.logits[:, [-1], :]

		else:
				_a					:				List[str]				=							3
				_a					:				List[Any]				=							8
				_a					:				Tuple				=							torch.randint(256						,      (batch_size, sequence_length, n_codes_total)						,      dtype=torch.int  )

				_a					:				Union[str, Any]				=							model(lowerCAmelCase_						,      lowerCAmelCase_  )
				_a					:				int				=							bark_model(lowerCAmelCase_						,      lowerCAmelCase_  )

				_a					:				List[str]				=							output_new_model_total.logits

		# output difference should come from the difference of self-attention implementation design
		if output_new_model.shape != output_old_model.shape:
				raise ValueError('initial and new outputs don\'t have the same shape'  )
		if (output_new_model - output_old_model).abs().max().item() > 1E-3:
				raise ValueError('initial and new outputs are not equal'  )

		Path(lowerCAmelCase_  ).mkdir(exist_ok=lowerCAmelCase_  )
		model.save_pretrained(lowerCAmelCase_  )



def 					__lowerCamelCase					(				lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_						,      )      ->				Any:
		_a					:				Any				=							os.path.join(lowerCAmelCase_						,      lowerCAmelCase_  )

		_a					:				int				=							BarkSemanticConfig.from_pretrained(os.path.join(lowerCAmelCase_						,      'config.json'  )  )
		_a					:				Any				=							BarkCoarseConfig.from_pretrained(os.path.join(lowerCAmelCase_						,      'config.json'  )  )
		_a					:				List[Any]				=							BarkFineConfig.from_pretrained(os.path.join(lowerCAmelCase_						,      'config.json'  )  )
		_a					:				List[str]				=							EncodecConfig.from_pretrained('facebook/encodec_24khz'  )

		_a					:				str				=							BarkSemanticModel.from_pretrained(lowerCAmelCase_  )
		_a					:				Dict				=							BarkCoarseModel.from_pretrained(lowerCAmelCase_  )
		_a					:				int				=							BarkFineModel.from_pretrained(lowerCAmelCase_  )
		_a					:				List[Any]				=							EncodecModel.from_pretrained('facebook/encodec_24khz'  )

		_a					:				Any				=							BarkConfig.from_sub_model_configs(
		    lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_						,      lowerCAmelCase_  )

		_a					:				List[str]				=							BarkGenerationConfig.from_sub_model_configs(
		    semantic.generation_config						,      coarseAcoustic.generation_config						,      fineAcoustic.generation_config  )

		_a					:				Optional[Any]				=							BarkModel(lowerCAmelCase_  )

		_a					:				List[str]				=							semantic
		_a					:				Union[str, Any]				=							coarseAcoustic
		_a					:				Optional[int]				=							fineAcoustic
		_a					:				Optional[Any]				=							codec

		_a					:				List[Any]				=							bark_generation_config

		Path(lowerCAmelCase_  ).mkdir(exist_ok=lowerCAmelCase_  )
		bark.save_pretrained(lowerCAmelCase_						,      repo_id=lowerCAmelCase_						,      push_to_hub=lowerCAmelCase_  )


if __name__ == "__main__":
							__lowerCAmelCase						      =       argparse.ArgumentParser()
							# Required parameters

							parser.add_argument('''model_type''', type=str, help='''text, coarse or fine.''')
							parser.add_argument('''pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
							parser.add_argument('''--is_small''', action='''store_true''', help='''convert the small version instead of the large.''')

							__lowerCAmelCase						      =       parser.parse_args()

							load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)




