'''simple docstring'''

import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import (
      XLMForMultipleChoice,
      XLMForQuestionAnswering,
      XLMForQuestionAnsweringSimple,
      XLMForSequenceClassification,
      XLMForTokenClassification,
      XLMModel,
      XLMWithLMHeadModel,
  )
  from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST



class      UpperCAmelCase    :





       '''simple docstring'''



       def __init__(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase=13      ,  __lowerCAmelCase=7      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=False      ,  __lowerCAmelCase=False      ,  __lowerCAmelCase=False      ,  __lowerCAmelCase=2      ,  __lowerCAmelCase=99      ,  __lowerCAmelCase=0      ,  __lowerCAmelCase=32      ,  __lowerCAmelCase=5      ,  __lowerCAmelCase=4      ,  __lowerCAmelCase=0.1      ,  __lowerCAmelCase=0.1      ,  __lowerCAmelCase=512      ,  __lowerCAmelCase=2      ,  __lowerCAmelCase=0.0_2      ,  __lowerCAmelCase=2      ,  __lowerCAmelCase=4      ,  __lowerCAmelCase="last"      ,  __lowerCAmelCase=True      ,  __lowerCAmelCase=None      ,  __lowerCAmelCase=0      ,  )    ->    Tuple:
        lowercase__							:       Tuple										=   parent
        lowercase__							:       List[str]										=   batch_size
        lowercase__							:       Optional[Any]										=   seq_length
        lowercase__							:       Optional[Any]										=   is_training
        lowercase__							:       Tuple										=   use_input_lengths
        lowercase__							:       Union[str, Any]										=   use_token_type_ids
        lowercase__							:       Dict										=   use_labels
        lowercase__							:       Any										=   gelu_activation
        lowercase__							:       Tuple										=   sinusoidal_embeddings
        lowercase__							:       str										=   causal
        lowercase__							:       List[str]										=   asm
        lowercase__							:       str										=   n_langs
        lowercase__							:       Tuple										=   vocab_size
        lowercase__							:       List[str]										=   n_special
        lowercase__							:       Union[str, Any]										=   hidden_size
        lowercase__							:       Any										=   num_hidden_layers
        lowercase__							:       int										=   num_attention_heads
        lowercase__							:       Optional[Any]										=   hidden_dropout_prob
        lowercase__							:       List[Any]										=   attention_probs_dropout_prob
        lowercase__							:       List[str]										=   max_position_embeddings
        lowercase__							:       Dict										=   type_sequence_label_size
        lowercase__							:       Any										=   initializer_range
        lowercase__							:       List[str]										=   num_labels
        lowercase__							:       List[Any]										=   num_choices
        lowercase__							:       Any										=   summary_type
        lowercase__							:       int										=   use_proj
        lowercase__							:       str										=   scope
        lowercase__							:       str										=   bos_token_id

       def _lowerCAmelCase(		self					)    ->    str:
        lowercase__							:       int										=   ids_tensor([self.batch_size, self.seq_length]      ,  self.vocab_size					)
        lowercase__							:       Optional[Any]										=   random_attention_mask([self.batch_size, self.seq_length]					)

        lowercase__							:       Dict										=   None
        if self.use_input_lengths:
         lowercase__							:       int										=   (
             ids_tensor([self.batch_size]      ,  vocab_size=2					) + self.seq_length - 2
         )  # small variation of seq_length

        lowercase__							:       str										=   None
        if self.use_token_type_ids:
         lowercase__							:       int										=   ids_tensor([self.batch_size, self.seq_length]      ,  self.n_langs					)

        lowercase__							:       List[Any]										=   None
        lowercase__							:       Union[str, Any]										=   None
        lowercase__							:       Dict										=   None
        if self.use_labels:
         lowercase__							:       Optional[int]										=   ids_tensor([self.batch_size]      ,  self.type_sequence_label_size					)
         lowercase__							:       Any										=   ids_tensor([self.batch_size, self.seq_length]      ,  self.num_labels					)
         lowercase__							:       Tuple										=   ids_tensor([self.batch_size]      ,  2					).float()
         lowercase__							:       Any										=   ids_tensor([self.batch_size]      ,  self.num_choices					)

        lowercase__							:       str										=   self.get_config()

        return (
            config,
            input_ids,
            token_type_ids,
            input_lengths,
            sequence_labels,
            token_labels,
            is_impossible_labels,
            choice_labels,
            input_mask,
        )

       def _lowerCAmelCase(		self					)    ->    Union[str, Any]:
        return XLMConfig(
            vocab_size=self.vocab_size      ,  n_special=self.n_special      ,  emb_dim=self.hidden_size      ,  n_layers=self.num_hidden_layers      ,  n_heads=self.num_attention_heads      ,  dropout=self.hidden_dropout_prob      ,  attention_dropout=self.attention_probs_dropout_prob      ,  gelu_activation=self.gelu_activation      ,  sinusoidal_embeddings=self.sinusoidal_embeddings      ,  asm=self.asm      ,  causal=self.causal      ,  n_langs=self.n_langs      ,  max_position_embeddings=self.max_position_embeddings      ,  initializer_range=self.initializer_range      ,  summary_type=self.summary_type      ,  use_proj=self.use_proj      ,  num_labels=self.num_labels      ,  bos_token_id=self.bos_token_id      ,  )

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    int:
        lowercase__							:       int										=   XLMModel(config=snake_case_					)
        model.to(snake_case_					)
        model.eval()
        lowercase__							:       Optional[int]										=   model(snake_case_      ,  lengths=snake_case_      ,  langs=snake_case_					)
        lowercase__							:       Dict										=   model(snake_case_      ,  langs=snake_case_					)
        lowercase__							:       Union[str, Any]										=   model(snake_case_					)
        self.parent.assertEqual(result.last_hidden_state.shape      ,  (self.batch_size, self.seq_length, self.hidden_size)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    List[Any]:
        lowercase__							:       List[Any]										=   XLMWithLMHeadModel(snake_case_					)
        model.to(snake_case_					)
        model.eval()

        lowercase__							:       Any										=   model(snake_case_      ,  token_type_ids=snake_case_      ,  labels=snake_case_					)
        self.parent.assertEqual(result.loss.shape      ,  ()					)
        self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.vocab_size)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    Any:
        lowercase__							:       Optional[int]										=   XLMForQuestionAnsweringSimple(snake_case_					)
        model.to(snake_case_					)
        model.eval()

        lowercase__							:       str										=   model(snake_case_					)

        lowercase__							:       Dict										=   model(snake_case_      ,  start_positions=snake_case_      ,  end_positions=snake_case_					)
        lowercase__							:       Optional[Any]										=   outputs
        self.parent.assertEqual(result.start_logits.shape      ,  (self.batch_size, self.seq_length)					)
        self.parent.assertEqual(result.end_logits.shape      ,  (self.batch_size, self.seq_length)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    List[str]:
        lowercase__							:       Tuple										=   XLMForQuestionAnswering(snake_case_					)
        model.to(snake_case_					)
        model.eval()

        lowercase__							:       Tuple										=   model(snake_case_					)

        lowercase__							:       Tuple										=   model(
            snake_case_      ,  start_positions=snake_case_      ,  end_positions=snake_case_      ,  cls_index=snake_case_      ,  is_impossible=snake_case_      ,  p_mask=snake_case_      ,  )

        lowercase__							:       List[str]										=   model(
            snake_case_      ,  start_positions=snake_case_      ,  end_positions=snake_case_      ,  cls_index=snake_case_      ,  is_impossible=snake_case_      ,  )

        (lowercase__							)							:       Tuple										=   result_with_labels.to_tuple()

        lowercase__							:       Dict										=   model(snake_case_      ,  start_positions=snake_case_      ,  end_positions=snake_case_					)

        (lowercase__							)							:       Dict										=   result_with_labels.to_tuple()

        self.parent.assertEqual(result_with_labels.loss.shape      ,  ()					)
        self.parent.assertEqual(result.start_top_log_probs.shape      ,  (self.batch_size, model.config.start_n_top)					)
        self.parent.assertEqual(result.start_top_index.shape      ,  (self.batch_size, model.config.start_n_top)					)
        self.parent.assertEqual(
            result.end_top_log_probs.shape      ,  (self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
        self.parent.assertEqual(
            result.end_top_index.shape      ,  (self.batch_size, model.config.start_n_top * model.config.end_n_top)					)
        self.parent.assertEqual(result.cls_logits.shape      ,  (self.batch_size,)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    Union[str, Any]:
        lowercase__							:       Tuple										=   XLMForSequenceClassification(snake_case_					)
        model.to(snake_case_					)
        model.eval()

        lowercase__							:       Dict										=   model(snake_case_					)
        lowercase__							:       Optional[int]										=   model(snake_case_      ,  labels=snake_case_					)
        self.parent.assertEqual(result.loss.shape      ,  ()					)
        self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.type_sequence_label_size)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    Tuple:
        lowercase__							:       List[str]										=   self.num_labels
        lowercase__							:       Union[str, Any]										=   XLMForTokenClassification(snake_case_					)
        model.to(snake_case_					)
        model.eval()

        lowercase__							:       List[Any]										=   model(snake_case_      ,  attention_mask=snake_case_      ,  labels=snake_case_					)
        self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.seq_length, self.num_labels)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  )    ->    Tuple:
        lowercase__							:       Tuple										=   self.num_choices
        lowercase__							:       Dict										=   XLMForMultipleChoice(config=snake_case_					)
        model.to(snake_case_					)
        model.eval()
        lowercase__							:       Optional[Any]										=   input_ids.unsqueeze(1					).expand(-1      ,  self.num_choices      ,  -1					).contiguous()
        lowercase__							:       Dict										=   token_type_ids.unsqueeze(1					).expand(-1      ,  self.num_choices      ,  -1					).contiguous()
        lowercase__							:       str										=   input_mask.unsqueeze(1					).expand(-1      ,  self.num_choices      ,  -1					).contiguous()
        lowercase__							:       Dict										=   model(
            snake_case_      ,  attention_mask=snake_case_      ,  token_type_ids=snake_case_      ,  labels=snake_case_      ,  )
        self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_choices)					)


       def _lowerCAmelCase(		self					)    ->    int:
        lowercase__							:       Optional[int]										=   self.prepare_config_and_inputs()
        (
            lowercase__							
        )							:       List[str]										=   config_and_inputs
        lowercase__							:       Optional[int]										=   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
        return config, inputs_dict



