





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()
lowercase_					     =	logging.get_logger(__name__)

lowercase_					     =	{
    """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""",
}
lowercase_					     =	[
    """lm_head""",
    """quantizer.weight_proj""",
    """quantizer.codevectors""",
    """project_q""",
    """project_hid""",
]




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     List[str]:
  for attribute in key.split('.'	):
    lowercase__			    =						getattr(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)

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

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




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Any:
  lowercase__			    =						[]
  lowercase__			    =						fairseq_model.state_dict()

  lowercase__			    =						hf_model.feature_extractor
  lowercase__			    =						hf_model.adapter

  for name, value in fairseq_dict.items():
    lowercase__			    =						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'	,		)
      lowercase__			    =						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	)
      lowercase__			    =						True
    else:
      for key, mapped_key in MAPPING.items():
        if key in name or key.split('w2v_model.'	)[-1] == name.split('.'	)[0]:
          lowercase__			    =						True
          if "*" in mapped_key:
            lowercase__			    =						name.split(_SCREAMING_SNAKE_CASE	)[0].split('.'	)[-2]
            lowercase__			    =						mapped_key.replace('*'	,		_SCREAMING_SNAKE_CASE	)
          if "weight_g" in name:
            lowercase__			    =						'weight_g'
          elif "weight_v" in name:
            lowercase__			    =						'weight_v'
          elif "bias" in name:
            lowercase__			    =						'bias'
          elif "weight" in name:
            lowercase__			    =						'weight'
          else:
            lowercase__			    =						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     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Any:
  lowercase__			    =						full_name.split('conv_layers.'	)[-1]
  lowercase__			    =						name.split('.'	)
  lowercase__			    =						int(items[0]	)
  lowercase__			    =						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__			    =						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__			    =						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__			    =						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__			    =						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     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     List[str]:
  lowercase__			    =						full_name.split('adaptor.'	)[-1]
  lowercase__			    =						name.split('.'	)

  if items[1].isdigit():
    lowercase__			    =						int(items[1]	)
  else:
    lowercase__			    =						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."""
        lowercase__			    =						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."""
        lowercase__			    =						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."""
        lowercase__			    =						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."""
        lowercase__			    =						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."""
      lowercase__			    =						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."""
      lowercase__			    =						value
      logger.info(F"""Adapter layer {layer_id} bias was initialized from {full_name}."""	)
  else:
    unused_weights.append(_SCREAMING_SNAKE_CASE	)




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	)       ->     Optional[Any]:
  lowercase__    ,					lowercase__			    =						emb.weight.shape
  lowercase__			    =						nn.Linear(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		bias=_SCREAMING_SNAKE_CASE	)
  lowercase__			    =						emb.weight.data
  return lin_layer




@torch.no_grad()
def     __UpperCamelCase						(_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	,		)       ->     List[str]:
  lowercase__			    =						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	,		)
  lowercase__			    =						MBartConfig.from_pretrained(_SCREAMING_SNAKE_CASE	)

  # load model
  lowercase__    ,					lowercase__    ,					lowercase__			    =						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,
      }	,		)
  lowercase__			    =						model[0].eval()

  # load feature extractor
  lowercase__			    =						WavaVecaFeatureExtractor.from_pretrained(_SCREAMING_SNAKE_CASE	,		use_auth_token=_SCREAMING_SNAKE_CASE	)

  # set weights for wav2vec2 encoder
  lowercase__			    =						WavaVecaModel(_SCREAMING_SNAKE_CASE	)

  recursively_load_weights_wavaveca(model.encoder	,		_SCREAMING_SNAKE_CASE	)

  # load decoder weights
  lowercase__			    =						MBartForCausalLM(_SCREAMING_SNAKE_CASE	)
  lowercase__    ,					lowercase__			    =						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}"""	)

  lowercase__			    =						SpeechEncoderDecoderModel(encoder=_SCREAMING_SNAKE_CASE	,		decoder=_SCREAMING_SNAKE_CASE	)
  lowercase__			    =						False

  lowercase__			    =						MBartaaTokenizer(_SCREAMING_SNAKE_CASE	)
  tokenizer.save_pretrained(_SCREAMING_SNAKE_CASE	)

  lowercase__			    =						hf_wavavec.config.to_dict()
  lowercase__			    =						tokenizer.pad_token_id
  lowercase__			    =						tokenizer.bos_token_id
  lowercase__			    =						tokenizer.eos_token_id
  lowercase__			    =						'mbart50'
  lowercase__			    =						'wav2vec2'

  lowercase__			    =						tokenizer.eos_token_id
  lowercase__			    =						250004
  lowercase__			    =						tokenizer.eos_token_id

  lowercase__			    =						SpeechEncoderDecoderConfig.from_dict(_SCREAMING_SNAKE_CASE	)

  hf_wavavec.save_pretrained(_SCREAMING_SNAKE_CASE	)
  feature_extractor.save_pretrained(_SCREAMING_SNAKE_CASE	)


if __name__ == "__main__":
      lowercase_					     =	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""")

      lowercase_					     =	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,
      )









