





import argparse

import torch
from saprl.hub import distilhubert

from transformers import HubertConfig, HubertModel, WavaVecaFeatureExtractor, logging


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

_lowerCamelCase							=     {
    '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 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							int   ,				UpperCamelCase__:							Tuple   ,				UpperCamelCase__:							str   ,				UpperCamelCase__:							Tuple   ,				UpperCamelCase__:							List[str]      ):
		for attribute in key.split("""."""      ):
				SCREAMING_SNAKE_CASE__								= getattr(UpperCamelCase__   ,				UpperCamelCase__      )

		if weight_type is not None:
				SCREAMING_SNAKE_CASE__								= getattr(UpperCamelCase__   ,				UpperCamelCase__      ).shape
		else:
				SCREAMING_SNAKE_CASE__								= 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":
				SCREAMING_SNAKE_CASE__								= value
		elif weight_type == "weight_g":
				SCREAMING_SNAKE_CASE__								= value
		elif weight_type == "weight_v":
				SCREAMING_SNAKE_CASE__								= value
		elif weight_type == "bias":
				SCREAMING_SNAKE_CASE__								= value
		else:
				SCREAMING_SNAKE_CASE__								= value

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






def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							Dict   ,				UpperCamelCase__:							Dict      ):
		SCREAMING_SNAKE_CASE__								= []
		SCREAMING_SNAKE_CASE__								= fairseq_model.state_dict()

		SCREAMING_SNAKE_CASE__								= hf_model.feature_extractor

		for name, value in fairseq_dict.items():
				SCREAMING_SNAKE_CASE__								= False
				if "conv_layers" in name:
						load_conv_layer(
						    UpperCamelCase__   ,				UpperCamelCase__   ,				UpperCamelCase__   ,				UpperCamelCase__   ,				hf_model.config.feat_extract_norm == """group"""   ,				)
						SCREAMING_SNAKE_CASE__								= True
				else:
						for key, mapped_key in MAPPING.items():
								SCREAMING_SNAKE_CASE__								= mapped_key

								if key in name:
										SCREAMING_SNAKE_CASE__								= True
										if "*" in mapped_key:
												SCREAMING_SNAKE_CASE__								= name.split(UpperCamelCase__      )[0].split("""."""      )[-2]
												SCREAMING_SNAKE_CASE__								= mapped_key.replace("""*"""   ,				UpperCamelCase__      )
										if "weight_g" in name:
												SCREAMING_SNAKE_CASE__								= """weight_g"""
										elif "weight_v" in name:
												SCREAMING_SNAKE_CASE__								= """weight_v"""
										elif "weight" in name:
												SCREAMING_SNAKE_CASE__								= """weight"""
										elif "bias" in name:
												SCREAMING_SNAKE_CASE__								= """bias"""
										else:
												SCREAMING_SNAKE_CASE__								= None
										set_recursively(UpperCamelCase__   ,				UpperCamelCase__   ,				UpperCamelCase__   ,				UpperCamelCase__   ,				UpperCamelCase__      )
								continue
				if not is_used:
						unused_weights.append(UpperCamelCase__      )

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






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






