






'''simple docstring'''







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase     :   List[str]	      =		{
    'configuration_x_clip': [
        'XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'XCLIPConfig',
        'XCLIPTextConfig',
        'XCLIPVisionConfig',
    ],
    'processing_x_clip': ['XCLIPProcessor'],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							UpperCamelCase     :   Union[str, Any]	      =		[
							    'XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST',
							    'XCLIPModel',
							    'XCLIPPreTrainedModel',
							    'XCLIPTextModel',
							    'XCLIPVisionModel',
							]

if TYPE_CHECKING:
							from .configuration_x_clip import (
							    XCLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
							    XCLIPConfig,
							    XCLIPTextConfig,
							    XCLIPVisionConfig,
							)
							from .processing_x_clip import XCLIPProcessor

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .modeling_x_clip import (
														    XCLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
														    XCLIPModel,
														    XCLIPPreTrainedModel,
														    XCLIPTextModel,
														    XCLIPVisionModel,
														)

else:
							import sys

							UpperCamelCase     :   Dict	      =		_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


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

UpperCAmelCase_		 =					{
    """post_extract_proj""": """feature_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.upsample.0""": """encoder.upsample.projection""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """layer_norm""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Union[str, Any]	,   _snake_case	:str	,   _snake_case	:Any	,   _snake_case	:int	,   _snake_case	:List[Any]				)  ->     Optional[int]:
   for attribute in key.split('''.'''				):
      _A      =       getattr(_snake_case	,   _snake_case				)

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

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

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Union[str, Any]	,   _snake_case	:Any	,   _snake_case	:int				)  ->     Any:
   _A      =       []
   _A      =       fairseq_model.state_dict()

   _A      =       hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

   for name, value in fairseq_dict.items():
      _A      =       False
      if "conv_layers" in name:
         load_conv_layer(
             _snake_case	,   _snake_case	,   _snake_case	,   _snake_case	,   hf_model.config.feat_extract_norm == '''group'''	,   )
         _A      =       True
      else:
         for key, mapped_key in MAPPING.items():
            _A      =       '''sew.''' + 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]:
               _A      =       True
               if "*" in mapped_key:
                  _A      =       name.split(_snake_case				)[0].split('''.'''				)[-2]
                  _A      =       mapped_key.replace('''*'''	,   _snake_case				)
               if "weight_g" in name:
                  _A      =       '''weight_g'''
               elif "weight_v" in name:
                  _A      =       '''weight_v'''
               elif "weight" in name:
                  _A      =       '''weight'''
               elif "bias" in name:
                  _A      =       '''bias'''
               else:
                  _A      =       None
               set_recursively(_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 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Tuple	,   _snake_case	:List[str]	,   _snake_case	:List[str]	,   _snake_case	:Optional[int]	,   _snake_case	:List[Any]				)  ->     Any:
   _A      =       full_name.split('''conv_layers.'''				)[-1]
   _A      =       name.split('''.'''				)
   _A      =       int(items[0]				)
   _A      =       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.'''
         )
         _A      =       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.'''
         )
         _A      =       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."
         )
         _A      =       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.'''
         )
         _A      =       value
         logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''				)
   else:
      unused_weights.append(_snake_case				)

def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:str	,   _snake_case	:Dict				)  ->     Tuple:
   _A      =       SEWConfig()
   if is_finetuned:
      _A      =       model.wav_encoder.wav_model.cfg
   else:
      _A      =       model.cfg

   _A      =       fs_config.conv_bias
   _A      =       eval(fs_config.conv_feature_layers				)
   _A      =       [x[0] for x in conv_layers]
   _A      =       [x[1] for x in conv_layers]
   _A      =       [x[2] for x in conv_layers]
   _A      =       '''gelu'''
   _A      =       '''layer''' if fs_config.extractor_mode == '''layer_norm''' else '''group'''
   _A      =       0.0
   _A      =       fs_config.activation_fn.name
   _A      =       fs_config.encoder_embed_dim
   _A      =       0.02
   _A      =       fs_config.encoder_ffn_embed_dim
   _A      =       1E-5
   _A      =       fs_config.encoder_layerdrop
   _A      =       fs_config.encoder_attention_heads
   _A      =       fs_config.conv_pos_groups
   _A      =       fs_config.conv_pos
   _A      =       len(_snake_case				)
   _A      =       fs_config.encoder_layers
   _A      =       fs_config.squeeze_factor

   # take care of any params that are overridden by the Wav2VecCtc model
   if is_finetuned:
      _A      =       model.cfg
      _A      =       fs_config.final_dropout
      _A      =       fs_config.layerdrop
   _A      =       fs_config.activation_dropout
   _A      =       fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
   _A      =       fs_config.attention_dropout
   _A      =       fs_config.dropout_input
   _A      =       fs_config.dropout
   _A      =       fs_config.mask_channel_length
   _A      =       fs_config.mask_channel_prob
   _A      =       fs_config.mask_length
   _A      =       fs_config.mask_prob

   _A      =       '''Wav2Vec2FeatureExtractor'''
   _A      =       '''Wav2Vec2CTCTokenizer'''

   return config

@torch.no_grad()
def 	SCREAMING_SNAKE_CASE_							(	_snake_case	:Optional[int]	,   _snake_case	:Union[str, Any]	,   _snake_case	:Optional[Any]=None	,   _snake_case	:Optional[int]=None	,   _snake_case	:Dict=True				)  ->     List[Any]:

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

   if config_path is not None:
      _A      =       SEWConfig.from_pretrained(_snake_case				)
   else:
      _A      =       convert_config(model[0]	,   _snake_case				)
   _A      =       model[0].eval()

   _A      =       True if config.feat_extract_norm == '''layer''' else False
   _A      =       WavaVecaFeatureExtractor(
       feature_size=1	,   sampling_rate=16_000	,   padding_value=0	,   do_normalize=_snake_case	,   return_attention_mask=_snake_case	,   )

   if is_finetuned:
      if dict_path:
         _A      =       Dictionary.load(_snake_case				)

         # important change bos & pad token id since CTC symbol is <pad> and
         # not <s> as in fairseq
         _A      =       target_dict.pad_index
         _A      =       target_dict.bos_index
         _A      =       target_dict.pad_index
         _A      =       target_dict.bos_index
         _A      =       target_dict.eos_index
         _A      =       len(target_dict.symbols				)
         _A      =       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				)
         with open(_snake_case	,   '''w'''	,   encoding='''utf-8'''				) as vocab_handle:
            json.dump(target_dict.indices	,   _snake_case				)
         _A      =       WavaVecaCTCTokenizer(
             _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	,   )
         _A      =       WavaVecaProcessor(feature_extractor=_snake_case	,   tokenizer=_snake_case				)
         processor.save_pretrained(_snake_case				)

      _A      =       SEWForCTC(_snake_case				)
   else:
      _A      =       SEWModel(_snake_case				)
      feature_extractor.save_pretrained(_snake_case				)

   recursively_load_weights(_snake_case	,   _snake_case	,   _snake_case				)

   hf_model.save_pretrained(_snake_case				)


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



