
'''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


__snake_case:  Union[str, 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 	_snake_case						(      A_       :    List[str]   ):
  """simple docstring"""







  a_      :	List[Any]      =							["""layers""", """blocks"""]
  for k in ignore_keys:
    state_dict.pop(A_     ,		A_   )


__snake_case:  Dict								=			{
    "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 	_snake_case						(      A_       :    Any   ):
  """simple docstring"""







  a_      :	Tuple      =							list(s_dict.keys()   )
  for key in keys:
    a_      :	Any      =							key
    for k, v in WHISPER_MAPPING.items():
      if k in key:
        a_      :	int      =							new_key.replace(A_     ,		A_   )

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

    a_      :	List[str]      =							s_dict.pop(A_   )
  return s_dict






def 	_snake_case						(      A_       :    Optional[Any]   ):
  """simple docstring"""







  a_      , a_      :	int      =							emb.weight.shape
  a_      :	Dict      =							nn.Linear(A_     ,		A_     ,		bias=A_   )
  a_      :	Optional[Any]      =							emb.weight.data
  return lin_layer






def 	_snake_case						(      A_       :    str     ,		A_       :    str   ):
  """simple docstring"""







  os.makedirs(A_     ,		exist_ok=A_   )
  a_      :	Optional[Any]      =							os.path.basename(A_   )

  a_      :	Union[str, Any]      =							url.split("""/"""   )[-2]
  a_      :	Tuple      =							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_   ):
    a_      :	Tuple      =							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:
        a_      :	int      =							source.read(8192   )
        if not buffer:
          break

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

  a_      :	str      =							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 	_snake_case						(      A_       :    Optional[int]     ,		A_       :    Dict   ):
  """simple docstring"""







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

  a_      :	int      =							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"""]     ,		)

  a_      :	Union[str, Any]      =							WhisperForConditionalGeneration(A_   )
  a_      , a_      :	Dict      =							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:
    a_      :	int      =							make_linear_from_emb(model.model.decoder.embed_tokens   )
  else:
    a_      :	Dict      =							proj_out_weights

  model.save_pretrained(A_   )


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

    convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)




'''simple docstring'''

import os
from shutil import copyfile
from typing import List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


__snake_case:  Optional[Any]								=			logging.get_logger(__name__)

__snake_case:  Tuple								=			{"vocab_file": "sentencepiece.model"}

__snake_case:  int								=			{
    "vocab_file": {
        "google/rembert": "https://huggingface.co/google/rembert/resolve/main/sentencepiece.model",
    },
}

__snake_case:  Union[str, Any]								=			{
    "google/rembert": 2_56,
}


class  _UpperCAmelCase ( lowerCAmelCase__    ):

     """simple docstring"""

     a_    	=					VOCAB_FILES_NAMES
     a_    	=					PRETRAINED_VOCAB_FILES_MAP
     a_    	=					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES


     def __init__(      self							,		lowerCAmelCase_							,		lowerCAmelCase_=False							,		lowerCAmelCase_=True							,		lowerCAmelCase_=True							,		lowerCAmelCase_="[CLS]"							,		lowerCAmelCase_="[SEP]"							,		lowerCAmelCase_="[UNK]"							,		lowerCAmelCase_="[SEP]"							,		lowerCAmelCase_="[PAD]"							,		lowerCAmelCase_="[CLS]"							,		lowerCAmelCase_="[MASK]"							,		**lowerCAmelCase_							,		):

       '''simple docstring'''

       super().__init__(
           do_lower_case=lowerCAmelCase_							,		remove_space=lowerCAmelCase_							,		keep_accents=lowerCAmelCase_							,		bos_token=lowerCAmelCase_							,		eos_token=lowerCAmelCase_							,		unk_token=lowerCAmelCase_							,		sep_token=lowerCAmelCase_							,		pad_token=lowerCAmelCase_							,		cls_token=lowerCAmelCase_							,		mask_token=lowerCAmelCase_							,		**lowerCAmelCase_							,		)

       a_      :	int      =							do_lower_case
       a_      :	int      =							remove_space
       a_      :	Dict      =							keep_accents
       a_      :	Union[str, Any]      =							vocab_file

       a_      :	List[str]      =							spm.SentencePieceProcessor()
       self.sp_model.Load(lowerCAmelCase_ )


     @property
     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       return len(self.sp_model )


     def 							_lowerCAmelCase  (      self ):

       '''simple docstring'''

       a_      :	Dict      =							{self.convert_ids_to_tokens(lowerCAmelCase_ ): i for i in range(self.vocab_size )}
       vocab.update(self.added_tokens_encoder )
       return vocab


     def __getstate__(      self ):

       '''simple docstring'''

       a_      :	Union[str, Any]      =							self.__dict__.copy()
       a_      :	Optional[Any]      =							None
       return state


     def __setstate__(      self							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	int      =							d
       a_      :	Optional[int]      =							spm.SentencePieceProcessor()
       self.sp_model.Load(self.vocab_file )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_=False ):

       '''simple docstring'''

       a_      :	Tuple      =							self.sp_model.EncodeAsPieces(lowerCAmelCase_ )
       return pieces


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_ ):

       '''simple docstring'''

       return self.sp_model.PieceToId(lowerCAmelCase_ )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_ ):

       '''simple docstring'''

       return self.sp_model.IdToPiece(lowerCAmelCase_ )


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_ ):

       '''simple docstring'''

       a_      :	Tuple      =							self.sp_model.decode_pieces(lowerCAmelCase_ )
       return out_string


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_ = None ):

       '''simple docstring'''

       a_      :	List[Any]      =							[self.sep_token_id]
       a_      :	Tuple      =							[self.cls_token_id]
       if token_ids_a is None:
         return cls + token_ids_a + sep
       return cls + token_ids_a + sep + token_ids_a + sep


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_ = None							,		lowerCAmelCase_ = False ):

       '''simple docstring'''

       if already_has_special_tokens:
         if token_ids_a is not None:
           raise ValueError(
               """You should not supply a second sequence if the provided sequence of """
               """ids is already formatted with special tokens for the model.""" )
         return [1 if x in [self.sep_token_id, self.cls_token_id] else 0 for x in token_ids_a]

       if token_ids_a is not None:
         return [1] + ([0] * len(lowerCAmelCase_ )) + [1] + ([0] * len(lowerCAmelCase_ )) + [1]
       return [1] + ([0] * len(lowerCAmelCase_ )) + [1]


     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_ = None ):

       '''simple docstring'''

       a_      :	Optional[int]      =							[self.sep_token_id]
       a_      :	Optional[Any]      =							[self.cls_token_id]

       if token_ids_a is None:
         return len(cls + token_ids_a + sep ) * [0]
       return len(cls + token_ids_a + sep ) * [0] + len(token_ids_a + sep ) * [1]






     def 							_lowerCAmelCase  (      self							,		lowerCAmelCase_							,		lowerCAmelCase_ = None ):

       '''simple docstring'''

       if not os.path.isdir(lowerCAmelCase_ ):
         logger.error("""Vocabulary path ({}) should be a directory""".format(lowerCAmelCase_ ) )
         return
       a_      :	List[str]      =							os.path.join(
           lowerCAmelCase_							,		(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""] )

       if os.path.abspath(self.vocab_file ) != os.path.abspath(lowerCAmelCase_ ):
         copyfile(self.vocab_file							,		lowerCAmelCase_ )

       return (out_vocab_file,)


