






"""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()
_UpperCamelCase							:       Tuple       					=		logging.get_logger(__name__)

_UpperCamelCase							:       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',
}
_UpperCamelCase							:       List[str]       					=		[
    'lm_head',
]




def      snake_case    (A_ :Optional[int]      ,     A_ :List[Any]      ,     A_ :List[Any]      ,     A_ :Dict      ,     A_ :Any ):







  '''simple docstring'''

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

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

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




def      snake_case    (A_ :str      ,     A_ :Optional[int]      ,     A_ :Tuple ):







  '''simple docstring'''

  a	:     Tuple							=  []
  a	:     str							=  fairseq_model.state_dict()

  if not is_headless:
    a	:     int							=  hf_model.dataavec_audio.feature_extractor
    a	:     Union[str, Any]							=  hf_model.dataavec_audio.encoder.pos_conv_embed

  else:
    a	:     Optional[Any]							=  hf_model.feature_extractor
    a	:     Tuple							=  hf_model.encoder.pos_conv_embed

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

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




def      snake_case    (A_ :List[str]      ,     A_ :Any ):







  '''simple docstring'''

  a	:     List[str]							=  path.split('.' )
  return reduce(A_      ,     A_      ,     A_ )




def      snake_case    (A_ :Any      ,     A_ :Any      ,     A_ :int      ,     A_ :Dict ):







  '''simple docstring'''

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




def      snake_case    (A_ :Optional[Any]      ,     A_ :List[str]      ,     A_ :str      ,     A_ :List[str] ):







  '''simple docstring'''

  a	:     Optional[int]							=  full_name.split('conv_layers.' )[-1]
  a	:     List[Any]							=  name.split('.' )
  a	:     Any							=  int(items[0] )
  a	:     int							=  int(items[1] )

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

  set_weights(A_      ,     A_      ,     A_      ,     f'''conv_layers.{layer_id}.{layer_type}.{weight_type}''' )




def      snake_case    (A_ :int      ,     A_ :Tuple      ,     A_ :Union[str, Any]      ,     A_ :str ):







  '''simple docstring'''

  a	:     List[Any]							=  full_name.split('pos_conv.' )[-1]
  a	:     List[str]							=  name.split('.' )
  a	:     str							=  int(items[0] )
  a	:     str							=  int(items[1] )

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

  set_weights(A_      ,     A_      ,     A_      ,     f'''layers.{layer_id}.{layer_type}.{weight_type}''' )




@torch.no_grad()
def      snake_case    (A_ :Optional[int]      ,     A_ :Dict      ,     A_ :Tuple=None      ,     A_ :int=None      ,     A_ :int=True ):







  '''simple docstring'''

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

  if not is_finetuned:
    # Modify final_proj layer name
    a	:     str							=  DataaVecAudioModel(A_ )
    a	:     Tuple							=  os.path.dirname(A_ )

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

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

  a	:     str							=  load_dataavec(A_ )

  recursively_load_weights(A_      ,     A_      ,     not is_finetuned )

  a	:     int							=  WavaVecaProcessor.from_pretrained('facebook/wav2vec2-large-lv60' )

  a	:     List[Any]							=  load_dataset('patrickvonplaten/librispeech_asr_dummy'      ,     'clean'      ,     split='validation' )
  a	:     Dict							=  [x['array'] for x in ds[:4]['audio']]

  a	:     Tuple							=  processor(A_      ,     return_tensors='pt'      ,     padding=A_ )

  a	:     Any							=  inputs.input_values
  a	:     str							=  inputs.attention_mask
  #    input_values = inputs.input_values[:, :-1]
  #    attention_mask = inputs.attention_mask[:, :-1]

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

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

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

  print(our_output.shape      ,     their_output.shape )
  a	:     int							=  torch.max(torch.abs(our_output - their_output ) ).item()
  print(f'''max_absolute_diff = {max_absolute_diff}''' )  # ~ 1e-7
  a	:     int							=  torch.allclose(A_      ,     A_      ,     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(A_ )

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


if __name__ == "__main__":
   _UpperCamelCase							:       Union[str, 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							:       int       					=		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 __future__ import annotations




def      snake_case    (A_ :list[int | float]      ,     A_ :int      ,     A_ :int ):







  '''simple docstring'''

  if len(A_ ) == 0:
    raise ValueError('find_min() arg is an empty sequence' )
  if (
      left >= len(A_ )
      or left < -len(A_ )
      or right >= len(A_ )
      or right < -len(A_ )
  ):
    raise IndexError('list index out of range' )
  if left == right:
    return nums[left]
  a	:     Optional[Any]							=  (left + right) >> 1  # the middle
  a	:     List[str]							=  find_min(A_      ,     A_      ,     A_ )  # find min in range[left, mid]
  a	:     Union[str, Any]							=  find_min(A_      ,     mid + 1      ,     A_ )  # find min in range[mid + 1, right]

  return left_min if left_min <= right_min else right_min


if __name__ == "__main__":
   import doctest

   doctest.testmod(verbose=True)




