



"""simple docstring"""







import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    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',
    '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',
}
A		=						[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
]

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							int     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							Union[str, Any] ):
						"""simple docstring"""






						for attribute in key.split("." ):
												snake_case   :			Dict						    = getattr(lowerCamelCase_     ,		lowerCamelCase_ )

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

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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Dict ):
						"""simple docstring"""






						snake_case   :			Optional[Any]						    = []
						snake_case   :			Tuple						    = fairseq_model.state_dict()

						snake_case   :			Union[str, Any]						    = hf_model.feature_extractor
						snake_case   :			Any						    = hf_model.adapter

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

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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							List[str]     ,		lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Optional[int] ):
						"""simple docstring"""






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

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Any ):
						"""simple docstring"""






						snake_case   :			str						    = full_name.split("adaptor." )[-1]
						snake_case   :			List[str]						    = name.split("." )

						if items[1].isdigit():
												snake_case   :			List[Any]						    = int(items[1] )
						else:
												snake_case   :			Optional[int]						    = None

						if "adaptor" not in full_name:
												if "proj_ln" in full_name:
																		# has to be layer norm
																		if "bias" in name:
																								assert (
																								    value.shape == adapter.proj_layer_norm.bias.data.shape
																								), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
																								snake_case   :			Optional[Any]						    = value
																								logger.info(f'''Adapter proj layer norm bias was initialized from {full_name}.''' )
																		if "weight" in name:
																								assert (
																								    value.shape == adapter.proj_layer_norm.weight.data.shape
																								), f'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
																								snake_case   :			Dict						    = value
												else:
																		# has to be projection layer
																		if "bias" in name:
																								assert (
																								    value.shape == adapter.proj.bias.data.shape
																								), f'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
																								snake_case   :			Optional[int]						    = value
																								logger.info(f'''Adapter proj layer bias was initialized from {full_name}.''' )
																		if "weight" in name:
																								assert (
																								    value.shape == adapter.proj.weight.data.shape
																								), f'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
																								snake_case   :			List[Any]						    = value
																								logger.info(f'''Adapter proj layer weight was initialized from {full_name}.''' )
						elif isinstance(lowerCamelCase_     ,		lowerCamelCase_ ):
												if "bias" in name:
																		assert (
																		    value.shape == adapter.layers[layer_id].conv.bias.data.shape
																		), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
																		snake_case   :			Optional[int]						    = value
																		logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
												elif "weight" in name:
																		assert (
																		    value.shape == adapter.layers[layer_id].conv.weight.data.shape
																		), f'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
																		snake_case   :			Optional[Any]						    = value
																		logger.info(f'''Adapter layer {layer_id} bias was initialized from {full_name}.''' )
						else:
												unused_weights.append(lowerCamelCase_ )

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Union[str, Any] ):
						"""simple docstring"""






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