@require_torch
class      UpperCAmelCase    (		__A       ,		__A       ,		__A       ,		unittest.TestCase   ):





       '''simple docstring'''



       SCREAMING_SNAKE_CASE   						=       (
           (
               XLMModel,
               XLMWithLMHeadModel,
               XLMForQuestionAnswering,
               XLMForSequenceClassification,
               XLMForQuestionAnsweringSimple,
               XLMForTokenClassification,
               XLMForMultipleChoice,
           )
           if is_torch_available()
           else ()
       )
       SCREAMING_SNAKE_CASE   						=       (
           (XLMWithLMHeadModel,) if is_torch_available() else ()
       )  # TODO (PVP): Check other models whether language generation is also applicable
       SCREAMING_SNAKE_CASE   						=       (
           {
               'feature-extraction': XLMModel,
               'fill-mask': XLMWithLMHeadModel,
               'question-answering': XLMForQuestionAnsweringSimple,
               'text-classification': XLMForSequenceClassification,
               'text-generation': XLMWithLMHeadModel,
               'token-classification': XLMForTokenClassification,
               'zero-shot': XLMForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase					)    ->    Optional[Any]:


        if (
            pipeline_test_casse_name == "QAPipelineTests"
            and tokenizer_name is not None
            and not tokenizer_name.endswith('''Fast'''					)
        ):
         # `QAPipelineTests` fails for a few models when the slower tokenizer are used.
         # (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
         # TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
         return True

        return False

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase=False					)    ->    Tuple:
        lowercase__							:       Optional[Any]										=   super()._prepare_for_class(snake_case_      ,  snake_case_      ,  return_labels=snake_case_					)

        if return_labels:
         if model_class.__name__ == "XLMForQuestionAnswering":
          lowercase__							:       List[str]										=   torch.zeros(
              self.model_tester.batch_size      ,  dtype=torch.long      ,  device=snake_case_					)
          lowercase__							:       List[Any]										=   torch.zeros(
              self.model_tester.batch_size      ,  dtype=torch.long      ,  device=snake_case_					)

        return inputs_dict

       def _lowerCAmelCase(		self					)    ->    Dict:
        lowercase__							:       Optional[Any]										=   XLMModelTester(self					)
        lowercase__							:       Tuple										=   ConfigTester(self      ,  config_class=snake_case_      ,  emb_dim=37					)

       def _lowerCAmelCase(		self					)    ->    Dict:
        self.config_tester.run_common_tests()

       def _lowerCAmelCase(		self					)    ->    int:
        lowercase__							:       Any										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_model(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    Tuple:
        lowercase__							:       Optional[Any]										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_lm_head(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    Optional[int]:
        lowercase__							:       Any										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_simple_qa(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    int:
        lowercase__							:       Tuple										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_qa(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    Optional[int]:
        lowercase__							:       Any										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_sequence_classif(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    Union[str, Any]:
        lowercase__							:       List[str]										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_token_classif(*snake_case_					)

       def _lowerCAmelCase(		self					)    ->    str:
        lowercase__							:       List[str]										=   self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_multiple_choice(*snake_case_					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase=False      ,  __lowerCAmelCase=1					)    ->    Union[str, Any]:
        self.assertIsInstance(snake_case_      ,  snake_case_					)
        self.assertListEqual(
            [isinstance(snake_case_      ,  snake_case_					) for iter_attentions in attentions]      ,  [True] * len(snake_case_					)					)
        self.assertEqual(len(snake_case_					)      ,  (max_length - min_length) * num_beam_groups					)

        for idx, iter_attentions in enumerate(snake_case_					):
         # adds PAD dummy token
         lowercase__							:       Tuple										=   min_length + idx + 1
         lowercase__							:       Tuple										=   min_length + idx + 1

         lowercase__							:       Optional[Any]										=   (
             batch_size * num_beam_groups,
             config.num_attention_heads,
             tgt_len,
             src_len,
         )
         # check attn size
         self.assertListEqual(
             [layer_attention.shape for layer_attention in iter_attentions]      ,  [expected_shape] * len(snake_case_					)					)

       def _lowerCAmelCase(		self      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase      ,  __lowerCAmelCase=False      ,  __lowerCAmelCase=1					)    ->    Union[str, Any]:
        self.assertIsInstance(snake_case_      ,  snake_case_					)
        self.assertListEqual(
            [isinstance(snake_case_      ,  snake_case_					) for iter_hidden_states in hidden_states]      ,  [True] * len(snake_case_					)      ,  )
        self.assertEqual(len(snake_case_					)      ,  (max_length - min_length) * num_beam_groups					)

        for idx, iter_hidden_states in enumerate(snake_case_					):
         # adds PAD dummy token
         lowercase__							:       List[str]										=   min_length + idx + 1
         lowercase__							:       Optional[int]										=   (batch_size * num_beam_groups, seq_len, config.hidden_size)
         # check hidden size
         self.assertListEqual(
             [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]      ,  [expected_shape] * len(snake_case_					)      ,  )
        pass


       @slow
       def _lowerCAmelCase(		self					)    ->    Tuple:


        for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         lowercase__							:       Union[str, Any]										=   XLMModel.from_pretrained(snake_case_					)
         self.assertIsNotNone(snake_case_					)



@require_torch
class      UpperCAmelCase    (		unittest.TestCase   ):





       '''simple docstring'''



       @slow
       def _lowerCAmelCase(		self					)    ->    List[Any]:
        lowercase__							:       int										=   XLMWithLMHeadModel.from_pretrained('''xlm-mlm-en-2048'''					)
        model.to(snake_case_					)
        lowercase__							:       Tuple										=   torch.tensor([[14, 447]]      ,  dtype=torch.long      ,  device=snake_case_					)  # the president
        lowercase__							:       Optional[int]										=   [
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
            14,
            447,
        ]  # the president the president the president the president the president the president the president the president the president the president
        # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
        lowercase__							:       int										=   model.generate(snake_case_      ,  do_sample=snake_case_					)
        self.assertListEqual(output_ids[0].cpu().numpy().tolist()      ,  snake_case_					)


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


_UpperCamelCase								=		{
    "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__       :     Any								=			["""layers""", """blocks"""]
							for k in ignore_keys:
														state_dict.pop(lowerCAmelCase			,  lowerCAmelCase							)


_UpperCamelCase								=		{
    "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__       :     str								=			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[Any]								=			new_key.replace(lowerCAmelCase			,  lowerCAmelCase							)

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

														A__       :     Any								=			s_dict.pop(lowerCAmelCase							)
							return s_dict

def  _A(		lowerCAmelCase							):
							A__				,							A__       :     Union[str, Any]								=			emb.weight.shape
							A__       :     Optional[Any]								=			nn.Linear(lowerCAmelCase			,  lowerCAmelCase			,  bias=lowerCAmelCase							)
							A__       :     int								=			emb.weight.data
							return lin_layer

def  _A(		lowerCAmelCase			,  lowerCAmelCase							):
							os.makedirs(lowerCAmelCase			,  exist_ok=lowerCAmelCase							)
							A__       :     Union[str, Any]								=			os.path.basename(lowerCAmelCase							)

							A__       :     str								=			url.split("""/"""							)[-2]
							A__       :     Union[str, Any]								=			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__       :     Tuple								=			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__       :     Dict								=			source.read(8192							)
																												if not buffer:
																																			break

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

							A__       :     List[str]								=			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__       :     Optional[int]								=			torch.load(lowerCAmelCase			,  map_location="""cpu"""							)
							A__       :     List[str]								=			original_checkpoint["""dims"""]
							A__       :     List[Any]								=			original_checkpoint["""model_state_dict"""]
							A__       :     Dict								=			state_dict["""decoder.token_embedding.weight"""]
							remove_ignore_keys_(lowerCAmelCase							)
							rename_keys(lowerCAmelCase							)
							A__       :     Any								=			True
							A__       :     List[str]								=			state_dict["""decoder.layers.0.fc1.weight"""].shape[0]

							A__       :     str								=			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__       :     int								=			WhisperForConditionalGeneration(lowerCAmelCase							)
							A__				,							A__       :     List[Any]								=			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__       :     Tuple								=			make_linear_from_emb(model.model.decoder.embed_tokens							)
							else:
														A__       :     Any								=			proj_out_weights

							model.save_pretrained(lowerCAmelCase							)


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

				convert_openai_whisper_to_tfms(args.checkpoint_path, args.pytorch_dump_folder_path)





