





"""simple docstring"""


import argparse
from pathlib import Path
from typing import Dict, OrderedDict, Tuple

import torch
from audiocraft.models import MusicGen

from transformers import (
    AutoFeatureExtractor,
    AutoTokenizer,
    EncodecModel,
    MusicgenDecoderConfig,
    MusicgenForConditionalGeneration,
    MusicgenProcessor,
    TaEncoderModel,
)
from transformers.models.musicgen.modeling_musicgen import MusicgenForCausalLM
from transformers.utils import logging


logging.set_verbosity_info()
lowerCamelCase__				  =   logging.get_logger(__name__)


lowerCamelCase__				  =   ["model.decoder.embed_positions.weights"]




def 					lowercase__    (					lowercase_			)					->		Optional[Any]:




      """simple docstring"""

      if "emb" in name:
            _UpperCamelCase      :	List[str]						=       name.replace("emb"		,"model.decoder.embed_tokens"			)
      if "transformer" in name:
            _UpperCamelCase      :	Optional[int]						=       name.replace("transformer"		,"model.decoder"			)
      if "cross_attention" in name:
            _UpperCamelCase      :	Optional[int]						=       name.replace("cross_attention"		,"encoder_attn"			)
      if "linear1" in name:
            _UpperCamelCase      :	Optional[Any]						=       name.replace("linear1"		,"fc1"			)
      if "linear2" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("linear2"		,"fc2"			)
      if "norm1" in name:
            _UpperCamelCase      :	Optional[Any]						=       name.replace("norm1"		,"self_attn_layer_norm"			)
      if "norm_cross" in name:
            _UpperCamelCase      :	Dict						=       name.replace("norm_cross"		,"encoder_attn_layer_norm"			)
      if "norm2" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("norm2"		,"final_layer_norm"			)
      if "out_norm" in name:
            _UpperCamelCase      :	Union[str, Any]						=       name.replace("out_norm"		,"model.decoder.layer_norm"			)
      if "linears" in name:
            _UpperCamelCase      :	List[str]						=       name.replace("linears"		,"lm_heads"			)
      if "condition_provider.conditioners.description.output_proj" in name:
            _UpperCamelCase      :	Any						=       name.replace("condition_provider.conditioners.description.output_proj"		,"enc_to_dec_proj"			)
      return name




def 					lowercase__    (					lowercase_		,lowercase_			)					->		Tuple[Dict, Dict]:




      """simple docstring"""

      _UpperCamelCase      :	str						=       list(state_dict.keys()			)
      _UpperCamelCase      :	Optional[Any]						=       {}
      for key in keys:
            _UpperCamelCase      :	Optional[int]						=       state_dict.pop(lowercase_			)
            _UpperCamelCase      :	List[Any]						=       rename_keys(lowercase_			)
            if "in_proj_weight" in key:
                  # split fused qkv proj
                  _UpperCamelCase      :	Tuple						=       val[:hidden_size, :]
                  _UpperCamelCase      :	Optional[Any]						=       val[hidden_size : 2 * hidden_size, :]
                  _UpperCamelCase      :	Optional[Any]						=       val[-hidden_size:, :]
            elif "enc_to_dec_proj" in key:
                  _UpperCamelCase      :	Optional[Any]						=       val
            else:
                  _UpperCamelCase      :	List[str]						=       val
      return state_dict, enc_dec_proj_state_dict




def 					lowercase__    (					lowercase_			)					->		MusicgenDecoderConfig:




      """simple docstring"""

      if checkpoint == "small":
            # default config values
            _UpperCamelCase      :	List[Any]						=       1_024
            _UpperCamelCase      :	List[str]						=       24
            _UpperCamelCase      :	Any						=       16
      elif checkpoint == "medium":
            _UpperCamelCase      :	Tuple						=       1_536
            _UpperCamelCase      :	Dict						=       48
            _UpperCamelCase      :	Tuple						=       24
      elif checkpoint == "large":
            _UpperCamelCase      :	int						=       2_048
            _UpperCamelCase      :	Optional[int]						=       48
            _UpperCamelCase      :	Dict						=       32
      else:
            raise ValueError(F'''Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'''			)
      _UpperCamelCase      :	str						=       MusicgenDecoderConfig(
          hidden_size=lowercase_		,ffn_dim=hidden_size * 4		,num_hidden_layers=lowercase_		,num_attention_heads=lowercase_		,)
      return config




