






"""simple docstring"""







import subprocess
import sys
import warnings
from argparse import ArgumentParser
from pathlib import Path

from packaging import version

from .. import AutoFeatureExtractor, AutoImageProcessor, AutoProcessor, AutoTokenizer
from ..utils import logging
from ..utils.import_utils import is_optimum_available
from .convert import export, validate_model_outputs
from .features import FeaturesManager
from .utils import get_preprocessor


_UpperCamelCase     			=  """1.5.0"""

_UpperCamelCase     			=  ["""vision-encoder-decoder"""]

def _a (		_snake_case							):







      """simple docstring"""






      if is_optimum_available():
            from optimum.version import __version__ as optimum_version

            UpperCAmelCase							=       version.parse(_snake_case							)
            if parsed_optimum_version < version.parse(_snake_case							):
                  raise RuntimeError(
                      F'''transformers.onnx requires optimum >= {MIN_OPTIMUM_VERSION} but {optimum_version} is installed. You '''
                      """can upgrade optimum by running: pip install -U optimum[exporters]"""							)
      else:
            raise RuntimeError(
                """transformers.onnx requires optimum to run, you can install the library by running: pip install """
                """optimum[exporters]"""							)
      UpperCAmelCase							=       [
          sys.executable,
          """-m""",
          """optimum.exporters.onnx""",
          F'''--model {args.model}''',
          F'''--task {args.feature}''',
          F'''--framework {args.framework}''' if args.framework is not None else """""",
          F'''{args.output}''',
      ]
      UpperCAmelCase							=       subprocess.Popen(""" """.join(_snake_case							)      ,       stdout=subprocess.PIPE      ,       shell=_snake_case							)
      proc.wait()

      logger.info(
          """The export was done by optimum.exporters.onnx. We recommend using to use this package directly in future, as """
          """transformers.onnx is deprecated, and will be removed in v5. You can find more information here: """
          """https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model."""							)

