
import argparse
import json
import os

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

from transformers import ConvNextImageProcessor, ConvNextVaConfig, ConvNextVaForImageClassification
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


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





def    _a (						lowerCamelCase							):
 lowerCamelCase		:     List[Any]          =				ConvNextVaConfig()

 if "atto" in checkpoint_url:
  lowerCamelCase		:     Any          =				[2, 2, 6, 2]
  lowerCamelCase		:     Optional[Any]          =				[40, 80, 160, 320]
 if "femto" in checkpoint_url:
  lowerCamelCase		:     Union[str, Any]          =				[2, 2, 6, 2]
  lowerCamelCase		:     str          =				[48, 96, 192, 384]
 if "pico" in checkpoint_url:
  lowerCamelCase		:     Optional[int]          =				[2, 2, 6, 2]
  lowerCamelCase		:     str          =				[64, 128, 256, 512]
 if "nano" in checkpoint_url:
  lowerCamelCase		:     str          =				[2, 2, 8, 2]
  lowerCamelCase		:     Union[str, Any]          =				[80, 160, 320, 640]
 if "tiny" in checkpoint_url:
  lowerCamelCase		:     Union[str, Any]          =				[3, 3, 9, 3]
  lowerCamelCase		:     Optional[Any]          =				[96, 192, 384, 768]
 if "base" in checkpoint_url:
  lowerCamelCase		:     List[str]          =				[3, 3, 27, 3]
  lowerCamelCase		:     Tuple          =				[128, 256, 512, 1024]
 if "large" in checkpoint_url:
  lowerCamelCase		:     Optional[int]          =				[3, 3, 27, 3]
  lowerCamelCase		:     Any          =				[192, 384, 768, 1536]
 if "huge" in checkpoint_url:
  lowerCamelCase		:     Dict          =				[3, 3, 27, 3]
  lowerCamelCase		:     Union[str, Any]          =				[352, 704, 1408, 2816]

 lowerCamelCase		:     Any          =				1000
 lowerCamelCase		:     Optional[int]          =				"""imagenet-1k-id2label.json"""
 lowerCamelCase		:     List[str]          =				(1, 1000)

 lowerCamelCase		:     int          =				"""huggingface/label-files"""
 lowerCamelCase		:     Optional[int]          =				num_labels
 lowerCamelCase		:     Tuple          =				json.load(open(hf_hub_download(lowerCamelCase,     lowerCamelCase,     repo_type="""dataset"""							),     """r"""							)							)
 lowerCamelCase		:     str          =				{int(lowerCamelCase							): v for k, v in idalabel.items()}

 lowerCamelCase		:     Optional[Any]          =				idalabel
 lowerCamelCase		:     str          =				{v: k for k, v in idalabel.items()}
 lowerCamelCase		:     Union[str, Any]          =				hidden_sizes
 lowerCamelCase		:     List[Any]          =				depths

 return config, expected_shape





def    _a (						lowerCamelCase							):
 if "downsample_layers.0.0" in name:
  lowerCamelCase		:     List[Any]          =				name.replace("""downsample_layers.0.0""",     """embeddings.patch_embeddings"""							)
 if "downsample_layers.0.1" in name:
  lowerCamelCase		:     Any          =				name.replace("""downsample_layers.0.1""",     """embeddings.norm"""							)  # we rename to layernorm later on
 if "downsample_layers.1.0" in name:
  lowerCamelCase		:     Dict          =				name.replace("""downsample_layers.1.0""",     """stages.1.downsampling_layer.0"""							)
 if "downsample_layers.1.1" in name:
  lowerCamelCase		:     Optional[Any]          =				name.replace("""downsample_layers.1.1""",     """stages.1.downsampling_layer.1"""							)
 if "downsample_layers.2.0" in name:
  lowerCamelCase		:     Tuple          =				name.replace("""downsample_layers.2.0""",     """stages.2.downsampling_layer.0"""							)
 if "downsample_layers.2.1" in name:
  lowerCamelCase		:     List[Any]          =				name.replace("""downsample_layers.2.1""",     """stages.2.downsampling_layer.1"""							)
 if "downsample_layers.3.0" in name:
  lowerCamelCase		:     Union[str, Any]          =				name.replace("""downsample_layers.3.0""",     """stages.3.downsampling_layer.0"""							)
 if "downsample_layers.3.1" in name:
  lowerCamelCase		:     int          =				name.replace("""downsample_layers.3.1""",     """stages.3.downsampling_layer.1"""							)
 if "stages" in name and "downsampling_layer" not in name:
  # stages.0.0. for instance should be renamed to stages.0.layers.0.
  lowerCamelCase		:     List[Any]          =				name[: len("""stages.0"""							)] + """.layers""" + name[len("""stages.0"""							) :]
 if "gamma" in name:
  lowerCamelCase		:     List[Any]          =				name.replace("""gamma""",     """weight"""							)
 if "beta" in name:
  lowerCamelCase		:     Dict          =				name.replace("""beta""",     """bias"""							)
 if "stages" in name:
  lowerCamelCase		:     List[Any]          =				name.replace("""stages""",     """encoder.stages"""							)
 if "norm" in name:
  lowerCamelCase		:     List[str]          =				name.replace("""norm""",     """layernorm"""							)
 if "head" in name:
  lowerCamelCase		:     str          =				name.replace("""head""",     """classifier"""							)

 return name





