


"""simple docstring"""






from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
							import tensorflow as tf

							from transformers import TFLEDForConditionalGeneration, TFLEDModel





@require_tf
class      lowerCamelCase__   :







					"""simple docstring"""


					__a	       =							LEDConfig
					__a	       =							{}
					__a	       =							"""gelu"""






					def __init__(       self    : List[str] ,       UpperCamelCase    : Optional[Any] ,       UpperCamelCase    : List[Any]=13 ,       UpperCamelCase    : Dict=7 ,       UpperCamelCase    : Any=True ,       UpperCamelCase    : int=False ,       UpperCamelCase    : Any=99 ,       UpperCamelCase    : str=32 ,       UpperCamelCase    : Any=2 ,       UpperCamelCase    : int=4 ,       UpperCamelCase    : List[str]=37 ,       UpperCamelCase    : List[Any]=0.1 ,       UpperCamelCase    : List[Any]=0.1 ,       UpperCamelCase    : Tuple=20 ,       UpperCamelCase    : str=2 ,       UpperCamelCase    : List[str]=1 ,       UpperCamelCase    : int=0 ,       UpperCamelCase    : Union[str, Any]=4 ,       ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[int]      							=					parent
												__UpperCAmelCase   :						Tuple      							=					batch_size
												__UpperCAmelCase   :						Dict      							=					seq_length
												__UpperCAmelCase   :						Optional[Any]      							=					is_training
												__UpperCAmelCase   :						Tuple      							=					use_labels
												__UpperCAmelCase   :						Optional[int]      							=					vocab_size
												__UpperCAmelCase   :						Dict      							=					hidden_size
												__UpperCAmelCase   :						List[Any]      							=					num_hidden_layers
												__UpperCAmelCase   :						int      							=					num_attention_heads
												__UpperCAmelCase   :						Optional[Any]      							=					intermediate_size

												__UpperCAmelCase   :						str      							=					hidden_dropout_prob
												__UpperCAmelCase   :						List[str]      							=					attention_probs_dropout_prob
												__UpperCAmelCase   :						int      							=					max_position_embeddings
												__UpperCAmelCase   :						Tuple      							=					eos_token_id
												__UpperCAmelCase   :						Tuple      							=					pad_token_id
												__UpperCAmelCase   :						Optional[int]      							=					bos_token_id
												__UpperCAmelCase   :						Any      							=					attention_window

												# `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
												# [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
												# returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
												# because its local attention only attends to `self.attention_window` and one before and one after
												__UpperCAmelCase   :						str      							=					self.attention_window + 2

												# because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
												# the `test_attention_outputs` and `test_hidden_states_output` tests
												__UpperCAmelCase   :						Union[str, Any]      							=					(
												    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
												)






					def lowerCamelCase__							(       self    : Optional[Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						str      							=					ids_tensor([self.batch_size, self.seq_length - 1] ,       self.vocab_size  )
												__UpperCAmelCase   :						Optional[int]      							=					tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size  ) ,       1  )
												__UpperCAmelCase   :						List[str]      							=					tf.concat([input_ids, eos_tensor] ,       axis=1  )

												__UpperCAmelCase   :						str      							=					ids_tensor([self.batch_size, self.seq_length] ,       self.vocab_size  )

												__UpperCAmelCase   :						List[Any]      							=					self.config_cls(
												    vocab_size=self.vocab_size ,       d_model=self.hidden_size ,       encoder_layers=self.num_hidden_layers ,       decoder_layers=self.num_hidden_layers ,       encoder_attention_heads=self.num_attention_heads ,       decoder_attention_heads=self.num_attention_heads ,       encoder_ffn_dim=self.intermediate_size ,       decoder_ffn_dim=self.intermediate_size ,       dropout=self.hidden_dropout_prob ,       attention_dropout=self.attention_probs_dropout_prob ,       max_position_embeddings=self.max_position_embeddings ,       eos_token_ids=[2] ,       bos_token_id=self.bos_token_id ,       pad_token_id=self.pad_token_id ,       decoder_start_token_id=self.pad_token_id ,       attention_window=self.attention_window ,       **self.config_updates ,       )
												__UpperCAmelCase   :						int      							=					prepare_led_inputs_dict(UpperCamelCase ,       UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						List[Any]      							=					tf.concat(
												    [tf.zeros_like(UpperCamelCase  )[:, :-1], tf.ones_like(UpperCamelCase  )[:, -1:]] ,       axis=-1 ,       )
												__UpperCAmelCase   :						List[Any]      							=					global_attention_mask
												return config, inputs_dict






					def lowerCamelCase__							(       self    : List[Any] ,       UpperCamelCase    : str ,       UpperCamelCase    : List[str]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Any      							=					TFLEDModel(config=UpperCamelCase  ).get_decoder()
												__UpperCAmelCase   :						List[str]      							=					inputs_dict["""input_ids"""]

												__UpperCAmelCase   :						int      							=					input_ids[:1, :]
												__UpperCAmelCase   :						List[str]      							=					inputs_dict["""attention_mask"""][:1, :]
												__UpperCAmelCase   :						Union[str, Any]      							=					1

												# first forward pass
												__UpperCAmelCase   :						List[str]      							=					model(UpperCamelCase ,       attention_mask=UpperCamelCase ,       use_cache=UpperCamelCase  )

												__UpperCAmelCase							,__UpperCAmelCase   :						List[Any]      							=					outputs.to_tuple()

												# create hypothetical next token and extent to next_input_ids
												__UpperCAmelCase   :						List[str]      							=					ids_tensor((self.batch_size, 3) ,       config.vocab_size  )
												__UpperCAmelCase   :						List[Any]      							=					tf.cast(ids_tensor((self.batch_size, 3) ,       2  ) ,       tf.inta  )

												# append to next input_ids and
												__UpperCAmelCase   :						str      							=					tf.concat([input_ids, next_tokens] ,       axis=-1  )
												__UpperCAmelCase   :						int      							=					tf.concat([attention_mask, next_attn_mask] ,       axis=-1  )

												__UpperCAmelCase   :						Tuple      							=					model(UpperCamelCase ,       attention_mask=UpperCamelCase  )[0]
												__UpperCAmelCase   :						Tuple      							=					model(UpperCamelCase ,       attention_mask=UpperCamelCase ,       past_key_values=UpperCamelCase  )[0]

												self.parent.assertEqual(next_tokens.shape[1] ,       output_from_past.shape[1]  )

												# select random slice
												__UpperCAmelCase   :						Dict      							=					int(ids_tensor((1,) ,       output_from_past.shape[-1]  )  )
												__UpperCAmelCase   :						Tuple      							=					output_from_no_past[:, -3:, random_slice_idx]
												__UpperCAmelCase   :						Tuple      							=					output_from_past[:, :, random_slice_idx]

												# test that outputs are equal for slice
												tf.debugging.assert_near(UpperCamelCase ,       UpperCamelCase ,       rtol=1e-3  )






def      lowerCamelCase       (     _UpperCamelCase		:      Optional[int]       ,			_UpperCamelCase		:      Dict       ,			_UpperCamelCase		:      List[str]       ,			_UpperCamelCase		:      Optional[Any]=None       ,			_UpperCamelCase		:      List[str]=None       ,			_UpperCamelCase		:      Tuple=None       ,			_UpperCamelCase		:      List[str]=None       ,			)		->       Union[str, Any]:

							'''simple docstring'''



							if attention_mask is None:
														__UpperCAmelCase   :						List[str]      							=					tf.cast(tf.math.not_equal(_UpperCamelCase       ,			config.pad_token_id							)       ,			tf.inta							)
							if decoder_attention_mask is None:
														__UpperCAmelCase   :						List[str]      							=					tf.concat(
														    [
														        tf.ones(decoder_input_ids[:, :1].shape       ,			dtype=tf.inta							),
														        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]       ,			config.pad_token_id							)       ,			tf.inta							),
														    ]       ,			axis=-1       ,			)
							if head_mask is None:
														__UpperCAmelCase   :						Any      							=					tf.ones((config.encoder_layers, config.encoder_attention_heads)							)
							if decoder_head_mask is None:
														__UpperCAmelCase   :						Optional[int]      							=					tf.ones((config.decoder_layers, config.decoder_attention_heads)							)
							return {
							    "input_ids": input_ids,
							    "attention_mask": attention_mask,
							    "decoder_input_ids": decoder_input_ids,
							    "decoder_attention_mask": decoder_attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							}





