


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


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







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				Union[str, Any]:
  if "emb" in name:
    snake_case__						= name.replace('''emb'''							,   '''model.decoder.embed_tokens'''				)
  if "transformer" in name:
    snake_case__						= name.replace('''transformer'''							,   '''model.decoder'''				)
  if "cross_attention" in name:
    snake_case__						= name.replace('''cross_attention'''							,   '''encoder_attn'''				)
  if "linear1" in name:
    snake_case__						= name.replace('''linear1'''							,   '''fc1'''				)
  if "linear2" in name:
    snake_case__						= name.replace('''linear2'''							,   '''fc2'''				)
  if "norm1" in name:
    snake_case__						= name.replace('''norm1'''							,   '''self_attn_layer_norm'''				)
  if "norm_cross" in name:
    snake_case__						= name.replace('''norm_cross'''							,   '''encoder_attn_layer_norm'''				)
  if "norm2" in name:
    snake_case__						= name.replace('''norm2'''							,   '''final_layer_norm'''				)
  if "out_norm" in name:
    snake_case__						= name.replace('''out_norm'''							,   '''model.decoder.layer_norm'''				)
  if "linears" in name:
    snake_case__						= name.replace('''linears'''							,   '''lm_heads'''				)
  if "condition_provider.conditioners.description.output_proj" in name:
    snake_case__						= name.replace('''condition_provider.conditioners.description.output_proj'''							,   '''enc_to_dec_proj'''				)
  return name







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







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase				)    ->				MusicgenDecoderConfig:
  if checkpoint == "small":
    # default config values
    snake_case__						= 1024
    snake_case__						= 24
    snake_case__						= 16
  elif checkpoint == "medium":
    snake_case__						= 1536
    snake_case__						= 48
    snake_case__						= 24
  elif checkpoint == "large":
    snake_case__						= 2048
    snake_case__						= 48
    snake_case__						= 32
  else:
    raise ValueError(F"""Checkpoint should be one of `['small', 'medium', 'large']`, got {checkpoint}."""				)
  snake_case__						= MusicgenDecoderConfig(
      hidden_size=__lowerCAmelCase							,   ffn_dim=hidden_size * 4							,   num_hidden_layers=__lowerCAmelCase							,   num_attention_heads=__lowerCAmelCase							,   )
  return config







@torch.no_grad()
def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   __lowerCAmelCase="cpu"				)    ->				Union[str, Any]:
  snake_case__						= MusicGen.get_pretrained(__lowerCAmelCase							,   device=__lowerCAmelCase				)
  snake_case__						= decoder_config_from_checkpoint(__lowerCAmelCase				)

  snake_case__						= fairseq_model.lm.state_dict()
  snake_case__				,				snake_case__						= rename_state_dict(
      __lowerCAmelCase							,   hidden_size=decoder_config.hidden_size				)

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

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

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

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

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

  # init the composite model
  snake_case__						= MusicgenForConditionalGeneration(text_encoder=__lowerCAmelCase							,   audio_encoder=__lowerCAmelCase							,   decoder=__lowerCAmelCase				)

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

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

  with torch.no_grad():
    snake_case__						= model(input_ids=__lowerCAmelCase							,   decoder_input_ids=__lowerCAmelCase				).logits

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

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

  snake_case__						= MusicgenProcessor(feature_extractor=__lowerCAmelCase							,   tokenizer=__lowerCAmelCase				)

  # set the appropriate bos/pad token ids
  snake_case__						= 2048
  snake_case__						= 2048

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

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

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


if __name__ == "__main__":
   lowerCamelCase__							:						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__							:						Optional[int]							     =   parser.parse_args()
   convert_musicgen_checkpoint(args.checkpoint, args.pytorch_dump_folder, args.push_to_hub)





from __future__ import annotations

import unittest

