




'''simple docstring'''






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()
_lowerCAmelCase      	=					logging.get_logger(__name__)

_lowerCAmelCase      	=					{
    "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",
}
_lowerCAmelCase      	=					[
    "ctc_proj",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]

def 			__lowerCAmelCase					(    snake_case__					,      snake_case__					,      snake_case__					,      snake_case__					,      snake_case__					,      snake_case__						):
			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
												__UpperCamelCase      :			Optional[Any]									=   'lm_head'

						__UpperCamelCase      :			int									=   getattr(snake_case__					,      snake_case__						)

			if weight_type is not None:
						__UpperCamelCase      :			Tuple									=   getattr(snake_case__					,      snake_case__						).shape
			else:
						__UpperCamelCase      :			List[Any]									=   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      :			Any									=   value
			elif weight_type == "weight_g":
						__UpperCamelCase      :			Union[str, Any]									=   value
			elif weight_type == "weight_v":
						__UpperCamelCase      :			Optional[Any]									=   value
			elif weight_type == "bias":
						__UpperCamelCase      :			Dict									=   value
			else:
						__UpperCamelCase      :			Optional[Any]									=   value

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

def 			__lowerCAmelCase					(    snake_case__					,      snake_case__					,      snake_case__						):
			__UpperCamelCase      :			Optional[int]									=   []
			__UpperCamelCase      :			Optional[Any]									=   fairseq_model.state_dict()

			__UpperCamelCase      :			Tuple									=   hf_model.unispeech.feature_extractor

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

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

def 			__lowerCAmelCase					(    snake_case__					,      snake_case__					,      snake_case__					,      snake_case__					,      snake_case__						):
			__UpperCamelCase      :			Optional[int]									=   full_name.split("conv_layers."						)[-1]
			__UpperCamelCase      :			str									=   name.split("."						)
			__UpperCamelCase      :			Union[str, Any]									=   int(items[0]						)
			__UpperCamelCase      :			Optional[Any]									=   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      :			int									=   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      :			List[str]									=   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      :			List[Any]									=   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      :			str									=   value
									logger.info(F"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."						)
			else:
						unused_weights.append(snake_case__						)

@torch.no_grad()
def 			__lowerCAmelCase					(    snake_case__					,      snake_case__					,      snake_case__=None					,      snake_case__=None					,      snake_case__=True						):
			if config_path is not None:
						__UpperCamelCase      :			Tuple									=   UniSpeechConfig.from_pretrained(snake_case__						)
			else:
						__UpperCamelCase      :			Tuple									=   UniSpeechConfig()

			if is_finetuned:
						if dict_path:
									__UpperCamelCase      :			Tuple									=   Dictionary.load_from_json(snake_case__						)

									# important change bos & pad token id since CTC symbol is <pad> and
									# not <s> as in fairseq
									__UpperCamelCase      :			Any									=   target_dict.pad_index
									__UpperCamelCase      :			List[str]									=   target_dict.bos_index
									__UpperCamelCase      :			List[Any]									=   target_dict.eos_index
									__UpperCamelCase      :			Optional[Any]									=   len(target_dict.symbols						)
									__UpperCamelCase      :			List[Any]									=   os.path.join(snake_case__					,      "vocab.json"						)
									if not os.path.isdir(snake_case__						):
												logger.error("--pytorch_dump_folder_path ({}) should be a directory".format(snake_case__						)						)
												return
									os.makedirs(snake_case__					,      exist_ok=snake_case__						)
									__UpperCamelCase      :			List[Any]									=   target_dict.indices

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

						__UpperCamelCase      :			int									=   UniSpeechForCTC(snake_case__						)
			else:
						__UpperCamelCase      :			Optional[int]									=   UniSpeechForPreTraining(snake_case__						)

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

			__UpperCamelCase      :			Tuple									=   model[0].eval()

			recursively_load_weights(snake_case__					,      snake_case__					,      snake_case__						)

			hf_unispeech.save_pretrained(snake_case__						)


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_unispeech_checkpoint(
				    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
				)








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


a_   :Optional[Any]        =      {"configuration_reformer": ["REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP", "ReformerConfig"]}

try:
 if not is_sentencepiece_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   :str        =      ["ReformerTokenizer"]

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   :int        =      ["ReformerTokenizerFast"]

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 a_   :List[str]        =      [
     "REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST",
     "ReformerAttention",
     "ReformerForMaskedLM",
     "ReformerForQuestionAnswering",
     "ReformerForSequenceClassification",
     "ReformerLayer",
     "ReformerModel",
     "ReformerModelWithLMHead",
     "ReformerPreTrainedModel",
 ]


if TYPE_CHECKING:
 from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig

 try:
  if not is_sentencepiece_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_reformer import ReformerTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_reformer_fast import ReformerTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_reformer import (
      REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
      ReformerAttention,
      ReformerForMaskedLM,
      ReformerForQuestionAnswering,
      ReformerForSequenceClassification,
      ReformerLayer,
      ReformerModel,
      ReformerModelWithLMHead,
      ReformerPreTrainedModel,
  )

else:
 import sys

 a_   :Optional[Any]        =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






