

import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConformerConfig,
    WavaVecaConformerForCTC,
    WavaVecaConformerForPreTraining,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
UpperCamelCase__  :      int			=    logging.get_logger(__name__)

UpperCamelCase__  :      int			=    {
    "post_extract_proj": "feature_projection.projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "self_attn.linear_k": "encoder.layers.*.self_attn.linear_k",
    "self_attn.linear_v": "encoder.layers.*.self_attn.linear_v",
    "self_attn.linear_q": "encoder.layers.*.self_attn.linear_q",
    "self_attn.pos_bias_u": "encoder.layers.*.self_attn.pos_bias_u",
    "self_attn.pos_bias_v": "encoder.layers.*.self_attn.pos_bias_v",
    "self_attn.linear_out": "encoder.layers.*.self_attn.linear_out",
    "self_attn.linear_pos": "encoder.layers.*.self_attn.linear_pos",
    "self_attn.rotary_emb": "encoder.embed_positions",
    "self_attn_layer_norm": "encoder.layers.*.self_attn_layer_norm",
    "conv_module.pointwise_conv1": "encoder.layers.*.conv_module.pointwise_conv1",
    "conv_module.pointwise_conv2": "encoder.layers.*.conv_module.pointwise_conv2",
    "conv_module.depthwise_conv": "encoder.layers.*.conv_module.depthwise_conv",
    "conv_module.batch_norm": "encoder.layers.*.conv_module.batch_norm",
    "conv_module.layer_norm": "encoder.layers.*.conv_module.layer_norm",
    "ffn1.w_1": "encoder.layers.*.ffn1.intermediate_dense",
    "ffn1.w_2": "encoder.layers.*.ffn1.output_dense",
    "ffn1.layer_norm": "encoder.layers.*.ffn1_layer_norm",
    "ffn2.w_1": "encoder.layers.*.ffn2.intermediate_dense",
    "ffn2.w_2": "encoder.layers.*.ffn2.output_dense",
    "ffn2.layer_norm": "encoder.layers.*.ffn2_layer_norm",
    "final_layer_norm": "encoder.layers.*.final_layer_norm",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "feature_projection.layer_norm",
    "quantizer.weight_proj": "quantizer.weight_proj",
    "quantizer.vars": "quantizer.codevectors",
    "project_q": "project_q",
    "final_proj": "project_hid",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
UpperCamelCase__  :      List[Any]			=    [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		int ,   _SCREAMING_SNAKE_CASE					:		Optional[Any] ,   _SCREAMING_SNAKE_CASE					:		Tuple ,   _SCREAMING_SNAKE_CASE					:		int ,   _SCREAMING_SNAKE_CASE					:		List[str]       ):





    """simple docstring"""


    for attribute in key.split('.'       ):
        SCREAMING_SNAKE_CASE_									=						getattr(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE       )

    if weight_type is not None:
        SCREAMING_SNAKE_CASE_									=						getattr(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE       ).shape
    else:
        SCREAMING_SNAKE_CASE_									=						hf_pointer.shape

    if hf_shape != value.shape:
        raise ValueError(
            f"""Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"""
            f""" {value.shape} for {full_name}"""       )

    if weight_type == "weight":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "weight_g":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "weight_v":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "bias":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "running_mean":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "running_var":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "num_batches_tracked":
        SCREAMING_SNAKE_CASE_									=						value
    elif weight_type == "inv_freq":
        SCREAMING_SNAKE_CASE_									=						value
    else:
        SCREAMING_SNAKE_CASE_									=						value

    logger.info(f"""{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}."""       )


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		List[str] ,   _SCREAMING_SNAKE_CASE					:		Optional[Any] ,   _SCREAMING_SNAKE_CASE					:		Optional[int]       ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						[]
    SCREAMING_SNAKE_CASE_									=						fairseq_model.state_dict()

    SCREAMING_SNAKE_CASE_									=						hf_model.wavaveca_conformer.feature_extractor

    for name, value in fairseq_dict.items():
        SCREAMING_SNAKE_CASE_									=						False
        if "conv_layers" in name:
            load_conv_layer(
                _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   hf_model.config.feat_extract_norm == 'group' ,   )
            SCREAMING_SNAKE_CASE_									=						True
        else:
            for key, mapped_key in MAPPING.items():
                SCREAMING_SNAKE_CASE_									=						'wav2vec2_conformer.' + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
                if key in name or key.split('w2v_model.'       )[-1] == name.split('.'       )[0]:
                    SCREAMING_SNAKE_CASE_									=						True
                    if "*" in mapped_key:
                        SCREAMING_SNAKE_CASE_									=						name.split(_SCREAMING_SNAKE_CASE       )[0].split('.'       )[-2]
                        SCREAMING_SNAKE_CASE_									=						mapped_key.replace('*' ,   _SCREAMING_SNAKE_CASE       )
                    if "pos_bias_u" in name:
                        SCREAMING_SNAKE_CASE_									=						None
                    elif "pos_bias_v" in name:
                        SCREAMING_SNAKE_CASE_									=						None
                    elif "weight_g" in name:
                        SCREAMING_SNAKE_CASE_									=						'weight_g'
                    elif "weight_v" in name:
                        SCREAMING_SNAKE_CASE_									=						'weight_v'
                    elif "bias" in name:
                        SCREAMING_SNAKE_CASE_									=						'bias'
                    elif "weight" in name:
                        # TODO: don't match quantizer.weight_proj
                        SCREAMING_SNAKE_CASE_									=						'weight'
                    elif "running_mean" in name:
                        SCREAMING_SNAKE_CASE_									=						'running_mean'
                    elif "inv_freq" in name:
                        SCREAMING_SNAKE_CASE_									=						'inv_freq'
                    elif "running_var" in name:
                        SCREAMING_SNAKE_CASE_									=						'running_var'
                    elif "num_batches_tracked" in name:
                        SCREAMING_SNAKE_CASE_									=						'num_batches_tracked'
                    else:
                        SCREAMING_SNAKE_CASE_									=						None
                    set_recursively(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE       )
                continue
        if not is_used:
            unused_weights.append(_SCREAMING_SNAKE_CASE       )

    logger.warning(f"""Unused weights: {unused_weights}"""       )


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		List[Any] ,   _SCREAMING_SNAKE_CASE					:		Optional[Any] ,   _SCREAMING_SNAKE_CASE					:		str ,   _SCREAMING_SNAKE_CASE					:		Dict ,   _SCREAMING_SNAKE_CASE					:		Dict       ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						full_name.split('conv_layers.'       )[-1]
    SCREAMING_SNAKE_CASE_									=						name.split('.'       )
    SCREAMING_SNAKE_CASE_									=						int(items[0]       )
    SCREAMING_SNAKE_CASE_									=						int(items[1]       )

    if type_id == 0:
        if "bias" in name:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                raise ValueError(
                    f"""{full_name} has size {value.shape}, but"""
                    f""" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."""       )
            SCREAMING_SNAKE_CASE_									=						value
            logger.info(f"""Feat extract conv layer {layer_id} was initialized from {full_name}."""       )
        elif "weight" in name:
            if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                raise ValueError(
                    f"""{full_name} has size {value.shape}, but"""
                    f""" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."""       )
            SCREAMING_SNAKE_CASE_									=						value
            logger.info(f"""Feat extract conv layer {layer_id} was initialized from {full_name}."""       )
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                raise ValueError(
                    f"""{full_name} has size {value.shape}, but"""
                    f""" {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found."""       )
            SCREAMING_SNAKE_CASE_									=						value
            logger.info(f"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""       )
        elif "weight" in name:
            if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                raise ValueError(
                    f"""{full_name} has size {value.shape}, but"""
                    f""" {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found."""       )
            SCREAMING_SNAKE_CASE_									=						value
            logger.info(f"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""       )
    else:
        unused_weights.append(_SCREAMING_SNAKE_CASE       )


@torch.no_grad()
def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		str ,   _SCREAMING_SNAKE_CASE					:		Optional[Any] ,   _SCREAMING_SNAKE_CASE					:		Tuple=None ,   _SCREAMING_SNAKE_CASE					:		List[str]=None ,   _SCREAMING_SNAKE_CASE					:		Any=True       ):





    """simple docstring"""


    if config_path is not None:
        SCREAMING_SNAKE_CASE_									=						WavaVecaConformerConfig.from_pretrained(_SCREAMING_SNAKE_CASE ,   hidden_act='swish'       )
    else:
        SCREAMING_SNAKE_CASE_									=						WavaVecaConformerConfig()

    if "rope" in checkpoint_path:
        SCREAMING_SNAKE_CASE_									=						'rotary'

    if is_finetuned:
        if dict_path:
            SCREAMING_SNAKE_CASE_									=						Dictionary.load(_SCREAMING_SNAKE_CASE       )

            # important change bos & pad token id since CTC symbol is <pad> and
            # not <s> as in fairseq
            SCREAMING_SNAKE_CASE_									=						target_dict.pad_index
            SCREAMING_SNAKE_CASE_									=						target_dict.bos_index
            SCREAMING_SNAKE_CASE_									=						target_dict.eos_index
            SCREAMING_SNAKE_CASE_									=						len(target_dict.symbols       )
            SCREAMING_SNAKE_CASE_									=						os.path.join(_SCREAMING_SNAKE_CASE ,   'vocab.json'       )
            if not os.path.isdir(_SCREAMING_SNAKE_CASE       ):
                logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(_SCREAMING_SNAKE_CASE       )       )
                return
            os.makedirs(_SCREAMING_SNAKE_CASE ,   exist_ok=_SCREAMING_SNAKE_CASE       )
            SCREAMING_SNAKE_CASE_									=						target_dict.indices

            # fairseq has the <pad> and <s> switched
            SCREAMING_SNAKE_CASE_									=						0
            SCREAMING_SNAKE_CASE_									=						1
            with open(_SCREAMING_SNAKE_CASE ,   'w' ,   encoding='utf-8'       ) as vocab_handle:
                json.dump(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE       )
            SCREAMING_SNAKE_CASE_									=						WavaVecaCTCTokenizer(
                _SCREAMING_SNAKE_CASE ,   unk_token=target_dict.unk_word ,   pad_token=target_dict.pad_word ,   bos_token=target_dict.bos_word ,   eos_token=target_dict.eos_word ,   word_delimiter_token='|' ,   do_lower_case=_SCREAMING_SNAKE_CASE ,   )
            SCREAMING_SNAKE_CASE_									=						True if config.feat_extract_norm == 'layer' else False
            SCREAMING_SNAKE_CASE_									=						WavaVecaFeatureExtractor(
                feature_size=1 ,   sampling_rate=16_000 ,   padding_value=0 ,   do_normalize=_SCREAMING_SNAKE_CASE ,   return_attention_mask=_SCREAMING_SNAKE_CASE ,   )
            SCREAMING_SNAKE_CASE_									=						WavaVecaProcessor(feature_extractor=_SCREAMING_SNAKE_CASE ,   tokenizer=_SCREAMING_SNAKE_CASE       )
            processor.save_pretrained(_SCREAMING_SNAKE_CASE       )

        SCREAMING_SNAKE_CASE_									=						WavaVecaConformerForCTC(_SCREAMING_SNAKE_CASE       )
    else:
        SCREAMING_SNAKE_CASE_									=						WavaVecaConformerForPreTraining(_SCREAMING_SNAKE_CASE       )

    if is_finetuned:
        SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path] ,   arg_overrides={'data': '/'.join(dict_path.split('/'       )[:-1]       )}       )
    else:
        SCREAMING_SNAKE_CASE_									=						argparse.Namespace(task='audio_pretraining'       )
        SCREAMING_SNAKE_CASE_									=						fairseq.tasks.setup_task(_SCREAMING_SNAKE_CASE       )

        SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path] ,   task=_SCREAMING_SNAKE_CASE       )

    SCREAMING_SNAKE_CASE_									=						model[0].eval()

    recursively_load_weights(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   not is_finetuned       )

    hf_wavavec.save_pretrained(_SCREAMING_SNAKE_CASE       )


