


# flake8: noqa
# Lint as: python3
__UpperCamelCase    :						Dict     =							[
    'VerificationMode',
    'Version',
    'disable_progress_bar',
    'enable_progress_bar',
    'is_progress_bar_enabled',
    'experimental',
]

from .info_utils import VerificationMode
from .logging import disable_progress_bar, enable_progress_bar, is_progress_bar_enabled
from .version import Version
from .experimental import experimental







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    :						Tuple     =							{
    '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   _UpperCAmelCase       (     UpperCAmelCase      :			Dict      ):
      """simple docstring"""




      __lowerCamelCase     :      int       						=						["""layers""", """blocks"""]
      for k in ignore_keys:
            state_dict.pop(UpperCAmelCase						,	UpperCAmelCase      )


__UpperCamelCase    :						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   _UpperCAmelCase       (     UpperCAmelCase      :			Dict      ):
      """simple docstring"""




      __lowerCamelCase     :      List[str]       						=						list(s_dict.keys()      )
      for key in keys:
            __lowerCamelCase     :      Union[str, Any]       						=						key
            for k, v in WHISPER_MAPPING.items():
                  if k in key:
                        __lowerCamelCase     :      Optional[int]       						=						new_key.replace(UpperCAmelCase						,	UpperCAmelCase      )

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

            __lowerCamelCase     :      Dict       						=						s_dict.pop(UpperCAmelCase      )
      return s_dict





def   _UpperCAmelCase       (     UpperCAmelCase      :			Optional[int]      ):
      """simple docstring"""




      __lowerCamelCase						, __lowerCamelCase     :      Dict       						=						emb.weight.shape
      __lowerCamelCase     :      Tuple       						=						nn.Linear(UpperCAmelCase						,	UpperCAmelCase						,	bias=UpperCAmelCase      )
      __lowerCamelCase     :      Union[str, Any]       						=						emb.weight.data
      return lin_layer





def   _UpperCAmelCase       (     UpperCAmelCase      :			str						,	UpperCAmelCase      :			str      ):
      """simple docstring"""




      os.makedirs(UpperCAmelCase						,	exist_ok=UpperCAmelCase      )
      __lowerCamelCase     :      Any       						=						os.path.basename(UpperCAmelCase      )

      __lowerCamelCase     :      int       						=						url.split("""/"""      )[-2]
      __lowerCamelCase     :      str       						=						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      ):
            __lowerCamelCase     :      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=1_024      ) as loop:
                  while True:
                        __lowerCamelCase     :      Union[str, Any]       						=						source.read(8_192      )
                        if not buffer:
                              break

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

      __lowerCamelCase     :      List[str]       						=						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   _UpperCAmelCase       (     UpperCAmelCase      :			List[Any]						,	UpperCAmelCase      :			Optional[int]      ):
      """simple docstring"""




      if ".pt" not in checkpoint_path:
            __lowerCamelCase     :      Optional[int]       						=						_download(_MODELS[checkpoint_path]      )
      else:
            __lowerCamelCase     :      Union[str, Any]       						=						torch.load(UpperCAmelCase						,	map_location="""cpu"""      )
      __lowerCamelCase     :      Any       						=						original_checkpoint["""dims"""]
      __lowerCamelCase     :      int       						=						original_checkpoint["""model_state_dict"""]
      __lowerCamelCase     :      Tuple       						=						state_dict["""decoder.token_embedding.weight"""]
      remove_ignore_keys_(UpperCAmelCase      )
      rename_keys(UpperCAmelCase      )
      __lowerCamelCase     :      int       						=						True
      __lowerCamelCase     :      Tuple       						=						state_dict["""decoder.layers.0.fc1.weight"""].shape[0]

      __lowerCamelCase     :      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"""]						,	)

      __lowerCamelCase     :      List[str]       						=						WhisperForConditionalGeneration(UpperCAmelCase      )
      __lowerCamelCase						, __lowerCamelCase     :      str       						=						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:
            __lowerCamelCase     :      Tuple       						=						make_linear_from_emb(model.model.decoder.embed_tokens      )
      else:
            __lowerCamelCase     :      Tuple       						=						proj_out_weights

      model.save_pretrained(UpperCAmelCase      )


if __name__ == "__main__":
    __UpperCamelCase    :						Union[str, Any]     =							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    :						Tuple     =							parser.parse_args()

    convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)



