
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				=			{
    """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   (       __lowerCamelCase	:					Optional[Any]       ):
       snake_case			:  Optional[int]     							= ["layers", "blocks"]
       for k in ignore_keys:
              state_dict.pop(__a							,   __a       )


__lowerCamelCase				=			{
    """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   (       __lowerCamelCase	:					Optional[int]       ):
       snake_case			:  List[Any]     							= list(s_dict.keys()       )
       for key in keys:
              snake_case			:  str     							= key
              for k, v in WHISPER_MAPPING.items():
                     if k in key:
                            snake_case			:  Optional[int]     							= new_key.replace(__a							,   __a       )

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

              snake_case			:  str     							= s_dict.pop(__a       )
       return s_dict


def      UpperCamelCase   (       __lowerCamelCase	:					Optional[int]       ):
       snake_case							,   snake_case			:  Dict     							= emb.weight.shape
       snake_case			:  Dict     							= nn.Linear(__a							,   __a							,   bias=__a       )
       snake_case			:  int     							= emb.weight.data
       return lin_layer


def      UpperCamelCase   (       __lowerCamelCase	:					Union[str, Any]							,   __lowerCamelCase	:					List[Any]       ):
       os.makedirs(__a							,   exist_ok=__a       )
       snake_case			:  Optional[int]     							= os.path.basename(__a       )

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

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

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

       snake_case			:  Optional[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"]							,   )

       snake_case			:  str     							= WhisperForConditionalGeneration(__a       )
       snake_case							,   snake_case			:  str     							= 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:
              snake_case			:  List[str]     							= make_linear_from_emb(model.model.decoder.embed_tokens       )
       else:
              snake_case			:  Optional[Any]     							= proj_out_weights

       model.save_pretrained(__a       )


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

   convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)







"""simple docstring"""




import unittest

from datasets import load_dataset

from transformers import BloomTokenizerFast
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin





@require_tokenizers
class 			snake_case  (				__lowercase							,      unittest.TestCase ):
							UpperCAmelCase__      =   None
							UpperCAmelCase__      =   BloomTokenizerFast
							UpperCAmelCase__      =   BloomTokenizerFast
							UpperCAmelCase__      =   True
							UpperCAmelCase__      =   False
							UpperCAmelCase__      =   '''tokenizer_file'''
							UpperCAmelCase__      =   {'''bos_token''': '''<s>''', '''eos_token''': '''</s>''', '''unk_token''': '''<unk>''', '''pad_token''': '''<pad>'''}

							def _lowercase       (self     ):

									"""simple docstring"""







									super().setUp()
									SCREAMING_SNAKE_CASE_    			=			BloomTokenizerFast.from_pretrained('''bigscience/tokenizer'''     )
									tokenizer.save_pretrained(self.tmpdirname     )

							def _lowercase       (self							,							**SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									kwargs.update(self.special_tokens_map     )
									return BloomTokenizerFast.from_pretrained(self.tmpdirname							,							**SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.get_rust_tokenizer()

									SCREAMING_SNAKE_CASE_    			=			['''The quick brown fox</s>''', '''jumps over the lazy dog</s>''']
									SCREAMING_SNAKE_CASE_    			=			[[21_75, 2_37_14, 7_31_73, 14_42_52, 2], [77, 13_26_19, 34_78, 3_68, 10_95_86, 3_54_33, 2]]

									SCREAMING_SNAKE_CASE_    			=			tokenizer.batch_encode_plus(SCREAMING_SNAKE_CASE_     )['''input_ids''']
									self.assertListEqual(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			tokenizer.batch_decode(SCREAMING_SNAKE_CASE_     )
									self.assertListEqual(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_=6     ):

									"""simple docstring"""







									for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
											with self.subTest(f'{tokenizer.__class__.__name__} ({pretrained_name})'     ):
													SCREAMING_SNAKE_CASE_    			=			self.rust_tokenizer_class.from_pretrained(SCREAMING_SNAKE_CASE_							,							**SCREAMING_SNAKE_CASE_     )
													# tokenizer_r.pad_token = None # Hotfixing padding = None
													# Simple input
													SCREAMING_SNAKE_CASE_    			=			'''This is a simple input'''
													SCREAMING_SNAKE_CASE_    			=			['''This is a simple input 1''', '''This is a simple input 2''']
													SCREAMING_SNAKE_CASE_    			=			('''This is a simple input''', '''This is a pair''')
													SCREAMING_SNAKE_CASE_    			=			[
													    ('''This is a simple input 1''', '''This is a simple input 2'''),
													    ('''This is a simple pair 1''', '''This is a simple pair 2'''),
													]

													# Simple input tests
													try:
															tokenizer_r.encode(SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_     )
															tokenizer_r.encode_plus(SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_     )

															tokenizer_r.batch_encode_plus(SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_     )
															tokenizer_r.encode(SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_     )
															tokenizer_r.batch_encode_plus(SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_     )
													except ValueError:
															self.fail('''Bloom Tokenizer should be able to deal with padding'''     )

													SCREAMING_SNAKE_CASE_    			=			None  # Hotfixing padding = None
													self.assertRaises(SCREAMING_SNAKE_CASE_							,							tokenizer_r.encode							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''     )

													# Simple input
													self.assertRaises(SCREAMING_SNAKE_CASE_							,							tokenizer_r.encode_plus							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''     )

													# Simple input
													self.assertRaises(
													    SCREAMING_SNAKE_CASE_							,							tokenizer_r.batch_encode_plus							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''							,							)

													# Pair input
													self.assertRaises(SCREAMING_SNAKE_CASE_							,							tokenizer_r.encode							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''     )

													# Pair input
													self.assertRaises(SCREAMING_SNAKE_CASE_							,							tokenizer_r.encode_plus							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''     )

													# Pair input
													self.assertRaises(
													    SCREAMING_SNAKE_CASE_							,							tokenizer_r.batch_encode_plus							,							SCREAMING_SNAKE_CASE_							,							max_length=SCREAMING_SNAKE_CASE_							,							padding='''max_length'''							,							)

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.get_rust_tokenizer()
									SCREAMING_SNAKE_CASE_    			=			load_dataset('''xnli'''							,							'''all_languages'''							,							split='''test'''							,							streaming=SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_    			=			next(iter(SCREAMING_SNAKE_CASE_     )     )['''premise''']  # pick up one data
									SCREAMING_SNAKE_CASE_    			=			list(sample_data.values()     )

									SCREAMING_SNAKE_CASE_    			=			list(map(tokenizer.encode							,							SCREAMING_SNAKE_CASE_     )     )
									SCREAMING_SNAKE_CASE_    			=			[tokenizer.decode(SCREAMING_SNAKE_CASE_							,							clean_up_tokenization_spaces=SCREAMING_SNAKE_CASE_     ) for x in output_tokens]
									self.assertListEqual(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )







							def _lowercase       (self     ):

									"""simple docstring"""







									self.assertGreaterEqual(len(self.tokenizer_class.pretrained_vocab_files_map     )							,							1     )
									self.assertGreaterEqual(len(list(self.tokenizer_class.pretrained_vocab_files_map.values()     )[0]     )							,							1     )