


'''simple docstring'''






from collections import OrderedDict
from typing import TYPE_CHECKING, Any, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxSeqaSeqConfigWithPast
from ...utils import logging


if TYPE_CHECKING:
  from ...feature_extraction_utils import FeatureExtractionMixin
  from ...tokenization_utils_base import PreTrainedTokenizerBase
  from ...utils import TensorType

__A :  str    							=				logging.get_logger(__name__)

__A :  str    							=				{
    "openai/whisper-base": "https://huggingface.co/openai/whisper-base/resolve/main/config.json",
}

# fmt: off
__A :  str    							=				[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 357, 366, 438, 532, 685,
    705, 796, 930, 1058, 1220, 1267, 1279, 1303, 1343, 1377,
    1391, 1635, 1782, 1875, 2162, 2361, 2488, 3467, 4008, 4211,
    4600, 4808, 5299, 5855, 6329, 7203, 9609, 9959, 1_0563, 1_0786,
    1_1420, 1_1709, 1_1907, 1_3163, 1_3697, 1_3700, 1_4808, 1_5306, 1_6410, 1_6791,
    1_7992, 1_9203, 1_9510, 2_0724, 2_2305, 2_2935, 2_7007, 3_0109, 3_0420, 3_3409,
    3_4949, 4_0283, 4_0493, 4_0549, 4_7282, 4_9146, 5_0257, 5_0359, 5_0360, 5_0361
]
__A :  Tuple    							=				[
    1, 2, 7, 8, 9, 10, 14, 25,
    26, 27, 28, 29, 31, 58, 59, 60, 61, 62,
    63, 90, 91, 92, 93, 359, 503, 522, 542, 873,
    893, 902, 918, 922, 931, 1350, 1853, 1982, 2460, 2627,
    3246, 3253, 3268, 3536, 3846, 3961, 4183, 4667, 6585, 6647,
    7273, 9061, 9383, 1_0428, 1_0929, 1_1938, 1_2033, 1_2331, 1_2562, 1_3793,
    1_4157, 1_4635, 1_5265, 1_5618, 1_6553, 1_6604, 1_8362, 1_8956, 2_0075, 2_1675,
    2_2520, 2_6130, 2_6161, 2_6435, 2_8279, 2_9464, 3_1650, 3_2302, 3_2470, 3_6865,
    4_2863, 4_7425, 4_9870, 5_0254, 5_0258, 5_0360, 5_0361, 5_0362
]






