




from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...utils import logging


if TYPE_CHECKING:
       from ...feature_extraction_utils import FeatureExtractionMixin
       from ...tokenization_utils_base import PreTrainedTokenizerBase
       from ...utils import TensorType

UpperCAmelCase_									=	logging.get_logger(__name__)

UpperCAmelCase_									=	{
    'openai/whisper-base': 'https://huggingface.co/openai/whisper-base/resolve/main/config.json',
}

# fmt: off
UpperCAmelCase_									=	[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 357, 366, 438, 532, 685,
    705, 796, 930, 1_058, 1_220, 1_267, 1_279, 1_303, 1_343, 1_377,
    1_391, 1_635, 1_782, 1_875, 2_162, 2_361, 2_488, 3_467, 4_008, 4_211,
    4_600, 4_808, 5_299, 5_855, 6_329, 7_203, 9_609, 9_959, 10_563, 10_786,
    11_420, 11_709, 11_907, 13_163, 13_697, 13_700, 14_808, 15_306, 16_410, 16_791,
    17_992, 19_203, 19_510, 20_724, 22_305, 22_935, 27_007, 30_109, 30_420, 33_409,
    34_949, 40_283, 40_493, 40_549, 47_282, 49_146, 50_257, 50_359, 50_360, 50_361
]
UpperCAmelCase_									=	[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 359, 503, 522, 542, 873,
    893, 902, 918, 922, 931, 1_350, 1_853, 1_982, 2_460, 2_627,
    3_246, 3_253, 3_268, 3_536, 3_846, 3_961, 4_183, 4_667, 6_585, 6_647,
    7_273, 9_061, 9_383, 10_428, 10_929, 11_938, 12_033, 12_331, 12_562, 13_793,
    14_157, 14_635, 15_265, 15_618, 16_553, 16_604, 18_362, 18_956, 20_075, 21_675,
    22_520, 26_130, 26_161, 26_435, 28_279, 29_464, 31_650, 32_302, 32_470, 36_865,
    42_863, 47_425, 49_870, 50_254, 50_258, 50_360, 50_361, 50_362
]





class     lowerCamelCase__(					UpperCamelCase_):
      UpperCAmelCase__      :	Optional[Any]    		=   'whisper'
      UpperCAmelCase__      :	List[Any]    		=   ['past_key_values']
      UpperCAmelCase__      :	List[Any]    		=   {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}
      def __init__(				self: Optional[int]    ,       UpperCamelCase_: Optional[Any]=5_18_65    ,       UpperCamelCase_: List[Any]=80    ,       UpperCamelCase_: Optional[Any]=6    ,       UpperCamelCase_: Tuple=4    ,       UpperCamelCase_: List[str]=6    ,       UpperCamelCase_: Union[str, Any]=4    ,       UpperCamelCase_: str=15_36    ,       UpperCamelCase_: List[Any]=15_36    ,       UpperCamelCase_: Any=0.0    ,       UpperCamelCase_: Tuple=0.0    ,       UpperCamelCase_: Dict=5_02_57    ,       UpperCamelCase_: Optional[Any]=True    ,       UpperCamelCase_: Tuple=True    ,       UpperCamelCase_: Union[str, Any]="gelu"    ,       UpperCamelCase_: str=2_56    ,       UpperCamelCase_: Optional[Any]=0.0    ,       UpperCamelCase_: List[str]=0.0    ,       UpperCamelCase_: Tuple=0.0    ,       UpperCamelCase_: Tuple=0.02    ,       UpperCamelCase_: Dict=False    ,       UpperCamelCase_: List[str]=15_00    ,       UpperCamelCase_: int=4_48    ,       UpperCamelCase_: Optional[Any]=5_02_56    ,       UpperCamelCase_: int=5_02_56    ,       UpperCamelCase_: Optional[int]=5_02_56    ,       UpperCamelCase_: str=None    ,       UpperCamelCase_: Optional[Any]=[2_20, 5_02_56]    ,       UpperCamelCase_: Optional[Any]=False    ,       UpperCamelCase_: str=2_56    ,       UpperCamelCase_: List[str]=False    ,       UpperCamelCase_: int=0.05    ,       UpperCamelCase_: List[str]=10    ,       UpperCamelCase_: int=2    ,       UpperCamelCase_: int=0.0    ,       UpperCamelCase_: List[str]=10    ,       UpperCamelCase_: Dict=0    ,       UpperCamelCase_: Dict=7    ,       **UpperCamelCase_: List[Any]    ,       ):
             __lowerCamelCase   =						vocab_size
             __lowerCamelCase   =						num_mel_bins
             __lowerCamelCase   =						d_model
             __lowerCamelCase   =						encoder_layers
             __lowerCamelCase   =						encoder_attention_heads
             __lowerCamelCase   =						decoder_layers
             __lowerCamelCase   =						decoder_attention_heads
             __lowerCamelCase   =						decoder_ffn_dim
             __lowerCamelCase   =						encoder_ffn_dim
             __lowerCamelCase   =						dropout
             __lowerCamelCase   =						attention_dropout
             __lowerCamelCase   =						activation_dropout
             __lowerCamelCase   =						activation_function
             __lowerCamelCase   =						init_std
             __lowerCamelCase   =						encoder_layerdrop
             __lowerCamelCase   =						decoder_layerdrop
             __lowerCamelCase   =						use_cache
             __lowerCamelCase   =						encoder_layers
             __lowerCamelCase   =						scale_embedding  # scale factor will be sqrt(d_model) if True
             __lowerCamelCase   =						max_source_positions
             __lowerCamelCase   =						max_target_positions

             # Audio Classification-specific parameters. Feel free to ignore for other classes.
             __lowerCamelCase   =						classifier_proj_size
             __lowerCamelCase   =						use_weighted_layer_sum

             # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
             __lowerCamelCase   =						apply_spec_augment
             __lowerCamelCase   =						mask_time_prob
             __lowerCamelCase   =						mask_time_length
             __lowerCamelCase   =						mask_time_min_masks
             __lowerCamelCase   =						mask_feature_prob
             __lowerCamelCase   =						mask_feature_length
             __lowerCamelCase   =						mask_feature_min_masks

             __lowerCamelCase   =						median_filter_width

             super().__init__(
                 pad_token_id=UpperCamelCase__    ,       bos_token_id=UpperCamelCase__    ,       eos_token_id=UpperCamelCase__    ,       is_encoder_decoder=UpperCamelCase__    ,       decoder_start_token_id=UpperCamelCase__    ,       suppress_tokens=UpperCamelCase__    ,       begin_suppress_tokens=UpperCamelCase__    ,       **UpperCamelCase__    ,       )





