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


lowerCAmelCase_	:		List[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 						__A							(     UpperCAmelCase					)     ->			int:




						'''simple docstring'''
						_UpperCamelCase			:			Optional[Any]   			=						["layers", "blocks"]
						for k in ignore_keys:
												state_dict.pop(UpperCAmelCase     ,UpperCAmelCase					)


lowerCAmelCase_	:		List[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 						__A							(     UpperCAmelCase					)     ->			Tuple:




						'''simple docstring'''
						_UpperCamelCase			:			List[str]   			=						list(s_dict.keys()					)
						for key in keys:
												_UpperCamelCase			:			Union[str, Any]   			=						key
												for k, v in WHISPER_MAPPING.items():
																		if k in key:
																								_UpperCamelCase			:			Any   			=						new_key.replace(UpperCAmelCase     ,UpperCAmelCase					)

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

												_UpperCamelCase			:			Union[str, Any]   			=						s_dict.pop(UpperCAmelCase					)
						return s_dict
def 						__A							(     UpperCAmelCase					)     ->			Optional[int]:




						'''simple docstring'''
						_UpperCamelCase			,       _UpperCamelCase			:			int   			=						emb.weight.shape
						_UpperCamelCase			:			Tuple   			=						nn.Linear(UpperCAmelCase     ,UpperCAmelCase     ,bias=UpperCAmelCase					)
						_UpperCamelCase			:			Any   			=						emb.weight.data
						return lin_layer
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			bytes:




						'''simple docstring'''
						os.makedirs(UpperCAmelCase     ,exist_ok=UpperCAmelCase					)
						_UpperCamelCase			:			Tuple   			=						os.path.basename(UpperCAmelCase					)

						_UpperCamelCase			:			List[Any]   			=						url.split("/"					)[-2]
						_UpperCamelCase			:			Optional[Any]   			=						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					):
												_UpperCamelCase			:			Dict   			=						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=8_0     ,unit="iB"     ,unit_scale=UpperCAmelCase     ,unit_divisor=1_0_2_4					) as loop:
																		while True:
																								_UpperCamelCase			:			List[str]   			=						source.read(8_1_9_2					)
																								if not buffer:
																														break

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

						_UpperCamelCase			:			Dict   			=						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 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Optional[Any]:




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

						_UpperCamelCase			:			Tuple   			=						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"]     ,)

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

						model.save_pretrained(UpperCAmelCase					)


if __name__ == "__main__":
			lowerCAmelCase_	:		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.""")
			lowerCAmelCase_	:		List[str]			=     parser.parse_args()

			convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)



'''simple docstring'''





from typing import List, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

lowerCAmelCase_	:		Tuple			=     {
    """huggingface/time-series-transformer-tourism-monthly""": (
        """https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"""
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}




class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE	):







				'''simple docstring'''
				UpperCAmelCase__							=							'''time_series_transformer'''
				UpperCAmelCase__							=							{
				    '''hidden_size''': '''d_model''',
				    '''num_attention_heads''': '''encoder_attention_heads''',
				    '''num_hidden_layers''': '''encoder_layers''',
				}





				def __init__(						self					:						int       ,		lowercase__					:						Optional[int] = None       ,		lowercase__					:						Optional[int] = None       ,		lowercase__					:						str = "student_t"       ,		lowercase__					:						str = "nll"       ,		lowercase__					:						int = 1       ,		lowercase__					:						List[int] = [1, 2, 3, 4, 5, 6, 7]       ,		lowercase__					:						Optional[Union[str, bool]] = "mean"       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						int = 0       ,		lowercase__					:						Optional[List[int]] = None       ,		lowercase__					:						Optional[List[int]] = None       ,		lowercase__					:						int = 32       ,		lowercase__					:						int = 32       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						int = 2       ,		lowercase__					:						bool = True       ,		lowercase__					:						str = "gelu"       ,		lowercase__					:						int = 64       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						float = 0.1       ,		lowercase__					:						int = 100       ,		lowercase__					:						float = 0.0_2       ,		lowercase__					:						Optional[int]=True       ,		**lowercase__					:						Optional[Any]       ,		)							->Optional[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			int   			=						prediction_length
										_UpperCamelCase			:			Optional[Any]   			=						context_length or prediction_length
										_UpperCamelCase			:			List[str]   			=						distribution_output
										_UpperCamelCase			:			Optional[Any]   			=						loss
										_UpperCamelCase			:			Tuple   			=						input_size
										_UpperCamelCase			:			Optional[int]   			=						num_time_features
										_UpperCamelCase			:			Union[str, Any]   			=						lags_sequence
										_UpperCamelCase			:			int   			=						scaling
										_UpperCamelCase			:			Dict   			=						num_dynamic_real_features
										_UpperCamelCase			:			str   			=						num_static_real_features
										_UpperCamelCase			:			Any   			=						num_static_categorical_features
										if cardinality and num_static_categorical_features > 0:
																if len(lowercase__					) != num_static_categorical_features:
																						raise ValueError(
																						    "The cardinality should be a list of the same length as `num_static_categorical_features`"					)
																_UpperCamelCase			:			Dict   			=						cardinality
										else:
																_UpperCamelCase			:			str   			=						[0]
										if embedding_dimension and num_static_categorical_features > 0:
																if len(lowercase__					) != num_static_categorical_features:
																						raise ValueError(
																						    "The embedding dimension should be a list of the same length as `num_static_categorical_features`"					)
																_UpperCamelCase			:			int   			=						embedding_dimension
										else:
																_UpperCamelCase			:			Optional[Any]   			=						[min(50       ,		(cat + 1) // 2					) for cat in self.cardinality]
										_UpperCamelCase			:			Union[str, Any]   			=						num_parallel_samples

										# Transformer architecture configuration
										_UpperCamelCase			:			int   			=						input_size * len(lowercase__					) + self._number_of_features
										_UpperCamelCase			:			Optional[Any]   			=						d_model
										_UpperCamelCase			:			str   			=						encoder_attention_heads
										_UpperCamelCase			:			List[Any]   			=						decoder_attention_heads
										_UpperCamelCase			:			str   			=						encoder_ffn_dim
										_UpperCamelCase			:			List[str]   			=						decoder_ffn_dim
										_UpperCamelCase			:			Union[str, Any]   			=						encoder_layers
										_UpperCamelCase			:			Optional[int]   			=						decoder_layers

										_UpperCamelCase			:			int   			=						dropout
										_UpperCamelCase			:			Optional[int]   			=						attention_dropout
										_UpperCamelCase			:			int   			=						activation_dropout
										_UpperCamelCase			:			List[Any]   			=						encoder_layerdrop
										_UpperCamelCase			:			int   			=						decoder_layerdrop

										_UpperCamelCase			:			List[Any]   			=						activation_function
										_UpperCamelCase			:			Any   			=						init_std

										_UpperCamelCase			:			Optional[Any]   			=						use_cache

										super().__init__(is_encoder_decoder=lowercase__       ,		**lowercase__					)







				@property
				def  snake_case__				(						self					:						Union[str, Any]					)							->int:


										'''simple docstring'''




										return (
										    sum(self.embedding_dimension					)
										    + self.num_dynamic_real_features
										    + self.num_time_features
										    + self.num_static_real_features
										    + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
										)


