
import argparse
import json
from pathlib import Path

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

from transformers import (
    BeitConfig,
    BeitForImageClassification,
    BeitForMaskedImageModeling,
    BeitForSemanticSegmentation,
    BeitImageProcessor,
)
from transformers.image_utils import PILImageResampling
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase_  :       str					      =     logging.get_logger(__name__)







def 				A__    (  lowerCamelCase   ,			lowerCamelCase=False   ,			lowerCamelCase=False  )     ->		Optional[Any]:
				UpperCamelCase_:							Optional[Any]		=  """backbone.""" if is_semantic else """"""

				UpperCamelCase_:							Optional[Any]		=  []
				for i in range(config.num_hidden_layers  ):
								# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
								rename_keys.append((F'''{prefix}blocks.{i}.norm1.weight''', F'''beit.encoder.layer.{i}.layernorm_before.weight''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.norm1.bias''', F'''beit.encoder.layer.{i}.layernorm_before.bias''')  )
								rename_keys.append(
								    (F'''{prefix}blocks.{i}.attn.proj.weight''', F'''beit.encoder.layer.{i}.attention.output.dense.weight''')  )
								rename_keys.append(
								    (F'''{prefix}blocks.{i}.attn.proj.bias''', F'''beit.encoder.layer.{i}.attention.output.dense.bias''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.norm2.weight''', F'''beit.encoder.layer.{i}.layernorm_after.weight''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.norm2.bias''', F'''beit.encoder.layer.{i}.layernorm_after.bias''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc1.weight''', F'''beit.encoder.layer.{i}.intermediate.dense.weight''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc1.bias''', F'''beit.encoder.layer.{i}.intermediate.dense.bias''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc2.weight''', F'''beit.encoder.layer.{i}.output.dense.weight''')  )
								rename_keys.append((F'''{prefix}blocks.{i}.mlp.fc2.bias''', F'''beit.encoder.layer.{i}.output.dense.bias''')  )

				# projection layer + position embeddings
				rename_keys.extend(
				    [
				        (F'''{prefix}cls_token''', """beit.embeddings.cls_token"""),
				        (F'''{prefix}patch_embed.proj.weight''', """beit.embeddings.patch_embeddings.projection.weight"""),
				        (F'''{prefix}patch_embed.proj.bias''', """beit.embeddings.patch_embeddings.projection.bias"""),
				    ]  )

				if has_lm_head:
								# mask token + shared relative position bias + layernorm
								rename_keys.extend(
								    [
								        ("""mask_token""", """beit.embeddings.mask_token"""),
								        (
								            """rel_pos_bias.relative_position_bias_table""",
								            """beit.encoder.relative_position_bias.relative_position_bias_table""",
								        ),
								        (
								            """rel_pos_bias.relative_position_index""",
								            """beit.encoder.relative_position_bias.relative_position_index""",
								        ),
								        ("""norm.weight""", """layernorm.weight"""),
								        ("""norm.bias""", """layernorm.bias"""),
								    ]  )
				elif is_semantic:
								# semantic segmentation classification heads
								rename_keys.extend(
								    [
								        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
								        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
								        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
								        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
								    ]  )
				else:
								# layernorm + classification head
								rename_keys.extend(
								    [
								        ("""fc_norm.weight""", """beit.pooler.layernorm.weight"""),
								        ("""fc_norm.bias""", """beit.pooler.layernorm.bias"""),
								        ("""head.weight""", """classifier.weight"""),
								        ("""head.bias""", """classifier.bias"""),
								    ]  )

				return rename_keys







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase=False   ,			lowerCamelCase=False  )     ->		int:
				for i in range(config.num_hidden_layers  ):
								UpperCamelCase_:							Any		=  """backbone.""" if is_semantic else """"""
								# queries, keys and values
								UpperCamelCase_:							Optional[Any]		=  state_dict.pop(F'''{prefix}blocks.{i}.attn.qkv.weight'''  )
								UpperCamelCase_:							Tuple		=  state_dict.pop(F'''{prefix}blocks.{i}.attn.q_bias'''  )
								UpperCamelCase_:							Any		=  state_dict.pop(F'''{prefix}blocks.{i}.attn.v_bias'''  )

								UpperCamelCase_:							List[Any]		=  in_proj_weight[
								    : config.hidden_size, :
								]
								UpperCamelCase_:							Optional[int]		=  q_bias
								UpperCamelCase_:							int		=  in_proj_weight[
								    config.hidden_size : config.hidden_size * 2, :
								]
								UpperCamelCase_:							Tuple		=  in_proj_weight[
								    -config.hidden_size :, :
								]
								UpperCamelCase_:							Union[str, Any]		=  v_bias

								# gamma_1 and gamma_2
								# we call them lambda because otherwise they are renamed when using .from_pretrained
								UpperCamelCase_:							Tuple		=  state_dict.pop(F'''{prefix}blocks.{i}.gamma_1'''  )
								UpperCamelCase_:							Tuple		=  state_dict.pop(F'''{prefix}blocks.{i}.gamma_2'''  )

								UpperCamelCase_:							Tuple		=  gamma_a
								UpperCamelCase_:							Union[str, Any]		=  gamma_a

								# relative_position bias table + index
								if not has_lm_head:
												# each layer has its own relative position bias
												UpperCamelCase_:							List[Any]		=  state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_bias_table'''  )
												UpperCamelCase_:							Dict		=  state_dict.pop(F'''{prefix}blocks.{i}.attn.relative_position_index'''  )

												UpperCamelCase_:							Union[str, Any]		=  table
												UpperCamelCase_:							int		=  index







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		int:
				UpperCamelCase_:							int		=  dct.pop(lowerCamelCase  )
				UpperCamelCase_:							Any		=  val







def 				A__    (  )     ->		Optional[Any]:
				UpperCamelCase_:							Union[str, Any]		=  """http://images.cocodataset.org/val2017/000000039769.jpg"""
				UpperCamelCase_:							Union[str, Any]		=  Image.open(requests.get(lowerCamelCase   ,			stream=lowerCamelCase  ).raw  )
				return im







@torch.no_grad()
def 				A__    (  lowerCamelCase   ,			lowerCamelCase  )     ->		Optional[Any]:
				UpperCamelCase_:							Any		=  BeitConfig()
				UpperCamelCase_:							List[str]		=  False
				UpperCamelCase_:							List[Any]		=  False
				UpperCamelCase_:							List[str]		=  """huggingface/label-files"""
				# set config parameters based on URL
				if checkpoint_url[-9:-4] == "pt22k":
								# masked image modeling
								UpperCamelCase_:							Any		=  True
								UpperCamelCase_:							str		=  True
								UpperCamelCase_:							Optional[int]		=  True
				elif checkpoint_url[-9:-4] == "ft22k":
								# intermediate fine-tuning on ImageNet-22k
								UpperCamelCase_:							Any		=  True
								UpperCamelCase_:							Dict		=  2_18_41
								UpperCamelCase_:							Optional[Any]		=  """imagenet-22k-id2label.json"""
								UpperCamelCase_:							Any		=  json.load(open(hf_hub_download(lowerCamelCase   ,			lowerCamelCase   ,			repo_type="""dataset"""  )   ,			"""r"""  )  )
								UpperCamelCase_:							Any		=  {int(lowerCamelCase  ): v for k, v in idalabel.items()}
								# this dataset contains 21843 labels but the model only has 21841
								# we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
								del idalabel[92_05]
								del idalabel[1_50_27]
								UpperCamelCase_:							int		=  idalabel
								UpperCamelCase_:							Any		=  {v: k for k, v in idalabel.items()}
				elif checkpoint_url[-8:-4] == "to1k":
								# fine-tuning on ImageNet-1k
								UpperCamelCase_:							Dict		=  True
								UpperCamelCase_:							str		=  10_00
								UpperCamelCase_:							Optional[int]		=  """imagenet-1k-id2label.json"""
								UpperCamelCase_:							Dict		=  json.load(open(hf_hub_download(lowerCamelCase   ,			lowerCamelCase   ,			repo_type="""dataset"""  )   ,			"""r"""  )  )
								UpperCamelCase_:							Any		=  {int(lowerCamelCase  ): v for k, v in idalabel.items()}
								UpperCamelCase_:							Union[str, Any]		=  idalabel
								UpperCamelCase_:							Tuple		=  {v: k for k, v in idalabel.items()}
								if "384" in checkpoint_url:
												UpperCamelCase_:							List[Any]		=  3_84
								if "512" in checkpoint_url:
												UpperCamelCase_:							List[str]		=  5_12
				elif "ade20k" in checkpoint_url:
								# fine-tuning
								UpperCamelCase_:							str		=  True
								UpperCamelCase_:							Dict		=  1_50
								UpperCamelCase_:							Any		=  """ade20k-id2label.json"""
								UpperCamelCase_:							Dict		=  json.load(open(hf_hub_download(lowerCamelCase   ,			lowerCamelCase   ,			repo_type="""dataset"""  )   ,			"""r"""  )  )
								UpperCamelCase_:							Any		=  {int(lowerCamelCase  ): v for k, v in idalabel.items()}
								UpperCamelCase_:							Union[str, Any]		=  idalabel
								UpperCamelCase_:							Any		=  {v: k for k, v in idalabel.items()}
								UpperCamelCase_:							Union[str, Any]		=  6_40
								UpperCamelCase_:							Any		=  True
				else:
								raise ValueError("""Checkpoint not supported, URL should either end with 'pt22k', 'ft22k', 'to1k' or 'ade20k'"""  )

				# size of the architecture
				if "base" in checkpoint_url:
								pass
				elif "large" in checkpoint_url:
								UpperCamelCase_:							Any		=  10_24
								UpperCamelCase_:							Dict		=  40_96
								UpperCamelCase_:							Optional[int]		=  24
								UpperCamelCase_:							Optional[Any]		=  16
								if "ade20k" in checkpoint_url:
												UpperCamelCase_:							List[str]		=  6_40
												UpperCamelCase_:							str		=  [7, 11, 15, 23]
				else:
								raise ValueError("""Should either find 'base' or 'large' in checkpoint URL"""  )

				# load state_dict of original model, remove and rename some keys
				UpperCamelCase_:							Union[str, Any]		=  torch.hub.load_state_dict_from_url(lowerCamelCase   ,			map_location="""cpu"""   ,			check_hash=lowerCamelCase  )
				UpperCamelCase_:							Optional[Any]		=  state_dict["""model"""] if """ade20k""" not in checkpoint_url else state_dict["""state_dict"""]

				UpperCamelCase_:							str		=  create_rename_keys(lowerCamelCase   ,			has_lm_head=lowerCamelCase   ,			is_semantic=lowerCamelCase  )
				for src, dest in rename_keys:
								rename_key(lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )
				read_in_q_k_v(lowerCamelCase   ,			lowerCamelCase   ,			has_lm_head=lowerCamelCase   ,			is_semantic=lowerCamelCase  )
				if is_semantic:
								# add prefix to decoder keys
								for key, val in state_dict.copy().items():
												UpperCamelCase_:							Dict		=  state_dict.pop(lowerCamelCase  )
												if key.startswith("""backbone.fpn"""  ):
																UpperCamelCase_:							Any		=  key.replace("""backbone.fpn"""   ,			"""fpn"""  )
												UpperCamelCase_:							Optional[Any]		=  val

    # load HuggingFace model
				if checkpoint_url[-9:-4] == "pt22k":
								UpperCamelCase_:							Tuple		=  BeitForMaskedImageModeling(lowerCamelCase  )
				elif "ade20k" in checkpoint_url:
								UpperCamelCase_:							str		=  BeitForSemanticSegmentation(lowerCamelCase  )
				else:
								UpperCamelCase_:							Union[str, Any]		=  BeitForImageClassification(lowerCamelCase  )
				model.eval()
				model.load_state_dict(lowerCamelCase  )

				# Check outputs on an image
				if is_semantic:
								UpperCamelCase_:							List[Any]		=  BeitImageProcessor(size=config.image_size   ,			do_center_crop=lowerCamelCase  )
								UpperCamelCase_:							Union[str, Any]		=  load_dataset("""hf-internal-testing/fixtures_ade20k"""   ,			split="""test"""  )
								UpperCamelCase_:							List[str]		=  Image.open(ds[0]["""file"""]  )
				else:
								UpperCamelCase_:							List[str]		=  BeitImageProcessor(
								    size=config.image_size   ,			resample=PILImageResampling.BILINEAR   ,			do_center_crop=lowerCamelCase  )
								UpperCamelCase_:							Tuple		=  prepare_img()

				UpperCamelCase_:							int		=  image_processor(images=lowerCamelCase   ,			return_tensors="""pt"""  )
				UpperCamelCase_:							List[str]		=  encoding["""pixel_values"""]

				UpperCamelCase_:							Union[str, Any]		=  model(lowerCamelCase  )
				UpperCamelCase_:							Any		=  outputs.logits

				# verify logits
				UpperCamelCase_:							Optional[int]		=  torch.Size([1, 10_00]  )
				if checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k"""  ):
								UpperCamelCase_:							str		=  torch.Size([1, 1_96, 81_92]  )
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k"""  ):
								UpperCamelCase_:							int		=  torch.Size([1, 1_96, 81_92]  )
				elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft22k"""  ):
								UpperCamelCase_:							str		=  torch.Size([1, 2_18_41]  )
								UpperCamelCase_:							str		=  torch.tensor([2.2288, 2.4671, 0.7395]  )
								UpperCamelCase_:							List[str]		=  23_97
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft22k"""  ):
								UpperCamelCase_:							List[Any]		=  torch.Size([1, 2_18_41]  )
								UpperCamelCase_:							int		=  torch.tensor([1.6881, -0.2787, 0.5901]  )
								UpperCamelCase_:							List[Any]		=  23_96
				elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft1k"""  ):
								UpperCamelCase_:							Tuple		=  torch.tensor([0.1241, 0.0798, -0.6569]  )
								UpperCamelCase_:							Any		=  2_85
				elif checkpoint_url[:-4].endswith("""beit_base_patch16_224_pt22k_ft22kto1k"""  ):
								UpperCamelCase_:							Optional[Any]		=  torch.tensor([-1.2385, -1.0987, -1.0108]  )
								UpperCamelCase_:							str		=  2_81
				elif checkpoint_url[:-4].endswith("""beit_base_patch16_384_pt22k_ft22kto1k"""  ):
								UpperCamelCase_:							Optional[int]		=  torch.tensor([-1.5303, -0.9484, -0.3147]  )
								UpperCamelCase_:							Optional[int]		=  7_61
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft1k"""  ):
								UpperCamelCase_:							Optional[Any]		=  torch.tensor([0.4610, -0.0928, 0.2086]  )
								UpperCamelCase_:							Any		=  7_61
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_224_pt22k_ft22kto1k"""  ):
								UpperCamelCase_:							Tuple		=  torch.tensor([-0.4804, 0.6257, -0.1837]  )
								UpperCamelCase_:							Any		=  7_61
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_384_pt22k_ft22kto1k"""  ):
								UpperCamelCase_:							Tuple		=  torch.tensor([[-0.5122, 0.5117, -0.2113]]  )
								UpperCamelCase_:							Optional[Any]		=  7_61
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_512_pt22k_ft22kto1k"""  ):
								UpperCamelCase_:							List[Any]		=  torch.tensor([-0.3062, 0.7261, 0.4852]  )
								UpperCamelCase_:							Optional[int]		=  7_61
				elif checkpoint_url[:-4].endswith("""beit_base_patch16_640_pt22k_ft22ktoade20k"""  ):
								UpperCamelCase_:							List[str]		=  (1, 1_50, 1_60, 1_60)
								UpperCamelCase_:							Any		=  torch.tensor(
								    [
								        [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
								        [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
								        [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
								    ]  )
				elif checkpoint_url[:-4].endswith("""beit_large_patch16_640_pt22k_ft22ktoade20k"""  ):
								UpperCamelCase_:							Tuple		=  (1, 1_50, 1_60, 1_60)
								UpperCamelCase_:							List[str]		=  torch.tensor(
								    [
								        [[-4.3305, -2.3049, -3.0161], [-2.9591, -1.5305, -2.2251], [-3.4198, -1.8004, -2.9062]],
								        [[-5.8922, -3.7435, -4.3978], [-4.2063, -2.7872, -3.4755], [-4.2791, -3.1874, -4.1681]],
								        [[0.9895, 4.3467, 4.7663], [4.2476, 5.6830, 6.1518], [4.5550, 6.2495, 6.5154]],
								    ]  )
				else:
								raise ValueError("""Can't verify logits as model is not supported"""  )

				assert logits.shape == expected_shape, "Shape of logits not as expected"
				if not has_lm_head:
								if is_semantic:
												assert torch.allclose(
												    logits[0, :3, :3, :3]   ,			lowerCamelCase   ,			atol=1E-3  ), "First elements of logits not as expected"
								else:
												print("""Predicted class idx:"""   ,			logits.argmax(-1  ).item()  )
												assert torch.allclose(
												    logits[0, :3]   ,			lowerCamelCase   ,			atol=1E-3  ), "First elements of logits not as expected"
												assert logits.argmax(-1  ).item() == expected_class_idx, "Predicted class index not as expected"

				Path(lowerCamelCase  ).mkdir(exist_ok=lowerCamelCase  )
				print(F'''Saving model to {pytorch_dump_folder_path}'''  )
				model.save_pretrained(lowerCamelCase  )
				print(F'''Saving image processor to {pytorch_dump_folder_path}'''  )
				image_processor.save_pretrained(lowerCamelCase  )


if __name__ == "__main__":
							lowerCamelCase_  :       Dict					      =     argparse.ArgumentParser()

							parser.add_argument(
							    """--checkpoint_url""",
							    default="""https://conversationhub.blob.core.windows.net/beit-share-public/beit/beit_base_patch16_224_pt22k_ft22kto1k.pth""",
							    type=str,
							    help="""URL to the original PyTorch checkpoint (.pth file).""",
							)
							parser.add_argument(
							    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
							)
							lowerCamelCase_  :       int					      =     parser.parse_args()
							convert_beit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)








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_  :       int					      =     logging.get_logger(__name__)

lowerCamelCase_  :       Optional[int]					      =     {
    """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_  :       Any					      =     [
    """lm_head""",
]







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		List[Any]:
				for attribute in key.split("""."""  ):
								UpperCamelCase_:							List[str]		=  getattr(lowerCamelCase   ,			lowerCamelCase  )

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

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







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		Union[str, Any]:
				UpperCamelCase_:							Optional[int]		=  []
				UpperCamelCase_:							Tuple		=  fairseq_model.state_dict()

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

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

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







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		str:
				UpperCamelCase_:							Union[str, Any]		=  access_by_string(lowerCamelCase   ,			lowerCamelCase  )
				UpperCamelCase_:							str		=  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.'''  )
				UpperCamelCase_:							Optional[int]		=  fsq_value
				logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''  )







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

				UpperCamelCase_:							Tuple		=  name.split("""."""  )[-1]
				if type_id == 0:
								UpperCamelCase_:							List[Any]		=  """conv"""
				elif type_id == 2:
								UpperCamelCase_:							Tuple		=  """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  )     ->		Optional[Any]:
				UpperCamelCase_:							Optional[int]		=  full_name.split("""pos_conv."""  )[-1]
				UpperCamelCase_:							Any		=  name.split("""."""  )
				UpperCamelCase_:							List[str]		=  int(items[0]  )
				UpperCamelCase_:							str		=  int(items[1]  )

				UpperCamelCase_:							List[str]		=  name.split("""."""  )[-1]
				if type_id != 0:
								unused_weights.append(lowerCamelCase  )
								return
				else:
								UpperCamelCase_:							Tuple		=  """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  )     ->		Dict:
				if config_path is not None:
								UpperCamelCase_:							str		=  DataaVecAudioConfig.from_pretrained(lowerCamelCase  )
				else:
								UpperCamelCase_:							Tuple		=  DataaVecAudioConfig()

				if not is_finetuned:
								# Modify final_proj layer name
								UpperCamelCase_:							Union[str, Any]		=  DataaVecAudioModel(lowerCamelCase  )
								UpperCamelCase_:							Optional[int]		=  os.path.dirname(lowerCamelCase  )

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

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

				UpperCamelCase_:							List[Any]		=  load_dataavec(lowerCamelCase  )

				recursively_load_weights(lowerCamelCase   ,			lowerCamelCase   ,			not is_finetuned  )

				UpperCamelCase_:							Tuple		=  WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""  )

				UpperCamelCase_:							Dict		=  load_dataset("""patrickvonplaten/librispeech_asr_dummy"""   ,			"""clean"""   ,			split="""validation"""  )
				UpperCamelCase_:							Optional[Any]		=  [x["""array"""] for x in ds[:4]["""audio"""]]

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

				UpperCamelCase_:							str		=  inputs.input_values
				UpperCamelCase_:							Dict		=  inputs.attention_mask
				#    input_values = inputs.input_values[:, :-1]
				#    attention_mask = inputs.attention_mask[:, :-1]

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

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

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

				print(our_output.shape   ,			their_output.shape  )
				UpperCamelCase_:							Optional[Any]		=  torch.max(torch.abs(our_output - their_output  )  ).item()
				print(F'''max_absolute_diff = {max_absolute_diff}'''  )  # ~ 1e-7
				UpperCamelCase_:							List[str]		=  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_  :       str					      =     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_  :       Dict					      =     parser.parse_args()
							convert_wavaveca_checkpoint(
							    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
							)






