





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()
A       :   List[str]					= logging.get_logger(__name__)


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







def   __lowerCamelCase   (				__a  :List[str]       )						->							List[Any]:






		"""simple docstring"""

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







def   __lowerCamelCase   (				__a  :OrderedDict       ,	__a  :int       )						->							Tuple[Dict, Dict]:






		"""simple docstring"""

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







def   __lowerCamelCase   (				__a  :str       )						->							MusicgenDecoderConfig:






		"""simple docstring"""

		if checkpoint == "small":
				# default config values
				A__     =  1_0_2_4
				A__     =  2_4
				A__     =  1_6
		elif checkpoint == "medium":
				A__     =  1_5_3_6
				A__     =  4_8
				A__     =  2_4
		elif checkpoint == "large":
				A__     =  2_0_4_8
				A__     =  4_8
				A__     =  3_2
		else:
				raise ValueError(F'Checkpoint should be one of `[\'small\', \'medium\', \'large\']`, got {checkpoint}.'       )
		A__     =  MusicgenDecoderConfig(
		    hidden_size=_snake_case       ,	ffn_dim=hidden_size * 4       ,	num_hidden_layers=_snake_case       ,	num_attention_heads=_snake_case       ,	)
		return config







@torch.no_grad()
def   __lowerCamelCase   (				__a  :Optional[int]       ,	__a  :int=None       ,	__a  :int=None       ,	__a  :Optional[int]="cpu"       )						->							List[str]:






		"""simple docstring"""

		A__     =  MusicGen.get_pretrained(_snake_case       ,	device=_snake_case       )
		A__     =  decoder_config_from_checkpoint(_snake_case       )

		A__     =  fairseq_model.lm.state_dict()
		A__ ,     A__     =  rename_state_dict(
		    _snake_case       ,	hidden_size=decoder_config.hidden_size       )

		A__     =  TaEncoderModel.from_pretrained("""t5-base"""       )
		A__     =  EncodecModel.from_pretrained("""facebook/encodec_32khz"""       )
		A__     =  MusicgenForCausalLM(_snake_case       ).eval()

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

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

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

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

		# init the composite model
		A__     =  MusicgenForConditionalGeneration(text_encoder=_snake_case       ,	audio_encoder=_snake_case       ,	decoder=_snake_case       )

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

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

		with torch.no_grad():
				A__     =  model(input_ids=_snake_case       ,	decoder_input_ids=_snake_case       ).logits

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

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

		A__     =  MusicgenProcessor(feature_extractor=_snake_case       ,	tokenizer=_snake_case       )

		# set the appropriate bos/pad token ids
		A__     =  2_0_4_8
		A__     =  2_0_4_8

		# set other default generation config params
		A__     =  int(3_0 * audio_encoder.config.frame_rate       )
		A__     =  True
		A__     =  3.0

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

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