if __name__ == "__main__":
       UpperCamelCase__  :      List[Any]			=    argparse.ArgumentParser()
       parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
       parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
       parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
       parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
       parser.add_argument(
           "--not_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not"
       )
       UpperCamelCase__  :      Tuple			=    parser.parse_args()
       convert_wavaveca_conformer_checkpoint(
           args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
       )





import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    SwiftFormerConfig,
    SwiftFormerForImageClassification,
    ViTImageProcessor,
)
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase__  :      Optional[int]			=    logging.get_logger(__name__)

UpperCamelCase__  :      List[Any]			=    torch.device("cpu")


def 					_UpperCAmelCase       (  ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						'http://images.cocodataset.org/val2017/000000039769.jpg'
    SCREAMING_SNAKE_CASE_									=						Image.open(requests.get(_SCREAMING_SNAKE_CASE ,   stream=_SCREAMING_SNAKE_CASE       ).raw       )
    return im


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		int       ):





    """simple docstring"""


    if swiftformer_name == "swiftformer_xs":
        return torch.tensor([-2.1_7_0_3E0_0, 2.1_1_0_7E0_0, -2.0_8_1_1E0_0, 8.8_6_8_5E-0_1, 2.4_3_6_0E-0_1]       )

    elif swiftformer_name == "swiftformer_s":
        return torch.tensor([3.9_6_3_6E-0_1, 2.3_4_7_8E-0_1, -1.6_9_6_3E0_0, -1.7_3_8_1E0_0, -8.6_3_3_7E-0_1]       )

    elif swiftformer_name == "swiftformer_l1":
        return torch.tensor([-4.2_7_6_8E-0_1, -4.7_4_2_9E-0_1, -1.0_8_9_7E0_0, -1.0_2_4_8E0_0, 3.5_5_2_3E-0_2]       )

    elif swiftformer_name == "swiftformer_l3":
        return torch.tensor([-2.5_3_3_0E-0_1, 2.4_2_1_1E-0_1, -6.0_1_8_5E-0_1, -8.2_7_8_9E-0_1, -6.0_4_4_6E-0_2]       )


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		Optional[int] ,   _SCREAMING_SNAKE_CASE					:		Union[str, Any] ,   _SCREAMING_SNAKE_CASE					:		Optional[Any]       ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						dct.pop(_SCREAMING_SNAKE_CASE       )
    SCREAMING_SNAKE_CASE_									=						val


