



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

snake_case    :						List[str]  =						logging.get_logger(__name__)

snake_case    :						Optional[Any]  =						{
    'openai/whisper-base': 'https://huggingface.co/openai/whisper-base/resolve/main/config.json',
}

# fmt: off
snake_case    :						Optional[int]  =						[
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_7, 3_6_6, 4_3_8, 5_3_2, 6_8_5,
    7_0_5, 7_9_6, 9_3_0, 1_0_5_8, 1_2_2_0, 1_2_6_7, 1_2_7_9, 1_3_0_3, 1_3_4_3, 1_3_7_7,
    1_3_9_1, 1_6_3_5, 1_7_8_2, 1_8_7_5, 2_1_6_2, 2_3_6_1, 2_4_8_8, 3_4_6_7, 4_0_0_8, 4_2_1_1,
    4_6_0_0, 4_8_0_8, 5_2_9_9, 5_8_5_5, 6_3_2_9, 7_2_0_3, 9_6_0_9, 9_9_5_9, 1_0_5_6_3, 1_0_7_8_6,
    1_1_4_2_0, 1_1_7_0_9, 1_1_9_0_7, 1_3_1_6_3, 1_3_6_9_7, 1_3_7_0_0, 1_4_8_0_8, 1_5_3_0_6, 1_6_4_1_0, 1_6_7_9_1,
    1_7_9_9_2, 1_9_2_0_3, 1_9_5_1_0, 2_0_7_2_4, 2_2_3_0_5, 2_2_9_3_5, 2_7_0_0_7, 3_0_1_0_9, 3_0_4_2_0, 3_3_4_0_9,
    3_4_9_4_9, 4_0_2_8_3, 4_0_4_9_3, 4_0_5_4_9, 4_7_2_8_2, 4_9_1_4_6, 5_0_2_5_7, 5_0_3_5_9, 5_0_3_6_0, 5_0_3_6_1
]
snake_case    :						Union[str, Any]  =						[
    1, 2, 7, 8, 9, 1_0, 1_4, 2_5,
    2_6, 2_7, 2_8, 2_9, 3_1, 5_8, 5_9, 6_0, 6_1, 6_2,
    6_3, 9_0, 9_1, 9_2, 9_3, 3_5_9, 5_0_3, 5_2_2, 5_4_2, 8_7_3,
    8_9_3, 9_0_2, 9_1_8, 9_2_2, 9_3_1, 1_3_5_0, 1_8_5_3, 1_9_8_2, 2_4_6_0, 2_6_2_7,
    3_2_4_6, 3_2_5_3, 3_2_6_8, 3_5_3_6, 3_8_4_6, 3_9_6_1, 4_1_8_3, 4_6_6_7, 6_5_8_5, 6_6_4_7,
    7_2_7_3, 9_0_6_1, 9_3_8_3, 1_0_4_2_8, 1_0_9_2_9, 1_1_9_3_8, 1_2_0_3_3, 1_2_3_3_1, 1_2_5_6_2, 1_3_7_9_3,
    1_4_1_5_7, 1_4_6_3_5, 1_5_2_6_5, 1_5_6_1_8, 1_6_5_5_3, 1_6_6_0_4, 1_8_3_6_2, 1_8_9_5_6, 2_0_0_7_5, 2_1_6_7_5,
    2_2_5_2_0, 2_6_1_3_0, 2_6_1_6_1, 2_6_4_3_5, 2_8_2_7_9, 2_9_4_6_4, 3_1_6_5_0, 3_2_3_0_2, 3_2_4_7_0, 3_6_8_6_5,
    4_2_8_6_3, 4_7_4_2_5, 4_9_8_7_0, 5_0_2_5_4, 5_0_2_5_8, 5_0_3_6_0, 5_0_3_6_1, 5_0_3_6_2
]