from transformers import BlenderbotConfig, BlenderbotTokenizer, is_tf_available
from transformers.testing_utils import require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import TFAutoModelForSeqaSeqLM, TFBlenderbotForConditionalGeneration, TFBlenderbotModel
@require_tf
class 		__magic_name__						:





     '''simple docstring'''



     __lowercase				:							int	   =       BlenderbotConfig
     __lowercase				:							Any	   =       {}
     __lowercase				:							Optional[Any]	   =       'gelu'

     def __init__(					self:Tuple     ,		_a:Optional[Any]     ,		_a:Optional[Any]=13     ,		_a:Tuple=7     ,		_a:Union[str, Any]=True     ,		_a:int=False     ,		_a:int=99     ,		_a:Optional[int]=32     ,		_a:List[str]=2     ,		_a:List[str]=4     ,		_a:List[Any]=37     ,		_a:Any=0.1     ,		_a:int=0.1     ,		_a:List[Any]=20     ,		_a:List[str]=2     ,		_a:int=1     ,		_a:Dict=0     ,		):
       snake_case__						= parent
       snake_case__						= batch_size
       snake_case__						= seq_length
       snake_case__						= is_training
       snake_case__						= use_labels
       snake_case__						= vocab_size
       snake_case__						= hidden_size
       snake_case__						= num_hidden_layers
       snake_case__						= num_attention_heads
       snake_case__						= intermediate_size

       snake_case__						= hidden_dropout_prob
       snake_case__						= attention_probs_dropout_prob
       snake_case__						= max_position_embeddings
       snake_case__						= eos_token_id
       snake_case__						= pad_token_id
       snake_case__						= bos_token_id

     def 	SCREAMING_SNAKE_CASE__						(					self:int   ):
       snake_case__						= ids_tensor([self.batch_size, self.seq_length - 1]     ,		self.vocab_size   )
       snake_case__						= tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size   )     ,		1   )
       snake_case__						= tf.concat([input_ids, eos_tensor]     ,		axis=1   )

       snake_case__						= ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size   )

       snake_case__						= self.config_cls(
           vocab_size=self.vocab_size     ,		d_model=self.hidden_size     ,		encoder_layers=self.num_hidden_layers     ,		decoder_layers=self.num_hidden_layers     ,		encoder_attention_heads=self.num_attention_heads     ,		decoder_attention_heads=self.num_attention_heads     ,		encoder_ffn_dim=self.intermediate_size     ,		decoder_ffn_dim=self.intermediate_size     ,		dropout=self.hidden_dropout_prob     ,		attention_dropout=self.attention_probs_dropout_prob     ,		max_position_embeddings=self.max_position_embeddings     ,		eos_token_ids=[2]     ,		bos_token_id=self.bos_token_id     ,		pad_token_id=self.pad_token_id     ,		decoder_start_token_id=self.pad_token_id     ,		**self.config_updates     ,		)
       snake_case__						= prepare_blenderbot_inputs_dict(_a     ,		_a     ,		_a   )
       return config, inputs_dict







     def 	SCREAMING_SNAKE_CASE__						(					self:int     ,		_a:Optional[Any]     ,		_a:int   ):
       snake_case__						= TFBlenderbotModel(config=_a   ).get_decoder()
       snake_case__						= inputs_dict['''input_ids''']

       snake_case__						= input_ids[:1, :]
       snake_case__						= inputs_dict['''attention_mask'''][:1, :]
       snake_case__						= inputs_dict['''head_mask''']
       snake_case__						= 1

       # first forward pass
       snake_case__						= model(_a     ,		attention_mask=_a     ,		head_mask=_a     ,		use_cache=_a   )

       snake_case__				,				snake_case__						= outputs.to_tuple()

       # create hypothetical next token and extent to next_input_ids
       snake_case__						= ids_tensor((self.batch_size, 3)     ,		config.vocab_size   )
       snake_case__						= tf.cast(ids_tensor((self.batch_size, 3)     ,		2   )     ,		tf.inta   )

       # append to next input_ids and
       snake_case__						= tf.concat([input_ids, next_tokens]     ,		axis=-1   )
       snake_case__						= tf.concat([attention_mask, next_attn_mask]     ,		axis=-1   )

       snake_case__						= model(_a     ,		attention_mask=_a   )[0]
       snake_case__						= model(_a     ,		attention_mask=_a     ,		past_key_values=_a   )[0]

       self.parent.assertEqual(next_tokens.shape[1]     ,		output_from_past.shape[1]   )

       # select random slice
       snake_case__						= int(ids_tensor((1,)     ,		output_from_past.shape[-1]   )   )
       snake_case__						= output_from_no_past[:, -3:, random_slice_idx]
       snake_case__						= output_from_past[:, :, random_slice_idx]

       # test that outputs are equal for slice
       tf.debugging.assert_near(_a     ,		_a     ,		rtol=1e-3   )







