





import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


logging.set_verbosity_info()
__A				:    List[Any]       =   logging.get_logger(__name__)

__A				:    str       =   {
    """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""",
}
__A				:    List[Any]       =   [
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__ )  ->				int:







			'''simple docstring'''

			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

			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   		=				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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__ )  ->				str:







			'''simple docstring'''

			UpperCAmelCase   		=				[]
			UpperCAmelCase   		=				fairseq_model.state_dict()

			UpperCAmelCase   		=				hf_model.feature_extractor
			UpperCAmelCase   		=				hf_model.adapter

			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
						elif any(x in name for x in ['''adaptor''', '''w2v_encoder.proj.''', '''w2v_proj_ln.'''] ):
									load_adapter(SCREAMING_SNAKE_CASE__			,			SCREAMING_SNAKE_CASE__			,			SCREAMING_SNAKE_CASE__			,			SCREAMING_SNAKE_CASE__ )
									UpperCAmelCase   		=				True
						else:
									for key, mapped_key in MAPPING.items():
												if key in name or key.split('''w2v_model.''' )[-1] == name.split('''.''' )[0]:
															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:
																		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 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__ )  ->				Any:







			'''simple docstring'''

			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:
									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   		=				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   		=				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   		=				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   		=				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__ )





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__ )  ->				Optional[Any]:







			'''simple docstring'''

			UpperCAmelCase   		=				full_name.split('''adaptor.''' )[-1]
			UpperCAmelCase   		=				name.split('''.''' )

			if items[1].isdigit():
						UpperCAmelCase   		=				int(items[1] )
			else:
						UpperCAmelCase   		=				None

			if "adaptor" not in full_name:
						if "proj_ln" in full_name:
									# has to be layer norm
									if "bias" in name:
												assert (
												    value.shape == adapter.proj_layer_norm.bias.data.shape
												), F"""{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found."""
												UpperCAmelCase   		=				value
												logger.info(F"""Adapter proj layer norm bias was initialized from {full_name}.""" )
									if "weight" in name:
												assert (
												    value.shape == adapter.proj_layer_norm.weight.data.shape
												), F"""{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found."""
												UpperCAmelCase   		=				value
						else:
									# has to be projection layer
									if "bias" in name:
												assert (
												    value.shape == adapter.proj.bias.data.shape
												), F"""{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found."""
												UpperCAmelCase   		=				value
												logger.info(F"""Adapter proj layer bias was initialized from {full_name}.""" )
									if "weight" in name:
												assert (
												    value.shape == adapter.proj.weight.data.shape
												), F"""{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found."""
												UpperCAmelCase   		=				value
												logger.info(F"""Adapter proj layer weight was initialized from {full_name}.""" )
			elif isinstance(SCREAMING_SNAKE_CASE__			,			SCREAMING_SNAKE_CASE__ ):
						if "bias" in name:
									assert (
									    value.shape == adapter.layers[layer_id].conv.bias.data.shape
									), F"""{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found."""
									UpperCAmelCase   		=				value
									logger.info(F"""Adapter layer {layer_id} bias was initialized from {full_name}.""" )
						elif "weight" in name:
									assert (
									    value.shape == adapter.layers[layer_id].conv.weight.data.shape
									), F"""{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found."""
									UpperCAmelCase   		=				value
									logger.info(F"""Adapter layer {layer_id} bias was initialized from {full_name}.""" )
			else:
						unused_weights.append(SCREAMING_SNAKE_CASE__ )





def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__ )  ->				Optional[int]:







			'''simple docstring'''

			UpperCAmelCase   		=				emb.weight.shape
			UpperCAmelCase   		=				nn.Linear(SCREAMING_SNAKE_CASE__			,			SCREAMING_SNAKE_CASE__			,			bias=SCREAMING_SNAKE_CASE__ )
			UpperCAmelCase   		=				emb.weight.data
			return lin_layer





