


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


lowerCamelCase							:   List[Any]  						=	['model.decoder.embed_positions.weights']

def _SCREAMING_SNAKE_CASE      (A       )       ->       Any:






       """simple docstring"""







       if "emb" in name:
              lowercase__      			= name.replace('''emb'''  , '''model.decoder.embed_tokens'''       )
       if "transformer" in name:
              lowercase__      			= name.replace('''transformer'''  , '''model.decoder'''       )
       if "cross_attention" in name:
              lowercase__      			= name.replace('''cross_attention'''  , '''encoder_attn'''       )
       if "linear1" in name:
              lowercase__      			= name.replace('''linear1'''  , '''fc1'''       )
       if "linear2" in name:
              lowercase__      			= name.replace('''linear2'''  , '''fc2'''       )
       if "norm1" in name:
              lowercase__      			= name.replace('''norm1'''  , '''self_attn_layer_norm'''       )
       if "norm_cross" in name:
              lowercase__      			= name.replace('''norm_cross'''  , '''encoder_attn_layer_norm'''       )
       if "norm2" in name:
              lowercase__      			= name.replace('''norm2'''  , '''final_layer_norm'''       )
       if "out_norm" in name:
              lowercase__      			= name.replace('''out_norm'''  , '''model.decoder.layer_norm'''       )
       if "linears" in name:
              lowercase__      			= name.replace('''linears'''  , '''lm_heads'''       )
       if "condition_provider.conditioners.description.output_proj" in name:
              lowercase__      			= name.replace('''condition_provider.conditioners.description.output_proj'''  , '''enc_to_dec_proj'''       )
       return name

def _SCREAMING_SNAKE_CASE      (A  , A       )       ->       Tuple[Dict, Dict]:






       """simple docstring"""







       lowercase__      			= list(state_dict.keys()       )
       lowercase__      			= {}
       for key in keys:
              lowercase__      			= state_dict.pop(A       )
              lowercase__      			= rename_keys(A       )
              if "in_proj_weight" in key:
                     # split fused qkv proj
                     lowercase__      			= val[:hidden_size, :]
                     lowercase__      			= val[hidden_size : 2 * hidden_size, :]
                     lowercase__      			= val[-hidden_size:, :]
              elif "enc_to_dec_proj" in key:
                     lowercase__      			= val
              else:
                     lowercase__      			= val
       return state_dict, enc_dec_proj_state_dict

def _SCREAMING_SNAKE_CASE      (A       )       ->       MusicgenDecoderConfig:






       """simple docstring"""







       if checkpoint == "small":
              # default config values
              lowercase__      			= 1_024
              lowercase__      			= 24
              lowercase__      			= 16
       elif checkpoint == "medium":
              lowercase__      			= 1_536
              lowercase__      			= 48
              lowercase__      			= 24
       elif checkpoint == "large":
              lowercase__      			= 2_048
              lowercase__      			= 48
              lowercase__      			= 32
       else:
              raise ValueError(f"Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."       )
       lowercase__      			= MusicgenDecoderConfig(
           hidden_size=A  , ffn_dim=hidden_size * 4  , num_hidden_layers=A  , num_attention_heads=A  , )
       return config

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






       """simple docstring"""







       lowercase__      			= MusicGen.get_pretrained(A  , device=A       )
       lowercase__      			= decoder_config_from_checkpoint(A       )

       lowercase__      			= fairseq_model.lm.state_dict()
       lowercase__			,lowercase__      			= rename_state_dict(
           A  , hidden_size=decoder_config.hidden_size       )

       lowercase__      			= TaEncoderModel.from_pretrained('''t5-base'''       )
       lowercase__      			= EncodecModel.from_pretrained('''facebook/encodec_32khz'''       )
       lowercase__      			= MusicgenForCausalLM(A       ).eval()

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

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

       if len(A       ) > 0:
              raise ValueError(f"Missing key(s) in state_dict: {missing_keys}"       )

       if len(A       ) > 0:
              raise ValueError(f"Unexpected key(s) in state_dict: {unexpected_keys}"       )

       # init the composite model
       lowercase__      			= MusicgenForConditionalGeneration(text_encoder=A  , audio_encoder=A  , decoder=A       )

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

       # check we can do a forward pass
       lowercase__      			= torch.arange(0  , 8  , dtype=torch.long       ).reshape(2  , -1       )
       lowercase__      			= input_ids.reshape(2 * 4  , -1       )

       with torch.no_grad():
              lowercase__      			= model(input_ids=A  , decoder_input_ids=A       ).logits

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

       # now construct the processor
       lowercase__      			= AutoTokenizer.from_pretrained('''t5-base'''       )
       lowercase__      			= AutoFeatureExtractor.from_pretrained('''facebook/encodec_32khz'''  , padding_side='''left'''       )

       lowercase__      			= MusicgenProcessor(feature_extractor=A  , tokenizer=A       )

       # set the appropriate bos/pad token ids
       lowercase__      			= 2_048
       lowercase__      			= 2_048

       # set other default generation config params
       lowercase__      			= int(30 * audio_encoder.config.frame_rate       )
       lowercase__      			= True
       lowercase__      			= 3.0

       if pytorch_dump_folder is not None:
              Path(A       ).mkdir(exist_ok=A       )
              logger.info(f"Saving model {checkpoint} to {pytorch_dump_folder}"       )
              model.save_pretrained(A       )
              processor.save_pretrained(A       )

       if repo_id:
              logger.info(f"Pushing model {checkpoint} to {repo_id}"       )
              model.push_to_hub(A       )
              processor.push_to_hub(A       )