def SCREAMING_SNAKE_CASE	(			__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   __lowerCAmelCase=None							,   )    ->				Tuple:
  if attention_mask is None:
    snake_case__						= tf.cast(tf.math.not_equal(__lowerCAmelCase							,   config.pad_token_id				)							,   tf.inta				)
  if decoder_attention_mask is None:
    snake_case__						= tf.concat(
        [
            tf.ones(decoder_input_ids[:, :1].shape							,   dtype=tf.inta				),
            tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]							,   config.pad_token_id				)							,   tf.inta				),
        ]							,   axis=-1							,   )
  if head_mask is None:
    snake_case__						= tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
  if decoder_head_mask is None:
    snake_case__						= tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
  if cross_attn_head_mask is None:
    snake_case__						= tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
  return {
      "input_ids": input_ids,
      "decoder_input_ids": decoder_input_ids,
      "attention_mask": attention_mask,
      "decoder_attention_mask": decoder_attention_mask,
      "head_mask": head_mask,
      "decoder_head_mask": decoder_head_mask,
      "cross_attn_head_mask": cross_attn_head_mask,
  }
@require_tf
class 		__magic_name__						(snake_case_   ,snake_case_   ,unittest.TestCase				):





     '''simple docstring'''



     __lowercase				:							List[str]	   =       (TFBlenderbotForConditionalGeneration, TFBlenderbotModel) if is_tf_available() else ()
     __lowercase				:							Any	   =       (TFBlenderbotForConditionalGeneration,) if is_tf_available() else ()
     __lowercase				:							Tuple	   =       (
         {
             'conversational': TFBlenderbotForConditionalGeneration,
             'feature-extraction': TFBlenderbotModel,
             'summarization': TFBlenderbotForConditionalGeneration,
             'text2text-generation': TFBlenderbotForConditionalGeneration,
             'translation': TFBlenderbotForConditionalGeneration,
         }
         if is_tf_available()
         else {}
     )
     __lowercase				:							Any	   =       True
     __lowercase				:							int	   =       False
     __lowercase				:							int	   =       False

     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]   ):
       snake_case__						= TFBlenderbotModelTester(self   )
       snake_case__						= ConfigTester(self     ,		config_class=_a   )

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       self.config_tester.run_common_tests()







     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[Any]   ):
       snake_case__						= self.model_tester.prepare_config_and_inputs_for_common()
       self.model_tester.check_decoder_model_past_large_inputs(*_a   )




@require_tokenizers
@require_tf
class 		__magic_name__						(unittest.TestCase				):





     '''simple docstring'''



     __lowercase				:							Optional[int]	   =       ['My friends are cool but they eat too many carbs.']
     __lowercase				:							Optional[int]	   =       'facebook/blenderbot-400M-distill'

     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       return BlenderbotTokenizer.from_pretrained(self.model_name   )

     @cached_property
     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]   ):
       snake_case__						= TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name   )
       return model







     @slow
     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]   ):
       snake_case__						= self.tokenizer(self.src_text     ,		return_tensors='''tf'''   )
       snake_case__						= self.model.generate(
           model_inputs.input_ids     ,		)
       snake_case__						= self.tokenizer.batch_decode(generated_ids.numpy()     ,		skip_special_tokens=_a   )[0]
       assert (
           generated_words
           == " That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?"
       )

