





import argparse

import fairseq
import torch
from torch import nn

from transformers import (
    MBartaaTokenizer,
    MBartConfig,
    MBartForCausalLM,
    SpeechEncoderDecoderConfig,
    SpeechEncoderDecoderModel,
    WavaVecaConfig,
    WavaVecaFeatureExtractor,
    WavaVecaModel,
    logging,
)


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

__UpperCAmelCase       :   Tuple				       =							{
    "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": "lm_head",
    "mask_emb": "masked_spec_embed",
}
__UpperCAmelCase       :   Dict				       =							[
    "lm_head",
    "quantizer.weight_proj",
    "quantizer.codevectors",
    "project_q",
    "project_hid",
]







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    Optional[int]:
  for attribute in key.split("""."""):
    __snake_case:						Union[str, Any]   =						getattr(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)

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

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







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    Any:
  __snake_case:						str   =						[]
  __snake_case:						Union[str, Any]   =						fairseq_model.state_dict()

  __snake_case:						Dict   =						hf_model.feature_extractor
  __snake_case:						Dict   =						hf_model.adapter

  for name, value in fairseq_dict.items():
    __snake_case:						List[str]   =						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:						Optional[int]   =						True
    elif any(x in name for x in ["""adaptor""", """w2v_encoder.proj.""", """w2v_proj_ln."""]):
      load_adapter(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)
      __snake_case:						Dict   =						True
    else:
      for key, mapped_key in MAPPING.items():
        if key in name or key.split("""w2v_model.""")[-1] == name.split(""".""")[0]:
          __snake_case:						Any   =						True
          if "*" in mapped_key:
            __snake_case:						Any   =						name.split(SCREAMING_SNAKE_CASE__)[0].split(""".""")[-2]
            __snake_case:						List[str]   =						mapped_key.replace("""*"""					,					SCREAMING_SNAKE_CASE__)
          if "weight_g" in name:
            __snake_case:						Optional[int]   =						"""weight_g"""
          elif "weight_v" in name:
            __snake_case:						Tuple   =						"""weight_v"""
          elif "bias" in name:
            __snake_case:						Optional[Any]   =						"""bias"""
          elif "weight" in name:
            __snake_case:						Optional[int]   =						"""weight"""
          else:
            __snake_case:						Dict   =						None
          set_recursively(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 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    int:
  __snake_case:						Dict   =						full_name.split("""conv_layers.""")[-1]
  __snake_case:						List[Any]   =						name.split(""".""")
  __snake_case:						List[str]   =						int(items[0])
  __snake_case:						Optional[Any]   =						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:						Union[str, Any]   =						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:						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:						Tuple   =						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:						Tuple   =						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__)