def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							Dict      ):
		SCREAMING_SNAKE_CASE__								= HubertConfig()
		SCREAMING_SNAKE_CASE__								= model.config

		SCREAMING_SNAKE_CASE__								= fs_config.activation_dropout
		SCREAMING_SNAKE_CASE__								= False
		SCREAMING_SNAKE_CASE__								= fs_config.attention_dropout
		SCREAMING_SNAKE_CASE__								= False
		SCREAMING_SNAKE_CASE__								= eval(fs_config.extractor_conv_feature_layers      )
		SCREAMING_SNAKE_CASE__								= [x[0] for x in conv_layers]
		SCREAMING_SNAKE_CASE__								= [x[1] for x in conv_layers]
		SCREAMING_SNAKE_CASE__								= [x[2] for x in conv_layers]
		SCREAMING_SNAKE_CASE__								= """gelu"""
		SCREAMING_SNAKE_CASE__								= """layer""" if fs_config.extractor_mode == """layer_norm""" else """group"""
		SCREAMING_SNAKE_CASE__								= False
		SCREAMING_SNAKE_CASE__								= 0.0
		SCREAMING_SNAKE_CASE__								= 0.0
		SCREAMING_SNAKE_CASE__								= fs_config.activation_fn
		SCREAMING_SNAKE_CASE__								= fs_config.dropout
		SCREAMING_SNAKE_CASE__								= fs_config.encoder_embed_dim
		SCREAMING_SNAKE_CASE__								= 0.0_2
		SCREAMING_SNAKE_CASE__								= fs_config.encoder_ffn_embed_dim
		SCREAMING_SNAKE_CASE__								= 1e-5
		SCREAMING_SNAKE_CASE__								= 0.0
		SCREAMING_SNAKE_CASE__								= fs_config.encoder_attention_heads
		SCREAMING_SNAKE_CASE__								= fs_config.conv_pos_groups
		SCREAMING_SNAKE_CASE__								= fs_config.conv_pos
		SCREAMING_SNAKE_CASE__								= len(UpperCamelCase__      )
		SCREAMING_SNAKE_CASE__								= fs_config.encoder_layers

		return config






@torch.no_grad()
def 			SCREAMING_SNAKE_CASE__							(    UpperCamelCase__:							Tuple   ,				UpperCamelCase__:							Union[str, Any]=None      ):
		SCREAMING_SNAKE_CASE__								= distilhubert().model.model

		if config_path is not None:
				SCREAMING_SNAKE_CASE__								= HubertConfig.from_pretrained(UpperCamelCase__      )
		else:
				SCREAMING_SNAKE_CASE__								= convert_config(UpperCamelCase__      )
		SCREAMING_SNAKE_CASE__								= model.eval()

		SCREAMING_SNAKE_CASE__								= WavaVecaFeatureExtractor(
		    feature_size=1   ,				sampling_rate=16_000   ,				padding_value=0   ,				do_normalize=UpperCamelCase__   ,				return_attention_mask=UpperCamelCase__   ,				)
		SCREAMING_SNAKE_CASE__								= HubertModel(UpperCamelCase__      )

		recursively_load_weights(UpperCamelCase__   ,				UpperCamelCase__      )

		feature_extractor.save_pretrained(UpperCamelCase__      )
		hf_model.save_pretrained(UpperCamelCase__      )


if __name__ == "__main__":
	_lowerCamelCase							=     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')
	_lowerCamelCase							=     parser.parse_args()
	convert_hubert_checkpoint(args.pytorch_dump_folder_path, args.config_path)






from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


