




"""simple docstring"""




from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class 					__magic_name__							( UpperCAmelCase__       ):






				'''simple docstring'''
				__UpperCamelCase    						=		42
				__UpperCamelCase    						=		42






class 					__magic_name__							( nn.Module       ):






				'''simple docstring'''
				__UpperCamelCase    						=		42
				__UpperCamelCase    						=		(16, 32, 96, 2_56)
				__UpperCamelCase    						=		jnp.floataa





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      nn.Conv(
							    self.block_out_channels[0]	,				kernel_size=(3, 3)	,				padding=((1, 1), (1, 1))	,				dtype=self.dtype	,				)

							lowerCamelCase       =      []
							for i in range(len(self.block_out_channels				) - 1				):
										lowerCamelCase       =      self.block_out_channels[i]
										lowerCamelCase       =      self.block_out_channels[i + 1]
										lowerCamelCase       =      nn.Conv(
										    _lowerCAmelCase	,				kernel_size=(3, 3)	,				padding=((1, 1), (1, 1))	,				dtype=self.dtype	,				)
										blocks.append(_lowerCAmelCase				)
										lowerCamelCase       =      nn.Conv(
										    _lowerCAmelCase	,				kernel_size=(3, 3)	,				strides=(2, 2)	,				padding=((1, 1), (1, 1))	,				dtype=self.dtype	,				)
										blocks.append(_lowerCAmelCase				)
							lowerCamelCase       =      blocks

							lowerCamelCase       =      nn.Conv(
							    self.conditioning_embedding_channels	,				kernel_size=(3, 3)	,				padding=((1, 1), (1, 1))	,				kernel_init=nn.initializers.zeros_init()	,				bias_init=nn.initializers.zeros_init()	,				dtype=self.dtype	,				)





				def __call__( self	,				_a				):





							"""simple docstring"""
							lowerCamelCase       =      self.conv_in(_lowerCAmelCase				)
							lowerCamelCase       =      nn.silu(_lowerCAmelCase				)

							for block in self.blocks:
										lowerCamelCase       =      block(_lowerCAmelCase				)
										lowerCamelCase       =      nn.silu(_lowerCAmelCase				)

							lowerCamelCase       =      self.conv_out(_lowerCAmelCase				)

							return embedding






