




'''simple docstring'''




import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    HubertConfig,
    HubertForCTC,
    HubertModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


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

UpperCAmelCase_				 =					{
    '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',
    'w2v_encoder.proj': 'lm_head',
    'mask_emb': 'masked_spec_embed',
}

def       _UpperCamelCase					( SCREAMING_SNAKE_CASE__	:					Dict				,  SCREAMING_SNAKE_CASE__	:					Optional[int]				,  SCREAMING_SNAKE_CASE__	:					int				,  SCREAMING_SNAKE_CASE__	:					Union[str, Any]				,  SCREAMING_SNAKE_CASE__	:					List[Any]				):
 '''simple docstring'''





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

 if weight_type is not None:
  UpperCAmelCase__								=			getattr(__lowerCamelCase				,  __lowerCamelCase				).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					( SCREAMING_SNAKE_CASE__	:					Optional[int]				,  SCREAMING_SNAKE_CASE__	:					Optional[Any]				,  SCREAMING_SNAKE_CASE__	:					List[Any]				):
 '''simple docstring'''





 UpperCAmelCase__								=			[]
 UpperCAmelCase__								=			fairseq_model.state_dict()

 UpperCAmelCase__								=			hf_model.hubert.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(
       __lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				,  hf_model.config.feat_extract_norm == """group"""				,  )
   UpperCAmelCase__								=			True
  else:
   for key, mapped_key in MAPPING.items():
    UpperCAmelCase__								=			"""hubert.""" + 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] and not is_finetuned):
     UpperCAmelCase__								=			True
     if "*" in mapped_key:
      UpperCAmelCase__								=			name.split(__lowerCamelCase				)[0].split("""."""				)[-2]
      UpperCAmelCase__								=			mapped_key.replace("""*"""				,  __lowerCamelCase				)
     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(__lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				)
    continue
  if not is_used:
   unused_weights.append(__lowerCamelCase				)

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

def       _UpperCamelCase					( SCREAMING_SNAKE_CASE__	:					Dict				,  SCREAMING_SNAKE_CASE__	:					Any				,  SCREAMING_SNAKE_CASE__	:					List[Any]				,  SCREAMING_SNAKE_CASE__	:					Optional[int]				,  SCREAMING_SNAKE_CASE__	:					Union[str, Any]				):
 '''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(__lowerCamelCase				)

@torch.no_grad()
def       _UpperCamelCase					( SCREAMING_SNAKE_CASE__	:					str				,  SCREAMING_SNAKE_CASE__	:					str				,  SCREAMING_SNAKE_CASE__	:					List[Any]=None				,  SCREAMING_SNAKE_CASE__	:					Optional[int]=None				,  SCREAMING_SNAKE_CASE__	:					Union[str, Any]=True				):
 '''simple docstring'''





 if config_path is not None:
  UpperCAmelCase__								=			HubertConfig.from_pretrained(__lowerCamelCase				)
 else:
  UpperCAmelCase__								=			HubertConfig()

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

   # 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.eos_index
   UpperCAmelCase__								=			len(target_dict.symbols				)
   UpperCAmelCase__								=			os.path.join(__lowerCamelCase				,  """vocab.json"""				)
   if not os.path.isdir(__lowerCamelCase				):
    logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(__lowerCamelCase				)				)
    return
   os.makedirs(__lowerCamelCase				,  exist_ok=__lowerCamelCase				)
   with open(__lowerCamelCase				,  """w"""				,  encoding="""utf-8"""				) as vocab_handle:
    json.dump(target_dict.indices				,  __lowerCamelCase				)
   UpperCAmelCase__								=			WavaVecaCTCTokenizer(
       __lowerCamelCase				,  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=__lowerCamelCase				,  )
   UpperCAmelCase__								=			True if config.feat_extract_norm == """layer""" else False
   UpperCAmelCase__								=			WavaVecaFeatureExtractor(
       feature_size=1				,  sampling_rate=16000				,  padding_value=0				,  do_normalize=__lowerCamelCase				,  return_attention_mask=__lowerCamelCase				,  )
   UpperCAmelCase__								=			WavaVecaProcessor(feature_extractor=__lowerCamelCase				,  tokenizer=__lowerCamelCase				)
   processor.save_pretrained(__lowerCamelCase				)

  UpperCAmelCase__								=			HubertForCTC(__lowerCamelCase				)
 else:
  UpperCAmelCase__								=			HubertModel(__lowerCamelCase				)

 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]				)

 UpperCAmelCase__								=			model[0].eval()

 recursively_load_weights(__lowerCamelCase				,  __lowerCamelCase				,  __lowerCamelCase				)

 hf_wavavec.save_pretrained(__lowerCamelCase				)


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







'''simple docstring'''

from typing import TYPE_CHECKING

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


lowercase_					= {"""configuration_mbart""": ["""MBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MBartConfig""", """MBartOnnxConfig"""]}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					= ["""MBartTokenizer"""]

try:
      if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					= ["""MBartTokenizerFast"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					= [
          """MBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """MBartForCausalLM""",
          """MBartForConditionalGeneration""",
          """MBartForQuestionAnswering""",
          """MBartForSequenceClassification""",
          """MBartModel""",
          """MBartPreTrainedModel""",
      ]

try:
      if not is_tf_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					= [
          """TFMBartForConditionalGeneration""",
          """TFMBartModel""",
          """TFMBartPreTrainedModel""",
      ]

try:
      if not is_flax_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      lowercase_					= [
          """FlaxMBartForConditionalGeneration""",
          """FlaxMBartForQuestionAnswering""",
          """FlaxMBartForSequenceClassification""",
          """FlaxMBartModel""",
          """FlaxMBartPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfig, MBartOnnxConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_mbart import MBartTokenizer

      try:
            if not is_tokenizers_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_mbart_fast import MBartTokenizerFast

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_mbart import (
                MBART_PRETRAINED_MODEL_ARCHIVE_LIST,
                MBartForCausalLM,
                MBartForConditionalGeneration,
                MBartForQuestionAnswering,
                MBartForSequenceClassification,
                MBartModel,
                MBartPreTrainedModel,
            )

      try:
            if not is_tf_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_tf_mbart import TFMBartForConditionalGeneration, TFMBartModel, TFMBartPreTrainedModel

      try:
            if not is_flax_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_flax_mbart import (
                FlaxMBartForConditionalGeneration,
                FlaxMBartForQuestionAnswering,
                FlaxMBartForSequenceClassification,
                FlaxMBartModel,
                FlaxMBartPreTrainedModel,
            )

else:
      import sys

      lowercase_					= _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