@torch.no_grad()
def 					lowercase__    (					lowercase_		,lowercase_=None		,lowercase_=None		,lowercase_="cpu"			)					->		List[str]:




      """simple docstring"""

      _UpperCamelCase      :	str						=       MusicGen.get_pretrained(lowercase_		,device=lowercase_			)
      _UpperCamelCase      :	Union[str, Any]						=       decoder_config_from_checkpoint(lowercase_			)

      _UpperCamelCase      :	Optional[int]						=       fairseq_model.lm.state_dict()
      _UpperCamelCase,  _UpperCamelCase      :	Optional[Any]						=       rename_state_dict(
          lowercase_		,hidden_size=decoder_config.hidden_size			)

      _UpperCamelCase      :	Tuple						=       TaEncoderModel.from_pretrained("t5-base"			)
      _UpperCamelCase      :	Union[str, Any]						=       EncodecModel.from_pretrained("facebook/encodec_32khz"			)
      _UpperCamelCase      :	str						=       MusicgenForCausalLM(lowercase_			).eval()

      # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
      _UpperCamelCase,  _UpperCamelCase      :	str						=       decoder.load_state_dict(lowercase_		,strict=lowercase_			)

      for key in missing_keys.copy():
            if key.startswith(("text_encoder", "audio_encoder")			) or key in EXPECTED_MISSING_KEYS:
                  missing_keys.remove(lowercase_			)

      if len(lowercase_			) > 0:
            raise ValueError(F'''Missing key(s) in state_dict: {missing_keys}'''			)

      if len(lowercase_			) > 0:
            raise ValueError(F'''Unexpected key(s) in state_dict: {unexpected_keys}'''			)

      # init the composite model
      _UpperCamelCase      :	str						=       MusicgenForConditionalGeneration(text_encoder=lowercase_		,audio_encoder=lowercase_		,decoder=lowercase_			)

      # load the pre-trained enc-dec projection (from the decoder state dict)
      model.enc_to_dec_proj.load_state_dict(lowercase_			)

      # check we can do a forward pass
      _UpperCamelCase      :	List[str]						=       torch.arange(0		,8		,dtype=torch.long			).reshape(2		,-1			)
      _UpperCamelCase      :	Dict						=       input_ids.reshape(2 * 4		,-1			)

      with torch.no_grad():
            _UpperCamelCase      :	Tuple						=       model(input_ids=lowercase_		,decoder_input_ids=lowercase_			).logits

      if logits.shape != (8, 1, 2_048):
            raise ValueError("Incorrect shape for logits"			)

      # now construct the processor
      _UpperCamelCase      :	int						=       AutoTokenizer.from_pretrained("t5-base"			)
      _UpperCamelCase      :	str						=       AutoFeatureExtractor.from_pretrained("facebook/encodec_32khz"		,padding_side="left"			)

      _UpperCamelCase      :	Optional[int]						=       MusicgenProcessor(feature_extractor=lowercase_		,tokenizer=lowercase_			)

      # set the appropriate bos/pad token ids
      _UpperCamelCase      :	str						=       2_048
      _UpperCamelCase      :	str						=       2_048

      # set other default generation config params
      _UpperCamelCase      :	Optional[Any]						=       int(30 * audio_encoder.config.frame_rate			)
      _UpperCamelCase      :	List[str]						=       True
      _UpperCamelCase      :	int						=       3.0

      if pytorch_dump_folder is not None:
            Path(lowercase_			).mkdir(exist_ok=lowercase_			)
            logger.info(F'''Saving model {checkpoint} to {pytorch_dump_folder}'''			)
            model.save_pretrained(lowercase_			)
            processor.save_pretrained(lowercase_			)

      if repo_id:
            logger.info(F'''Pushing model {checkpoint} to {repo_id}'''			)
            model.push_to_hub(lowercase_			)
            processor.push_to_hub(lowercase_			)


