
import argparse
import json
import os

import fairseq
import torch
from fairseq.data import Dictionary

# Register SEW's fairseq modules
from sew_asapp import tasks  # noqa: F401

from transformers import (
    SEWConfig,
    SEWForCTC,
    SEWModel,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaProcessor,
    logging,
)


logging.set_verbosity_info()
lowerCamelCase__     :		List[Any]        = logging.get_logger(__name__)

lowerCamelCase__     :		Dict        = {
    """post_extract_proj""": """feature_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.upsample.0""": """encoder.upsample.projection""",
    """encoder.layer_norm""": """encoder.layer_norm""",
    """w2v_model.layer_norm""": """layer_norm""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )		->						Dict:


  '''simple docstring'''






  for attribute in key.split(""".""" ):
    lowercase__	: Union[str, Any]     		=   getattr(lowercase_			,			lowercase_ )

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

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

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_ )		->						List[Any]:


  '''simple docstring'''






  lowercase__	: List[Any]     		=   []
  lowercase__	: Union[str, Any]     		=   fairseq_model.state_dict()

  lowercase__	: Any     		=   hf_model.sew.feature_extractor if is_finetuned else hf_model.feature_extractor

  for name, value in fairseq_dict.items():
    lowercase__	: Optional[int]     		=   False
    if "conv_layers" in name:
      load_conv_layer(
          lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			hf_model.config.feat_extract_norm == """group"""			,			)
      lowercase__	: str     		=   True
    else:
      for key, mapped_key in MAPPING.items():
        lowercase__	: Optional[int]     		=   """sew.""" + mapped_key if (is_finetuned and mapped_key != """lm_head""") else mapped_key

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

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

def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_			,			lowercase_			,			lowercase_ )		->						List[Any]:


  '''simple docstring'''






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

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						Optional[Any]:


  '''simple docstring'''






  lowercase__	: Optional[Any]     		=   SEWConfig()
  if is_finetuned:
    lowercase__	: Any     		=   model.wav_encoder.wav_model.cfg
  else:
    lowercase__	: Dict     		=   model.cfg

  lowercase__	: Optional[Any]     		=   fs_config.conv_bias
  lowercase__	: Tuple     		=   eval(fs_config.conv_feature_layers )
  lowercase__	: List[str]     		=   [x[0] for x in conv_layers]
  lowercase__	: Dict     		=   [x[1] for x in conv_layers]
  lowercase__	: Tuple     		=   [x[2] for x in conv_layers]
  lowercase__	: List[str]     		=   """gelu"""
  lowercase__	: Union[str, Any]     		=   """layer""" if fs_config.extractor_mode == """layer_norm""" else """group"""
  lowercase__	: Union[str, Any]     		=   0.0
  lowercase__	: Tuple     		=   fs_config.activation_fn.name
  lowercase__	: Tuple     		=   fs_config.encoder_embed_dim
  lowercase__	: List[str]     		=   0.02
  lowercase__	: Optional[Any]     		=   fs_config.encoder_ffn_embed_dim
  lowercase__	: Optional[Any]     		=   1E-5
  lowercase__	: List[Any]     		=   fs_config.encoder_layerdrop
  lowercase__	: Any     		=   fs_config.encoder_attention_heads
  lowercase__	: Any     		=   fs_config.conv_pos_groups
  lowercase__	: Dict     		=   fs_config.conv_pos
  lowercase__	: List[Any]     		=   len(lowercase_ )
  lowercase__	: Union[str, Any]     		=   fs_config.encoder_layers
  lowercase__	: Optional[Any]     		=   fs_config.squeeze_factor

  # take care of any params that are overridden by the Wav2VecCtc model
  if is_finetuned:
    lowercase__	: Optional[Any]     		=   model.cfg
    lowercase__	: Union[str, Any]     		=   fs_config.final_dropout
    lowercase__	: int     		=   fs_config.layerdrop
  lowercase__	: int     		=   fs_config.activation_dropout
  lowercase__	: Optional[Any]     		=   fs_config.mask_prob > 0 or fs_config.mask_channel_prob > 0
  lowercase__	: Tuple     		=   fs_config.attention_dropout
  lowercase__	: Any     		=   fs_config.dropout_input
  lowercase__	: Any     		=   fs_config.dropout
  lowercase__	: Dict     		=   fs_config.mask_channel_length
  lowercase__	: Optional[int]     		=   fs_config.mask_channel_prob
  lowercase__	: Any     		=   fs_config.mask_length
  lowercase__	: Dict     		=   fs_config.mask_prob

  lowercase__	: Dict     		=   """Wav2Vec2FeatureExtractor"""
  lowercase__	: List[Any]     		=   """Wav2Vec2CTCTokenizer"""

  return config

@torch.no_grad()
def 		UpperCamelCase						(   lowercase_			,			lowercase_			,			lowercase_=None			,			lowercase_=None			,			lowercase_=True )		->						int:


  '''simple docstring'''






  if is_finetuned:
    lowercase__		,     lowercase__		,     lowercase__	: Optional[Any]     		=   fairseq.checkpoint_utils.load_model_ensemble_and_task(
        [checkpoint_path]			,			arg_overrides={"""data""": """/""".join(dict_path.split("""/""" )[:-1] )} )
  else:
    lowercase__		,     lowercase__		,     lowercase__	: Optional[int]     		=   fairseq.checkpoint_utils.load_model_ensemble_and_task([checkpoint_path] )

  if config_path is not None:
    lowercase__	: Optional[Any]     		=   SEWConfig.from_pretrained(lowercase_ )
  else:
    lowercase__	: Optional[int]     		=   convert_config(model[0]			,			lowercase_ )
  lowercase__	: Dict     		=   model[0].eval()

  lowercase__	: Any     		=   True if config.feat_extract_norm == """layer""" else False
  lowercase__	: List[Any]     		=   WavaVecaFeatureExtractor(
      feature_size=1			,			sampling_rate=1_60_00			,			padding_value=0			,			do_normalize=lowercase_			,			return_attention_mask=lowercase_			,			)

  if is_finetuned:
    if dict_path:
      lowercase__	: Optional[int]     		=   Dictionary.load(lowercase_ )

      # important change bos & pad token id since CTC symbol is <pad> and
      # not <s> as in fairseq
      lowercase__	: List[Any]     		=   target_dict.pad_index
      lowercase__	: int     		=   target_dict.bos_index
      lowercase__	: Optional[Any]     		=   target_dict.pad_index
      lowercase__	: List[str]     		=   target_dict.bos_index
      lowercase__	: Optional[int]     		=   target_dict.eos_index
      lowercase__	: Union[str, Any]     		=   len(target_dict.symbols )
      lowercase__	: Optional[int]     		=   os.path.join(lowercase_			,			"""vocab.json""" )
      if not os.path.isdir(lowercase_ ):
        logger.error("""--pytorch_dump_folder_path ({}) should be a directory""".format(lowercase_ ) )
        return
      os.makedirs(lowercase_			,			exist_ok=lowercase_ )
      with open(lowercase_			,			"""w"""			,			encoding="""utf-8""" ) as vocab_handle:
        json.dump(target_dict.indices			,			lowercase_ )
      lowercase__	: Optional[Any]     		=   WavaVecaCTCTokenizer(
          lowercase_			,			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=lowercase_			,			)
      lowercase__	: Union[str, Any]     		=   WavaVecaProcessor(feature_extractor=lowercase_			,			tokenizer=lowercase_ )
      processor.save_pretrained(lowercase_ )

    lowercase__	: List[str]     		=   SEWForCTC(lowercase_ )
  else:
    lowercase__	: str     		=   SEWModel(lowercase_ )
    feature_extractor.save_pretrained(lowercase_ )

  recursively_load_weights(lowercase_			,			lowercase_			,			lowercase_ )

  hf_model.save_pretrained(lowercase_ )


if __name__ == "__main__":
     lowerCamelCase__     :		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(
         """--is_finetuned""", action="""store_true""", help="""Whether the model to convert is a fine-tuned model or not"""
     )
     lowerCamelCase__     :		Optional[int]        = parser.parse_args()
     convert_sew_checkpoint(
         args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, args.is_finetuned
     )













from collections import defaultdict
from math import ceil, sqrt





def 				__lowerCAmelCase(       _SCREAMING_SNAKE_CASE = 1_000_000							,     _SCREAMING_SNAKE_CASE = 10				)							-> int:




    """simple docstring"""




    _A     =							defaultdict(_SCREAMING_SNAKE_CASE				)

    for outer_width in range(3							,     (t_limit // 4) + 2				):
        if outer_width * outer_width > t_limit:
            _A     =							max(
                ceil(sqrt(outer_width * outer_width - t_limit				)				)							,     1				)
        else:
            _A     =							1

        hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

        for hole_width in range(_SCREAMING_SNAKE_CASE							,     outer_width - 1							,     2				):
            count[outer_width * outer_width - hole_width * hole_width] += 1

    return sum(1 for n in count.values() if 1 <= n <= 10				)


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





