





import math

from ...configuration_utils import PretrainedConfig
from ...utils import logging


_snake_case           =			logging.get_logger(__name__)

_snake_case           =			{
    """facebook/data2vec-base-960h""": """https://huggingface.co/facebook/data2vec-audio-base-960h/resolve/main/config.json""",
    # See all Data2VecAudio models at https://huggingface.co/models?filter=data2vec-audio
}





class  lowerCAmelCase							(       lowercase_ ):
  __lowerCamelCase		=      'data2vec-audio'







  def __init__(      self  :Dict   ,							_lowercase  :str=32   ,							_lowercase  :Optional[int]=7_68   ,							_lowercase  :List[str]=12   ,							_lowercase  :List[str]=12   ,							_lowercase  :List[str]=30_72   ,							_lowercase  :List[str]="gelu"   ,							_lowercase  :List[Any]=0.1   ,							_lowercase  :Tuple=0.1   ,							_lowercase  :Union[str, Any]=0.1   ,							_lowercase  :Optional[int]=0.0   ,							_lowercase  :List[str]=0.1   ,							_lowercase  :Tuple=0.1   ,							_lowercase  :Any=0.02   ,							_lowercase  :Optional[Any]=1e-5   ,							_lowercase  :List[Any]="gelu"   ,							_lowercase  :Optional[int]=(5_12, 5_12, 5_12, 5_12, 5_12, 5_12, 5_12)   ,							_lowercase  :Optional[Any]=(5, 2, 2, 2, 2, 2, 2)   ,							_lowercase  :List[str]=(10, 3, 3, 3, 3, 2, 2)   ,							_lowercase  :Any=False   ,							_lowercase  :Union[str, Any]=16   ,							_lowercase  :Union[str, Any]=19   ,							_lowercase  :List[str]=5   ,							_lowercase  :int=0.05   ,							_lowercase  :List[Any]=10   ,							_lowercase  :str=2   ,							_lowercase  :Union[str, Any]=0.0   ,							_lowercase  :Tuple=10   ,							_lowercase  :Tuple=0   ,							_lowercase  :Tuple="sum"   ,							_lowercase  :Optional[int]=False   ,							_lowercase  :str=False   ,							_lowercase  :int=2_56   ,							_lowercase  :int=(5_12, 5_12, 5_12, 5_12, 15_00)   ,							_lowercase  :List[Any]=(5, 3, 3, 1, 1)   ,							_lowercase  :int=(1, 2, 3, 1, 1)   ,							_lowercase  :Union[str, Any]=5_12   ,							_lowercase  :List[str]=0   ,							_lowercase  :str=1   ,							_lowercase  :Optional[int]=2   ,							_lowercase  :Dict=False   ,							_lowercase  :Any=3   ,							_lowercase  :Union[str, Any]=2   ,							_lowercase  :Tuple=3   ,							_lowercase  :int=None   ,							**_lowercase  :Tuple   ,							):






     '''simple docstring'''





     super().__init__(**_lowercase   ,							pad_token_id=_lowercase   ,							bos_token_id=_lowercase   ,							eos_token_id=_lowercase )
     lowercase__			       =  hidden_size
     lowercase__			       =  feat_extract_activation
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  conv_bias
     lowercase__			       =  num_conv_pos_embeddings
     lowercase__			       =  num_conv_pos_embedding_groups
     lowercase__			       =  conv_pos_kernel_size
     lowercase__			       =  len(self.conv_dim )
     lowercase__			       =  num_hidden_layers
     lowercase__			       =  intermediate_size
     lowercase__			       =  hidden_act
     lowercase__			       =  num_attention_heads
     lowercase__			       =  hidden_dropout
     lowercase__			       =  attention_dropout
     lowercase__			       =  activation_dropout
     lowercase__			       =  feat_proj_dropout
     lowercase__			       =  final_dropout
     lowercase__			       =  layerdrop
     lowercase__			       =  layer_norm_eps
     lowercase__			       =  initializer_range
     lowercase__			       =  vocab_size
     lowercase__			       =  use_weighted_layer_sum

     if (
         (len(self.conv_stride ) != self.num_feat_extract_layers)
         or (len(self.conv_kernel ) != self.num_feat_extract_layers)
         or (len(self.conv_dim ) != self.num_feat_extract_layers)
     ):
        raise ValueError(
            "Configuration for convolutional layers is incorrect. It is required that `len(config.conv_dim)` =="
            " `len(config.conv_stride)` == `len(config.conv_kernel)`, but is `len(config.conv_dim) ="
            f''' {len(self.conv_dim )}`, `len(config.conv_stride) = {len(self.conv_stride )}`,'''
            f''' `len(config.conv_kernel) = {len(self.conv_kernel )}`.''' )

     # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
     lowercase__			       =  mask_time_prob
     lowercase__			       =  mask_time_length
     lowercase__			       =  mask_time_min_masks
     lowercase__			       =  mask_feature_prob
     lowercase__			       =  mask_feature_length
     lowercase__			       =  mask_feature_min_masks

     # ctc loss
     lowercase__			       =  ctc_loss_reduction
     lowercase__			       =  ctc_zero_infinity

     # adapter
     lowercase__			       =  add_adapter
     lowercase__			       =  adapter_kernel_size
     lowercase__			       =  adapter_stride
     lowercase__			       =  num_adapter_layers
     lowercase__			       =  output_hidden_size or hidden_size

     # SequenceClassification-specific parameter. Feel free to ignore for other classes.
     lowercase__			       =  classifier_proj_size

     # XVector-specific parameters. Feel free to ignore for other classes.
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  list(_lowercase )
     lowercase__			       =  xvector_output_dim







  @property
  def   UpperCAmelCase					(      self  :List[Any] ):






     '''simple docstring'''





     return math.prod(self.conv_stride )













