



"""simple docstring"""






import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


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

__UpperCAmelCase      :      Dict   					=					{
    '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',
    'mask_emb': 'masked_spec_embed',
}
def      A      (			_A,     _A,     _A,     _A,     _A							):




					"""simple docstring"""
					for attribute in key.split("."							):
										snake_case_						:str								=						getattr(_A,     _A							)

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

					logger.info(F'''{key + '.' + weight_type if weight_type is not None else ''} was initialized from {full_name}.'''							)
def      A      (			_A,     _A							):




					"""simple docstring"""
					snake_case_						:List[Any]								=						[]
					snake_case_						:List[str]								=						fairseq_model.state_dict()

					snake_case_						:Union[str, Any]								=						hf_model.feature_extractor

					for name, value in fairseq_dict.items():
										snake_case_						:Union[str, Any]								=						False
										if "conv_layers" in name:
															load_conv_layer(
															    _A,     _A,     _A,     _A,     hf_model.config.feat_extract_norm == "group",     )
															snake_case_						:str								=						True
										else:
															for key, mapped_key in MAPPING.items():
																				snake_case_						:str								=						mapped_key

																				if key in name:
																									snake_case_						:Dict								=						True
																									if "*" in mapped_key:
																														snake_case_						:Dict								=						name.split(_A							)[0].split("."							)[-2]
																														snake_case_						:List[Any]								=						mapped_key.replace("*",     _A							)
																									if "weight_g" in name:
																														snake_case_						:Any								=						"weight_g"
																									elif "weight_v" in name:
																														snake_case_						:Union[str, Any]								=						"weight_v"
																									elif "weight" in name:
																														snake_case_						:List[str]								=						"weight"
																									elif "bias" in name:
																														snake_case_						:Optional[int]								=						"bias"
																									else:
																														snake_case_						:Any								=						None
																									set_recursively(_A,     _A,     _A,     _A,     _A							)
																				continue
										if not is_used:
															unused_weights.append(_A							)

					logger.warning(F'''Unused weights: {unused_weights}'''							)
def      A      (			_A,     _A,     _A,     _A,     _A							):




					"""simple docstring"""
					snake_case_						:Optional[int]								=						full_name.split("conv_layers."							)[-1]
					snake_case_						:Union[str, Any]								=						name.split("."							)
					snake_case_						:Dict								=						int(items[0]							)
					snake_case_						:str								=						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_						:List[str]								=						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_						: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."
															)
															snake_case_						:str								=						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_						:Any								=						value
															logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''							)
					else:
										unused_weights.append(_A							)
def      A      (			_A							):




					"""simple docstring"""
					snake_case_						:List[str]								=						HubertConfig()
					snake_case_						:str								=						model.config

					snake_case_						:Optional[int]								=						fs_config.activation_dropout
					snake_case_						:Dict								=						False
					snake_case_						:Optional[Any]								=						fs_config.attention_dropout
					snake_case_						:Optional[Any]								=						False
					snake_case_						:Tuple								=						eval(fs_config.extractor_conv_feature_layers							)
					snake_case_						:Dict								=						[x[0] for x in conv_layers]
					snake_case_						:Dict								=						[x[1] for x in conv_layers]
					snake_case_						:Optional[Any]								=						[x[2] for x in conv_layers]
					snake_case_						:int								=						"gelu"
					snake_case_						:Dict								=						"layer" if fs_config.extractor_mode == "layer_norm" else "group"
					snake_case_						:Dict								=						False
					snake_case_						:List[Any]								=						0.0
					snake_case_						:Tuple								=						0.0
					snake_case_						:Optional[int]								=						fs_config.activation_fn
					snake_case_						:Any								=						fs_config.dropout
					snake_case_						:int								=						fs_config.encoder_embed_dim
					snake_case_						:List[str]								=						0.02
					snake_case_						:Optional[int]								=						fs_config.encoder_ffn_embed_dim
					snake_case_						:Union[str, Any]								=						1e-5
					snake_case_						:Tuple								=						0.0
					snake_case_						:List[str]								=						fs_config.encoder_attention_heads
					snake_case_						:Optional[Any]								=						fs_config.conv_pos_groups
					snake_case_						:Tuple								=						fs_config.conv_pos
					snake_case_						:Dict								=						len(_A							)
					snake_case_						:Optional[int]								=						fs_config.encoder_layers

					return config