class 							lowerCAmelCase__  (	UpperCamelCase						):
					__A     :    List[Any]           =						'whisper'
					__A     :    int           =						['past_key_values']
					__A     :    Dict           =						{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}







					def __init__(       self      :     List[Any]			,		_A      :     List[str]=5_1865			,		_A      :     int=80			,		_A      :     List[str]=6			,		_A      :     List[Any]=4			,		_A      :     Tuple=6			,		_A      :     List[str]=4			,		_A      :     Union[str, Any]=1536			,		_A      :     Tuple=1536			,		_A      :     Tuple=0.0			,		_A      :     Optional[int]=0.0			,		_A      :     Dict=5_0257			,		_A      :     Union[str, Any]=True			,		_A      :     Any=True			,		_A      :     Union[str, Any]="gelu"			,		_A      :     Union[str, Any]=256			,		_A      :     Optional[int]=0.0			,		_A      :     Tuple=0.0			,		_A      :     Any=0.0			,		_A      :     Tuple=0.02			,		_A      :     str=False			,		_A      :     Tuple=1500			,		_A      :     str=448			,		_A      :     Optional[Any]=5_0256			,		_A      :     str=5_0256			,		_A      :     Optional[int]=5_0256			,		_A      :     List[str]=None			,		_A      :     Optional[Any]=[220, 5_0256]			,		_A      :     Optional[int]=False			,		_A      :     Optional[Any]=256			,		_A      :     Tuple=False			,		_A      :     Optional[int]=0.05			,		_A      :     Tuple=10			,		_A      :     str=2			,		_A      :     str=0.0			,		_A      :     Optional[Any]=10			,		_A      :     Tuple=0			,		_A      :     List[str]=7			,		**_A      :     List[str]			,		):
											A__	:							List[Any]						    = vocab_size
											A__	:							Any						    = num_mel_bins
											A__	:							List[Any]						    = d_model
											A__	:							int						    = encoder_layers
											A__	:							List[Any]						    = encoder_attention_heads
											A__	:							int						    = decoder_layers
											A__	:							List[str]						    = decoder_attention_heads
											A__	:							List[str]						    = decoder_ffn_dim
											A__	:							List[str]						    = encoder_ffn_dim
											A__	:							List[Any]						    = dropout
											A__	:							List[str]						    = attention_dropout
											A__	:							Dict						    = activation_dropout
											A__	:							List[str]						    = activation_function
											A__	:							str						    = init_std
											A__	:							Any						    = encoder_layerdrop
											A__	:							Dict						    = decoder_layerdrop
											A__	:							Tuple						    = use_cache
											A__	:							Tuple						    = encoder_layers
											A__	:							Any						    = scale_embedding  # scale factor will be sqrt(d_model) if True
											A__	:							int						    = max_source_positions
											A__	:							Dict						    = max_target_positions

											# Audio Classification-specific parameters. Feel free to ignore for other classes.
											A__	:							Optional[int]						    = classifier_proj_size
											A__	:							int						    = use_weighted_layer_sum

											# fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
											A__	:							Dict						    = apply_spec_augment
											A__	:							str						    = mask_time_prob
											A__	:							Optional[int]						    = mask_time_length
											A__	:							Any						    = mask_time_min_masks
											A__	:							Tuple						    = mask_feature_prob
											A__	:							Union[str, Any]						    = mask_feature_length
											A__	:							List[Any]						    = mask_feature_min_masks

											A__	:							Any						    = median_filter_width

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




