





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__       :			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       _A  (      lowerCamelCase   ):
   a__	:						Optional[int]      		=  ["layers", "blocks"]
   for k in ignore_keys:
      state_dict.pop(lowerCamelCase							, lowerCamelCase   )


SCREAMING_SNAKE_CASE__       :			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  (      lowerCamelCase   ):
   a__	:						Tuple      		=  list(s_dict.keys()   )
   for key in keys:
      a__	:						Optional[Any]      		=  key
      for k, v in WHISPER_MAPPING.items():
         if k in key:
            a__	:						Optional[int]      		=  new_key.replace(lowerCamelCase							, lowerCamelCase   )

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

      a__	:						Dict      		=  s_dict.pop(lowerCamelCase   )
   return s_dict






def       _A  (      lowerCamelCase   ):
   a__		,     a__	:						Any      		=  emb.weight.shape
   a__	:						Optional[Any]      		=  nn.Linear(lowerCamelCase							, lowerCamelCase							, bias=lowerCamelCase   )
   a__	:						Optional[Any]      		=  emb.weight.data
   return lin_layer






def       _A  (      lowerCamelCase							, lowerCamelCase   ):
   os.makedirs(lowerCamelCase							, exist_ok=lowerCamelCase   )
   a__	:						Optional[Any]      		=  os.path.basename(lowerCamelCase   )

   a__	:						List[Any]      		=  url.split("/"   )[-2]
   a__	:						Tuple      		=  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   ):
      a__	:						Any      		=  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=1024   ) as loop:
         while True:
            a__	:						Optional[Any]      		=  source.read(8192   )
            if not buffer:
               break

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

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

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

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

   model.save_pretrained(lowerCamelCase   )


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

  convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)








import logging
import math
from functools import partial
from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union

import torch

from .tensor_utils import tensor_tree_map, tree_map






def       _A  (      lowerCamelCase   ):
   a__	:						List[str]      		=  []
   if isinstance(lowerCamelCase							, lowerCamelCase   ):
      for v in tree.values():
         shapes.extend(_fetch_dims(lowerCamelCase   )   )
   elif isinstance(lowerCamelCase							, (list, tuple)   ):
      for t in tree:
         shapes.extend(_fetch_dims(lowerCamelCase   )   )
   elif isinstance(lowerCamelCase							, torch.Tensor   ):
      shapes.append(tree.shape   )
   else:
      raise ValueError("Not supported"   )

   return shapes






@torch.jit.ignore
def       _A  (      lowerCamelCase							, lowerCamelCase   ):
   a__	:						List[str]      		=  []
   for d in reversed(lowerCamelCase   ):
      idx.append(flat_idx % d   )
      a__	:						Union[str, Any]      		=  flat_idx // d

   return tuple(reversed(lowerCamelCase   )   )






@torch.jit.ignore
def       _A  (      lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase = None							, lowerCamelCase = None							, ):

   # start_edges and end_edges both indicate whether, starting from any given
   # dimension, the start/end index is at the top/bottom edge of the
   # corresponding tensor, modeled as a tree
   def reduce_edge_list(lowerCamelCase   ) -> None:
      a__	:						int      		=  True
      for i in range(len(lowerCamelCase   )   ):
         a__	:						Optional[Any]      		=  -1 * (i + 1)
         l[reversed_idx] &= tally
         a__	:						Tuple      		=  l[reversed_idx]

   if start_edges is None:
      a__	:						Optional[int]      		=  [s == 0 for s in start]
      reduce_edge_list(lowerCamelCase   )
   if end_edges is None:
      a__	:						Union[str, Any]      		=  [e == (d - 1) for e, d in zip(lowerCamelCase							, lowerCamelCase   )]
      reduce_edge_list(lowerCamelCase   )

   # Base cases. Either start/end are empty and we're done, or the final,
   # one-dimensional tensor can be simply sliced
   if len(lowerCamelCase   ) == 0:
      return [()]
   elif len(lowerCamelCase   ) == 1:
      return [(slice(start[0]							, end[0] + 1   ),)]

   a__	:						List[Tuple[slice, ...]]      		=  []
   a__	:						List[slice]      		=  []

   # Dimensions common to start and end can be selected directly
   for s, e in zip(lowerCamelCase							, lowerCamelCase   ):
      if s == e:
         path_list.append(slice(lowerCamelCase							, s + 1   )   )
      else:
         break

   a__	:						Tuple[slice, ...]      		=  tuple(lowerCamelCase   )
   a__	:						Optional[Any]      		=  len(lowerCamelCase   )

   # start == end, and we're done
   if divergence_idx == len(lowerCamelCase   ):
      return [path]

   def upper() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      a__	:						Optional[Any]      		=  start[divergence_idx]
      return tuple(
          path + (slice(lowerCamelCase							, sdi + 1   ),) + s
          for s in _get_minimal_slice_set(
              start[divergence_idx + 1 :]							, [d - 1 for d in dims[divergence_idx + 1 :]]							, dims[divergence_idx + 1 :]							, start_edges=start_edges[divergence_idx + 1 :]							, end_edges=[True for _ in end_edges[divergence_idx + 1 :]]							, )   )

   def lower() -> Tuple[Tuple[slice, ...], ...]:
      assert start_edges is not None
      assert end_edges is not None

      a__	:						List[str]      		=  end[divergence_idx]
      return tuple(
          path + (slice(lowerCamelCase							, edi + 1   ),) + s
          for s in _get_minimal_slice_set(
              [0 for _ in start[divergence_idx + 1 :]]							, end[divergence_idx + 1 :]							, dims[divergence_idx + 1 :]							, start_edges=[True for _ in start_edges[divergence_idx + 1 :]]							, end_edges=end_edges[divergence_idx + 1 :]							, )   )

   # If both start and end are at the edges of the subtree rooted at
   # divergence_idx, we can just select the whole subtree at once
   if start_edges[divergence_idx] and end_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]							, end[divergence_idx] + 1   ),)   )
   # If just start is at the edge, we can grab almost all of the subtree,
   # treating only the ragged bottom edge as an edge case
   elif start_edges[divergence_idx]:
      slices.append(path + (slice(start[divergence_idx]							, end[divergence_idx]   ),)   )
      slices.extend(lower()   )
   # Analogous to the previous case, but the top is ragged this time
   elif end_edges[divergence_idx]:
      slices.extend(upper()   )
      slices.append(path + (slice(start[divergence_idx] + 1							, end[divergence_idx] + 1   ),)   )
   # If both sides of the range are ragged, we need to handle both sides
   # separately. If there's contiguous meat in between them, we can index it
   # in one big chunk
   else:
      slices.extend(upper()   )
      a__	:						Optional[int]      		=  end[divergence_idx] - start[divergence_idx]
      if middle_ground > 1:
         slices.append(path + (slice(start[divergence_idx] + 1							, end[divergence_idx]   ),)   )
      slices.extend(lower()   )

   return slices