@torch.no_grad()
def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							List[Any]     ,		lowerCamelCase_:							List[Any]     ,		lowerCamelCase_:							Optional[Any]     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							List[Any]     ,		lowerCamelCase_:							Tuple     ,		lowerCamelCase_:							List[str]     ,		lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							int     ,		lowerCamelCase_:							List[str]     ,		lowerCamelCase_:							int     ,		):
						"""simple docstring"""






						snake_case   :			Union[str, Any]						    = WavaVecaConfig.from_pretrained(
						    lowerCamelCase_     ,		add_adapter=lowerCamelCase_     ,		adapter_stride=lowerCamelCase_     ,		adapter_kernel_size=lowerCamelCase_     ,		use_auth_token=lowerCamelCase_     ,		output_hidden_size=lowerCamelCase_     ,		)
						snake_case   :			int						    = MBartConfig.from_pretrained(lowerCamelCase_ )

						# load model
						snake_case    ,			snake_case    ,			snake_case   :			int						    = fairseq.checkpoint_utils.load_model_ensemble_and_task(
						    [checkpoint_path]     ,		arg_overrides={
						        "config_yaml": config_yaml_path,
						        "data": "/".join(dict_path.split("/" )[:-1] ),
						        "w2v_path": checkpoint_path,
						        "load_pretrained_decoder_from": None,
						    }     ,		)
						snake_case   :			Dict						    = model[0].eval()

						# load feature extractor
						snake_case   :			List[str]						    = WavaVecaFeatureExtractor.from_pretrained(lowerCamelCase_     ,		use_auth_token=lowerCamelCase_ )

						# set weights for wav2vec2 encoder
						snake_case   :			Any						    = WavaVecaModel(lowerCamelCase_ )

						recursively_load_weights_wavaveca(model.encoder     ,		lowerCamelCase_ )

						# load decoder weights
						snake_case   :			Tuple						    = MBartForCausalLM(lowerCamelCase_ )
						snake_case    ,			snake_case   :			Union[str, Any]						    = hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()     ,		strict=lowerCamelCase_ )
						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   :			int						    = SpeechEncoderDecoderModel(encoder=lowerCamelCase_     ,		decoder=lowerCamelCase_ )
						snake_case   :			Tuple						    = False

						snake_case   :			Optional[Any]						    = MBartaaTokenizer(lowerCamelCase_ )
						tokenizer.save_pretrained(lowerCamelCase_ )

						snake_case   :			Union[str, Any]						    = hf_wavavec.config.to_dict()
						snake_case   :			Union[str, Any]						    = tokenizer.pad_token_id
						snake_case   :			Tuple						    = tokenizer.bos_token_id
						snake_case   :			List[Any]						    = tokenizer.eos_token_id
						snake_case   :			Dict						    = "mbart50"
						snake_case   :			Any						    = "wav2vec2"

						snake_case   :			Optional[int]						    = tokenizer.eos_token_id
						snake_case   :			Optional[Any]						    = 2_5_0_0_0_4
						snake_case   :			int						    = tokenizer.eos_token_id

						snake_case   :			int						    = SpeechEncoderDecoderConfig.from_dict(lowerCamelCase_ )

						hf_wavavec.save_pretrained(lowerCamelCase_ )
						feature_extractor.save_pretrained(lowerCamelCase_ )


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_yaml_path', default=None, type=str, help='Path to yaml file of fine-tuned model')
			parser.add_argument(
			    '--encoder_config_path',
			    default='facebook/wav2vec2-xls-r-1b',
			    type=str,
			    help='Path to hf encoder wav2vec2 checkpoint config',
			)
			parser.add_argument(
			    '--decoder_config_path',
			    default='facebook/mbart-large-50-one-to-many-mmt',
			    type=str,
			    help='Path to hf decoder checkpoint config',
			)
			parser.add_argument('--add_adapter', default=True, type=bool, help='whethere to add model adapter layers')
			parser.add_argument('--adapter_stride', default=2, type=int, help='stride of adapter layers')
			parser.add_argument('--adapter_kernel_size', default=3, type=int, help='kernel size of adapter layers')
			parser.add_argument('--encoder_output_dim', default=1_0_2_4, type=int, help='encoder output dim')
			parser.add_argument('--start_token_id', default=2_5_0_0_0_4, type=int, help='`decoder_start_token_id` of model config')

			A		=						parser.parse_args()
			convert_wavaveca_checkpoint(
			    args.checkpoint_path,
			    args.pytorch_dump_folder_path,
			    args.dict_path,
			    args.config_yaml_path,
			    encoder_config_path=args.encoder_config_path,
			    decoder_config_path=args.decoder_config_path,
			    add_adapter=args.add_adapter,
			    adapter_kernel_size=args.adapter_kernel_size,
			    adapter_stride=args.adapter_stride,
			    decoder_start_token_id=args.start_token_id,
			    encoder_output_dim=args.encoder_output_dim,
			)








"""simple docstring"""







import argparse
import json
import os
import time
import zipfile

from get_ci_error_statistics import download_artifact, get_artifacts_links

from transformers import logging