if __name__ == "__main__":
  lowerCamelCase__				  =   argparse.ArgumentParser()
  # Required parameters
  parser.add_argument(
      "--checkpoint",
      default="small",
      type=str,
      help="Checkpoint size of the MusicGen model you'd like to convert. Can be one of: `['small', 'medium', 'large']`.",
  )
  parser.add_argument(
      "--pytorch_dump_folder",
      required=True,
      default=None,
      type=str,
      help="Path to the output PyTorch model directory.",
  )
  parser.add_argument(
      "--push_to_hub", default=None, type=str, help="Where to upload the converted model on the 🤗 hub."
  )
  parser.add_argument(
      "--device", default="cpu", type=str, help="Torch device to run the conversion, either cpu or cuda."
  )

  lowerCamelCase__				  =   parser.parse_args()
  convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)







"""simple docstring"""


import argparse
import collections
import os
import re

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
lowerCamelCase__				  =   "src/transformers"
lowerCamelCase__				  =   "docs/source/en"
lowerCamelCase__				  =   "."




def 					lowercase__    (					lowercase_		,lowercase_		,lowercase_			)					->		List[str]:




      """simple docstring"""

      with open(lowercase_		,"r"		,encoding="utf-8"		,newline="\n"			) as f:
            _UpperCamelCase      :	Union[str, Any]						=       f.readlines()
      # Find the start prompt.
      _UpperCamelCase      :	Dict						=       0
      while not lines[start_index].startswith(lowercase_			):
            start_index += 1
      start_index += 1

      _UpperCamelCase      :	Optional[int]						=       start_index
      while not lines[end_index].startswith(lowercase_			):
            end_index += 1
      end_index -= 1

      while len(lines[start_index]			) <= 1:
            start_index += 1
      while len(lines[end_index]			) <= 1:
            end_index -= 1
      end_index += 1
      return "".join(lines[start_index:end_index]			), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