class     lowerCamelCase__(					UpperCamelCase_):
      @property
      def 			lowerCAmelCase__						(				self: Dict   ):
             __lowerCamelCase   =						OrderedDict(
                 [
                     ("""input_features""", {0: """batch""", 1: """feature_size""", 2: """encoder_sequence"""}),
                 ]   )
             if self.use_past:
                    __lowerCamelCase   =						{0: '''batch'''}
             else:
                    __lowerCamelCase   =						{0: '''batch''', 1: '''decoder_sequence'''}

             if self.use_past:
                    self.fill_with_past_key_values_(UpperCamelCase__    ,       direction="""inputs"""   )

             return common_inputs
      def 			lowerCAmelCase__						(				self: Union[str, Any]    ,       UpperCamelCase_: Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]    ,       UpperCamelCase_: int = -1    ,       UpperCamelCase_: int = -1    ,       UpperCamelCase_: bool = False    ,       UpperCamelCase_: Optional["TensorType"] = None    ,       UpperCamelCase_: int = 2_20_50    ,       UpperCamelCase_: float = 5.0    ,       UpperCamelCase_: int = 2_20    ,       ):
             __lowerCamelCase   =						OrderedDict()
             __lowerCamelCase   =						OnnxConfig.generate_dummy_inputs(
                 self    ,       preprocessor=preprocessor.feature_extractor    ,       batch_size=UpperCamelCase__    ,       framework=UpperCamelCase__    ,       sampling_rate=UpperCamelCase__    ,       time_duration=UpperCamelCase__    ,       frequency=UpperCamelCase__    ,       )
             __lowerCamelCase   =						encoder_inputs['''input_features'''].shape[2]
             __lowerCamelCase   =						encoder_sequence_length // 2 if self.use_past else seq_length

             __lowerCamelCase   =						super().generate_dummy_inputs(
                 preprocessor.tokenizer    ,       UpperCamelCase__    ,       UpperCamelCase__    ,       UpperCamelCase__    ,       UpperCamelCase__   )

             __lowerCamelCase   =						encoder_inputs.pop("""input_features"""   )
             __lowerCamelCase   =						decoder_inputs.pop("""decoder_input_ids"""   )

             if "past_key_values" in decoder_inputs:
                    __lowerCamelCase   =						decoder_inputs.pop("""past_key_values"""   )

             return dummy_inputs




      @property
      def 			lowerCAmelCase__						(				self: Optional[Any]   ):
             return 1E-3


import argparse

import torch

