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


SCREAMING_SNAKE_CASE_    					=						{
    '''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__						(					lowerCAmelCase				:					Tuple ) ->      Optional[int]:






    """simple docstring"""




    UpperCAmelCase									=					['layers', 'blocks']
    for k in ignore_keys:
        state_dict.pop(lowerCAmelCase ,				lowerCAmelCase )


SCREAMING_SNAKE_CASE_    					=						{
    '''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__						(					lowerCAmelCase				:					Tuple ) ->      Any:






    """simple docstring"""




    UpperCAmelCase									=					list(s_dict.keys() )
    for key in keys:
        UpperCAmelCase									=					key
        for k, v in WHISPER_MAPPING.items():
            if k in key:
                UpperCAmelCase									=					new_key.replace(lowerCAmelCase ,				lowerCAmelCase )

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

        UpperCAmelCase									=					s_dict.pop(lowerCAmelCase )
    return s_dict




def     lowercase__						(					lowerCAmelCase				:					str ) ->      List[Any]:






    """simple docstring"""




    UpperCAmelCase			,							UpperCAmelCase									=					emb.weight.shape
    UpperCAmelCase									=					nn.Linear(lowerCAmelCase ,				lowerCAmelCase ,				bias=lowerCAmelCase )
    UpperCAmelCase									=					emb.weight.data
    return lin_layer




def     lowercase__						(					lowerCAmelCase				:					str ,				lowerCAmelCase				:					str ) ->      bytes:






    """simple docstring"""




    os.makedirs(lowerCAmelCase ,				exist_ok=lowerCAmelCase )
    UpperCAmelCase									=					os.path.basename(lowerCAmelCase )

    UpperCAmelCase									=					url.split('/' )[-2]
    UpperCAmelCase									=					os.path.join(lowerCAmelCase ,				lowerCAmelCase )

    if os.path.exists(lowerCAmelCase ) and not os.path.isfile(lowerCAmelCase ):
        raise RuntimeError(F"{download_target} exists and is not a regular file" )

    if os.path.isfile(lowerCAmelCase ):
        UpperCAmelCase									=					open(lowerCAmelCase ,				'rb' ).read()
        if hashlib.shaaaa(lowerCAmelCase ).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(lowerCAmelCase ) as source, open(lowerCAmelCase ,				'wb' ) as output:
        with tqdm(
            total=int(source.info().get('Content-Length' ) ) ,				ncols=80 ,				unit='iB' ,				unit_scale=lowerCAmelCase ,				unit_divisor=1_024 ) as loop:
            while True:
                UpperCAmelCase									=					source.read(8_192 )
                if not buffer:
                    break

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

    UpperCAmelCase									=					open(lowerCAmelCase ,				'rb' ).read()
    if hashlib.shaaaa(lowerCAmelCase ).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__						(					lowerCAmelCase				:					List[str] ,				lowerCAmelCase				:					Dict ) ->      List[str]:






    """simple docstring"""




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

    UpperCAmelCase									=					WhisperConfig(
        vocab_size=dimensions['n_vocab'] ,				encoder_ffn_dim=lowerCAmelCase ,				decoder_ffn_dim=lowerCAmelCase ,				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									=					WhisperForConditionalGeneration(lowerCAmelCase )
    UpperCAmelCase			,							UpperCAmelCase									=					model.model.load_state_dict(lowerCAmelCase ,				strict=lowerCAmelCase )
    if len(lowerCAmelCase ) > 0 and not set(lowerCAmelCase ) <= {
        "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									=					make_linear_from_emb(model.model.decoder.embed_tokens )
    else:
        UpperCAmelCase									=					proj_out_weights

    model.save_pretrained(lowerCAmelCase )


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

   convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)





"""simple docstring"""

import json
import os
from typing import Optional, Tuple

import regex as re

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


SCREAMING_SNAKE_CASE_    					=						logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_    					=						{
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
}

SCREAMING_SNAKE_CASE_    					=						{
    '''vocab_file''': {'''ctrl''': '''https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-vocab.json'''},
    '''merges_file''': {'''ctrl''': '''https://raw.githubusercontent.com/salesforce/ctrl/master/ctrl-merges.txt'''},
}

SCREAMING_SNAKE_CASE_    					=						{
    '''ctrl''': 256,
}

SCREAMING_SNAKE_CASE_    					=						{
    '''Pregnancy''': 16_8629,
    '''Christianity''': 7675,
    '''Explain''': 10_6423,
    '''Fitness''': 6_3440,
    '''Saving''': 6_3163,
    '''Ask''': 2_7171,
    '''Ass''': 9_5985,
    '''Joke''': 16_3509,
    '''Questions''': 4_5622,
    '''Thoughts''': 4_9605,
    '''Retail''': 5_2342,
    '''Feminism''': 16_4338,
    '''Writing''': 1_1992,
    '''Atheism''': 19_2263,
    '''Netflix''': 4_8616,
    '''Computing''': 3_9639,
    '''Opinion''': 4_3213,
    '''Alone''': 4_4967,
    '''Funny''': 5_8917,
    '''Gaming''': 4_0358,
    '''Human''': 4088,
    '''India''': 1331,
    '''Joker''': 7_7138,
    '''Diet''': 3_6206,
    '''Legal''': 1_1859,
    '''Norman''': 4939,
    '''Tip''': 7_2689,
    '''Weight''': 5_2343,
    '''Movies''': 4_6273,
    '''Running''': 2_3425,
    '''Science''': 2090,
    '''Horror''': 3_7793,
    '''Confession''': 6_0572,
    '''Finance''': 1_2250,
    '''Politics''': 1_6360,
    '''Scary''': 19_1985,
    '''Support''': 1_2654,
    '''Technologies''': 3_2516,
    '''Teenage''': 6_6160,
    '''Event''': 3_2769,
    '''Learned''': 6_7460,
    '''Notion''': 18_2770,
    '''Wikipedia''': 3_7583,
    '''Books''': 6665,
    '''Extract''': 7_6050,
    '''Confessions''': 10_2701,
    '''Conspiracy''': 7_5932,
    '''Links''': 6_3674,
    '''Narcissus''': 15_0425,
    '''Relationship''': 5_4766,
    '''Relationships''': 13_4796,
    '''Reviews''': 4_1671,
    '''News''': 4256,
    '''Translation''': 2_6820,
    '''multilingual''': 12_8406,
}




def     lowercase__						(					lowerCAmelCase				:					int ) ->      int:






    """simple docstring"""




    UpperCAmelCase									=					set()
    UpperCAmelCase									=					word[0]
    for char in word[1:]:
        pairs.add((prev_char, char) )
        UpperCAmelCase									=					char

    UpperCAmelCase									=					set(lowerCAmelCase )
    return pairs







class      _UpperCAmelCase							(		SCREAMING_SNAKE_CASE_							):
      __SCREAMING_SNAKE_CASE   :							Optional[Any]							      =	VOCAB_FILES_NAMES
      __SCREAMING_SNAKE_CASE   :							List[str]							      =	PRETRAINED_VOCAB_FILES_MAP
      __SCREAMING_SNAKE_CASE   :							List[Any]							      =	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      __SCREAMING_SNAKE_CASE   :							Union[str, Any]							      =	CONTROL_CODES




      def __init__(							self   ,       lowercase_   ,       lowercase_   ,       lowercase_="<unk>"   ,       **lowercase_							)						->					Optional[int]:
          super().__init__(unk_token=lowercase_   ,       **lowercase_							)

          with open(lowercase_   ,       encoding='utf-8'							) as vocab_handle:
              UpperCAmelCase									=					json.load(lowercase_							)
          UpperCAmelCase									=					{v: k for k, v in self.encoder.items()}
          with open(lowercase_   ,       encoding='utf-8'							) as merges_handle:
              UpperCAmelCase									=					merges_handle.read().split('\n'							)[1:-1]
          UpperCAmelCase									=					[tuple(merge.split()							) for merge in merges]
          UpperCAmelCase									=					dict(zip(lowercase_   ,       range(len(lowercase_							)							)							)							)
          UpperCAmelCase									=					{}




      @property
      def     a_      (							self							)						->					Union[str, Any]:
          return len(self.encoder							)




      def     a_      (							self							)						->					Optional[Any]:
          return dict(self.encoder   ,       **self.added_tokens_encoder							)




      def     a_      (							self   ,       lowercase_							)						->					int:
          if token in self.cache:
              return self.cache[token]
          UpperCAmelCase									=					tuple(lowercase_							)
          UpperCAmelCase									=					tuple(list(word[:-1]							) + [word[-1] + '</w>']							)
          UpperCAmelCase									=					get_pairs(lowercase_							)

          if not pairs:
              return token

          while True:
              UpperCAmelCase									=					min(lowercase_   ,       key=lambda lowercase_							: self.bpe_ranks.get(lowercase_   ,       float('inf'							)							)							)
              if bigram not in self.bpe_ranks:
                  break
              UpperCAmelCase			,							UpperCAmelCase									=					bigram
              UpperCAmelCase									=					[]
              UpperCAmelCase									=					0
              while i < len(lowercase_							):
                  try:
                      UpperCAmelCase									=					word.index(lowercase_   ,       lowercase_							)
                  except ValueError:
                      new_word.extend(word[i:]							)
                      break
                  else:
                      new_word.extend(word[i:j]							)
                      UpperCAmelCase									=					j

                  if word[i] == first and i < len(lowercase_							) - 1 and word[i + 1] == second:
                      new_word.append(first + second							)
                      i += 2
                  else:
                      new_word.append(word[i]							)
                      i += 1
              UpperCAmelCase									=					tuple(lowercase_							)
              UpperCAmelCase									=					new_word
              if len(lowercase_							) == 1:
                  break
              else:
                  UpperCAmelCase									=					get_pairs(lowercase_							)
          UpperCAmelCase									=					'@@ '.join(lowercase_							)
          UpperCAmelCase									=					word[:-4]
          UpperCAmelCase									=					word
          return word




      def     a_      (							self   ,       lowercase_							)						->					List[Any]:
          UpperCAmelCase									=					[]

          UpperCAmelCase									=					re.findall(R'\S+\n?'   ,       lowercase_							)

          for token in words:
              split_tokens.extend(list(self.bpe(lowercase_							).split(' '							)							)							)
          return split_tokens




      def     a_      (							self   ,       lowercase_							)						->					Tuple:
          return self.encoder.get(lowercase_   ,       self.encoder.get(self.unk_token							)							)




      def     a_      (							self   ,       lowercase_							)						->					Optional[Any]:
          return self.decoder.get(lowercase_   ,       self.unk_token							)




      def     a_      (							self   ,       lowercase_							)						->					str:
          UpperCAmelCase									=					' '.join(lowercase_							).replace('@@ '   ,       ''							).strip()
          return out_string




      def     a_      (							self   ,       lowercase_   ,       lowercase_ = None							)						->					Tuple[str]:
          if not os.path.isdir(lowercase_							):
              logger.error(F"Vocabulary path ({save_directory}) should be a directory"							)
              return
          UpperCAmelCase									=					os.path.join(
              lowercase_   ,       (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']							)
          UpperCAmelCase									=					os.path.join(
              lowercase_   ,       (filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['merges_file']							)

          with open(lowercase_   ,       'w'   ,       encoding='utf-8'							) as f:
              f.write(json.dumps(self.encoder   ,       indent=2   ,       sort_keys=lowercase_   ,       ensure_ascii=lowercase_							) + '\n'							)

          UpperCAmelCase									=					0
          with open(lowercase_   ,       'w'   ,       encoding='utf-8'							) as writer:
              writer.write('#version: 0.2\n'							)
              for bpe_tokens, token_index in sorted(self.bpe_ranks.items()   ,       key=lambda lowercase_							: kv[1]							):
                  if index != token_index:
                      logger.warning(
                          F"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
                          ' Please check that the tokenizer is not corrupted!'							)
                      UpperCAmelCase									=					token_index
                  writer.write(' '.join(lowercase_							) + '\n'							)
                  index += 1

          return vocab_file, merge_file

      # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
      #     filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
      #     tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
      #     tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
      #     return ''.join(tokens_generated_so_far)