@torch.jit.ignore
def       _A  (      lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase   ):
   a__	:						Optional[int]      		=  t.shape[:no_batch_dims]
   a__	:						List[str]      		=  list(_flat_idx_to_idx(lowerCamelCase							, lowerCamelCase   )   )
   # _get_minimal_slice_set is inclusive
   a__	:						Dict      		=  list(_flat_idx_to_idx(flat_end - 1							, lowerCamelCase   )   )

   # Get an ordered list of slices to perform
   a__	:						str      		=  _get_minimal_slice_set(
       lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, )

   a__	:						Any      		=  [t[s] for s in slices]

   return torch.cat([s.view((-1,) + t.shape[no_batch_dims:]   ) for s in sliced_tensors]   )






def       _A  (      lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase							, lowerCamelCase = False							, lowerCamelCase = None							, lowerCamelCase = False							, ):
   if not (len(lowerCamelCase   ) > 0):
      raise ValueError("Must provide at least one input"   )

   a__	:						str      		=  [shape[:no_batch_dims] for shape in _fetch_dims(lowerCamelCase   )]
   a__	:						Dict      		=  tuple([max(lowerCamelCase   ) for s in zip(*lowerCamelCase   )]   )

   def _prep_inputs(lowerCamelCase   ) -> torch.Tensor:
      if not low_mem:
         if not sum(t.shape[:no_batch_dims]   ) == no_batch_dims:
            a__	:						Any      		=  t.expand(orig_batch_dims + t.shape[no_batch_dims:]   )
         a__	:						Optional[Any]      		=  t.reshape(-1							, *t.shape[no_batch_dims:]   )
      else:
         a__	:						Dict      		=  t.expand(orig_batch_dims + t.shape[no_batch_dims:]   )
      return t

   a__	:						Dict[str, Any]      		=  tensor_tree_map(_prep_inputs							, lowerCamelCase   )
   a__	:						str      		=  None
   if _out is not None:
      a__	:						Optional[int]      		=  tensor_tree_map(lambda lowerCamelCase   : t.view([-1] + list(t.shape[no_batch_dims:]   )   )							, _out   )

   a__	:						Optional[Any]      		=  1
   for d in orig_batch_dims:
      flat_batch_dim *= d

   a__	:						Tuple      		=  flat_batch_dim // chunk_size + (flat_batch_dim % chunk_size != 0)

   def _select_chunk(lowerCamelCase   ) -> torch.Tensor:
      return t[i : i + chunk_size] if t.shape[0] != 1 else t

   a__	:						str      		=  0
   a__	:						Any      		=  prepped_outputs
   for _ in range(lowerCamelCase   ):
      # Chunk the input
      if not low_mem:
         a__	:						str      		=  _select_chunk
      else:
         a__	:						Tuple      		=  partial(
             _chunk_slice							, flat_start=lowerCamelCase							, flat_end=min(lowerCamelCase							, i + chunk_size   )							, no_batch_dims=len(lowerCamelCase   )							, )

      a__	:						Dict[str, Any]      		=  tensor_tree_map(lowerCamelCase							, lowerCamelCase   )

      # Run the layer on the chunk
      a__	:						Any      		=  layer(**lowerCamelCase   )

      # Allocate space for the output
      if out is None:
         a__	:						Optional[Any]      		=  tensor_tree_map(lambda lowerCamelCase   : t.new_zeros((flat_batch_dim,) + t.shape[1:]   )							, lowerCamelCase   )

      # Put the chunk in its pre-allocated space
      if isinstance(lowerCamelCase							, lowerCamelCase   ):

         def assign(lowerCamelCase							, lowerCamelCase   ) -> None:
            for k, v in da.items():
               if isinstance(lowerCamelCase							, lowerCamelCase   ):
                  assign(lowerCamelCase							, da[k]   )
               else:
                  if _add_into_out:
                     v[i : i + chunk_size] += da[k]
                  else:
                     a__	:						Dict      		=  da[k]

         assign(lowerCamelCase							, lowerCamelCase   )
      elif isinstance(lowerCamelCase							, lowerCamelCase   ):
         for xa, xa in zip(lowerCamelCase							, lowerCamelCase   ):
            if _add_into_out:
               xa[i : i + chunk_size] += xa
            else:
               a__	:						Dict      		=  xa
      elif isinstance(lowerCamelCase							, torch.Tensor   ):
         if _add_into_out:
            out[i : i + chunk_size] += output_chunk
         else:
            a__	:						Dict      		=  output_chunk
      else:
         raise ValueError("Not supported"   )

      i += chunk_size

   a__	:						Any      		=  tensor_tree_map(lambda lowerCamelCase   : t.view(orig_batch_dims + t.shape[1:]   )							, lowerCamelCase   )

   return out







