


'''simple docstring'''



import argparse
import json
import os

import fairseq
import torch
from torch import nn

from transformers import (
    SpeechaTextaConfig,
    SpeechaTextaForCausalLM,
    SpeechaTextaTokenizer,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
__UpperCAmelCase    :Dict				       =     logging.get_logger(__name__)

__UpperCAmelCase    :Tuple				       =     {
    "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": "lm_head",
    "mask_emb": "masked_spec_embed",
}
__UpperCAmelCase    :Tuple				       =     [
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]




def 			_a			(		_lowercase      :			List[Any]    ,      _lowercase      :			Tuple    ,      _lowercase      :			Union[str, Any]    ,      _lowercase      :			Tuple    ,      _lowercase      :			Tuple						):





     '''simple docstring'''






     for attribute in key.split('''.'''						):
          __UpperCAmelCase			:				Optional[int]													=    getattr(_lowercase    ,      _lowercase						)

     if weight_type is not None:
          __UpperCAmelCase			:				int													=    getattr(_lowercase    ,      _lowercase						).shape
     else:
          __UpperCAmelCase			:				List[str]													=    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			:				Dict													=    value
     elif weight_type == "weight_v":
          __UpperCAmelCase			:				List[str]													=    value
     elif weight_type == "bias":
          __UpperCAmelCase			:				List[str]													=    value
     else:
          __UpperCAmelCase			:				List[str]													=    value

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




def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			Union[str, Any]						):





     '''simple docstring'''






     __UpperCAmelCase			:				List[str]													=    []
     __UpperCAmelCase			:				Tuple													=    fairseq_model.state_dict()

     __UpperCAmelCase			:				Any													=    hf_model.feature_extractor

     # if encoder has different dim to decoder -> use proj_weight
     __UpperCAmelCase			:				Any													=    None

     for name, value in fairseq_dict.items():
          __UpperCAmelCase			:				int													=    False
          if "conv_layers" in name:
               load_conv_layer(
                   _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase    ,      hf_model.config.feat_extract_norm == '''group'''    ,      )
               __UpperCAmelCase			:				str													=    True
          elif name.split('''.'''						)[0] == "proj":
               __UpperCAmelCase			:				Dict													=    fairseq_model.proj
               __UpperCAmelCase			:				str													=    True
          else:
               for key, mapped_key in MAPPING.items():
                    if key in name or key.split('''w2v_model.'''						)[-1] == name.split('''.'''						)[0]:
                         __UpperCAmelCase			:				Dict													=    True
                         if "*" in mapped_key:
                              __UpperCAmelCase			:				Optional[int]													=    name.split(_lowercase						)[0].split('''.'''						)[-2]
                              __UpperCAmelCase			:				List[str]													=    mapped_key.replace('''*'''    ,      _lowercase						)
                         if "weight_g" in name:
                              __UpperCAmelCase			:				int													=    '''weight_g'''
                         elif "weight_v" in name:
                              __UpperCAmelCase			:				Dict													=    '''weight_v'''
                         elif "bias" in name:
                              __UpperCAmelCase			:				Tuple													=    '''bias'''
                         elif "weight" in name:
                              __UpperCAmelCase			:				List[Any]													=    '''weight'''
                         else:
                              __UpperCAmelCase			:				List[str]													=    None
                         set_recursively(_lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase    ,      _lowercase						)
                    continue
          if not is_used:
               unused_weights.append(_lowercase						)

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

     return proj_weight




def 			_a			(		_lowercase      :			str    ,      _lowercase      :			Any    ,      _lowercase      :			Optional[Any]    ,      _lowercase      :			Union[str, Any]    ,      _lowercase      :			str						):





     '''simple docstring'''






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




def 			_a			(		_lowercase      :			Any						):





     '''simple docstring'''






     __UpperCAmelCase							,     __UpperCAmelCase			:				List[Any]													=    emb.weight.shape
     __UpperCAmelCase			:				List[str]													=    nn.Linear(_lowercase    ,      _lowercase    ,      bias=_lowercase						)
     __UpperCAmelCase			:				int													=    emb.weight.data
     return lin_layer




def 			_a			(		_lowercase      :			Any						):





     '''simple docstring'''






     with open(_lowercase    ,      '''r'''    ,      encoding='''utf-8'''						) as f:
          __UpperCAmelCase			:				List[Any]													=    f.readlines()
          __UpperCAmelCase			:				Any													=    [line.split(''' '''						)[0] for line in lines]

     __UpperCAmelCase			:				Optional[int]													=    len(_lowercase						)

     __UpperCAmelCase			:				int													=    {
         '''<s>''': 0,
         '''<pad>''': 1,
         '''</s>''': 2,
         '''<unk>''': 3,
     }

     vocab_dict.update(dict(zip(_lowercase    ,      range(4    ,      num_words + 4						)						)						)						)
     return vocab_dict




