



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[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      __lowercase       (			a__       )    ->			Dict:
							__SCREAMING_SNAKE_CASE									= ['layers', 'blocks']
							for k in ignore_keys:
														state_dict.pop(a__							,       a__       )


lowerCAmelCase__		:					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      __lowercase       (			a__       )    ->			int:
							__SCREAMING_SNAKE_CASE									= list(s_dict.keys()       )
							for key in keys:
														__SCREAMING_SNAKE_CASE									= key
														for k, v in WHISPER_MAPPING.items():
																					if k in key:
																												__SCREAMING_SNAKE_CASE									= new_key.replace(a__							,       a__       )

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

														__SCREAMING_SNAKE_CASE									= s_dict.pop(a__       )
							return s_dict






def      __lowercase       (			a__       )    ->			Optional[int]:
							__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= emb.weight.shape
							__SCREAMING_SNAKE_CASE									= nn.Linear(a__							,       a__							,       bias=a__       )
							__SCREAMING_SNAKE_CASE									= emb.weight.data
							return lin_layer






def      __lowercase       (			a__							,       a__       )    ->			bytes:
							os.makedirs(a__							,       exist_ok=a__       )
							__SCREAMING_SNAKE_CASE									= os.path.basename(a__       )

							__SCREAMING_SNAKE_CASE									= url.split('/'       )[-2]
							__SCREAMING_SNAKE_CASE									= 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__       ):
														__SCREAMING_SNAKE_CASE									= 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=10_24       ) as loop:
																					while True:
																												__SCREAMING_SNAKE_CASE									= source.read(81_92       )
																												if not buffer:
																																			break

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

							__SCREAMING_SNAKE_CASE									= 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      __lowercase       (			a__							,       a__       )    ->			Union[str, Any]:
							if ".pt" not in checkpoint_path:
														__SCREAMING_SNAKE_CASE									= _download(_MODELS[checkpoint_path]       )
							else:
														__SCREAMING_SNAKE_CASE									= torch.load(a__							,       map_location='cpu'       )
							__SCREAMING_SNAKE_CASE									= original_checkpoint['dims']
							__SCREAMING_SNAKE_CASE									= original_checkpoint['model_state_dict']
							__SCREAMING_SNAKE_CASE									= state_dict['decoder.token_embedding.weight']
							remove_ignore_keys_(a__       )
							rename_keys(a__       )
							__SCREAMING_SNAKE_CASE									= True
							__SCREAMING_SNAKE_CASE									= state_dict['decoder.layers.0.fc1.weight'].shape[0]

							__SCREAMING_SNAKE_CASE									= 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']							,       )

							__SCREAMING_SNAKE_CASE									= WhisperForConditionalGeneration(a__       )
							__SCREAMING_SNAKE_CASE							,	__SCREAMING_SNAKE_CASE									= 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:
														__SCREAMING_SNAKE_CASE									= make_linear_from_emb(model.model.decoder.embed_tokens       )
							else:
														__SCREAMING_SNAKE_CASE									= proj_out_weights

							model.save_pretrained(a__       )


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

			convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)











import warnings
from contextlib import contextmanager

from ...processing_utils import ProcessorMixin
from .feature_extraction_wavaveca import WavaVecaFeatureExtractor
from .tokenization_wavaveca import WavaVecaCTCTokenizer