@flax_register_to_config
class 					__magic_name__							( nn.Module		,      UpperCAmelCase__		,      UpperCAmelCase__       ):






				'''simple docstring'''
				__UpperCamelCase    						=		32
				__UpperCamelCase    						=		4
				__UpperCamelCase    						=		(
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "CrossAttnDownBlock2D",
				    "DownBlock2D",
				)
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		(3_20, 6_40, 12_80, 12_80)
				__UpperCamelCase    						=		2
				__UpperCamelCase    						=		8
				__UpperCamelCase    						=		None
				__UpperCamelCase    						=		12_80
				__UpperCamelCase    						=		0.0
				__UpperCamelCase    						=		False
				__UpperCamelCase    						=		jnp.floataa
				__UpperCamelCase    						=		True
				__UpperCamelCase    						=		0
				__UpperCamelCase    						=		"rgb"
				__UpperCamelCase    						=		(16, 32, 96, 2_56)





				def      _lowerCAmelCase       ( self	,				_a				):





							"""simple docstring"""
							# init input tensors
							lowerCamelCase       =      (1, self.in_channels, self.sample_size, self.sample_size)
							lowerCamelCase       =      jnp.zeros(_lowerCAmelCase	,				dtype=jnp.floataa				)
							lowerCamelCase       =      jnp.ones((1,)	,				dtype=jnp.intaa				)
							lowerCamelCase       =      jnp.zeros((1, 1, self.cross_attention_dim)	,				dtype=jnp.floataa				)
							lowerCamelCase       =      (1, 3, self.sample_size * 8, self.sample_size * 8)
							lowerCamelCase       =      jnp.zeros(_lowerCAmelCase	,				dtype=jnp.floataa				)

							lowerCamelCase   ,     lowerCamelCase       =      jax.random.split(_lowerCAmelCase				)
							lowerCamelCase       =      {"""params""": params_rng, """dropout""": dropout_rng}

							return self.init(_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase				)["params"]





				def      _lowerCAmelCase       ( self				):





							"""simple docstring"""
							lowerCamelCase       =      self.block_out_channels
							lowerCamelCase       =      block_out_channels[0] * 4

							# If `num_attention_heads` is not defined (which is the case for most models)
							# it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
							# The reason for this behavior is to correct for incorrectly named variables that were introduced
							# when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
							# Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
							# which is why we correct for the naming here.
							lowerCamelCase       =      self.num_attention_heads or self.attention_head_dim

							# input
							lowerCamelCase       =      nn.Conv(
							    block_out_channels[0]	,				kernel_size=(3, 3)	,				strides=(1, 1)	,				padding=((1, 1), (1, 1))	,				dtype=self.dtype	,				)

							# time
							lowerCamelCase       =      FlaxTimesteps(
							    block_out_channels[0]	,				flip_sin_to_cos=self.flip_sin_to_cos	,				freq_shift=self.config.freq_shift				)
							lowerCamelCase       =      FlaxTimestepEmbedding(_lowerCAmelCase	,				dtype=self.dtype				)

							lowerCamelCase       =      FlaxControlNetConditioningEmbedding(
							    conditioning_embedding_channels=block_out_channels[0]	,				block_out_channels=self.conditioning_embedding_out_channels	,				)

							lowerCamelCase       =      self.only_cross_attention
							if isinstance(_lowerCAmelCase	,				_lowerCAmelCase				):
										lowerCamelCase       =      (only_cross_attention,) * len(self.down_block_types				)

							if isinstance(_lowerCAmelCase	,				_lowerCAmelCase				):
										lowerCamelCase       =      (num_attention_heads,) * len(self.down_block_types				)

							# down
							lowerCamelCase       =      []
							lowerCamelCase       =      []

							lowerCamelCase       =      block_out_channels[0]

							lowerCamelCase       =      nn.Conv(
							    _lowerCAmelCase	,				kernel_size=(1, 1)	,				padding="""VALID"""	,				kernel_init=nn.initializers.zeros_init()	,				bias_init=nn.initializers.zeros_init()	,				dtype=self.dtype	,				)
							controlnet_down_blocks.append(_lowerCAmelCase				)

							for i, down_block_type in enumerate(self.down_block_types				):
										lowerCamelCase       =      output_channel
										lowerCamelCase       =      block_out_channels[i]
										lowerCamelCase       =      i == len(_lowerCAmelCase				) - 1

										if down_block_type == "CrossAttnDownBlock2D":
													lowerCamelCase       =      FlaxCrossAttnDownBlockaD(
													    in_channels=_lowerCAmelCase	,				out_channels=_lowerCAmelCase	,				dropout=self.dropout	,				num_layers=self.layers_per_block	,				num_attention_heads=num_attention_heads[i]	,				add_downsample=not is_final_block	,				use_linear_projection=self.use_linear_projection	,				only_cross_attention=only_cross_attention[i]	,				dtype=self.dtype	,				)
										else:
													lowerCamelCase       =      FlaxDownBlockaD(
													    in_channels=_lowerCAmelCase	,				out_channels=_lowerCAmelCase	,				dropout=self.dropout	,				num_layers=self.layers_per_block	,				add_downsample=not is_final_block	,				dtype=self.dtype	,				)

										down_blocks.append(_lowerCAmelCase				)

										for _ in range(self.layers_per_block				):
													lowerCamelCase       =      nn.Conv(
													    _lowerCAmelCase	,				kernel_size=(1, 1)	,				padding="""VALID"""	,				kernel_init=nn.initializers.zeros_init()	,				bias_init=nn.initializers.zeros_init()	,				dtype=self.dtype	,				)
													controlnet_down_blocks.append(_lowerCAmelCase				)

										if not is_final_block:
													lowerCamelCase       =      nn.Conv(
													    _lowerCAmelCase	,				kernel_size=(1, 1)	,				padding="""VALID"""	,				kernel_init=nn.initializers.zeros_init()	,				bias_init=nn.initializers.zeros_init()	,				dtype=self.dtype	,				)
													controlnet_down_blocks.append(_lowerCAmelCase				)

							lowerCamelCase       =      down_blocks
							lowerCamelCase       =      controlnet_down_blocks

							# mid
							lowerCamelCase       =      block_out_channels[-1]
							lowerCamelCase       =      FlaxUNetMidBlockaDCrossAttn(
							    in_channels=_lowerCAmelCase	,				dropout=self.dropout	,				num_attention_heads=num_attention_heads[-1]	,				use_linear_projection=self.use_linear_projection	,				dtype=self.dtype	,				)

							lowerCamelCase       =      nn.Conv(
							    _lowerCAmelCase	,				kernel_size=(1, 1)	,				padding="""VALID"""	,				kernel_init=nn.initializers.zeros_init()	,				bias_init=nn.initializers.zeros_init()	,				dtype=self.dtype	,				)





				def __call__( self	,				_a	,				_a	,				_a	,				_a	,				_a = 1.0	,				_a = True	,				_a = False	,				):





							"""simple docstring"""
							lowerCamelCase       =      self.controlnet_conditioning_channel_order
							if channel_order == "bgr":
										lowerCamelCase       =      jnp.flip(_lowerCAmelCase	,				axis=1				)

							# 1. time
							if not isinstance(_lowerCAmelCase	,				jnp.ndarray				):
										lowerCamelCase       =      jnp.array([timesteps]	,				dtype=jnp.intaa				)
							elif isinstance(_lowerCAmelCase	,				jnp.ndarray				) and len(timesteps.shape				) == 0:
										lowerCamelCase       =      timesteps.astype(dtype=jnp.floataa				)
										lowerCamelCase       =      jnp.expand_dims(_lowerCAmelCase	,				0				)

							lowerCamelCase       =      self.time_proj(_lowerCAmelCase				)
							lowerCamelCase       =      self.time_embedding(_lowerCAmelCase				)

							# 2. pre-process
							lowerCamelCase       =      jnp.transpose(_lowerCAmelCase	,				(0, 2, 3, 1)				)
							lowerCamelCase       =      self.conv_in(_lowerCAmelCase				)

							lowerCamelCase       =      jnp.transpose(_lowerCAmelCase	,				(0, 2, 3, 1)				)
							lowerCamelCase       =      self.controlnet_cond_embedding(_lowerCAmelCase				)
							sample += controlnet_cond

							# 3. down
							lowerCamelCase       =      (sample,)
							for down_block in self.down_blocks:
										if isinstance(_lowerCAmelCase	,				_lowerCAmelCase				):
													lowerCamelCase   ,     lowerCamelCase       =      down_block(_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				deterministic=not train				)
										else:
													lowerCamelCase   ,     lowerCamelCase       =      down_block(_lowerCAmelCase	,				_lowerCAmelCase	,				deterministic=not train				)
										down_block_res_samples += res_samples

							# 4. mid
							lowerCamelCase       =      self.mid_block(_lowerCAmelCase	,				_lowerCAmelCase	,				_lowerCAmelCase	,				deterministic=not train				)

							# 5. contronet blocks
							lowerCamelCase       =      ()
							for down_block_res_sample, controlnet_block in zip(_lowerCAmelCase	,				self.controlnet_down_blocks				):
										lowerCamelCase       =      controlnet_block(_lowerCAmelCase				)
										controlnet_down_block_res_samples += (down_block_res_sample,)

							lowerCamelCase       =      controlnet_down_block_res_samples

							lowerCamelCase       =      self.controlnet_mid_block(_lowerCAmelCase				)

							# 6. scaling
							lowerCamelCase       =      [sample * conditioning_scale for sample in down_block_res_samples]
							mid_block_res_sample *= conditioning_scale

							if not return_dict:
										return (down_block_res_samples, mid_block_res_sample)

							return FlaxControlNetOutput(
							    down_block_res_samples=_lowerCAmelCase	,				mid_block_res_sample=_lowerCAmelCase				)









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

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



