


import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, 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',
    'encoder.layer_norm_for_extract': 'layer_norm_for_extract',
    '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',
    'label_embs_concat': 'label_embeddings_concat',
    'mask_emb': 'masked_spec_embed',
    'spk_proj': 'speaker_proj',
}
_A			=						[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'label_embeddings_concat',
    'speaker_proj',
    'layer_norm_for_extract',
]

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : List[str]      ,					SCREAMING_SNAKE_CASE__      : Dict      ,					SCREAMING_SNAKE_CASE__      : List[str]      ,					SCREAMING_SNAKE_CASE__      : str      ,					SCREAMING_SNAKE_CASE__      : Optional[Any]    ):
     for attribute in key.split('.'    ):
          __UpperCamelCase   							=getattr(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )

     if weight_type is not None:
          __UpperCamelCase   							=getattr(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    ).shape
     else:
          __UpperCamelCase   							=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":
          __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[Any]      ,					SCREAMING_SNAKE_CASE__      : str    ):
     __UpperCamelCase   							=[]
     __UpperCamelCase   							=fairseq_model.state_dict()

     __UpperCamelCase   							=hf_model.unispeech_sat.feature_extractor

     for name, value in fairseq_dict.items():
          __UpperCamelCase   							=False
          if "conv_layers" in name:
               load_conv_layer(
                   SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					hf_model.config.feat_extract_norm == 'group'      ,					)
               __UpperCamelCase   							=True
          else:
               for key, mapped_key in MAPPING.items():
                    __UpperCamelCase   							='unispeech_sat.' + 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]:
                         if "layer_norm_for_extract" in name and (".".join(name.split('.'    )[:-1]    ) != key):
                              # special case since naming is very similar
                              continue
                         __UpperCamelCase   							=True
                         if "*" in mapped_key:
                              __UpperCamelCase   							=name.split(SCREAMING_SNAKE_CASE__    )[0].split('.'    )[-2]
                              __UpperCamelCase   							=mapped_key.replace('*'      ,					SCREAMING_SNAKE_CASE__    )
                         if "weight_g" in name:
                              __UpperCamelCase   							='weight_g'
                         elif "weight_v" in name:
                              __UpperCamelCase   							='weight_v'
                         elif "bias" in name:
                              __UpperCamelCase   							='bias'
                         elif "weight" in name:
                              # TODO: don't match quantizer.weight_proj
                              __UpperCamelCase   							='weight'
                         else:
                              __UpperCamelCase   							=None
                         set_recursively(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )
                    continue
          if not is_used:
               unused_weights.append(SCREAMING_SNAKE_CASE__    )

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

def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Optional[Any]      ,					SCREAMING_SNAKE_CASE__      : List[Any]      ,					SCREAMING_SNAKE_CASE__      : Union[str, Any]      ,					SCREAMING_SNAKE_CASE__      : List[str]      ,					SCREAMING_SNAKE_CASE__      : Dict    ):
     __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:
               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.'    )
               __UpperCamelCase   							=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.'    )
               __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:
               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[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:
               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[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(SCREAMING_SNAKE_CASE__    )

@torch.no_grad()
def 						_UpperCAmelCase    (      SCREAMING_SNAKE_CASE__      : Any      ,					SCREAMING_SNAKE_CASE__      : Any      ,					SCREAMING_SNAKE_CASE__      : List[str]=None      ,					SCREAMING_SNAKE_CASE__      : Dict=None      ,					SCREAMING_SNAKE_CASE__      : Dict=True    ):
     if config_path is not None:
          __UpperCamelCase   							=UniSpeechSatConfig.from_pretrained(SCREAMING_SNAKE_CASE__    )
     else:
          __UpperCamelCase   							=UniSpeechSatConfig()

     __UpperCamelCase   							=''

     if is_finetuned:
          __UpperCamelCase   							=UniSpeechSatForCTC(SCREAMING_SNAKE_CASE__    )
     else:
          __UpperCamelCase   							=UniSpeechSatForPreTraining(SCREAMING_SNAKE_CASE__    )

     __UpperCamelCase							,				__UpperCamelCase							,				__UpperCamelCase   							=fairseq.checkpoint_utils.load_model_ensemble_and_task(
         [checkpoint_path]      ,					arg_overrides={'data': '/'.join(dict_path.split('/'    )[:-1]    )}    )
     __UpperCamelCase   							=model[0].eval()

     recursively_load_weights(SCREAMING_SNAKE_CASE__      ,					SCREAMING_SNAKE_CASE__    )

     hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE__    )


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









import warnings
from contextlib import contextmanager

from ....processing_utils import ProcessorMixin






class    UpperCAmelCase__	( A_		):



    """simple docstring"""






    UpperCAmelCase__						:		Optional[int]              =  "MCTCTFeatureExtractor"
    UpperCAmelCase__						:		str              =  "AutoTokenizer"






    def __init__(      self ,       A_ ,       A_   )  ->	Dict:
         super().__init__(A_ ,       A_   )
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False






    def __call__(      self ,       *A_ ,       **A_   )  ->	Optional[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor(*A_ ,       **A_   )

         if "raw_speech" in kwargs:
              warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'   )
              __UpperCamelCase   							=kwargs.pop('raw_speech'   )
         else:
              __UpperCamelCase   							=kwargs.pop('audio' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('sampling_rate' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('text' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if audio is None and text is None:
              raise ValueError('You need to specify either an `audio` or `text` input to process.'   )

         if audio is not None:
              __UpperCamelCase   							=self.feature_extractor(A_ ,       *A_ ,       sampling_rate=A_ ,       **A_   )
         if text is not None:
              __UpperCamelCase   							=self.tokenizer(A_ ,       **A_   )

         if text is None:
              return inputs
         elif audio is None:
              return encodings
         else:
              __UpperCamelCase   							=encodings['input_ids']
              return inputs






    def _a (      self ,       *A_ ,       **A_   )  ->	str:
         return self.tokenizer.batch_decode(*A_ ,       **A_   )






    def _a (      self ,       *A_ ,       **A_   )  ->	List[Any]:
         # For backward compatibility
         if self._in_target_context_manager:
              return self.current_processor.pad(*A_ ,       **A_   )

         __UpperCamelCase   							=kwargs.pop('input_features' ,       A_   )
         __UpperCamelCase   							=kwargs.pop('labels' ,       A_   )
         if len(A_   ) > 0:
              __UpperCamelCase   							=args[0]
              __UpperCamelCase   							=args[1:]

         if input_features is not None:
              __UpperCamelCase   							=self.feature_extractor.pad(A_ ,       *A_ ,       **A_   )
         if labels is not None:
              __UpperCamelCase   							=self.tokenizer.pad(A_ ,       **A_   )

         if labels is None:
              return input_features
         elif input_features is None:
              return labels
         else:
              __UpperCamelCase   							=labels['input_ids']
              return input_features






    def _a (      self ,       *A_ ,       **A_   )  ->	Optional[int]:
         return self.tokenizer.decode(*A_ ,       **A_   )






    @contextmanager
    def _a (      self   )  ->	str:
         warnings.warn(
             '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
             'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
             'your audio inputs, or in a separate call.'   )
         __UpperCamelCase   							=True
         __UpperCamelCase   							=self.tokenizer
         yield
         __UpperCamelCase   							=self.feature_extractor
         __UpperCamelCase   							=False