@torch.no_grad()
def 						__SCREAMING_SNAKE_CASE  (				UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			UpperCamelCase__			,			)  ->				Union[str, Any]:







			'''simple docstring'''

			UpperCAmelCase   		=				WavaVecaConfig.from_pretrained(
			    SCREAMING_SNAKE_CASE__			,			add_adapter=SCREAMING_SNAKE_CASE__			,			adapter_stride=SCREAMING_SNAKE_CASE__			,			adapter_kernel_size=SCREAMING_SNAKE_CASE__			,			use_auth_token=SCREAMING_SNAKE_CASE__			,			output_hidden_size=SCREAMING_SNAKE_CASE__			,			)
			UpperCAmelCase   		=				MBartConfig.from_pretrained(SCREAMING_SNAKE_CASE__ )

			# load model
			UpperCAmelCase   		=				fairseq.checkpoint_utils.load_model_ensemble_and_task(
			    [checkpoint_path]			,			arg_overrides={
			        '''config_yaml''': config_yaml_path,
			        '''data''': '''/'''.join(dict_path.split('''/''' )[:-1] ),
			        '''w2v_path''': checkpoint_path,
			        '''load_pretrained_decoder_from''': None,
			    }			,			)
			UpperCAmelCase   		=				model[0].eval()

			# load feature extractor
			UpperCAmelCase   		=				WavaVecaFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE__			,			use_auth_token=SCREAMING_SNAKE_CASE__ )

			# set weights for wav2vec2 encoder
			UpperCAmelCase   		=				WavaVecaModel(SCREAMING_SNAKE_CASE__ )

			recursively_load_weights_wavaveca(model.encoder			,			SCREAMING_SNAKE_CASE__ )

			# load decoder weights
			UpperCAmelCase   		=				MBartForCausalLM(SCREAMING_SNAKE_CASE__ )
			UpperCAmelCase   		=				hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()			,			strict=SCREAMING_SNAKE_CASE__ )
			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   		=				SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE__			,			decoder=SCREAMING_SNAKE_CASE__ )
			UpperCAmelCase   		=				False

			UpperCAmelCase   		=				MBartaaTokenizer(SCREAMING_SNAKE_CASE__ )
			tokenizer.save_pretrained(SCREAMING_SNAKE_CASE__ )

			UpperCAmelCase   		=				hf_wavavec.config.to_dict()
			UpperCAmelCase   		=				tokenizer.pad_token_id
			UpperCAmelCase   		=				tokenizer.bos_token_id
			UpperCAmelCase   		=				tokenizer.eos_token_id
			UpperCAmelCase   		=				"""mbart50"""
			UpperCAmelCase   		=				"""wav2vec2"""

			UpperCAmelCase   		=				tokenizer.eos_token_id
			UpperCAmelCase   		=				25_0004
			UpperCAmelCase   		=				tokenizer.eos_token_id

			UpperCAmelCase   		=				SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE__ )

			hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE__ )
			feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE__ )


if __name__ == "__main__":
				__A				:    Any       =   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_yaml_path", default=None, type=str, help="Path to yaml file of fine-tuned model")
				parser.add_argument(
				    "--encoder_config_path",
				    default="facebook/wav2vec2-xls-r-1b",
				    type=str,
				    help="Path to hf encoder wav2vec2 checkpoint config",
				)
				parser.add_argument(
				    "--decoder_config_path",
				    default="facebook/mbart-large-50-one-to-many-mmt",
				    type=str,
				    help="Path to hf decoder checkpoint config",
				)
				parser.add_argument("--add_adapter", default=True, type=bool, help="whethere to add model adapter layers")
				parser.add_argument("--adapter_stride", default=2, type=int, help="stride of adapter layers")
				parser.add_argument("--adapter_kernel_size", default=3, type=int, help="kernel size of adapter layers")
				parser.add_argument("--encoder_output_dim", default=1_024, type=int, help="encoder output dim")
				parser.add_argument("--start_token_id", default=250_004, type=int, help="`decoder_start_token_id` of model config")

				__A				:    Union[str, Any]       =   parser.parse_args()
				convert_wavaveca_checkpoint(
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.dict_path,
				    args.config_yaml_path,
				    encoder_config_path=args.encoder_config_path,
				    decoder_config_path=args.decoder_config_path,
				    add_adapter=args.add_adapter,
				    adapter_kernel_size=args.adapter_kernel_size,
				    adapter_stride=args.adapter_stride,
				    decoder_start_token_id=args.start_token_id,
				    encoder_output_dim=args.encoder_output_dim,
				)





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


lowercase  :					str						    =							{
    """configuration_funnel""": ["""FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FunnelConfig"""],
    """convert_funnel_original_tf_checkpoint_to_pytorch""": [],
    """tokenization_funnel""": ["""FunnelTokenizer"""],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase  :					Tuple						    =							["""FunnelTokenizerFast"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase  :					List[str]						    =							[
      """FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """FunnelBaseModel""",
      """FunnelForMaskedLM""",
      """FunnelForMultipleChoice""",
      """FunnelForPreTraining""",
      """FunnelForQuestionAnswering""",
      """FunnelForSequenceClassification""",
      """FunnelForTokenClassification""",
      """FunnelModel""",
      """FunnelPreTrainedModel""",
      """load_tf_weights_in_funnel""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase  :					Optional[Any]						    =							[
      """TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFFunnelBaseModel""",
      """TFFunnelForMaskedLM""",
      """TFFunnelForMultipleChoice""",
      """TFFunnelForPreTraining""",
      """TFFunnelForQuestionAnswering""",
      """TFFunnelForSequenceClassification""",
      """TFFunnelForTokenClassification""",
      """TFFunnelModel""",
      """TFFunnelPreTrainedModel""",
  ]


if TYPE_CHECKING:
  from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
  from .tokenization_funnel import FunnelTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_funnel_fast import FunnelTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_funnel import (
        FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        FunnelBaseModel,
        FunnelForMaskedLM,
        FunnelForMultipleChoice,
        FunnelForPreTraining,
        FunnelForQuestionAnswering,
        FunnelForSequenceClassification,
        FunnelForTokenClassification,
        FunnelModel,
        FunnelPreTrainedModel,
        load_tf_weights_in_funnel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_funnel import (
        TF_FUNNEL_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFFunnelBaseModel,
        TFFunnelForMaskedLM,
        TFFunnelForMultipleChoice,
        TFFunnelForPreTraining,
        TFFunnelForQuestionAnswering,
        TFFunnelForSequenceClassification,
        TFFunnelForTokenClassification,
        TFFunnelModel,
        TFFunnelPreTrainedModel,
    )

else:
  import sys

  lowercase  :					Any						    =							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