def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		Dict       ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						[]
    for k in state_dict.keys():
        SCREAMING_SNAKE_CASE_									=						k
        if ".pwconv" in k:
            SCREAMING_SNAKE_CASE_									=						k_new.replace('.pwconv' ,   '.point_wise_conv'       )
        if ".dwconv" in k:
            SCREAMING_SNAKE_CASE_									=						k_new.replace('.dwconv' ,   '.depth_wise_conv'       )
        if ".Proj." in k:
            SCREAMING_SNAKE_CASE_									=						k_new.replace('.Proj.' ,   '.proj.'       )
        if "patch_embed" in k_new:
            SCREAMING_SNAKE_CASE_									=						k_new.replace('patch_embed' ,   'swiftformer.patch_embed.patch_embedding'       )
        if "network" in k_new:
            SCREAMING_SNAKE_CASE_									=						k_new.split('.'       )
            if ls[2].isdigit():
                SCREAMING_SNAKE_CASE_									=						'swiftformer.encoder.network.' + ls[1] + '.blocks.' + ls[2] + '.' + '.'.join(ls[3:]       )
            else:
                SCREAMING_SNAKE_CASE_									=						k_new.replace('network' ,   'swiftformer.encoder.network'       )
        rename_keys.append((k, k_new)       )
    return rename_keys