@require_tf
class      lowerCamelCase__   (     A     ,    A     ,    unittest.TestCase				):







					"""simple docstring"""


					__a	       =							(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
					__a	       =							(TFLEDForConditionalGeneration,) if is_tf_available() else ()
					__a	       =							(
					    {
					        """conversational""": TFLEDForConditionalGeneration,
					        """feature-extraction""": TFLEDModel,
					        """summarization""": TFLEDForConditionalGeneration,
					        """text2text-generation""": TFLEDForConditionalGeneration,
					        """translation""": TFLEDForConditionalGeneration,
					    }
					    if is_tf_available()
					    else {}
					)
					__a	       =							True
					__a	       =							False
					__a	       =							False
					__a	       =							False






					def lowerCamelCase__							(       self    : List[str]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						List[Any]      							=					TFLEDModelTester(self  )
												__UpperCAmelCase   :						Any      							=					ConfigTester(self ,       config_class=UpperCamelCase  )






					def lowerCamelCase__							(       self    : str  ):
												'''simple docstring'''






												self.config_tester.run_common_tests()






					def lowerCamelCase__							(       self    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[Any]      							=					self.model_tester.prepare_config_and_inputs_for_common()
												self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase  )






					def lowerCamelCase__							(       self    : Tuple  ):
												'''simple docstring'''






												__UpperCAmelCase							,__UpperCAmelCase   :						List[str]      							=					self.model_tester.prepare_config_and_inputs_for_common()
												__UpperCAmelCase   :						List[Any]      							=					tf.zeros_like(inputs_dict["""attention_mask"""]  )
												__UpperCAmelCase   :						Optional[int]      							=					2
												__UpperCAmelCase   :						Optional[Any]      							=					tf.where(
												    tf.range(self.model_tester.seq_length  )[None, :] < num_global_attn_indices ,       1 ,       inputs_dict["""global_attention_mask"""] ,       )

												__UpperCAmelCase   :						List[Any]      							=					True
												__UpperCAmelCase   :						str      							=					self.model_tester.seq_length
												__UpperCAmelCase   :						List[Any]      							=					self.model_tester.encoder_seq_length

												def check_decoder_attentions_output(UpperCamelCase    : int  ):
																			__UpperCAmelCase   :						Dict      							=					outputs.decoder_attentions
																			self.assertEqual(len(UpperCamelCase  ) ,       self.model_tester.num_hidden_layers  )
																			self.assertListEqual(
																			    list(decoder_attentions[0].shape[-3:]  ) ,       [self.model_tester.num_attention_heads, seq_length, seq_length] ,       )

												def check_encoder_attentions_output(UpperCamelCase    : Optional[int]  ):
																			__UpperCAmelCase   :						Tuple      							=					[t.numpy() for t in outputs.encoder_attentions]
																			__UpperCAmelCase   :						Dict      							=					[t.numpy() for t in outputs.encoder_global_attentions]
																			self.assertEqual(len(UpperCamelCase  ) ,       self.model_tester.num_hidden_layers  )
																			self.assertEqual(len(UpperCamelCase  ) ,       self.model_tester.num_hidden_layers  )
																			self.assertListEqual(
																			    list(attentions[0].shape[-3:]  ) ,       [self.model_tester.num_attention_heads, seq_length, seq_length] ,       )
																			self.assertListEqual(
																			    list(global_attentions[0].shape[-3:]  ) ,       [self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices] ,       )

												for model_class in self.all_model_classes:
																			__UpperCAmelCase   :						Optional[Any]      							=					True
																			__UpperCAmelCase   :						Union[str, Any]      							=					False
																			__UpperCAmelCase   :						str      							=					False
																			__UpperCAmelCase   :						Dict      							=					model_class(UpperCamelCase  )
																			__UpperCAmelCase   :						int      							=					model(self._prepare_for_class(UpperCamelCase ,       UpperCamelCase  )  )
																			__UpperCAmelCase   :						List[Any]      							=					len(UpperCamelCase  )
																			self.assertEqual(config.output_hidden_states ,       UpperCamelCase  )
																			check_encoder_attentions_output(UpperCamelCase  )

																			if self.is_encoder_decoder:
																										__UpperCAmelCase   :						Dict      							=					model_class(UpperCamelCase  )
																										__UpperCAmelCase   :						Any      							=					model(self._prepare_for_class(UpperCamelCase ,       UpperCamelCase  )  )
																										self.assertEqual(config.output_hidden_states ,       UpperCamelCase  )
																										check_decoder_attentions_output(UpperCamelCase  )

																			# Check that output attentions can also be changed via the config
																			del inputs_dict["output_attentions"]
																			__UpperCAmelCase   :						List[str]      							=					True
																			__UpperCAmelCase   :						List[str]      							=					model_class(UpperCamelCase  )
																			__UpperCAmelCase   :						Tuple      							=					model(self._prepare_for_class(UpperCamelCase ,       UpperCamelCase  )  )
																			self.assertEqual(config.output_hidden_states ,       UpperCamelCase  )
																			check_encoder_attentions_output(UpperCamelCase  )

																			# Check attention is always last and order is fine
																			__UpperCAmelCase   :						str      							=					True
																			__UpperCAmelCase   :						Union[str, Any]      							=					True
																			__UpperCAmelCase   :						Tuple      							=					model_class(UpperCamelCase  )
																			__UpperCAmelCase   :						Optional[Any]      							=					model(self._prepare_for_class(UpperCamelCase ,       UpperCamelCase  )  )

																			self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1) ,       len(UpperCamelCase  )  )
																			self.assertEqual(model.config.output_hidden_states ,       UpperCamelCase  )
																			check_encoder_attentions_output(UpperCamelCase  )






					@unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing."""  )
					def lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												pass






					def lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												pass






def      lowerCamelCase       (     _UpperCamelCase		:      List[Any]							)		->       List[str]:

							'''simple docstring'''



							return tf.constant(_UpperCamelCase       ,			dtype=tf.intaa							)


UpperCAmelCase    :     List[str]   						=		1E-4





@slow
@require_tf
class      lowerCamelCase__   (     unittest.TestCase				):







					"""simple docstring"""


					def lowerCamelCase__							(       self    : Any  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Optional[Any]      							=					TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""  ).led

												# change to intended input here
												__UpperCAmelCase   :						List[str]      							=					_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]  )
												__UpperCAmelCase   :						List[Any]      							=					_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]  )
												__UpperCAmelCase   :						int      							=					prepare_led_inputs_dict(model.config ,       UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						List[Any]      							=					model(**UpperCamelCase  )[0]
												__UpperCAmelCase   :						Tuple      							=					(1, 1_024, 768)
												self.assertEqual(output.shape ,       UpperCamelCase  )
												# change to expected output here
												__UpperCAmelCase   :						Optional[int]      							=					tf.convert_to_tensor(
												    [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]] ,       )
												tf.debugging.assert_near(output[:, :3, :3] ,       UpperCamelCase ,       atol=1e-3  )






					def lowerCamelCase__							(       self    : Union[str, Any]  ):
												'''simple docstring'''






												__UpperCAmelCase   :						Union[str, Any]      							=					TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""  )

												# change to intended input here
												__UpperCAmelCase   :						Optional[Any]      							=					_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]  )
												__UpperCAmelCase   :						int      							=					_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]  )
												__UpperCAmelCase   :						Any      							=					prepare_led_inputs_dict(model.config ,       UpperCamelCase ,       UpperCamelCase  )
												__UpperCAmelCase   :						Any      							=					model(**UpperCamelCase  )[0]
												__UpperCAmelCase   :						str      							=					(1, 1_024, model.config.vocab_size)
												self.assertEqual(output.shape ,       UpperCamelCase  )
												# change to expected output here
												__UpperCAmelCase   :						Dict      							=					tf.convert_to_tensor(
												    [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]] ,       )
												tf.debugging.assert_near(output[:, :3, :3] ,       UpperCamelCase ,       atol=1e-3 ,       rtol=1e-3  )




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