




'''simple docstring'''


import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    WavaVecaConfig,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaForPreTraining,
    WavaVecaProcessor,
    logging,
)
from transformers.models.wavaveca.modeling_wavaveca import WavaVecaForSequenceClassification


logging.set_verbosity_info()
_snake_case : Union[str, Any]  				=						logging.get_logger(__name__)

_snake_case : 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',
    'adapter_layer': 'encoder.layers.*.adapter_layer',
    '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',
    'pooling_layer.linear': 'projector',
    'pooling_layer.projection': 'classifier',
}
_snake_case : Union[str, Any]  				=						[
    'lm_head',
    'quantizer.weight_proj',
    'quantizer.codevectors',
    'project_q',
    'project_hid',
    'projector',
    'classifier',
]
def 						snake_case_       (UpperCamelCase			:    str ):





							'''simple docstring'''


							_a      	=							{}
							with open(UpperCamelCase					,		'''r''' ) as file:
														for line_number, line in enumerate(UpperCamelCase ):
																					_a      	=							line.strip()
																					if line:
																												_a      	=							line.split()
																												_a      	=							line_number
																												_a      	=							words[0]
																												_a      	=							value
							return result
def 						snake_case_       (UpperCamelCase			:    Optional[Any]					,		UpperCamelCase			:    str					,		UpperCamelCase			:    Tuple					,		UpperCamelCase			:    int					,		UpperCamelCase			:    int ):





							'''simple docstring'''


							for attribute in key.split('''.''' ):
														_a      	=							getattr(UpperCamelCase					,		UpperCamelCase )

							_a      	=							None
							for param_key in PARAM_MAPPING.keys():
														if full_name.endswith(UpperCamelCase ):
																					_a      	=							PARAM_MAPPING[full_name.split('''.''' )[-1]]
																					_a      	=							'''param'''

							if weight_type is not None and weight_type != "param":
														_a      	=							getattr(UpperCamelCase					,		UpperCamelCase ).shape
							elif weight_type is not None and weight_type == "param":
														_a      	=							hf_pointer
														for attribute in hf_param_name.split('''.''' ):
																					_a      	=							getattr(UpperCamelCase					,		UpperCamelCase )
														_a      	=							shape_pointer.shape

														# let's reduce dimension
														_a      	=							value[0]
							else:
														_a      	=							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":
														_a      	=							value
							elif weight_type == "weight_g":
														_a      	=							value
							elif weight_type == "weight_v":
														_a      	=							value
							elif weight_type == "bias":
														_a      	=							value
							elif weight_type == "param":
														for attribute in hf_param_name.split('''.''' ):
																					_a      	=							getattr(UpperCamelCase					,		UpperCamelCase )
														_a      	=							value
							else:
														_a      	=							value

							logger.info(f'{key + "." + weight_type if weight_type is not None else ""} was initialized from {full_name}.' )
def 						snake_case_       (UpperCamelCase			:    Dict					,		UpperCamelCase			:    Tuple					,		UpperCamelCase			:    Any					,		UpperCamelCase			:    Optional[int]					,		UpperCamelCase			:    int ):





							'''simple docstring'''


							_a      	=							None
							for param_key in PARAM_MAPPING.keys():
														if full_name.endswith(UpperCamelCase ):
																					_a      	=							PARAM_MAPPING[full_name.split('''.''' )[-1]]
																					_a      	=							'''param'''

							if weight_type is not None and weight_type != "param":
														_a      	=							'''.'''.join([key, weight_type] )
							elif weight_type is not None and weight_type == "param":
														_a      	=							'''.'''.join([key, hf_param_name] )
							else:
														_a      	=							key

							_a      	=							value if '''lm_head''' in full_key else value[0]


