





import os
from pathlib import Path

import numpy as np
import pytest
from pack_dataset import pack_data_dir
from parameterized import parameterized
from save_len_file import save_len_file
from torch.utils.data import DataLoader

from transformers import AutoTokenizer
from transformers.models.mbart.modeling_mbart import shift_tokens_right
from transformers.testing_utils import TestCasePlus, slow
from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeqaSeqDataset, SeqaSeqDataset


SCREAMING_SNAKE_CASE__       :			Union[str, Any]			      =						"""bert-base-cased"""
SCREAMING_SNAKE_CASE__       :			Optional[Any]			      =						"""google/pegasus-xsum"""
SCREAMING_SNAKE_CASE__       :			Union[str, Any]			      =						[""" Sam ate lunch today.""", """Sams lunch ingredients."""]
SCREAMING_SNAKE_CASE__       :			Dict			      =						["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
SCREAMING_SNAKE_CASE__       :			int			      =						"""patrickvonplaten/t5-tiny-random"""
SCREAMING_SNAKE_CASE__       :			str			      =						"""sshleifer/bart-tiny-random"""
SCREAMING_SNAKE_CASE__       :			Tuple			      =						"""sshleifer/tiny-mbart"""
SCREAMING_SNAKE_CASE__       :			str			      =						"""sshleifer/tiny-marian-en-de"""






def       _A  (      lowerCamelCase							, lowerCamelCase   ):
   a__	:						Union[str, Any]      		=  "\n".join(lowerCamelCase   )
   Path(lowerCamelCase   ).open("w"   ).writelines(lowerCamelCase   )






def       _A  (      lowerCamelCase   ):
   for split in ["train", "val", "test"]:
      _dump_articles(os.path.join(lowerCamelCase							, F"""{split}.source"""   )							, lowerCamelCase   )
      _dump_articles(os.path.join(lowerCamelCase							, F"""{split}.target"""   )							, lowerCamelCase   )
   return tmp_dir







class   __lowerCAmelCase			(					_UpperCamelCase  ):





   @parameterized.expand(
       [
           MBART_TINY,
           MARIAN_TINY,
           T5_TINY,
           BART_TINY,
           PEGASUS_XSUM,
       ]					,			)
   @slow
   def 	_snake_case ( self					,			snake_case      )				->  Dict:
      """simple docstring"""





      a__	:						Optional[int]      		=  AutoTokenizer.from_pretrained(snake_case      )
      a__	:						str      		=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )
      a__	:						int      		=  max(len(tokenizer.encode(snake_case      )      ) for a in ARTICLES      )
      a__	:						List[Any]      		=  max(len(tokenizer.encode(snake_case      )      ) for a in SUMMARIES      )
      a__	:						Tuple      		=  4
      a__	:						Optional[Any]      		=  8
      assert max_len_target > max_src_len  # Will be truncated
      assert max_len_source > max_src_len  # Will be truncated
      a__		,     a__	:						Union[str, Any]      		=  "ro_RO", "de_DE"  # ignored for all but mbart, but never causes error.
      a__	:						Tuple      		=  SeqaSeqDataset(
          snake_case					,			data_dir=snake_case					,			type_path="train"					,			max_source_length=snake_case					,			max_target_length=snake_case					,			src_lang=snake_case					,			tgt_lang=snake_case					,			)
      a__	:						int      		=  DataLoader(snake_case					,			batch_size=2					,			collate_fn=train_dataset.collate_fn      )
      for batch in dataloader:
         assert isinstance(snake_case					,			snake_case      )
         assert batch["attention_mask"].shape == batch["input_ids"].shape
         # show that articles were trimmed.
         assert batch["input_ids"].shape[1] == max_src_len
         # show that targets are the same len
         assert batch["labels"].shape[1] == max_tgt_len
         if tok_name != MBART_TINY:
            continue
         # check language codes in correct place
         a__	:						str      		=  shift_tokens_right(batch["labels"]					,			tokenizer.pad_token_id      )
         assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
         assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
         assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
         assert batch["input_ids"][0, -1].item() == tokenizer.lang_code_to_id[src_lang]

         break  # No need to test every batch





   @parameterized.expand([BART_TINY, BERT_BASE_CASED]      )
   def 	_snake_case ( self					,			snake_case      )				->  List[str]:
      """simple docstring"""





      a__	:						Optional[Any]      		=  AutoTokenizer.from_pretrained(snake_case      )
      a__	:						Union[str, Any]      		=  make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )
      a__	:						List[Any]      		=  max(len(tokenizer.encode(snake_case      )      ) for a in ARTICLES      )
      a__	:						List[str]      		=  max(len(tokenizer.encode(snake_case      )      ) for a in SUMMARIES      )
      a__	:						Optional[int]      		=  4
      a__	:						Union[str, Any]      		=  LegacySeqaSeqDataset(
          snake_case					,			data_dir=snake_case					,			type_path="train"					,			max_source_length=20					,			max_target_length=snake_case					,			)
      a__	:						Union[str, Any]      		=  DataLoader(snake_case					,			batch_size=2					,			collate_fn=train_dataset.collate_fn      )
      for batch in dataloader:
         assert batch["attention_mask"].shape == batch["input_ids"].shape
         # show that articles were trimmed.
         assert batch["input_ids"].shape[1] == max_len_source
         assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
         # show that targets were truncated
         assert batch["labels"].shape[1] == trunc_target  # Truncated
         assert max_len_target > trunc_target  # Truncated
         break  # No need to test every batch





   def 	_snake_case ( self      )				->  Dict:
      """simple docstring"""





      a__	:						List[str]      		=  AutoTokenizer.from_pretrained("facebook/mbart-large-cc25"      )

      a__	:						str      		=  Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )      )
      a__	:						Tuple      		=  tmp_dir.joinpath("train.source"      ).open().readlines()
      a__	:						int      		=  Path(make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )      )
      pack_data_dir(snake_case					,			snake_case					,			128					,			snake_case      )
      a__	:						str      		=  {x.name for x in tmp_dir.iterdir()}
      a__	:						Dict      		=  {x.name for x in save_dir.iterdir()}
      a__	:						Dict      		=  save_dir.joinpath("train.source"      ).open().readlines()
      # orig: [' Sam ate lunch today.\n', 'Sams lunch ingredients.']
      # desired_packed: [' Sam ate lunch today.\n Sams lunch ingredients.']
      assert len(snake_case      ) < len(snake_case      )
      assert len(snake_case      ) == 1
      assert len(packed_examples[0]      ) == sum(len(snake_case      ) for x in orig_examples      )
      assert orig_paths == new_paths





   @pytest.mark.skipif(not FAIRSEQ_AVAILABLE					,			reason="This test requires fairseq"      )
   def 	_snake_case ( self      )				->  Any:
      """simple docstring"""





      if not FAIRSEQ_AVAILABLE:
         return
      a__		,     a__		,     a__	:						Optional[int]      		=  self._get_dataset(max_len=64      )
      a__	:						List[Any]      		=  64
      a__	:						Optional[int]      		=  ds.make_dynamic_sampler(snake_case					,			required_batch_size_multiple=snake_case      )
      a__	:						Union[str, Any]      		=  [len(snake_case      ) for x in batch_sampler]
      assert len(set(snake_case      )      ) > 1  # it's not dynamic batch size if every batch is the same length
      assert sum(snake_case      ) == len(snake_case      )  # no dropped or added examples
      a__	:						Tuple      		=  DataLoader(snake_case					,			batch_sampler=snake_case					,			collate_fn=ds.collate_fn					,			num_workers=2      )
      a__	:						List[Any]      		=  []
      a__	:						int      		=  []
      for batch in data_loader:
         a__	:						Tuple      		=  batch["input_ids"].shape
         a__	:						Tuple      		=  src_shape[0]
         assert bs % required_batch_size_multiple == 0 or bs < required_batch_size_multiple
         a__	:						Any      		=  np.product(batch["input_ids"].shape      )
         num_src_per_batch.append(snake_case      )
         if num_src_tokens > (max_tokens * 1.1):
            failures.append(snake_case      )
      assert num_src_per_batch[0] == max(snake_case      )
      if failures:
         raise AssertionError(F"""too many tokens in {len(snake_case      )} batches"""      )





   def 	_snake_case ( self      )				->  str:
      """simple docstring"""





      a__		,     a__		,     a__	:						Any      		=  self._get_dataset(max_len=512      )
      a__	:						Dict      		=  2
      a__	:						Optional[int]      		=  ds.make_sortish_sampler(snake_case					,			shuffle=snake_case      )

      a__	:						int      		=  DataLoader(snake_case					,			batch_size=snake_case					,			collate_fn=ds.collate_fn					,			num_workers=2      )
      a__	:						Optional[Any]      		=  DataLoader(snake_case					,			batch_size=snake_case					,			collate_fn=ds.collate_fn					,			num_workers=2					,			sampler=snake_case      )

      a__	:						Optional[Any]      		=  tokenizer.pad_token_id

      def count_pad_tokens(snake_case					,			snake_case="input_ids"      ):
         return [batch[k].eq(snake_case      ).sum().item() for batch in data_loader]

      assert sum(count_pad_tokens(snake_case					,			k="labels"      )      ) < sum(count_pad_tokens(snake_case					,			k="labels"      )      )
      assert sum(count_pad_tokens(snake_case      )      ) < sum(count_pad_tokens(snake_case      )      )
      assert len(snake_case      ) == len(snake_case      )





   def 	_snake_case ( self					,			snake_case=1_000					,			snake_case=128      )				->  Tuple:
      """simple docstring"""





      if os.getenv("USE_REAL_DATA"					,			snake_case      ):
         a__	:						Tuple      		=  "examples/seq2seq/wmt_en_ro"
         a__	:						Tuple      		=  max_len * 2 * 64
         if not Path(snake_case      ).joinpath("train.len"      ).exists():
            save_len_file(snake_case					,			snake_case      )
      else:
         a__	:						str      		=  "examples/seq2seq/test_data/wmt_en_ro"
         a__	:						Dict      		=  max_len * 4
         save_len_file(snake_case					,			snake_case      )

      a__	:						str      		=  AutoTokenizer.from_pretrained(snake_case      )
      a__	:						List[Any]      		=  SeqaSeqDataset(
          snake_case					,			data_dir=snake_case					,			type_path="train"					,			max_source_length=snake_case					,			max_target_length=snake_case					,			n_obs=snake_case					,			)
      return ds, max_tokens, tokenizer





   def 	_snake_case ( self      )				->  Tuple:
      """simple docstring"""





      a__		,     a__		,     a__	:						str      		=  self._get_dataset()
      a__	:						Union[str, Any]      		=  set(DistributedSortishSampler(snake_case					,			256					,			num_replicas=2					,			rank=0					,			add_extra_examples=snake_case      )      )
      a__	:						int      		=  set(DistributedSortishSampler(snake_case					,			256					,			num_replicas=2					,			rank=1					,			add_extra_examples=snake_case      )      )
      assert idsa.intersection(snake_case      ) == set()





   @parameterized.expand(
       [
           MBART_TINY,
           MARIAN_TINY,
           T5_TINY,
           BART_TINY,
           PEGASUS_XSUM,
       ]					,			)
   def 	_snake_case ( self					,			snake_case      )				->  Any:
      """simple docstring"""





      a__	:						int      		=  AutoTokenizer.from_pretrained(snake_case					,			use_fast=snake_case      )
      if tok_name == MBART_TINY:
         a__	:						Optional[Any]      		=  SeqaSeqDataset(
             snake_case					,			data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )					,			type_path="train"					,			max_source_length=4					,			max_target_length=8					,			src_lang="EN"					,			tgt_lang="FR"					,			)
         a__	:						str      		=  train_dataset.dataset_kwargs
         assert "src_lang" in kwargs and "tgt_lang" in kwargs
      else:
         a__	:						Tuple      		=  SeqaSeqDataset(
             snake_case					,			data_dir=make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()      )					,			type_path="train"					,			max_source_length=4					,			max_target_length=8					,			)
         a__	:						Any      		=  train_dataset.dataset_kwargs
         assert "add_prefix_space" not in kwargs if tok_name != BART_TINY else "add_prefix_space" in kwargs
         assert len(snake_case      ) == 1 if tok_name == BART_TINY else len(snake_case      ) == 0








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)