@torch.no_grad()
def 					_UpperCAmelCase       (  _SCREAMING_SNAKE_CASE					:		Optional[Any] ,   _SCREAMING_SNAKE_CASE					:		str ,   _SCREAMING_SNAKE_CASE					:		Optional[Any]       ):





    """simple docstring"""


    SCREAMING_SNAKE_CASE_									=						SwiftFormerConfig()

    # dataset (ImageNet-21k only or also fine-tuned on ImageNet 2012), patch_size and image_size
    SCREAMING_SNAKE_CASE_									=						1_000
    SCREAMING_SNAKE_CASE_									=						'huggingface/label-files'
    SCREAMING_SNAKE_CASE_									=						'imagenet-1k-id2label.json'
    SCREAMING_SNAKE_CASE_									=						json.load(open(hf_hub_download(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   repo_type='dataset'       ) ,   'r'       )       )
    SCREAMING_SNAKE_CASE_									=						{int(_SCREAMING_SNAKE_CASE       ): v for k, v in idalabel.items()}
    SCREAMING_SNAKE_CASE_									=						idalabel
    SCREAMING_SNAKE_CASE_									=						{v: k for k, v in idalabel.items()}

    # size of the architecture
    if swiftformer_name == "swiftformer_xs":
        SCREAMING_SNAKE_CASE_									=						[3, 3, 6, 4]
        SCREAMING_SNAKE_CASE_									=						[48, 56, 112, 220]

    elif swiftformer_name == "swiftformer_s":
        SCREAMING_SNAKE_CASE_									=						[3, 3, 9, 6]
        SCREAMING_SNAKE_CASE_									=						[48, 64, 168, 224]

    elif swiftformer_name == "swiftformer_l1":
        SCREAMING_SNAKE_CASE_									=						[4, 3, 10, 5]
        SCREAMING_SNAKE_CASE_									=						[48, 96, 192, 384]

    elif swiftformer_name == "swiftformer_l3":
        SCREAMING_SNAKE_CASE_									=						[4, 4, 12, 6]
        SCREAMING_SNAKE_CASE_									=						[64, 128, 320, 512]

    # load state_dict of original model, remove and rename some keys
    if original_ckpt:
        if original_ckpt.startswith('https'       ):
            SCREAMING_SNAKE_CASE_									=						torch.hub.load_state_dict_from_url(_SCREAMING_SNAKE_CASE ,   map_location='cpu' ,   check_hash=_SCREAMING_SNAKE_CASE       )
        else:
            SCREAMING_SNAKE_CASE_									=						torch.load(_SCREAMING_SNAKE_CASE ,   map_location='cpu'       )
    SCREAMING_SNAKE_CASE_									=						checkpoint

    SCREAMING_SNAKE_CASE_									=						create_rename_keys(_SCREAMING_SNAKE_CASE       )
    for rename_key_src, rename_key_dest in rename_keys:
        rename_key(_SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE ,   _SCREAMING_SNAKE_CASE       )

    # load HuggingFace model
    SCREAMING_SNAKE_CASE_									=						SwiftFormerForImageClassification(_SCREAMING_SNAKE_CASE       ).eval()
    hf_model.load_state_dict(_SCREAMING_SNAKE_CASE       )

    # prepare test inputs
    SCREAMING_SNAKE_CASE_									=						prepare_img()
    SCREAMING_SNAKE_CASE_									=						ViTImageProcessor.from_pretrained('preprocessor_config'       )
    SCREAMING_SNAKE_CASE_									=						processor(images=_SCREAMING_SNAKE_CASE ,   return_tensors='pt'       )

    # compare outputs from both models
    SCREAMING_SNAKE_CASE_									=						get_expected_output(_SCREAMING_SNAKE_CASE       )
    SCREAMING_SNAKE_CASE_									=						hf_model(inputs['pixel_values']       ).logits

    assert hf_logits.shape == torch.Size([1, 1_000]       )
    assert torch.allclose(hf_logits[0, 0:5] ,   _SCREAMING_SNAKE_CASE ,   atol=1E-3       )

    Path(_SCREAMING_SNAKE_CASE       ).mkdir(exist_ok=_SCREAMING_SNAKE_CASE       )
    print(f"""Saving model {swiftformer_name} to {pytorch_dump_folder_path}"""       )
    hf_model.save_pretrained(_SCREAMING_SNAKE_CASE       )


if __name__ == "__main__":
       UpperCamelCase__  :      str			=    argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           "--swiftformer_name",
           default="swiftformer_xs",
           choices=["swiftformer_xs", "swiftformer_s", "swiftformer_l1", "swiftformer_l3"],
           type=str,
           help="Name of the SwiftFormer model you'd like to convert.",
       )
       parser.add_argument(
           "--pytorch_dump_folder_path",
           default="./converted_outputs/",
           type=str,
           help="Path to the output PyTorch model directory.",
       )
       parser.add_argument("--original_ckpt", default=None, type=str, help="Path to the original model checkpoint.")

       UpperCamelCase__  :      Union[str, Any]			=    parser.parse_args()
       convert_swiftformer_checkpoint(args.swiftformer_name, args.pytorch_dump_folder_path, args.original_ckpt)


