





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase						:   List[str]			  =     {
    'configuration_tapas': ['TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TapasConfig'],
    'tokenization_tapas': ['TapasTokenizer'],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase						:   List[Any]			  =     [
				    'TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'TapasForMaskedLM',
				    'TapasForQuestionAnswering',
				    'TapasForSequenceClassification',
				    'TapasModel',
				    'TapasPreTrainedModel',
				    'load_tf_weights_in_tapas',
				]
try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase						:   Union[str, Any]			  =     [
				    'TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'TFTapasForMaskedLM',
				    'TFTapasForQuestionAnswering',
				    'TFTapasForSequenceClassification',
				    'TFTapasModel',
				    'TFTapasPreTrainedModel',
				]


if TYPE_CHECKING:
				from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
				from .tokenization_tapas import TapasTokenizer

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tapas import (
								    TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TapasForMaskedLM,
								    TapasForQuestionAnswering,
								    TapasForSequenceClassification,
								    TapasModel,
								    TapasPreTrainedModel,
								    load_tf_weights_in_tapas,
								)

				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_tapas import (
								    TF_TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFTapasForMaskedLM,
								    TFTapasForQuestionAnswering,
								    TFTapasForSequenceClassification,
								    TFTapasModel,
								    TFTapasPreTrainedModel,
								)


else:
				import sys

				UpperCAmelCase						:   Optional[int]			  =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)