@torch.no_grad()
def 			_a			(		_lowercase      :			Optional[Any]    ,      _lowercase      :			Optional[Any]    ,      _lowercase      :			Tuple    ,      _lowercase      :			Any    ,      _lowercase      :			List[str]    ,      _lowercase      :			Any    ,      _lowercase      :			str    ,      ):





     '''simple docstring'''






     __UpperCAmelCase			:				Any													=    WavaVecaConfig.from_pretrained(_lowercase						)
     __UpperCAmelCase			:				List[str]													=    SpeechaTextaConfig.from_pretrained(
         _lowercase    ,      vocab_size=_lowercase    ,      decoder_layers=_lowercase    ,      do_stable_layer_norm=_lowercase						)

     __UpperCAmelCase			:				Optional[int]													=    WavaVecaFeatureExtractor(
         feature_size=1    ,      sampling_rate=16000    ,      padding_value=0    ,      do_normalize=_lowercase    ,      return_attention_mask=_lowercase    ,      )

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

     # set weights for wav2vec2 encoder
     __UpperCAmelCase			:				Tuple													=    WavaVecaModel(_lowercase						)
     __UpperCAmelCase			:				Optional[Any]													=    recursively_load_weights_wavaveca(model.encoder    ,      _lowercase						)

     __UpperCAmelCase			:				str													=    SpeechaTextaForCausalLM(_lowercase						)
     __UpperCAmelCase							,     __UpperCAmelCase			:				Dict													=    hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()    ,      strict=_lowercase						)

     # set output linear layer
     unexpected_keys.remove('''embed_out'''						)
     __UpperCAmelCase			:				str													=    nn.Parameter(model.decoder.embed_out.detach()						)

     # layer norm is init to identity matrix so leaving it is fine
     logger.warning(F'The following keys are missing when loading the decoder weights: {missing_keys}'						)
     logger.warning(F'The following keys are unexpected when loading the decoder weights: {unexpected_keys}'						)

     __UpperCAmelCase			:				Any													=    SpeechEncoderDecoderModel(encoder=_lowercase    ,      decoder=_lowercase						)
     __UpperCAmelCase			:				List[Any]													=    False

     # add projection layer
     __UpperCAmelCase			:				List[Any]													=    nn.Parameter(projection_layer.weight						)
     __UpperCAmelCase			:				Tuple													=    nn.Parameter(projection_layer.bias						)

     __UpperCAmelCase			:				str													=    create_vocab_dict(_lowercase						)

     with open(os.path.join(_lowercase    ,      '''vocab.json'''						)    ,      '''w'''						) as fp:
          json.dump(_lowercase    ,      _lowercase						)

     __UpperCAmelCase			:				Tuple													=    SpeechaTextaTokenizer(os.path.join(_lowercase    ,      '''vocab.json'''						)						)
     tokenizer.save_pretrained(_lowercase						)

     __UpperCAmelCase			:				Optional[Any]													=    hf_wavavec.config.to_dict()
     __UpperCAmelCase			:				List[str]													=    tokenizer.pad_token_id
     __UpperCAmelCase			:				Any													=    tokenizer.bos_token_id
     __UpperCAmelCase			:				str													=    tokenizer.eos_token_id
     __UpperCAmelCase			:				Any													=    '''speech_to_text_2'''
     __UpperCAmelCase			:				Any													=    '''wav2vec2'''

     __UpperCAmelCase			:				List[Any]													=    SpeechEncoderDecoderConfig.from_dict(_lowercase						)

     hf_wavavec.save_pretrained(_lowercase						)
     feature_extractor.save_pretrained(_lowercase						)


if __name__ == "__main__":
      __UpperCAmelCase    :Tuple				       =     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(
          "--encoder_config_path",
          default="facebook/wav2vec2-large-lv60",
          type=str,
          help="Path to hf encoder wav2vec2 checkpoint config",
      )
      parser.add_argument(
          "--decoder_config_path",
          default="facebook/s2t-small-mustc-en-fr-st",
          type=str,
          help="Path to hf decoder s2t checkpoint config",
      )
      parser.add_argument("--vocab_size", default=1_0_2_2_4, type=int, help="Vocab size of decoder")
      parser.add_argument("--num_decoder_layers", default=7, type=int, help="Number of decoder layers")

      __UpperCAmelCase    :int				       =     parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path,
          args.pytorch_dump_folder_path,
          args.dict_path,
          encoder_config_path=args.encoder_config_path,
          decoder_config_path=args.decoder_config_path,
          vocab_size=args.vocab_size,
          num_decoder_layers=args.num_decoder_layers,
      )



'''simple docstring'''



from collections.abc import Iterator, MutableMapping
from dataclasses import dataclass
from typing import Generic, TypeVar

__UpperCAmelCase    :Optional[int]				       =     TypeVar("KEY")
__UpperCAmelCase    :Tuple				       =     TypeVar("VAL")



@dataclass(frozen=_a			,						slots=_a )
class 			a       (	Generic[KEY, VAL] ):


     """simple docstring"""






     SCREAMING_SNAKE_CASE :				KEY
     SCREAMING_SNAKE_CASE :				VAL



class 			a       (	_Item ):


     """simple docstring"""






     def __init__(       self  :	Optional[int]    )		->							None:
          super().__init__(snake_case							,				snake_case    )






     def __bool__(       self  :	Dict    )		->							bool:
          return False