def _a (		_snake_case							):







      """simple docstring"""






      UpperCAmelCase							=       args.output if args.output.is_file() else args.output.joinpath("""model.onnx"""							)
      if not args.output.parent.exists():
            args.output.parent.mkdir(parents=_snake_case							)

      # Allocate the model
      UpperCAmelCase							=       FeaturesManager.get_model_from_feature(
          args.feature      ,       args.model      ,       framework=args.framework      ,       cache_dir=args.cache_dir							)

      UpperCAmelCase							=       FeaturesManager.check_supported_model_or_raise(_snake_case      ,       feature=args.feature							)
      UpperCAmelCase							=       model_onnx_config(model.config							)

      if model_kind in ENCODER_DECODER_MODELS:
            UpperCAmelCase							=       model.get_encoder()
            UpperCAmelCase							=       model.get_decoder()

            UpperCAmelCase							=       onnx_config.get_encoder_config(encoder_model.config							)
            UpperCAmelCase							=       onnx_config.get_decoder_config(
                encoder_model.config      ,       decoder_model.config      ,       feature=args.feature							)

            if args.opset is None:
                  UpperCAmelCase							=       max(encoder_onnx_config.default_onnx_opset      ,       decoder_onnx_config.default_onnx_opset							)

            if args.opset < min(encoder_onnx_config.default_onnx_opset      ,       decoder_onnx_config.default_onnx_opset							):
                  raise ValueError(
                      F'''Opset {args.opset} is not sufficient to export {model_kind}. At least '''
                      F''' {min(encoder_onnx_config.default_onnx_opset      ,       decoder_onnx_config.default_onnx_opset							)} is required.'''							)

            UpperCAmelCase							=       AutoFeatureExtractor.from_pretrained(args.model							)

            UpperCAmelCase							=       export(
                _snake_case      ,       _snake_case      ,       _snake_case      ,       args.opset      ,       args.output.parent.joinpath("""encoder_model.onnx"""							)      ,       )

            validate_model_outputs(
                _snake_case      ,       _snake_case      ,       _snake_case      ,       args.output.parent.joinpath("""encoder_model.onnx"""							)      ,       _snake_case      ,       args.atol if args.atol else encoder_onnx_config.atol_for_validation      ,       )

            UpperCAmelCase							=       AutoTokenizer.from_pretrained(args.model							)

            UpperCAmelCase							=       export(
                _snake_case      ,       _snake_case      ,       _snake_case      ,       args.opset      ,       args.output.parent.joinpath("""decoder_model.onnx"""							)      ,       )

            validate_model_outputs(
                _snake_case      ,       _snake_case      ,       _snake_case      ,       args.output.parent.joinpath("""decoder_model.onnx"""							)      ,       _snake_case      ,       args.atol if args.atol else decoder_onnx_config.atol_for_validation      ,       )
            logger.info(
                F'''All good, model saved at: {args.output.parent.joinpath('encoder_model.onnx'							).as_posix()},'''
                F''' {args.output.parent.joinpath('decoder_model.onnx'							).as_posix()}'''							)

      else:
            # Instantiate the appropriate preprocessor
            if args.preprocessor == "auto":
                  UpperCAmelCase							=       get_preprocessor(args.model							)
            elif args.preprocessor == "tokenizer":
                  UpperCAmelCase							=       AutoTokenizer.from_pretrained(args.model							)
            elif args.preprocessor == "image_processor":
                  UpperCAmelCase							=       AutoImageProcessor.from_pretrained(args.model							)
            elif args.preprocessor == "feature_extractor":
                  UpperCAmelCase							=       AutoFeatureExtractor.from_pretrained(args.model							)
            elif args.preprocessor == "processor":
                  UpperCAmelCase							=       AutoProcessor.from_pretrained(args.model							)
            else:
                  raise ValueError(F'''Unknown preprocessor type \'{args.preprocessor}\''''							)

            # Ensure the requested opset is sufficient
            if args.opset is None:
                  UpperCAmelCase							=       onnx_config.default_onnx_opset

            if args.opset < onnx_config.default_onnx_opset:
                  raise ValueError(
                      F'''Opset {args.opset} is not sufficient to export {model_kind}. '''
                      F'''At least  {onnx_config.default_onnx_opset} is required.'''							)

            UpperCAmelCase							=       export(
                _snake_case      ,       _snake_case      ,       _snake_case      ,       args.opset      ,       args.output      ,       )

            if args.atol is None:
                  UpperCAmelCase							=       onnx_config.atol_for_validation

            validate_model_outputs(_snake_case      ,       _snake_case      ,       _snake_case      ,       args.output      ,       _snake_case      ,       args.atol							)
            logger.info(F'''All good, model saved at: {args.output.as_posix()}'''							)
            warnings.warn(
                """The export was done by transformers.onnx which is deprecated and will be removed in v5. We recommend"""
                """ using optimum.exporters.onnx in future. You can find more information here:"""
                """ https://huggingface.co/docs/optimum/exporters/onnx/usage_guides/export_a_model."""      ,       _snake_case      ,       )