A		=						logging.get_logger(__name__)

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Optional[int]     ,		lowerCamelCase_:							List[Any] ):
						"""simple docstring"""






						snake_case   :			Optional[Any]						    = set()
						snake_case   :			Optional[int]						    = []

						def parse_line(lowerCamelCase_:							Union[str, Any] ):
												for line in fp:
																		if isinstance(lowerCamelCase_     ,		lowerCamelCase_ ):
																								snake_case   :			Optional[Any]						    = line.decode("UTF-8" )
																		if "warnings summary (final)" in line:
																								continue
																		# This means we are outside the body of a warning
																		elif not line.startswith(" " ):
																								# process a single warning and move it to `selected_warnings`.
																								if len(lowerCamelCase_ ) > 0:
																														snake_case   :			str						    = "\n".join(lowerCamelCase_ )
																														# Only keep the warnings specified in `targets`
																														if any(f''': {x}: ''' in warning for x in targets ):
																																				selected_warnings.add(lowerCamelCase_ )
																														buffer.clear()
																								continue
																		else:
																								snake_case   :			Dict						    = line.strip()
																								buffer.append(lowerCamelCase_ )

						if from_gh:
												for filename in os.listdir(lowerCamelCase_ ):
																		snake_case   :			Tuple						    = os.path.join(lowerCamelCase_     ,		lowerCamelCase_ )
																		if not os.path.isdir(lowerCamelCase_ ):
																								# read the file
																								if filename != "warnings.txt":
																														continue
																								with open(lowerCamelCase_ ) as fp:
																														parse_line(lowerCamelCase_ )
						else:
												try:
																		with zipfile.ZipFile(lowerCamelCase_ ) as z:
																								for filename in z.namelist():
																														if not os.path.isdir(lowerCamelCase_ ):
																																				# read the file
																																				if filename != "warnings.txt":
																																										continue
																																				with z.open(lowerCamelCase_ ) as fp:
																																										parse_line(lowerCamelCase_ )
												except Exception:
																		logger.warning(
																		    f'''{artifact_path} is either an invalid zip file or something else wrong. This file is skipped.''' )

						return selected_warnings

def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Union[str, Any]     ,		lowerCamelCase_:							Any ):
						"""simple docstring"""






						snake_case   :			Dict						    = set()

						snake_case   :			str						    = [os.path.join(lowerCamelCase_     ,		lowerCamelCase_ ) for p in os.listdir(lowerCamelCase_ ) if (p.endswith(".zip" ) or from_gh)]
						for p in paths:
												selected_warnings.update(extract_warnings_from_single_artifact(lowerCamelCase_     ,		lowerCamelCase_ ) )

						return selected_warnings


if __name__ == "__main__":

			def  __SCREAMING_SNAKE_CASE		(			lowerCamelCase_:							Dict ):
									"""simple docstring"""






									return values.split("," )

			A		=						argparse.ArgumentParser()
			# Required parameters
			parser.add_argument('--workflow_run_id', type=str, required=True, help='A GitHub Actions workflow run id.')
			parser.add_argument(
			    '--output_dir',
			    type=str,
			    required=True,
			    help='Where to store the downloaded artifacts and other result files.',
			)
			parser.add_argument('--token', default=None, type=str, help='A token that has actions:read permission.')
			# optional parameters
			parser.add_argument(
			    '--targets',
			    default='DeprecationWarning,UserWarning,FutureWarning',
			    type=list_str,
			    help='Comma-separated list of target warning(s) which we want to extract.',
			)
			parser.add_argument(
			    '--from_gh',
			    action='store_true',
			    help='If running from a GitHub action workflow and collecting warnings from its artifacts.',
			)

			A		=						parser.parse_args()

			A		=						args.from_gh
			if from_gh:
						# The artifacts have to be downloaded using `actions/download-artifact@v3`
						pass
			else:
						os.makedirs(args.output_dir, exist_ok=True)

						# get download links
						A		=						get_artifacts_links(args.workflow_run_id, token=args.token)
						with open(os.path.join(args.output_dir, 'artifacts.json'), 'w', encoding='UTF-8') as fp:
									json.dump(artifacts, fp, ensure_ascii=False, indent=4)

						# download artifacts
						for idx, (name, url) in enumerate(artifacts.items()):
									print(name)
									print(url)
									print('=' * 8_0)
									download_artifact(name, url, args.output_dir, args.token)
									# Be gentle to GitHub
									time.sleep(1)

    # extract warnings from artifacts
			A		=						extract_warnings(args.output_dir, args.targets)
			A		=						sorted(selected_warnings)
			with open(os.path.join(args.output_dir, 'selected_warnings.json'), 'w', encoding='UTF-8') as fp:
						json.dump(selected_warnings, fp, ensure_ascii=False, indent=4)