def 					a__				(       snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__       )		->      List[Any]:
			for attribute in key.split("""."""       ):
						lowerCamelCase       =      getattr(snake_case__ ,				snake_case__       )

			if weight_type is not None:
						lowerCamelCase       =      getattr(snake_case__ ,				snake_case__       ).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 					a__				(       snake_case__ ,				snake_case__       )		->      Tuple:
			lowerCamelCase       =      []
			lowerCamelCase       =      fairseq_model.state_dict()

			lowerCamelCase       =      hf_model.feature_extractor

			# if encoder has different dim to decoder -> use proj_weight
			lowerCamelCase       =      None

			for name, value in fairseq_dict.items():
						lowerCamelCase       =      False
						if "conv_layers" in name:
									load_conv_layer(
									    snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				hf_model.config.feat_extract_norm == """group""" ,				)
									lowerCamelCase       =      True
						elif name.split("""."""       )[0] == "proj":
									lowerCamelCase       =      fairseq_model.proj
									lowerCamelCase       =      True
						else:
									for key, mapped_key in MAPPING.items():
												if key in name or key.split("""w2v_model."""       )[-1] == name.split("""."""       )[0]:
															lowerCamelCase       =      True
															if "*" in mapped_key:
																		lowerCamelCase       =      name.split(snake_case__       )[0].split("""."""       )[-2]
																		lowerCamelCase       =      mapped_key.replace("""*""" ,				snake_case__       )
															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:
																		lowerCamelCase       =      """weight"""
															else:
																		lowerCamelCase       =      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 					a__				(       snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__       )		->      Optional[int]:
			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(snake_case__       )



