


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

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__UpperCamelCase    :						Optional[int]     =							logging.get_logger(__name__)

__UpperCamelCase    :						str     =							'▁'

__UpperCamelCase    :						List[str]     =							{'vocab_file': 'sentencepiece.bpe.model'}

__UpperCamelCase    :						Optional[Any]     =							{
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    }
}

__UpperCamelCase    :						List[Any]     =							{
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
__UpperCamelCase    :						Optional[int]     =							['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']



class 							_UpperCamelCase (						A						):

			'''simple docstring'''







			a_				:							Union[str, Any]    = VOCAB_FILES_NAMES
			a_				:							int    = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			a_				:							Tuple    = PRETRAINED_VOCAB_FILES_MAP
			a_				:							List[str]    = ["input_ids", "attention_mask"]

			a_				:							str    = []
			a_				:							int    = []

			def __init__(  self  : List[str]       ,  _lowerCamelCase  : Dict       ,  _lowerCamelCase  : Dict=None       ,  _lowerCamelCase  : List[str]=None       ,  _lowerCamelCase  : List[Any]="</s>"       ,  _lowerCamelCase  : Any="</s>"       ,  _lowerCamelCase  : Optional[Any]="<s>"       ,  _lowerCamelCase  : List[str]="<unk>"       ,  _lowerCamelCase  : Optional[Any]="<pad>"       ,  _lowerCamelCase  : Tuple="<mask>"       ,  _lowerCamelCase  : Optional[Dict[str, Any]] = None       ,  **_lowerCamelCase  : Union[str, Any]       ,  ):






									'''simple docstring'''



									__lowerCamelCase     :      int       						=						AddedToken(__A       ,  lstrip=__A       ,  rstrip=__A							) if isinstance(__A       ,  __A							) else mask_token

									__lowerCamelCase     :      int       						=						{} if sp_model_kwargs is None else sp_model_kwargs

									__lowerCamelCase     :      Optional[Any]       						=						kwargs.get("""additional_special_tokens"""       ,  []							)
									kwargs["additional_special_tokens"] += [
									    code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
									]

									super().__init__(
									    src_lang=__A       ,  tgt_lang=__A       ,  eos_token=__A       ,  unk_token=__A       ,  sep_token=__A       ,  cls_token=__A       ,  pad_token=__A       ,  mask_token=__A       ,  sp_model_kwargs=self.sp_model_kwargs       ,  **__A       ,  )

									__lowerCamelCase     :      Optional[Any]       						=						spm.SentencePieceProcessor(**self.sp_model_kwargs							)
									self.sp_model.Load(str(__A							)							)
									__lowerCamelCase     :      Tuple       						=						vocab_file

									# Original fairseq vocab and spm vocab must be "aligned":
									# Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
									# -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
									# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
									# spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

									# Mimic fairseq token-to-id alignment for the first 4 token
									__lowerCamelCase     :      Optional[int]       						=						{"""<s>""": 0, """<pad>""": 1, """</s>""": 2, """<unk>""": 3}

									# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
									__lowerCamelCase     :      Union[str, Any]       						=						1

									__lowerCamelCase     :      Any       						=						len(self.sp_model							)
									__lowerCamelCase     :      Union[str, Any]       						=						{
									    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(__A							)
									}
									__lowerCamelCase     :      Any       						=						{v: k for k, v in self.lang_code_to_id.items()}
									__lowerCamelCase     :      Optional[int]       						=						len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset

									self.fairseq_tokens_to_ids.update(self.lang_code_to_id							)
									__lowerCamelCase     :      int       						=						{v: k for k, v in self.fairseq_tokens_to_ids.items()}

									__lowerCamelCase     :      Dict       						=						src_lang if src_lang is not None else """en_XX"""
									__lowerCamelCase     :      Any       						=						self.lang_code_to_id[self._src_lang]
									__lowerCamelCase     :      Union[str, Any]       						=						tgt_lang
									self.set_src_lang_special_tokens(self._src_lang							)

			@property
			def   _snake_case    (  self  : str							):






									'''simple docstring'''



									return len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset + 1  # Plus 1 for the mask token

			@property
			def   _snake_case    (  self  : Any							):






									'''simple docstring'''



									return self._src_lang

			@src_lang.setter
			def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									__lowerCamelCase     :      int       						=						new_src_lang
									self.set_src_lang_special_tokens(self._src_lang							)

			def __getstate__(  self  : str							):






									'''simple docstring'''



									__lowerCamelCase     :      Any       						=						self.__dict__.copy()
									__lowerCamelCase     :      List[str]       						=						None
									return state

			def __setstate__(  self  : int       ,  _lowerCamelCase  : Dict							):






									'''simple docstring'''



									__lowerCamelCase     :      Tuple       						=						d

									# for backward compatibility
									if not hasattr(self       ,  """sp_model_kwargs"""							):
															__lowerCamelCase     :      List[Any]       						=						{}

									__lowerCamelCase     :      Any       						=						spm.SentencePieceProcessor(**self.sp_model_kwargs							)
									self.sp_model.Load(self.vocab_file							)

			def   _snake_case    (  self  : List[Any]							):






									'''simple docstring'''



									__lowerCamelCase     :      List[str]       						=						{self.convert_ids_to_tokens(__A							): i for i in range(self.vocab_size							)}
									vocab.update(self.added_tokens_encoder							)
									return vocab

			def   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									return self.sp_model.encode(__A       ,  out_type=__A							)

			def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									if token in self.fairseq_tokens_to_ids:
															return self.fairseq_tokens_to_ids[token]
									__lowerCamelCase     :      List[Any]       						=						self.sp_model.PieceToId(__A							)

									# Need to return unknown token if the SP model returned 0
									return spm_id + self.fairseq_offset if spm_id else self.unk_token_id

			def   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : int							):






									'''simple docstring'''



									if index in self.fairseq_ids_to_tokens:
															return self.fairseq_ids_to_tokens[index]
									return self.sp_model.IdToPiece(index - self.fairseq_offset							)

			def   _snake_case    (  self  : List[Any]       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									__lowerCamelCase     :      Optional[int]       						=						[]
									__lowerCamelCase     :      Dict       						=						""""""
									__lowerCamelCase     :      List[str]       						=						False
									for token in tokens:
															# make sure that special tokens are not decoded using sentencepiece model
															if token in self.all_special_tokens:
																					if not prev_is_special:
																											out_string += " "
																					out_string += self.sp_model.decode(__A							) + token
																					__lowerCamelCase     :      Any       						=						True
																					__lowerCamelCase     :      Optional[int]       						=						[]
															else:
																					current_sub_tokens.append(__A							)
																					__lowerCamelCase     :      str       						=						False
									out_string += self.sp_model.decode(__A							)
									return out_string.strip()

			def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : Optional[str] = None							):






									'''simple docstring'''



									if not os.path.isdir(__A							):
															logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""							)
															return
									__lowerCamelCase     :      Optional[int]       						=						os.path.join(
									    __A       ,  (filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

									if os.path.abspath(self.vocab_file							) != os.path.abspath(__A							) and os.path.isfile(self.vocab_file							):
															copyfile(self.vocab_file       ,  __A							)
									elif not os.path.isfile(self.vocab_file							):
															with open(__A       ,  """wb"""							) as fi:
																					__lowerCamelCase     :      List[str]       						=						self.sp_model.serialized_model_proto()
																					fi.write(__A							)

									return (out_vocab_file,)

			def   _snake_case    (  self  : Optional[int]       ,  _lowerCamelCase  : List[int]       ,  _lowerCamelCase  : Optional[List[int]] = None       ,  _lowerCamelCase  : bool = False							):






									'''simple docstring'''



									if already_has_special_tokens:
															return super().get_special_tokens_mask(
															    token_ids_a=__A       ,  token_ids_a=__A       ,  already_has_special_tokens=__A							)

									__lowerCamelCase     :      List[Any]       						=						[1] * len(self.prefix_tokens							)
									__lowerCamelCase     :      Optional[int]       						=						[1] * len(self.suffix_tokens							)
									if token_ids_a is None:
															return prefix_ones + ([0] * len(__A							)) + suffix_ones
									return prefix_ones + ([0] * len(__A							)) + ([0] * len(__A							)) + suffix_ones

			def   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : List[int]       ,  _lowerCamelCase  : Optional[List[int]] = None							):






									'''simple docstring'''



									if token_ids_a is None:
															return self.prefix_tokens + token_ids_a + self.suffix_tokens
									# We don't expect to process pairs, but leave the pair logic for API consistency
									return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens

			def   _snake_case    (  self  : Tuple       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : str       ,  _lowerCamelCase  : Optional[str]       ,  _lowerCamelCase  : Optional[str]       ,  **_lowerCamelCase  : Any							):






									'''simple docstring'''



									if src_lang is None or tgt_lang is None:
															raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""							)
									__lowerCamelCase     :      Any       						=						src_lang
									__lowerCamelCase     :      Optional[int]       						=						self(__A       ,  add_special_tokens=__A       ,  return_tensors=__A       ,  **__A							)
									__lowerCamelCase     :      List[str]       						=						self.convert_tokens_to_ids(__A							)
									__lowerCamelCase     :      List[Any]       						=						tgt_lang_id
									return inputs

			def   _snake_case    (  self  : Optional[int]       ,  _lowerCamelCase  : List[str]       ,  _lowerCamelCase  : str = "en_XX"       ,  _lowerCamelCase  : Optional[List[str]] = None       ,  _lowerCamelCase  : str = "ro_RO"       ,  **_lowerCamelCase  : Optional[Any]       ,  ):






									'''simple docstring'''



									__lowerCamelCase     :      List[str]       						=						src_lang
									__lowerCamelCase     :      Tuple       						=						tgt_lang
									return super().prepare_seqaseq_batch(__A       ,  __A       ,  **__A							)

			def   _snake_case    (  self  : List[str]							):






									'''simple docstring'''



									return self.set_src_lang_special_tokens(self.src_lang							)

			def   _snake_case    (  self  : Tuple							):






									'''simple docstring'''



									return self.set_tgt_lang_special_tokens(self.tgt_lang							)

			def   _snake_case    (  self  : Dict       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									__lowerCamelCase     :      str       						=						self.lang_code_to_id[src_lang]
									__lowerCamelCase     :      Dict       						=						[self.cur_lang_code_id]
									__lowerCamelCase     :      Optional[Any]       						=						[self.eos_token_id]


			def   _snake_case    (  self  : List[str]       ,  _lowerCamelCase  : str							):






									'''simple docstring'''



									__lowerCamelCase     :      str       						=						self.lang_code_to_id[tgt_lang]
									__lowerCamelCase     :      Any       						=						[self.cur_lang_code_id]
									__lowerCamelCase     :      int       						=						[self.eos_token_id]







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)



