






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


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

a_							=    {
    """post_extract_proj""": """feature_projection.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.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """feature_projection.layer_norm""",
    """quantizer.weight_proj""": """quantizer.weight_proj""",
    """quantizer.vars""": """quantizer.codevectors""",
    """project_q""": """project_q""",
    """final_proj""": """project_hid""",
    """w2v_encoder.proj""": """ctc_proj""",
    """mask_emb""": """masked_spec_embed""",
}
a_							=    [
    """ctc_proj""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]


def        _a  (     UpperCamelCase_					:				Any    ,   UpperCamelCase_					:				Optional[int]    ,   UpperCamelCase_					:				Any    ,   UpperCamelCase_					:				Dict    ,   UpperCamelCase_					:				List[str]    ,   UpperCamelCase_					:				Any   ) ->   Optional[Any]:
						"""simple docstring"""







						for attribute in key.split("."   ):
												if is_finetuned:
																		if attribute in ["quantizer", "project_q", "project_hid"]:
																								# those layers are only relevant for pretraining and should be dropped
																								return

																		if attribute == "ctc_proj":
																								# we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
																								lowerCAmelCase__       =						"lm_head"

												lowerCAmelCase__       =						getattr(UpperCamelCase_    ,   UpperCamelCase_   )

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

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


def        _a  (     UpperCamelCase_					:				str    ,   UpperCamelCase_					:				Optional[int]    ,   UpperCamelCase_					:				Any   ) ->   Any:
						"""simple docstring"""







						lowerCAmelCase__       =						[]
						lowerCAmelCase__       =						fairseq_model.state_dict()

						lowerCAmelCase__       =						hf_model.unispeech.feature_extractor

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

						logger.warning(F"Unused weights: {unused_weights}"   )


def        _a  (     UpperCamelCase_					:				int    ,   UpperCamelCase_					:				Any    ,   UpperCamelCase_					:				int    ,   UpperCamelCase_					:				str    ,   UpperCamelCase_					:				str   ) ->   Tuple:
						"""simple docstring"""







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


@torch.no_grad()
def        _a  (     UpperCamelCase_					:				Any    ,   UpperCamelCase_					:				List[str]    ,   UpperCamelCase_					:				int=None    ,   UpperCamelCase_					:				Union[str, Any]=None    ,   UpperCamelCase_					:				str=True   ) ->   int:
						"""simple docstring"""







						if config_path is not None:
												lowerCAmelCase__       =						UniSpeechConfig.from_pretrained(UpperCamelCase_   )
						else:
												lowerCAmelCase__       =						UniSpeechConfig()

						if is_finetuned:
												if dict_path:
																		lowerCAmelCase__       =						Dictionary.load_from_json(UpperCamelCase_   )

																		# important change bos & pad token id since CTC symbol is <pad> and
																		# not <s> as in fairseq
																		lowerCAmelCase__       =						target_dict.pad_index
																		lowerCAmelCase__       =						target_dict.bos_index
																		lowerCAmelCase__       =						target_dict.eos_index
																		lowerCAmelCase__       =						len(target_dict.symbols   )
																		lowerCAmelCase__       =						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_   )
																		lowerCAmelCase__       =						target_dict.indices

																		# fairseq has the <pad> and <s> switched
																		lowerCAmelCase__       =						42
																		lowerCAmelCase__       =						43
																		with open(UpperCamelCase_    ,   "w"    ,   encoding="utf-8"   ) as vocab_handle:
																								json.dump(UpperCamelCase_    ,   UpperCamelCase_   )
																		lowerCAmelCase__       =						WavaVecaPhonemeCTCTokenizer(
																		    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_    ,   )
																		lowerCAmelCase__       =						True if config.feat_extract_norm == "layer" else False
																		lowerCAmelCase__       =						WavaVecaFeatureExtractor(
																		    feature_size=1    ,   sampling_rate=16_000    ,   padding_value=0    ,   do_normalize=UpperCamelCase_    ,   return_attention_mask=UpperCamelCase_    ,   )
																		lowerCAmelCase__       =						WavaVecaProcessor(feature_extractor=UpperCamelCase_    ,   tokenizer=UpperCamelCase_   )
																		processor.save_pretrained(UpperCamelCase_   )

												lowerCAmelCase__       =						UniSpeechForCTC(UpperCamelCase_   )
						else:
												lowerCAmelCase__       =						UniSpeechForPreTraining(UpperCamelCase_   )

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

						lowerCAmelCase__       =						model[0].eval()

						recursively_load_weights(UpperCamelCase_    ,   UpperCamelCase_    ,   UpperCamelCase_   )

						hf_unispeech.save_pretrained(UpperCamelCase_   )