class   __lowerCAmelCase			:





   def __init__( self					,			snake_case = 512					,			)				->  List[str]:
      """simple docstring"""





      a__	:						int      		=  max_chunk_size
      a__	:						Optional[int]      		=  None
      a__	:						Optional[tuple]      		=  None





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case      )				->  int:
      """simple docstring"""





      logging.info("Tuning chunk size..."      )

      if min_chunk_size >= self.max_chunk_size:
         return min_chunk_size

      a__	:						List[int]      		=  [2**l for l in range(int(math.log(self.max_chunk_size					,			2      )      ) + 1      )]
      a__	:						List[str]      		=  [c for c in candidates if c > min_chunk_size]
      a__	:						Optional[int]      		=  [min_chunk_size] + candidates
      candidates[-1] += 4

      def test_chunk_size(snake_case      ) -> bool:
         try:
            with torch.no_grad():
               fn(*snake_case					,			chunk_size=snake_case      )
            return True
         except RuntimeError:
            return False

      a__	:						Union[str, Any]      		=  0
      a__	:						Dict      		=  len(snake_case      ) - 1
      while i > min_viable_chunk_size_index:
         a__	:						Any      		=  test_chunk_size(candidates[i]      )
         if not viable:
            a__	:						List[Any]      		=  (min_viable_chunk_size_index + i) // 2
         else:
            a__	:						Tuple      		=  i
            a__	:						Any      		=  (i + len(snake_case      ) - 1) // 2

      return candidates[min_viable_chunk_size_index]





   def 	_snake_case ( self					,			snake_case					,			snake_case      )				->  bool:
      """simple docstring"""





      a__	:						str      		=  True
      for aa, aa in zip(snake_case					,			snake_case      ):
         assert type(snake_case      ) == type(snake_case      )
         if isinstance(snake_case					,			(list, tuple)      ):
            consistent &= self._compare_arg_caches(snake_case					,			snake_case      )
         elif isinstance(snake_case					,			snake_case      ):
            a__	:						Union[str, Any]      		=  [v for _, v in sorted(aa.items()					,			key=lambda snake_case      : x[0]      )]
            a__	:						List[Any]      		=  [v for _, v in sorted(aa.items()					,			key=lambda snake_case      : x[0]      )]
            consistent &= self._compare_arg_caches(snake_case					,			snake_case      )
         else:
            consistent &= aa == aa

      return consistent





   def 	_snake_case ( self					,			snake_case					,			snake_case					,			snake_case					,			)				->  int:
      """simple docstring"""





      a__	:						List[Any]      		=  True
      a__	:						tuple      		=  tree_map(lambda snake_case      : a.shape if isinstance(snake_case					,			torch.Tensor      ) else a					,			snake_case					,			snake_case      )
      if self.cached_arg_data is not None:
         # If args have changed shape/value, we need to re-tune
         assert len(self.cached_arg_data      ) == len(snake_case      )
         a__	:						Union[str, Any]      		=  self._compare_arg_caches(self.cached_arg_data					,			snake_case      )
      else:
         # Otherwise, we can reuse the precomputed value
         a__	:						Optional[int]      		=  False

      if not consistent:
         a__	:						List[str]      		=  self._determine_favorable_chunk_size(
             snake_case					,			snake_case					,			snake_case					,			)
         a__	:						List[str]      		=  arg_data

      assert self.cached_chunk_size is not None

      return self.cached_chunk_size