__UpperCAmelCase    :Optional[Any]				       =     _DeletedItem()



class 			a       (	MutableMapping[KEY, VAL] ):


     """simple docstring"""






     def __init__(       self  :	str							,				snake_case  :	int = 8							,				snake_case  :	float = 0.75    )		->							None:
          __UpperCAmelCase			:				Dict													=    initial_block_size
          __UpperCAmelCase			:				list[_Item | None]													=    [None] * initial_block_size
          assert 0.0 < capacity_factor < 1.0
          __UpperCAmelCase			:				str													=    capacity_factor
          __UpperCAmelCase			:				Optional[int]													=    0

     def 				lowerCamelCase__				(       self  :	Optional[Any]							,				snake_case  :	KEY    )		->							int:
          return hash(snake_case    ) % len(self._buckets    )

     def 				lowerCamelCase__				(       self  :	Union[str, Any]							,				snake_case  :	int    )		->							int:
          return (ind + 1) % len(self._buckets    )

     def 				lowerCamelCase__				(       self  :	List[str]							,				snake_case  :	int							,				snake_case  :	KEY							,				snake_case  :	VAL    )		->							bool:
          __UpperCAmelCase			:				int													=    self._buckets[ind]
          if not stored:
               __UpperCAmelCase			:				List[str]													=    _Item(snake_case							,				snake_case    )
               self._len += 1
               return True
          elif stored.key == key:
               __UpperCAmelCase			:				int													=    _Item(snake_case							,				snake_case    )
               return True
          else:
               return False

     def 				lowerCamelCase__				(       self  :	str    )		->							bool:
          __UpperCAmelCase			:				Union[str, Any]													=    len(self._buckets    ) * self._capacity_factor
          return len(self    ) >= int(snake_case    )

     def 				lowerCamelCase__				(       self  :	List[str]    )		->							bool:
          if len(self._buckets    ) <= self._initial_block_size:
               return False
          __UpperCAmelCase			:				List[str]													=    len(self._buckets    ) * self._capacity_factor / 2
          return len(self    ) < limit

     def 				lowerCamelCase__				(       self  :	Optional[int]							,				snake_case  :	int    )		->							None:
          __UpperCAmelCase			:				Union[str, Any]													=    self._buckets
          __UpperCAmelCase			:				Union[str, Any]													=    [None] * new_size
          __UpperCAmelCase			:				List[str]													=    0
          for item in old_buckets:
               if item:
                    self._add_item(item.key							,				item.val    )

     def 				lowerCamelCase__				(       self  :	Optional[int]    )		->							None:
          self._resize(len(self._buckets    ) * 2    )

     def 				lowerCamelCase__				(       self  :	List[str]    )		->							None:
          self._resize(len(self._buckets    ) // 2    )

     def 				lowerCamelCase__				(       self  :	Union[str, Any]							,				snake_case  :	KEY    )		->							Iterator[int]:
          __UpperCAmelCase			:				Tuple													=    self._get_bucket_index(snake_case    )
          for _ in range(len(self._buckets    )    ):
               yield ind
               __UpperCAmelCase			:				Dict													=    self._get_next_ind(snake_case    )

     def 				lowerCamelCase__				(       self  :	str							,				snake_case  :	KEY							,				snake_case  :	VAL    )		->							None:
          for ind in self._iterate_buckets(snake_case    ):
               if self._try_set(snake_case							,				snake_case							,				snake_case    ):
                    break

     def __setitem__(       self  :	List[str]							,				snake_case  :	KEY							,				snake_case  :	VAL    )		->							None:
          if self._is_full():
               self._size_up()

          self._add_item(snake_case							,				snake_case    )

     def __delitem__(       self  :	Any							,				snake_case  :	KEY    )		->							None:
          for ind in self._iterate_buckets(snake_case    ):
               __UpperCAmelCase			:				List[Any]													=    self._buckets[ind]
               if item is None:
                    raise KeyError(snake_case    )
               if item is _deleted:
                    continue
               if item.key == key:
                    __UpperCAmelCase			:				Optional[Any]													=    _deleted
                    self._len -= 1
                    break
          if self._is_sparse():
               self._size_down()

     def __getitem__(       self  :	List[str]							,				snake_case  :	KEY    )		->							VAL:
          for ind in self._iterate_buckets(snake_case    ):
               __UpperCAmelCase			:				Optional[int]													=    self._buckets[ind]
               if item is None:
                    break
               if item is _deleted:
                    continue
               if item.key == key:
                    return item.val
          raise KeyError(snake_case    )

     def __len__(       self  :	List[Any]    )		->							int:
          return self._len

     def __iter__(       self  :	str    )		->							Iterator[KEY]:
          yield from (item.key for item in self._buckets if item)






     def __repr__(       self  :	Optional[Any]    )		->							str:
          __UpperCAmelCase			:				str													=    ''' ,'''.join(
              f'{item.key}: {item.val}' for item in self._buckets if item    )
          return f'HashMap({val_string})'