def _a (		):







      """simple docstring"""






      UpperCAmelCase							=       ArgumentParser("""Hugging Face Transformers ONNX exporter"""							)
      parser.add_argument(
          """-m"""      ,       """--model"""      ,       type=_snake_case      ,       required=_snake_case      ,       help="""Model ID on huggingface.co or path on disk to load model from."""							)
      parser.add_argument(
          """--feature"""      ,       default="""default"""      ,       help="""The type of features to export the model with."""      ,       )
      parser.add_argument("""--opset"""      ,       type=_snake_case      ,       default=_snake_case      ,       help="""ONNX opset version to export the model with."""							)
      parser.add_argument(
          """--atol"""      ,       type=_snake_case      ,       default=_snake_case      ,       help="""Absolute difference tolerance when validating the model."""							)
      parser.add_argument(
          """--framework"""      ,       type=_snake_case      ,       choices=["""pt""", """tf"""]      ,       default=_snake_case      ,       help=(
              """The framework to use for the ONNX export."""
              """ If not provided, will attempt to use the local checkpoint's original framework"""
              """ or what is available in the environment."""
          )      ,       )
      parser.add_argument("""output"""      ,       type=_snake_case      ,       help="""Path indicating where to store generated ONNX model."""							)
      parser.add_argument("""--cache_dir"""      ,       type=_snake_case      ,       default=_snake_case      ,       help="""Path indicating where to store cache."""							)
      parser.add_argument(
          """--preprocessor"""      ,       type=_snake_case      ,       choices=["""auto""", """tokenizer""", """feature_extractor""", """image_processor""", """processor"""]      ,       default="""auto"""      ,       help="""Which type of preprocessor to use. 'auto' tries to automatically detect it."""      ,       )
      parser.add_argument(
          """--export_with_transformers"""      ,       action="""store_true"""      ,       help=(
              """Whether to use transformers.onnx instead of optimum.exporters.onnx to perform the ONNX export. It can be """
              """useful when exporting a model supported in transformers but not in optimum, otherwise it is not """
              """recommended."""
          )      ,       )

      UpperCAmelCase							=       parser.parse_args()
      if args.export_with_transformers or not is_optimum_available():
            export_with_transformers(_snake_case							)
      else:
            export_with_optimum(_snake_case							)


if __name__ == "__main__":
  _UpperCamelCase     			=  logging.get_logger("""transformers.onnx""")  # pylint: disable=invalid-name
  logger.setLevel(logging.INFO)
  main()





import argparse
import os
from functools import reduce

import fairseq
import torch
from datasets import load_dataset

from transformers import WavaVecaProcessor, logging
from transformers.models.dataavec.configuration_dataavec_audio import DataaVecAudioConfig

# Copied from https://github.com/pytorch/fairseq/blob/main/examples/data2vec/models/data2vec_audio.py
from transformers.models.dataavec.dataavec_audio import DataaVecAudioModel as Dummy  # noqa: F401
from transformers.models.dataavec.modeling_dataavec_audio import DataaVecAudioForCTC, DataaVecAudioModel


logging.set_verbosity_info()
lowerCamelCase_  :       int					      =     logging.get_logger(__name__)