_lowerCamelCase							=     {
    'configuration_utils': ['GenerationConfig'],
    'streamers': ['TextIteratorStreamer', 'TextStreamer'],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase							=     [
	    'Constraint',
	    'ConstraintListState',
	    'DisjunctiveConstraint',
	    'PhrasalConstraint',
	]
	_lowerCamelCase							=     [
	    'BeamHypotheses',
	    'BeamScorer',
	    'BeamSearchScorer',
	    'ConstrainedBeamSearchScorer',
	]
	_lowerCamelCase							=     [
	    'EpsilonLogitsWarper',
	    'EtaLogitsWarper',
	    'ForcedBOSTokenLogitsProcessor',
	    'ForcedEOSTokenLogitsProcessor',
	    'HammingDiversityLogitsProcessor',
	    'InfNanRemoveLogitsProcessor',
	    'LogitsProcessor',
	    'LogitsProcessorList',
	    'LogitsWarper',
	    'MinLengthLogitsProcessor',
	    'MinNewTokensLengthLogitsProcessor',
	    'NoBadWordsLogitsProcessor',
	    'NoRepeatNGramLogitsProcessor',
	    'PrefixConstrainedLogitsProcessor',
	    'RepetitionPenaltyLogitsProcessor',
	    'SequenceBiasLogitsProcessor',
	    'EncoderRepetitionPenaltyLogitsProcessor',
	    'TemperatureLogitsWarper',
	    'TopKLogitsWarper',
	    'TopPLogitsWarper',
	    'TypicalLogitsWarper',
	    'EncoderNoRepeatNGramLogitsProcessor',
	    'ExponentialDecayLengthPenalty',
	    'LogitNormalization',
	]
	_lowerCamelCase							=     [
	    'MaxNewTokensCriteria',
	    'MaxLengthCriteria',
	    'MaxTimeCriteria',
	    'StoppingCriteria',
	    'StoppingCriteriaList',
	    'validate_stopping_criteria',
	]
	_lowerCamelCase							=     [
	    'GenerationMixin',
	    'top_k_top_p_filtering',
	    'GreedySearchEncoderDecoderOutput',
	    'GreedySearchDecoderOnlyOutput',
	    'SampleEncoderDecoderOutput',
	    'SampleDecoderOnlyOutput',
	    'BeamSearchEncoderDecoderOutput',
	    'BeamSearchDecoderOnlyOutput',
	    'BeamSampleEncoderDecoderOutput',
	    'BeamSampleDecoderOnlyOutput',
	    'ContrastiveSearchEncoderDecoderOutput',
	    'ContrastiveSearchDecoderOnlyOutput',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase							=     [
	    'TFForcedBOSTokenLogitsProcessor',
	    'TFForcedEOSTokenLogitsProcessor',
	    'TFLogitsProcessor',
	    'TFLogitsProcessorList',
	    'TFLogitsWarper',
	    'TFMinLengthLogitsProcessor',
	    'TFNoBadWordsLogitsProcessor',
	    'TFNoRepeatNGramLogitsProcessor',
	    'TFRepetitionPenaltyLogitsProcessor',
	    'TFTemperatureLogitsWarper',
	    'TFTopKLogitsWarper',
	    'TFTopPLogitsWarper',
	    'TFForceTokensLogitsProcessor',
	    'TFSuppressTokensAtBeginLogitsProcessor',
	    'TFSuppressTokensLogitsProcessor',
	]
	_lowerCamelCase							=     [
	    'TFGenerationMixin',
	    'tf_top_k_top_p_filtering',
	    'TFGreedySearchDecoderOnlyOutput',
	    'TFGreedySearchEncoderDecoderOutput',
	    'TFSampleEncoderDecoderOutput',
	    'TFSampleDecoderOnlyOutput',
	    'TFBeamSearchEncoderDecoderOutput',
	    'TFBeamSearchDecoderOnlyOutput',
	    'TFBeamSampleEncoderDecoderOutput',
	    'TFBeamSampleDecoderOnlyOutput',
	    'TFContrastiveSearchEncoderDecoderOutput',
	    'TFContrastiveSearchDecoderOnlyOutput',
	]

try:
	if not is_flax_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	_lowerCamelCase							=     [
	    'FlaxForcedBOSTokenLogitsProcessor',
	    'FlaxForcedEOSTokenLogitsProcessor',
	    'FlaxLogitsProcessor',
	    'FlaxLogitsProcessorList',
	    'FlaxLogitsWarper',
	    'FlaxMinLengthLogitsProcessor',
	    'FlaxTemperatureLogitsWarper',
	    'FlaxTopKLogitsWarper',
	    'FlaxTopPLogitsWarper',
	]
	_lowerCamelCase							=     [
	    'FlaxGenerationMixin',
	    'FlaxGreedySearchOutput',
	    'FlaxSampleOutput',
	    'FlaxBeamSearchOutput',
	]

if TYPE_CHECKING:
	from .configuration_utils import GenerationConfig
	from .streamers import TextIteratorStreamer, TextStreamer

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
		from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
		from .logits_process import (
		    EncoderNoRepeatNGramLogitsProcessor,
		    EncoderRepetitionPenaltyLogitsProcessor,
		    EpsilonLogitsWarper,
		    EtaLogitsWarper,
		    ExponentialDecayLengthPenalty,
		    ForcedBOSTokenLogitsProcessor,
		    ForcedEOSTokenLogitsProcessor,
		    HammingDiversityLogitsProcessor,
		    InfNanRemoveLogitsProcessor,
		    LogitNormalization,
		    LogitsProcessor,
		    LogitsProcessorList,
		    LogitsWarper,
		    MinLengthLogitsProcessor,
		    MinNewTokensLengthLogitsProcessor,
		    NoBadWordsLogitsProcessor,
		    NoRepeatNGramLogitsProcessor,
		    PrefixConstrainedLogitsProcessor,
		    RepetitionPenaltyLogitsProcessor,
		    SequenceBiasLogitsProcessor,
		    TemperatureLogitsWarper,
		    TopKLogitsWarper,
		    TopPLogitsWarper,
		    TypicalLogitsWarper,
		)
		from .stopping_criteria import (
		    MaxLengthCriteria,
		    MaxNewTokensCriteria,
		    MaxTimeCriteria,
		    StoppingCriteria,
		    StoppingCriteriaList,
		    validate_stopping_criteria,
		)
		from .utils import (
		    BeamSampleDecoderOnlyOutput,
		    BeamSampleEncoderDecoderOutput,
		    BeamSearchDecoderOnlyOutput,
		    BeamSearchEncoderDecoderOutput,
		    ContrastiveSearchDecoderOnlyOutput,
		    ContrastiveSearchEncoderDecoderOutput,
		    GenerationMixin,
		    GreedySearchDecoderOnlyOutput,
		    GreedySearchEncoderDecoderOutput,
		    SampleDecoderOnlyOutput,
		    SampleEncoderDecoderOutput,
		    top_k_top_p_filtering,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tf_logits_process import (
		    TFForcedBOSTokenLogitsProcessor,
		    TFForcedEOSTokenLogitsProcessor,
		    TFForceTokensLogitsProcessor,
		    TFLogitsProcessor,
		    TFLogitsProcessorList,
		    TFLogitsWarper,
		    TFMinLengthLogitsProcessor,
		    TFNoBadWordsLogitsProcessor,
		    TFNoRepeatNGramLogitsProcessor,
		    TFRepetitionPenaltyLogitsProcessor,
		    TFSuppressTokensAtBeginLogitsProcessor,
		    TFSuppressTokensLogitsProcessor,
		    TFTemperatureLogitsWarper,
		    TFTopKLogitsWarper,
		    TFTopPLogitsWarper,
		)
		from .tf_utils import (
		    TFBeamSampleDecoderOnlyOutput,
		    TFBeamSampleEncoderDecoderOutput,
		    TFBeamSearchDecoderOnlyOutput,
		    TFBeamSearchEncoderDecoderOutput,
		    TFContrastiveSearchDecoderOnlyOutput,
		    TFContrastiveSearchEncoderDecoderOutput,
		    TFGenerationMixin,
		    TFGreedySearchDecoderOnlyOutput,
		    TFGreedySearchEncoderDecoderOutput,
		    TFSampleDecoderOnlyOutput,
		    TFSampleEncoderDecoderOutput,
		    tf_top_k_top_p_filtering,
		)

	try:
		if not is_flax_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .flax_logits_process import (
		    FlaxForcedBOSTokenLogitsProcessor,
		    FlaxForcedEOSTokenLogitsProcessor,
		    FlaxLogitsProcessor,
		    FlaxLogitsProcessorList,
		    FlaxLogitsWarper,
		    FlaxMinLengthLogitsProcessor,
		    FlaxTemperatureLogitsWarper,
		    FlaxTopKLogitsWarper,
		    FlaxTopPLogitsWarper,
		)
		from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
	import sys

	_lowerCamelCase							=     _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)