class 							lowerCAmelCase__  (	UpperCamelCase						):







					@property
					def  _lowercase			(       self      :     Optional[Any]):
											A__	:							Optional[Any]						    = OrderedDict(
											    [
											        ("input_features", {0: "batch", 1: "feature_size", 2: "encoder_sequence"}),
											    ])
											if self.use_past:
																	A__	:							str						    = {0: "batch"}
											else:
																	A__	:							Union[str, Any]						    = {0: "batch", 1: "decoder_sequence"}

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

											return common_inputs







					def  _lowercase			(       self      :     Union[str, Any]			,		_A      :     Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]			,		_A      :     int = -1			,		_A      :     int = -1			,		_A      :     bool = False			,		_A      :     Optional["TensorType"] = None			,		_A      :     int = 2_2050			,		_A      :     float = 5.0			,		_A      :     int = 220			,		):
											A__	:							List[Any]						    = OrderedDict()
											A__	:							Dict						    = OnnxConfig.generate_dummy_inputs(
											    self			,		preprocessor=preprocessor.feature_extractor			,		batch_size=_A			,		framework=_A			,		sampling_rate=_A			,		time_duration=_A			,		frequency=_A			,		)
											A__	:							Dict						    = encoder_inputs["input_features"].shape[2]
											A__	:							List[Any]						    = encoder_sequence_length // 2 if self.use_past else seq_length

											A__	:							List[Any]						    = super().generate_dummy_inputs(
											    preprocessor.tokenizer			,		_A			,		_A			,		_A			,		_A)

											A__	:							Tuple						    = encoder_inputs.pop("input_features")
											A__	:							Optional[Any]						    = decoder_inputs.pop("decoder_input_ids")

											if "past_key_values" in decoder_inputs:
																	A__	:							Optional[int]						    = decoder_inputs.pop("past_key_values")

											return dummy_inputs







					@property
					def  _lowercase			(       self      :     List[Any]):
											return 1e-3




import argparse

import tensorflow as tf
import torch

from transformers import BertConfig, BertForMaskedLM
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertPooler,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


