






'''simple docstring'''





import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWDConfig,
    SEWDForCTC,
    SEWDModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
lowerCamelCase										=  logging.get_logger(__name__)

lowerCamelCase										=  {
    "post_extract_proj": "feature_projection",
    "encoder.pos_conv.0": "encoder.pos_conv_embed.conv",
    "attention.self.query_proj": "encoder.encoder.layer.*.attention.self.query_proj",
    "attention.self.key_proj": "encoder.encoder.layer.*.attention.self.key_proj",
    "attention.self.value_proj": "encoder.encoder.layer.*.attention.self.value_proj",
    "attention.output.dense": "encoder.encoder.layer.*.attention.output.dense",
    "attention.output.LayerNorm": "encoder.encoder.layer.*.attention.output.LayerNorm",
    "intermediate.dense": "encoder.encoder.layer.*.intermediate.dense",
    "output.dense": "encoder.encoder.layer.*.output.dense",
    "output.LayerNorm": "encoder.encoder.layer.*.output.LayerNorm",
    "encoder.encoder.rel_embeddings": "encoder.encoder.rel_embeddings",
    "encoder.encoder.LayerNorm": "encoder.encoder.LayerNorm",
    "encoder.upsample.0": "encoder.upsample.projection",
    "encoder.layer_norm": "encoder.layer_norm",
    "w2v_model.layer_norm": "layer_norm",
    "w2v_encoder.proj": "lm_head",
    "mask_emb": "masked_spec_embed",
}
def        UpperCAmelCase      (   _lowercase,      _lowercase,      _lowercase,      _lowercase,      _lowercase      ):


    '''simple docstring'''



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

    if weight_type is not None:
        _lowercase     =	getattr(UpperCamelCase_,      UpperCamelCase_      ).shape
    else:
        _lowercase     =	hf_pointer.shape

    assert hf_shape == value.shape, (
        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()

    _lowercase     =	hf_model.sew_d.feature_extractor if is_finetuned else hf_model.feature_extractor

    for name, value in fairseq_dict.items():
        _lowercase     =	False
        if "conv_layers" in name:
            load_conv_layer(
                UpperCamelCase_,      UpperCamelCase_,      UpperCamelCase_,      UpperCamelCase_,      hf_model.config.feat_extract_norm == 'group',      )
            _lowercase     =	True
        else:
            for key, mapped_key in MAPPING.items():
                _lowercase     =	'sew_d.' + mapped_key if (is_finetuned and mapped_key != 'lm_head') 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(UpperCamelCase_      )[0].split('.'      )[-2]
                        if not layer_index.isnumeric():
                            continue
                        _lowercase     =	mapped_key.replace('*',      UpperCamelCase_      )
                    if "weight_g" in name:
                        _lowercase     =	'weight_g'
                    elif "weight_v" in name:
                        _lowercase     =	'weight_v'
                    elif "weight" in name:
                        _lowercase     =	'weight'
                    elif "bias" in name:
                        _lowercase     =	'bias'
                    else:
                        _lowercase     =	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      (   _lowercase,      _lowercase,      _lowercase,      _lowercase,      _lowercase      ):


    '''simple docstring'''



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

    if type_id == 0:
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.bias.data.shape, (
                f'''{full_name} has size {value.shape}, but'''
                f''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''
            )
            _lowercase     =	value
            logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''      )
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].conv.weight.data.shape, (
                f'''{full_name} has size {value.shape}, but'''
                f''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''
            )
            _lowercase     =	value
            logger.info(f'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''      )
    elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
        if "bias" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape, (
                f'''{full_name} has size {value.shape}, but {feature_extractor[layer_id].layer_norm.bias.data.shape} was'''
                " found."
            )
            _lowercase     =	value
            logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''      )
        elif "weight" in name:
            assert value.shape == feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape, (
                f'''{full_name} has size {value.shape}, but'''
                f''' {feature_extractor[layer_id].layer_norm.weight.data.shape} was found.'''
            )
            _lowercase     =	value
            logger.info(f'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''      )
    else:
        unused_weights.append(UpperCamelCase_      )
def        UpperCAmelCase      (   _lowercase,      _lowercase      ):


    '''simple docstring'''



    _lowercase     =	SEWDConfig()
    if is_finetuned:
        _lowercase     =	model.wav_encoder.wav_model.cfg
    else:
        _lowercase     =	model.cfg

    _lowercase     =	fs_config.conv_bias
    _lowercase     =	eval(fs_config.conv_feature_layers      )
    _lowercase     =	[x[0] for x in conv_layers]
    _lowercase     =	[x[1] for x in conv_layers]
    _lowercase     =	[x[2] for x in conv_layers]
    _lowercase     =	'gelu'
    _lowercase     =	'layer' if fs_config.extractor_mode == 'layer_norm' else 'group'
    _lowercase     =	0.0
    _lowercase     =	fs_config.activation_fn.name
    _lowercase     =	fs_config.encoder_embed_dim
    _lowercase     =	0.02
    _lowercase     =	fs_config.encoder_ffn_embed_dim
    _lowercase     =	1e-5
    _lowercase     =	fs_config.encoder_layerdrop
    _lowercase     =	fs_config.encoder_attention_heads
    _lowercase     =	fs_config.conv_pos_groups
    _lowercase     =	fs_config.conv_pos
    _lowercase     =	len(UpperCamelCase_      )
    _lowercase     =	fs_config.encoder_layers
    _lowercase     =	fs_config.squeeze_factor
    # DeBERTa-specific parameters:
    _lowercase     =	fs_config.max_position_embeddings
    _lowercase     =	fs_config.position_buckets
    _lowercase     =	fs_config.share_att_key
    _lowercase     =	fs_config.relative_attention
    _lowercase     =	fs_config.position_biased_input
    _lowercase     =	tuple(fs_config.pos_att_type.split('|'      )      )
    _lowercase     =	fs_config.norm_rel_ebd

    # take care of any params that are overridden by the Wav2VecCtc model
    if is_finetuned:
        _lowercase     =	model.cfg
        _lowercase     =	fs_config.final_dropout
        _lowercase     =	fs_config.layerdrop
    _lowercase     =	fs_config.activation_dropout
    _lowercase     =	fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
    _lowercase     =	fs_config.attention_dropout
    _lowercase     =	fs_config.dropout_input
    _lowercase     =	fs_config.dropout
    _lowercase     =	fs_config.mask_channel_length
    _lowercase     =	fs_config.mask_channel_prob
    _lowercase     =	fs_config.mask_length
    _lowercase     =	fs_config.mask_prob

    _lowercase     =	'Wav2Vec2FeatureExtractor'
    _lowercase     =	'Wav2Vec2CTCTokenizer'

    return config
@torch.no_grad()
def        UpperCAmelCase      (   _lowercase,      _lowercase,      _lowercase=None,      _lowercase=None,      _lowercase=True      ):


    '''simple docstring'''



    if is_finetuned:
        _lowercase ,    _lowercase ,    _lowercase     =	fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path],      arg_overrides={'data': '/'.join(dict_path.split('/'      )[:-1]      )}      )
    else:
        _lowercase ,    _lowercase ,    _lowercase     =	fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]      )

    if config_path is not None:
        _lowercase     =	SEWDConfig.from_pretrained(UpperCamelCase_      )
    else:
        _lowercase     =	convert_config(model[0],      UpperCamelCase_      )
    _lowercase     =	model[0].eval()

    _lowercase     =	True if config.feat_extract_norm == 'layer' else False
    _lowercase     =	WavaVecaFeatureExtractor(
        feature_size=1,      sampling_rate=16000,      padding_value=0,      do_normalize=UpperCamelCase_,      return_attention_mask=UpperCamelCase_,      )

    if is_finetuned:
        if dict_path:
            _lowercase     =	Dictionary.load(UpperCamelCase_      )

            # important change bos & pad token id since CTC symbol is <pad> and
            # not <s> as in fairseq
            _lowercase     =	target_dict.pad_index
            _lowercase     =	target_dict.bos_index
            _lowercase     =	target_dict.pad_index
            _lowercase     =	target_dict.bos_index
            _lowercase     =	target_dict.eos_index
            _lowercase     =	len(target_dict.symbols      )
            _lowercase     =	os.path.join(UpperCamelCase_,      'vocab.json'      )
            if not os.path.isdir(UpperCamelCase_      ):
                logger.error('--pytorch_dump_folder_path ({}) should be a directory'.format(UpperCamelCase_      )      )
                return
            os.makedirs(UpperCamelCase_,      exist_ok=UpperCamelCase_      )
            with open(UpperCamelCase_,      'w',      encoding='utf-8'      ) as vocab_handle:
                json.dump(target_dict.indices,      UpperCamelCase_      )
            _lowercase     =	WavaVecaCTCTokenizer(
                UpperCamelCase_,      unk_token=target_dict.unk_word,      pad_token=target_dict.pad_word,      bos_token=target_dict.bos_word,      eos_token=target_dict.eos_word,      word_delimiter_token='|',      do_lower_case=UpperCamelCase_,      )
            _lowercase     =	WavaVecaProcessor(feature_extractor=UpperCamelCase_,      tokenizer=UpperCamelCase_      )
            processor.save_pretrained(UpperCamelCase_      )

        _lowercase     =	SEWDForCTC(UpperCamelCase_      )
    else:
        _lowercase     =	SEWDModel(UpperCamelCase_      )
        feature_extractor.save_pretrained(UpperCamelCase_      )

    recursively_load_weights(UpperCamelCase_,      UpperCamelCase_,      UpperCamelCase_      )

    hf_model.save_pretrained(UpperCamelCase_      )


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(
       '''--is_finetuned''', action='''store_true''', help='''Whether the model to convert is a fine-tuned model or not'''
   )
   lowerCamelCase										=  parser.parse_args()
   convert_sew_checkpoint(
       args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
   )