_snake_case : Tuple  				=						{
    'W_a': 'linear_1.weight',
    'W_b': 'linear_2.weight',
    'b_a': 'linear_1.bias',
    'b_b': 'linear_2.bias',
    'ln_W': 'norm.weight',
    'ln_b': 'norm.bias',
}
def 						snake_case_       (UpperCamelCase			:    Optional[int]					,		UpperCamelCase			:    Union[str, Any]					,		UpperCamelCase			:    Union[str, Any]=None					,		UpperCamelCase			:    str=None ):





							'''simple docstring'''


							_a      	=							False
							for key, mapped_key in MAPPING.items():
														_a      	=							'''wav2vec2.''' + 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]:
																					_a      	=							True
																					if "*" in mapped_key:
																												_a      	=							name.split(UpperCamelCase )[0].split('''.''' )[-2]
																												_a      	=							mapped_key.replace('''*'''					,		UpperCamelCase )
																					if "weight_g" in name:
																												_a      	=							'''weight_g'''
																					elif "weight_v" in name:
																												_a      	=							'''weight_v'''
																					elif "bias" in name:
																												_a      	=							'''bias'''
																					elif "weight" in name:
																												# TODO: don't match quantizer.weight_proj
																												_a      	=							'''weight'''
																					else:
																												_a      	=							None
																					if hf_dict is not None:
																												rename_dict(UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase )
																					else:
																												set_recursively(UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase )
																					return is_used
							return is_used
def 						snake_case_       (UpperCamelCase			:    Tuple					,		UpperCamelCase			:    Optional[Any]					,		UpperCamelCase			:    Any ):





							'''simple docstring'''


							_a      	=							[]
							_a      	=							fairseq_model.state_dict()

							_a      	=							hf_model.wavaveca.feature_extractor

							for name, value in fairseq_dict.items():
														_a      	=							False
														if "conv_layers" in name:
																					load_conv_layer(
																					    UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase					,		hf_model.config.feat_extract_norm == '''group'''					,		)
																					_a      	=							True
														else:
																					_a      	=							load_wavaveca_layer(UpperCamelCase					,		UpperCamelCase					,		UpperCamelCase )
														if not is_used:
																					unused_weights.append(UpperCamelCase )

							logger.warning(f'Unused weights: {unused_weights}' )
def 						snake_case_       (UpperCamelCase			:    Optional[Any]					,		UpperCamelCase			:    Optional[int]					,		UpperCamelCase			:    Any					,		UpperCamelCase			:    Tuple					,		UpperCamelCase			:    int ):





							'''simple docstring'''


							_a      	=							full_name.split('''conv_layers.''' )[-1]
							_a      	=							name.split('''.''' )
							_a      	=							int(items[0] )
							_a      	=							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.' )
																					_a      	=							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.' )
																					_a      	=							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.conv_layers[layer_id].layer_norm.bias.data.shape} was found.' )
																					_a      	=							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.conv_layers[layer_id].layer_norm.weight.data.shape} was found.' )
																					_a      	=							value
																					logger.info(f'Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.' )
							else:
														unused_weights.append(UpperCamelCase )
@torch.no_grad()
def 						snake_case_       (UpperCamelCase			:    Any					,		UpperCamelCase			:    Dict					,		UpperCamelCase			:    Dict=None					,		UpperCamelCase			:    Dict=None					,		UpperCamelCase			:    Any=True					,		UpperCamelCase			:    Any=False ):





							'''simple docstring'''


							if config_path is not None:
														_a      	=							WavaVecaConfig.from_pretrained(UpperCamelCase )
							else:
														_a      	=							WavaVecaConfig()

							if is_seq_class:
														_a      	=							read_txt_into_dict(UpperCamelCase )
														_a      	=							idalabel
														_a      	=							WavaVecaForSequenceClassification(UpperCamelCase )
														_a      	=							WavaVecaFeatureExtractor(
														    feature_size=1					,		sampling_rate=1_6000					,		padding_value=0					,		do_normalize=UpperCamelCase					,		return_attention_mask=UpperCamelCase					,		)
														feature_extractor.save_pretrained(UpperCamelCase )

							elif is_finetuned:
														if dict_path:
																					_a      	=							Dictionary.load(UpperCamelCase )

																					# important change bos & pad token id since CTC symbol is <pad> and
																					# not <s> as in fairseq
																					_a      	=							target_dict.pad_index
																					_a      	=							target_dict.bos_index
																					_a      	=							target_dict.eos_index
																					_a      	=							len(target_dict.symbols )
																					_a      	=							os.path.join(UpperCamelCase					,		'''vocab.json''' )
																					if not os.path.isdir(UpperCamelCase ):
																												logger.error('''--pytorch_dump_folder_path ({}) should be a directory'''.format(UpperCamelCase ) )
																												return
																					os.makedirs(UpperCamelCase					,		exist_ok=UpperCamelCase )
																					_a      	=							target_dict.indices

																					# fairseq has the <pad> and <s> switched
																					_a      	=							0
																					_a      	=							1
																					with open(UpperCamelCase					,		'''w'''					,		encoding='''utf-8''' ) as vocab_handle:
																												json.dump(UpperCamelCase					,		UpperCamelCase )
																					_a      	=							WavaVecaCTCTokenizer(
																					    UpperCamelCase					,		unk_token=target_dict.unk_word					,		pad_token=target_dict.pad_word					,		bos_token=target_dict.bos_word					,		eos_token=target_dict.eos_word					,		word_delimiter_token='''|'''					,		do_lower_case=UpperCamelCase					,		)
																					_a      	=							True if config.feat_extract_norm == '''layer''' else False
																					_a      	=							WavaVecaFeatureExtractor(
																					    feature_size=1					,		sampling_rate=1_6000					,		padding_value=0					,		do_normalize=UpperCamelCase					,		return_attention_mask=UpperCamelCase					,		)
																					_a      	=							WavaVecaProcessor(feature_extractor=UpperCamelCase					,		tokenizer=UpperCamelCase )
																					processor.save_pretrained(UpperCamelCase )

														_a      	=							WavaVecaForCTC(UpperCamelCase )
							else:
														_a      	=							WavaVecaForPreTraining(UpperCamelCase )

							if is_finetuned or is_seq_class:
														_a	,				_a	,				_a      	=							fairseq.checkpoint_utils.load_model_ensemble_and_task(
														    [checkpoint_path]					,		arg_overrides={'''data''': '''/'''.join(dict_path.split('''/''' )[:-1] )} )
							else:
														_a      	=							argparse.Namespace(task='''audio_pretraining''' )
														_a      	=							fairseq.tasks.setup_task(UpperCamelCase )

														_a	,				_a	,				_a      	=							fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]					,		task=UpperCamelCase )

							_a      	=							model[0].eval()

							recursively_load_weights(UpperCamelCase					,		UpperCamelCase					,		not is_finetuned )

							hf_wavavec.save_pretrained(UpperCamelCase )


