






import shutil
import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import SpeechaTextTokenizer, is_speech_available
from transformers.models.speech_to_text.tokenization_speech_to_text import VOCAB_FILES_NAMES, save_json
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_torch, require_torchaudio
from transformers.utils import FEATURE_EXTRACTOR_NAME

from .test_feature_extraction_speech_to_text import floats_list


if is_speech_available():
   from transformers import SpeechaTextFeatureExtractor, SpeechaTextProcessor


lowerCamelCase					    =	get_tests_dir('''fixtures/test_sentencepiece.model''')

@require_torch
@require_torchaudio
@require_sentencepiece
class        snake_case_				(							unittest.TestCase   ):







   """simple docstring"""



   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[Any]						)	->		Dict:


       """simple docstring"""



       _lowercase     			=						tempfile.mkdtemp()

       _lowercase     			=						['<s>', '<pad>', '</s>', '<unk>', '▁This', '▁is', '▁a', '▁t', 'est']
       _lowercase     			=						dict(zip(snake_case__ ,							range(len(snake_case__						)						)						)						)
       _lowercase     			=						Path(self.tmpdirname						)
       save_json(snake_case__ ,							save_dir / VOCAB_FILES_NAMES['vocab_file']						)
       if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
           copyfile(snake_case__ ,							save_dir / VOCAB_FILES_NAMES['spm_file']						)

       _lowercase     			=						SpeechaTextTokenizer.from_pretrained(self.tmpdirname						)
       tokenizer.save_pretrained(self.tmpdirname						)

       _lowercase     			=						{
           'feature_size': 24,
           'num_mel_bins': 24,
           'padding_value': 0.0,
           'sampling_rate': 16_000,
           'return_attention_mask': False,
           'do_normalize': True,
       }
       save_json(snake_case__ ,							save_dir / FEATURE_EXTRACTOR_NAME						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Any ,							**snake_case__:						Optional[Any]						)	->		List[Any]:


       """simple docstring"""



       return SpeechaTextTokenizer.from_pretrained(self.tmpdirname ,							**snake_case__						)

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


       """simple docstring"""



       return SpeechaTextFeatureExtractor.from_pretrained(self.tmpdirname ,							**snake_case__						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Union[str, Any]						)	->		List[Any]:


       """simple docstring"""



       shutil.rmtree(self.tmpdirname						)

   def       SCREAMING_SNAKE_CASE__					(     self:						str						)	->		Dict:


       """simple docstring"""



       _lowercase     			=						self.get_tokenizer()
       _lowercase     			=						self.get_feature_extractor()

       _lowercase     			=						SpeechaTextProcessor(tokenizer=snake_case__ ,							feature_extractor=snake_case__						)

       processor.save_pretrained(self.tmpdirname						)
       _lowercase     			=						SpeechaTextProcessor.from_pretrained(self.tmpdirname						)

       self.assertEqual(processor.tokenizer.get_vocab() ,							tokenizer.get_vocab()						)
       self.assertIsInstance(processor.tokenizer ,							snake_case__						)

       self.assertEqual(processor.feature_extractor.to_json_string() ,							feature_extractor.to_json_string()						)
       self.assertIsInstance(processor.feature_extractor ,							snake_case__						)

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


       """simple docstring"""



       _lowercase     			=						SpeechaTextProcessor(
           tokenizer=self.get_tokenizer() ,							feature_extractor=self.get_feature_extractor()						)
       processor.save_pretrained(self.tmpdirname						)

       _lowercase     			=						self.get_tokenizer(bos_token='(BOS)' ,							eos_token='(EOS)'						)
       _lowercase     			=						self.get_feature_extractor(do_normalize=snake_case__ ,							padding_value=1.0						)

       _lowercase     			=						SpeechaTextProcessor.from_pretrained(
           self.tmpdirname ,							bos_token='(BOS)' ,							eos_token='(EOS)' ,							do_normalize=snake_case__ ,							padding_value=1.0						)

       self.assertEqual(processor.tokenizer.get_vocab() ,							tokenizer_add_kwargs.get_vocab()						)
       self.assertIsInstance(processor.tokenizer ,							snake_case__						)

       self.assertEqual(processor.feature_extractor.to_json_string() ,							feature_extractor_add_kwargs.to_json_string()						)
       self.assertIsInstance(processor.feature_extractor ,							snake_case__						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Union[str, Any]						)	->		Union[str, Any]:


       """simple docstring"""



       _lowercase     			=						self.get_feature_extractor()
       _lowercase     			=						self.get_tokenizer()

       _lowercase     			=						SpeechaTextProcessor(tokenizer=snake_case__ ,							feature_extractor=snake_case__						)

       _lowercase     			=						floats_list((3, 1_000)						)

       _lowercase     			=						feature_extractor(snake_case__ ,							return_tensors='np'						)
       _lowercase     			=						processor(snake_case__ ,							return_tensors='np'						)

       for key in input_feat_extract.keys():
           self.assertAlmostEqual(input_feat_extract[key].sum() ,							input_processor[key].sum() ,							delta=1E-2						)

   def       SCREAMING_SNAKE_CASE__					(     self:						List[Any]						)	->		int:


       """simple docstring"""



       _lowercase     			=						self.get_feature_extractor()
       _lowercase     			=						self.get_tokenizer()

       _lowercase     			=						SpeechaTextProcessor(tokenizer=snake_case__ ,							feature_extractor=snake_case__						)

       _lowercase     			=						'This is a test string'

       _lowercase     			=						processor(text=snake_case__						)

       _lowercase     			=						tokenizer(snake_case__						)

       for key in encoded_tok.keys():
           self.assertListEqual(encoded_tok[key] ,							encoded_processor[key]						)

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


       """simple docstring"""



       _lowercase     			=						self.get_feature_extractor()
       _lowercase     			=						self.get_tokenizer()

       _lowercase     			=						SpeechaTextProcessor(tokenizer=snake_case__ ,							feature_extractor=snake_case__						)

       _lowercase     			=						[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

       _lowercase     			=						processor.batch_decode(snake_case__						)
       _lowercase     			=						tokenizer.batch_decode(snake_case__						)

       self.assertListEqual(snake_case__ ,							snake_case__						)

   def       SCREAMING_SNAKE_CASE__					(     self:						Optional[int]						)	->		Tuple:


       """simple docstring"""



       _lowercase     			=						self.get_feature_extractor()
       _lowercase     			=						self.get_tokenizer()

       _lowercase     			=						SpeechaTextProcessor(tokenizer=snake_case__ ,							feature_extractor=snake_case__						)

       self.assertListEqual(
           processor.model_input_names ,							feature_extractor.model_input_names ,							msg='`processor` and `feature_extractor` model input names do not match' ,							)








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
   )
