





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


__lowerCamelCase	:					int			  =					{
    '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_    (					_lowerCAmelCase					)						->			int:
 UpperCamelCase     : Optional[int]	     =       ["layers", "blocks"]
 for k in ignore_keys:
  state_dict.pop(UpperCamelCase__							,       UpperCamelCase__					)


__lowerCamelCase	:					Any			  =					{
    '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_    (					_lowerCAmelCase					)						->			List[str]:
 UpperCamelCase     : Any	     =       list(s_dict.keys()					)
 for key in keys:
  UpperCamelCase     : str	     =       key
  for k, v in WHISPER_MAPPING.items():
   if k in key:
    UpperCamelCase     : Any	     =       new_key.replace(UpperCamelCase__							,       UpperCamelCase__					)

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

  UpperCamelCase     : Optional[Any]	     =       s_dict.pop(UpperCamelCase__					)
 return s_dict



def 		A_    (					_lowerCAmelCase					)						->			Union[str, Any]:
 UpperCamelCase     ,   UpperCamelCase     : List[Any]	     =       emb.weight.shape
 UpperCamelCase     : List[Any]	     =       nn.Linear(UpperCamelCase__							,       UpperCamelCase__							,       bias=UpperCamelCase__					)
 UpperCamelCase     : Any	     =       emb.weight.data
 return lin_layer



def 		A_    (					_lowerCAmelCase							,       _lowerCAmelCase					)						->			bytes:
 os.makedirs(UpperCamelCase__							,       exist_ok=UpperCamelCase__					)
 UpperCamelCase     : int	     =       os.path.basename(UpperCamelCase__					)

 UpperCamelCase     : Optional[int]	     =       url.split("/"					)[-2]
 UpperCamelCase     : List[Any]	     =       os.path.join(UpperCamelCase__							,       UpperCamelCase__					)

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

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

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

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

 UpperCamelCase     : Union[str, Any]	     =       WhisperConfig(
     vocab_size=dimensions["n_vocab"]							,       encoder_ffn_dim=UpperCamelCase__							,       decoder_ffn_dim=UpperCamelCase__							,       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     : int	     =       WhisperForConditionalGeneration(UpperCamelCase__					)
 UpperCamelCase     ,   UpperCamelCase     : Union[str, Any]	     =       model.model.load_state_dict(UpperCamelCase__							,       strict=UpperCamelCase__					)
 if len(UpperCamelCase__					) > 0 and not set(UpperCamelCase__					) <= {
     "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     : List[str]	     =       make_linear_from_emb(model.model.decoder.embed_tokens					)
 else:
  UpperCamelCase     : Dict	     =       proj_out_weights

 model.save_pretrained(UpperCamelCase__					)


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

    convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)










from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
)


__lowerCamelCase	:					str			  =					{"""configuration_encoder_decoder""": ["""EncoderDecoderConfig"""]}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase	:					Optional[Any]			  =					["""EncoderDecoderModel"""]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase	:					Optional[Any]			  =					["""TFEncoderDecoderModel"""]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    __lowerCamelCase	:					Union[str, Any]			  =					["""FlaxEncoderDecoderModel"""]

if TYPE_CHECKING:
    from .configuration_encoder_decoder import EncoderDecoderConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_encoder_decoder import EncoderDecoderModel

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_encoder_decoder import TFEncoderDecoderModel

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_encoder_decoder import FlaxEncoderDecoderModel

else:
    import sys

    __lowerCamelCase	:					Dict			  =					_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