import argparse
import json

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

from transformers import ViTImageProcessor, ViTMSNConfig, ViTMSNModel
from transformers.image_utils import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD


torch.set_grad_enabled(False)


def _A       (   __magic_name__  , __magic_name__=False	):
   lowercase__			       =  []
   for i in range(config.num_hidden_layers	):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append((f'''module.blocks.{i}.norm1.weight''', f'''vit.encoder.layer.{i}.layernorm_before.weight''')	)
      rename_keys.append((f'''module.blocks.{i}.norm1.bias''', f'''vit.encoder.layer.{i}.layernorm_before.bias''')	)
      rename_keys.append(
          (f'''module.blocks.{i}.attn.proj.weight''', f'''vit.encoder.layer.{i}.attention.output.dense.weight''')	)
      rename_keys.append((f'''module.blocks.{i}.attn.proj.bias''', f'''vit.encoder.layer.{i}.attention.output.dense.bias''')	)
      rename_keys.append((f'''module.blocks.{i}.norm2.weight''', f'''vit.encoder.layer.{i}.layernorm_after.weight''')	)
      rename_keys.append((f'''module.blocks.{i}.norm2.bias''', f'''vit.encoder.layer.{i}.layernorm_after.bias''')	)
      rename_keys.append((f'''module.blocks.{i}.mlp.fc1.weight''', f'''vit.encoder.layer.{i}.intermediate.dense.weight''')	)
      rename_keys.append((f'''module.blocks.{i}.mlp.fc1.bias''', f'''vit.encoder.layer.{i}.intermediate.dense.bias''')	)
      rename_keys.append((f'''module.blocks.{i}.mlp.fc2.weight''', f'''vit.encoder.layer.{i}.output.dense.weight''')	)
      rename_keys.append((f'''module.blocks.{i}.mlp.fc2.bias''', f'''vit.encoder.layer.{i}.output.dense.bias''')	)

   # projection layer + position embeddings
   rename_keys.extend(
       [
           ("module.cls_token", "vit.embeddings.cls_token"),
           ("module.patch_embed.proj.weight", "vit.embeddings.patch_embeddings.projection.weight"),
           ("module.patch_embed.proj.bias", "vit.embeddings.patch_embeddings.projection.bias"),
           ("module.pos_embed", "vit.embeddings.position_embeddings"),
       ]	)

   if base_model:
      # layernorm + pooler
      rename_keys.extend(
          [
              ("module.norm.weight", "layernorm.weight"),
              ("module.norm.bias", "layernorm.bias"),
          ]	)

      # if just the base model, we should remove "vit" from all keys that start with "vit"
      lowercase__			       =  [(pair[0], pair[1][4:]) if pair[1].startswith("vit"	) else pair for pair in rename_keys]
   else:
      # layernorm + classification head
      rename_keys.extend(
          [
              ("norm.weight", "vit.layernorm.weight"),
              ("norm.bias", "vit.layernorm.bias"),
              ("head.weight", "classifier.weight"),
              ("head.bias", "classifier.bias"),
          ]	)

   return rename_keys