class        UpperCAmelCase_  (       UpperCamelCase_	):

						'''simple docstring'''




						UpperCamelCase__      :			Union[str, Any]      =  '''Wav2Vec2FeatureExtractor'''
						UpperCamelCase__      :			Union[str, Any]      =  '''AutoTokenizer'''






						def __init__(   self				,							_A				,							_A		):







													'''simple docstring'''

													super().__init__(_A				,							_A		)
													__SCREAMING_SNAKE_CASE									= self.feature_extractor
													__SCREAMING_SNAKE_CASE									= False






						@classmethod
						def 	_A     (   cls				,							_A				,							**_A		):







													'''simple docstring'''

													try:
																				return super().from_pretrained(_A				,							**_A		)
													except OSError:
																				warnings.warn(
																				    f"""Loading a tokenizer inside {cls.__name__} from a config that does not"""
																				    ' include a `tokenizer_class` attribute is deprecated and will be '
																				    'removed in v5. Please add `\'tokenizer_class\': \'Wav2Vec2CTCTokenizer\'`'
																				    ' attribute to either your `config.json` or `tokenizer_config.json` '
																				    'file to suppress this warning: '				,							_A				,							)

																				__SCREAMING_SNAKE_CASE									= WavaVecaFeatureExtractor.from_pretrained(_A				,							**_A		)
																				__SCREAMING_SNAKE_CASE									= WavaVecaCTCTokenizer.from_pretrained(_A				,							**_A		)

																				return cls(feature_extractor=_A				,							tokenizer=_A		)






						def __call__(   self				,							*_A				,							**_A		):







													'''simple docstring'''

													if self._in_target_context_manager:
																				return self.current_processor(*_A				,							**_A		)

													if "raw_speech" in kwargs:
																				warnings.warn('Using `raw_speech` as a keyword argument is deprecated. Use `audio` instead.'		)
																				__SCREAMING_SNAKE_CASE									= kwargs.pop('raw_speech'		)
													else:
																				__SCREAMING_SNAKE_CASE									= kwargs.pop('audio'				,							_A		)
													__SCREAMING_SNAKE_CASE									= kwargs.pop('sampling_rate'				,							_A		)
													__SCREAMING_SNAKE_CASE									= kwargs.pop('text'				,							_A		)
													if len(_A		) > 0:
																				__SCREAMING_SNAKE_CASE									= args[0]
																				__SCREAMING_SNAKE_CASE									= args[1:]

													if audio is None and text is None:
																				raise ValueError('You need to specify either an `audio` or `text` input to process.'		)

													if audio is not None:
																				__SCREAMING_SNAKE_CASE									= self.feature_extractor(_A				,							*_A				,							sampling_rate=_A				,							**_A		)
													if text is not None:
																				__SCREAMING_SNAKE_CASE									= self.tokenizer(_A				,							**_A		)

													if text is None:
																				return inputs
													elif audio is None:
																				return encodings
													else:
																				__SCREAMING_SNAKE_CASE									= encodings['input_ids']
																				return inputs






						def 	_A     (   self				,							*_A				,							**_A		):







													'''simple docstring'''

													if self._in_target_context_manager:
																				return self.current_processor.pad(*_A				,							**_A		)

													__SCREAMING_SNAKE_CASE									= kwargs.pop('input_features'				,							_A		)
													__SCREAMING_SNAKE_CASE									= kwargs.pop('labels'				,							_A		)
													if len(_A		) > 0:
																				__SCREAMING_SNAKE_CASE									= args[0]
																				__SCREAMING_SNAKE_CASE									= args[1:]

													if input_features is not None:
																				__SCREAMING_SNAKE_CASE									= self.feature_extractor.pad(_A				,							*_A				,							**_A		)
													if labels is not None:
																				__SCREAMING_SNAKE_CASE									= self.tokenizer.pad(_A				,							**_A		)

													if labels is None:
																				return input_features
													elif input_features is None:
																				return labels
													else:
																				__SCREAMING_SNAKE_CASE									= labels['input_ids']
																				return input_features






						def 	_A     (   self				,							*_A				,							**_A		):







													'''simple docstring'''

													return self.tokenizer.batch_decode(*_A				,							**_A		)






						def 	_A     (   self				,							*_A				,							**_A		):







													'''simple docstring'''

													return self.tokenizer.decode(*_A				,							**_A		)






						@contextmanager
						def 	_A     (   self		):







													'''simple docstring'''

													warnings.warn(
													    '`as_target_processor` is deprecated and will be removed in v5 of Transformers. You can process your '
													    'labels by using the argument `text` of the regular `__call__` method (either in the same call as '
													    'your audio inputs, or in a separate call.'		)
													__SCREAMING_SNAKE_CASE									= True
													__SCREAMING_SNAKE_CASE									= self.tokenizer
													yield
													__SCREAMING_SNAKE_CASE									= self.feature_extractor
													__SCREAMING_SNAKE_CASE									= False