def 					a__				(       snake_case__       )		->      Tuple:
			lowerCamelCase   ,     lowerCamelCase       =      emb.weight.shape
			lowerCamelCase       =      nn.Linear(snake_case__ ,				snake_case__ ,				bias=snake_case__       )
			lowerCamelCase       =      emb.weight.data
			return lin_layer



def 					a__				(       snake_case__       )		->      Optional[int]:
			with open(snake_case__ ,				"""r""" ,				encoding="""utf-8"""       ) as f:
						lowerCamelCase       =      f.readlines()
						lowerCamelCase       =      [line.split(""" """       )[0] for line in lines]

			lowerCamelCase       =      len(snake_case__       )

			lowerCamelCase       =      {
			    """<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 					a__				(       snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				snake_case__ ,				)		->      Optional[int]:
			lowerCamelCase       =      WavaVecaConfig.from_pretrained(snake_case__       )
			lowerCamelCase       =      SpeechaTextaConfig.from_pretrained(
			    snake_case__ ,				vocab_size=snake_case__ ,				decoder_layers=snake_case__ ,				do_stable_layer_norm=snake_case__       )

			lowerCamelCase       =      WavaVecaFeatureExtractor(
			    feature_size=1 ,				sampling_rate=1_60_00 ,				padding_value=0 ,				do_normalize=snake_case__ ,				return_attention_mask=snake_case__ ,				)

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

			# set weights for wav2vec2 encoder
			lowerCamelCase       =      WavaVecaModel(snake_case__       )
			lowerCamelCase       =      recursively_load_weights_wavaveca(model.encoder ,				snake_case__       )

			lowerCamelCase       =      SpeechaTextaForCausalLM(snake_case__       )
			lowerCamelCase   ,     lowerCamelCase       =      hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict() ,				strict=snake_case__       )

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

			lowerCamelCase       =      SpeechEncoderDecoderModel(encoder=snake_case__ ,				decoder=snake_case__       )
			lowerCamelCase       =      False

			# add projection layer
			lowerCamelCase       =      nn.Parameter(projection_layer.weight       )
			lowerCamelCase       =      nn.Parameter(projection_layer.bias       )

			lowerCamelCase       =      create_vocab_dict(snake_case__       )

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

			lowerCamelCase       =      SpeechaTextaTokenizer(os.path.join(snake_case__ ,				"""vocab.json"""       )       )
			tokenizer.save_pretrained(snake_case__       )

			lowerCamelCase       =      hf_wavavec.config.to_dict()
			lowerCamelCase       =      tokenizer.pad_token_id
			lowerCamelCase       =      tokenizer.bos_token_id
			lowerCamelCase       =      tokenizer.eos_token_id
			lowerCamelCase       =      """speech_to_text_2"""
			lowerCamelCase       =      """wav2vec2"""

			lowerCamelCase       =      SpeechEncoderDecoderConfig.from_dict(snake_case__       )

			hf_wavavec.save_pretrained(snake_case__       )
			feature_extractor.save_pretrained(snake_case__       )


if __name__ == "__main__":
					lowerCAmelCase							:    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_0224, type=int, help="""Vocab size of decoder""")
					parser.add_argument("""--num_decoder_layers""", default=7, type=int, help="""Number of decoder layers""")

					lowerCAmelCase							:    str  							=							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,
					)