"""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    :     Optional[Any]   						=		logging.get_logger(__name__)

UpperCAmelCase    :     Optional[int]   						=		{
    '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    :     Dict   						=		[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]
def      lowerCamelCase       (     _UpperCamelCase		:      Optional[int]       ,			_UpperCamelCase		:      Optional[int]       ,			_UpperCamelCase		:      Union[str, Any]       ,			_UpperCamelCase		:      List[str]       ,			_UpperCamelCase		:      int							)		->       List[str]:

							'''simple docstring'''



							for attribute in key.split("""."""							):
														__UpperCAmelCase   :						Tuple      							=					getattr(_UpperCamelCase       ,			_UpperCamelCase							)

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

							'''simple docstring'''



							__UpperCAmelCase   :						Union[str, Any]      							=					[]
							__UpperCAmelCase   :						int      							=					fairseq_model.state_dict()

							__UpperCAmelCase   :						Union[str, Any]      							=					hf_model.feature_extractor

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

							for name, value in fairseq_dict.items():
														__UpperCAmelCase   :						Union[str, Any]      							=					False
														if "conv_layers" in name:
																					load_conv_layer(
																					    _UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			hf_model.config.feat_extract_norm == """group"""       ,			)
																					__UpperCAmelCase   :						int      							=					True
														elif name.split("""."""							)[0] == "proj":
																					__UpperCAmelCase   :						Optional[Any]      							=					fairseq_model.proj
																					__UpperCAmelCase   :						Union[str, Any]      							=					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   :						Any      							=					name.split(_UpperCamelCase							)[0].split("""."""							)[-2]
																																										__UpperCAmelCase   :						Tuple      							=					mapped_key.replace("""*"""       ,			_UpperCamelCase							)
																																			if "weight_g" in name:
																																										__UpperCAmelCase   :						Any      							=					"""weight_g"""
																																			elif "weight_v" in name:
																																										__UpperCAmelCase   :						Union[str, Any]      							=					"""weight_v"""
																																			elif "bias" in name:
																																										__UpperCAmelCase   :						List[Any]      							=					"""bias"""
																																			elif "weight" in name:
																																										__UpperCAmelCase   :						List[Any]      							=					"""weight"""
																																			else:
																																										__UpperCAmelCase   :						Dict      							=					None
																																			set_recursively(_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase       ,			_UpperCamelCase							)
																												continue
														if not is_used:
																					unused_weights.append(_UpperCamelCase							)

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

							return proj_weight
def      lowerCamelCase       (     _UpperCamelCase		:      str       ,			_UpperCamelCase		:      str       ,			_UpperCamelCase		:      List[Any]       ,			_UpperCamelCase		:      List[str]       ,			_UpperCamelCase		:      Dict							)		->       Tuple:

							'''simple docstring'''



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

							'''simple docstring'''



							__UpperCAmelCase							,__UpperCAmelCase   :						Tuple      							=					emb.weight.shape
							__UpperCAmelCase   :						Optional[int]      							=					nn.Linear(_UpperCamelCase       ,			_UpperCamelCase       ,			bias=_UpperCamelCase							)
							__UpperCAmelCase   :						List[str]      							=					emb.weight.data
							return lin_layer
def      lowerCamelCase       (     _UpperCamelCase		:      Any							)		->       Any:

							'''simple docstring'''



							with open(_UpperCamelCase       ,			"""r"""       ,			encoding="""utf-8"""							) as f:
														__UpperCAmelCase   :						Optional[int]      							=					f.readlines()
														__UpperCAmelCase   :						Optional[Any]      							=					[line.split(""" """							)[0] for line in lines]

							__UpperCAmelCase   :						List[Any]      							=					len(_UpperCamelCase							)

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

							vocab_dict.update(dict(zip(_UpperCamelCase       ,			range(4       ,			num_words + 4							)							)							)							)
							return vocab_dict
@torch.no_grad()
def      lowerCamelCase       (     _UpperCamelCase		:      List[Any]       ,			_UpperCamelCase		:      Optional[Any]       ,			_UpperCamelCase		:      Optional[Any]       ,			_UpperCamelCase		:      List[Any]       ,			_UpperCamelCase		:      Dict       ,			_UpperCamelCase		:      Tuple       ,			_UpperCamelCase		:      List[str]       ,			)		->       Any:

							'''simple docstring'''



							__UpperCAmelCase   :						Optional[Any]      							=					WavaVecaConfig.from_pretrained(_UpperCamelCase							)
							__UpperCAmelCase   :						Optional[int]      							=					SpeechaTextaConfig.from_pretrained(
							    _UpperCamelCase       ,			vocab_size=_UpperCamelCase       ,			decoder_layers=_UpperCamelCase       ,			do_stable_layer_norm=_UpperCamelCase							)

							__UpperCAmelCase   :						Any      							=					WavaVecaFeatureExtractor(
							    feature_size=1       ,			sampling_rate=1_6_0_0_0       ,			padding_value=0       ,			do_normalize=_UpperCamelCase       ,			return_attention_mask=_UpperCamelCase       ,			)

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

							# set weights for wav2vec2 encoder
							__UpperCAmelCase   :						int      							=					WavaVecaModel(_UpperCamelCase							)
							__UpperCAmelCase   :						List[str]      							=					recursively_load_weights_wavaveca(model.encoder       ,			_UpperCamelCase							)

							__UpperCAmelCase   :						List[str]      							=					SpeechaTextaForCausalLM(_UpperCamelCase							)
							__UpperCAmelCase							,__UpperCAmelCase   :						Optional[Any]      							=					hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()       ,			strict=_UpperCamelCase							)

							# set output linear layer
							unexpected_keys.remove("""embed_out"""							)
							__UpperCAmelCase   :						List[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   :						str      							=					SpeechEncoderDecoderModel(encoder=_UpperCamelCase       ,			decoder=_UpperCamelCase							)
							__UpperCAmelCase   :						Any      							=					False

							# add projection layer
							__UpperCAmelCase   :						str      							=					nn.Parameter(projection_layer.weight							)
							__UpperCAmelCase   :						str      							=					nn.Parameter(projection_layer.bias							)

							__UpperCAmelCase   :						List[Any]      							=					create_vocab_dict(_UpperCamelCase							)

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

							__UpperCAmelCase   :						List[str]      							=					SpeechaTextaTokenizer(os.path.join(_UpperCamelCase       ,			"""vocab.json"""							)							)
							tokenizer.save_pretrained(_UpperCamelCase							)

							__UpperCAmelCase   :						Dict      							=					hf_wavavec.config.to_dict()
							__UpperCAmelCase   :						int      							=					tokenizer.pad_token_id
							__UpperCAmelCase   :						List[str]      							=					tokenizer.bos_token_id
							__UpperCAmelCase   :						Tuple      							=					tokenizer.eos_token_id
							__UpperCAmelCase   :						int      							=					"""speech_to_text_2"""
							__UpperCAmelCase   :						Optional[int]      							=					"""wav2vec2"""

							__UpperCAmelCase   :						Union[str, Any]      							=					SpeechEncoderDecoderConfig.from_dict(_UpperCamelCase							)

							hf_wavavec.save_pretrained(_UpperCamelCase							)
							feature_extractor.save_pretrained(_UpperCamelCase							)


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_0224, 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,
							)