lowerCamelCase_  :       Optional[int]					      =     {
    """post_extract_proj""": """feature_projection.projection""",
    """models.0.layer_norm""": """feature_projection.layer_norm""",
    """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""",
    """w2v_encoder.proj""": """lm_head""",
    """mask_emb""": """masked_spec_embed""",
}
lowerCamelCase_  :       Any					      =     [
    """lm_head""",
]







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		List[Any]:
				for attribute in key.split("""."""  ):
								UpperCamelCase_:							List[str]		=  getattr(lowerCamelCase   ,			lowerCamelCase  )

				if weight_type is not None:
								UpperCamelCase_:							Tuple		=  getattr(lowerCamelCase   ,			lowerCamelCase  ).shape
				else:
								UpperCamelCase_:							List[Any]		=  hf_pointer.shape

				if hf_shape != value.shape:
								raise ValueError(
								    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":
								UpperCamelCase_:							str		=  value
				elif weight_type == "weight_g":
								UpperCamelCase_:							List[str]		=  value
				elif weight_type == "weight_v":
								UpperCamelCase_:							Tuple		=  value
				elif weight_type == "bias":
								UpperCamelCase_:							Any		=  value
				else:
								UpperCamelCase_:							Dict		=  value

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







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		Union[str, Any]:
				UpperCamelCase_:							Optional[int]		=  []
				UpperCamelCase_:							Tuple		=  fairseq_model.state_dict()

				if not is_headless:
								UpperCamelCase_:							Dict		=  hf_model.dataavec_audio.feature_extractor
								UpperCamelCase_:							Tuple		=  hf_model.dataavec_audio.encoder.pos_conv_embed

				else:
								UpperCamelCase_:							Tuple		=  hf_model.feature_extractor
								UpperCamelCase_:							int		=  hf_model.encoder.pos_conv_embed

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

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







def 				A__    (  lowerCamelCase   ,			lowerCamelCase  )     ->		Dict:
				UpperCamelCase_:							Union[str, Any]		=  path.split("""."""  )
				return reduce(lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		str:
				UpperCamelCase_:							Union[str, Any]		=  access_by_string(lowerCamelCase   ,			lowerCamelCase  )
				UpperCamelCase_:							str		=  hf_weight.data

				if fsq_value.shape != hf_value.shape:
								raise ValueError(F'''{full_name} has size {fsq_value.shape}, but {hf_value.shape} was found.'''  )
				UpperCamelCase_:							Optional[int]		=  fsq_value
				logger.info(F'''{full_name} was correctly initialized from {hf_weight_path}.'''  )







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		Optional[int]:
				UpperCamelCase_:							str		=  full_name.split("""conv_layers."""  )[-1]
				UpperCamelCase_:							Union[str, Any]		=  name.split("""."""  )
				UpperCamelCase_:							Union[str, Any]		=  int(items[0]  )
				UpperCamelCase_:							Union[str, Any]		=  int(items[1]  )

				UpperCamelCase_:							Tuple		=  name.split("""."""  )[-1]
				if type_id == 0:
								UpperCamelCase_:							List[Any]		=  """conv"""
				elif type_id == 2:
								UpperCamelCase_:							Tuple		=  """layer_norm"""
				else:
								unused_weights.append(lowerCamelCase  )
								return

				set_weights(lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			F'''conv_layers.{layer_id}.{layer_type}.{weight_type}'''  )







def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase  )     ->		Optional[Any]:
				UpperCamelCase_:							Optional[int]		=  full_name.split("""pos_conv."""  )[-1]
				UpperCamelCase_:							Any		=  name.split("""."""  )
				UpperCamelCase_:							List[str]		=  int(items[0]  )
				UpperCamelCase_:							str		=  int(items[1]  )

				UpperCamelCase_:							List[str]		=  name.split("""."""  )[-1]
				if type_id != 0:
								unused_weights.append(lowerCamelCase  )
								return
				else:
								UpperCamelCase_:							Tuple		=  """conv"""

				set_weights(lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase   ,			F'''layers.{layer_id}.{layer_type}.{weight_type}'''  )







@torch.no_grad()
def 				A__    (  lowerCamelCase   ,			lowerCamelCase   ,			lowerCamelCase=None   ,			lowerCamelCase=None   ,			lowerCamelCase=True  )     ->		Dict:
				if config_path is not None:
								UpperCamelCase_:							str		=  DataaVecAudioConfig.from_pretrained(lowerCamelCase  )
				else:
								UpperCamelCase_:							Tuple		=  DataaVecAudioConfig()

				if not is_finetuned:
								# Modify final_proj layer name
								UpperCamelCase_:							Union[str, Any]		=  DataaVecAudioModel(lowerCamelCase  )
								UpperCamelCase_:							Optional[int]		=  os.path.dirname(lowerCamelCase  )

								UpperCamelCase_:							Optional[Any]		=  torch.load(lowerCamelCase  )
								UpperCamelCase_:							List[str]		=  state_dict["""model"""].pop("""final_proj.0.weight"""  )
								UpperCamelCase_:							int		=  state_dict["""model"""].pop("""final_proj.0.bias"""  )
								UpperCamelCase_:							Optional[int]		=  os.path.join(lowerCamelCase   ,			"""converted.pt"""  )
								torch.save(lowerCamelCase   ,			lowerCamelCase  )
				else:
								UpperCamelCase_:							int		=  DataaVecAudioForCTC(lowerCamelCase  )
								UpperCamelCase_:							Optional[int]		=  checkpoint_path

				def load_dataavec(lowerCamelCase  ):
								UpperCamelCase_,      UpperCamelCase_,      UpperCamelCase_:							Optional[int]		=  fairseq.checkpoint_utils.load_model_ensemble_and_task([path]  )
								return model[0].eval()

				UpperCamelCase_:							List[Any]		=  load_dataavec(lowerCamelCase  )

				recursively_load_weights(lowerCamelCase   ,			lowerCamelCase   ,			not is_finetuned  )

				UpperCamelCase_:							Tuple		=  WavaVecaProcessor.from_pretrained("""facebook/wav2vec2-large-lv60"""  )

				UpperCamelCase_:							Dict		=  load_dataset("""patrickvonplaten/librispeech_asr_dummy"""   ,			"""clean"""   ,			split="""validation"""  )
				UpperCamelCase_:							Optional[Any]		=  [x["""array"""] for x in ds[:4]["""audio"""]]

				UpperCamelCase_:							List[Any]		=  processor(lowerCamelCase   ,			return_tensors="""pt"""   ,			padding=lowerCamelCase  )

				UpperCamelCase_:							str		=  inputs.input_values
				UpperCamelCase_:							Dict		=  inputs.attention_mask
				#    input_values = inputs.input_values[:, :-1]
				#    attention_mask = inputs.attention_mask[:, :-1]

				hf_wavavec.eval()
				model.eval()
				if is_finetuned:
								UpperCamelCase_:							Optional[Any]		=  model(source=lowerCamelCase   ,			padding_mask=(1 - attention_mask)   ,			mask=lowerCamelCase   ,			features_only=lowerCamelCase  )[
								    """encoder_out"""
								].transpose(0   ,			1  )
								UpperCamelCase_:							Tuple		=  hf_wavavec(lowerCamelCase   ,			attention_mask=lowerCamelCase  )["""logits"""]

								UpperCamelCase_:							Optional[Any]		=  torch.argmax(lowerCamelCase   ,			dim=-1  )
								UpperCamelCase_:							Any		=  processor.batch_decode(lowerCamelCase  )

								print(F'''Expected Output: {ds[:4]["text"]}, Pred: {output_string}'''  )
				else:
								UpperCamelCase_:							Optional[Any]		=  model(source=lowerCamelCase   ,			padding_mask=(1 - attention_mask)   ,			mask=lowerCamelCase   ,			features_only=lowerCamelCase  )[
								    """layer_results"""
								][-1][0].transpose(0   ,			1  )
								UpperCamelCase_:							List[str]		=  hf_wavavec(lowerCamelCase   ,			attention_mask=lowerCamelCase  )["""last_hidden_state"""]

				print(our_output.shape   ,			their_output.shape  )
				UpperCamelCase_:							Optional[Any]		=  torch.max(torch.abs(our_output - their_output  )  ).item()
				print(F'''max_absolute_diff = {max_absolute_diff}'''  )  # ~ 1e-7
				UpperCamelCase_:							List[str]		=  torch.allclose(lowerCamelCase   ,			lowerCamelCase   ,			atol=1E-3  )
				print("""Do both models output the same tensors?"""   ,			"""🔥""" if success else """💩"""  )
				if not success:
								raise Exception("""Something went wRoNg"""  )

				hf_wavavec.save_pretrained(lowerCamelCase  )

				if is_finetuned:
								processor.save_pretrained(lowerCamelCase  )
				else:
								processor.feature_extractor.save_pretrained(lowerCamelCase  )


if __name__ == "__main__":
							lowerCamelCase_  :       str					      =     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"""
							)
							lowerCamelCase_  :       Dict					      =     parser.parse_args()
							convert_wavaveca_checkpoint(
							    args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.dict_path, not args.not_finetuned
							)