import os
import zipfile

import requests
from get_ci_error_statistics import download_artifact, get_artifacts_links




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE=7	)       ->     List[Any]:
  lowercase__			    =						None
  if token is not None:
    lowercase__			    =						{'Accept': 'application/vnd.github+json', 'Authorization': F"""Bearer {token}"""}

  # The id of a workflow (not of a workflow run)
  lowercase__			    =						'636036'

  lowercase__			    =						F"""https://api.github.com/repos/huggingface/transformers/actions/workflows/{workflow_id}/runs"""
  # On `main` branch + event being `schedule` + not returning PRs + only `num_runs` results
  url += F"""?branch=main&event=schedule&exclude_pull_requests=true&per_page={num_runs}"""

  lowercase__			    =						requests.get(_SCREAMING_SNAKE_CASE	,		headers=_SCREAMING_SNAKE_CASE	).json()

  return result["workflow_runs"]




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	)       ->     Union[str, Any]:
  lowercase__			    =						get_daily_ci_runs(_SCREAMING_SNAKE_CASE	)
  lowercase__			    =						None
  for workflow_run in workflow_runs:
    if workflow_run["status"] == "completed":
      lowercase__			    =						workflow_run['id']
      break

  return workflow_run_id




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Tuple:
  lowercase__			    =						get_last_daily_ci_runs(_SCREAMING_SNAKE_CASE	)
  if workflow_run_id is not None:
    lowercase__			    =						get_artifacts_links(worflow_run_id=_SCREAMING_SNAKE_CASE	,		token=_SCREAMING_SNAKE_CASE	)
    for artifact_name in artifact_names:
      if artifact_name in artifacts_links:
        lowercase__			    =						artifacts_links[artifact_name]
        download_artifact(
            artifact_name=_SCREAMING_SNAKE_CASE	,		artifact_url=_SCREAMING_SNAKE_CASE	,		output_dir=_SCREAMING_SNAKE_CASE	,		token=_SCREAMING_SNAKE_CASE	)




def     __UpperCamelCase						(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)       ->     Dict:
  get_last_daily_ci_artifacts(_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	,		_SCREAMING_SNAKE_CASE	)

  lowercase__			    =						{}
  for artifact_name in artifact_names:
    lowercase__			    =						os.path.join(_SCREAMING_SNAKE_CASE	,		F"""{artifact_name}.zip"""	)
    if os.path.isfile(_SCREAMING_SNAKE_CASE	):
      lowercase__			    =						{}
      with zipfile.ZipFile(_SCREAMING_SNAKE_CASE	) as z:
        for filename in z.namelist():
          if not os.path.isdir(_SCREAMING_SNAKE_CASE	):
            # read the file
            with z.open(_SCREAMING_SNAKE_CASE	) as f:
              lowercase__			    =						f.read().decode('UTF-8'	)

  return results


