




import collections
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_flax_cross_test,
    require_flax,
    require_torch,
    require_vision,
    slow,
    torch_device,
)
from transformers.utils import is_flax_available, is_torch_available, is_vision_available

from ...test_modeling_flax_common import floats_tensor, ids_tensor, random_attention_mask
from ..bert.test_modeling_flax_bert import FlaxBertModelTester
from ..clip.test_modeling_flax_clip import FlaxCLIPVisionModelTester
from ..vit.test_modeling_flax_vit import FlaxViTModelTester


if is_flax_available():
				from transformers import (
				    FlaxBertModel,
				    FlaxCLIPVisionModel,
				    FlaxVisionTextDualEncoderModel,
				    FlaxViTModel,
				    VisionTextDualEncoderConfig,
				    VisionTextDualEncoderProcessor,
				)
				from transformers.modeling_flax_pytorch_utils import (
				    convert_pytorch_state_dict_to_flax,
				    load_flax_weights_in_pytorch_model,
				)


if is_torch_available():
				import torch

				from transformers import VisionTextDualEncoderModel

if is_vision_available():
				from PIL import Image




def        _SCREAMING_SNAKE_CASE      (		snake_case       )						->				Any:
						if isinstance(snake_case			,     collections.abc.Iterable       ):
												return x
						return (x, x)




