


'''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()
__UpperCAmelCase      :Tuple			    =       logging.get_logger(__name__)


__UpperCAmelCase      :Optional[Any]			    =       ["model.decoder.embed_positions.weights"]




def      _a     (						_lowercase :     Optional[Any]  ):





     '''simple docstring'''






     if "emb" in name:
          __UpperCAmelCase     :			Optional[Any]  							=      name.replace('''emb'''      , '''model.decoder.embed_tokens'''  )
     if "transformer" in name:
          __UpperCAmelCase     :			Any  							=      name.replace('''transformer'''      , '''model.decoder'''  )
     if "cross_attention" in name:
          __UpperCAmelCase     :			List[str]  							=      name.replace('''cross_attention'''      , '''encoder_attn'''  )
     if "linear1" in name:
          __UpperCAmelCase     :			List[Any]  							=      name.replace('''linear1'''      , '''fc1'''  )
     if "linear2" in name:
          __UpperCAmelCase     :			Optional[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     :			str  							=      name.replace('''norm_cross'''      , '''encoder_attn_layer_norm'''  )
     if "norm2" in name:
          __UpperCAmelCase     :			Tuple  							=      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     :			str  							=      name.replace('''linears'''      , '''lm_heads'''  )
     if "condition_provider.conditioners.description.output_proj" in name:
          __UpperCAmelCase     :			str  							=      name.replace('''condition_provider.conditioners.description.output_proj'''      , '''enc_to_dec_proj'''  )
     return name




def      _a     (						_lowercase :     OrderedDict      , _lowercase :     int  ):





     '''simple docstring'''






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




def      _a     (						_lowercase :     str  ):





     '''simple docstring'''






     if checkpoint == "small":
          # default config values
          __UpperCAmelCase     :			Optional[int]  							=      1024
          __UpperCAmelCase     :			Tuple  							=      24
          __UpperCAmelCase     :			Optional[int]  							=      16
     elif checkpoint == "medium":
          __UpperCAmelCase     :			Optional[Any]  							=      1536
          __UpperCAmelCase     :			int  							=      48
          __UpperCAmelCase     :			List[Any]  							=      24
     elif checkpoint == "large":
          __UpperCAmelCase     :			List[str]  							=      2048
          __UpperCAmelCase     :			Union[str, Any]  							=      48
          __UpperCAmelCase     :			Optional[Any]  							=      32
     else:
          raise ValueError(F'Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'  )
     __UpperCAmelCase     :			Union[str, Any]  							=      MusicgenDecoderConfig(
         hidden_size=_lowercase      , ffn_dim=hidden_size * 4      , num_hidden_layers=_lowercase      , num_attention_heads=_lowercase      , )
     return config




@torch.no_grad()
def      _a     (						_lowercase :     Any      , _lowercase :     Optional[Any]=None      , _lowercase :     Optional[Any]=None      , _lowercase :     Dict="cpu"  ):





     '''simple docstring'''






     __UpperCAmelCase     :			str  							=      MusicGen.get_pretrained(_lowercase      , device=_lowercase  )
     __UpperCAmelCase     :			List[str]  							=      decoder_config_from_checkpoint(_lowercase  )

     __UpperCAmelCase     :			Tuple  							=      fairseq_model.lm.state_dict()
     __UpperCAmelCase     :			int  							=      rename_state_dict(
         _lowercase      , hidden_size=decoder_config.hidden_size  )

     __UpperCAmelCase     :			str  							=      TaEncoderModel.from_pretrained('''t5-base'''  )
     __UpperCAmelCase     :			Dict  							=      EncodecModel.from_pretrained('''facebook/encodec_32khz'''  )
     __UpperCAmelCase     :			Tuple  							=      MusicgenForCausalLM(_lowercase  ).eval()

     # load all decoder weights - expect that we'll be missing embeddings and enc-dec projection
     __UpperCAmelCase     :			int  							=      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     :			Any  							=      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     :			int  							=      torch.arange(0      , 8      , dtype=torch.long  ).reshape(2      , -1  )
     __UpperCAmelCase     :			Tuple  							=      input_ids.reshape(2 * 4      , -1  )

     with torch.no_grad():
          __UpperCAmelCase     :			Union[str, Any]  							=      model(input_ids=_lowercase      , decoder_input_ids=_lowercase  ).logits

     if logits.shape != (8, 1, 2048):
          raise ValueError('''Incorrect shape for logits'''  )

     # now construct the processor
     __UpperCAmelCase     :			Union[str, Any]  							=      AutoTokenizer.from_pretrained('''t5-base'''  )
     __UpperCAmelCase     :			Any  							=      AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''      , padding_side='''left'''  )

     __UpperCAmelCase     :			Optional[Any]  							=      MusicgenProcessor(feature_extractor=_lowercase      , tokenizer=_lowercase  )

     # set the appropriate bos/pad token ids
     __UpperCAmelCase     :			List[str]  							=      2048
     __UpperCAmelCase     :			Optional[Any]  							=      2048

     # set other default generation config params
     __UpperCAmelCase     :			int  							=      int(30 * audio_encoder.config.frame_rate  )
     __UpperCAmelCase     :			str  							=      True
     __UpperCAmelCase     :			List[Any]  							=      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__":
      __UpperCAmelCase      :Dict			    =       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."
      )

      __UpperCAmelCase      :int			    =       parser.parse_args()
      convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)




'''simple docstring'''



import json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_bert import BertTokenizer


__UpperCAmelCase      :Dict			    =       logging.get_logger(__name__)

__UpperCAmelCase      :str			    =       {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}

__UpperCAmelCase      :Optional[int]			    =       {
    "vocab_file": {
        "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt",
        "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/vocab.txt",
        "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/vocab.txt",
        "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/vocab.txt",
        "bert-base-multilingual-uncased": (
            "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/vocab.txt"
        ),
        "bert-base-multilingual-cased": "https://huggingface.co/bert-base-multilingual-cased/resolve/main/vocab.txt",
        "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/vocab.txt",
        "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/vocab.txt",
        "bert-large-uncased-whole-word-masking": (
            "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/vocab.txt"
        ),
        "bert-large-cased-whole-word-masking": (
            "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/vocab.txt"
        ),
        "bert-large-uncased-whole-word-masking-finetuned-squad": (
            "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt"
        ),
        "bert-large-cased-whole-word-masking-finetuned-squad": (
            "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/vocab.txt"
        ),
        "bert-base-cased-finetuned-mrpc": (
            "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/vocab.txt"
        ),
        "bert-base-german-dbmdz-cased": "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/vocab.txt",
        "bert-base-german-dbmdz-uncased": (
            "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/vocab.txt"
        ),
        "TurkuNLP/bert-base-finnish-cased-v1": (
            "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/vocab.txt"
        ),
        "TurkuNLP/bert-base-finnish-uncased-v1": (
            "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/vocab.txt"
        ),
        "wietsedv/bert-base-dutch-cased": (
            "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/vocab.txt"
        ),
    },
    "tokenizer_file": {
        "bert-base-uncased": "https://huggingface.co/bert-base-uncased/resolve/main/tokenizer.json",
        "bert-large-uncased": "https://huggingface.co/bert-large-uncased/resolve/main/tokenizer.json",
        "bert-base-cased": "https://huggingface.co/bert-base-cased/resolve/main/tokenizer.json",
        "bert-large-cased": "https://huggingface.co/bert-large-cased/resolve/main/tokenizer.json",
        "bert-base-multilingual-uncased": (
            "https://huggingface.co/bert-base-multilingual-uncased/resolve/main/tokenizer.json"
        ),
        "bert-base-multilingual-cased": (
            "https://huggingface.co/bert-base-multilingual-cased/resolve/main/tokenizer.json"
        ),
        "bert-base-chinese": "https://huggingface.co/bert-base-chinese/resolve/main/tokenizer.json",
        "bert-base-german-cased": "https://huggingface.co/bert-base-german-cased/resolve/main/tokenizer.json",
        "bert-large-uncased-whole-word-masking": (
            "https://huggingface.co/bert-large-uncased-whole-word-masking/resolve/main/tokenizer.json"
        ),
        "bert-large-cased-whole-word-masking": (
            "https://huggingface.co/bert-large-cased-whole-word-masking/resolve/main/tokenizer.json"
        ),
        "bert-large-uncased-whole-word-masking-finetuned-squad": (
            "https://huggingface.co/bert-large-uncased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json"
        ),
        "bert-large-cased-whole-word-masking-finetuned-squad": (
            "https://huggingface.co/bert-large-cased-whole-word-masking-finetuned-squad/resolve/main/tokenizer.json"
        ),
        "bert-base-cased-finetuned-mrpc": (
            "https://huggingface.co/bert-base-cased-finetuned-mrpc/resolve/main/tokenizer.json"
        ),
        "bert-base-german-dbmdz-cased": (
            "https://huggingface.co/bert-base-german-dbmdz-cased/resolve/main/tokenizer.json"
        ),
        "bert-base-german-dbmdz-uncased": (
            "https://huggingface.co/bert-base-german-dbmdz-uncased/resolve/main/tokenizer.json"
        ),
        "TurkuNLP/bert-base-finnish-cased-v1": (
            "https://huggingface.co/TurkuNLP/bert-base-finnish-cased-v1/resolve/main/tokenizer.json"
        ),
        "TurkuNLP/bert-base-finnish-uncased-v1": (
            "https://huggingface.co/TurkuNLP/bert-base-finnish-uncased-v1/resolve/main/tokenizer.json"
        ),
        "wietsedv/bert-base-dutch-cased": (
            "https://huggingface.co/wietsedv/bert-base-dutch-cased/resolve/main/tokenizer.json"
        ),
    },
}

__UpperCAmelCase      :Union[str, Any]			    =       {
    "bert-base-uncased": 5_1_2,
    "bert-large-uncased": 5_1_2,
    "bert-base-cased": 5_1_2,
    "bert-large-cased": 5_1_2,
    "bert-base-multilingual-uncased": 5_1_2,
    "bert-base-multilingual-cased": 5_1_2,
    "bert-base-chinese": 5_1_2,
    "bert-base-german-cased": 5_1_2,
    "bert-large-uncased-whole-word-masking": 5_1_2,
    "bert-large-cased-whole-word-masking": 5_1_2,
    "bert-large-uncased-whole-word-masking-finetuned-squad": 5_1_2,
    "bert-large-cased-whole-word-masking-finetuned-squad": 5_1_2,
    "bert-base-cased-finetuned-mrpc": 5_1_2,
    "bert-base-german-dbmdz-cased": 5_1_2,
    "bert-base-german-dbmdz-uncased": 5_1_2,
    "TurkuNLP/bert-base-finnish-cased-v1": 5_1_2,
    "TurkuNLP/bert-base-finnish-uncased-v1": 5_1_2,
    "wietsedv/bert-base-dutch-cased": 5_1_2,
}

__UpperCAmelCase      :str			    =       {
    "bert-base-uncased": {"do_lower_case": True},
    "bert-large-uncased": {"do_lower_case": True},
    "bert-base-cased": {"do_lower_case": False},
    "bert-large-cased": {"do_lower_case": False},
    "bert-base-multilingual-uncased": {"do_lower_case": True},
    "bert-base-multilingual-cased": {"do_lower_case": False},
    "bert-base-chinese": {"do_lower_case": False},
    "bert-base-german-cased": {"do_lower_case": False},
    "bert-large-uncased-whole-word-masking": {"do_lower_case": True},
    "bert-large-cased-whole-word-masking": {"do_lower_case": False},
    "bert-large-uncased-whole-word-masking-finetuned-squad": {"do_lower_case": True},
    "bert-large-cased-whole-word-masking-finetuned-squad": {"do_lower_case": False},
    "bert-base-cased-finetuned-mrpc": {"do_lower_case": False},
    "bert-base-german-dbmdz-cased": {"do_lower_case": False},
    "bert-base-german-dbmdz-uncased": {"do_lower_case": True},
    "TurkuNLP/bert-base-finnish-cased-v1": {"do_lower_case": False},
    "TurkuNLP/bert-base-finnish-uncased-v1": {"do_lower_case": True},
    "wietsedv/bert-base-dutch-cased": {"do_lower_case": False},
}



class      a    (		_a	):


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       VOCAB_FILES_NAMES
     SCREAMING_SNAKE_CASE	:			List[Any]							       =       PRETRAINED_VOCAB_FILES_MAP
     SCREAMING_SNAKE_CASE	:			Tuple							       =       PRETRAINED_INIT_CONFIGURATION
     SCREAMING_SNAKE_CASE	:			Optional[int]							       =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     SCREAMING_SNAKE_CASE	:			Optional[Any]							       =       BertTokenizer

     def __init__(    self					:		Optional[int]							,			snake_case					:		Union[str, Any]=None							,			snake_case					:		str=None							,			snake_case					:		Any=True							,			snake_case					:		Tuple="[UNK]"							,			snake_case					:		int="[SEP]"							,			snake_case					:		Optional[Any]="[PAD]"							,			snake_case					:		int="[CLS]"							,			snake_case					:		Optional[Any]="[MASK]"							,			snake_case					:		Union[str, Any]=True							,			snake_case					:		List[Any]=None							,			**snake_case					:		int							,			)						->							str:
          super().__init__(
              snake_case							,			tokenizer_file=snake_case							,			do_lower_case=snake_case							,			unk_token=snake_case							,			sep_token=snake_case							,			pad_token=snake_case							,			cls_token=snake_case							,			mask_token=snake_case							,			tokenize_chinese_chars=snake_case							,			strip_accents=snake_case							,			**snake_case							,			)

          __UpperCAmelCase     :			List[str]  							=      json.loads(self.backend_tokenizer.normalizer.__getstate__()      )
          if (
              normalizer_state.get('''lowercase'''							,			snake_case      ) != do_lower_case
              or normalizer_state.get('''strip_accents'''							,			snake_case      ) != strip_accents
              or normalizer_state.get('''handle_chinese_chars'''							,			snake_case      ) != tokenize_chinese_chars
          ):
               __UpperCAmelCase     :			List[Any]  							=      getattr(snake_case							,			normalizer_state.pop('''type'''      )      )
               __UpperCAmelCase     :			List[Any]  							=      do_lower_case
               __UpperCAmelCase     :			List[Any]  							=      strip_accents
               __UpperCAmelCase     :			str  							=      tokenize_chinese_chars
               __UpperCAmelCase     :			List[str]  							=      normalizer_class(**snake_case      )

          __UpperCAmelCase     :			Optional[int]  							=      do_lower_case

     def 			lowerCamelCase__			(    self					:		List[Any]							,			snake_case					:		Tuple							,			snake_case					:		List[str]=None      )						->							str:
          __UpperCAmelCase     :			Optional[int]  							=      [self.cls_token_id] + token_ids_a + [self.sep_token_id]

          if token_ids_a:
               output += token_ids_a + [self.sep_token_id]

          return output

     def 			lowerCamelCase__			(    self					:		str							,			snake_case					:		List[int]							,			snake_case					:		Optional[List[int]] = None      )						->							List[int]:
          __UpperCAmelCase     :			List[str]  							=      [self.sep_token_id]
          __UpperCAmelCase     :			Tuple  							=      [self.cls_token_id]
          if token_ids_a is None:
               return len(cls + token_ids_a + sep      ) * [0]
          return len(cls + token_ids_a + sep      ) * [0] + len(token_ids_a + sep      ) * [1]






     def 			lowerCamelCase__			(    self					:		List[Any]							,			snake_case					:		str							,			snake_case					:		Optional[str] = None      )						->							Tuple[str]:
          __UpperCAmelCase     :			Optional[int]  							=      self._tokenizer.model.save(snake_case							,			name=snake_case      )
          return tuple(snake_case      )