
import argparse
import hashlib
import os
import urllib
import warnings

import torch
from torch import nn
from tqdm import tqdm

from transformers import WhisperConfig, WhisperForConditionalGeneration


_A     =  {
    "tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
    "tiny": "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
    "base.en": "https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt",
    "base": "https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt",
    "small.en": "https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt",
    "small": "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt",
    "medium.en": "https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt",
    "medium": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
    "large": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large.pt",
    "large-v2": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
}






def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[Any]       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			["layers", "blocks"]
			for k in ignore_keys:
						state_dict.pop(__lowerCAmelCase					,							__lowerCAmelCase       )


_A     =  {
    "blocks": "layers",
    "mlp.0": "fc1",
    "mlp.2": "fc2",
    "mlp_ln": "final_layer_norm",
    ".attn.query": ".self_attn.q_proj",
    ".attn.key": ".self_attn.k_proj",
    ".attn.value": ".self_attn.v_proj",
    ".attn_ln": ".self_attn_layer_norm",
    ".attn.out": ".self_attn.out_proj",
    ".cross_attn.query": ".encoder_attn.q_proj",
    ".cross_attn.key": ".encoder_attn.k_proj",
    ".cross_attn.value": ".encoder_attn.v_proj",
    ".cross_attn_ln": ".encoder_attn_layer_norm",
    ".cross_attn.out": ".encoder_attn.out_proj",
    "decoder.ln.": "decoder.layer_norm.",
    "encoder.ln.": "encoder.layer_norm.",
    "token_embedding": "embed_tokens",
    "encoder.positional_embedding": "encoder.embed_positions.weight",
    "decoder.positional_embedding": "decoder.embed_positions.weight",
    "ln_post": "layer_norm",
}






def      lowerCamelCase__       (       __lowerCAmelCase						:						str       ):




			"""simple docstring"""


			lowerCAmelCase_      			=			list(s_dict.keys()       )
			for key in keys:
						lowerCAmelCase_      			=			key
						for k, v in WHISPER_MAPPING.items():
									if k in key:
												lowerCAmelCase_      			=			new_key.replace(__lowerCAmelCase					,							__lowerCAmelCase       )

						print(F"""{key} -> {new_key}"""       )

						lowerCAmelCase_      			=			s_dict.pop(__lowerCAmelCase       )
			return s_dict






def      lowerCamelCase__       (       __lowerCAmelCase						:						Union[str, Any]       ):




			"""simple docstring"""


			lowerCAmelCase_			,						lowerCAmelCase_      			=			emb.weight.shape
			lowerCAmelCase_      			=			nn.Linear(__lowerCAmelCase					,							__lowerCAmelCase					,							bias=__lowerCAmelCase       )
			lowerCAmelCase_      			=			emb.weight.data
			return lin_layer






def      lowerCamelCase__       (       __lowerCAmelCase						:						str					,							__lowerCAmelCase						:						str       ):




			"""simple docstring"""


			os.makedirs(__lowerCAmelCase					,							exist_ok=__lowerCAmelCase       )
			lowerCAmelCase_      			=			os.path.basename(__lowerCAmelCase       )

			lowerCAmelCase_      			=			url.split("/"       )[-2]
			lowerCAmelCase_      			=			os.path.join(__lowerCAmelCase					,							__lowerCAmelCase       )

			if os.path.exists(__lowerCAmelCase       ) and not os.path.isfile(__lowerCAmelCase       ):
						raise RuntimeError(F"""{download_target} exists and is not a regular file"""       )

			if os.path.isfile(__lowerCAmelCase       ):
						lowerCAmelCase_      			=			open(__lowerCAmelCase					,							"rb"       ).read()
						if hashlib.shaaaa(__lowerCAmelCase       ).hexdigest() == expected_shaaaa:
									return model_bytes
						else:
									warnings.warn(F"""{download_target} exists, but the SHA256 checksum does not match; re-downloading the file"""       )

			with urllib.request.urlopen(__lowerCAmelCase       ) as source, open(__lowerCAmelCase					,							"wb"       ) as output:
						with tqdm(
						    total=int(source.info().get("Content-Length"       )       )					,							ncols=80					,							unit="iB"					,							unit_scale=__lowerCAmelCase					,							unit_divisor=1024       ) as loop:
									while True:
												lowerCAmelCase_      			=			source.read(8192       )
												if not buffer:
															break

												output.write(__lowerCAmelCase       )
												loop.update(len(__lowerCAmelCase       )       )

			lowerCAmelCase_      			=			open(__lowerCAmelCase					,							"rb"       ).read()
			if hashlib.shaaaa(__lowerCAmelCase       ).hexdigest() != expected_shaaaa:
						raise RuntimeError(
						    "Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model."       )

			return model_bytes