from transformers import YosoConfig, YosoForMaskedLM
def       A	(    _lowercase						):
						if "model" in orig_key:
												SCREAMING_SNAKE_CASE							:    int							=							orig_key.replace('''model.'''		,  ''''''						)
						if "norm1" in orig_key:
												SCREAMING_SNAKE_CASE							:    List[Any]							=							orig_key.replace('''norm1'''		,  '''attention.output.LayerNorm'''						)
						if "norm2" in orig_key:
												SCREAMING_SNAKE_CASE							:    str							=							orig_key.replace('''norm2'''		,  '''output.LayerNorm'''						)
						if "norm" in orig_key:
												SCREAMING_SNAKE_CASE							:    Tuple							=							orig_key.replace('''norm'''		,  '''LayerNorm'''						)
						if "transformer" in orig_key:
												SCREAMING_SNAKE_CASE							:    int							=							orig_key.split('''.'''						)[0].split('''_'''						)[-1]
												SCREAMING_SNAKE_CASE							:    List[str]							=							orig_key.replace(f"""transformer_{layer_num}"""		,  f"""encoder.layer.{layer_num}"""						)
						if "mha.attn" in orig_key:
												SCREAMING_SNAKE_CASE							:    Any							=							orig_key.replace('''mha.attn'''		,  '''attention.self'''						)
						if "mha" in orig_key:
												SCREAMING_SNAKE_CASE							:    Optional[int]							=							orig_key.replace('''mha'''		,  '''attention'''						)
						if "W_q" in orig_key:
												SCREAMING_SNAKE_CASE							:    List[Any]							=							orig_key.replace('''W_q'''		,  '''self.query'''						)
						if "W_k" in orig_key:
												SCREAMING_SNAKE_CASE							:    Optional[int]							=							orig_key.replace('''W_k'''		,  '''self.key'''						)
						if "W_v" in orig_key:
												SCREAMING_SNAKE_CASE							:    str							=							orig_key.replace('''W_v'''		,  '''self.value'''						)
						if "ff1" in orig_key:
												SCREAMING_SNAKE_CASE							:    Any							=							orig_key.replace('''ff1'''		,  '''intermediate.dense'''						)
						if "ff2" in orig_key:
												SCREAMING_SNAKE_CASE							:    List[Any]							=							orig_key.replace('''ff2'''		,  '''output.dense'''						)
						if "ff" in orig_key:
												SCREAMING_SNAKE_CASE							:    Dict							=							orig_key.replace('''ff'''		,  '''output.dense'''						)
						if "mlm_class" in orig_key:
												SCREAMING_SNAKE_CASE							:    Optional[int]							=							orig_key.replace('''mlm.mlm_class'''		,  '''cls.predictions.decoder'''						)
						if "mlm" in orig_key:
												SCREAMING_SNAKE_CASE							:    str							=							orig_key.replace('''mlm'''		,  '''cls.predictions.transform'''						)
						if "cls" not in orig_key:
												SCREAMING_SNAKE_CASE							:    List[str]							=							'''yoso.''' + orig_key

						return orig_key
def       A	(    _lowercase		,  _lowercase						):
						for key in orig_state_dict.copy().keys():
												SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							orig_state_dict.pop(_lowercase						)

												if ("pooler" in key) or ("sen_class" in key):
																		continue
												else:
																		SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							val

						SCREAMING_SNAKE_CASE							:    List[str]							=							orig_state_dict['''cls.predictions.decoder.bias''']
						SCREAMING_SNAKE_CASE							:    Dict							=							torch.arange(_lowercase						).expand((1, -1)						) + 2

						return orig_state_dict
def       A	(    _lowercase		,  _lowercase		,  _lowercase						):
						SCREAMING_SNAKE_CASE							:    Tuple							=							torch.load(_lowercase		,  map_location='''cpu'''						)['''model_state_dict''']
						SCREAMING_SNAKE_CASE							:    List[Any]							=							YosoConfig.from_json_file(_lowercase						)
						SCREAMING_SNAKE_CASE							:    str							=							YosoForMaskedLM(_lowercase						)

						SCREAMING_SNAKE_CASE							:    Union[str, Any]							=							convert_checkpoint_helper(config.max_position_embeddings		,  _lowercase						)

						print(model.load_state_dict(_lowercase						)						)
						model.eval()
						model.save_pretrained(_lowercase						)

						print(f"""Checkpoint successfuly converted. Model saved at {pytorch_dump_path}"""						)


if __name__ == "__main__":
					__UpperCamelCase		:       Tuple									=     argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    '--pytorch_model_path', default=None, type=str, required=True, help='Path to YOSO pytorch checkpoint.'
					)
					parser.add_argument(
					    '--config_file',
					    default=None,
					    type=str,
					    required=True,
					    help='The json file for YOSO model config.',
					)
					parser.add_argument(
					    '--pytorch_dump_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
					)
					__UpperCamelCase		:       Optional[Any]									=     parser.parse_args()
					convert_yoso_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path)





