




'''simple docstring'''





import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


logging.set_verbosity_info()
_SCREAMING_SNAKE_CASE							:       List[Any]						=			logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE							:       List[Any]						=			{
    """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""",
    """adapter_layer""": """encoder.layers.*.adapter_layer""",
    """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""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
    """pooling_layer.linear""": """projector""",
    """pooling_layer.projection""": """classifier""",
}
_SCREAMING_SNAKE_CASE							:       Tuple						=			[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
    """projector""",
    """classifier""",
]


def 					UpperCamelCase_(     snake_case							:	int					):

							'''simple docstring'''






							snake_case_      =				{}
							with open(_UpperCAmelCase							,    "r"					) as file:
														for line_number, line in enumerate(_UpperCAmelCase					):
																					snake_case_      =				line.strip()
																					if line:
																												snake_case_      =				line.split()
																												snake_case_      =				line_number
																												snake_case_      =				words[0]
																												snake_case_      =				value
							return result


def 					UpperCamelCase_(     snake_case							:	Optional[Any]							,    snake_case							:	Optional[int]							,    snake_case							:	Tuple							,    snake_case							:	List[str]							,    snake_case							:	Optional[Any]					):

							'''simple docstring'''






							for attribute in key.split("."					):
														snake_case_      =				getattr(_UpperCAmelCase							,    _UpperCAmelCase					)

							snake_case_      =				None
							for param_key in PARAM_MAPPING.keys():
														if full_name.endswith(_UpperCAmelCase					):
																					snake_case_      =				PARAM_MAPPING[full_name.split("."					)[-1]]
																					snake_case_      =				'param'

							if weight_type is not None and weight_type != "param":
														snake_case_      =				getattr(_UpperCAmelCase							,    _UpperCAmelCase					).shape
							elif weight_type is not None and weight_type == "param":
														snake_case_      =				hf_pointer
														for attribute in hf_param_name.split("."					):
																					snake_case_      =				getattr(_UpperCAmelCase							,    _UpperCAmelCase					)
														snake_case_      =				shape_pointer.shape

														# let's reduce dimension
														snake_case_      =				value[0]
							else:
														snake_case_      =				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":
														snake_case_      =				value
							elif weight_type == "weight_g":
														snake_case_      =				value
							elif weight_type == "weight_v":
														snake_case_      =				value
							elif weight_type == "bias":
														snake_case_      =				value
							elif weight_type == "param":
														for attribute in hf_param_name.split("."					):
																					snake_case_      =				getattr(_UpperCAmelCase							,    _UpperCAmelCase					)
														snake_case_      =				value
							else:
														snake_case_      =				value

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


def 					UpperCamelCase_(     snake_case							:	Optional[int]							,    snake_case							:	Optional[int]							,    snake_case							:	Optional[int]							,    snake_case							:	Any							,    snake_case							:	Any					):

							'''simple docstring'''






							snake_case_      =				None
							for param_key in PARAM_MAPPING.keys():
														if full_name.endswith(_UpperCAmelCase					):
																					snake_case_      =				PARAM_MAPPING[full_name.split("."					)[-1]]
																					snake_case_      =				'param'

							if weight_type is not None and weight_type != "param":
														snake_case_      =				'.'.join([key, weight_type]					)
							elif weight_type is not None and weight_type == "param":
														snake_case_      =				'.'.join([key, hf_param_name]					)
							else:
														snake_case_      =				key

							snake_case_      =				value if 'lm_head' in full_key else value[0]


_SCREAMING_SNAKE_CASE							:       List[Any]						=			{
    """W_a""": """linear_1.weight""",
    """W_b""": """linear_2.weight""",
    """b_a""": """linear_1.bias""",
    """b_b""": """linear_2.bias""",
    """ln_W""": """norm.weight""",
    """ln_b""": """norm.bias""",
}


def 					UpperCamelCase_(     snake_case							:	List[str]							,    snake_case							:	Any							,    snake_case							:	Tuple=None							,    snake_case							:	Dict=None					):

							'''simple docstring'''






							snake_case_      =				False
							for key, mapped_key in MAPPING.items():
														snake_case_      =				'wav2vec2.' + 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]:
																					snake_case_      =				True
																					if "*" in mapped_key:
																												snake_case_      =				name.split(_UpperCAmelCase					)[0].split("."					)[-2]
																												snake_case_      =				mapped_key.replace("*"							,    _UpperCAmelCase					)
																					if "weight_g" in name:
																												snake_case_      =				'weight_g'
																					elif "weight_v" in name:
																												snake_case_      =				'weight_v'
																					elif "bias" in name:
																												snake_case_      =				'bias'
																					elif "weight" in name:
																												# TODO: don't match quantizer.weight_proj
																												snake_case_      =				'weight'
																					else:
																												snake_case_      =				None
																					if hf_dict is not None:
																												rename_dict(_UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase					)
																					else:
																												set_recursively(_UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase					)
																					return is_used
							return is_used


def 					UpperCamelCase_(     snake_case							:	Optional[int]							,    snake_case							:	Dict							,    snake_case							:	Optional[int]					):

							'''simple docstring'''






							snake_case_      =				[]
							snake_case_      =				fairseq_model.state_dict()

							snake_case_      =				hf_model.wavaveca.feature_extractor

							for name, value in fairseq_dict.items():
														snake_case_      =				False
														if "conv_layers" in name:
																					load_conv_layer(
																					    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase							,    hf_model.config.feat_extract_norm == "group"							,    )
																					snake_case_      =				True
														else:
																					snake_case_      =				load_wavaveca_layer(_UpperCAmelCase							,    _UpperCAmelCase							,    _UpperCAmelCase					)
														if not is_used:
																					unused_weights.append(_UpperCAmelCase					)

							logger.warning(f'Unused weights: {unused_weights}'					)


def 					UpperCamelCase_(     snake_case							:	Tuple							,    snake_case							:	Dict							,    snake_case							:	Tuple							,    snake_case							:	int							,    snake_case							:	List[str]					):

							'''simple docstring'''






							snake_case_      =				full_name.split("conv_layers."					)[-1]
							snake_case_      =				name.split("."					)
							snake_case_      =				int(items[0]					)
							snake_case_      =				int(items[1]					)

							if type_id == 0:
														if "bias" in name:
																					if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
																												raise ValueError(
																												    f'{full_name} has size {value.shape}, but'
																												    f' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'					)
																					snake_case_      =				value
																					logger.info(f'Feat extract conv layer {layer_id} was initialized from {full_name}.'					)
														elif "weight" in name:
																					if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
																												raise ValueError(
																												    f'{full_name} has size {value.shape}, but'
																												    f' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'					)
																					snake_case_      =				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:
																					if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
																												raise ValueError(
																												    f'{full_name} has size {value.shape}, but'
																												    f' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'					)
																					snake_case_      =				value
																					logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'					)
														elif "weight" in name:
																					if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
																												raise ValueError(
																												    f'{full_name} has size {value.shape}, but'
																												    f' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'					)
																					snake_case_      =				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 					UpperCamelCase_(     snake_case							:	List[Any]							,    snake_case							:	Union[str, Any]							,    snake_case							:	List[Any]=None							,    snake_case							:	Dict=None							,    snake_case							:	List[str]=True							,    snake_case							:	Optional[int]=False					):

							'''simple docstring'''






							if config_path is not None:
														snake_case_      =				WavaVecaConfig.from_pretrained(_UpperCAmelCase					)
							else:
														snake_case_      =				WavaVecaConfig()

							if is_seq_class:
														snake_case_      =				read_txt_into_dict(_UpperCAmelCase					)
														snake_case_      =				idalabel
														snake_case_      =				WavaVecaForSequenceClassification(_UpperCAmelCase					)
														snake_case_      =				WavaVecaFeatureExtractor(
														    feature_size=1							,    sampling_rate=1_6_0_0_0							,    padding_value=0							,    do_normalize=_UpperCAmelCase							,    return_attention_mask=_UpperCAmelCase							,    )
														feature_extractor.save_pretrained(_UpperCAmelCase					)

							elif is_finetuned:
														if dict_path:
																					snake_case_      =				Dictionary.load(_UpperCAmelCase					)

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

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

														snake_case_      =				WavaVecaForCTC(_UpperCAmelCase					)
							else:
														snake_case_      =				WavaVecaForPreTraining(_UpperCAmelCase					)

							if is_finetuned or is_seq_class:
														snake_case_      =				fairseq.checkpoint_utils.load_model_ensemble_and_task(
														    [checkpoint_path]							,    arg_overrides={"data": "/".join(dict_path.split("/"					)[:-1]					)}					)
							else:
														snake_case_      =				argparse.Namespace(task="audio_pretraining"					)
														snake_case_      =				fairseq.tasks.setup_task(_UpperCAmelCase					)

														snake_case_      =				fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]							,    task=_UpperCAmelCase					)

							snake_case_      =				model[0].eval()

							recursively_load_weights(_UpperCAmelCase							,    _UpperCAmelCase							,    not is_finetuned					)

							hf_wavavec.save_pretrained(_UpperCAmelCase					)


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE							:       List[Any]						=			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"
		)
		parser.add_argument(
		    "--is_seq_class",
		    action="store_true",
		    help="Whether the model to convert is a fine-tuned sequence classification model or not",
		)
		_SCREAMING_SNAKE_CASE							:       Optional[int]						=			parser.parse_args()

		_SCREAMING_SNAKE_CASE							:       Optional[int]						=			not args.not_finetuned and not args.is_seq_class
		convert_wavaveca_checkpoint(
		    args.checkpoint_path,
		    args.pytorch_dump_folder_path,
		    args.config_path,
		    args.dict_path,
		    is_finetuned,
		    args.is_seq_class,
		)











'''simple docstring'''





import argparse
from collections import defaultdict

import yaml


_SCREAMING_SNAKE_CASE							:       Optional[Any]						=			"docs/source/en/_toctree.yml"


def 					UpperCamelCase_(     snake_case							:	Optional[Any]					):

							'''simple docstring'''






							snake_case_      =				defaultdict(snake_case					)
							for doc in model_doc:
														counts[doc["local"]] += 1
							snake_case_      =				[key for key, value in counts.items() if value > 1]

							snake_case_      =				[]
							for duplicate_key in duplicates:
														snake_case_      =				list({doc["title"] for doc in model_doc if doc["local"] == duplicate_key}					)
														if len(snake_case					) > 1:
																					raise ValueError(
																					    f'{duplicate_key} is present several times in the documentation table of content at '
																					    "`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the "
																					    "others."					)
														# Only add this once
														new_doc.append({"local": duplicate_key, "title": titles[0]}					)

							# Add none duplicate-keys
							new_doc.extend([doc for doc in model_doc if counts[doc["local"]] == 1]					)

							# Sort
							return sorted(snake_case							,    key=lambda snake_case					: s["title"].lower()					)


def 					UpperCamelCase_(     snake_case							:	Optional[int]=False					):

							'''simple docstring'''






							with open(snake_case							,    encoding="utf-8"					) as f:
														snake_case_      =				yaml.safe_load(f.read()					)

							# Get to the API doc
							snake_case_      =				0
							while content[api_idx]["title"] != "API":
														api_idx += 1
							snake_case_      =				content[api_idx]["sections"]

							# Then to the model doc
							snake_case_      =				0
							while api_doc[model_idx]["title"] != "Models":
														model_idx += 1

							snake_case_      =				api_doc[model_idx]["sections"]

							snake_case_      =				[(idx, section) for idx, section in enumerate(snake_case					) if "sections" in section]
							snake_case_      =				False
							for idx, modality_doc in modalities_docs:
														snake_case_      =				modality_doc["sections"]
														snake_case_      =				clean_model_doc_toc(snake_case					)

														if old_modality_doc != new_modality_doc:
																					snake_case_      =				True
																					if overwrite:
																												snake_case_      =				new_modality_doc

							if diff:
														if overwrite:
																					snake_case_      =				model_doc
																					snake_case_      =				api_doc
																					with open(snake_case							,    "w"							,    encoding="utf-8"					) as f:
																												f.write(yaml.dump(snake_case							,    allow_unicode=snake_case					)					)
														else:
																					raise ValueError(
																					    "The model doc part of the table of content is not properly sorted, run `make style` to fix this."					)


if __name__ == "__main__":
		_SCREAMING_SNAKE_CASE							:       List[str]						=			argparse.ArgumentParser()
		parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
		_SCREAMING_SNAKE_CASE							:       Any						=			parser.parse_args()

		check_model_doc(args.fix_and_overwrite)




