




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

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





def 					lowercase_ (    __snake_case     :							Optional[int]	,	__snake_case     :							Union[str, Any]	,	__snake_case     :							Optional[int]	,	__snake_case     :							int	,	__snake_case     :							Union[str, Any]      )     -> str:





				'''simple docstring'''


				for attribute in key.split("."      ):
								snake_case__      :Union[str, Any]         =   getattr(__snake_case	,	__snake_case      )

				if weight_type is not None:
								snake_case__      :Optional[Any]         =   getattr(__snake_case	,	__snake_case      ).shape
				else:
								snake_case__      :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":
								snake_case__      :Optional[int]         =   value
				elif weight_type == "weight_g":
								snake_case__      :Union[str, Any]         =   value
				elif weight_type == "weight_v":
								snake_case__      :Any         =   value
				elif weight_type == "bias":
								snake_case__      :int         =   value
				else:
								snake_case__      :Optional[int]         =   value

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





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Optional[int]      )     -> List[str]:





				'''simple docstring'''


				snake_case__      :Any         =   []
				snake_case__      :Any         =   fairseq_model.state_dict()

				snake_case__      :Tuple         =   hf_model.feature_extractor

				# if encoder has different dim to decoder -> use proj_weight
				snake_case__      :str         =   None

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

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

				return proj_weight





def 					lowercase_ (    __snake_case     :							Any	,	__snake_case     :							Any	,	__snake_case     :							List[Any]	,	__snake_case     :							List[str]	,	__snake_case     :							Dict      )     -> List[str]:





				'''simple docstring'''


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





def 					lowercase_ (    __snake_case     :							int      )     -> Dict:





				'''simple docstring'''


				snake_case__      :Optional[int]         =   emb.weight.shape
				snake_case__      :Union[str, Any]         =   nn.Linear(__snake_case	,	__snake_case	,	bias=__snake_case      )
				snake_case__      :str         =   emb.weight.data
				return lin_layer





def 					lowercase_ (    __snake_case     :							Any      )     -> Optional[Any]:





				'''simple docstring'''


				with open(__snake_case	,	"r"	,	encoding="utf-8"      ) as f:
								snake_case__      :Union[str, Any]         =   f.readlines()
								snake_case__      :Dict         =   [line.split(" "      )[0] for line in lines]

				snake_case__      :Union[str, Any]         =   len(__snake_case      )

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

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





@torch.no_grad()
def 					lowercase_ (    __snake_case     :							Dict	,	__snake_case     :							Tuple	,	__snake_case     :							List[str]	,	__snake_case     :							Any	,	__snake_case     :							Tuple	,	__snake_case     :							int	,	__snake_case     :							List[str]	,	)     -> str:





				'''simple docstring'''


				snake_case__      :List[str]         =   WavaVecaConfig.from_pretrained(__snake_case      )
				snake_case__      :Optional[int]         =   SpeechaTextaConfig.from_pretrained(
				    __snake_case	,	vocab_size=__snake_case	,	decoder_layers=__snake_case	,	do_stable_layer_norm=__snake_case      )

				snake_case__      :Dict         =   WavaVecaFeatureExtractor(
				    feature_size=1	,	sampling_rate=1_60_00	,	padding_value=0	,	do_normalize=__snake_case	,	return_attention_mask=__snake_case	,	)

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

				# set weights for wav2vec2 encoder
				snake_case__      :str         =   WavaVecaModel(__snake_case      )
				snake_case__      :Optional[int]         =   recursively_load_weights_wavaveca(model.encoder	,	__snake_case      )

				snake_case__      :Union[str, Any]         =   SpeechaTextaForCausalLM(__snake_case      )
				snake_case__      :Optional[int]         =   hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()	,	strict=__snake_case      )

				# set output linear layer
				unexpected_keys.remove("embed_out"      )
				snake_case__      :int         =   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}'      )

				snake_case__      :List[Any]         =   SpeechEncoderDecoderModel(encoder=__snake_case	,	decoder=__snake_case      )
				snake_case__      :List[Any]         =   False

				# add projection layer
				snake_case__      :int         =   nn.Parameter(projection_layer.weight      )
				snake_case__      :Optional[int]         =   nn.Parameter(projection_layer.bias      )

				snake_case__      :Union[str, Any]         =   create_vocab_dict(__snake_case      )

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

				snake_case__      :Tuple         =   SpeechaTextaTokenizer(os.path.join(__snake_case	,	"vocab.json"      )      )
				tokenizer.save_pretrained(__snake_case      )

				snake_case__      :Dict         =   hf_wavavec.config.to_dict()
				snake_case__      :int         =   tokenizer.pad_token_id
				snake_case__      :Union[str, Any]         =   tokenizer.bos_token_id
				snake_case__      :Tuple         =   tokenizer.eos_token_id
				snake_case__      :Optional[Any]         =   "speech_to_text_2"
				snake_case__      :Optional[int]         =   "wav2vec2"

				snake_case__      :Union[str, Any]         =   SpeechEncoderDecoderConfig.from_dict(__snake_case      )

				hf_wavavec.save_pretrained(__snake_case      )
				feature_extractor.save_pretrained(__snake_case      )


