





'''simple docstring'''

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


lowercase__						:      str	=    {
    '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__    (    lowercase :					Optional[Any]				)		->		str:





     """simple docstring"""

     _UpperCamelCase       =	['''layers''', '''blocks''']
     for k in ignore_keys:
          state_dict.pop(lowercase, lowercase				)


lowercase__						:      List[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__    (    lowercase :					List[str]				)		->		Union[str, Any]:





     """simple docstring"""

     _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(lowercase, lowercase				)

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

          _UpperCamelCase       =	s_dict.pop(lowercase				)
     return s_dict







def  a__    (    lowercase :					Optional[int]				)		->		Optional[Any]:





     """simple docstring"""

     _UpperCamelCase   ,    _UpperCamelCase       =	emb.weight.shape
     _UpperCamelCase       =	nn.Linear(lowercase, lowercase, bias=lowercase				)
     _UpperCamelCase       =	emb.weight.data
     return lin_layer







def  a__    (    lowercase :					str, lowercase :					str				)		->		bytes:





     """simple docstring"""

     os.makedirs(lowercase, exist_ok=lowercase				)
     _UpperCamelCase       =	os.path.basename(lowercase				)

     _UpperCamelCase       =	url.split('''/'''				)[-2]
     _UpperCamelCase       =	os.path.join(lowercase, lowercase				)

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

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

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

     _UpperCamelCase       =	open(lowercase, '''rb'''				).read()
     if hashlib.shaaaa(lowercase				).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__    (    lowercase :					Tuple, lowercase :					List[str]				)		->		List[str]:





     """simple docstring"""

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

     _UpperCamelCase       =	WhisperConfig(
         vocab_size=dimensions['''n_vocab'''], encoder_ffn_dim=lowercase, decoder_ffn_dim=lowercase, 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(lowercase				)
     _UpperCamelCase   ,    _UpperCamelCase       =	model.model.load_state_dict(lowercase, strict=lowercase				)
     if len(lowercase				) > 0 and not set(lowercase				) <= {
         "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(lowercase				)


if __name__ == "__main__":
   lowercase__						:      Optional[int]	=    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.')
   lowercase__						:      Dict	=    parser.parse_args()

   convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)










'''simple docstring'''

from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowercase__						:      Optional[Any]	=    {
    'configuration_autoformer': [
        'AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'AutoformerConfig',
    ],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   lowercase__						:      List[Any]	=    [
       'AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST',
       'AutoformerForPrediction',
       'AutoformerModel',
       'AutoformerPreTrainedModel',
   ]


if TYPE_CHECKING:
   from .configuration_autoformer import (
       AUTOFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP,
       AutoformerConfig,
   )

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_autoformer import (
          AUTOFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
          AutoformerForPrediction,
          AutoformerModel,
          AutoformerPreTrainedModel,
      )

else:
   import sys

   lowercase__						:      Tuple	=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)