if __name__ == "__main__":
					A       :   List[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.'''
					)

					A       :   Dict					= parser.parse_args()
					convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)








import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


UpperCAmelCase_		 =					logging.get_logger(__name__)

UpperCAmelCase_		 =					"""▁"""

UpperCAmelCase_		 =					{"""vocab_file""": """sentencepiece.bpe.model""", """monolingual_vocab_file""": """dict.txt"""}

UpperCAmelCase_		 =					{
    """vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/sentencepiece.bpe.model""",
    },
    """monolingual_vocab_file""": {
        """vinai/bartpho-syllable""": """https://huggingface.co/vinai/bartpho-syllable/resolve/main/dict.txt""",
    },
}

UpperCAmelCase_		 =					{"""vinai/bartpho-syllable""": 1_0_2_4}


class 	lowerCamelCase__	(       _A):



 """simple docstring"""

 a__					: int =  VOCAB_FILES_NAMES
 a__					: Tuple =  PRETRAINED_VOCAB_FILES_MAP
 a__					: Optional[int] =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 a__					: Tuple =  ["input_ids", "attention_mask"]

 def __init__(    self					:					Union[str, Any]			,						__lowerCAmelCase					:					Dict			,						__lowerCAmelCase					:					List[str]			,						__lowerCAmelCase					:					Union[str, Any]="<s>"			,						__lowerCAmelCase					:					Dict="</s>"			,						__lowerCAmelCase					:					List[Any]="</s>"			,						__lowerCAmelCase					:					Optional[Any]="<s>"			,						__lowerCAmelCase					:					Tuple="<unk>"			,						__lowerCAmelCase					:					int="<pad>"			,						__lowerCAmelCase					:					Optional[Any]="<mask>"			,						__lowerCAmelCase					:					Optional[Dict[str, Any]] = None			,						**__lowerCAmelCase					:					Tuple			,						)   ->	None:
    # Mask token behave like a normal word, i.e. include the space before it
    _A      =       AddedToken(__lowerCAmelCase			,						lstrip=__lowerCAmelCase			,						rstrip=__lowerCAmelCase      ) if isinstance(__lowerCAmelCase			,						__lowerCAmelCase      ) else mask_token

    _A      =       {} if sp_model_kwargs is None else sp_model_kwargs

    super().__init__(
        bos_token=__lowerCAmelCase			,						eos_token=__lowerCAmelCase			,						unk_token=__lowerCAmelCase			,						sep_token=__lowerCAmelCase			,						cls_token=__lowerCAmelCase			,						pad_token=__lowerCAmelCase			,						mask_token=__lowerCAmelCase			,						sp_model_kwargs=self.sp_model_kwargs			,						**__lowerCAmelCase			,						)

    _A      =       vocab_file
    _A      =       monolingual_vocab_file
    _A      =       spm.SentencePieceProcessor(**self.sp_model_kwargs      )
    self.sp_model.Load(str(__lowerCAmelCase      )      )

    # Load the reduced vocab

    # Keep order of special tokens for backward compatibility
    _A      =       {}
    _A      =       0
    for token in [bos_token, pad_token, eos_token, unk_token, sep_token, cls_token]:
       if str(__lowerCAmelCase      ) not in self.fairseq_tokens_to_ids:
          _A      =       cnt
          cnt += 1
    with open(__lowerCAmelCase			,						'''r'''			,						encoding='''utf-8'''      ) as f:
       for line in f.readlines():
          _A      =       line.strip().split()[0]
          _A      =       len(self.fairseq_tokens_to_ids      )
    if str(__lowerCAmelCase      ) not in self.fairseq_tokens_to_ids:
       _A      =       len(self.fairseq_tokens_to_ids      )

    _A      =       {v: k for k, v in self.fairseq_tokens_to_ids.items()}

 def __getstate__(    self					:					Any      )   ->	List[Any]:
    _A      =       self.__dict__.copy()
    _A      =       None
    _A      =       self.sp_model.serialized_model_proto()
    return state

 def __setstate__(    self					:					Union[str, Any]			,						__lowerCAmelCase					:					Dict      )   ->	List[Any]:
    _A      =       d

    # for backward compatibility
    if not hasattr(self			,						'''sp_model_kwargs'''      ):
       _A      =       {}

    _A      =       spm.SentencePieceProcessor(**self.sp_model_kwargs      )
    self.sp_model.LoadFromSerializedProto(self.sp_model_proto      )

 def     snake_case_      (    self					:					Optional[Any]			,						__lowerCAmelCase					:					List[int]			,						__lowerCAmelCase					:					Optional[List[int]] = None      )   ->	List[int]:

    if token_ids_a is None:
       return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
    _A      =       [self.cls_token_id]
    _A      =       [self.sep_token_id]
    return cls + token_ids_a + sep + sep + token_ids_a + sep

 def     snake_case_      (    self					:					List[Any]			,						__lowerCAmelCase					:					List[int]			,						__lowerCAmelCase					:					Optional[List[int]] = None			,						__lowerCAmelCase					:					bool = False      )   ->	List[int]:

    if already_has_special_tokens:
       return super().get_special_tokens_mask(
           token_ids_a=__lowerCAmelCase			,						token_ids_a=__lowerCAmelCase			,						already_has_special_tokens=__lowerCAmelCase      )

    if token_ids_a is None:
       return [1] + ([0] * len(__lowerCAmelCase      )) + [1]
    return [1] + ([0] * len(__lowerCAmelCase      )) + [1, 1] + ([0] * len(__lowerCAmelCase      )) + [1]

 def     snake_case_      (    self					:					Any			,						__lowerCAmelCase					:					List[int]			,						__lowerCAmelCase					:					Optional[List[int]] = None      )   ->	List[int]:
    _A      =       [self.sep_token_id]
    _A      =       [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 + sep + token_ids_a + sep      ) * [0]

 @property
 def     snake_case_      (    self					:					Optional[int]      )   ->	Union[str, Any]:
    return len(self.fairseq_ids_to_tokens      )

 def     snake_case_      (    self					:					Dict      )   ->	Optional[Any]:
    _A      =       {self.convert_ids_to_tokens(__lowerCAmelCase      ): i for i in range(self.vocab_size      )}
    vocab.update(self.added_tokens_encoder      )
    return vocab

 def     snake_case_      (    self					:					List[str]			,						__lowerCAmelCase					:					str      )   ->	List[str]:
    return self.sp_model.encode(__lowerCAmelCase			,						out_type=__lowerCAmelCase      )

 def     snake_case_      (    self					:					str			,						__lowerCAmelCase					:					Optional[Any]      )   ->	Dict:
    if token in self.fairseq_tokens_to_ids:
       return self.fairseq_tokens_to_ids[token]
    else:
       return self.unk_token_id

 def     snake_case_      (    self					:					int			,						__lowerCAmelCase					:					Optional[int]      )   ->	List[str]:
    return self.fairseq_ids_to_tokens[index]

 def     snake_case_      (    self					:					List[str]			,						__lowerCAmelCase					:					Union[str, Any]      )   ->	Tuple:
    _A      =       ''''''.join(__lowerCAmelCase      ).replace(__lowerCAmelCase			,						''' '''      ).strip()
    return out_string

 def     snake_case_      (    self					:					str			,						__lowerCAmelCase					:					str			,						__lowerCAmelCase					:					Optional[str] = None      )   ->	Tuple[str]:
    if not os.path.isdir(__lowerCAmelCase      ):
       logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''      )
       return
    _A      =       os.path.join(
        __lowerCAmelCase			,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )
    _A      =       os.path.join(
        __lowerCAmelCase			,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''monolingual_vocab_file''']			,						)

    if os.path.abspath(self.vocab_file      ) != os.path.abspath(__lowerCAmelCase      ) and os.path.isfile(self.vocab_file      ):
       copyfile(self.vocab_file			,						__lowerCAmelCase      )
    elif not os.path.isfile(self.vocab_file      ):
       with open(__lowerCAmelCase			,						'''wb'''      ) as fi:
          _A      =       self.sp_model.serialized_model_proto()
          fi.write(__lowerCAmelCase      )

    if os.path.abspath(self.monolingual_vocab_file      ) != os.path.abspath(
        __lowerCAmelCase      ) and os.path.isfile(self.monolingual_vocab_file      ):
       copyfile(self.monolingual_vocab_file			,						__lowerCAmelCase      )
    elif not os.path.isfile(self.monolingual_vocab_file      ):
       with open(__lowerCAmelCase			,						'''w'''			,						encoding='''utf-8'''      ) as fp:
          for token in self.fairseq_tokens_to_ids:
             if token not in self.all_special_tokens:
                fp.write(f'''{str(__lowerCAmelCase      )} \n'''      )

    return out_vocab_file, out_monolingual_vocab_file