if __name__ == "__main__":
 lowerCamelCase							:   Union[str, Any]  						=	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							:   List[str]  						=	parser.parse_args()
 convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)








'''simple docstring'''

import os
import re
import warnings
from shutil import copyfile
from typing import List, Optional, Tuple

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
 from .tokenization_ta import TaTokenizer
else:
 lowerCamelCase							:   Any  						=	None


lowerCamelCase							:   Dict  						=	logging.get_logger(__name__)

lowerCamelCase							:   Optional[int]  						=	{'vocab_file': 'spiece.model', 'tokenizer_file': 'tokenizer.json'}

lowerCamelCase							:   List[str]  						=	{
    'vocab_file': {
        't5-small': 'https://huggingface.co/t5-small/resolve/main/spiece.model',
        't5-base': 'https://huggingface.co/t5-base/resolve/main/spiece.model',
        't5-large': 'https://huggingface.co/t5-large/resolve/main/spiece.model',
        't5-3b': 'https://huggingface.co/t5-3b/resolve/main/spiece.model',
        't5-11b': 'https://huggingface.co/t5-11b/resolve/main/spiece.model',
    },
    'tokenizer_file': {
        't5-small': 'https://huggingface.co/t5-small/resolve/main/tokenizer.json',
        't5-base': 'https://huggingface.co/t5-base/resolve/main/tokenizer.json',
        't5-large': 'https://huggingface.co/t5-large/resolve/main/tokenizer.json',
        't5-3b': 'https://huggingface.co/t5-3b/resolve/main/tokenizer.json',
        't5-11b': 'https://huggingface.co/t5-11b/resolve/main/tokenizer.json',
    },
}


