





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


UpperCamelCase__   : Optional[Any]					  = {
    """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    A_(			A   ):
	UpperCAmelCase_								=     ["""layers""", """blocks"""]
	for k in ignore_keys:
		state_dict.pop(A				,      A   )


UpperCamelCase__   : List[str]					  = {
    """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    A_(			A   ):
	UpperCAmelCase_								=     list(s_dict.keys()   )
	for key in keys:
		UpperCAmelCase_								=     key
		for k, v in WHISPER_MAPPING.items():
			if k in key:
				UpperCAmelCase_								=     new_key.replace(A				,      A   )

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

		UpperCAmelCase_								=     s_dict.pop(A   )
	return s_dict




def    A_(			A   ):
	UpperCAmelCase_						,	UpperCAmelCase_								=     emb.weight.shape
	UpperCAmelCase_								=     nn.Linear(A				,      A				,      bias=A   )
	UpperCAmelCase_								=     emb.weight.data
	return lin_layer




def    A_(			A				,      A   ):
	os.makedirs(A				,      exist_ok=A   )
	UpperCAmelCase_								=     os.path.basename(A   )

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

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

	if os.path.isfile(A   ):
		UpperCAmelCase_								=     open(A				,      """rb"""   ).read()
		if hashlib.shaaaa(A   ).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(A   ) as source, open(A				,      """wb"""   ) as output:
		with tqdm(
		    total=int(source.info().get("""Content-Length"""   )   )				,      ncols=80				,      unit="""iB"""				,      unit_scale=A				,      unit_divisor=1024   ) as loop:
			while True:
				UpperCAmelCase_								=     source.read(8192   )
				if not buffer:
					break

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

	UpperCAmelCase_								=     open(A				,      """rb"""   ).read()
	if hashlib.shaaaa(A   ).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    A_(			A				,      A   ):
	if ".pt" not in checkpoint_path:
		UpperCAmelCase_								=     _download(_MODELS[checkpoint_path]   )
	else:
		UpperCAmelCase_								=     torch.load(A				,      map_location="""cpu"""   )
	UpperCAmelCase_								=     original_checkpoint["""dims"""]
	UpperCAmelCase_								=     original_checkpoint["""model_state_dict"""]
	UpperCAmelCase_								=     state_dict["""decoder.token_embedding.weight"""]
	remove_ignore_keys_(A   )
	rename_keys(A   )
	UpperCAmelCase_								=     True
	UpperCAmelCase_								=     state_dict["""decoder.layers.0.fc1.weight"""].shape[0]

	UpperCAmelCase_								=     WhisperConfig(
	    vocab_size=dimensions["""n_vocab"""]				,      encoder_ffn_dim=A				,      decoder_ffn_dim=A				,      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"""]				,      )

	UpperCAmelCase_								=     WhisperForConditionalGeneration(A   )
	UpperCAmelCase_						,	UpperCAmelCase_								=     model.model.load_state_dict(A				,      strict=A   )
	if len(A   ) > 0 and not set(A   ) <= {
	    "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:
		UpperCAmelCase_								=     make_linear_from_emb(model.model.decoder.embed_tokens   )
	else:
		UpperCAmelCase_								=     proj_out_weights

	model.save_pretrained(A   )


if __name__ == "__main__":
				UpperCamelCase__   : List[str]					  = 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.""")
				UpperCamelCase__   : Any					  = parser.parse_args()

				convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)









from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand




def    A_(			A   ):
	return ConvertCommand(
	    args.model_type				,      args.tf_checkpoint				,      args.pytorch_dump_output				,      args.config				,      args.finetuning_task_name   )


UpperCamelCase__   : Any					  = """
transformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires
TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.
"""







class   _UpperCamelCase  (       A_  ):







					'''simple docstring'''

					@staticmethod
					def        SCREAMING_SNAKE_CASE							(						__lowercase    :  ArgumentParser		):






						'''simple docstring'''







						UpperCAmelCase_								=     parser.add_parser(
						    """convert""" ,					help="""CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints.""" ,					)
						train_parser.add_argument("""--model_type""" ,					type=__lowercase ,					required=__lowercase ,					help="""Model's type."""		)
						train_parser.add_argument(
						    """--tf_checkpoint""" ,					type=__lowercase ,					required=__lowercase ,					help="""TensorFlow checkpoint path or folder."""		)
						train_parser.add_argument(
						    """--pytorch_dump_output""" ,					type=__lowercase ,					required=__lowercase ,					help="""Path to the PyTorch saved model output."""		)
						train_parser.add_argument("""--config""" ,					type=__lowercase ,					default="""""" ,					help="""Configuration file path or folder."""		)
						train_parser.add_argument(
						    """--finetuning_task_name""" ,					type=__lowercase ,					default=__lowercase ,					help="""Optional fine-tuning task name if the TF model was a finetuned model.""" ,					)
						train_parser.set_defaults(func=__lowercase		)





					def __init__(						self    :  Optional[int] ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					__lowercase    :  str ,					*__lowercase    :  Tuple ,					):






						'''simple docstring'''







						UpperCAmelCase_								=     logging.get_logger("""transformers-cli/converting"""		)

						self._logger.info(F"""Loading model {model_type}"""		)
						UpperCAmelCase_								=     model_type
						UpperCAmelCase_								=     tf_checkpoint
						UpperCAmelCase_								=     pytorch_dump_output
						UpperCAmelCase_								=     config
						UpperCAmelCase_								=     finetuning_task_name





					def        SCREAMING_SNAKE_CASE							(						self    :  List[Any]		):






						'''simple docstring'''







						if self._model_type == "albert":
							try:
								from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "bert":
							try:
								from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "funnel":
							try:
								from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
								    convert_tf_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "t5":
							try:
								from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
							except ImportError:
								raise ImportError(__lowercase		)

							convert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "gpt":
							from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
							    convert_openai_checkpoint_to_pytorch,
							)

							convert_openai_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "transfo_xl":
							try:
								from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
								    convert_transfo_xl_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							if "ckpt" in self._tf_checkpoint.lower():
								UpperCAmelCase_								=     self._tf_checkpoint
								UpperCAmelCase_								=     """"""
							else:
								UpperCAmelCase_								=     self._tf_checkpoint
								UpperCAmelCase_								=     """"""
							convert_transfo_xl_checkpoint_to_pytorch(
							    __lowercase ,					self._config ,					self._pytorch_dump_output ,					__lowercase		)
						elif self._model_type == "gpt2":
							try:
								from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
								    convert_gpta_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						elif self._model_type == "xlnet":
							try:
								from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
								    convert_xlnet_checkpoint_to_pytorch,
								)
							except ImportError:
								raise ImportError(__lowercase		)

							convert_xlnet_checkpoint_to_pytorch(
							    self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output ,					self._finetuning_task_name		)
						elif self._model_type == "xlm":
							from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
							    convert_xlm_checkpoint_to_pytorch,
							)

							convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint ,					self._pytorch_dump_output		)
						elif self._model_type == "lxmert":
							from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
							    convert_lxmert_checkpoint_to_pytorch,
							)

							convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint ,					self._pytorch_dump_output		)
						elif self._model_type == "rembert":
							from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
							    convert_rembert_tf_checkpoint_to_pytorch,
							)

							convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint ,					self._config ,					self._pytorch_dump_output		)
						else:
							raise ValueError(
							    """--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]"""		)


