
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__     :		int        = {
    """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						(   lowercase_ )		->						Union[str, Any]:


  '''simple docstring'''






  lowercase__	: int     		=   ["""layers""", """blocks"""]
  for k in ignore_keys:
    state_dict.pop(lowercase_			,			lowercase_ )


lowerCamelCase__     :		Optional[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 		UpperCamelCase						(   lowercase_ )		->						int:


  '''simple docstring'''






  lowercase__	: str     		=   list(s_dict.keys() )
  for key in keys:
    lowercase__	: int     		=   key
    for k, v in WHISPER_MAPPING.items():
      if k in key:
        lowercase__	: Union[str, Any]     		=   new_key.replace(lowercase_			,			lowercase_ )

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

    lowercase__	: List[str]     		=   s_dict.pop(lowercase_ )
  return s_dict

def 		UpperCamelCase						(   lowercase_ )		->						List[str]:


  '''simple docstring'''






  lowercase__		,     lowercase__	: int     		=   emb.weight.shape
  lowercase__	: List[Any]     		=   nn.Linear(lowercase_			,			lowercase_			,			bias=lowercase_ )
  lowercase__	: Any     		=   emb.weight.data
  return lin_layer

def 		UpperCamelCase						(   lowercase_			,			lowercase_ )		->						bytes:


  '''simple docstring'''






  os.makedirs(lowercase_			,			exist_ok=lowercase_ )
  lowercase__	: List[str]     		=   os.path.basename(lowercase_ )

  lowercase__	: Union[str, Any]     		=   url.split("""/""" )[-2]
  lowercase__	: int     		=   os.path.join(lowercase_			,			lowercase_ )

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

  if os.path.isfile(lowercase_ ):
    lowercase__	: Tuple     		=   open(lowercase_			,			"""rb""" ).read()
    if hashlib.shaaaa(lowercase_ ).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(lowercase_ ) as source, open(lowercase_			,			"""wb""" ) as output:
    with tqdm(
        total=int(source.info().get("""Content-Length""" ) )			,			ncols=80			,			unit="""iB"""			,			unit_scale=lowercase_			,			unit_divisor=10_24 ) as loop:
      while True:
        lowercase__	: str     		=   source.read(81_92 )
        if not buffer:
          break

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

  lowercase__	: Tuple     		=   open(lowercase_			,			"""rb""" ).read()
  if hashlib.shaaaa(lowercase_ ).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						(   lowercase_			,			lowercase_ )		->						List[Any]:


  '''simple docstring'''






  if ".pt" not in checkpoint_path:
    lowercase__	: Dict     		=   _download(_MODELS[checkpoint_path] )
  else:
    lowercase__	: str     		=   torch.load(lowercase_			,			map_location="""cpu""" )
  lowercase__	: List[Any]     		=   original_checkpoint["""dims"""]
  lowercase__	: Optional[Any]     		=   original_checkpoint["""model_state_dict"""]
  lowercase__	: str     		=   state_dict["""decoder.token_embedding.weight"""]
  remove_ignore_keys_(lowercase_ )
  rename_keys(lowercase_ )
  lowercase__	: Any     		=   True
  lowercase__	: Any     		=   state_dict["""decoder.layers.0.fc1.weight"""].shape[0]

  lowercase__	: int     		=   WhisperConfig(
      vocab_size=dimensions["""n_vocab"""]			,			encoder_ffn_dim=lowercase_			,			decoder_ffn_dim=lowercase_			,			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"""]			,			)

  lowercase__	: Optional[Any]     		=   WhisperForConditionalGeneration(lowercase_ )
  lowercase__		,     lowercase__	: int     		=   model.model.load_state_dict(lowercase_			,			strict=lowercase_ )
  if len(lowercase_ ) > 0 and not set(lowercase_ ) <= {
      "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:
    lowercase__	: List[Any]     		=   make_linear_from_emb(model.model.decoder.embed_tokens )
  else:
    lowercase__	: Tuple     		=   proj_out_weights

  model.save_pretrained(lowercase_ )


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

     convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)











import unittest

from transformers import MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING, is_vision_available, pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
							from PIL import Image
else:


							class   __lowerCAmelCase  :
												"""simple docstring"""






												@staticmethod
												def      _a				(  *_snake_case       :							int					,       **_snake_case       :							List[str]       ):






															"""simple docstring"""


															pass


