



import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    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',
    'quantizer.weight_proj': 'quantizer.weight_proj',
    'quantizer.vars': 'quantizer.codevectors',
    'project_q': 'project_q',
    'final_proj': 'project_hid',
    'w2v_encoder.proj': 'ctc_proj',
    'mask_emb': 'masked_spec_embed',
}
UpperCAmelCase_									=							[
    'ctc_proj',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	Dict     ,				_SCREAMING_SNAKE_CASE   :	Union[str, Any]     ,				_SCREAMING_SNAKE_CASE   :	int     ,				_SCREAMING_SNAKE_CASE   :	int     ,				_SCREAMING_SNAKE_CASE   :	Optional[Any]     ,				_SCREAMING_SNAKE_CASE   :	Optional[int]				)->Dict:

       for attribute in key.split('''.'''				):
              if is_finetuned:
                     if attribute in ["quantizer", "project_q", "project_hid"]:
                            # those layers are only relevant for pretraining and should be dropped
                            return

                     if attribute == "ctc_proj":
                            # we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
                            _lowerCAmelCase 			=  '''lm_head'''

              _lowerCAmelCase 			=  getattr(lowerCamelCase_     ,				lowerCamelCase_				)

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

       logger.info(f'''{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.'''				)
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	List[Any]     ,				_SCREAMING_SNAKE_CASE   :	Dict     ,				_SCREAMING_SNAKE_CASE   :	Optional[int]				)->List[Any]:
       _lowerCAmelCase 			=  []
       _lowerCAmelCase 			=  fairseq_model.state_dict()

       _lowerCAmelCase 			=  hf_model.unispeech.feature_extractor

       for name, value in fairseq_dict.items():
              _lowerCAmelCase 			=  False
              if "conv_layers" in name:
                     load_conv_layer(
                         lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				hf_model.config.feat_extract_norm == '''group'''     ,				)
                     _lowerCAmelCase 			=  True
              else:
                     for key, mapped_key in MAPPING.items():
                            _lowerCAmelCase 			=  '''unispeech.''' + 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]:
                                   _lowerCAmelCase 			=  True
                                   if "*" in mapped_key:
                                          _lowerCAmelCase 			=  name.split(lowerCamelCase_				)[0].split('''.'''				)[-2]
                                          _lowerCAmelCase 			=  mapped_key.replace('''*'''     ,				lowerCamelCase_				)
                                   if "weight_g" in name:
                                          _lowerCAmelCase 			=  '''weight_g'''
                                   elif "weight_v" in name:
                                          _lowerCAmelCase 			=  '''weight_v'''
                                   elif "bias" in name:
                                          _lowerCAmelCase 			=  '''bias'''
                                   elif "weight" in name:
                                          # TODO: don't match quantizer.weight_proj
                                          _lowerCAmelCase 			=  '''weight'''
                                   else:
                                          _lowerCAmelCase 			=  None
                                   set_recursively(lowerCamelCase_     ,				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   :	Tuple     ,				_SCREAMING_SNAKE_CASE   :	Dict     ,				_SCREAMING_SNAKE_CASE   :	List[Any]     ,				_SCREAMING_SNAKE_CASE   :	Tuple     ,				_SCREAMING_SNAKE_CASE   :	Optional[int]				)->Union[str, Any]:
       _lowerCAmelCase 			=  full_name.split('''conv_layers.'''				)[-1]
       _lowerCAmelCase 			=  name.split('''.'''				)
       _lowerCAmelCase 			=  int(items[0]				)
       _lowerCAmelCase 			=  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.'''
                     )
                     _lowerCAmelCase 			=  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.'''
                     )
                     _lowerCAmelCase 			=  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."
                     )
                     _lowerCAmelCase 			=  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.'''
                     )
                     _lowerCAmelCase 			=  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   :	List[str]     ,				_SCREAMING_SNAKE_CASE   :	Union[str, Any]     ,				_SCREAMING_SNAKE_CASE   :	str=None     ,				_SCREAMING_SNAKE_CASE   :	List[Any]=None     ,				_SCREAMING_SNAKE_CASE   :	List[str]=True				)->Optional[int]:

       if config_path is not None:
              _lowerCAmelCase 			=  UniSpeechConfig.from_pretrained(lowerCamelCase_				)
       else:
              _lowerCAmelCase 			=  UniSpeechConfig()

       if is_finetuned:
              if dict_path:
                     _lowerCAmelCase 			=  Dictionary.load_from_json(lowerCamelCase_				)

                     # important change bos & pad token id since CTC symbol is <pad> and
                     # not <s> as in fairseq
                     _lowerCAmelCase 			=  target_dict.pad_index
                     _lowerCAmelCase 			=  target_dict.bos_index
                     _lowerCAmelCase 			=  target_dict.eos_index
                     _lowerCAmelCase 			=  len(target_dict.symbols				)
                     _lowerCAmelCase 			=  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_				)
                     _lowerCAmelCase 			=  target_dict.indices

                     # fairseq has the <pad> and <s> switched
                     _lowerCAmelCase 			=  4_2
                     _lowerCAmelCase 			=  4_3
                     with open(lowerCamelCase_     ,				'''w'''     ,				encoding='''utf-8'''				) as vocab_handle:
                            json.dump(lowerCamelCase_     ,				lowerCamelCase_				)
                     _lowerCAmelCase 			=  WavaVecaPhonemeCTCTokenizer(
                         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_     ,				)
                     _lowerCAmelCase 			=  True if config.feat_extract_norm == '''layer''' else False
                     _lowerCAmelCase 			=  WavaVecaFeatureExtractor(
                         feature_size=1     ,				sampling_rate=1_6_0_0_0     ,				padding_value=0     ,				do_normalize=lowerCamelCase_     ,				return_attention_mask=lowerCamelCase_     ,				)
                     _lowerCAmelCase 			=  WavaVecaProcessor(feature_extractor=lowerCamelCase_     ,				tokenizer=lowerCamelCase_				)
                     processor.save_pretrained(lowerCamelCase_				)

              _lowerCAmelCase 			=  UniSpeechForCTC(lowerCamelCase_				)
       else:
              _lowerCAmelCase 			=  UniSpeechForPreTraining(lowerCamelCase_				)

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

       _lowerCAmelCase 			=  model[0].eval()

       recursively_load_weights(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_				)

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




import argparse
import pathlib

import fairseq
import torch
from fairseq.models.roberta import RobertaModel as FairseqRobertaModel
from fairseq.modules import TransformerSentenceEncoderLayer
from packaging import version

from transformers import XLMRobertaConfig, XLMRobertaXLForMaskedLM, XLMRobertaXLForSequenceClassification
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.models.roberta.modeling_roberta import RobertaAttention
from transformers.utils import logging


if version.parse(fairseq.__version__) < version.parse("1.0.0a"):
    raise Exception("requires fairseq >= 1.0.0a")

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

UpperCAmelCase_									=							"Hello world! cécé herlolip"
def 							UpperCAmelCase__		(				_SCREAMING_SNAKE_CASE   :	str     ,				_SCREAMING_SNAKE_CASE   :	str     ,				_SCREAMING_SNAKE_CASE   :	bool				)->List[Any]:
      _lowerCAmelCase 			=  FairseqRobertaModel.from_pretrained(_SCREAMING_SNAKE_CASE				)
      roberta.eval()  # disable dropout
      _lowerCAmelCase 			=  roberta.model.encoder.sentence_encoder
      _lowerCAmelCase 			=  XLMRobertaConfig(
          vocab_size=roberta_sent_encoder.embed_tokens.num_embeddings     ,				hidden_size=roberta.cfg.model.encoder_embed_dim     ,				num_hidden_layers=roberta.cfg.model.encoder_layers     ,				num_attention_heads=roberta.cfg.model.encoder_attention_heads     ,				intermediate_size=roberta.cfg.model.encoder_ffn_embed_dim     ,				max_position_embeddings=5_1_4     ,				type_vocab_size=1     ,				layer_norm_eps=1e-5     ,				)
      if classification_head:
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''].out_proj.weight.shape[0]

      print('''Our RoBERTa config:'''     ,				_SCREAMING_SNAKE_CASE				)

      _lowerCAmelCase 			=  XLMRobertaXLForSequenceClassification(_SCREAMING_SNAKE_CASE				) if classification_head else XLMRobertaXLForMaskedLM(_SCREAMING_SNAKE_CASE				)
      model.eval()

      # Now let's copy all the weights.
      # Embeddings
      _lowerCAmelCase 			=  roberta_sent_encoder.embed_tokens.weight
      _lowerCAmelCase 			=  roberta_sent_encoder.embed_positions.weight
      _lowerCAmelCase 			=  torch.zeros_like(
          model.roberta.embeddings.token_type_embeddings.weight				)  # just zero them out b/c RoBERTa doesn't use them.

      _lowerCAmelCase 			=  roberta_sent_encoder.layer_norm.weight
      _lowerCAmelCase 			=  roberta_sent_encoder.layer_norm.bias

      for i in range(config.num_hidden_layers				):
            # Encoder: start of layer
            _lowerCAmelCase 			=  model.roberta.encoder.layer[i]
            _lowerCAmelCase 			=  roberta_sent_encoder.layers[i]

            _lowerCAmelCase 			=  layer.attention
            _lowerCAmelCase 			=  roberta_layer.self_attn_layer_norm.weight
            _lowerCAmelCase 			=  roberta_layer.self_attn_layer_norm.bias

            # self attention
            _lowerCAmelCase 			=  layer.attention.self
            assert (
                roberta_layer.self_attn.k_proj.weight.data.shape
                == roberta_layer.self_attn.q_proj.weight.data.shape
                == roberta_layer.self_attn.v_proj.weight.data.shape
                == torch.Size((config.hidden_size, config.hidden_size)				)
            )

            _lowerCAmelCase 			=  roberta_layer.self_attn.q_proj.weight
            _lowerCAmelCase 			=  roberta_layer.self_attn.q_proj.bias
            _lowerCAmelCase 			=  roberta_layer.self_attn.k_proj.weight
            _lowerCAmelCase 			=  roberta_layer.self_attn.k_proj.bias
            _lowerCAmelCase 			=  roberta_layer.self_attn.v_proj.weight
            _lowerCAmelCase 			=  roberta_layer.self_attn.v_proj.bias

            # self-attention output
            _lowerCAmelCase 			=  layer.attention.output
            assert self_output.dense.weight.shape == roberta_layer.self_attn.out_proj.weight.shape
            _lowerCAmelCase 			=  roberta_layer.self_attn.out_proj.weight
            _lowerCAmelCase 			=  roberta_layer.self_attn.out_proj.bias

            # this one is final layer norm
            _lowerCAmelCase 			=  roberta_layer.final_layer_norm.weight
            _lowerCAmelCase 			=  roberta_layer.final_layer_norm.bias

            # intermediate
            _lowerCAmelCase 			=  layer.intermediate
            assert intermediate.dense.weight.shape == roberta_layer.fca.weight.shape
            _lowerCAmelCase 			=  roberta_layer.fca.weight
            _lowerCAmelCase 			=  roberta_layer.fca.bias

            # output
            _lowerCAmelCase 			=  layer.output
            assert bert_output.dense.weight.shape == roberta_layer.fca.weight.shape
            _lowerCAmelCase 			=  roberta_layer.fca.weight
            _lowerCAmelCase 			=  roberta_layer.fca.bias
            # end of layer

      if classification_head:
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''].dense.weight
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''].dense.bias
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''].out_proj.weight
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''].out_proj.bias
      else:
            # LM Head
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.dense.weight
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.dense.bias
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.layer_norm.weight
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.layer_norm.bias
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.weight
            _lowerCAmelCase 			=  roberta.model.encoder.lm_head.bias

      # Let's check that we get the same results.
      _lowerCAmelCase 			=  roberta.encode(_SCREAMING_SNAKE_CASE				).unsqueeze(0				)  # batch of size 1

      _lowerCAmelCase 			=  model(_SCREAMING_SNAKE_CASE				)[0]
      if classification_head:
            _lowerCAmelCase 			=  roberta.model.classification_heads['''mnli'''](roberta.extract_features(_SCREAMING_SNAKE_CASE				)				)
      else:
            _lowerCAmelCase 			=  roberta.model(_SCREAMING_SNAKE_CASE				)[0]
      print(our_output.shape     ,				their_output.shape				)
      _lowerCAmelCase 			=  torch.max(torch.abs(our_output - their_output				)				).item()
      print(f'''max_absolute_diff = {max_absolute_diff}'''				)  # ~ 1e-7
      _lowerCAmelCase 			=  torch.allclose(_SCREAMING_SNAKE_CASE     ,				_SCREAMING_SNAKE_CASE     ,				atol=1e-3				)
      print('''Do both models output the same tensors?'''     ,				'''🔥''' if success else '''💩'''				)
      if not success:
            raise Exception('''Something went wRoNg'''				)

      pathlib.Path(_SCREAMING_SNAKE_CASE				).mkdir(parents=_SCREAMING_SNAKE_CASE     ,				exist_ok=_SCREAMING_SNAKE_CASE				)
      print(f'''Saving model to {pytorch_dump_folder_path}'''				)
      model.save_pretrained(_SCREAMING_SNAKE_CASE				)


if __name__ == "__main__":
    UpperCAmelCase_									=							argparse.ArgumentParser()
    # Required parameters
    parser.add_argument(
        "--roberta_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump."
    )
    parser.add_argument(
        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
    )
    parser.add_argument(
        "--classification_head", action="store_true", help="Whether to convert a final classification head."
    )
    UpperCAmelCase_									=							parser.parse_args()
    convert_xlm_roberta_xl_checkpoint_to_pytorch(
        args.roberta_checkpoint_path, args.pytorch_dump_folder_path, args.classification_head
    )