@require_flax
class       _A				:






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								pass






		def 				UpperCAmelCase (    self			):
								pass






		def 				UpperCAmelCase (    self			):
								pass






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			np.abs((a - b)			).max()
								self.assertLessEqual(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		F"Difference between torch and flax is {diff} (>= {tol})."			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=None						,		**_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(input_ids=_SCREAMING_SNAKE_CASE						,		pixel_values=_SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE			)

								self.assertEqual(output["""text_embeds"""].shape						,		(input_ids.shape[0], config.projection_dim)			)
								self.assertEqual(output["""image_embeds"""].shape						,		(pixel_values.shape[0], config.projection_dim)			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=None						,		**_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase    ,					_UpperCAmelCase      =			self.get_vision_text_model(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			{"""vision_model""": vision_model, """text_model""": text_model}
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(input_ids=_SCREAMING_SNAKE_CASE						,		pixel_values=_SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE			)

								self.assertEqual(output["""text_embeds"""].shape						,		(input_ids.shape[0], model.config.projection_dim)			)
								self.assertEqual(output["""image_embeds"""].shape						,		(pixel_values.shape[0], model.config.projection_dim)			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=None						,		**_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase    ,					_UpperCAmelCase      =			self.get_vision_text_model(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			{"""vision_model""": vision_model, """text_model""": text_model}
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(input_ids=_SCREAMING_SNAKE_CASE						,		pixel_values=_SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			output[0]

								with tempfile.TemporaryDirectory() as tmpdirname:
														model.save_pretrained(_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE			)

														_UpperCAmelCase      =			model(input_ids=_SCREAMING_SNAKE_CASE						,		pixel_values=_SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			after_output[0]
														_UpperCAmelCase      =			np.amax(np.abs(out_a - out_a			)			)
														self.assertLessEqual(_SCREAMING_SNAKE_CASE						,		1e-3			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=None						,		**_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase    ,					_UpperCAmelCase      =			self.get_vision_text_model(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			{"""vision_model""": vision_model, """text_model""": text_model}
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(**_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(
								    input_ids=_SCREAMING_SNAKE_CASE						,		pixel_values=_SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE						,		output_attentions=_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			output.vision_model_output.attentions
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		vision_config.num_hidden_layers			)

								# in ViT, the seq_len equals the number of patches + 1 (we add 1 for the [CLS] token)
								_UpperCAmelCase      =			to_atuple(vision_model.config.image_size			)
								_UpperCAmelCase      =			to_atuple(vision_model.config.patch_size			)
								_UpperCAmelCase      =			(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
								_UpperCAmelCase      =			num_patches + 1
								self.assertEqual(vision_attentions[0].shape[-3:]						,		(vision_config.num_attention_heads, seq_len, seq_len)			)

								_UpperCAmelCase      =			output.text_model_output.attentions
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		text_config.num_hidden_layers			)

								self.assertEqual(
								    text_attentions[0].shape[-3:]						,		(text_config.num_attention_heads, input_ids.shape[-1], input_ids.shape[-1])						,		)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								pt_model.to(_SCREAMING_SNAKE_CASE			)
								pt_model.eval()

								# prepare inputs
								_UpperCAmelCase      =			inputs_dict
								_UpperCAmelCase      =			{k: torch.tensor(v.tolist()			) for k, v in flax_inputs.items()}

								with torch.no_grad():
														_UpperCAmelCase      =			pt_model(**_SCREAMING_SNAKE_CASE			).to_tuple()

								_UpperCAmelCase      =			fx_model(**_SCREAMING_SNAKE_CASE			).to_tuple()
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		len(_SCREAMING_SNAKE_CASE			)						,		"""Output lengths differ between Flax and PyTorch"""			)
								for fx_output, pt_output in zip(fx_outputs[:4]						,		pt_outputs[:4]			):
														self.assert_almost_equals(_SCREAMING_SNAKE_CASE						,		pt_output.numpy()						,		4e-2			)

								# PT -> Flax
								with tempfile.TemporaryDirectory() as tmpdirname:
														pt_model.save_pretrained(_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE						,		from_pt=_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			fx_model_loaded(**_SCREAMING_SNAKE_CASE			).to_tuple()
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		len(_SCREAMING_SNAKE_CASE			)						,		"""Output lengths differ between Flax and PyTorch"""			)
								for fx_output_loaded, pt_output in zip(fx_outputs_loaded[:4]						,		pt_outputs[:4]			):
														self.assert_almost_equals(_SCREAMING_SNAKE_CASE						,		pt_output.numpy()						,		4e-2			)

								# Flax -> PT
								with tempfile.TemporaryDirectory() as tmpdirname:
														fx_model.save_pretrained(_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			VisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE						,		from_flax=_SCREAMING_SNAKE_CASE			)

								pt_model_loaded.to(_SCREAMING_SNAKE_CASE			)
								pt_model_loaded.eval()

								with torch.no_grad():
														_UpperCAmelCase      =			pt_model_loaded(**_SCREAMING_SNAKE_CASE			).to_tuple()

								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		len(_SCREAMING_SNAKE_CASE			)						,		"""Output lengths differ between Flax and PyTorch"""			)
								for fx_output, pt_output_loaded in zip(fx_outputs[:4]						,		pt_outputs_loaded[:4]			):
														self.assert_almost_equals(_SCREAMING_SNAKE_CASE						,		pt_output_loaded.numpy()						,		4e-2			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,		_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			fx_state

								self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			VisionTextDualEncoderConfig.from_vision_text_configs(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			VisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel(_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			load_flax_weights_in_pytorch_model(_SCREAMING_SNAKE_CASE						,		fx_model.params			)

								self.check_pt_flax_equivalence(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.prepare_config_and_inputs()
								self.check_model_from_pretrained_configs(**_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.prepare_config_and_inputs()
								self.check_vision_text_dual_encoder_from_pretrained(**_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.prepare_config_and_inputs()
								self.check_save_load(**_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.prepare_config_and_inputs()
								self.check_vision_text_output_attention(**_SCREAMING_SNAKE_CASE			)






		@is_pt_flax_cross_test
		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.prepare_config_and_inputs()
								_UpperCAmelCase      =			config_inputs_dict.pop("""vision_config"""			)
								_UpperCAmelCase      =			config_inputs_dict.pop("""text_config"""			)

								_UpperCAmelCase      =			config_inputs_dict

								self.check_equivalence_pt_to_flax(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								self.check_equivalence_flax_to_pt(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)






		@slow
		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase    ,					_UpperCAmelCase      =			self.get_pretrained_model_and_inputs()

								_UpperCAmelCase      =			model_a(**_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			outputs[0]

								with tempfile.TemporaryDirectory() as tmp_dirname:
														model_a.save_pretrained(_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_pretrained(_SCREAMING_SNAKE_CASE			)

														_UpperCAmelCase      =			model_a(**_SCREAMING_SNAKE_CASE			)
														_UpperCAmelCase      =			after_outputs[0]
														_UpperCAmelCase      =			np.amax(np.abs(out_a - out_a			)			)
														self.assertLessEqual(_SCREAMING_SNAKE_CASE						,		1e-5			)




@require_flax
class       _A				(		__lowercase       ,     unittest.TestCase     ):






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
								    """hf-internal-testing/tiny-random-vit"""						,		"""hf-internal-testing/tiny-bert"""						,		vision_from_pt=_SCREAMING_SNAKE_CASE						,		text_from_pt=_SCREAMING_SNAKE_CASE						,		)
								_UpperCAmelCase      =			13
								_UpperCAmelCase      =			floats_tensor(
								    [
								        batch_size,
								        model.config.vision_config.num_channels,
								        model.config.vision_config.image_size,
								        model.config.vision_config.image_size,
								    ]			)
								_UpperCAmelCase      =			ids_tensor([batch_size, 4]						,		model.config.text_config.vocab_size			)
								_UpperCAmelCase      =			random_attention_mask([batch_size, 4]			)
								_UpperCAmelCase      =			{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

								return model, inputs






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			FlaxViTModel(_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			FlaxBertModel(_SCREAMING_SNAKE_CASE			)
								return vision_model, text_model






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			FlaxViTModelTester(self			)
								_UpperCAmelCase      =			FlaxBertModelTester(self			)
								_UpperCAmelCase      =			vit_model_tester.prepare_config_and_inputs()
								_UpperCAmelCase      =			bert_model_tester.prepare_config_and_inputs()

								_UpperCAmelCase    ,					_UpperCAmelCase      =			vision_config_and_inputs

								_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase      =			text_config_and_inputs

								# make sure that cross attention layers are added
								return {
								    "text_config": text_config,
								    "vision_config": vision_config,
								    "pixel_values": pixel_values,
								    "attention_mask": attention_mask,
								    "input_ids": input_ids,
								    "token_type_ids": token_type_ids,
								}




@require_torch
class       _A				(		__lowercase       ,     unittest.TestCase     ):






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_vision_text_pretrained(
								    """hf-internal-testing/tiny-random-clip"""						,		"""hf-internal-testing/tiny-bert"""						,		vision_from_pt=_SCREAMING_SNAKE_CASE						,		text_from_pt=_SCREAMING_SNAKE_CASE						,		)
								_UpperCAmelCase      =			13
								_UpperCAmelCase      =			floats_tensor(
								    [
								        batch_size,
								        model.config.vision_config.num_channels,
								        model.config.vision_config.image_size,
								        model.config.vision_config.image_size,
								    ]			)
								_UpperCAmelCase      =			ids_tensor([batch_size, 4]						,		model.config.text_config.vocab_size			)
								_UpperCAmelCase      =			random_attention_mask([batch_size, 4]			)
								_UpperCAmelCase      =			{"""pixel_values""": pixel_values, """input_ids""": input_ids, """attention_mask""": attention_mask}

								return model, inputs






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								_UpperCAmelCase      =			FlaxCLIPVisionModel(_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			FlaxBertModel(_SCREAMING_SNAKE_CASE			)
								return vision_model, text_model






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			FlaxCLIPVisionModelTester(self			)
								_UpperCAmelCase      =			FlaxBertModelTester(self			)
								_UpperCAmelCase      =			clip_model_tester.prepare_config_and_inputs()
								_UpperCAmelCase      =			bert_model_tester.prepare_config_and_inputs()

								_UpperCAmelCase    ,					_UpperCAmelCase      =			vision_config_and_inputs

								_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase    ,					_UpperCAmelCase      =			text_config_and_inputs

								# make sure that cross attention layers are added
								return {
								    "text_config": text_config,
								    "vision_config": vision_config,
								    "pixel_values": pixel_values,
								    "attention_mask": attention_mask,
								    "input_ids": input_ids,
								    "token_type_ids": token_type_ids,
								}




@require_flax
@require_vision
class       _A				(		unittest.TestCase     ):






		@slow
		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			FlaxVisionTextDualEncoderModel.from_pretrained("""clip-italian/clip-italian"""						,		logit_scale_init_value=1.0			)
								_UpperCAmelCase      =			VisionTextDualEncoderProcessor.from_pretrained("""clip-italian/clip-italian"""			)

								_UpperCAmelCase      =			Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""			)
								_UpperCAmelCase      =			processor(
								    text=["""una foto di un gatto""", """una foto di un cane"""]						,		images=_SCREAMING_SNAKE_CASE						,		padding=_SCREAMING_SNAKE_CASE						,		return_tensors="""np"""			)

								_UpperCAmelCase      =			model(**_SCREAMING_SNAKE_CASE			)

								# verify the logits
								self.assertEqual(outputs.logits_per_image.shape						,		(inputs.pixel_values.shape[0], inputs.input_ids.shape[0])			)
								self.assertEqual(
								    outputs.logits_per_text.shape						,		(inputs.input_ids.shape[0], inputs.pixel_values.shape[0])						,		)

								_UpperCAmelCase      =			np.array([[1.228_4727, 0.310_4122]]			)

								self.assertTrue(np.allclose(outputs.logits_per_image						,		_SCREAMING_SNAKE_CASE						,		atol=1e-3			)			)





import argparse

import fairseq
import torch

from transformers import UniSpeechSatConfig, UniSpeechSatForCTC, UniSpeechSatForPreTraining, logging


logging.set_verbosity_info()
a													=			logging.get_logger(__name__)

a													=			{
    "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",
    "encoder.layer_norm_for_extract": "layer_norm_for_extract",
    "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",
    "label_embs_concat": "label_embeddings_concat",
    "mask_emb": "masked_spec_embed",
    "spk_proj": "speaker_proj",
}
a													=			[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
    "label_embeddings_concat",
    "speaker_proj",
    "layer_norm_for_extract",
]




def        _SCREAMING_SNAKE_CASE      (		snake_case			,     snake_case			,     snake_case			,     snake_case			,     snake_case       )						->				List[str]:
						for attribute in key.split("""."""       ):
												_UpperCAmelCase      =			getattr(snake_case			,     snake_case       )

						if weight_type is not None:
												_UpperCAmelCase      =			getattr(snake_case			,     snake_case       ).shape
						else:
												_UpperCAmelCase      =			hf_pointer.shape

						if hf_shape != value.shape:
												raise ValueError(
												    f"Shape of hf {key + '.' + weight_type if weight_type is not None else ''} is {hf_shape}, but should be"
												    f" {value.shape} for {full_name}"       )

						if weight_type == "weight":
												_UpperCAmelCase      =			value
						elif weight_type == "weight_g":
												_UpperCAmelCase      =			value
						elif weight_type == "weight_v":
												_UpperCAmelCase      =			value
						elif weight_type == "bias":
												_UpperCAmelCase      =			value
						else:
												_UpperCAmelCase      =			value

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




def        _SCREAMING_SNAKE_CASE      (		snake_case			,     snake_case       )						->				Optional[int]:
						_UpperCAmelCase      =			[]
						_UpperCAmelCase      =			fairseq_model.state_dict()

						_UpperCAmelCase      =			hf_model.unispeech_sat.feature_extractor

						for name, value in fairseq_dict.items():
												_UpperCAmelCase      =			False
												if "conv_layers" in name:
																		load_conv_layer(
																		    snake_case			,     snake_case			,     snake_case			,     snake_case			,     hf_model.config.feat_extract_norm == """group"""			,     )
																		_UpperCAmelCase      =			True
												else:
																		for key, mapped_key in MAPPING.items():
																								_UpperCAmelCase      =			"""unispeech_sat.""" + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key
																								if key in name or key.split("""w2v_model."""       )[-1] == name.split("""."""       )[0]:
																														if "layer_norm_for_extract" in name and (".".join(name.split("""."""       )[:-1]       ) != key):
																																				# special case since naming is very similar
																																				continue
																														_UpperCAmelCase      =			True
																														if "*" in mapped_key:
																																				_UpperCAmelCase      =			name.split(snake_case       )[0].split("""."""       )[-2]
																																				_UpperCAmelCase      =			mapped_key.replace("""*"""			,     snake_case       )
																														if "weight_g" in name:
																																				_UpperCAmelCase      =			"""weight_g"""
																														elif "weight_v" in name:
																																				_UpperCAmelCase      =			"""weight_v"""
																														elif "bias" in name:
																																				_UpperCAmelCase      =			"""bias"""
																														elif "weight" in name:
																																				# TODO: don't match quantizer.weight_proj
																																				_UpperCAmelCase      =			"""weight"""
																														else:
																																				_UpperCAmelCase      =			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}"       )




def        _SCREAMING_SNAKE_CASE      (		snake_case			,     snake_case			,     snake_case			,     snake_case			,     snake_case       )						->				str:
						_UpperCAmelCase      =			full_name.split("""conv_layers."""       )[-1]
						_UpperCAmelCase      =			name.split("""."""       )
						_UpperCAmelCase      =			int(items[0]       )
						_UpperCAmelCase      =			int(items[1]       )

						if type_id == 0:
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found."       )
																		_UpperCAmelCase      =			value
																		logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."       )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found."       )
																		_UpperCAmelCase      =			value
																		logger.info(f"Feat extract conv layer {layer_id} was initialized from {full_name}."       )
						elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
												if "bias" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor[layer_id].layer_norm.bias.data.shape} was found."       )
																		_UpperCAmelCase      =			value
																		logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."       )
												elif "weight" in name:
																		if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
																								raise ValueError(
																								    f"{full_name} has size {value.shape}, but"
																								    f" {feature_extractor[layer_id].layer_norm.weight.data.shape} was found."       )
																		_UpperCAmelCase      =			value
																		logger.info(f"Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}."       )
						else:
												unused_weights.append(snake_case       )




@torch.no_grad()
def        _SCREAMING_SNAKE_CASE      (		snake_case			,     snake_case			,     snake_case=None			,     snake_case=None			,     snake_case=True       )						->				List[Any]:
						if config_path is not None:
												_UpperCAmelCase      =			UniSpeechSatConfig.from_pretrained(snake_case       )
						else:
												_UpperCAmelCase      =			UniSpeechSatConfig()

						_UpperCAmelCase      =			""""""

						if is_finetuned:
												_UpperCAmelCase      =			UniSpeechSatForCTC(snake_case       )
						else:
												_UpperCAmelCase      =			UniSpeechSatForPreTraining(snake_case       )

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

						recursively_load_weights(snake_case			,     snake_case       )

						hf_wavavec.save_pretrained(snake_case       )


if __name__ == "__main__":
				a													=			argparse.ArgumentParser()
				parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
				parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
				parser.add_argument("--dict_path", default=None, type=str, help="Path to dict of fine-tuned model")
				parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
				parser.add_argument(
				    "--not_finetuned", action="store_true", help="Whether the model to convert is a fine-tuned model or not"
				)
				a													=			parser.parse_args()
				convert_unispeech_sat_checkpoint(
				    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
				)