






import timeit
from typing import Callable, Optional

from ..configuration_utils import PretrainedConfig
from ..models.auto.modeling_auto import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING
from ..utils import is_pyanvml_available, is_torch_available, logging
from .benchmark_utils import (
    Benchmark,
    Memory,
    MemorySummary,
    measure_peak_memory_cpu,
    start_memory_tracing,
    stop_memory_tracing,
)


if is_torch_available():
   import torch

   from .benchmark_args import PyTorchBenchmarkArguments


if is_pyanvml_available():
   import pyanvml.pyanvml as nvml


lowerCamelCase					    =	logging.get_logger(__name__)

class        snake_case_				(							lowercase_   ):







   """simple docstring"""



   __lowerCAmelCase  =    42
   __lowerCAmelCase  =    42
   __lowerCAmelCase  =    "PyTorch"

   @property
   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		List[str]:


       """simple docstring"""



       return torch.__version__

   def       SCREAMING_SNAKE_CASE__					(     self:						str ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		float:


       """simple docstring"""



       _lowercase     			=						self._prepare_inference_func(snake_case__ ,							snake_case__ ,							snake_case__						)
       return self._measure_speed(_inference						)

   def       SCREAMING_SNAKE_CASE__					(     self:						int ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		[Memory, Optional[MemorySummary]]:


       """simple docstring"""



       _lowercase     			=						self._prepare_inference_func(snake_case__ ,							snake_case__ ,							snake_case__						)
       return self._measure_memory(_inference						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Tuple ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		float:


       """simple docstring"""



       _lowercase     			=						self._prepare_train_func(snake_case__ ,							snake_case__ ,							snake_case__						)
       return self._measure_speed(_train						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Any ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		[Memory, Optional[MemorySummary]]:


       """simple docstring"""



       _lowercase     			=						self._prepare_train_func(snake_case__ ,							snake_case__ ,							snake_case__						)
       return self._measure_memory(_train						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Union[str, Any] ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		Callable[[], None]:


       """simple docstring"""



       _lowercase     			=						self.config_dict[model_name]

       if self.args.torchscript:
           _lowercase     			=						True

       _lowercase     			=						(
           hasattr(snake_case__ ,							'architectures'						)
           and isinstance(config.architectures ,							snake_case__						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
           try:
               _lowercase     			=						config.architectures[0]
               _lowercase     			=						__import__('transformers' ,							fromlist=[model_class]						)
               _lowercase     			=						getattr(snake_case__ ,							snake_case__						)
               _lowercase     			=						model_cls(snake_case__						)
           except ImportError:
               raise ImportError(
                   f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                   ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'						)
       else:
           _lowercase     			=						MODEL_MAPPING[config.__class__](snake_case__						)

       model.eval()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       _lowercase     			=						config.vocab_size if hasattr(snake_case__ ,							'vocab_size'						) else config.encoder.vocab_size
       _lowercase     			=						torch.randint(snake_case__ ,							(batch_size, sequence_length) ,							dtype=torch.long ,							device=self.args.device						)

       if self.args.fpaa:
           logger.info('Running training in Mixed Precision...'						)
           if not self.args.is_gpu:
               raise ValueError('Mixed precision is possible only for GPU.'						)
           # amp seems to have memory leaks so that memory usage
           # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
           model.half()

       if self.args.torchscript:
           with torch.no_grad():
               _lowercase     			=						torch.jit.trace(snake_case__ ,							snake_case__						)
       else:
           _lowercase     			=						model

       def encoder_decoder_forward():
           with torch.no_grad():
               _lowercase     			=						inference_model(snake_case__ ,							decoder_input_ids=snake_case__						)
           return outputs

       def encoder_forward():
           with torch.no_grad():
               _lowercase     			=						inference_model(snake_case__						)
           return outputs

       _lowercase     			=						encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
       return _forward

   def       SCREAMING_SNAKE_CASE__					(     self:						List[Any] ,							snake_case__:						str ,							snake_case__:						int ,							snake_case__:						int						)	->		Callable[[], None]:


       """simple docstring"""



       _lowercase     			=						self.config_dict[model_name]

       _lowercase     			=						(
           hasattr(snake_case__ ,							'architectures'						)
           and isinstance(config.architectures ,							snake_case__						)
           and len(config.architectures						) > 0
       )
       if not self.args.only_pretrain_model and has_model_class_in_config:
           try:
               _lowercase     			=						config.architectures[0]
               _lowercase     			=						__import__('transformers' ,							fromlist=[model_class]						)
               _lowercase     			=						getattr(snake_case__ ,							snake_case__						)
               _lowercase     			=						model_cls(snake_case__						)
           except ImportError:
               raise ImportError(
                   f'''{model_class} does not exist. If you just want to test the pretrained model, you might want to'''
                   ' set `--only_pretrain_model` or `args.only_pretrain_model=True`.'						)
       else:
           _lowercase     			=						MODEL_WITH_LM_HEAD_MAPPING[config.__class__](snake_case__						)

       if self.args.torchscript:
           raise NotImplementedError('Training for torchscript is currently not implemented'						)
       else:
           _lowercase     			=						model

       model.train()
       model.to(self.args.device						)

       # encoder-decoder has vocab size saved differently
       _lowercase     			=						config.vocab_size if hasattr(snake_case__ ,							'vocab_size'						) else config.encoder.vocab_size
       _lowercase     			=						torch.randint(snake_case__ ,							(batch_size, sequence_length) ,							dtype=torch.long ,							device=self.args.device						)

       if self.args.fpaa:
           logger.info('Running training in Mixed Precision...'						)
           if not self.args.is_gpu:
               raise ValueError('Mixed precision is possible only for GPU.'						)

           # amp seems to have memory leaks so that memory usage
           # is measured using .half() for now https://github.com/NVIDIA/apex/issues/439
           model.half()

       def compute_loss_and_backprob_encoder():
           _lowercase     			=						train_model(snake_case__ ,							labels=snake_case__						)[0]
           loss.backward()
           return loss

       def compute_loss_and_backprob_encoder_decoder():
           _lowercase     			=						train_model(snake_case__ ,							decoder_input_ids=snake_case__ ,							labels=snake_case__						)[0]
           loss.backward()
           return loss

       _lowercase     			=						(
           compute_loss_and_backprob_encoder_decoder
           if config.is_encoder_decoder
           else compute_loss_and_backprob_encoder
       )
       return _train

   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[int] ,							snake_case__:						Optional[int]						)	->		float:


       """simple docstring"""



       try:
           if self.args.is_tpu or self.args.torchscript:
               # run additional 10 times to stabilize compilation for tpu and torchscript
               logger.info('Do inference on TPU or torchscript. Running model 5 times to stabilize compilation'						)
               timeit.repeat(
                   snake_case__ ,							repeat=1 ,							number=5 ,							)

           # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
           _lowercase     			=						timeit.repeat(
               snake_case__ ,							repeat=self.args.repeat ,							number=10 ,							)

           if self.args.is_tpu and self.args.torch_xla_tpu_print_metrics:
               import torch_xla.debug.metrics as met

               self.print_fn(met.metrics_report()						)

           return min(snake_case__						) / 10.0
       except RuntimeError as e:
           self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)
           return "N/A"

   def       SCREAMING_SNAKE_CASE__					(     self:						Tuple ,							snake_case__:						Callable[[], None]						)	->		[Memory, MemorySummary]:


       """simple docstring"""



       try:
           if self.args.trace_memory_line_by_line:
               _lowercase     			=						start_memory_tracing('transformers'						)

           if self.args.is_tpu:
               # tpu
               raise NotImplementedError(
                   'Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with'
                   ' `--no-memory` or `args.memory=False`'						)
           elif self.args.is_gpu:
               if not is_pyanvml_available():
                   logger.warning(
                       'py3nvml not installed, we won\'t log GPU memory usage. '
                       'Install py3nvml (pip install py3nvml) to log information about GPU.'						)
                   _lowercase     			=						'N/A'
               else:
                   logger.info(
                       'Measuring total GPU usage on GPU device. Make sure to not have additional processes running'
                       ' on the same GPU.'						)
                   # init nvml
                   nvml.nvmlInit()
                   func()
                   _lowercase     			=						nvml.nvmlDeviceGetHandleByIndex(self.args.device_idx						)
                   _lowercase     			=						nvml.nvmlDeviceGetMemoryInfo(snake_case__						)
                   _lowercase     			=						meminfo.used
                   _lowercase     			=						Memory(snake_case__						)
                   # shutdown nvml
                   nvml.nvmlShutdown()
           else:
               # cpu
               _lowercase     			=						measure_peak_memory_cpu(snake_case__						)
               _lowercase     			=						Memory(snake_case__						) if isinstance(snake_case__ ,							snake_case__						) else memory_bytes

           if self.args.trace_memory_line_by_line:
               _lowercase     			=						stop_memory_tracing(snake_case__						)
           else:
               _lowercase     			=						None

           return memory, summary
       except RuntimeError as e:
           self.print_fn(f'''Doesn\'t fit on GPU. {e}'''						)
           return "N/A", None








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 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



    for attribute in key.split('.'   ):
        _lowercase     			=						getattr(_lowercase,   _lowercase   )

    if weight_type is not None:
        _lowercase     			=						getattr(_lowercase,   _lowercase   ).shape
    else:
        _lowercase     			=						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":
        _lowercase     			=						value
    elif weight_type == "weight_g":
        _lowercase     			=						value
    elif weight_type == "weight_v":
        _lowercase     			=						value
    elif weight_type == "bias":
        _lowercase     			=						value
    else:
        _lowercase     			=						value

    logger.info(f'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''   )
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



    _lowercase     			=						[]
    _lowercase     			=						fairseq_model.state_dict()

    if not is_headless:
        _lowercase     			=						hf_model.dataavec_audio.feature_extractor
        _lowercase     			=						hf_model.dataavec_audio.encoder.pos_conv_embed

    else:
        _lowercase     			=						hf_model.feature_extractor
        _lowercase     			=						hf_model.encoder.pos_conv_embed

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

    logger.warning(f'''Unused weights: {unused_weights}'''   )
def 							UpperCAmelCase   ( _lowercase,   _lowercase   ):


    '''simple docstring'''



    _lowercase     			=						path.split('.'   )
    return reduce(_lowercase,   _lowercase,   _lowercase   )
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



    _lowercase     			=						access_by_string(_lowercase,   _lowercase   )
    _lowercase     			=						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.'''   )
    _lowercase     			=						fsq_value
    logger.info(f'''{full_name} was correctly initialized from {hf_weight_path}.'''   )
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



    _lowercase     			=						full_name.split('conv_layers.'   )[-1]
    _lowercase     			=						name.split('.'   )
    _lowercase     			=						int(items[0]   )
    _lowercase     			=						int(items[1]   )

    _lowercase     			=						name.split('.'   )[-1]
    if type_id == 0:
        _lowercase     			=						'conv'
    elif type_id == 2:
        _lowercase     			=						'layer_norm'
    else:
        unused_weights.append(_lowercase   )
        return

    set_weights(_lowercase,   _lowercase,   _lowercase,   f'''conv_layers.{layer_id}.{layer_type}.{weight_type}'''   )
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase,   _lowercase   ):


    '''simple docstring'''



    _lowercase     			=						full_name.split('pos_conv.'   )[-1]
    _lowercase     			=						name.split('.'   )
    _lowercase     			=						int(items[0]   )
    _lowercase     			=						int(items[1]   )

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

    set_weights(_lowercase,   _lowercase,   _lowercase,   f'''layers.{layer_id}.{layer_type}.{weight_type}'''   )
@torch.no_grad()
def 							UpperCAmelCase   ( _lowercase,   _lowercase,   _lowercase=None,   _lowercase=None,   _lowercase=True   ):


    '''simple docstring'''



    if config_path is not None:
        _lowercase     			=						DataaVecAudioConfig.from_pretrained(_lowercase   )
    else:
        _lowercase     			=						DataaVecAudioConfig()

    if not is_finetuned:
        # Modify final_proj layer name
        _lowercase     			=						DataaVecAudioModel(_lowercase   )
        _lowercase     			=						os.path.dirname(_lowercase   )

        _lowercase     			=						torch.load(_lowercase   )
        _lowercase     			=						state_dict['model'].pop('final_proj.0.weight'   )
        _lowercase     			=						state_dict['model'].pop('final_proj.0.bias'   )
        _lowercase     			=						os.path.join(_lowercase,   'converted.pt'   )
        torch.save(_lowercase,   _lowercase   )
    else:
        _lowercase     			=						DataaVecAudioForCTC(_lowercase   )
        _lowercase     			=						checkpoint_path

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

    _lowercase     			=						load_dataavec(_lowercase   )

    recursively_load_weights(_lowercase,   _lowercase,   not is_finetuned   )

    _lowercase     			=						WavaVecaProcessor.from_pretrained('facebook/wav2vec2-large-lv60'   )

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

    _lowercase     			=						processor(_lowercase,   return_tensors='pt',   padding=_lowercase   )

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

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

        _lowercase     			=						torch.argmax(_lowercase,   dim=-1   )
        _lowercase     			=						processor.batch_decode(_lowercase   )

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

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

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


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
   )