def _A       (   __magic_name__  , __magic_name__  , __magic_name__=False	):
   for i in range(config.num_hidden_layers	):
      if base_model:
         lowercase__			       =  ""
      else:
         lowercase__			       =  "vit."
      # read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
      lowercase__			       =  state_dict.pop(f'''module.blocks.{i}.attn.qkv.weight'''	)
      lowercase__			       =  state_dict.pop(f'''module.blocks.{i}.attn.qkv.bias'''	)
      # next, add query, keys and values (in that order) to the state dict
      lowercase__			       =  in_proj_weight[
          : config.hidden_size, :
      ]
      lowercase__			       =  in_proj_bias[: config.hidden_size]
      lowercase__			       =  in_proj_weight[
          config.hidden_size : config.hidden_size * 2, :
      ]
      lowercase__			       =  in_proj_bias[
          config.hidden_size : config.hidden_size * 2
      ]
      lowercase__			       =  in_proj_weight[
          -config.hidden_size :, :
      ]
      lowercase__			       =  in_proj_bias[-config.hidden_size :]


def _A       (   __magic_name__	):
   lowercase__			       =  ["head.weight", "head.bias"]
   for k in ignore_keys:
      state_dict.pop(__magic_name__  , __magic_name__	)


def _A       (   __magic_name__	):
   # projection head is used in the self-supervised pre-training in MSN,
   # for downstream task it's not needed.
   lowercase__			       =  [
       "module.fc.fc1.weight",
       "module.fc.fc1.bias",
       "module.fc.bn1.weight",
       "module.fc.bn1.bias",
       "module.fc.bn1.running_mean",
       "module.fc.bn1.running_var",
       "module.fc.bn1.num_batches_tracked",
       "module.fc.fc2.weight",
       "module.fc.fc2.bias",
       "module.fc.bn2.weight",
       "module.fc.bn2.bias",
       "module.fc.bn2.running_mean",
       "module.fc.bn2.running_var",
       "module.fc.bn2.num_batches_tracked",
       "module.fc.fc3.weight",
       "module.fc.fc3.bias",
   ]
   for k in ignore_keys:
      state_dict.pop(__magic_name__  , __magic_name__	)


def _A       (   __magic_name__  , __magic_name__  , __magic_name__	):
   lowercase__			       =  dct.pop(__magic_name__	)
   lowercase__			       =  val