if __name__ == "__main__":
	a_							=    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'''
	)
	a_							=    parser.parse_args()
	convert_unispeech_checkpoint(
	    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
	)














import os
from dataclasses import dataclass, field
from io import BytesIO
from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Union

import numpy as np
import pyarrow as pa

from .. import config
from ..download.streaming_download_manager import xopen, xsplitext
from ..table import array_cast
from ..utils.py_utils import no_op_if_value_is_null, string_to_dict


if TYPE_CHECKING:
	from .features import FeatureType


a_,						a_,						a_							=    False, False, False


@dataclass
class      lowercase__     :
	a_         =None
	a_         =True
	a_         =True
	a_         =None
	# Automatically constructed
	a_         ="dict"
	a_         =pa.struct({"""bytes""": pa.binary(), """path""": pa.string()}	)
	a_         =field(default="""Audio""",					init=_UpperCAmelCase,					repr=_UpperCAmelCase	)

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







							'''simple docstring'''





							return self.pa_type

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					dict:







							'''simple docstring'''





							try:
													import soundfile as sf  # soundfile is a dependency of librosa, needed to decode audio files.
							except ImportError as err:
													raise ImportError("To support encoding audio data, please install 'soundfile'."       ) from err
							if isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
													return {"bytes": None, "path": value}
							elif isinstance(__UpperCAmelCase		,    __UpperCAmelCase       ):
													return {"bytes": value, "path": None}
							elif "array" in value:
													# convert the audio array to wav bytes
													lowerCAmelCase__       =						BytesIO()
													sf.write(__UpperCAmelCase		,    value["array"]		,    value["sampling_rate"]		,    format="wav"       )
													return {"bytes": buffer.getvalue(), "path": None}
							elif value.get("path"       ) is not None and os.path.isfile(value["path"]       ):
													# we set "bytes": None to not duplicate the data if they're already available locally
													if value["path"].endswith("pcm"       ):
																			# "PCM" only has raw audio bytes
																			if value.get("sampling_rate"       ) is None:
																									# At least, If you want to convert "PCM-byte" to "WAV-byte", you have to know sampling rate
																									raise KeyError("To use PCM files, please specify a 'sampling_rate' in Audio object"       )
																			if value.get("bytes"       ):
																									# If we already had PCM-byte, we don`t have to make "read file, make bytes" (just use it!)
																									lowerCAmelCase__       =						np.frombuffer(value["bytes"]		,    dtype=np.intaa       ).astype(np.floataa       ) / 32767
																			else:
																									lowerCAmelCase__       =						np.memmap(value["path"]		,    dtype="h"		,    mode="r"       ).astype(np.floataa       ) / 32767

																			lowerCAmelCase__       =						BytesIO(bytes()       )
																			sf.write(__UpperCAmelCase		,    __UpperCAmelCase		,    value["sampling_rate"]		,    format="wav"       )
																			return {"bytes": buffer.getvalue(), "path": None}
													else:
																			return {"bytes": None, "path": value.get("path"       )}
							elif value.get("bytes"       ) is not None or value.get("path"       ) is not None:
													# store the audio bytes, and path is used to infer the audio format using the file extension
													return {"bytes": value.get("bytes"       ), "path": value.get("path"       )}
							else:
													raise ValueError(
													    F"An audio sample should have one of 'path' or 'bytes' but they are missing or None in {value}."       )

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase		,    __UpperCAmelCase = None       )->					dict:







							'''simple docstring'''





							if not self.decode:
													raise RuntimeError("Decoding is disabled for this feature. Please use Audio(decode=True) instead."       )

							lowerCAmelCase__      ,							lowerCAmelCase__       =						(value["path"], BytesIO(value["bytes"]       )) if value["bytes"] is not None else (value["path"], None)
							if path is None and file is None:
													raise ValueError(F"An audio sample should have one of 'path' or 'bytes' but both are None in {value}."       )

							try:
													import librosa
													import soundfile as sf
							except ImportError as err:
													raise ImportError("To support decoding audio files, please install 'librosa' and 'soundfile'."       ) from err

							lowerCAmelCase__       =						xsplitext(__UpperCAmelCase       )[1][1:].lower() if path is not None else None
							if not config.IS_OPUS_SUPPORTED and audio_format == "opus":
													raise RuntimeError(
													    "Decoding 'opus' files requires system library 'libsndfile'>=1.0.31, "
													    "You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. "       )
							elif not config.IS_MP3_SUPPORTED and audio_format == "mp3":
													raise RuntimeError(
													    "Decoding 'mp3' files requires system library 'libsndfile'>=1.1.0, "
													    "You can try to update `soundfile` python library: `pip install \"soundfile>=0.12.1\"`. "       )

							if file is None:
													lowerCAmelCase__       =						token_per_repo_id or {}
													lowerCAmelCase__       =						path.split("::"       )[-1]
													try:
																			lowerCAmelCase__       =						string_to_dict(__UpperCAmelCase		,    config.HUB_DATASETS_URL       )["repo_id"]
																			lowerCAmelCase__       =						token_per_repo_id[repo_id]
													except (ValueError, KeyError):
																			lowerCAmelCase__       =						None

													with xopen(__UpperCAmelCase		,    "rb"		,    use_auth_token=__UpperCAmelCase       ) as f:
																			lowerCAmelCase__      ,							lowerCAmelCase__       =						sf.read(__UpperCAmelCase       )

							else:
													lowerCAmelCase__      ,							lowerCAmelCase__       =						sf.read(__UpperCAmelCase       )

							lowerCAmelCase__       =						array.T
							if self.mono:
													lowerCAmelCase__       =						librosa.to_mono(__UpperCAmelCase       )
							if self.sampling_rate and self.sampling_rate != sampling_rate:
													lowerCAmelCase__       =						librosa.resample(__UpperCAmelCase		,    orig_sr=__UpperCAmelCase		,    target_sr=self.sampling_rate       )
													lowerCAmelCase__       =						self.sampling_rate

							return {"path": path, "array": array, "sampling_rate": sampling_rate}

	def 		UpperCAmelCase  (							self       )->					Union["FeatureType", Dict[str, "FeatureType"]]:







							'''simple docstring'''





							from .features import Value

							if self.decode:
													raise ValueError("Cannot flatten a decoded Audio feature."       )
							return {
							    "bytes": Value("binary"       ),
							    "path": Value("string"       ),
							}

	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					pa.StructArray:







							'''simple docstring'''





							if pa.types.is_string(storage.type       ):
													lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.binary()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, storage]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							elif pa.types.is_binary(storage.type       ):
													lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.string()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([storage, path_array]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							elif pa.types.is_struct(storage.type       ) and storage.type.get_all_field_indices("array"       ):
													lowerCAmelCase__       =						pa.array([Audio().encode_example(__UpperCAmelCase       ) if x is not None else None for x in storage.to_pylist()]       )
							elif pa.types.is_struct(storage.type       ):
													if storage.type.get_field_index("bytes"       ) >= 0:
																			lowerCAmelCase__       =						storage.field("bytes"       )
													else:
																			lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.binary()       )
													if storage.type.get_field_index("path"       ) >= 0:
																			lowerCAmelCase__       =						storage.field("path"       )
													else:
																			lowerCAmelCase__       =						pa.array([None] * len(__UpperCAmelCase       )		,    type=pa.string()       )
													lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, path_array]		,    ["bytes", "path"]		,    mask=storage.is_null()       )
							return array_cast(__UpperCAmelCase		,    self.pa_type       )


	def 		UpperCAmelCase  (							self		,    __UpperCAmelCase       )->					pa.StructArray:







							'''simple docstring'''





							@no_op_if_value_is_null
							def path_to_bytes(__UpperCAmelCase       ):
													with xopen(__UpperCAmelCase		,    "rb"       ) as f:
																			lowerCAmelCase__       =						f.read()
													return bytes_

							lowerCAmelCase__       =						pa.array(
							    [
							        (path_to_bytes(x["path"]       ) if x["bytes"] is None else x["bytes"]) if x is not None else None
							        for x in storage.to_pylist()
							    ]		,    type=pa.binary()		,    )
							lowerCAmelCase__       =						pa.array(
							    [os.path.basename(__UpperCAmelCase       ) if path is not None else None for path in storage.field("path"       ).to_pylist()]		,    type=pa.string()		,    )
							lowerCAmelCase__       =						pa.StructArray.from_arrays([bytes_array, path_array]		,    ["bytes", "path"]		,    mask=bytes_array.is_null()       )
							return array_cast(__UpperCAmelCase		,    self.pa_type       )