if __name__ == "__main__":
	__UpperCAmelCase							:   Union[str, 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(
	    "--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							:   Union[str, Any]							=  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,
	)







import os
import sys
import unittest


__UpperCAmelCase							:   str							=  os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
__UpperCAmelCase							:   Tuple							=  os.path.join(git_repo_path, "src", "diffusers")


class 					_snake_case						(			unittest.TestCase ):






	def    lowerCAmelCase_							(		self       )						->    Union[str, Any]:
					snake_case__      :Tuple         =   find_backend("    if not is_torch_available():"       )
					self.assertEqual(UpperCamelCase					,"torch"       )

					# backend_with_underscore = find_backend("    if not is_tensorflow_text_available():")
					# self.assertEqual(backend_with_underscore, "tensorflow_text")

					snake_case__      :Tuple         =   find_backend("    if not (is_torch_available() and is_transformers_available()):"       )
					self.assertEqual(UpperCamelCase					,"torch_and_transformers"       )

					# double_backend_with_underscore = find_backend(
					#    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"
					# )
					# self.assertEqual(double_backend_with_underscore, "sentencepiece_and_tensorflow_text")

					snake_case__      :str         =   find_backend(
					    "    if not (is_torch_available() and is_transformers_available() and is_onnx_available()):"       )
					self.assertEqual(UpperCamelCase					,"torch_and_transformers_and_onnx"       )






	def    lowerCAmelCase_							(		self       )						->    str:
					snake_case__      :int         =   read_init()
					# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
					self.assertIn("torch"					,UpperCamelCase       )
					self.assertIn("torch_and_transformers"					,UpperCamelCase       )
					self.assertIn("flax_and_transformers"					,UpperCamelCase       )
					self.assertIn("torch_and_transformers_and_onnx"					,UpperCamelCase       )

					# Likewise, we can't assert on the exact content of a key
					self.assertIn("UNet2DModel"					,objects["torch"]       )
					self.assertIn("FlaxUNet2DConditionModel"					,objects["flax"]       )
					self.assertIn("StableDiffusionPipeline"					,objects["torch_and_transformers"]       )
					self.assertIn("FlaxStableDiffusionPipeline"					,objects["flax_and_transformers"]       )
					self.assertIn("LMSDiscreteScheduler"					,objects["torch_and_scipy"]       )
					self.assertIn("OnnxStableDiffusionPipeline"					,objects["torch_and_transformers_and_onnx"]       )






	def    lowerCAmelCase_							(		self       )						->    Any:
					snake_case__      :Union[str, Any]         =   create_dummy_object("CONSTANT"					,"'torch'"       )
					self.assertEqual(UpperCamelCase					,"\nCONSTANT = None\n"       )

					snake_case__      :Optional[Any]         =   create_dummy_object("function"					,"'torch'"       )
					self.assertEqual(
					    UpperCamelCase					,"\ndef function(*args, **kwargs):\n    requires_backends(function, 'torch')\n"       )

					snake_case__      :str         =   "\nclass FakeClass(metaclass=DummyObject):\n    _backends = 'torch'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, 'torch')\n\n    @classmethod\n    def from_config(cls, *args, **kwargs):\n        requires_backends(cls, 'torch')\n\n    @classmethod\n    def from_pretrained(cls, *args, **kwargs):\n        requires_backends(cls, 'torch')\n"
					snake_case__      :List[str]         =   create_dummy_object("FakeClass"					,"'torch'"       )
					self.assertEqual(UpperCamelCase					,UpperCamelCase       )






	def    lowerCAmelCase_							(		self       )						->    List[Any]:
					snake_case__      :Tuple         =   "# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, [\"torch\"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = [\"torch\"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, [\"torch\"])\n\n    @classmethod\n    def from_config(cls, *args, **kwargs):\n        requires_backends(cls, [\"torch\"])\n\n    @classmethod\n    def from_pretrained(cls, *args, **kwargs):\n        requires_backends(cls, [\"torch\"])\n"
					snake_case__      :int         =   create_dummy_files({"torch": ["CONSTANT", "function", "FakeClass"]}       )
					self.assertEqual(dummy_files["torch"]					,UpperCamelCase       )