





import itertools
import os
import random
import tempfile
import unittest

import numpy as np
from datasets import load_dataset

from transformers import is_speech_available
from transformers.testing_utils import check_json_file_has_correct_format, require_torch, require_torchaudio
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
		from transformers import WhisperFeatureExtractor

if is_torch_available():
		import torch

__a      : Optional[int] 				=				random.Random()
def 		__magic_name__				( lowercase_						,		lowercase_=1.0						,		lowercase_=None						,		lowercase_=None				)		->					str:


						'''simple docstring'''


						if rng is None:
												UpperCamelCase  							= global_rng

						UpperCamelCase  							= []
						for batch_idx in range(shape[0]				):
												values.append([]				)
												for _ in range(shape[1]				):
																		values[-1].append(rng.random() * scale				)

						return values


@require_torch
@require_torchaudio
class 							__UpperCAmelCase							(						unittest.TestCase   ):



						"""simple docstring"""







						def __init__(   self	,		SCREAMING_SNAKE_CASE	,		SCREAMING_SNAKE_CASE=7	,		SCREAMING_SNAKE_CASE=400	,		SCREAMING_SNAKE_CASE=2000	,		SCREAMING_SNAKE_CASE=10	,		SCREAMING_SNAKE_CASE=160	,		SCREAMING_SNAKE_CASE=8	,		SCREAMING_SNAKE_CASE=0.0	,		SCREAMING_SNAKE_CASE=4000	,		SCREAMING_SNAKE_CASE=False	,		SCREAMING_SNAKE_CASE=True	,		) ->      List[str]:





												"""simple docstring"""
												UpperCamelCase  							= parent
												UpperCamelCase  							= batch_size
												UpperCamelCase  							= min_seq_length
												UpperCamelCase  							= max_seq_length
												UpperCamelCase  							= (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
												UpperCamelCase  							= padding_value
												UpperCamelCase  							= sampling_rate
												UpperCamelCase  							= return_attention_mask
												UpperCamelCase  							= do_normalize
												UpperCamelCase  							= feature_size
												UpperCamelCase  							= chunk_length
												UpperCamelCase  							= hop_length




						def 		__lowerCAmelCase			(   self				) ->      Optional[int]:





												"""simple docstring"""
												return {
												    "feature_size": self.feature_size,
												    "hop_length": self.hop_length,
												    "chunk_length": self.chunk_length,
												    "padding_value": self.padding_value,
												    "sampling_rate": self.sampling_rate,
												    "return_attention_mask": self.return_attention_mask,
												    "do_normalize": self.do_normalize,
												}




						def 		__lowerCAmelCase			(   self	,		SCREAMING_SNAKE_CASE=False	,		SCREAMING_SNAKE_CASE=False				) ->      List[str]:





												"""simple docstring"""
												def _flatten(SCREAMING_SNAKE_CASE				):
																		return list(itertools.chain(*A_				)				)

												if equal_length:
																		UpperCamelCase  							= [floats_list((self.max_seq_length, self.feature_size)				) for _ in range(self.batch_size				)]
												else:
																		# make sure that inputs increase in size
																		UpperCamelCase  							= [
																		    floats_list((x, self.feature_size)				)
																		    for x in range(self.min_seq_length	,		self.max_seq_length	,		self.seq_length_diff				)
																		]
												if numpify:
																		UpperCamelCase  							= [np.asarray(A_				) for x in speech_inputs]
												return speech_inputs


@require_torch
@require_torchaudio
class 							__UpperCAmelCase							(						snake_case__    ,  unittest.TestCase   ):



						"""simple docstring"""







						lowercase							  =    WhisperFeatureExtractor if is_speech_available() else None




						def 		__lowerCAmelCase			(   self				) ->      List[Any]:





												"""simple docstring"""
												UpperCamelCase  							= WhisperFeatureExtractionTester(self				)




						def 		__lowerCAmelCase			(   self				) ->      List[Any]:





												"""simple docstring"""
												UpperCamelCase  							= self.feature_extraction_class(**self.feat_extract_dict				)

												with tempfile.TemporaryDirectory() as tmpdirname:
																		UpperCamelCase  							= feat_extract_first.save_pretrained(A_				)[0]
																		check_json_file_has_correct_format(A_				)
																		UpperCamelCase  							= self.feature_extraction_class.from_pretrained(A_				)

												UpperCamelCase  							= feat_extract_first.to_dict()
												UpperCamelCase  							= feat_extract_second.to_dict()
												UpperCamelCase  							= feat_extract_first.mel_filters
												UpperCamelCase  							= feat_extract_second.mel_filters
												self.assertTrue(np.allclose(A_	,		A_				)				)
												self.assertEqual(A_	,		A_				)




						def 		__lowerCAmelCase			(   self				) ->      Optional[int]:





												"""simple docstring"""
												UpperCamelCase  							= self.feature_extraction_class(**self.feat_extract_dict				)

												with tempfile.TemporaryDirectory() as tmpdirname:
																		UpperCamelCase  							= os.path.join(A_	,		"feat_extract.json"				)
																		feat_extract_first.to_json_file(A_				)
																		UpperCamelCase  							= self.feature_extraction_class.from_json_file(A_				)

												UpperCamelCase  							= feat_extract_first.to_dict()
												UpperCamelCase  							= feat_extract_second.to_dict()
												UpperCamelCase  							= feat_extract_first.mel_filters
												UpperCamelCase  							= feat_extract_second.mel_filters
												self.assertTrue(np.allclose(A_	,		A_				)				)
												self.assertEqual(A_	,		A_				)




						def 		__lowerCAmelCase			(   self				) ->      List[str]:





												"""simple docstring"""
												# Tests that all call wrap to encode_plus and batch_encode_plus
												UpperCamelCase  							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
												# create three inputs of length 800, 1000, and 1200
												UpperCamelCase  							= [floats_list((1, x)				)[0] for x in range(800	,		1400	,		200				)]
												UpperCamelCase  							= [np.asarray(A_				) for speech_input in speech_inputs]

												# Test feature size
												UpperCamelCase  							= feature_extractor(A_	,		padding="max_length"	,		return_tensors="np"				).input_features
												self.assertTrue(input_features.ndim == 3				)
												self.assertTrue(input_features.shape[-1] == feature_extractor.nb_max_frames				)
												self.assertTrue(input_features.shape[-2] == feature_extractor.feature_size				)

												# Test not batched input
												UpperCamelCase  							= feature_extractor(speech_inputs[0]	,		return_tensors="np"				).input_features
												UpperCamelCase  							= feature_extractor(np_speech_inputs[0]	,		return_tensors="np"				).input_features
												self.assertTrue(np.allclose(A_	,		A_	,		atol=1e-3				)				)

												# Test batched
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												for enc_seq_a, enc_seq_a in zip(A_	,		A_				):
																		self.assertTrue(np.allclose(A_	,		A_	,		atol=1e-3				)				)

												# Test 2-D numpy arrays are batched.
												UpperCamelCase  							= [floats_list((1, x)				)[0] for x in (800, 800, 800)]
												UpperCamelCase  							= np.asarray(A_				)
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												for enc_seq_a, enc_seq_a in zip(A_	,		A_				):
																		self.assertTrue(np.allclose(A_	,		A_	,		atol=1e-3				)				)

												# Test truncation required
												UpperCamelCase  							= [floats_list((1, x)				)[0] for x in range(200	,		(feature_extractor.n_samples + 500)	,		200				)]
												UpperCamelCase  							= [np.asarray(A_				) for speech_input in speech_inputs]

												UpperCamelCase  							= [x[: feature_extractor.n_samples] for x in speech_inputs]
												UpperCamelCase  							= [np.asarray(A_				) for speech_input in speech_inputs_truncated]

												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="np"				).input_features
												for enc_seq_a, enc_seq_a in zip(A_	,		A_				):
																		self.assertTrue(np.allclose(A_	,		A_	,		atol=1e-3				)				)




						def 		__lowerCAmelCase			(   self				) ->      List[Any]:





												"""simple docstring"""
												import torch

												UpperCamelCase  							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
												UpperCamelCase  							= np.random.rand(100	,		32				).astype(np.floataa				)
												UpperCamelCase  							= np_speech_inputs.tolist()

												for inputs in [py_speech_inputs, np_speech_inputs]:
																		UpperCamelCase  							= feature_extractor.pad([{"input_features": inputs}]	,		return_tensors="np"				)
																		self.assertTrue(np_processed.input_features.dtype == np.floataa				)
																		UpperCamelCase  							= feature_extractor.pad([{"input_features": inputs}]	,		return_tensors="pt"				)
																		self.assertTrue(pt_processed.input_features.dtype == torch.floataa				)




						def 		__lowerCAmelCase			(   self	,		SCREAMING_SNAKE_CASE				) ->      Any:





												"""simple docstring"""
												UpperCamelCase  							= load_dataset("hf-internal-testing/librispeech_asr_dummy"	,		"clean"	,		split="validation"				)
												# automatic decoding with librispeech
												UpperCamelCase  							= ds.sort("id"				).select(range(A_				)				)[:num_samples]["audio"]

												return [x["array"] for x in speech_samples]




						def 		__lowerCAmelCase			(   self				) ->      List[str]:





												"""simple docstring"""
												# fmt: off
												UpperCamelCase  							= torch.tensor(
												    [
												        0.1_193, -0.0_946, -0.1_098, -0.0_196, 0.0_225, -0.0_690, -0.1_736, 0.0_951,
												        0.0_971, -0.0_817, -0.0_702, 0.0_162, 0.0_260, 0.0_017, -0.0_192, -0.1_678,
												        0.0_709, -0.1_867, -0.0_655, -0.0_274, -0.0_234, -0.1_884, -0.0_516, -0.0_554,
												        -0.0_274, -0.1_425, -0.1_423, 0.0_837, 0.0_377, -0.0_854
												    ]				)
												# fmt: on

												UpperCamelCase  							= self._load_datasamples(1				)
												UpperCamelCase  							= WhisperFeatureExtractor()
												UpperCamelCase  							= feature_extractor(A_	,		return_tensors="pt"				).input_features
												self.assertEqual(input_features.shape	,		(1, 80, 3000)				)
												self.assertTrue(torch.allclose(input_features[0, 0, :30]	,		A_	,		atol=1e-4				)				)




						def 		__lowerCAmelCase			(   self				) ->      Optional[Any]:





												"""simple docstring"""
												UpperCamelCase  							= self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()				)
												UpperCamelCase  							= self._load_datasamples(1				)[0]
												UpperCamelCase  							= ((audio - audio.min()) / (audio.max() - audio.min())) * 65535  # Rescale to [0, 65535] to show issue
												UpperCamelCase  							= feat_extract.zero_mean_unit_var_norm([audio]	,		attention_mask=A_				)[0]

												self.assertTrue(np.all(np.mean(A_				) < 1e-3				)				)
												self.assertTrue(np.all(np.abs(np.var(A_				) - 1				) < 1e-3				)				)











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 (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


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

UpperCamelCase_												=	{
    'post_extract_proj': 'feature_projection',
    'encoder.pos_conv.0': 'encoder.pos_conv_embed.conv',
    '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.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			(    A				,							A				,							A				,							A				,							A     ):




   '''simple docstring'''

   for attribute in key.split("."     ):
      UpperCAmelCase__						=getattr(A				,							A     )

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

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






def 				_UpperCAmelCase			(    A				,							A				,							A     ):




   '''simple docstring'''

   UpperCAmelCase__						=[]
   UpperCAmelCase__						=fairseq_model.state_dict()

   UpperCAmelCase__						=hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

   for name, value in fairseq_dict.items():
      UpperCAmelCase__						=False
      if "conv_layers" in name:
         load_conv_layer(
             A				,							A				,							A				,							A				,							hf_model.config.feat_extract_norm == "group"				,							)
         UpperCAmelCase__						=True
      else:
         for key, mapped_key in MAPPING.items():
            UpperCAmelCase__						="sew." + 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]:
               UpperCAmelCase__						=True
               if "*" in mapped_key:
                  UpperCAmelCase__						=name.split(A     )[0].split("."     )[-2]
                  UpperCAmelCase__						=mapped_key.replace("*"				,							A     )
               if "weight_g" in name:
                  UpperCAmelCase__						="weight_g"
               elif "weight_v" in name:
                  UpperCAmelCase__						="weight_v"
               elif "weight" in name:
                  UpperCAmelCase__						="weight"
               elif "bias" in name:
                  UpperCAmelCase__						="bias"
               else:
                  UpperCAmelCase__						=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 				_UpperCAmelCase			(    A				,							A				,							A				,							A				,							A     ):




   '''simple docstring'''

   UpperCAmelCase__						=full_name.split("conv_layers."     )[-1]
   UpperCAmelCase__						=name.split("."     )
   UpperCAmelCase__						=int(items[0]     )
   UpperCAmelCase__						=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."""
         )
         UpperCAmelCase__						=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."""
         )
         UpperCAmelCase__						=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."
         )
         UpperCAmelCase__						=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."""
         )
         UpperCAmelCase__						=value
         logger.info(F"""Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."""     )
   else:
      unused_weights.append(A     )






def 				_UpperCAmelCase			(    A				,							A     ):




   '''simple docstring'''

   UpperCAmelCase__						=SEWConfig()
   if is_finetuned:
      UpperCAmelCase__						=model.wav_encoder.wav_model.cfg
   else:
      UpperCAmelCase__						=model.cfg

   UpperCAmelCase__						=fs_config.conv_bias
   UpperCAmelCase__						=eval(fs_config.conv_feature_layers     )
   UpperCAmelCase__						=[x[0] for x in conv_layers]
   UpperCAmelCase__						=[x[1] for x in conv_layers]
   UpperCAmelCase__						=[x[2] for x in conv_layers]
   UpperCAmelCase__						="gelu"
   UpperCAmelCase__						="layer" if fs_config.extractor_mode == "layer_norm" else "group"
   UpperCAmelCase__						=0.0
   UpperCAmelCase__						=fs_config.activation_fn.name
   UpperCAmelCase__						=fs_config.encoder_embed_dim
   UpperCAmelCase__						=0.02
   UpperCAmelCase__						=fs_config.encoder_ffn_embed_dim
   UpperCAmelCase__						=1e-5
   UpperCAmelCase__						=fs_config.encoder_layerdrop
   UpperCAmelCase__						=fs_config.encoder_attention_heads
   UpperCAmelCase__						=fs_config.conv_pos_groups
   UpperCAmelCase__						=fs_config.conv_pos
   UpperCAmelCase__						=len(A     )
   UpperCAmelCase__						=fs_config.encoder_layers
   UpperCAmelCase__						=fs_config.squeeze_factor

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

   UpperCAmelCase__						="Wav2Vec2FeatureExtractor"
   UpperCAmelCase__						="Wav2Vec2CTCTokenizer"

   return config






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




   '''simple docstring'''

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

   if config_path is not None:
      UpperCAmelCase__						=SEWConfig.from_pretrained(A     )
   else:
      UpperCAmelCase__						=convert_config(model[0]				,							A     )
   UpperCAmelCase__						=model[0].eval()

   UpperCAmelCase__						=True if config.feat_extract_norm == "layer" else False
   UpperCAmelCase__						=WavaVecaFeatureExtractor(
       feature_size=1				,							sampling_rate=16000				,							padding_value=0				,							do_normalize=A				,							return_attention_mask=A				,							)

   if is_finetuned:
      if dict_path:
         UpperCAmelCase__						=Dictionary.load(A     )

         # important change bos & pad token id since CTC symbol is <pad> and
         # not <s> as in fairseq
         UpperCAmelCase__						=target_dict.pad_index
         UpperCAmelCase__						=target_dict.bos_index
         UpperCAmelCase__						=target_dict.pad_index
         UpperCAmelCase__						=target_dict.bos_index
         UpperCAmelCase__						=target_dict.eos_index
         UpperCAmelCase__						=len(target_dict.symbols     )
         UpperCAmelCase__						=os.path.join(A				,							"vocab.json"     )
         if not os.path.isdir(A     ):
            logger.error("--pytorch_dump_folder_path ({}) should be a directory".format(A     )     )
            return
         os.makedirs(A				,							exist_ok=A     )
         with open(A				,							"w"				,							encoding="utf-8"     ) as vocab_handle:
            json.dump(target_dict.indices				,							A     )
         UpperCAmelCase__						=WavaVecaCTCTokenizer(
             A				,							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=A				,							)
         UpperCAmelCase__						=WavaVecaProcessor(feature_extractor=A				,							tokenizer=A     )
         processor.save_pretrained(A     )

      UpperCAmelCase__						=SEWForCTC(A     )
   else:
      UpperCAmelCase__						=SEWModel(A     )
      feature_extractor.save_pretrained(A     )

   recursively_load_weights(A				,							A				,							A     )

   hf_model.save_pretrained(A     )


if __name__ == "__main__":
    UpperCamelCase_												=	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'
    )
    UpperCamelCase_												=	parser.parse_args()
    convert_sew_checkpoint(
        args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
    )



