





"""simple docstring"""




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_					:						List[str]      = 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_					:						Union[str, Any]      = [
    'lm_head',
]





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  for attribute in key.split('.' ):
    A_  :	Dict    						=  getattr(_UpperCAmelCase			, _UpperCAmelCase )

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

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





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  []
  A_  :	Dict    						=  fairseq_model.state_dict()

  if not is_headless:
    A_  :	List[str]    						=  hf_model.dataavec_audio.feature_extractor
    A_  :	Optional[int]    						=  hf_model.dataavec_audio.encoder.pos_conv_embed

  else:
    A_  :	str    						=  hf_model.feature_extractor
    A_  :	Dict    						=  hf_model.encoder.pos_conv_embed

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

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





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Optional[int]    						=  path.split('.' )
  return reduce(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  access_by_string(_UpperCAmelCase			, _UpperCAmelCase )
  A_  :	List[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.""" )
  A_  :	int    						=  fsq_value
  logger.info(f"""{full_name} was correctly initialized from {hf_weight_path}.""" )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  full_name.split('conv_layers.' )[-1]
  A_  :	Union[str, Any]    						=  name.split('.' )
  A_  :	int    						=  int(items[0] )
  A_  :	Optional[Any]    						=  int(items[1] )

  A_  :	int    						=  name.split('.' )[-1]
  if type_id == 0:
    A_  :	int    						=  'conv'
  elif type_id == 2:
    A_  :	List[Any]    						=  'layer_norm'
  else:
    unused_weights.append(_UpperCAmelCase )
    return

  set_weights(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, f"""conv_layers.{layer_id}.{layer_type}.{weight_type}""" )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  full_name.split('pos_conv.' )[-1]
  A_  :	List[Any]    						=  name.split('.' )
  A_  :	Union[str, Any]    						=  int(items[0] )
  A_  :	Optional[int]    						=  int(items[1] )

  A_  :	Any    						=  name.split('.' )[-1]
  if type_id != 0:
    unused_weights.append(_UpperCAmelCase )
    return
  else:
    A_  :	Any    						=  'conv'

  set_weights(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, f"""layers.{layer_id}.{layer_type}.{weight_type}""" )





@torch.no_grad()
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase=None			, _UpperCAmelCase=None			, _UpperCAmelCase=True ):






  """simple docstring"""





  if config_path is not None:
    A_  :	Tuple    						=  DataaVecAudioConfig.from_pretrained(_UpperCAmelCase )
  else:
    A_  :	Optional[int]    						=  DataaVecAudioConfig()

  if not is_finetuned:
    # Modify final_proj layer name
    A_  :	int    						=  DataaVecAudioModel(_UpperCAmelCase )
    A_  :	Any    						=  os.path.dirname(_UpperCAmelCase )

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

  def load_dataavec(_UpperCAmelCase ):
    A_						,	A_						,	A_  :	Any    						=  fairseq.checkpoint_utils.load_model_ensemble_and_task([path] )
    return model[0].eval()

  A_  :	Any    						=  load_dataavec(_UpperCAmelCase )

  recursively_load_weights(_UpperCAmelCase			, _UpperCAmelCase			, not is_finetuned )

  A_  :	Optional[int]    						=  WavaVecaProcessor.from_pretrained('facebook/wav2vec2-large-lv60' )

  A_  :	Optional[int]    						=  load_dataset('patrickvonplaten/librispeech_asr_dummy'			, 'clean'			, split='validation' )
  A_  :	Dict    						=  [x['array'] for x in ds[:4]['audio']]

  A_  :	Optional[Any]    						=  processor(_UpperCAmelCase			, return_tensors='pt'			, padding=_UpperCAmelCase )

  A_  :	Tuple    						=  inputs.input_values
  A_  :	Tuple    						=  inputs.attention_mask
  #    input_values = inputs.input_values[:, :-1]
  #    attention_mask = inputs.attention_mask[:, :-1]

  hf_wavavec.eval()
  model.eval()
  if is_finetuned:
    A_  :	int    						=  model(source=_UpperCAmelCase			, padding_mask=(1 - attention_mask)			, mask=_UpperCAmelCase			, features_only=_UpperCAmelCase )[
        'encoder_out'
    ].transpose(0			, 1 )
    A_  :	Tuple    						=  hf_wavavec(_UpperCAmelCase			, attention_mask=_UpperCAmelCase )['logits']

    A_  :	int    						=  torch.argmax(_UpperCAmelCase			, dim=-1 )
    A_  :	Tuple    						=  processor.batch_decode(_UpperCAmelCase )

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

  print(our_output.shape			, their_output.shape )
  A_  :	List[str]    						=  torch.max(torch.abs(our_output - their_output ) ).item()
  print(f"""max_absolute_diff = {max_absolute_diff}""" )  # ~ 1e-7
  A_  :	str    						=  torch.allclose(_UpperCAmelCase			, _UpperCAmelCase			, 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(_UpperCAmelCase )

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


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






"""simple docstring"""




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if not inplace:
    A_  :	Optional[int]    						=  ma + ma
  else:
    ma += ma

  return ma





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  -1 * len(_UpperCAmelCase )
  A_  :	List[str]    						=  list(range(len(tensor.shape[:zero_index] ) ) )
  return tensor.permute(first_inds + [zero_index + i for i in inds] )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  return t.reshape(t.shape[:-no_dims] + (-1,) )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 1E-4 ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  mask.expand(*value.shape )
  return torch.sum(mask * value			, dim=_UpperCAmelCase ) / (eps + torch.sum(_UpperCAmelCase			, dim=_UpperCAmelCase ))





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase = 2.3_125			, _UpperCAmelCase = 21.6_875			, _UpperCAmelCase = 64 ):






  """simple docstring"""





  A_  :	List[str]    						=  torch.linspace(_UpperCAmelCase			, _UpperCAmelCase			, no_bins - 1			, device=pts.device )
  A_  :	Optional[Any]    						=  torch.sqrt(torch.sum((pts.unsqueeze(-2 ) - pts.unsqueeze(-3 )) ** 2			, dim=-1 ) )
  return torch.bucketize(_UpperCAmelCase			, _UpperCAmelCase )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Any    						=  dicts[0]
  A_  :	int    						=  {}
  for k, v in first.items():
    A_  :	Dict    						=  [d[k] for d in dicts]
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	List[Any]    						=  dict_multimap(_UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Union[str, Any]    						=  fn(_UpperCAmelCase )

  return new_dict





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  v_bins.view(((1,) * len(x.shape )) + (len(_UpperCAmelCase ),) )
  A_  :	Any    						=  x[..., None] - reshaped_bins
  A_  :	List[str]    						=  torch.argmin(torch.abs(_UpperCAmelCase )			, dim=-1 )
  return nn.functional.one_hot(_UpperCAmelCase			, num_classes=len(_UpperCAmelCase ) ).float()





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 0			, _UpperCAmelCase = 0 ):






  """simple docstring"""





  A_  :	List[Union[slice, torch.Tensor]]    						=  []
  for i, s in enumerate(data.shape[:no_batch_dims] ):
    A_  :	List[Any]    						=  torch.arange(_UpperCAmelCase )
    A_  :	Dict    						=  r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape ) - i - 1))) )
    ranges.append(_UpperCAmelCase )

  A_  :	List[Union[slice, torch.Tensor]]    						=  [slice(_UpperCAmelCase ) for _ in range(len(data.shape ) - no_batch_dims )]
  A_  :	Optional[Any]    						=  inds
  ranges.extend(_UpperCAmelCase )
  # Matt note: Editing this to get around the behaviour of using a list as an array index changing
  # in recent Numpy versions
  return data[tuple(_UpperCAmelCase )]


lowerCamelCase_					:						Union[str, Any]      = TypeVar('T')





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Dict[Any, Union[dict, list, tuple, Any]]    						=  {}
  for k, v in dic.items():
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	int    						=  dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Any    						=  tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )

  return new_dict





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return [tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree]
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return tuple(tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return fn(_UpperCAmelCase )
  else:
    print(type(_UpperCAmelCase ) )
    raise ValueError('Not supported' )





lowerCamelCase_					:						Optional[int]      = partial(tree_map, leaf_type=torch.Tensor)