@is_pipeline_test
@require_vision
@require_torch
class   __lowerCAmelCase  (		unittest.TestCase   ):
					"""simple docstring"""






					A__					:		List[str]      							=     MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING



					def      _a				(  self       :							Any					,       _snake_case       :							Union[str, Any]					,       _snake_case       :							Tuple					,       _snake_case       :							Optional[Any]       ):






								"""simple docstring"""


								A__	  =     pipeline(
								    'zero-shot-object-detection'					,       model='hf-internal-testing/tiny-random-owlvit-object-detection'       )

								A__	  =     [
								    {
								        'image': './tests/fixtures/tests_samples/COCO/000000039769.png',
								        'candidate_labels': ['cat', 'remote', 'couch'],
								    }
								]
								return object_detector, examples



					def      _a				(  self       :							int					,       _snake_case       :							int					,       _snake_case       :							List[str]       ):






								"""simple docstring"""


								A__	  =     object_detector(examples[0]					,       threshold=0.0       )

								A__	  =     len(_snake_case       )
								self.assertGreater(_snake_case					,       0       )
								self.assertEqual(
								    _snake_case					,       [
								        {
								            'score': ANY(_snake_case       ),
								            'label': ANY(_snake_case       ),
								            'box': {'xmin': ANY(_snake_case       ), 'ymin': ANY(_snake_case       ), 'xmax': ANY(_snake_case       ), 'ymax': ANY(_snake_case       )},
								        }
								        for i in range(_snake_case       )
								    ]					,       )



					@require_tf
					@unittest.skip('Zero Shot Object Detection not implemented in TF'       )
					def      _a				(  self       :							List[str]       ):






								"""simple docstring"""


								pass



					@require_torch
					def      _a				(  self       :							Optional[int]       ):






								"""simple docstring"""


								A__	  =     pipeline(
								    'zero-shot-object-detection'					,       model='hf-internal-testing/tiny-random-owlvit-object-detection'       )

								A__	  =     object_detector(
								    './tests/fixtures/tests_samples/COCO/000000039769.png'					,       candidate_labels=['cat', 'remote', 'couch']					,       threshold=0.64					,       )

								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        {'score': 0.7235, 'label': 'cat', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								        {'score': 0.7218, 'label': 'remote', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								        {'score': 0.7184, 'label': 'couch', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								        {'score': 0.6748, 'label': 'remote', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								        {'score': 0.6656, 'label': 'cat', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								        {'score': 0.6614, 'label': 'couch', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								        {'score': 0.6456, 'label': 'remote', 'box': {'xmin': 4_94, 'ymin': 1_05, 'xmax': 5_21, 'ymax': 1_27}},
								        {'score': 0.642, 'label': 'remote', 'box': {'xmin': 67, 'ymin': 2_74, 'xmax': 93, 'ymax': 2_97}},
								        {'score': 0.6419, 'label': 'cat', 'box': {'xmin': 4_94, 'ymin': 1_05, 'xmax': 5_21, 'ymax': 1_27}},
								    ]					,       )

								A__	  =     object_detector(
								    [
								        {
								            'image': './tests/fixtures/tests_samples/COCO/000000039769.png',
								            'candidate_labels': ['cat', 'remote', 'couch'],
								        }
								    ]					,       threshold=0.64					,       )

								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        [
								            {'score': 0.7235, 'label': 'cat', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								            {'score': 0.7218, 'label': 'remote', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								            {'score': 0.7184, 'label': 'couch', 'box': {'xmin': 2_04, 'ymin': 1_67, 'xmax': 2_32, 'ymax': 1_90}},
								            {'score': 0.6748, 'label': 'remote', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								            {'score': 0.6656, 'label': 'cat', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								            {'score': 0.6614, 'label': 'couch', 'box': {'xmin': 5_71, 'ymin': 83, 'xmax': 5_98, 'ymax': 1_03}},
								            {'score': 0.6456, 'label': 'remote', 'box': {'xmin': 4_94, 'ymin': 1_05, 'xmax': 5_21, 'ymax': 1_27}},
								            {'score': 0.642, 'label': 'remote', 'box': {'xmin': 67, 'ymin': 2_74, 'xmax': 93, 'ymax': 2_97}},
								            {'score': 0.6419, 'label': 'cat', 'box': {'xmin': 4_94, 'ymin': 1_05, 'xmax': 5_21, 'ymax': 1_27}},
								        ]
								    ]					,       )



					@require_torch
					@slow
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								A__	  =     pipeline('zero-shot-object-detection'       )

								A__	  =     object_detector(
								    'http://images.cocodataset.org/val2017/000000039769.jpg'					,       candidate_labels=['cat', 'remote', 'couch']					,       )
								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        {'score': 0.2868, 'label': 'cat', 'box': {'xmin': 3_24, 'ymin': 20, 'xmax': 6_40, 'ymax': 3_73}},
								        {'score': 0.277, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 1_77, 'ymax': 1_15}},
								        {'score': 0.2537, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 3_15, 'ymax': 4_72}},
								        {'score': 0.1474, 'label': 'remote', 'box': {'xmin': 3_35, 'ymin': 74, 'xmax': 3_71, 'ymax': 1_87}},
								        {'score': 0.1208, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 6_42, 'ymax': 4_76}},
								    ]					,       )

								A__	  =     object_detector(
								    [
								        {
								            'image': 'http://images.cocodataset.org/val2017/000000039769.jpg',
								            'candidate_labels': ['cat', 'remote', 'couch'],
								        },
								        {
								            'image': 'http://images.cocodataset.org/val2017/000000039769.jpg',
								            'candidate_labels': ['cat', 'remote', 'couch'],
								        },
								    ]					,       )
								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        [
								            {'score': 0.2868, 'label': 'cat', 'box': {'xmin': 3_24, 'ymin': 20, 'xmax': 6_40, 'ymax': 3_73}},
								            {'score': 0.277, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 1_77, 'ymax': 1_15}},
								            {'score': 0.2537, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 3_15, 'ymax': 4_72}},
								            {'score': 0.1474, 'label': 'remote', 'box': {'xmin': 3_35, 'ymin': 74, 'xmax': 3_71, 'ymax': 1_87}},
								            {'score': 0.1208, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 6_42, 'ymax': 4_76}},
								        ],
								        [
								            {'score': 0.2868, 'label': 'cat', 'box': {'xmin': 3_24, 'ymin': 20, 'xmax': 6_40, 'ymax': 3_73}},
								            {'score': 0.277, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 1_77, 'ymax': 1_15}},
								            {'score': 0.2537, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 3_15, 'ymax': 4_72}},
								            {'score': 0.1474, 'label': 'remote', 'box': {'xmin': 3_35, 'ymin': 74, 'xmax': 3_71, 'ymax': 1_87}},
								            {'score': 0.1208, 'label': 'couch', 'box': {'xmin': 4, 'ymin': 0, 'xmax': 6_42, 'ymax': 4_76}},
								        ],
								    ]					,       )



					@require_tf
					@unittest.skip('Zero Shot Object Detection not implemented in TF'       )
					def      _a				(  self       :							int       ):






								"""simple docstring"""


								pass



					@require_torch
					@slow
					def      _a				(  self       :							str       ):






								"""simple docstring"""


								A__	  =     0.2
								A__	  =     pipeline('zero-shot-object-detection'       )

								A__	  =     object_detector(
								    'http://images.cocodataset.org/val2017/000000039769.jpg'					,       candidate_labels=['cat', 'remote', 'couch']					,       threshold=_snake_case					,       )
								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        {'score': 0.2868, 'label': 'cat', 'box': {'xmin': 3_24, 'ymin': 20, 'xmax': 6_40, 'ymax': 3_73}},
								        {'score': 0.277, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 1_77, 'ymax': 1_15}},
								        {'score': 0.2537, 'label': 'cat', 'box': {'xmin': 1, 'ymin': 55, 'xmax': 3_15, 'ymax': 4_72}},
								    ]					,       )



					@require_torch
					@slow
					def      _a				(  self       :							Any       ):






								"""simple docstring"""


								A__	  =     2
								A__	  =     pipeline('zero-shot-object-detection'       )

								A__	  =     object_detector(
								    'http://images.cocodataset.org/val2017/000000039769.jpg'					,       candidate_labels=['cat', 'remote', 'couch']					,       top_k=_snake_case					,       )
								self.assertEqual(
								    nested_simplify(_snake_case					,       decimals=4       )					,       [
								        {'score': 0.2868, 'label': 'cat', 'box': {'xmin': 3_24, 'ymin': 20, 'xmax': 6_40, 'ymax': 3_73}},
								        {'score': 0.277, 'label': 'remote', 'box': {'xmin': 40, 'ymin': 72, 'xmax': 1_77, 'ymax': 1_15}},
								    ]					,       )