class        __snake_case	(       _SCREAMING_SNAKE_CASE):







    """simple docstring"""



    lowercase       			= 'whisper'
    lowercase       			= ['past_key_values']
    lowercase       			= {'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}

    def __init__(       self   :  Optional[Any]   ,						lowerCamelCase   :  List[str]=5_18_65   ,						lowerCamelCase   :  Union[str, Any]=80   ,						lowerCamelCase   :  Optional[int]=6   ,						lowerCamelCase   :  List[str]=4   ,						lowerCamelCase   :  Optional[int]=6   ,						lowerCamelCase   :  List[Any]=4   ,						lowerCamelCase   :  Union[str, Any]=15_36   ,						lowerCamelCase   :  str=15_36   ,						lowerCamelCase   :  Any=0.0   ,						lowerCamelCase   :  List[Any]=0.0   ,						lowerCamelCase   :  int=5_02_57   ,						lowerCamelCase   :  List[str]=True   ,						lowerCamelCase   :  str=True   ,						lowerCamelCase   :  Union[str, Any]="gelu"   ,						lowerCamelCase   :  Tuple=2_56   ,						lowerCamelCase   :  Union[str, Any]=0.0   ,						lowerCamelCase   :  Optional[int]=0.0   ,						lowerCamelCase   :  Optional[Any]=0.0   ,						lowerCamelCase   :  Optional[Any]=0.02   ,						lowerCamelCase   :  Optional[int]=False   ,						lowerCamelCase   :  Any=15_00   ,						lowerCamelCase   :  Union[str, Any]=4_48   ,						lowerCamelCase   :  Any=5_02_56   ,						lowerCamelCase   :  Union[str, Any]=5_02_56   ,						lowerCamelCase   :  List[Any]=5_02_56   ,						lowerCamelCase   :  List[str]=None   ,						lowerCamelCase   :  List[Any]=[2_20, 5_02_56]   ,						lowerCamelCase   :  Union[str, Any]=False   ,						lowerCamelCase   :  int=2_56   ,						lowerCamelCase   :  int=False   ,						lowerCamelCase   :  int=0.05   ,						lowerCamelCase   :  Union[str, Any]=10   ,						lowerCamelCase   :  Union[str, Any]=2   ,						lowerCamelCase   :  Tuple=0.0   ,						lowerCamelCase   :  List[str]=10   ,						lowerCamelCase   :  int=0   ,						lowerCamelCase   :  int=7   ,						**lowerCamelCase   :  Optional[Any]   ,						)       ->  int:
      lowerCAmelCase_      :				Tuple      =       vocab_size
      lowerCAmelCase_      :				str      =       num_mel_bins
      lowerCAmelCase_      :				Union[str, Any]      =       d_model
      lowerCAmelCase_      :				Optional[Any]      =       encoder_layers
      lowerCAmelCase_      :				List[str]      =       encoder_attention_heads
      lowerCAmelCase_      :				List[Any]      =       decoder_layers
      lowerCAmelCase_      :				Optional[Any]      =       decoder_attention_heads
      lowerCAmelCase_      :				int      =       decoder_ffn_dim
      lowerCAmelCase_      :				int      =       encoder_ffn_dim
      lowerCAmelCase_      :				Dict      =       dropout
      lowerCAmelCase_      :				Any      =       attention_dropout
      lowerCAmelCase_      :				int      =       activation_dropout
      lowerCAmelCase_      :				List[Any]      =       activation_function
      lowerCAmelCase_      :				Optional[Any]      =       init_std
      lowerCAmelCase_      :				int      =       encoder_layerdrop
      lowerCAmelCase_      :				Any      =       decoder_layerdrop
      lowerCAmelCase_      :				Union[str, Any]      =       use_cache
      lowerCAmelCase_      :				str      =       encoder_layers
      lowerCAmelCase_      :				str      =       scale_embedding  # scale factor will be sqrt(d_model) if True
      lowerCAmelCase_      :				Tuple      =       max_source_positions
      lowerCAmelCase_      :				Optional[int]      =       max_target_positions

      # Audio Classification-specific parameters. Feel free to ignore for other classes.
      lowerCAmelCase_      :				Any      =       classifier_proj_size
      lowerCAmelCase_      :				Tuple      =       use_weighted_layer_sum

      # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
      lowerCAmelCase_      :				Optional[Any]      =       apply_spec_augment
      lowerCAmelCase_      :				Union[str, Any]      =       mask_time_prob
      lowerCAmelCase_      :				Dict      =       mask_time_length
      lowerCAmelCase_      :				List[Any]      =       mask_time_min_masks
      lowerCAmelCase_      :				int      =       mask_feature_prob
      lowerCAmelCase_      :				int      =       mask_feature_length
      lowerCAmelCase_      :				str      =       mask_feature_min_masks

      lowerCAmelCase_      :				Union[str, Any]      =       median_filter_width

      super().__init__(
          pad_token_id=lowerCamelCase   ,						bos_token_id=lowerCamelCase   ,						eos_token_id=lowerCamelCase   ,						is_encoder_decoder=lowerCamelCase   ,						decoder_start_token_id=lowerCamelCase   ,						suppress_tokens=lowerCamelCase   ,						begin_suppress_tokens=lowerCamelCase   ,						**lowerCamelCase   ,						)







class        __snake_case	(       _SCREAMING_SNAKE_CASE):







    """simple docstring"""



    @property
    def __lowercase       (       self   :  Optional[int]		)       ->  Mapping[str, Mapping[int, str]]:
      lowerCAmelCase_      :				List[Any]      =       OrderedDict(
          [
              ("""input_features""", {0: """batch""", 1: """feature_size""", 2: """encoder_sequence"""}),
          ]		)
      if self.use_past:
        lowerCAmelCase_      :				Any      =       {0: """batch"""}
      else:
        lowerCAmelCase_      :				str      =       {0: """batch""", 1: """decoder_sequence"""}

      if self.use_past:
        self.fill_with_past_key_values_(lowerCamelCase   ,						direction="""inputs"""		)

      return common_inputs

    def __lowercase       (       self   :  int   ,						lowerCamelCase   :  Union["PreTrainedTokenizerBase", "FeatureExtractionMixin"]   ,						lowerCamelCase   :  int = -1   ,						lowerCamelCase   :  int = -1   ,						lowerCamelCase   :  bool = False   ,						lowerCamelCase   :  Optional["TensorType"] = None   ,						lowerCamelCase   :  int = 2_20_50   ,						lowerCamelCase   :  float = 5.0   ,						lowerCamelCase   :  int = 2_20   ,						)       ->  Mapping[str, Any]:
      lowerCAmelCase_      :				int      =       OrderedDict()
      lowerCAmelCase_      :				List[str]      =       OnnxConfig.generate_dummy_inputs(
          self   ,						preprocessor=preprocessor.feature_extractor   ,						batch_size=lowerCamelCase   ,						framework=lowerCamelCase   ,						sampling_rate=lowerCamelCase   ,						time_duration=lowerCamelCase   ,						frequency=lowerCamelCase   ,						)
      lowerCAmelCase_      :				Union[str, Any]      =       encoder_inputs["""input_features"""].shape[2]
      lowerCAmelCase_      :				Union[str, Any]      =       encoder_sequence_length // 2 if self.use_past else seq_length

      lowerCAmelCase_      :				str      =       super().generate_dummy_inputs(
          preprocessor.tokenizer   ,						lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase   ,						lowerCamelCase		)

      lowerCAmelCase_      :				Any      =       encoder_inputs.pop("""input_features"""		)
      lowerCAmelCase_      :				Optional[int]      =       decoder_inputs.pop("""decoder_input_ids"""		)

      if "past_key_values" in decoder_inputs:
        lowerCAmelCase_      :				List[str]      =       decoder_inputs.pop("""past_key_values"""		)

      return dummy_inputs





    @property
    def __lowercase       (       self   :  Optional[Any]		)       ->  float:
      return 1E-3







'''simple docstring'''






import argparse
import json
import os
from pathlib import Path

import requests
import torch

from transformers import JukeboxConfig, JukeboxModel
from transformers.utils import logging


logging.set_verbosity_info()
__A :  List[str]    							=				logging.get_logger(__name__)


__A :  str    							=				"https://openaipublic.azureedge.net/jukebox/models/"
__A :  Any    							=				{
    "jukebox-1b-lyrics": [
        "5b/vqvae.pth.tar",
        "5b/prior_level_0.pth.tar",
        "5b/prior_level_1.pth.tar",
        "1b_lyrics/prior_level_2.pth.tar",
    ],
    "jukebox-5b-lyrics": [
        "5b/vqvae.pth.tar",
        "5b/prior_level_0.pth.tar",
        "5b/prior_level_1.pth.tar",
        "5b_lyrics/prior_level_2.pth.tar",
    ],
}







def     UpperCamelCase_      (   A__ :       int ):
  '''simple docstring'''





  if key.endswith(""".model.1.bias""" ) and len(key.split(""".""" ) ) > 10:
    lowerCAmelCase_      :				Union[str, Any]      =       key.replace(""".model.1.bias"""							,   """.conv1d_1.bias""" )
  elif key.endswith(""".model.1.weight""" ) and len(key.split(""".""" ) ) > 10:
    lowerCAmelCase_      :				Union[str, Any]      =       key.replace(""".model.1.weight"""							,   """.conv1d_1.weight""" )
  elif key.endswith(""".model.3.bias""" ) and len(key.split(""".""" ) ) > 10:
    lowerCAmelCase_      :				Union[str, Any]      =       key.replace(""".model.3.bias"""							,   """.conv1d_2.bias""" )
  elif key.endswith(""".model.3.weight""" ) and len(key.split(""".""" ) ) > 10:
    lowerCAmelCase_      :				int      =       key.replace(""".model.3.weight"""							,   """.conv1d_2.weight""" )

  if "conditioner_blocks.0." in key:
    lowerCAmelCase_      :				Any      =       key.replace("""conditioner_blocks.0"""							,   """conditioner_blocks""" )

  if "prime_prior" in key:
    lowerCAmelCase_      :				str      =       key.replace("""prime_prior"""							,   """encoder""" )

  if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
    lowerCAmelCase_      :				int      =       key.replace(""".emb."""							,   """.""" )

  if key.endswith("""k""" ):  # replace vqvae.X.k with vqvae.X.codebook
    return key.replace(""".k"""							,   """.codebook""" )
  if "y_emb." in key:
    return key.replace("""y_emb."""							,   """metadata_embedding.""" )

  if "x_emb.emb." in key:
    lowerCAmelCase_      :				List[Any]      =       key.replace("""0.x_emb.emb"""							,   """embed_tokens""" )

  if "prime_state_ln" in key:
    return key.replace("""prime_state_ln"""							,   """encoder.final_layer_norm""" )
  if ".ln" in key:
    return key.replace(""".ln"""							,   """.layer_norm""" )
  if "_ln" in key:
    return key.replace("""_ln"""							,   """_layer_norm""" )

  if "prime_state_proj" in key:
    return key.replace("""prime_state_proj"""							,   """encoder.proj_in""" )
  if "prime_x_out" in key:
    return key.replace("""prime_x_out"""							,   """encoder.lm_head""" )
  if "prior.x_out" in key:
    return key.replace("""x_out"""							,   """fc_proj_out""" )
  if "x_emb" in key:
    return key.replace("""x_emb"""							,   """embed_tokens""" )

  return key







def     UpperCamelCase_      (   A__ :       Dict							,   A__ :       Optional[Any]							,   A__ :       Tuple							,   A__ :       Any ):
  '''simple docstring'''





  lowerCAmelCase_      :				Union[str, Any]      =       {}
  import re

  lowerCAmelCase_      :				Union[str, Any]      =       re.compile(R"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)""" )
  lowerCAmelCase_      :				str      =       re.compile(
      R"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)""" )
  lowerCAmelCase_      :				Dict      =       re.compile(R"""encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)""" )

  lowerCAmelCase_      :				int      =       re.compile(R"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)""" )
  lowerCAmelCase_      :				Optional[int]      =       re.compile(
      R"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)""" )
  lowerCAmelCase_      :				Union[str, Any]      =       re.compile(R"""decoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)""" )

  lowerCAmelCase_      :				Any      =       re.compile(R"""conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)""" )
  lowerCAmelCase_      :				Dict      =       re.compile(
      R"""conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)""" )
  lowerCAmelCase_      :				str      =       re.compile(R"""conditioner_blocks.(\d*).cond.model.(\d*).(bias|weight)""" )

  for original_key, value in state_dict.items():
    # rename vqvae.encoder keys
    if re_encoder_block_conv_in.fullmatch(A__ ):
      lowerCAmelCase_      :				Dict      =       re_encoder_block_conv_in.match(A__ )
      lowerCAmelCase_      :				Optional[int]      =       regex_match.groups()
      lowerCAmelCase_      :				Tuple      =       int(groups[2] ) * 2 + int(groups[3] )
      lowerCAmelCase_      :				Optional[Any]      =       f'encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}'
      lowerCAmelCase_      :				List[str]      =       re_encoder_block_conv_in.sub(A__							,   A__ )

    elif re_encoder_block_resnet.fullmatch(A__ ):
      lowerCAmelCase_      :				Tuple      =       re_encoder_block_resnet.match(A__ )
      lowerCAmelCase_      :				Tuple      =       regex_match.groups()
      lowerCAmelCase_      :				List[Any]      =       int(groups[2] ) * 2 + int(groups[3] )
      lowerCAmelCase_      :				str      =       {"""1""": 1, """3""": 2}[groups[-2]]
      lowerCAmelCase_      :				int      =       f'encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.'
      lowerCAmelCase_      :				Optional[Any]      =       f'resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}'
      lowerCAmelCase_      :				List[str]      =       prefix + resnet_block
      lowerCAmelCase_      :				Tuple      =       re_encoder_block_resnet.sub(A__							,   A__ )

    elif re_encoder_block_proj_out.fullmatch(A__ ):
      lowerCAmelCase_      :				int      =       re_encoder_block_proj_out.match(A__ )
      lowerCAmelCase_      :				Tuple      =       regex_match.groups()
      lowerCAmelCase_      :				Optional[Any]      =       f'encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}'
      lowerCAmelCase_      :				str      =       re_encoder_block_proj_out.sub(A__							,   A__ )

    # rename vqvae.decoder keys
    elif re_decoder_block_conv_out.fullmatch(A__ ):
      lowerCAmelCase_      :				List[Any]      =       re_decoder_block_conv_out.match(A__ )
      lowerCAmelCase_      :				Tuple      =       regex_match.groups()
      lowerCAmelCase_      :				str      =       int(groups[2] ) * 2 + int(groups[3] ) - 2
      lowerCAmelCase_      :				Any      =       f'decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}'
      lowerCAmelCase_      :				Tuple      =       re_decoder_block_conv_out.sub(A__							,   A__ )

    elif re_decoder_block_resnet.fullmatch(A__ ):
      lowerCAmelCase_      :				Optional[Any]      =       re_decoder_block_resnet.match(A__ )
      lowerCAmelCase_      :				Optional[Any]      =       regex_match.groups()
      lowerCAmelCase_      :				Optional[Any]      =       int(groups[2] ) * 2 + int(groups[3] ) - 2
      lowerCAmelCase_      :				Any      =       {"""1""": 1, """3""": 2}[groups[-2]]
      lowerCAmelCase_      :				Optional[int]      =       f'decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.'
      lowerCAmelCase_      :				Union[str, Any]      =       f'resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}'
      lowerCAmelCase_      :				Dict      =       prefix + resnet_block
      lowerCAmelCase_      :				Any      =       re_decoder_block_resnet.sub(A__							,   A__ )

    elif re_decoder_block_proj_in.fullmatch(A__ ):
      lowerCAmelCase_      :				str      =       re_decoder_block_proj_in.match(A__ )
      lowerCAmelCase_      :				Optional[int]      =       regex_match.groups()
      lowerCAmelCase_      :				str      =       f'decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}'
      lowerCAmelCase_      :				Union[str, Any]      =       re_decoder_block_proj_in.sub(A__							,   A__ )

    # rename prior cond.model to upsampler.upsample_block and resnet
    elif re_prior_cond_conv_out.fullmatch(A__ ):
      lowerCAmelCase_      :				Optional[Any]      =       re_prior_cond_conv_out.match(A__ )
      lowerCAmelCase_      :				Union[str, Any]      =       regex_match.groups()
      lowerCAmelCase_      :				Union[str, Any]      =       int(groups[1] ) * 2 + int(groups[2] ) - 2
      lowerCAmelCase_      :				Optional[int]      =       f'conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}'
      lowerCAmelCase_      :				Dict      =       re_prior_cond_conv_out.sub(A__							,   A__ )

    elif re_prior_cond_resnet.fullmatch(A__ ):
      lowerCAmelCase_      :				Any      =       re_prior_cond_resnet.match(A__ )
      lowerCAmelCase_      :				int      =       regex_match.groups()
      lowerCAmelCase_      :				Dict      =       int(groups[1] ) * 2 + int(groups[2] ) - 2
      lowerCAmelCase_      :				Tuple      =       {"""1""": 1, """3""": 2}[groups[-2]]
      lowerCAmelCase_      :				Optional[Any]      =       f'conditioner_blocks.upsampler.upsample_block.{block_index}.'
      lowerCAmelCase_      :				List[Any]      =       f'resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}'
      lowerCAmelCase_      :				Optional[int]      =       prefix + resnet_block
      lowerCAmelCase_      :				Dict      =       re_prior_cond_resnet.sub(A__							,   A__ )

    elif re_prior_cond_proj_in.fullmatch(A__ ):
      lowerCAmelCase_      :				List[str]      =       re_prior_cond_proj_in.match(A__ )
      lowerCAmelCase_      :				Optional[Any]      =       regex_match.groups()
      lowerCAmelCase_      :				Any      =       f'conditioner_blocks.upsampler.proj_in.{groups[-1]}'
      lowerCAmelCase_      :				List[str]      =       re_prior_cond_proj_in.sub(A__							,   A__ )

    # keep original key
    else:
      lowerCAmelCase_      :				Optional[Any]      =       original_key

    lowerCAmelCase_      :				Optional[Any]      =       replace_key(A__ )

    if f'{key_prefix}.{key}' not in model_state_dict or key is None:
      print(f'failed converting {original_key} to {key}, does not match' )

    # handle missmatched shape
    elif value.shape != model_state_dict[f'{key_prefix}.{key}'].shape:
      lowerCAmelCase_      :				Dict      =       model_state_dict[f'{key_prefix}.{key}']
      print(f'{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match' )
      lowerCAmelCase_      :				str      =       original_key

    lowerCAmelCase_      :				Dict      =       original_key
    lowerCAmelCase_      :				Optional[int]      =       value

  return new_dict







@torch.no_grad()
def     UpperCamelCase_      (   A__ :       Optional[Any]=None							,   A__ :       str=None ):
  '''simple docstring'''





  for file in MODEL_MAPPING[model_name]:
    if not os.path.isfile(f'{pytorch_dump_folder_path}/{file.split("/" )[-1]}' ):
      lowerCAmelCase_      :				List[Any]      =       requests.get(f'{PREFIX}{file}'							,   allow_redirects=A__ )
      os.makedirs(f'{pytorch_dump_folder_path}/'							,   exist_ok=A__ )
      open(f'{pytorch_dump_folder_path}/{file.split("/" )[-1]}'							,   """wb""" ).write(r.content )

  lowerCAmelCase_      :				Optional[int]      =       MODEL_MAPPING[model_name.split("""/""" )[-1]]

  lowerCAmelCase_      :				List[str]      =       JukeboxConfig.from_pretrained(A__ )
  lowerCAmelCase_      :				Dict      =       JukeboxModel(A__ )

  lowerCAmelCase_      :				Optional[Any]      =       []
  lowerCAmelCase_      :				Dict      =       {}
  for i, dict_name in enumerate(A__ ):
    lowerCAmelCase_      :				List[Any]      =       torch.load(f'{pytorch_dump_folder_path}/{dict_name.split("/" )[-1]}' )["""model"""]

    lowerCAmelCase_      :				Dict      =       {}
    for k in old_dic.keys():
      if k.endswith(""".b""" ):
        lowerCAmelCase_      :				Optional[int]      =       old_dic[k]
      elif k.endswith(""".w""" ):
        lowerCAmelCase_      :				Optional[int]      =       old_dic[k]
      elif "level_2" not in dict_name and "cond.model." in k:
        lowerCAmelCase_      :				Tuple      =       old_dic[k]
      else:
        lowerCAmelCase_      :				List[Any]      =       old_dic[k]

    lowerCAmelCase_      :				Union[str, Any]      =       """vqvae""" if i == 0 else f'priors.{3 - i}'
    lowerCAmelCase_      :				str      =       fix_jukebox_keys(A__							,   model.state_dict()							,   A__							,   A__ )
    weight_dict.append(A__ )

  lowerCAmelCase_      :				int      =       weight_dict.pop(0 )
  model.vqvae.load_state_dict(A__ )
  for i in range(len(A__ ) ):
    model.priors[i].load_state_dict(weight_dict[2 - i] )

  Path(A__ ).mkdir(exist_ok=A__ )
  with open(f'{pytorch_dump_folder_path}/mapping.json'							,   """w""" ) as txtfile:
    json.dump(A__							,   A__ )

  print(f'Saving model {model_name} to {pytorch_dump_folder_path}' )
  model.save_pretrained(A__ )

  return weight_dict


if __name__ == "__main__":
  __A :  str    							=				argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--model_name",
      default="jukebox-5b-lyrics",
      type=str,
      help="Name of the model you'd like to convert.",
  )
  parser.add_argument(
      "--pytorch_dump_folder_path",
      default="jukebox-5b-lyrics-converted",
      type=str,
      help="Path to the output PyTorch model directory.",
  )
  __A :  Optional[Any]    							=				parser.parse_args()
  convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)