def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__)   ->    Optional[Any]:
  __snake_case:						Optional[Any]   =						full_name.split("""adaptor.""")[-1]
  __snake_case:						Tuple   =						name.split(""".""")

  if items[1].isdigit():
    __snake_case:						Tuple   =						int(items[1])
  else:
    __snake_case:						Tuple   =						None

  if "adaptor" not in full_name:
    if "proj_ln" in full_name:
      # has to be layer norm
      if "bias" in name:
        assert (
            value.shape == adapter.proj_layer_norm.bias.data.shape
        ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.bias.data.shape} was found.'''
        __snake_case:						int   =						value
        logger.info(F'''Adapter proj layer norm bias was initialized from {full_name}.''')
      if "weight" in name:
        assert (
            value.shape == adapter.proj_layer_norm.weight.data.shape
        ), F'''{full_name} has size {value.shape}, but {adapter.proj_layer_norm.weight.data.shape} was found.'''
        __snake_case:						str   =						value
    else:
      # has to be projection layer
      if "bias" in name:
        assert (
            value.shape == adapter.proj.bias.data.shape
        ), F'''{full_name} has size {value.shape}, but {adapter.proj.bias.data.shape} was found.'''
        __snake_case:						Union[str, Any]   =						value
        logger.info(F'''Adapter proj layer bias was initialized from {full_name}.''')
      if "weight" in name:
        assert (
            value.shape == adapter.proj.weight.data.shape
        ), F'''{full_name} has size {value.shape}, but {adapter.proj.weight.data.shape} was found.'''
        __snake_case:						Union[str, Any]   =						value
        logger.info(F'''Adapter proj layer weight was initialized from {full_name}.''')
  elif isinstance(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__):
    if "bias" in name:
      assert (
          value.shape == adapter.layers[layer_id].conv.bias.data.shape
      ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.bias.data.shape} was found.'''
      __snake_case:						Optional[Any]   =						value
      logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.''')
    elif "weight" in name:
      assert (
          value.shape == adapter.layers[layer_id].conv.weight.data.shape
      ), F'''{full_name} has size {value.shape}, but {adapter.layers[layer_id].conv.weight.data.shape} was found.'''
      __snake_case:						int   =						value
      logger.info(F'''Adapter layer {layer_id} bias was initialized from {full_name}.''')
  else:
    unused_weights.append(SCREAMING_SNAKE_CASE__)







def 			A__							(    SCREAMING_SNAKE_CASE__)   ->    Optional[Any]:
  __snake_case  ,   __snake_case:						Optional[Any]   =						emb.weight.shape
  __snake_case:						List[str]   =						nn.Linear(SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					bias=SCREAMING_SNAKE_CASE__)
  __snake_case:						str   =						emb.weight.data
  return lin_layer







@torch.no_grad()
def 			A__							(    SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					SCREAMING_SNAKE_CASE__					,					)   ->    Dict:
  __snake_case:						Tuple   =						WavaVecaConfig.from_pretrained(
      SCREAMING_SNAKE_CASE__					,					add_adapter=SCREAMING_SNAKE_CASE__					,					adapter_stride=SCREAMING_SNAKE_CASE__					,					adapter_kernel_size=SCREAMING_SNAKE_CASE__					,					use_auth_token=SCREAMING_SNAKE_CASE__					,					output_hidden_size=SCREAMING_SNAKE_CASE__					,					)
  __snake_case:						int   =						MBartConfig.from_pretrained(SCREAMING_SNAKE_CASE__)

  # load model
  __snake_case  ,   __snake_case  ,   __snake_case:						List[str]   =						fairseq.checkpoint_utils.load_model_ensemble_and_task(
      [checkpoint_path]					,					arg_overrides={
          """config_yaml""": config_yaml_path,
          """data""": """/""".join(dict_path.split("""/""")[:-1]),
          """w2v_path""": checkpoint_path,
          """load_pretrained_decoder_from""": None,
      }					,					)
  __snake_case:						List[Any]   =						model[0].eval()

  # load feature extractor
  __snake_case:						str   =						WavaVecaFeatureExtractor.from_pretrained(SCREAMING_SNAKE_CASE__					,					use_auth_token=SCREAMING_SNAKE_CASE__)

  # set weights for wav2vec2 encoder
  __snake_case:						Dict   =						WavaVecaModel(SCREAMING_SNAKE_CASE__)

  recursively_load_weights_wavaveca(model.encoder					,					SCREAMING_SNAKE_CASE__)

  # load decoder weights
  __snake_case:						Union[str, Any]   =						MBartForCausalLM(SCREAMING_SNAKE_CASE__)
  __snake_case  ,   __snake_case:						Optional[int]   =						hf_decoder.model.decoder.load_state_dict(model.decoder.state_dict()					,					strict=SCREAMING_SNAKE_CASE__)
  logger.warning(F'''The following keys are missing when loading the decoder weights: {missing_keys}''')
  logger.warning(F'''The following keys are unexpected when loading the decoder weights: {unexpected_keys}''')

  __snake_case:						Dict   =						SpeechEncoderDecoderModel(encoder=SCREAMING_SNAKE_CASE__					,					decoder=SCREAMING_SNAKE_CASE__)
  __snake_case:						Optional[int]   =						False

  __snake_case:						Union[str, Any]   =						MBartaaTokenizer(SCREAMING_SNAKE_CASE__)
  tokenizer.save_pretrained(SCREAMING_SNAKE_CASE__)

  __snake_case:						Optional[int]   =						hf_wavavec.config.to_dict()
  __snake_case:						int   =						tokenizer.pad_token_id
  __snake_case:						Union[str, Any]   =						tokenizer.bos_token_id
  __snake_case:						Tuple   =						tokenizer.eos_token_id
  __snake_case:						List[Any]   =						"""mbart50"""
  __snake_case:						Optional[int]   =						"""wav2vec2"""

  __snake_case:						Union[str, Any]   =						tokenizer.eos_token_id
  __snake_case:						str   =						25_0004
  __snake_case:						Optional[Any]   =						tokenizer.eos_token_id

  __snake_case:						List[Any]   =						SpeechEncoderDecoderConfig.from_dict(SCREAMING_SNAKE_CASE__)

  hf_wavavec.save_pretrained(SCREAMING_SNAKE_CASE__)
  feature_extractor.save_pretrained(SCREAMING_SNAKE_CASE__)


if __name__ == "__main__":
      __UpperCAmelCase       :   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_yaml_path", default=None, type=str, help="Path to yaml file of fine-tuned model")
      parser.add_argument(
          "--encoder_config_path",
          default="facebook/wav2vec2-xls-r-1b",
          type=str,
          help="Path to hf encoder wav2vec2 checkpoint config",
      )
      parser.add_argument(
          "--decoder_config_path",
          default="facebook/mbart-large-50-one-to-many-mmt",
          type=str,
          help="Path to hf decoder checkpoint config",
      )
      parser.add_argument("--add_adapter", default=True, type=bool, help="whethere to add model adapter layers")
      parser.add_argument("--adapter_stride", default=2, type=int, help="stride of adapter layers")
      parser.add_argument("--adapter_kernel_size", default=3, type=int, help="kernel size of adapter layers")
      parser.add_argument("--encoder_output_dim", default=1_024, type=int, help="encoder output dim")
      parser.add_argument("--start_token_id", default=250_004, type=int, help="`decoder_start_token_id` of model config")

      __UpperCAmelCase       :   Any				       =							parser.parse_args()
      convert_wavaveca_checkpoint(
          args.checkpoint_path,
          args.pytorch_dump_folder_path,
          args.dict_path,
          args.config_yaml_path,
          encoder_config_path=args.encoder_config_path,
          decoder_config_path=args.decoder_config_path,
          add_adapter=args.add_adapter,
          adapter_kernel_size=args.adapter_kernel_size,
          adapter_stride=args.adapter_stride,
          decoder_start_token_id=args.start_token_id,
          encoder_output_dim=args.encoder_output_dim,
      )













from __future__ import annotations

import typing
from collections import Counter







def 			A__							(    SCREAMING_SNAKE_CASE__)   ->    typing.Counter[int]:
  __snake_case:						typing.Counter[int]   =						Counter()
  for base in range(1					,					max_perimeter + 1):
    for perpendicular in range(SCREAMING_SNAKE_CASE__					,					max_perimeter + 1):
      __snake_case:						Dict   =						(base * base + perpendicular * perpendicular) ** 0.5
      if hypotenuse == int(SCREAMING_SNAKE_CASE__):
        __snake_case:						Any   =						int(base + perpendicular + hypotenuse)
        if perimeter > max_perimeter:
          continue
        triplets[perimeter] += 1
  return triplets







def 			A__							(    SCREAMING_SNAKE_CASE__ = 1000)   ->    int:
  __snake_case:						List[str]   =						pythagorean_triple(SCREAMING_SNAKE_CASE__)
  return triplets.most_common(1)[0][0]


if __name__ == "__main__":
      print(f'Perimeter {solution()} has maximum solutions')