def    _a (						):
 lowerCamelCase		:     List[str]          =				"""http://images.cocodataset.org/val2017/000000039769.jpg"""
 lowerCamelCase		:     List[str]          =				Image.open(requests.get(lowerCamelCase,     stream=lowerCamelCase							).raw							)
 return im





def    _a (						lowerCamelCase							):
 if "224" in checkpoint_url:
  lowerCamelCase		:     str          =				224
  lowerCamelCase		:     Any          =				224 / 256
 elif "384" in checkpoint_url:
  lowerCamelCase		:     Optional[Any]          =				384
  lowerCamelCase		:     List[Any]          =				None
 else:
  lowerCamelCase		:     Union[str, Any]          =				512
  lowerCamelCase		:     Any          =				None

 return ConvNextImageProcessor(
     size=lowerCamelCase,     crop_pct=lowerCamelCase,     image_mean=[0.4_8_5, 0.4_5_6, 0.4_0_6],     image_std=[0.2_2_9, 0.2_2_4, 0.2_2_5],     resample=PILImageResampling.BICUBIC,     )





@torch.no_grad()
def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 print("""Downloading original model from checkpoint..."""							)
 # define ConvNeXTV2 configuration based on URL
 lowerCamelCase		:     Any          =				get_convnextva_config(lowerCamelCase							)
 # load original state_dict from URL
 lowerCamelCase		:     int          =				torch.hub.load_state_dict_from_url(lowerCamelCase							)["""model"""]

 print("""Converting model parameters..."""							)
 # rename keys
 for key in state_dict.copy().keys():
  lowerCamelCase		:     List[str]          =				state_dict.pop(lowerCamelCase							)
  lowerCamelCase		:     List[Any]          =				val
 # add prefix to all keys expect classifier head
 for key in state_dict.copy().keys():
  lowerCamelCase		:     int          =				state_dict.pop(lowerCamelCase							)
  if not key.startswith("""classifier"""							):
   lowerCamelCase		:     List[Any]          =				"""convnextv2.""" + key
  lowerCamelCase		:     List[Any]          =				val

 # load HuggingFace model
 lowerCamelCase		:     List[Any]          =				ConvNextVaForImageClassification(lowerCamelCase							)
 model.load_state_dict(lowerCamelCase							)
 model.eval()

 # Check outputs on an image, prepared by ConvNextImageProcessor
 lowerCamelCase		:     Any          =				convert_preprocessor(lowerCamelCase							)
 lowerCamelCase		:     Any          =				preprocessor(images=prepare_img(),     return_tensors="""pt"""							)
 lowerCamelCase		:     List[str]          =				model(**lowerCamelCase							).logits

 # note: the logits below were obtained without center cropping
 if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
  lowerCamelCase		:     List[str]          =				torch.tensor([-0.3_9_3_0, 0.1_7_4_7, -0.5_2_4_6, 0.4_1_7_7, 0.4_2_9_5]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
  lowerCamelCase		:     List[Any]          =				torch.tensor([-0.1_7_2_7, -0.5_3_4_1, -0.7_8_1_8, -0.4_7_4_5, -0.6_5_6_6]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
  lowerCamelCase		:     List[Any]          =				torch.tensor([-0.0_3_3_3, 0.1_5_6_3, -0.9_1_3_7, 0.1_0_5_4, 0.0_3_8_1]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
  lowerCamelCase		:     Optional[Any]          =				torch.tensor([-0.1_7_4_4, -0.1_5_5_5, -0.0_7_1_3, 0.0_9_5_0, -0.1_4_3_1]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
  lowerCamelCase		:     Any          =				torch.tensor([0.9_9_9_6, 0.1_9_6_6, -0.4_3_8_6, -0.3_4_7_2, 0.6_6_6_1]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
  lowerCamelCase		:     Optional[int]          =				torch.tensor([-0.2_5_5_3, -0.6_7_0_8, -0.1_3_5_9, 0.2_5_1_8, -0.2_4_8_8]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
  lowerCamelCase		:     Tuple          =				torch.tensor([-0.0_6_7_3, -0.5_6_2_7, -0.3_7_5_3, -0.2_7_2_2, 0.0_1_7_8]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
  lowerCamelCase		:     str          =				torch.tensor([-0.6_3_7_7, -0.7_4_5_8, -0.2_1_5_0, 0.1_1_8_4, -0.0_5_9_7]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
  lowerCamelCase		:     str          =				torch.tensor([1.0_7_9_9, 0.2_3_2_2, -0.8_8_6_0, 1.0_2_1_9, 0.6_2_3_1]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
  lowerCamelCase		:     Tuple          =				torch.tensor([0.3_7_6_6, 0.4_9_1_7, -1.1_4_2_6, 0.9_9_4_2, 0.6_0_2_4]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
  lowerCamelCase		:     Union[str, Any]          =				torch.tensor([0.4_2_2_0, -0.6_9_1_9, -0.4_3_1_7, -0.2_8_8_1, -0.6_6_0_9]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
  lowerCamelCase		:     Any          =				torch.tensor([0.1_0_8_2, -0.8_2_8_6, -0.5_0_9_5, 0.4_6_8_1, -0.8_0_8_5]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
  lowerCamelCase		:     Tuple          =				torch.tensor([-0.2_4_1_9, -0.6_2_2_1, 0.2_1_7_6, -0.0_9_8_0, -0.7_5_2_7]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
  lowerCamelCase		:     Dict          =				torch.tensor([0.0_3_9_1, -0.4_3_7_1, 0.3_7_8_6, 0.1_2_5_1, -0.2_7_8_4]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
  lowerCamelCase		:     Any          =				torch.tensor([-0.0_5_0_4, 0.5_6_3_6, -0.1_7_2_9, -0.6_5_0_7, -0.3_9_4_9]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
  lowerCamelCase		:     int          =				torch.tensor([0.3_5_6_0, 0.9_4_8_6, 0.3_1_4_9, -0.2_6_6_7, -0.5_1_3_8]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
  lowerCamelCase		:     Tuple          =				torch.tensor([-0.2_4_6_9, -0.4_5_5_0, -0.5_8_5_3, -0.0_8_1_0, 0.0_3_0_9]							)
 elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
  lowerCamelCase		:     int          =				torch.tensor([-0.3_0_9_0, 0.0_8_0_2, -0.0_6_8_2, -0.1_9_7_9, -0.2_8_2_6]							)
 else:
  raise ValueError(F'''Unknown URL: {checkpoint_url}'''							)

 assert torch.allclose(logits[0, :5],     lowerCamelCase,     atol=1e-3							)
 assert logits.shape == expected_shape
 print("""Model outputs match the original results!"""							)

 if save_model:
  print("""Saving model to local..."""							)
  # Create folder to save model
  if not os.path.isdir(lowerCamelCase							):
   os.mkdir(lowerCamelCase							)

  model.save_pretrained(lowerCamelCase							)
  preprocessor.save_pretrained(lowerCamelCase							)

 lowerCamelCase		:     Dict          =				"""convnextv2"""
 if "atto" in checkpoint_url:
  model_name += "-atto"
 if "femto" in checkpoint_url:
  model_name += "-femto"
 if "pico" in checkpoint_url:
  model_name += "-pico"
 if "nano" in checkpoint_url:
  model_name += "-nano"
 elif "tiny" in checkpoint_url:
  model_name += "-tiny"
 elif "base" in checkpoint_url:
  model_name += "-base"
 elif "large" in checkpoint_url:
  model_name += "-large"
 elif "huge" in checkpoint_url:
  model_name += "-huge"
 if "22k" in checkpoint_url and "1k" not in checkpoint_url:
  model_name += "-22k"
 elif "22k" in checkpoint_url and "1k" in checkpoint_url:
  model_name += "-22k-1k"
 elif "1k" in checkpoint_url:
  model_name += "-1k"
 if "224" in checkpoint_url:
  model_name += "-224"
 elif "384" in checkpoint_url:
  model_name += "-384"
 elif "512" in checkpoint_url:
  model_name += "-512"

 if push_to_hub:
  print(F'''Pushing {model_name} to the hub...'''							)
  model.push_to_hub(lowerCamelCase							)
  preprocessor.push_to_hub(lowerCamelCase							)


if __name__ == "__main__":
       _lowerCamelCase        =argparse.ArgumentParser()
       # Required parameters
       parser.add_argument(
           """--checkpoint_url""",
           default="""https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt""",
           type=str,
           help="""URL of the original ConvNeXTV2 checkpoint you'd like to convert.""",
       )
       parser.add_argument(
           """--pytorch_dump_folder_path""",
           default="""model""",
           type=str,
           help="""Path to the output PyTorch model directory.""",
       )
       parser.add_argument("""--save_model""", action="""store_true""", help="""Save model to local""")
       parser.add_argument("""--push_to_hub""", action="""store_true""", help="""Push model and image preprocessor to the hub""")

       _lowerCamelCase        =parser.parse_args()
       convert_convnextva_checkpoint(
           args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
       )






import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


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

_lowerCamelCase        ={
    """post_extract_proj""": """feature_projection.projection""",
    """models.0.layer_norm""": """feature_projection.layer_norm""",
    """self_attn.k_proj""": """encoder.layers.*.attention.k_proj""",
    """self_attn.v_proj""": """encoder.layers.*.attention.v_proj""",
    """self_attn.q_proj""": """encoder.layers.*.attention.q_proj""",
    """self_attn.out_proj""": """encoder.layers.*.attention.out_proj""",
    """self_attn_layer_norm""": """encoder.layers.*.layer_norm""",
    """fc1""": """encoder.layers.*.feed_forward.intermediate_dense""",
    """fc2""": """encoder.layers.*.feed_forward.output_dense""",
    """final_layer_norm""": """encoder.layers.*.final_layer_norm""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
_lowerCamelCase        =[
    """lm_head""",
]





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 for attribute in key.split("""."""							):
  lowerCamelCase		:     str          =				getattr(lowerCamelCase,     lowerCamelCase							)

 if weight_type is not None:
  lowerCamelCase		:     Union[str, Any]          =				getattr(lowerCamelCase,     lowerCamelCase							).shape
 else:
  lowerCamelCase		:     Optional[int]          =				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":
  lowerCamelCase		:     Union[str, Any]          =				value
 elif weight_type == "weight_g":
  lowerCamelCase		:     Optional[Any]          =				value
 elif weight_type == "weight_v":
  lowerCamelCase		:     List[str]          =				value
 elif weight_type == "bias":
  lowerCamelCase		:     Any          =				value
 else:
  lowerCamelCase		:     str          =				value

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





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     List[str]          =				[]
 lowerCamelCase		:     Union[str, Any]          =				fairseq_model.state_dict()

 if not is_headless:
  lowerCamelCase		:     Tuple          =				hf_model.dataavec_audio.feature_extractor
  lowerCamelCase		:     Tuple          =				hf_model.dataavec_audio.encoder.pos_conv_embed

 else:
  lowerCamelCase		:     Tuple          =				hf_model.feature_extractor
  lowerCamelCase		:     int          =				hf_model.encoder.pos_conv_embed

 for name, value in fairseq_dict.items():
  lowerCamelCase		:     Any          =				False
  if "conv_layers" in name:
   load_conv_layer(
       lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     )
   lowerCamelCase		:     str          =				True
  elif "pos_conv" in name:
   load_pos_conv_layer(
       lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     )
   lowerCamelCase		:     Union[str, Any]          =				True
  else:
   for key, mapped_key in MAPPING.items():
    if not is_headless:
     lowerCamelCase		:     List[Any]          =				"""data2vec_audio.""" + 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]:
     lowerCamelCase		:     Union[str, Any]          =				True
     if "*" in mapped_key:
      lowerCamelCase		:     Tuple          =				name.split(lowerCamelCase							)[0].split("""."""							)[-2]
      lowerCamelCase		:     Union[str, Any]          =				mapped_key.replace("""*""",     lowerCamelCase							)
     if "weight_g" in name:
      lowerCamelCase		:     Tuple          =				"""weight_g"""
     elif "weight_v" in name:
      lowerCamelCase		:     Tuple          =				"""weight_v"""
     elif "bias" in name:
      lowerCamelCase		:     List[str]          =				"""bias"""
     elif "weight" in name:
      # TODO: don't match quantizer.weight_proj
      lowerCamelCase		:     List[Any]          =				"""weight"""
     else:
      lowerCamelCase		:     Any          =				None
     set_recursively(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							)
    continue
  if not is_used:
   unused_weights.append(lowerCamelCase							)

 logger.warning(F'''Unused weights: {unused_weights}'''							)





def    _a (						lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Union[str, Any]          =				path.split("""."""							)
 return reduce(lowerCamelCase,     lowerCamelCase,     lowerCamelCase							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Tuple          =				access_by_string(lowerCamelCase,     lowerCamelCase							)
 lowerCamelCase		:     List[Any]          =				hf_weight.data

 if fsq_value.shape != hf_value.shape:
  raise ValueError(F'''{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.'''							)
 lowerCamelCase		:     Dict          =				fsq_value
 logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Optional[int]          =				full_name.split("""conv_layers."""							)[-1]
 lowerCamelCase		:     int          =				name.split("""."""							)
 lowerCamelCase		:     Any          =				int(items[0]							)
 lowerCamelCase		:     Union[str, Any]          =				int(items[1]							)

 lowerCamelCase		:     Dict          =				name.split("""."""							)[-1]
 if type_id == 0:
  lowerCamelCase		:     Optional[int]          =				"""conv"""
 elif type_id == 2:
  lowerCamelCase		:     Dict          =				"""layer_norm"""
 else:
  unused_weights.append(lowerCamelCase							)
  return

 set_weights(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     F'''conv_layers.{layer_id}.{layer_type}.{weight_type}'''							)





def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     lowerCamelCase							):
 lowerCamelCase		:     Dict          =				full_name.split("""pos_conv."""							)[-1]
 lowerCamelCase		:     Any          =				name.split("""."""							)
 lowerCamelCase		:     List[str]          =				int(items[0]							)
 lowerCamelCase		:     List[str]          =				int(items[1]							)

 lowerCamelCase		:     List[Any]          =				name.split("""."""							)[-1]
 if type_id != 0:
  unused_weights.append(lowerCamelCase							)
  return
 else:
  lowerCamelCase		:     Union[str, Any]          =				"""conv"""

 set_weights(lowerCamelCase,     lowerCamelCase,     lowerCamelCase,     F'''layers.{layer_id}.{layer_type}.{weight_type}'''							)





@torch.no_grad()
def    _a (						lowerCamelCase,     lowerCamelCase,     lowerCamelCase=None,     lowerCamelCase=None,     lowerCamelCase=True							):
 if config_path is not None:
  lowerCamelCase		:     Optional[Any]          =				DataaVecAudioConfig.from_pretrained(lowerCamelCase							)
 else:
  lowerCamelCase		:     Union[str, Any]          =				DataaVecAudioConfig()

 if not is_finetuned:
  # Modify final_proj layer name
  lowerCamelCase		:     Optional[Any]          =				DataaVecAudioModel(lowerCamelCase							)
  lowerCamelCase		:     List[Any]          =				os.path.dirname(lowerCamelCase							)

  lowerCamelCase		:     Union[str, Any]          =				torch.load(lowerCamelCase							)
  lowerCamelCase		:     Optional[int]          =				state_dict["""model"""].pop("""final_proj.0.weight"""							)
  lowerCamelCase		:     Optional[int]          =				state_dict["""model"""].pop("""final_proj.0.bias"""							)
  lowerCamelCase		:     Tuple          =				os.path.join(lowerCamelCase,     """converted.pt"""							)
  torch.save(lowerCamelCase,     lowerCamelCase							)
 else:
  lowerCamelCase		:     Optional[Any]          =				DataaVecAudioForCTC(lowerCamelCase							)
  lowerCamelCase		:     Union[str, Any]          =				checkpoint_path

 def load_dataavec(lowerCamelCase							):
  lowerCamelCase      ,   lowerCamelCase      ,   lowerCamelCase		:     int          =				fairseq.checkpoint_utils.load_model_ensemble_and_task([path]							)
  return model[0].eval()

 lowerCamelCase		:     Optional[int]          =				load_dataavec(lowerCamelCase							)

 recursively_load_weights(lowerCamelCase,     lowerCamelCase,     not is_finetuned							)

 lowerCamelCase		:     int          =				WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""							)

 lowerCamelCase		:     Union[str, Any]          =				load_dataset("""patrickvonplaten/librispeech_asr_dummy""",     """clean""",     split="""validation"""							)
 lowerCamelCase		:     Any          =				[x["""array"""] for x in ds[:4]["""audio"""]]

 lowerCamelCase		:     List[Any]          =				processor(lowerCamelCase,     return_tensors="""pt""",     padding=lowerCamelCase							)

 lowerCamelCase		:     Optional[Any]          =				inputs.input_values
 lowerCamelCase		:     int          =				inputs.attention_mask
 #    input_values = inputs.input_values[:, :-1]
 #    attention_mask = inputs.attention_mask[:, :-1]

 hf_wavavec.eval()
 model.eval()
 if is_finetuned:
  lowerCamelCase		:     List[Any]          =				model(source=lowerCamelCase,     padding_mask=(1 - attention_mask),     mask=lowerCamelCase,     features_only=lowerCamelCase							)[
      """encoder_out"""
  ].transpose(0,     1							)
  lowerCamelCase		:     Tuple          =				hf_wavavec(lowerCamelCase,     attention_mask=lowerCamelCase							)["""logits"""]

  lowerCamelCase		:     Any          =				torch.argmax(lowerCamelCase,     dim=-1							)
  lowerCamelCase		:     List[str]          =				processor.batch_decode(lowerCamelCase							)

  print(F'''Expected Output: {ds[:4]["text"]}, Pred: {output_string}'''							)
 else:
  lowerCamelCase		:     Union[str, Any]          =				model(source=lowerCamelCase,     padding_mask=(1 - attention_mask),     mask=lowerCamelCase,     features_only=lowerCamelCase							)[
      """layer_results"""
  ][-1][0].transpose(0,     1							)
  lowerCamelCase		:     Optional[Any]          =				hf_wavavec(lowerCamelCase,     attention_mask=lowerCamelCase							)["""last_hidden_state"""]

 print(our_output.shape,     their_output.shape							)
 lowerCamelCase		:     str          =				torch.max(torch.abs(our_output - their_output							)							).item()
 print(F'''max_absolute_diff = {max_absolute_diff}'''							)  # ~ 1e-7
 lowerCamelCase		:     int          =				torch.allclose(lowerCamelCase,     lowerCamelCase,     atol=1e-3							)
 print("""Do both models output the same tensors?""",     """🔥""" if success else """💩"""							)
 if not success:
  raise Exception("""Something went wRoNg"""							)

 hf_wavavec.save_pretrained(lowerCamelCase							)

 if is_finetuned:
  processor.save_pretrained(lowerCamelCase							)
 else:
  processor.feature_extractor.save_pretrained(lowerCamelCase							)


if __name__ == "__main__":
      _lowerCamelCase        =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"""
      )
      _lowerCamelCase        =parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
      )