lowerCamelCase__				  =   "Model|Encoder|Decoder|ForConditionalGeneration"
# Regexes that match TF/Flax/PT model names.
lowerCamelCase__				  =   re.compile(R"TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
lowerCamelCase__				  =   re.compile(R"Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
lowerCamelCase__				  =   re.compile(R"(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)")


# This is to make sure the transformers module imported is the one in the repo.
lowerCamelCase__				  =   direct_transformers_import(TRANSFORMERS_PATH)




def 					lowercase__    (					lowercase_			)					->		Any:




      """simple docstring"""

      _UpperCamelCase      :	Tuple						=       re.finditer(".+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)"		,lowercase_			)
      return [m.group(0			) for m in matches]




def 					lowercase__    (					lowercase_		,lowercase_			)					->		Union[str, Any]:




      """simple docstring"""

      _UpperCamelCase      :	str						=       2 if text == "✅" or text == "❌" else len(lowercase_			)
      _UpperCamelCase      :	Union[str, Any]						=       (width - text_length) // 2
      _UpperCamelCase      :	Dict						=       width - text_length - left_indent
      return " " * left_indent + text + " " * right_indent




def 					lowercase__    (					)					->		str:




      """simple docstring"""

      _UpperCamelCase      :	Optional[Any]						=       transformers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
      _UpperCamelCase      :	str						=       {
          name: config_maping_names[code]
          for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
          if code in config_maping_names
      }
      _UpperCamelCase      :	Dict						=       {name: config.replace("Config"		,""			) for name, config in model_name_to_config.items()}

      # Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
      _UpperCamelCase      :	int						=       collections.defaultdict(lowercase_			)
      _UpperCamelCase      :	Dict						=       collections.defaultdict(lowercase_			)
      _UpperCamelCase      :	Dict						=       collections.defaultdict(lowercase_			)
      _UpperCamelCase      :	int						=       collections.defaultdict(lowercase_			)
      _UpperCamelCase      :	str						=       collections.defaultdict(lowercase_			)

      # Let's lookup through all transformers object (once).
      for attr_name in dir(lowercase_			):
            _UpperCamelCase      :	List[str]						=       None
            if attr_name.endswith("Tokenizer"			):
                  _UpperCamelCase      :	Tuple						=       slow_tokenizers
                  _UpperCamelCase      :	Any						=       attr_name[:-9]
            elif attr_name.endswith("TokenizerFast"			):
                  _UpperCamelCase      :	Optional[Any]						=       fast_tokenizers
                  _UpperCamelCase      :	List[str]						=       attr_name[:-13]
            elif _re_tf_models.match(lowercase_			) is not None:
                  _UpperCamelCase      :	List[Any]						=       tf_models
                  _UpperCamelCase      :	Dict						=       _re_tf_models.match(lowercase_			).groups()[0]
            elif _re_flax_models.match(lowercase_			) is not None:
                  _UpperCamelCase      :	Dict						=       flax_models
                  _UpperCamelCase      :	Union[str, Any]						=       _re_flax_models.match(lowercase_			).groups()[0]
            elif _re_pt_models.match(lowercase_			) is not None:
                  _UpperCamelCase      :	Optional[int]						=       pt_models
                  _UpperCamelCase      :	Any						=       _re_pt_models.match(lowercase_			).groups()[0]

            if lookup_dict is not None:
                  while len(lowercase_			) > 0:
                        if attr_name in model_name_to_prefix.values():
                              _UpperCamelCase      :	Dict						=       True
                              break
                        # Try again after removing the last word in the name
                        _UpperCamelCase      :	List[str]						=       "".join(camel_case_split(lowercase_			)[:-1]			)

    # Let's build that table!
      _UpperCamelCase      :	Any						=       list(model_name_to_config.keys()			)
      model_names.sort(key=str.lower			)
      _UpperCamelCase      :	List[str]						=       ["Model", "Tokenizer slow", "Tokenizer fast", "PyTorch support", "TensorFlow support", "Flax Support"]
      # We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
      _UpperCamelCase      :	Union[str, Any]						=       [len(lowercase_			) + 2 for c in columns]
      _UpperCamelCase      :	Any						=       max([len(lowercase_			) for name in model_names]			) + 2

      # Build the table per se
      _UpperCamelCase      :	Tuple						=       "|" + "|".join([_center_text(lowercase_		,lowercase_			) for c, w in zip(lowercase_		,lowercase_			)]			) + "|\n"
      # Use ":-----:" format to center-aligned table cell texts
      table += "|" + "|".join([":" + "-" * (w - 2) + ":" for w in widths]			) + "|\n"

      _UpperCamelCase      :	Union[str, Any]						=       {True: "✅", False: "❌"}
      for name in model_names:
            _UpperCamelCase      :	Optional[int]						=       model_name_to_prefix[name]
            _UpperCamelCase      :	Tuple						=       [
                name,
                check[slow_tokenizers[prefix]],
                check[fast_tokenizers[prefix]],
                check[pt_models[prefix]],
                check[tf_models[prefix]],
                check[flax_models[prefix]],
            ]
            table += "|" + "|".join([_center_text(lowercase_		,lowercase_			) for l, w in zip(lowercase_		,lowercase_			)]			) + "|\n"
      return table




def 					lowercase__    (					lowercase_=False			)					->		List[Any]:




      """simple docstring"""

      _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase,  _UpperCamelCase      :	str						=       _find_text_in_file(
          filename=os.path.join(lowercase_		,"index.md"			)		,start_prompt="<!--This table is updated automatically from the auto modules"		,end_prompt="<!-- End table-->"		,)
      _UpperCamelCase      :	Any						=       get_model_table_from_auto_modules()

      if current_table != new_table:
            if overwrite:
                  with open(os.path.join(lowercase_		,"index.md"			)		,"w"		,encoding="utf-8"		,newline="\n"			) as f:
                        f.writelines(lines[:start_index] + [new_table] + lines[end_index:]			)
            else:
                  raise ValueError(
                      "The model table in the `index.md` has not been updated. Run `make fix-copies` to fix this."			)


if __name__ == "__main__":
  lowerCamelCase__				  =   argparse.ArgumentParser()
  parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
  lowerCamelCase__				  =   parser.parse_args()

  check_model_table(args.fix_and_overwrite)
