



from __future__ import annotations

from decimal import Decimal
from math import *  # noqa: F403

from sympy import diff






def 	_lowerCAmelCase						(				__lowerCamelCase   :     str		,	__lowerCamelCase   :     float | Decimal		,	__lowerCamelCase   :     float = 10**-10     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Dict       =	a
 while True:
  __SCREAMING_SNAKE_CASE       :    Optional[int]       =	Decimal(__lowerCamelCase     ) - (
      Decimal(eval(__lowerCamelCase     )     ) / Decimal(eval(str(diff(__lowerCamelCase     )     )     )     )  # noqa: S307
  )
  # This number dictates the accuracy of the answer
  if abs(eval(__lowerCamelCase     )     ) < precision:  # noqa: S307
   return float(__lowerCamelCase     )


# Let's Execute
if __name__ == "__main__":
      # Find root of trigonometric function
      # Find value of pi
      print(f'''The root of sin(x) = 0 is {newton_raphson('sin(x)', 2)}''')
      # Find root of polynomial
      print(f'''The root of x**2 - 5*x + 2 = 0 is {newton_raphson('x**2 - 5*x + 2', 0.4)}''')
      # Find Square Root of 5
      print(f'''The root of log(x) - 1 = 0 is {newton_raphson('log(x) - 1', 2)}''')
      # Exponential Roots
      print(f'''The root of exp(x) - 1 = 0 is {newton_raphson('exp(x) - 1', 0)}''')











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()
_lowerCamelCase						    =  logging.get_logger(__name__)


_lowerCamelCase						    =  """https://openaipublic.azureedge.net/jukebox/models/"""
_lowerCamelCase						    =  {
    """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 	_lowerCAmelCase						(				__lowerCamelCase   :     Any     ):



 """simple docstring"""





 if key.endswith(".model.1.bias"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace(".model.1.bias"		,	".conv1d_1.bias"     )
 elif key.endswith(".model.1.weight"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Tuple       =	key.replace(".model.1.weight"		,	".conv1d_1.weight"     )
 elif key.endswith(".model.3.bias"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace(".model.3.bias"		,	".conv1d_2.bias"     )
 elif key.endswith(".model.3.weight"     ) and len(key.split("."     )     ) > 10:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	key.replace(".model.3.weight"		,	".conv1d_2.weight"     )

 if "conditioner_blocks.0." in key:
  __SCREAMING_SNAKE_CASE       :    List[Any]       =	key.replace("conditioner_blocks.0"		,	"conditioner_blocks"     )

 if "prime_prior" in key:
  __SCREAMING_SNAKE_CASE       :    Dict       =	key.replace("prime_prior"		,	"encoder"     )

 if ".emb." in key and "total" not in key and "absolute" not in key and "relative" not in key:
  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	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:
  __SCREAMING_SNAKE_CASE       :    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 	_lowerCAmelCase						(				__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Union[str, Any]		,	__lowerCamelCase   :     List[str]		,	__lowerCamelCase   :     Optional[Any]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    str       =	{}
 import re

 __SCREAMING_SNAKE_CASE       :    str       =	re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Tuple       =	re.compile(
     r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Tuple       =	re.compile(r"encoders.(\d*).level_blocks.(\d*).model.(\d*).(bias|weight)"     )

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

 __SCREAMING_SNAKE_CASE       :    int       =	re.compile(r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    List[str]       =	re.compile(
     r"conditioner_blocks.(\d*).cond.model.(\d*).(\d).model.(\d*).model.(\d*).(bias|weight)"     )
 __SCREAMING_SNAKE_CASE       :    Dict       =	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(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	re_encoder_block_conv_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    List[str]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	int(groups[2]     ) * 2 + int(groups[3]     )
   __SCREAMING_SNAKE_CASE       :    int       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_encoder_block_conv_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_encoder_block_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Dict       =	re_encoder_block_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    List[str]       =	int(groups[2]     ) * 2 + int(groups[3]     )
   __SCREAMING_SNAKE_CASE       :    Any       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.downsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    List[str]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    List[str]       =	re_encoder_block_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_encoder_block_proj_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Any       =	re_encoder_block_proj_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    int       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	F"""encoders.{groups[0]}.level_blocks.{groups[1]}.proj_out.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Dict       =	re_encoder_block_proj_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # rename vqvae.decoder keys
  elif re_decoder_block_conv_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[str]       =	re_decoder_block_conv_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    str       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	int(groups[2]     ) * 2 + int(groups[3]     ) - 2
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	re_decoder_block_conv_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_decoder_block_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_decoder_block_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    int       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	int(groups[2]     ) * 2 + int(groups[3]     ) - 2
   __SCREAMING_SNAKE_CASE       :    Any       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Tuple       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.upsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_decoder_block_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_decoder_block_proj_in.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	re_decoder_block_proj_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    List[str]       =	F"""decoders.{groups[0]}.level_blocks.{groups[1]}.proj_in.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    int       =	re_decoder_block_proj_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # rename prior cond.model to upsampler.upsample_block and resnet
  elif re_prior_cond_conv_out.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    List[Any]       =	re_prior_cond_conv_out.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    str       =	int(groups[1]     ) * 2 + int(groups[2]     ) - 2
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""conditioner_blocks.upsampler.upsample_block.{block_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	re_prior_cond_conv_out.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_prior_cond_resnet.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_prior_cond_resnet.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Optional[int]       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Any       =	int(groups[1]     ) * 2 + int(groups[2]     ) - 2
   __SCREAMING_SNAKE_CASE       :    str       =	{"1": 1, "3": 2}[groups[-2]]
   __SCREAMING_SNAKE_CASE       :    Any       =	F"""conditioner_blocks.upsampler.upsample_block.{block_index}."""
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	F"""resnet_block.{groups[-3]}.conv1d_{conv_index}.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    Any       =	prefix + resnet_block
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	re_prior_cond_resnet.sub(__lowerCamelCase		,	__lowerCamelCase     )

  elif re_prior_cond_proj_in.fullmatch(__lowerCamelCase     ):
   __SCREAMING_SNAKE_CASE       :    Tuple       =	re_prior_cond_proj_in.match(__lowerCamelCase     )
   __SCREAMING_SNAKE_CASE       :    Dict       =	regex_match.groups()
   __SCREAMING_SNAKE_CASE       :    Dict       =	F"""conditioner_blocks.upsampler.proj_in.{groups[-1]}"""
   __SCREAMING_SNAKE_CASE       :    int       =	re_prior_cond_proj_in.sub(__lowerCamelCase		,	__lowerCamelCase     )

  # keep original key
  else:
   __SCREAMING_SNAKE_CASE       :    Tuple       =	original_key

  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	replace_key(__lowerCamelCase     )

  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:
   __SCREAMING_SNAKE_CASE       :    List[str]       =	model_state_dict[F"""{key_prefix}.{key}"""]
   print(F"""{original_key}-> {key} : \nshape {val.shape} and { value.shape}, do not match"""     )
   __SCREAMING_SNAKE_CASE       :    str       =	original_key

  __SCREAMING_SNAKE_CASE       :    List[str]       =	original_key
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	value

 return new_dict






@torch.no_grad()
def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]=None		,	__lowerCamelCase   :     List[Any]=None     ):



 """simple docstring"""





 for file in MODEL_MAPPING[model_name]:
  if not os.path.isfile(F"""{pytorch_dump_folder_path}/{file.split("/"     )[-1]}"""     ):
   __SCREAMING_SNAKE_CASE       :    Dict       =	requests.get(F"""{PREFIX}{file}"""		,	allow_redirects=__lowerCamelCase     )
   os.makedirs(F"""{pytorch_dump_folder_path}/"""		,	exist_ok=__lowerCamelCase     )
   open(F"""{pytorch_dump_folder_path}/{file.split("/"     )[-1]}"""		,	"wb"     ).write(r.content     )

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

 __SCREAMING_SNAKE_CASE       :    List[str]       =	JukeboxConfig.from_pretrained(__lowerCamelCase     )
 __SCREAMING_SNAKE_CASE       :    str       =	JukeboxModel(__lowerCamelCase     )

 __SCREAMING_SNAKE_CASE       :    int       =	[]
 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	{}
 for i, dict_name in enumerate(__lowerCamelCase     ):
  __SCREAMING_SNAKE_CASE       :    str       =	torch.load(F"""{pytorch_dump_folder_path}/{dict_name.split("/"     )[-1]}"""     )["model"]

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

  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	"vqvae" if i == 0 else F"""priors.{3 - i}"""
  __SCREAMING_SNAKE_CASE       :    int       =	fix_jukebox_keys(__lowerCamelCase		,	model.state_dict()		,	__lowerCamelCase		,	__lowerCamelCase     )
  weight_dict.append(__lowerCamelCase     )

 __SCREAMING_SNAKE_CASE       :    Dict       =	weight_dict.pop(0     )
 model.vqvae.load_state_dict(__lowerCamelCase     )
 for i in range(len(__lowerCamelCase     )     ):
  model.priors[i].load_state_dict(weight_dict[2 - i]     )

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

 print(F"""Saving model {model_name} to {pytorch_dump_folder_path}"""     )
 model.save_pretrained(__lowerCamelCase     )

 return weight_dict


if __name__ == "__main__":
      _lowerCamelCase						    =  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.""",
      )
      _lowerCamelCase						    =  parser.parse_args()
      convert_openai_checkpoint(args.model_name, args.pytorch_dump_folder_path)