# TODO(PVP) - this should be removed in Transformers v5
lowerCamelCase							:   Any  						=	{
    't5-small': 512,
    't5-base': 512,
    't5-large': 512,
    't5-3b': 512,
    't5-11b': 512,
}
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






      lowerCAmelCase__					:     List[str]  					=							VOCAB_FILES_NAMES
      lowerCAmelCase__					:     str  					=							PRETRAINED_VOCAB_FILES_MAP
      lowerCAmelCase__					:     Optional[Any]  					=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      lowerCAmelCase__					:     int  					=							["""input_ids""", """attention_mask"""]
      lowerCAmelCase__					:     Optional[int]  					=							TaTokenizer

      lowerCAmelCase__					:     List[int]  					=							[]



      def __init__(self    :    Dict   ,						UpperCamelCase    :    str=None   ,						UpperCamelCase    :    Optional[Any]=None   ,						UpperCamelCase    :    Any="</s>"   ,						UpperCamelCase    :    str="<unk>"   ,						UpperCamelCase    :    List[str]="<pad>"   ,						UpperCamelCase    :    List[str]=100   ,						UpperCamelCase    :    Tuple=None   ,						**UpperCamelCase    :    List[str]   ,						):

             '''simple docstring'''







             if extra_ids > 0 and additional_special_tokens is None:
                    lowercase__      			= [f"<extra_id_{i}>" for i in range(UpperCamelCase )]
             elif extra_ids > 0 and additional_special_tokens is not None:
                    # Check that we have the right number of extra special tokens
                    lowercase__      			= len(set(filter(lambda UpperCamelCase : bool('''extra_id_''' in str(UpperCamelCase ) )   ,						UpperCamelCase ) ) )
                    if extra_tokens != extra_ids:
                           raise ValueError(
                               f"Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are"
                               ''' provided to T5Tokenizer. In this case the additional_special_tokens must include the extra_ids'''
                               ''' tokens''' )

             super().__init__(
                 UpperCamelCase   ,						tokenizer_file=UpperCamelCase   ,						eos_token=UpperCamelCase   ,						unk_token=UpperCamelCase   ,						pad_token=UpperCamelCase   ,						extra_ids=UpperCamelCase   ,						additional_special_tokens=UpperCamelCase   ,						**UpperCamelCase   ,						)

             lowercase__      			= vocab_file
             lowercase__      			= False if not self.vocab_file else True
             lowercase__      			= extra_ids



      @staticmethod
      def       UpperCamelCase__		(UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[Any] ):

             '''simple docstring'''







             if pretrained_model_name_or_path in TaTokenizerFast.max_model_input_sizes:
                    lowercase__      			= TaTokenizerFast.max_model_input_sizes[pretrained_model_name_or_path]
                    if init_max_model_length is not None and init_max_model_length != max_model_length:
                           return init_max_model_length
                    elif init_max_model_length is None:
                           warnings.warn(
                               '''This tokenizer was incorrectly instantiated with a model max length of'''
                               f" {deprecated_max_model_length} which will be corrected in Transformers v5.\nFor now, this"
                               ''' behavior is kept to avoid breaking backwards compatibility when padding/encoding with'''
                               ''' `truncation is True`.\n- Be aware that you SHOULD NOT rely on'''
                               f" {pretrained_model_name_or_path} automatically truncating your input to"
                               f" {deprecated_max_model_length} when padding/encoding.\n- If you want to encode/pad to sequences"
                               f" longer than {deprecated_max_model_length} you can either instantiate this tokenizer with"
                               ''' `model_max_length` or pass `max_length` when encoding/padding.\n- To avoid this warning, please'''
                               ''' instantiate this tokenizer with `model_max_length` set to your preferred value.'''   ,						UpperCamelCase   ,						)

             return max_model_length



      def       UpperCamelCase__		(self    :    Any   ,						UpperCamelCase    :    str   ,						UpperCamelCase    :    Optional[str] = None ):

             '''simple docstring'''







             if not self.can_save_slow_tokenizer:
                    raise ValueError(
                        '''Your fast tokenizer does not have the necessary information to save the vocabulary for a slow '''
                        '''tokenizer.''' )

             if not os.path.isdir(UpperCamelCase ):
                    logger.error(f"Vocabulary path ({save_directory}) should be a directory" )
                    return
             lowercase__      			= os.path.join(
                 UpperCamelCase   ,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''] )

             if os.path.abspath(self.vocab_file ) != os.path.abspath(UpperCamelCase ):
                    copyfile(self.vocab_file   ,						UpperCamelCase )
                    logger.info(f"Copy vocab file to {out_vocab_file}" )

             return (out_vocab_file,)



      def       UpperCamelCase__		(self    :    Any   ,						UpperCamelCase    :    List[int]   ,						UpperCamelCase    :    Optional[List[int]] = None ):

             '''simple docstring'''







             lowercase__      			= token_ids_a + [self.eos_token_id]
             if token_ids_a is None:
                    return self.prefix_tokens + token_ids_a
             else:
                    lowercase__      			= token_ids_a + [self.eos_token_id]
                    return self.prefix_tokens + token_ids_a + token_ids_a



      def       UpperCamelCase__		(self    :    Optional[Any]   ,						UpperCamelCase    :    List[int]   ,						UpperCamelCase    :    Optional[List[int]] = None ):

             '''simple docstring'''







             lowercase__      			= [self.eos_token_id]

             if token_ids_a is None:
                    return len(token_ids_a + eos ) * [0]
             return len(token_ids_a + eos + token_ids_a + eos ) * [0]



      def       UpperCamelCase__		(self    :    List[Any] ):

             '''simple docstring'''







             return list(
                 set(filter(lambda UpperCamelCase : bool(re.search(R'''<extra_id_\d+>'''   ,						UpperCamelCase ) ) is not None   ,						self.additional_special_tokens ) ) )






      def       UpperCamelCase__		(self    :    Union[str, Any] ):

             '''simple docstring'''







             return [self.convert_tokens_to_ids(UpperCamelCase ) for token in self.get_sentinel_tokens()]




