





"""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()
_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''',
    '''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''',
}
_lowerCamelCase    	=      [
    '''lm_head''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
    '''projector''',
    '''classifier''',
]




def 						lowerCAmelCase_      (    lowercase_  :				List[Any]						):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE					:     Dict   				=			{}
 with open(_lowercase					, '''r'''						) as file:
  for line_number, line in enumerate(_lowercase						):
   __SCREAMING_SNAKE_CASE					:     Any   				=			line.strip()
   if line:
    __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			line.split()
    __SCREAMING_SNAKE_CASE					:     Dict   				=			line_number
    __SCREAMING_SNAKE_CASE					:     Optional[int]   				=			words[0]
    __SCREAMING_SNAKE_CASE					:     str   				=			value
 return result




def 						lowerCAmelCase_      (    lowercase_  :				List[str]					, lowercase_  :				Optional[Any]					, lowercase_  :				Optional[Any]					, lowercase_  :				str					, lowercase_  :				Any						):

 '''simple docstring'''






 for attribute in key.split('''.'''						):
  __SCREAMING_SNAKE_CASE					:     int   				=			getattr(_lowercase					, _lowercase						)

 __SCREAMING_SNAKE_CASE					:     int   				=			None
 for param_key in PARAM_MAPPING.keys():
  if full_name.endswith(_lowercase						):
   __SCREAMING_SNAKE_CASE					:     List[str]   				=			PARAM_MAPPING[full_name.split('''.'''						)[-1]]
   __SCREAMING_SNAKE_CASE					:     Dict   				=			'''param'''

 if weight_type is not None and weight_type != "param":
  __SCREAMING_SNAKE_CASE					:     int   				=			getattr(_lowercase					, _lowercase						).shape
 elif weight_type is not None and weight_type == "param":
  __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			hf_pointer
  for attribute in hf_param_name.split('''.'''						):
   __SCREAMING_SNAKE_CASE					:     Optional[int]   				=			getattr(_lowercase					, _lowercase						)
  __SCREAMING_SNAKE_CASE					:     Optional[int]   				=			shape_pointer.shape

  # let's reduce dimension
  __SCREAMING_SNAKE_CASE					:     Optional[int]   				=			value[0]
 else:
  __SCREAMING_SNAKE_CASE					:     int   				=			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":
  __SCREAMING_SNAKE_CASE					:     Any   				=			value
 elif weight_type == "weight_g":
  __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			value
 elif weight_type == "weight_v":
  __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			value
 elif weight_type == "bias":
  __SCREAMING_SNAKE_CASE					:     Optional[Any]   				=			value
 elif weight_type == "param":
  for attribute in hf_param_name.split('''.'''						):
   __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			getattr(_lowercase					, _lowercase						)
  __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			value
 else:
  __SCREAMING_SNAKE_CASE					:     List[str]   				=			value

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




def 						lowerCAmelCase_      (    lowercase_  :				Tuple					, lowercase_  :				str					, lowercase_  :				Optional[int]					, lowercase_  :				Tuple					, lowercase_  :				Tuple						):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE					:     Tuple   				=			None
 for param_key in PARAM_MAPPING.keys():
  if full_name.endswith(_lowercase						):
   __SCREAMING_SNAKE_CASE					:     int   				=			PARAM_MAPPING[full_name.split('''.'''						)[-1]]
   __SCREAMING_SNAKE_CASE					:     Dict   				=			'''param'''

 if weight_type is not None and weight_type != "param":
  __SCREAMING_SNAKE_CASE					:     Dict   				=			'''.'''.join([key, weight_type]						)
 elif weight_type is not None and weight_type == "param":
  __SCREAMING_SNAKE_CASE					:     str   				=			'''.'''.join([key, hf_param_name]						)
 else:
  __SCREAMING_SNAKE_CASE					:     List[Any]   				=			key

 __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			value if '''lm_head''' in full_key else value[0]


_lowerCamelCase    	=      {
    '''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 						lowerCAmelCase_      (    lowercase_  :				Tuple					, lowercase_  :				int					, lowercase_  :				int=None					, lowercase_  :				Union[str, Any]=None						):

 '''simple docstring'''






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




def 						lowerCAmelCase_      (    lowercase_  :				Optional[Any]					, lowercase_  :				Optional[Any]					, lowercase_  :				Optional[int]						):

 '''simple docstring'''






 __SCREAMING_SNAKE_CASE					:     int   				=			[]
 __SCREAMING_SNAKE_CASE					:     Any   				=			fairseq_model.state_dict()

 __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			hf_model.wavaveca.feature_extractor

 for name, value in fairseq_dict.items():
  __SCREAMING_SNAKE_CASE					:     List[Any]   				=			False
  if "conv_layers" in name:
   load_conv_layer(
       _lowercase					, _lowercase					, _lowercase					, _lowercase					, hf_model.config.feat_extract_norm == '''group'''					, )
   __SCREAMING_SNAKE_CASE					:     List[str]   				=			True
  else:
   __SCREAMING_SNAKE_CASE					:     Optional[int]   				=			load_wavaveca_layer(_lowercase					, _lowercase					, _lowercase						)
  if not is_used:
   unused_weights.append(_lowercase						)

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




def 						lowerCAmelCase_      (    lowercase_  :				Optional[Any]					, lowercase_  :				Optional[Any]					, lowercase_  :				str					, lowercase_  :				str					, lowercase_  :				Optional[int]						):

 '''simple docstring'''






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




@torch.no_grad()
def 						lowerCAmelCase_      (    lowercase_  :				List[str]					, lowercase_  :				List[Any]					, lowercase_  :				Optional[int]=None					, lowercase_  :				str=None					, lowercase_  :				List[Any]=True					, lowercase_  :				Tuple=False						):

 '''simple docstring'''






 if config_path is not None:
  __SCREAMING_SNAKE_CASE					:     Tuple   				=			WavaVecaConfig.from_pretrained(_lowercase						)
 else:
  __SCREAMING_SNAKE_CASE					:     Dict   				=			WavaVecaConfig()

 if is_seq_class:
  __SCREAMING_SNAKE_CASE					:     str   				=			read_txt_into_dict(_lowercase						)
  __SCREAMING_SNAKE_CASE					:     str   				=			idalabel
  __SCREAMING_SNAKE_CASE					:     int   				=			WavaVecaForSequenceClassification(_lowercase						)
  __SCREAMING_SNAKE_CASE					:     Optional[Any]   				=			WavaVecaFeatureExtractor(
      feature_size=1					, sampling_rate=1_6000					, padding_value=0					, do_normalize=_lowercase					, return_attention_mask=_lowercase					, )
  feature_extractor.save_pretrained(_lowercase						)

 elif is_finetuned:
  if dict_path:
   __SCREAMING_SNAKE_CASE					:     str   				=			Dictionary.load(_lowercase						)

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

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

  __SCREAMING_SNAKE_CASE					:     Dict   				=			WavaVecaForCTC(_lowercase						)
 else:
  __SCREAMING_SNAKE_CASE					:     str   				=			WavaVecaForPreTraining(_lowercase						)

 if is_finetuned or is_seq_class:
  __SCREAMING_SNAKE_CASE				, __SCREAMING_SNAKE_CASE				, __SCREAMING_SNAKE_CASE					:     Union[str, Any]   				=			fairseq.checkpoint_utils.load_model_ensemble_and_task(
      [checkpoint_path]					, arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''						)[:-1]						)}						)
 else:
  __SCREAMING_SNAKE_CASE					:     str   				=			argparse.Namespace(task='''audio_pretraining'''						)
  __SCREAMING_SNAKE_CASE					:     Optional[Any]   				=			fairseq.tasks.setup_task(_lowercase						)

  __SCREAMING_SNAKE_CASE				, __SCREAMING_SNAKE_CASE				, __SCREAMING_SNAKE_CASE					:     int   				=			fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]					, task=_lowercase						)

 __SCREAMING_SNAKE_CASE					:     List[Any]   				=			model[0].eval()

 recursively_load_weights(_lowercase					, _lowercase					, not is_finetuned						)

 hf_wavavec.save_pretrained(_lowercase						)


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'''
    )
    parser.add_argument(
        '''--is_seq_class''',
        action='''store_true''',
        help='''Whether the model to convert is a fine-tuned sequence classification model or not''',
    )
    _lowerCamelCase    	=      parser.parse_args()

    _lowerCamelCase    	=      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,
    )



# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


_snake_case      = {
    '''configuration_efficientnet''': [
        '''EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''EfficientNetConfig''',
        '''EfficientNetOnnxConfig''',
    ]
}

try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case      = ['''EfficientNetImageProcessor''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _snake_case      = [
     '''EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''EfficientNetForImageClassification''',
     '''EfficientNetModel''',
     '''EfficientNetPreTrainedModel''',
 ]

if TYPE_CHECKING:
 from .configuration_efficientnet import (
     EFFICIENTNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
     EfficientNetConfig,
     EfficientNetOnnxConfig,
 )

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .image_processing_efficientnet import EfficientNetImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_efficientnet import (
      EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST,
      EfficientNetForImageClassification,
      EfficientNetModel,
      EfficientNetPreTrainedModel,
  )

else:
 import sys

 _snake_case      = _LazyModule(__name__, globals()['''__file__'''], _import_structure)