def _A       (   __magic_name__  , __magic_name__	):
   lowercase__			       =  ViTMSNConfig()
   lowercase__			       =  1000

   lowercase__			       =  "datasets/huggingface/label-files"
   lowercase__			       =  "imagenet-1k-id2label.json"
   lowercase__			       =  json.load(open(hf_hub_download(__magic_name__  , __magic_name__	)  , "r"	)	)
   lowercase__			       =  {int(__magic_name__	): v for k, v in idalabel.items()}
   lowercase__			       =  idalabel
   lowercase__			       =  {v: k for k, v in idalabel.items()}

   if "s16" in checkpoint_url:
      lowercase__			       =  384
      lowercase__			       =  1536
      lowercase__			       =  6
   elif "l16" in checkpoint_url:
      lowercase__			       =  1024
      lowercase__			       =  4096
      lowercase__			       =  24
      lowercase__			       =  16
      lowercase__			       =  0.1
   elif "b4" in checkpoint_url:
      lowercase__			       =  4
   elif "l7" in checkpoint_url:
      lowercase__			       =  7
      lowercase__			       =  1024
      lowercase__			       =  4096
      lowercase__			       =  24
      lowercase__			       =  16
      lowercase__			       =  0.1

   lowercase__			       =  ViTMSNModel(__magic_name__	)

   lowercase__			       =  torch.hub.load_state_dict_from_url(__magic_name__  , map_location="cpu"	)["target_encoder"]

   lowercase__			       =  ViTImageProcessor(size=config.image_size	)

   remove_projection_head(__magic_name__	)
   lowercase__			       =  create_rename_keys(__magic_name__  , base_model=__magic_name__	)

   for src, dest in rename_keys:
      rename_key(__magic_name__  , __magic_name__  , __magic_name__	)
   read_in_q_k_v(__magic_name__  , __magic_name__  , base_model=__magic_name__	)

   model.load_state_dict(__magic_name__	)
   model.eval()

   lowercase__			       =  "http://images.cocodataset.org/val2017/000000039769.jpg"

   lowercase__			       =  Image.open(requests.get(__magic_name__  , stream=__magic_name__	).raw	)
   lowercase__			       =  ViTImageProcessor(
       size=config.image_size  , image_mean=__magic_name__  , image_std=__magic_name__	)
   lowercase__			       =  image_processor(images=__magic_name__  , return_tensors="pt"	)

   # forward pass
   torch.manual_seed(2	)
   lowercase__			       =  model(**__magic_name__	)
   lowercase__			       =  outputs.last_hidden_state

   # The following Colab Notebook was used to generate these outputs:
   # https://colab.research.google.com/gist/sayakpaul/3672419a04f5997827503fd84079bdd1/scratchpad.ipynb
   if "s16" in checkpoint_url:
      lowercase__			       =  torch.tensor([[-1.0_915, -1.4_876, -1.1_809]]	)
   elif "b16" in checkpoint_url:
      lowercase__			       =  torch.tensor([[14.2_889, -18.9_045, 11.7_281]]	)
   elif "l16" in checkpoint_url:
      lowercase__			       =  torch.tensor([[41.5_028, -22.8_681, 45.6_475]]	)
   elif "b4" in checkpoint_url:
      lowercase__			       =  torch.tensor([[-4.3_868, 5.2_932, -0.4_137]]	)
   else:
      lowercase__			       =  torch.tensor([[-0.1_792, -0.6_465, 2.4_263]]	)

   # verify logits
   assert torch.allclose(last_hidden_state[:, 0, :3]  , __magic_name__  , atol=1e-4	)

   print(f'''Saving model to {pytorch_dump_folder_path}'''	)
   model.save_pretrained(__magic_name__	)

   print(f'''Saving image processor to {pytorch_dump_folder_path}'''	)
   image_processor.save_pretrained(__magic_name__	)


if __name__ == "__main__":
  _snake_case           =			argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      """--checkpoint_url""",
      default="""https://dl.fbaipublicfiles.com/msn/vits16_800ep.pth.tar""",
      type=str,
      help="""URL of the checkpoint you'd like to convert.""",
  )
  parser.add_argument(
      """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
  )

  _snake_case           =			parser.parse_args()
  convert_vit_msn_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)