if __name__ == "__main__":
				_snake_case : Union[str, Any]  				=						argparse.ArgumentParser()
				parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
				parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to fairseq checkpoint')
				parser.add_argument('--dict_path', default=None, type=str, help='Path to dict of fine-tuned model')
				parser.add_argument('--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'
				)
				parser.add_argument(
				    '--is_seq_class',
				    action='store_true',
				    help='Whether the model to convert is a fine-tuned sequence classification model or not',
				)
				_snake_case : str  				=						parser.parse_args()

				_snake_case : List[str]  				=						not args.not_finetuned and not args.is_seq_class
				convert_wavaveca_checkpoint(
				    args.checkpoint_path,
				    args.pytorch_dump_folder_path,
				    args.config_path,
				    args.dict_path,
				    is_finetuned,
				    args.is_seq_class,
				)







'''simple docstring'''


import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse('9.1.0'):
				_snake_case : Tuple  				=						{
				    'linear': PIL.Image.Resampling.BILINEAR,
				    'bilinear': PIL.Image.Resampling.BILINEAR,
				    'bicubic': PIL.Image.Resampling.BICUBIC,
				    'lanczos': PIL.Image.Resampling.LANCZOS,
				    'nearest': PIL.Image.Resampling.NEAREST,
				}
else:
				_snake_case : Any  				=						{
				    'linear': PIL.Image.LINEAR,
				    'bilinear': PIL.Image.BILINEAR,
				    'bicubic': PIL.Image.BICUBIC,
				    'lanczos': PIL.Image.LANCZOS,
				    'nearest': PIL.Image.NEAREST,
				}
def 						snake_case_       (UpperCamelCase			:    Optional[int] ):





							'''simple docstring'''


							_a      	=							(images / 2 + 0.5).clamp(0					,		1 )
							_a      	=							images.cpu().permute(0					,		2					,		3					,		1 ).float().numpy()
							_a      	=							numpy_to_pil(UpperCamelCase )
							return images

def 						snake_case_       (UpperCamelCase			:    str ):





							'''simple docstring'''


							if images.ndim == 3:
														_a      	=							images[None, ...]
							_a      	=							(images * 255).round().astype('''uint8''' )
							if images.shape[-1] == 1:
														# special case for grayscale (single channel) images
														_a      	=							[Image.fromarray(image.squeeze()					,		mode='''L''' ) for image in images]
							else:
														_a      	=							[Image.fromarray(UpperCamelCase ) for image in images]

							return pil_images