@torch.no_grad()
def      A      (			_A,     _A=None							):




					"""simple docstring"""
					snake_case_						:Optional[Any]								=						distilhubert().model.model

					if config_path is not None:
										snake_case_						:Dict								=						HubertConfig.from_pretrained(_A							)
					else:
										snake_case_						:Union[str, Any]								=						convert_config(_A							)
					snake_case_						:str								=						model.eval()

					snake_case_						:Optional[int]								=						WavaVecaFeatureExtractor(
					    feature_size=1,     sampling_rate=16_000,     padding_value=0,     do_normalize=_A,     return_attention_mask=_A,     )
					snake_case_						:List[str]								=						HubertModel(_A							)

					recursively_load_weights(_A,     _A							)

					feature_extractor.save_pretrained(_A							)
					hf_model.save_pretrained(_A							)


if __name__ == "__main__":
				__UpperCAmelCase      :      List[str]   					=					argparse.ArgumentParser()
				parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
				parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
				__UpperCAmelCase      :      Any   					=					parser.parse_args()
				convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)






"""simple docstring"""






from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


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

__UpperCAmelCase      :      str   					=					{
    'google/bigbird-pegasus-large-arxiv': (
        'https://huggingface.co/google/bigbird-pegasus-large-arxiv/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-pubmed': (
        'https://huggingface.co/google/bigbird-pegasus-large-pubmed/resolve/main/config.json'
    ),
    'google/bigbird-pegasus-large-bigpatent': (
        'https://huggingface.co/google/bigbird-pegasus-large-bigpatent/resolve/main/config.json'
    ),
    # See all BigBirdPegasus models at https://huggingface.co/models?filter=bigbird_pegasus
}

class 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					a__        =     'bigbird_pegasus'
					a__        =     ['past_key_values']
					a__        =     {
					    'num_attention_heads': 'encoder_attention_heads',
					    'hidden_size': 'd_model',
					    'attention_probs_dropout_prob': 'attention_dropout',
					}







					def __init__(      self    ,       a=9_61_03    ,       a=40_96    ,       a=16    ,       a=40_96    ,       a=16    ,       a=16    ,       a=40_96    ,       a=16    ,       a=0.0    ,       a=0.0    ,       a=True    ,       a=True    ,       a="gelu_new"    ,       a=10_24    ,       a=0.1    ,       a=0.0    ,       a=0.0    ,       a=0.02    ,       a=2    ,       a=0.0    ,       a=True    ,       a=0    ,       a=2    ,       a=1    ,       a="block_sparse"    ,       a=64    ,       a=3    ,       a=False    ,       **a    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						vocab_size
										snake_case_						:Any								=						max_position_embeddings
										snake_case_						:List[Any]								=						d_model
										snake_case_						:Any								=						encoder_ffn_dim
										snake_case_						:str								=						encoder_layers
										snake_case_						:Union[str, Any]								=						encoder_attention_heads
										snake_case_						:Any								=						decoder_ffn_dim
										snake_case_						:Dict								=						decoder_layers
										snake_case_						:List[Any]								=						decoder_attention_heads
										snake_case_						:Any								=						dropout
										snake_case_						:Union[str, Any]								=						attention_dropout
										snake_case_						:Optional[int]								=						activation_dropout
										snake_case_						:int								=						activation_function
										snake_case_						:int								=						init_std
										snake_case_						:Union[str, Any]								=						encoder_layerdrop
										snake_case_						:Dict								=						decoder_layerdrop
										snake_case_						:Any								=						classifier_dropout
										snake_case_						:Any								=						use_cache
										snake_case_						:Any								=						encoder_layers
										snake_case_						:int								=						scale_embedding  # scale factor will be sqrt(d_model) if True

										# extra config
										snake_case_						:Optional[int]								=						attention_type
										snake_case_						:int								=						block_size
										snake_case_						:Optional[int]								=						num_random_blocks
										snake_case_						:List[str]								=						use_bias

										super().__init__(
										    pad_token_id=a    ,       bos_token_id=a    ,       eos_token_id=a    ,       is_encoder_decoder=a    ,       decoder_start_token_id=a    ,       **a    ,       )



class 	__lowerCAmelCase							(__UpperCamelCase  ):


					'''simple docstring'''


					@property
					def 		_a (      self   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:int								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															    ]   )

															if self.use_past:
																				snake_case_						:int								=						{0: "batch"}
																				snake_case_						:List[str]								=						{0: "batch", 1: "past_decoder_sequence + sequence"}
															else:
																				snake_case_						:Optional[int]								=						{0: "batch", 1: "decoder_sequence"}
																				snake_case_						:Optional[Any]								=						{0: "batch", 1: "decoder_sequence"}

															if self.use_past:
																				self.fill_with_past_key_values_(a    ,       direction="inputs"   )
										elif self.task == "causal-lm":
															# TODO: figure this case out.
															snake_case_						:List[str]								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															    ]   )
															if self.use_past:
																				snake_case_	,  snake_case_						:Tuple								=						self.num_layers
																				for i in range(a   ):
																									snake_case_						:int								=						{0: "batch", 2: "past_sequence + sequence"}
																									snake_case_						:List[Any]								=						{0: "batch", 2: "past_sequence + sequence"}
										else:
															snake_case_						:List[Any]								=						OrderedDict(
															    [
															        ("input_ids", {0: "batch", 1: "encoder_sequence"}),
															        ("attention_mask", {0: "batch", 1: "encoder_sequence"}),
															        ("decoder_input_ids", {0: "batch", 1: "decoder_sequence"}),
															        ("decoder_attention_mask", {0: "batch", 1: "decoder_sequence"}),
															    ]   )

										return common_inputs







					@property
					def 		_a (      self   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:Optional[int]								=						super().outputs
										else:
															snake_case_						:List[str]								=						super(a    ,       self   ).outputs
															if self.use_past:
																				snake_case_	,  snake_case_						:int								=						self.num_layers
																				for i in range(a   ):
																									snake_case_						:Optional[Any]								=						{0: "batch", 2: "past_sequence + sequence"}
																									snake_case_						:Tuple								=						{0: "batch", 2: "past_sequence + sequence"}
										return common_outputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:int								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )

										# Generate decoder inputs
										snake_case_						:List[str]								=						seq_length if not self.use_past else 1
										snake_case_						:str								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )
										snake_case_						:Union[str, Any]								=						{F'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
										snake_case_						:Tuple								=						dict(**a    ,       **a   )

										if self.use_past:
															if not is_torch_available():
																				raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
															else:
																				import torch
															snake_case_	,  snake_case_						:Optional[int]								=						common_inputs["input_ids"].shape
															snake_case_						:Optional[int]								=						common_inputs["decoder_input_ids"].shape[1]
															snake_case_	,  snake_case_						:Optional[int]								=						self.num_attention_heads
															snake_case_						:str								=						(
															    batch,
															    num_encoder_attention_heads,
															    encoder_seq_length,
															    self._config.hidden_size // num_encoder_attention_heads,
															)
															snake_case_						:Any								=						decoder_seq_length + 3
															snake_case_						:int								=						(
															    batch,
															    num_decoder_attention_heads,
															    decoder_past_length,
															    self._config.hidden_size // num_decoder_attention_heads,
															)

															snake_case_						:Union[str, Any]								=						torch.cat(
															    [common_inputs["decoder_attention_mask"], torch.ones(a    ,       a   )]    ,       dim=1   )

															snake_case_						:Union[str, Any]								=						[]
															# If the number of encoder and decoder layers are present in the model configuration, both are considered
															snake_case_	,  snake_case_						:Tuple								=						self.num_layers
															snake_case_						:int								=						min(a    ,       a   )
															snake_case_						:Optional[Any]								=						max(a    ,       a   ) - min_num_layers
															snake_case_						:List[Any]								=						"encoder" if num_encoder_layers > num_decoder_layers else "decoder"

															for _ in range(a   ):
																				common_inputs["past_key_values"].append(
																				    (
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				        torch.zeros(a   ),
																				    )   )
															# TODO: test this.
															snake_case_						:Any								=						encoder_shape if remaining_side_name == "encoder" else decoder_shape
															for _ in range(a    ,       a   ):
																				common_inputs["past_key_values"].append((torch.zeros(a   ), torch.zeros(a   ))   )
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:List[str]								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
										    a    ,       a    ,       a    ,       a    ,       a   )

										if self.use_past:
															if not is_torch_available():
																				raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."   )
															else:
																				import torch
															snake_case_	,  snake_case_						:Union[str, Any]								=						common_inputs["input_ids"].shape
															# Not using the same length for past_key_values
															snake_case_						:str								=						seqlen + 2
															snake_case_	,  snake_case_						:Union[str, Any]								=						self.num_layers
															snake_case_	,  snake_case_						:Any								=						self.num_attention_heads
															snake_case_						:List[str]								=						(
															    batch,
															    num_encoder_attention_heads,
															    past_key_values_length,
															    self._config.hidden_size // num_encoder_attention_heads,
															)

															snake_case_						:Optional[Any]								=						common_inputs["attention_mask"].dtype
															snake_case_						:Tuple								=						torch.cat(
															    [common_inputs["attention_mask"], torch.ones(a    ,       a    ,       dtype=a   )]    ,       dim=1   )
															snake_case_						:Optional[int]								=						[
															    (torch.zeros(a   ), torch.zeros(a   )) for _ in range(a   )
															]
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										snake_case_						:Optional[int]								=						compute_effective_axis_dimension(
										    a    ,       fixed_dimension=OnnxConfig.default_fixed_batch    ,       num_token_to_add=0   )

										# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
										snake_case_						:List[Any]								=						tokenizer.num_special_tokens_to_add(a   )
										snake_case_						:Tuple								=						compute_effective_axis_dimension(
										    a    ,       fixed_dimension=OnnxConfig.default_fixed_sequence    ,       num_token_to_add=a   )

										# Generate dummy inputs according to compute batch and sequence
										snake_case_						:Any								=						[" ".join([tokenizer.unk_token]   ) * seq_length] * batch_size
										snake_case_						:List[str]								=						dict(tokenizer(a    ,       return_tensors=a   )   )
										return common_inputs







					def 		_a (      self    ,       a    ,       a = -1    ,       a = -1    ,       a = False    ,       a = None    ,       ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:Dict								=						self._generate_dummy_inputs_for_default_and_seqaseq_lm(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )

										elif self.task == "causal-lm":
															snake_case_						:Optional[Any]								=						self._generate_dummy_inputs_for_causal_lm(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )
										else:
															snake_case_						:List[Any]								=						self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
															    a    ,       batch_size=a    ,       seq_length=a    ,       is_pair=a    ,       framework=a   )

										return common_inputs







					def 		_a (      self    ,       a    ,       a    ,       a    ,       a   ):




										"""simple docstring"""





										if self.task in ["default", "seq2seq-lm"]:
															snake_case_						:str								=						super()._flatten_past_key_values_(a    ,       a    ,       a    ,       a   )
										else:
															snake_case_						:Any								=						super(a    ,       self   )._flatten_past_key_values_(
															    a    ,       a    ,       a    ,       a   )