def      lowerCamelCase__       (       __lowerCAmelCase						:						Optional[int]					,							__lowerCAmelCase						:						Any       ):




			"""simple docstring"""


			if ".pt" not in checkpoint_path:
						lowerCAmelCase_      			=			_download(_MODELS[checkpoint_path]       )
			else:
						lowerCAmelCase_      			=			torch.load(__lowerCAmelCase					,							map_location="cpu"       )
			lowerCAmelCase_      			=			original_checkpoint["dims"]
			lowerCAmelCase_      			=			original_checkpoint["model_state_dict"]
			lowerCAmelCase_      			=			state_dict["decoder.token_embedding.weight"]
			remove_ignore_keys_(__lowerCAmelCase       )
			rename_keys(__lowerCAmelCase       )
			lowerCAmelCase_      			=			True
			lowerCAmelCase_      			=			state_dict["decoder.layers.0.fc1.weight"].shape[0]

			lowerCAmelCase_      			=			WhisperConfig(
			    vocab_size=dimensions["n_vocab"]					,							encoder_ffn_dim=__lowerCAmelCase					,							decoder_ffn_dim=__lowerCAmelCase					,							num_mel_bins=dimensions["n_mels"]					,							d_model=dimensions["n_audio_state"]					,							max_target_positions=dimensions["n_text_ctx"]					,							encoder_layers=dimensions["n_audio_layer"]					,							encoder_attention_heads=dimensions["n_audio_head"]					,							decoder_layers=dimensions["n_text_layer"]					,							decoder_attention_heads=dimensions["n_text_state"]					,							max_source_positions=dimensions["n_audio_ctx"]					,							)

			lowerCAmelCase_      			=			WhisperForConditionalGeneration(__lowerCAmelCase       )
			lowerCAmelCase_			,						lowerCAmelCase_      			=			model.model.load_state_dict(__lowerCAmelCase					,							strict=__lowerCAmelCase       )
			if len(__lowerCAmelCase       ) > 0 and not set(__lowerCAmelCase       ) <= {
			    "encoder.embed_positions.weights",
			    "decoder.embed_positions.weights",
			}:
						raise ValueError(
						    "Only `encoder.embed_positions.weights` and `decoder.embed_positions.weights`  are allowed to be missing,"
						    F""" but all the following weights are missing {missing}"""       )

			if tie_embeds:
						lowerCAmelCase_      			=			make_linear_from_emb(model.model.decoder.embed_tokens       )
			else:
						lowerCAmelCase_      			=			proj_out_weights

			model.save_pretrained(__lowerCAmelCase       )


if __name__ == "__main__":
			_A     =  argparse.ArgumentParser()
			# # Required parameters
			parser.add_argument("--checkpoint_path", type=str, help="Patht to the downloaded checkpoints")
			parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
			_A     =  parser.parse_args()

			convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)





from string import ascii_lowercase, ascii_uppercase






def      lowerCamelCase__       (       __lowerCAmelCase						:						str       ):




			"""simple docstring"""


			if not sentence:
						return ""
			lowerCAmelCase_      			=			dict(zip(__lowerCAmelCase					,							__lowerCAmelCase       )       )
			return lower_to_upper.get(sentence[0]					,							sentence[0]       ) + sentence[1:]


if __name__ == "__main__":
			from doctest import testmod

			testmod()



