



"""simple docstring"""

import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

from transformers import (
    UniSpeechConfig,
    UniSpeechForCTC,
    UniSpeechForPreTraining,
    WavaVecaFeatureExtractor,
    WavaVecaPhonemeCTCTokenizer,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
lowercase__    :  List[str]							    =	logging.get_logger(__name__)

lowercase__    :  str							    =	{
    '''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''': '''ctc_proj''',
    '''mask_emb''': '''masked_spec_embed''',
}
lowercase__    :  Tuple							    =	[
    '''ctc_proj''',
    '''quantizer.weight_proj''',
    '''quantizer.codevectors''',
    '''project_q''',
    '''project_hid''',
]




def __lowercase			(							_a		,			_a		,			_a		,			_a		,			_a		,			_a    ):
    for attribute in key.split('''.'''    ):
        if is_finetuned:
            if attribute in ["quantizer", "project_q", "project_hid"]:
                # those layers are only relevant for pretraining and should be dropped
                return

            if attribute == "ctc_proj":
                # we should rename `ctc_proj` to `lm_head` for fine-tuned phoneme models
                snake_case_ :     str				=	'lm_head'

        snake_case_ :     List[str]				=	getattr(_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE    )

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

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




def __lowercase			(							_a		,			_a		,			_a    ):
    snake_case_ :     str				=	[]
    snake_case_ :     Dict				=	fairseq_model.state_dict()

    snake_case_ :     Any				=	hf_model.unispeech.feature_extractor

    for name, value in fairseq_dict.items():
        snake_case_ :     Tuple				=	False
        if "conv_layers" in name:
            load_conv_layer(
                _SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			hf_model.config.feat_extract_norm == '''group'''		,			)
            snake_case_ :     List[str]				=	True
        else:
            for key, mapped_key in MAPPING.items():
                snake_case_ :     Tuple				=	'unispeech.' + 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]:
                    snake_case_ :     Optional[Any]				=	True
                    if "*" in mapped_key:
                        snake_case_ :     List[Any]				=	name.split(_SCREAMING_SNAKE_CASE    )[0].split('''.'''    )[-2]
                        snake_case_ :     Dict				=	mapped_key.replace('''*'''		,			_SCREAMING_SNAKE_CASE    )
                    if "weight_g" in name:
                        snake_case_ :     Tuple				=	'weight_g'
                    elif "weight_v" in name:
                        snake_case_ :     str				=	'weight_v'
                    elif "bias" in name:
                        snake_case_ :     Optional[int]				=	'bias'
                    elif "weight" in name:
                        # TODO: don't match quantizer.weight_proj
                        snake_case_ :     List[Any]				=	'weight'
                    else:
                        snake_case_ :     List[Any]				=	None
                    set_recursively(_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE    )
                continue
        if not is_used:
            unused_weights.append(_SCREAMING_SNAKE_CASE    )

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




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




@torch.no_grad()
def __lowercase			(							_a		,			_a		,			_a=None		,			_a=None		,			_a=True    ):
    if config_path is not None:
        snake_case_ :     Optional[int]				=	UniSpeechConfig.from_pretrained(_SCREAMING_SNAKE_CASE    )
    else:
        snake_case_ :     List[str]				=	UniSpeechConfig()

    if is_finetuned:
        if dict_path:
            snake_case_ :     Tuple				=	Dictionary.load_from_json(_SCREAMING_SNAKE_CASE    )

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

            # fairseq has the <pad> and <s> switched
            snake_case_ :     List[Any]				=	42
            snake_case_ :     Optional[Any]				=	43
            with open(_SCREAMING_SNAKE_CASE		,			'''w'''		,			encoding='''utf-8'''    ) as vocab_handle:
                json.dump(_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE    )
            snake_case_ :     Dict				=	WavaVecaPhonemeCTCTokenizer(
                _SCREAMING_SNAKE_CASE		,			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=_SCREAMING_SNAKE_CASE		,			)
            snake_case_ :     Union[str, Any]				=	True if config.feat_extract_norm == 'layer' else False
            snake_case_ :     str				=	WavaVecaFeatureExtractor(
                feature_size=1		,			sampling_rate=16_000		,			padding_value=0		,			do_normalize=_SCREAMING_SNAKE_CASE		,			return_attention_mask=_SCREAMING_SNAKE_CASE		,			)
            snake_case_ :     List[Any]				=	WavaVecaProcessor(feature_extractor=_SCREAMING_SNAKE_CASE		,			tokenizer=_SCREAMING_SNAKE_CASE    )
            processor.save_pretrained(_SCREAMING_SNAKE_CASE    )

        snake_case_ :     Dict				=	UniSpeechForCTC(_SCREAMING_SNAKE_CASE    )
    else:
        snake_case_ :     List[Any]				=	UniSpeechForPreTraining(_SCREAMING_SNAKE_CASE    )

    if is_finetuned:
        snake_case_ :     List[str]				=	fairseq.checkpoint_utils.load_model_ensemble_and_task(
            [checkpoint_path]		,			arg_overrides={'''data''': '''/'''.join(dict_path.split('''/'''    )[:-1]    ), '''w2v_path''': checkpoint_path}    )
    else:
        snake_case_ :     List[str]				=	fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path]    )

    snake_case_ :     str				=	model[0].eval()

    recursively_load_weights(_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE		,			_SCREAMING_SNAKE_CASE    )

    hf_unispeech.save_pretrained(_SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
     lowercase__    :  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'''
     )
     lowercase__    :  Tuple							    =	parser.parse_args()
     convert_unispeech_checkpoint(
         args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
     )






"""simple docstring"""
def 						__a   (	_SCREAMING_SNAKE_CASE = 1000000   )	->int:
  a__:     int										=	limit + 1
  a__:     Optional[int]										=	[0] * limit
  for first_term in range(1		,     _SCREAMING_SNAKE_CASE   ):
    for n in range(_SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE		,     _SCREAMING_SNAKE_CASE   ):
      a__:     List[Any]										=	first_term + n / first_term
      if common_difference % 4:  # d must be divisble by 4
        continue
      else:
        common_difference /= 4
        if (
            first_term > common_difference
            and first_term < 4 * common_difference
        ):  # since x,y,z are positive integers
          frequency[n] += 1  # so z>0 and a>d ,also 4d<a

  a__:     Any										=	sum(1 for x in frequency[1:limit] if x == 10   )

  return count


if __name__ == "__main__":
 print(f"{solution() = }")