import json
import os
import shutil
import tempfile
import unittest

from transformers import is_speech_available, is_torch_available
from transformers.models.speechta import SpeechTaTokenizer
from transformers.testing_utils import get_tests_dir, require_torch
from transformers.utils import FEATURE_EXTRACTOR_NAME


if is_speech_available() and is_torch_available():
   from transformers import SpeechTaFeatureExtractor, SpeechTaProcessor

   from .test_feature_extraction_speechta import floats_list


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

@require_torch
class 					snake_case_							(       unittest.TestCase      ):







   """simple docstring"""



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


       """simple docstring"""



       _lowercase     =	tempfile.mkdtemp()

       _lowercase     =	SpeechTaTokenizer(snake_case__	)
       tokenizer.save_pretrained(self.tmpdirname	)

       _lowercase     =	{
           'feature_size': 1,
           'padding_value': 0.0,
           'sampling_rate': 16_000,
           'do_normalize': False,
           'num_mel_bins': 80,
           'hop_length': 16,
           'win_length': 64,
           'win_function': 'hann_window',
           'fmin': 80,
           'fmax': 7_600,
           'mel_floor': 1E-10,
           'reduction_factor': 2,
           'return_attention_mask': True,
       }

       _lowercase     =	os.path.join(self.tmpdirname   ,       snake_case__	)
       with open(self.feature_extraction_file   ,       'w'   ,       encoding='utf-8'	) as fp:
           fp.write(json.dumps(snake_case__	) + '\n'	)

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


       """simple docstring"""



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

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


       """simple docstring"""



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

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


       """simple docstring"""



       shutil.rmtree(self.tmpdirname	)

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


       """simple docstring"""



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

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

       processor.save_pretrained(self.tmpdirname	)
       _lowercase     =	SpeechTaProcessor.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:	int	)  ->		Any:


       """simple docstring"""



       _lowercase     =	SpeechTaProcessor(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     =	SpeechTaProcessor.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]	)  ->		int:


       """simple docstring"""



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

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

       _lowercase     =	floats_list((3, 1_000)	)

       _lowercase     =	feature_extractor(audio=snake_case__   ,       return_tensors='np'	)
       _lowercase     =	processor(audio=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:	Tuple	)  ->		List[str]:


       """simple docstring"""



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

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

       _lowercase     =	floats_list((3, 1_000)	)

       _lowercase     =	feature_extractor(audio_target=snake_case__   ,       return_tensors='np'	)
       _lowercase     =	processor(audio_target=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:	Dict	)  ->		Dict:


       """simple docstring"""



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

       _lowercase     =	SpeechTaProcessor(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:	Optional[int]	)  ->		int:


       """simple docstring"""



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

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

       _lowercase     =	'This is a test string'

       _lowercase     =	processor(text_target=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	)  ->		int:


       """simple docstring"""



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

       _lowercase     =	SpeechTaProcessor(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[Any]	)  ->		List[Any]:


       """simple docstring"""



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

       _lowercase     =	SpeechTaProcessor(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'   ,       )