logging.set_verbosity_info()
def 	snake_case__       (					__lowercase							,     __lowercase							,     __lowercase       )		->		List[Any]:







						"""simple docstring"""

						def get_masked_lm_array(__lowercase       ):
												A__	:							int						    = F'masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE'
												A__	:							Dict						    = tf.train.load_variable(__lowercase							,     __lowercase       )

												if "kernel" in name:
																		A__	:							Any						    = array.transpose()

												return torch.from_numpy(__lowercase       )

						def get_encoder_array(__lowercase       ):
												A__	:							Dict						    = F'encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE'
												A__	:							Tuple						    = tf.train.load_variable(__lowercase							,     __lowercase       )

												if "kernel" in name:
																		A__	:							str						    = array.transpose()

												return torch.from_numpy(__lowercase       )

						def get_encoder_layer_array(__lowercase							,     __lowercase       ):
												A__	:							int						    = F'encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE'
												A__	:							Dict						    = tf.train.load_variable(__lowercase							,     __lowercase       )

												if "kernel" in name:
																		A__	:							Any						    = array.transpose()

												return torch.from_numpy(__lowercase       )

						def get_encoder_attention_layer_array(__lowercase							,     __lowercase							,     __lowercase       ):
												A__	:							List[str]						    = F'encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE'
												A__	:							List[Any]						    = tf.train.load_variable(__lowercase							,     __lowercase       )
												A__	:							int						    = array.reshape(__lowercase       )

												if "kernel" in name:
																		A__	:							Union[str, Any]						    = array.transpose()

												return torch.from_numpy(__lowercase       )

						print(F'Loading model based on config from {config_path}...'       )
						A__	:							Tuple						    = BertConfig.from_json_file(__lowercase       )
						A__	:							List[str]						    = BertForMaskedLM(__lowercase       )

						# Layers
						for layer_index in range(0							,     config.num_hidden_layers       ):
												A__	:							BertLayer						    = model.bert.encoder.layer[layer_index]

												# Self-attention
												A__	:							BertSelfAttention						    = layer.attention.self

												A__	:							Optional[Any]						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_query_dense/kernel"							,     self_attn.query.weight.data.shape       )
												A__	:							List[Any]						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_query_dense/bias"							,     self_attn.query.bias.data.shape       )
												A__	:							int						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_key_dense/kernel"							,     self_attn.key.weight.data.shape       )
												A__	:							Tuple						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_key_dense/bias"							,     self_attn.key.bias.data.shape       )
												A__	:							List[Any]						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_value_dense/kernel"							,     self_attn.value.weight.data.shape       )
												A__	:							Optional[Any]						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_value_dense/bias"							,     self_attn.value.bias.data.shape       )

												# Self-attention Output
												A__	:							BertSelfOutput						    = layer.attention.output

												A__	:							Tuple						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_output_dense/kernel"							,     self_output.dense.weight.data.shape       )
												A__	:							Any						    = get_encoder_attention_layer_array(
												    __lowercase							,     "_output_dense/bias"							,     self_output.dense.bias.data.shape       )

												A__	:							Dict						    = get_encoder_layer_array(__lowercase							,     "_attention_layer_norm/gamma"       )
												A__	:							Optional[int]						    = get_encoder_layer_array(__lowercase							,     "_attention_layer_norm/beta"       )

												# Intermediate
												A__	:							BertIntermediate						    = layer.intermediate

												A__	:							Optional[Any]						    = get_encoder_layer_array(__lowercase							,     "_intermediate_dense/kernel"       )
												A__	:							Any						    = get_encoder_layer_array(__lowercase							,     "_intermediate_dense/bias"       )

												# Output
												A__	:							BertOutput						    = layer.output

												A__	:							Dict						    = get_encoder_layer_array(__lowercase							,     "_output_dense/kernel"       )
												A__	:							List[Any]						    = get_encoder_layer_array(__lowercase							,     "_output_dense/bias"       )

												A__	:							Optional[Any]						    = get_encoder_layer_array(__lowercase							,     "_output_layer_norm/gamma"       )
												A__	:							Dict						    = get_encoder_layer_array(__lowercase							,     "_output_layer_norm/beta"       )

						# Embeddings
						A__	:							str						    = get_encoder_array("_position_embedding_layer/embeddings"       )
						A__	:							List[str]						    = get_encoder_array("_type_embedding_layer/embeddings"       )
						A__	:							List[Any]						    = get_encoder_array("_embedding_norm_layer/gamma"       )
						A__	:							Optional[int]						    = get_encoder_array("_embedding_norm_layer/beta"       )

						# LM Head
						A__	:							Optional[Any]						    = model.cls.predictions.transform

						A__	:							int						    = get_masked_lm_array("dense/kernel"       )
						A__	:							Tuple						    = get_masked_lm_array("dense/bias"       )

						A__	:							Optional[Any]						    = get_masked_lm_array("layer_norm/gamma"       )
						A__	:							Tuple						    = get_masked_lm_array("layer_norm/beta"       )

						A__	:							Any						    = get_masked_lm_array("embedding_table"       )

						# Pooling
						A__	:							Optional[Any]						    = BertPooler(config=__lowercase       )
						A__	:							BertPooler						    = get_encoder_array("_pooler_layer/kernel"       )
						A__	:							BertPooler						    = get_encoder_array("_pooler_layer/bias"       )

						# Export final model
						model.save_pretrained(__lowercase       )

						# Integration test - should load without any errors ;)
						A__	:							Union[str, Any]						    = BertForMaskedLM.from_pretrained(__lowercase       )
						print(new_model.eval()       )

						print("Model conversion was done sucessfully!"       )


if __name__ == "__main__":
		snake_case    :						str  =						argparse.ArgumentParser()
		parser.add_argument(
		    '--tf_checkpoint_path', type=str, required=True, help='Path to the TensorFlow Token Dropping checkpoint path.'
		)
		parser.add_argument(
		    '--bert_config_file',
		    type=str,
		    required=True,
		    help='The config json file corresponding to the BERT model. This specifies the model architecture.',
		)
		parser.add_argument(
		    '--pytorch_dump_path',
		    type=str,
		    required=True,
		    help='Path to the output PyTorch model.',
		)
		snake_case    :						Dict  =						parser.parse_args()
		convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)